From 5ea8aa65d3ea3306c3c21b3eec5bc2dcbfffdaaf Mon Sep 17 00:00:00 2001
From: Tal Ben-Nun <tbennun@users.noreply.github.com>
Date: Fri, 18 Aug 2023 12:57:22 -0700
Subject: [PATCH] Dropout: make seed and states_num kernel arguments (#2277)

---
 src/kernels/MIOpenDropout.cl |  6 +++---
 src/ocl/dropoutocl.cpp       | 16 ++++++----------
 2 files changed, 9 insertions(+), 13 deletions(-)

diff --git a/src/kernels/MIOpenDropout.cl b/src/kernels/MIOpenDropout.cl
index 75e29fc4fb..596423418f 100644
--- a/src/kernels/MIOpenDropout.cl
+++ b/src/kernels/MIOpenDropout.cl
@@ -273,13 +273,13 @@ void xorwow_lite_init(prngStates* cur_state,
     cur_state->d += (uint)(offset)*362437;
 }
 
-__kernel void InitKernelState(__global prngStates* state)
+__kernel void InitKernelState(__global prngStates* state, ulong prng_seed, ulong states_num)
 {
-    for(uint gid = get_global_id(0); gid < STATES_NUM; gid += get_global_size(0))
+    for(uint gid = get_global_id(0); gid < states_num; gid += get_global_size(0))
     {
         prngStates state_gid;
         xorwow_lite_init(&state_gid,
-                         (unsigned long long)PRNG_SEED,
+                         (unsigned long long)prng_seed,
                          (unsigned long long)gid,
                          (unsigned long long)0);
 
diff --git a/src/ocl/dropoutocl.cpp b/src/ocl/dropoutocl.cpp
index c52f21da49..caf44d9569 100644
--- a/src/ocl/dropoutocl.cpp
+++ b/src/ocl/dropoutocl.cpp
@@ -132,18 +132,16 @@ void DropoutDescriptor::InitPRNGState(Handle& handle,
         MIOPEN_THROW("PRNG state size should not exceed system maximum memory allocation size.");
     }
 
-    size_t states_num = prng_stateSizeInBytes / sizeof(prngStates);
-    size_t wk_grp_num = std::min(size_t(MAX_PRNG_STATE / 256), (states_num + 255) / 256);
+    unsigned long long states_num = prng_stateSizeInBytes / sizeof(prngStates);
+    size_t wk_grp_num             = std::min(MAX_PRNG_STATE / 256ULL, (states_num + 255) / 256);
 
-    std::string network_config = "initprngs-" + std::to_string(states_num) + "x" +
-                                 std::to_string(sizeof(prngStates)) + "x" +
-                                 std::to_string(rng_mode) + "x" + std::to_string(prng_seed) + "x" +
-                                 std::to_string(wk_grp_num);
+    std::string network_config = "initprngs-" + std::to_string(sizeof(prngStates)) + "x" +
+                                 std::to_string(rng_mode) + "x" + std::to_string(wk_grp_num);
 
     auto&& kernels = handle.GetKernels(kernel_name, network_config);
     if(!kernels.empty())
     {
-        kernels.front()(prng_states);
+        kernels.front()(prng_states, prng_seed, states_num);
     }
     else
     {
@@ -152,14 +150,12 @@ void DropoutDescriptor::InitPRNGState(Handle& handle,
 
         std::string params;
         params += " -DRUN_INIT_PRNG=1";
-        params += " -DPRNG_SEED=" + std::to_string(prng_seed);
-        params += " -DSTATES_NUM=" + std::to_string(states_num);
 #if DROPOUT_DEBUG
         std::cout << "Threads allocated for PRNG states: " << vgd[0] << std::endl;
         std::cout << "Memory allocated for PRNG states: " << stateSizeInBytes << std::endl;
 #endif
         handle.AddKernel(kernel_name, network_config, program_name, kernel_name, vld, vgd, params)(
-            prng_states);
+            prng_states, prng_seed, states_num);
 #if DROPOUT_DEBUG
         std::cout << "Succeeded in launching InitPRNGState()." << stateSizeInBytes << std::endl;
 #endif