Free curand states before the thread is destroyed

Author: no1d

include/ctranslate2/devices.h

@@ -26,6 +26,8 @@ namespace ctranslate2 {
   void synchronize_device(Device device, int index);
   void synchronize_stream(Device device);
 
+  void destroy_context(Device device);
+
   class ScopedDeviceSetter {
   public:
     ScopedDeviceSetter(Device device, int index)

include/ctranslate2/replica_pool.h

@@ -354,6 +354,8 @@ namespace ctranslate2 {
 
     void finalize() override {
       _replica.reset();
+
+      destroy_context(_device);
     }
 
   private:

src/cuda/random.cu

@@ -48,12 +48,17 @@ namespace ctranslate2 {
       curandState* _states;
     };
 
+    static thread_local std::unique_ptr<ScopedCurandStates<curandStatePhilox4_32_10_t>> states;
+
     curandStatePhilox4_32_10_t* get_curand_states(size_t num_states) {
-      static thread_local std::unique_ptr<ScopedCurandStates<curandStatePhilox4_32_10_t>> states;
       if (!states || num_states > states->num_states())
         states = std::make_unique<ScopedCurandStates<curandStatePhilox4_32_10_t>>(num_states);
       return states->states();
     }
 
+    void free_curand_states() {
+      states.reset();
+    }
+
   }
 }

src/cuda/random.h

@@ -6,6 +6,7 @@ namespace ctranslate2 {
   namespace cuda {
 
     curandStatePhilox4_32_10_t* get_curand_states(size_t num_states);
+    void free_curand_states();
 
   }
 }

src/devices.cc

@@ -2,6 +2,7 @@
 
 #ifdef CT2_WITH_CUDA
 #  include "cuda/utils.h"
+#  include "cuda/random.h"
 #endif
 #ifdef CT2_WITH_TENSOR_PARALLEL
 #  include <unistd.h>
@@ -118,6 +119,17 @@ namespace ctranslate2 {
     (void)device;
 #endif
   }
+
+  void destroy_context(Device device) {
+#ifdef CT2_WITH_CUDA
+      if (device == Device::CUDA) {
+          cuda::free_curand_states();
+      }
+#else
+      (void)device;
+#endif
+  }
+
   // Initialize the static member variable
 #ifdef CT2_WITH_TENSOR_PARALLEL
     std::vector<ncclComm_t*> ScopedMPISetter::_nccl_comms;