include/ck/utility/random_gen.hpp Source File

include/ck/utility/random_gen.hpp Source File#

Composable Kernel: include/ck/utility/random_gen.hpp Source File
random_gen.hpp
Go to the documentation of this file.
1 // SPDX-License-Identifier: MIT
2 // Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved.
3 
4 #pragma once
5 #include <ck/utility/ignore.hpp>
6 #include "ck/ck.hpp"
7 
8 #ifdef CK_CODE_GEN_RTC
9 using uint8_t = unsigned char;
10 using uint16_t = unsigned short;
11 using uint32_t = unsigned int;
12 #endif
13 namespace ck {
14 
15 // Pseudo random number generator
16 // version for fp32
17 template <typename T, uint32_t seed_t, ck::enable_if_t<std::is_same<float, T>{}, bool> = false>
18 __host__ __device__ uint32_t prand_generator(index_t id, T val, uint32_t seed = seed_t)
19 {
20  uint32_t x = *(reinterpret_cast<uint32_t*>(&val));
21  uint32_t drop_bits = uint32_t(x) & 0xFFFFu;
22  drop_bits ^= x >> 16;
23  drop_bits = ((drop_bits & 31) << 11) | (drop_bits >> 5);
24  drop_bits *= 0x7000149;
25  // NOTE: If id is in 64 bit, we are only using lower 32 bit.
26  // So, it can have an effect of using same id for multiple elements when the id is very
27  // large!
28  uint32_t rng = (drop_bits ^ 0x13371337 ^ (id * 229791) ^ seed);
29  return rng;
30 }
31 
32 // version for fp16
33 template <typename T, uint32_t seed_t, ck::enable_if_t<std::is_same<_Float16, T>{}, bool> = false>
34 __host__ __device__ uint32_t prand_generator(index_t id, T val, uint32_t seed = seed_t)
35 {
36  uint16_t x = *(reinterpret_cast<uint16_t*>(&val));
37  uint32_t drop_bits = uint32_t(x) & 0xFFFFu;
38  drop_bits = ((drop_bits & 31) << 11) | (drop_bits >> 5);
39  drop_bits *= 0x7000149;
40  // NOTE: If id is in 64 bit, we are only using lower 32 bit.
41  // So, it can have an effect of using same id for multiple elements when the id is very
42  // large!
43  uint32_t rng = (drop_bits ^ 0x13371337 ^ (id * 229791) ^ seed);
44  return rng;
45 }
46 
47 // return 0 if data is not fp16 or fp32
48 template <typename T,
49  uint32_t seed_t,
50  ck::enable_if_t<!(std::is_same<float, T>{} || std::is_same<_Float16, T>{}), bool> = false>
51 __host__ __device__ uint32_t prand_generator(int id, T val, uint32_t seed = seed_t)
52 {
53  ck::ignore = id;
54  ck::ignore = val;
55  ck::ignore = seed;
56 
57  return 0;
58 }
59 
60 } // namespace ck
ck
Definition: ck.hpp:264
ck::ignore
constexpr detail::ignore_t ignore
Definition: ignore.hpp:20
ck::prand_generator
__host__ __device__ uint32_t prand_generator(index_t id, T val, uint32_t seed=seed_t)
Definition: random_gen.hpp:18
ck::index_t
int32_t index_t
Definition: ck.hpp:289
ck::enable_if_t
typename std::enable_if< B, T >::type enable_if_t
Definition: enable_if.hpp:13