Move device reactivation to corecel

celeritas-project · Jun 1, 2023 · e71fa0c · e71fa0c
1 parent dae5633
commit e71fa0c
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 6 deletions.
diff --git a/app/celer-sim/celer-sim.cc b/app/celer-sim/celer-sim.cc
@@ -127,12 +127,7 @@ void run(std::istream* is, std::shared_ptr<OutputRegistry> output)
 #endif
         for (size_type event = 0; event < run_stream.num_events(); ++event)
         {
-            if (device())
-            {
-                // See
-                // https://developer.nvidia.com/blog/cuda-pro-tip-always-set-current-device-avoid-multithreading-bugs/
-                CELER_DEVICE_CALL_PREFIX(SetDevice(device().device_id()));
-            }
+            activate_device_local();
 
             // Run a single event on a single thread
             CELER_TRY_HANDLE(result.events[event] = run_stream(

diff --git a/src/corecel/sys/Device.cc b/src/corecel/sys/Device.cc
@@ -329,6 +329,23 @@ void activate_device(MpiCommunicator const& comm)
     }
 }
 
+//---------------------------------------------------------------------------//
+/*!
+ * Call cudaSetDevice using the existing device, for thread-local safety.
+ *
+ * See
+ * https://developer.nvidia.com/blog/cuda-pro-tip-always-set-current-device-avoid-multithreading-bugs
+ *
+ * \pre activate_device was called to set \c device()
+ */
+void activate_device_local()
+{
+    if (device())
+    {
+        CELER_DEVICE_CALL_PREFIX(SetDevice(device().device_id()));
+    }
+}
+
 //---------------------------------------------------------------------------//
 /*!
  * Print device info.

diff --git a/src/corecel/sys/Device.hh b/src/corecel/sys/Device.hh
@@ -151,6 +151,9 @@ void activate_device();
 // Initialize a device in a round-robin fashion from a communicator.
 void activate_device(MpiCommunicator const&);
 
+// Call cudaSetDevice using the existing device, for thread-local safety
+void activate_device_local();
+
 // Print device info
 std::ostream& operator<<(std::ostream&, Device const&);