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

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

Composable Kernel: /home/docs/checkouts/readthedocs.org/user_builds/advanced-micro-devices-composable-kernel/checkouts/develop/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/utility/type.hpp>
7 #include "ck/ck.hpp"
8 
9 #ifdef CK_CODE_GEN_RTC
10 using uint8_t = unsigned char;
11 using uint16_t = unsigned short;
12 using uint32_t = unsigned int;
13 #endif
14 namespace ck {
15 
16 // Pseudo random number generator
17 // version for fp32
18 template <typename T, uint32_t seed_t, ck::enable_if_t<is_same<float, T>{}, bool> = false>
19 __host__ __device__ uint32_t prand_generator(index_t id, T val, uint32_t seed = seed_t)
20 {
21  uint32_t x = bit_cast<uint32_t>(val);
22  uint32_t drop_bits = uint32_t(x) & 0xFFFFu;
23  drop_bits ^= x >> 16;
24  drop_bits = ((drop_bits & 31) << 11) | (drop_bits >> 5);
25  drop_bits *= 0x7000149;
26  // NOTE: If id is in 64 bit, we are only using lower 32 bit.
27  // So, it can have an effect of using same id for multiple elements when the id is very
28  // large!
29  uint32_t rng = (drop_bits ^ 0x13371337 ^ (id * 229791) ^ seed);
30  return rng;
31 }
32 
33 // version for fp16
34 template <typename T, uint32_t seed_t, ck::enable_if_t<is_same<_Float16, T>{}, bool> = false>
35 __host__ __device__ uint32_t prand_generator(index_t id, T val, uint32_t seed = seed_t)
36 {
37  uint16_t x = bit_cast<uint16_t>(val);
38  uint32_t drop_bits = uint32_t(x) & 0xFFFFu;
39  drop_bits = ((drop_bits & 31) << 11) | (drop_bits >> 5);
40  drop_bits *= 0x7000149;
41  // NOTE: If id is in 64 bit, we are only using lower 32 bit.
42  // So, it can have an effect of using same id for multiple elements when the id is very
43  // large!
44  uint32_t rng = (drop_bits ^ 0x13371337 ^ (id * 229791) ^ seed);
45  return rng;
46 }
47 
48 // return 0 if data is not fp16 or fp32
49 template <typename T,
50  uint32_t seed_t,
51  ck::enable_if_t<!(is_same<float, T>{} || is_same<_Float16, T>{}), bool> = false>
52 __host__ __device__ uint32_t prand_generator(int id, T val, uint32_t seed = seed_t)
53 {
54  ck::ignore = id;
55  ck::ignore = val;
56  ck::ignore = seed;
57 
58  return 0;
59 }
60 
61 } // namespace ck
ck
Definition: ck.hpp:268
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:19
ck::index_t
int32_t index_t
Definition: ck.hpp:299
ck::enable_if_t
typename std::enable_if< B, T >::type enable_if_t
Definition: enable_if.hpp:27
uint16_t
unsigned short uint16_t
Definition: stdint.h:125
uint32_t
unsigned int uint32_t
Definition: stdint.h:126
uint8_t
unsigned char uint8_t
Definition: stdint.h:124