/home/docs/checkouts/readthedocs.org/user_builds/advanced-micro-devices-composable-kernel/checkouts/develop/include/ck/library/utility/host_tensor_generator.hpp Source File

/home/docs/checkouts/readthedocs.org/user_builds/advanced-micro-devices-composable-kernel/checkouts/develop/include/ck/library/utility/host_tensor_generator.hpp Source File#

Composable Kernel: /home/docs/checkouts/readthedocs.org/user_builds/advanced-micro-devices-composable-kernel/checkouts/develop/include/ck/library/utility/host_tensor_generator.hpp Source File
Go to the documentation of this file.
 // SPDX-License-Identifier: MIT
 // Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved.
  
 #pragma once
  
 #include <cmath>
 #include <numeric>
 #include <random>
  
 #include "ck/ck.hpp"
  
 template <typename T>
 struct GeneratorTensor_0
 {
     template <typename... Is>
     T operator()(Is...)
     {
         return T{0};
     }
 };
  
 template <typename T>
 struct GeneratorTensor_1
 {
     T value = 1;
  
     template <typename... Is>
     T operator()(Is...)
     {
         return value;
     }
 };
  
 template <>
 struct GeneratorTensor_1<ck::half_t>
 {
     float value = 1.0;
  
     template <typename... Is>
     ck::half_t operator()(Is...)
     {
         return ck::type_convert<ck::half_t>(value);
     }
 };
  
 template <>
 struct GeneratorTensor_1<ck::bhalf_t>
 {
     float value = 1.0;
  
     template <typename... Is>
     ck::bhalf_t operator()(Is...)
     {
         return ck::type_convert<ck::bhalf_t>(value);
     }
 };
  
 #if defined CK_ENABLE_FP8
 template <>
 struct GeneratorTensor_1<ck::f8_t>
 {
     float value = 1.0;
  
     template <typename... Is>
     ck::f8_t operator()(Is...)
     {
         return ck::type_convert<ck::f8_t>(value);
     }
 };
  
 template <>
 struct GeneratorTensor_1<ck::bf8_t>
 {
     float value = 1.0;
  
     template <typename... Is>
     ck::bf8_t operator()(Is...)
     {
         return ck::type_convert<ck::bf8_t>(value);
     }
 };
 #endif
  
 template <>
 struct GeneratorTensor_1<ck::f4_t>
 {
     float value = 1.0;
  
     template <typename... Is>
     ck::f4_t operator()(Is...)
     {
         return ck::type_convert<ck::f4_t>(value);
     }
 };
  
 template <>
 struct GeneratorTensor_1<ck::f4x2_pk_t>
 {
     float value = 1.0;
  
     template <typename... Is>
     ck::f4x2_pk_t operator()(Is...)
     {
         return ck::f4x2_pk_t{ck::type_convert<ck::f4x2_t>(ck::float2_t{value, value})};
     }
 };
  
 template <>
 struct GeneratorTensor_1<ck::f6x32_pk_t>
 {
     float value = 1.0;
  
     template <typename... Is>
     ck::f6x32_pk_t operator()(Is...)
     {
         ck::f6x32_pk_t r;
         ck::static_for<0, 32, 1>{}([&](auto i) {
             r.pack(ck::type_convert<ck::f6_t>(value), static_cast<ck::index_t>(i));
         });
         return r;
     }
 };
  
 template <>
 struct GeneratorTensor_1<ck::bf6x32_pk_t>
 {
     float value = 1.0;
  
     template <typename... Is>
     ck::bf6x32_pk_t operator()(Is...)
     {
         ck::bf6x32_pk_t r;
         ck::static_for<0, 32, 1>{}([&](auto i) {
             r.pack(ck::type_convert<ck::bf6_t>(value), static_cast<ck::index_t>(i));
         });
         return r;
     }
 };
  
 template <>
 struct GeneratorTensor_1<int8_t>
 {
     int8_t value = 1;
  
     template <typename... Is>
     int8_t operator()(Is...)
     {
         return value;
     }
 };
  
 template <>
 struct GeneratorTensor_1<ck::pk_i4_t>
 {
     int8_t value = 1;
  
     template <typename... Is>
     ck::pk_i4_t operator()(Is...)
     {
         int t         = value + 8;
         ck::pk_i4_t r = ((t << 4) + t) & 0xff;
         return r;
     }
 };
  
 template <>
 struct GeneratorTensor_1<ck::e8m0_bexp_t>
 {
     float value = 1;
  
     template <typename... Is>
     ck::e8m0_bexp_t operator()(Is...)
     {
         return ck::type_convert<ck::e8m0_bexp_t>(value);
     }
 };
  
 template <typename T>
 struct GeneratorTensor_2
 {
     int min_value = 0;
     int max_value = 1;
  
     template <typename... Is>
     T operator()(Is...)
     {
         return static_cast<T>((std::rand() % (max_value - min_value)) + min_value);
     }
 };
  
 template <>
 struct GeneratorTensor_2<ck::f6x32_pk_t>
 {
     int min_value = 0;
     int max_value = 1;
  
     template <typename... Is>
     ck::f6x32_pk_t operator()(Is...)
     {
         ck::f6x32_pk_t r;
         ck::static_for<0, 32, 1>{}([&](auto i) {
             float tmp = (std::rand() % (max_value - min_value)) + min_value;
             r.pack(ck::type_convert<ck::f6_t>(tmp), static_cast<ck::index_t>(i));
         });
  
         return r;
     }
 };
  
 template <>
 struct GeneratorTensor_2<ck::bf6x32_pk_t>
 {
     int min_value = 0;
     int max_value = 1;
  
     template <typename... Is>
     ck::bf6x32_pk_t operator()(Is...)
     {
         ck::bf6x32_pk_t r;
         ck::static_for<0, 32, 1>{}([&](auto i) {
             float tmp = (std::rand() % (max_value - min_value)) + min_value;
             r.pack(ck::type_convert<ck::bf6_t>(tmp), static_cast<ck::index_t>(i));
         });
  
         return r;
     }
 };
  
 template <>
 struct GeneratorTensor_2<ck::bhalf_t>
 {
     int min_value = 0;
     int max_value = 1;
  
     template <typename... Is>
     ck::bhalf_t operator()(Is...)
     {
         float tmp = (std::rand() % (max_value - min_value)) + min_value;
         return ck::type_convert<ck::bhalf_t>(tmp);
     }
 };
  
 template <>
 struct GeneratorTensor_2<int8_t>
 {
     int min_value = 0;
     int max_value = 1;
  
     template <typename... Is>
     int8_t operator()(Is...)
     {
         return (std::rand() % (max_value - min_value)) + min_value;
     }
 };
  
 template <>
 struct GeneratorTensor_2<ck::pk_i4_t>
 {
     int min_value = 0;
     int max_value = 1;
  
     template <typename... Is>
     ck::pk_i4_t operator()(Is...)
     {
         int hi        = std::rand() % (max_value - min_value) + min_value + 8;
         int lo        = std::rand() % (max_value - min_value) + min_value + 8;
         ck::pk_i4_t r = ((hi << 4) + lo) & 0xff;
         return r;
     }
 };
  
 #if defined CK_ENABLE_FP8
 template <>
 struct GeneratorTensor_2<ck::f8_t>
 {
     int min_value = 0;
     int max_value = 1;
  
     template <typename... Is>
     ck::f8_t operator()(Is...)
     {
         float tmp = (std::rand() % (max_value - min_value)) + min_value;
         return ck::type_convert<ck::f8_t>(tmp);
     }
 };
 #endif
  
 #if defined CK_ENABLE_BF8
 template <>
 struct GeneratorTensor_2<ck::bf8_t>
 {
     int min_value = 0;
     int max_value = 1;
  
     template <typename... Is>
     ck::bf8_t operator()(Is...)
     {
         float tmp = (std::rand() % (max_value - min_value)) + min_value;
         return ck::type_convert<ck::bf8_t>(tmp);
     }
 };
 #endif
  
 template <>
 struct GeneratorTensor_2<ck::f4_t>
 {
     int min_value = 0;
     int max_value = 1;
  
     template <typename... Is>
     ck::f4_t operator()(Is...)
     {
         float tmp = (std::rand() % (max_value - min_value)) + min_value;
         return ck::type_convert<ck::f4_t>(tmp);
     }
 };
  
 template <>
 struct GeneratorTensor_2<ck::f4x2_pk_t>
 {
     int min_value = 0;
     int max_value = 1;
  
     template <typename... Is>
     ck::f4x2_pk_t operator()(Is...)
     {
         float tmp0 = (std::rand() % (max_value - min_value)) + min_value;
         float tmp1 = (std::rand() % (max_value - min_value)) + min_value;
         return ck::f4x2_pk_t{ck::type_convert<ck::f4x2_t>(ck::float2_t{tmp0, tmp1})};
     }
 };
  
 template <typename T>
 struct GeneratorTensor_3
 {
     float min_value = 0;
     float max_value = 1;
  
     template <typename... Is>
     T operator()(Is...)
     {
         float tmp = float(std::rand()) / float(RAND_MAX);
  
         return static_cast<T>(min_value + tmp * (max_value - min_value));
     }
 };
  
 template <>
 struct GeneratorTensor_3<ck::bhalf_t>
 {
     float min_value = 0;
     float max_value = 1;
  
     template <typename... Is>
     ck::bhalf_t operator()(Is...)
     {
         float tmp = float(std::rand()) / float(RAND_MAX);
  
         float fp32_tmp = min_value + tmp * (max_value - min_value);
  
         return ck::type_convert<ck::bhalf_t>(fp32_tmp);
     }
 };
  
 #if defined CK_ENABLE_FP8
 template <>
 struct GeneratorTensor_3<ck::f8_t>
 {
     float min_value = 0;
     float max_value = 1;
  
     template <typename... Is>
     ck::f8_t operator()(Is...)
     {
         float tmp = float(std::rand()) / float(RAND_MAX);
  
         float fp32_tmp = min_value + tmp * (max_value - min_value);
  
         return ck::type_convert<ck::f8_t>(fp32_tmp);
     }
 };
 #endif
  
 #if defined CK_ENABLE_BF8
 template <>
 struct GeneratorTensor_3<ck::bf8_t>
 {
     float min_value = 0;
     float max_value = 1;
  
     template <typename... Is>
     ck::bf8_t operator()(Is...)
     {
         float tmp = float(std::rand()) / float(RAND_MAX);
  
         float fp32_tmp = min_value + tmp * (max_value - min_value);
  
         return ck::type_convert<ck::bf8_t>(fp32_tmp);
     }
 };
 #endif
  
 template <>
 struct GeneratorTensor_3<ck::f4_t>
 {
     float min_value = 0;
     float max_value = 1;
  
     template <typename... Is>
     ck::f4_t operator()(Is...)
     {
         float tmp = float(std::rand()) / float(RAND_MAX);
  
         float fp32_tmp = min_value + tmp * (max_value - min_value);
  
         return ck::type_convert<ck::f4_t>(fp32_tmp);
     }
 };
  
 template <>
 struct GeneratorTensor_3<ck::f4x2_pk_t>
 {
     float min_value = 0;
     float max_value = 1;
  
     template <typename... Is>
     ck::f4x2_pk_t operator()(Is...)
     {
         float tmp0 = float(std::rand()) / float(RAND_MAX);
         float tmp1 = float(std::rand()) / float(RAND_MAX);
  
         float fp32_tmp0 = min_value + tmp0 * (max_value - min_value);
         float fp32_tmp1 = min_value + tmp1 * (max_value - min_value);
  
         return ck::f4x2_pk_t{ck::type_convert<ck::f4x2_t>(ck::float2_t{fp32_tmp0, fp32_tmp1})};
     }
 };
  
 template <>
 struct GeneratorTensor_3<ck::f6x32_pk_t>
 {
     float min_value = 0;
     float max_value = 1;
  
     template <typename... Is>
     ck::f6x32_pk_t operator()(Is...)
     {
         ck::f6x32_pk_t r;
         ck::static_for<0, 32, 1>{}([&](auto i) {
             float rnd  = float(std::rand()) / float(RAND_MAX);
             float fp32 = min_value + rnd * (max_value - min_value);
             r.pack(ck::type_convert<ck::f6_t>(fp32), static_cast<ck::index_t>(i));
         });
  
         return r;
     }
 };
  
 template <>
 struct GeneratorTensor_3<ck::bf6x32_pk_t>
 {
     float min_value = 0;
     float max_value = 1;
  
     template <typename... Is>
     ck::bf6x32_pk_t operator()(Is...)
     {
         ck::bf6x32_pk_t r;
         ck::static_for<0, 32, 1>{}([&](auto i) {
             float rnd  = float(std::rand()) / float(RAND_MAX);
             float fp32 = min_value + rnd * (max_value - min_value);
             r.pack(ck::type_convert<ck::bf6_t>(fp32), static_cast<ck::index_t>(i));
         });
  
         return r;
     }
 };
  
 template <typename T>
 struct GeneratorTensor_4
 {
     std::mt19937 generator;
     std::normal_distribution<float> distribution;
  
     GeneratorTensor_4(float mean, float stddev, unsigned int seed = 1)
         : generator(seed), distribution(mean, stddev){};
  
     template <typename... Is>
     T operator()(Is...)
     {
         float tmp = distribution(generator);
  
         return ck::type_convert<T>(tmp);
     }
 };
  
 template <>
 struct GeneratorTensor_4<ck::f4x2_pk_t>
 {
     std::mt19937 generator;
     std::normal_distribution<float> distribution;
  
     GeneratorTensor_4(float mean, float stddev, unsigned int seed = 1)
         : generator(seed), distribution(mean, stddev){};
  
     template <typename... Is>
     ck::f4x2_pk_t operator()(Is...)
     {
         float fp32_tmp0 = distribution(generator);
         float fp32_tmp1 = distribution(generator);
  
         return ck::f4x2_pk_t{ck::type_convert<ck::f4x2_t>(ck::float2_t{fp32_tmp0, fp32_tmp1})};
     }
 };
  
 template <>
 struct GeneratorTensor_4<ck::f6x32_pk_t>
 {
     std::mt19937 generator;
     std::normal_distribution<float> distribution;
  
     GeneratorTensor_4(float mean, float stddev, unsigned int seed = 1)
         : generator(seed), distribution(mean, stddev){};
  
     template <typename... Is>
     ck::f6x32_pk_t operator()(Is...)
     {
         ck::f6x32_pk_t r;
         ck::static_for<0, 32, 1>{}([&](auto i) {
             r.pack(ck::type_convert<ck::f6_t>(distribution(generator)),
                    static_cast<ck::index_t>(i));
         });
  
         return r;
     }
 };
  
 template <>
 struct GeneratorTensor_4<ck::bf6x32_pk_t>
 {
     std::mt19937 generator;
     std::normal_distribution<float> distribution;
  
     GeneratorTensor_4(float mean, float stddev, unsigned int seed = 1)
         : generator(seed), distribution(mean, stddev){};
  
     template <typename... Is>
     ck::bf6x32_pk_t operator()(Is...)
     {
         ck::bf6x32_pk_t r;
         ck::static_for<0, 32, 1>{}([&](auto i) {
             r.pack(ck::type_convert<ck::bf6_t>(distribution(generator)),
                    static_cast<ck::index_t>(i));
         });
  
         return r;
     }
 };
  
 struct GeneratorTensor_Checkboard
 {
     template <typename... Ts>
     float operator()(Ts... Xs) const
     {
         std::array<ck::index_t, sizeof...(Ts)> dims = {static_cast<ck::index_t>(Xs)...};
         return std::accumulate(dims.begin(),
                                dims.end(),
                                true,
                                [](bool init, ck::index_t x) -> int { return init != (x % 2); })
                    ? 1
                    : -1;
     }
 };
  
 template <typename T, ck::index_t Dim>
 struct GeneratorTensor_Sequential
 {
     template <typename... Ts>
     T operator()(Ts... Xs) const
     {
         std::array<ck::index_t, sizeof...(Ts)> dims = {{static_cast<ck::index_t>(Xs)...}};
  
         float tmp = dims[Dim];
         return ck::type_convert<T>(tmp);
     }
 };
  
 template <ck::index_t Dim>
 struct GeneratorTensor_Sequential<ck::f4x2_pk_t, Dim>
 {
     template <typename... Ts>
     ck::f4x2_pk_t operator()(Ts... Xs) const
     {
         std::array<ck::index_t, sizeof...(Ts)> dims = {{static_cast<ck::index_t>(Xs)...}};
  
         float tmp = dims[Dim];
         return ck::type_convert<ck::f4x2_t>(ck::float2_t(tmp));
     }
 };
  
 template <ck::index_t Dim>
 struct GeneratorTensor_Sequential<ck::f6x32_pk_t, Dim>
 {
     template <typename... Ts>
     ck::f6x32_pk_t operator()(Ts... Xs) const
     {
         std::array<ck::index_t, sizeof...(Ts)> dims = {{static_cast<ck::index_t>(Xs)...}};
  
         float tmp = dims[Dim];
  
         ck::f6x32_pk_t r;
         ck::static_for<0, 32, 1>{}(
             [&](auto i) { r.pack(ck::type_convert<ck::f6_t>(tmp), static_cast<ck::index_t>(i)); });
         return r;
     }
 };
  
 template <ck::index_t Dim>
 struct GeneratorTensor_Sequential<ck::bf6x32_pk_t, Dim>
 {
     template <typename... Ts>
     ck::bf6x32_pk_t operator()(Ts... Xs) const
     {
         std::array<ck::index_t, sizeof...(Ts)> dims = {{static_cast<ck::index_t>(Xs)...}};
  
         float tmp = dims[Dim];
  
         ck::bf6x32_pk_t r;
         ck::static_for<0, 32, 1>{}(
             [&](auto i) { r.pack(ck::type_convert<ck::bf6_t>(tmp), static_cast<ck::index_t>(i)); });
         return r;
     }
 };
  
 template <typename T, size_t NumEffectiveDim = 2>
 struct GeneratorTensor_Diagonal
 {
     T value{1};
  
     template <typename... Ts>
     T operator()(Ts... Xs) const
     {
         std::array<ck::index_t, sizeof...(Ts)> dims = {{static_cast<ck::index_t>(Xs)...}};
         size_t start_dim                            = dims.size() - NumEffectiveDim;
         bool pred                                   = true;
         for(size_t i = start_dim + 1; i < dims.size(); i++)
         {
             pred &= (dims[start_dim] == dims[i]);
         }
         return pred ? value : T{0};
     }
 };