From 8688dea8e00149847dff2e839635a7e85a632bac Mon Sep 17 00:00:00 2001
From: hanhuifeng2020 <hanhuifeng1@huawei.com>
Date: Sat, 2 Apr 2022 16:59:41 +0800
Subject: [PATCH] [DynamicShape][GPU]fix a bug that the dropout op continued to
 create new Generators under dynamic shape

---
 .../gpu/kernel/nn/dropout3d_gpu_kernel.h      | 24 ++++++-------
 .../device/gpu/kernel/nn/dropout_gpu_kernel.h | 35 +++++++++----------
 2 files changed, 27 insertions(+), 32 deletions(-)

diff --git a/mindspore/ccsrc/plugin/device/gpu/kernel/nn/dropout3d_gpu_kernel.h b/mindspore/ccsrc/plugin/device/gpu/kernel/nn/dropout3d_gpu_kernel.h
index 8f4aa7df9f..9357dc2e70 100644
--- a/mindspore/ccsrc/plugin/device/gpu/kernel/nn/dropout3d_gpu_kernel.h
+++ b/mindspore/ccsrc/plugin/device/gpu/kernel/nn/dropout3d_gpu_kernel.h
@@ -43,15 +43,6 @@ class Dropout3DFwdGpuKernelMod : public NativeGpuKernelMod {
     bool *mask_addr = GetDeviceAddress<bool>(outputs, 1);
     float *rand_f = GetDeviceAddress<float>(workspace, 0);
 
-    if (!states_init_) {
-      CHECK_CURAND_RET_WITH_EXCEPT(curandCreateGenerator(&curand_generator_, CURAND_RNG_PSEUDO_DEFAULT),
-                                   "Failed to create generator");
-      CHECK_CURAND_RET_WITH_EXCEPT(curandSetPseudoRandomGeneratorSeed(curand_generator_, time(NULL)),
-                                   "Failed to SetPseudoRandomGeneratorSeed");
-      MS_EXCEPTION_IF_NULL(curand_generator_);
-      states_init_ = true;
-    }
-
     CHECK_CURAND_RET_WITH_EXCEPT(curandSetStream(curand_generator_, reinterpret_cast<cudaStream_t>(stream_ptr)),
                                  "Failed to set stream for generator");
     // curandGen only supports float or double.
@@ -109,6 +100,15 @@ class Dropout3DFwdGpuKernelMod : public NativeGpuKernelMod {
                         << "but got " << keep_prob_;
     }
 
+    if (!states_init_) {
+      CHECK_CURAND_RET_WITH_EXCEPT(curandCreateGenerator(&curand_generator_, CURAND_RNG_PSEUDO_DEFAULT),
+                                   "Failed to create generator");
+      CHECK_CURAND_RET_WITH_EXCEPT(curandSetPseudoRandomGeneratorSeed(curand_generator_, time(NULL)),
+                                   "Failed to SetPseudoRandomGeneratorSeed");
+      MS_EXCEPTION_IF_NULL(curand_generator_);
+      states_init_ = true;
+    }
+
     InitSizeLists();
     return true;
   }
@@ -119,8 +119,6 @@ class Dropout3DFwdGpuKernelMod : public NativeGpuKernelMod {
     kernel_name_ = "Dropout3D";
     num_count_ = 0;
     keep_prob_ = 0.0;
-    states_init_ = false;
-    curand_generator_ = nullptr;
     n_ = 0;
     c_ = 0;
     num_chan_ = 0;
@@ -144,10 +142,10 @@ class Dropout3DFwdGpuKernelMod : public NativeGpuKernelMod {
 
  private:
   cudnnHandle_t cudnn_handle_;
-  curandGenerator_t curand_generator_;
+  curandGenerator_t curand_generator_{nullptr};
   bool is_null_input_;
   std::string kernel_name_;
-  bool states_init_;
+  bool states_init_{false};
   size_t num_count_;
   size_t n_;
   size_t c_;
diff --git a/mindspore/ccsrc/plugin/device/gpu/kernel/nn/dropout_gpu_kernel.h b/mindspore/ccsrc/plugin/device/gpu/kernel/nn/dropout_gpu_kernel.h
index ebd8650b3c..88c60db628 100644
--- a/mindspore/ccsrc/plugin/device/gpu/kernel/nn/dropout_gpu_kernel.h
+++ b/mindspore/ccsrc/plugin/device/gpu/kernel/nn/dropout_gpu_kernel.h
@@ -43,14 +43,6 @@ class DropoutFwdGpuKernelMod : public NativeGpuKernelMod {
     T *mask = GetDeviceAddress<T>(outputs, 1);
     float *mask_f = GetDeviceAddress<float>(workspace, 0);
 
-    if (!states_init_) {
-      CHECK_CURAND_RET_WITH_EXCEPT(curandCreateGenerator(&mask_generator_, CURAND_RNG_PSEUDO_DEFAULT),
-                                   "Failed to create generator");
-      CHECK_CURAND_RET_WITH_EXCEPT(curandSetPseudoRandomGeneratorSeed(mask_generator_, seed_),
-                                   "Failed to SetPseudoRandomGeneratorSeed");
-      MS_EXCEPTION_IF_NULL(mask_generator_);
-      states_init_ = true;
-    }
     CHECK_CURAND_RET_WITH_EXCEPT(curandSetStream(mask_generator_, reinterpret_cast<cudaStream_t>(stream_ptr)),
                                  "Failed to set stream for generator");
     // curandGen only support float or double for mask.
@@ -83,14 +75,22 @@ class DropoutFwdGpuKernelMod : public NativeGpuKernelMod {
       num_count_ *= x;
     }
     keep_prob_ = GetAttr<float>(kernel_node, "keep_prob");
-    int64_t seed = GetAttr<int64_t>(kernel_node, "Seed0");
-    if (seed == 0) {
-      seed = GetAttr<int64_t>(kernel_node, "Seed1");
+    if (!states_init_) {
+      int64_t seed = GetAttr<int64_t>(kernel_node, "Seed0");
       if (seed == 0) {
-        seed = time(NULL);
+        seed = GetAttr<int64_t>(kernel_node, "Seed1");
+        if (seed == 0) {
+          seed = time(NULL);
+        }
       }
+      seed_ = static_cast<uint64_t>(seed);
+      CHECK_CURAND_RET_WITH_EXCEPT(curandCreateGenerator(&mask_generator_, CURAND_RNG_PSEUDO_DEFAULT),
+                                   "Failed to create generator");
+      CHECK_CURAND_RET_WITH_EXCEPT(curandSetPseudoRandomGeneratorSeed(mask_generator_, seed_),
+                                   "Failed to SetPseudoRandomGeneratorSeed");
+      MS_EXCEPTION_IF_NULL(mask_generator_);
+      states_init_ = true;
     }
-    seed_ = static_cast<uint64_t>(seed);
 
     InitSizeLists();
     return true;
@@ -102,9 +102,6 @@ class DropoutFwdGpuKernelMod : public NativeGpuKernelMod {
     kernel_name_ = "Dropout";
     num_count_ = 0;
     keep_prob_ = 0.0;
-    seed_ = 0;
-    states_init_ = false;
-    mask_generator_ = nullptr;
     input_size_list_.clear();
     output_size_list_.clear();
     workspace_size_list_.clear();
@@ -127,9 +124,9 @@ class DropoutFwdGpuKernelMod : public NativeGpuKernelMod {
   std::string kernel_name_;
   size_t num_count_;
   float keep_prob_;
-  bool states_init_;
-  uint64_t seed_;
-  curandGenerator_t mask_generator_;
+  bool states_init_{false};
+  uint64_t seed_{0};
+  curandGenerator_t mask_generator_{nullptr};
 };
 }  // namespace kernel
 }  // namespace mindspore