Replace use of cuda::device::current::scoped_override_t with plain CU…

…DA (#407) Remove cuda::device::current::scoped_override_t from CUDAScopedContext . Replace cuda::device::current::scoped_override_t wth cudautils::ScopedSetDevice .
cms-patatrack · Nov 8, 2019 · 614ee0b · 614ee0b
1 parent 20aff22
commit 614ee0b
Show file tree

Hide file tree

Showing 4 changed files with 18 additions and 7 deletions.
diff --git a/HeterogeneousCore/CUDACore/interface/CUDAScopedContext.h b/HeterogeneousCore/CUDACore/interface/CUDAScopedContext.h
@@ -33,6 +33,12 @@ namespace impl {
     const cudautils::SharedStreamPtr& streamPtr() const { return stream_; }
 
   protected:
+    // The constructors set the current device device, but the device
+    // is not set back to the previous value at the destructor. This
+    // should be sufficient (and tiny bit faster) as all CUDA API
+    // functions relying on the current device should be called from
+    // the scope where this context is. The current device doesn't
+    // really matter between modules (or across TBB tasks).
     explicit CUDAScopedContextBase(edm::StreamID streamID);
 
     explicit CUDAScopedContextBase(const CUDAProductBase& data);
@@ -41,7 +47,6 @@ namespace impl {
 
   private:
     int currentDevice_;
-    cuda::device::current::scoped_override_t<> setDeviceForThisScope_;
     cudautils::SharedStreamPtr stream_;
   };
 

diff --git a/HeterogeneousCore/CUDACore/src/CUDAScopedContext.cc b/HeterogeneousCore/CUDACore/src/CUDAScopedContext.cc
@@ -39,12 +39,14 @@ namespace {
 
 namespace impl {
   CUDAScopedContextBase::CUDAScopedContextBase(edm::StreamID streamID)
-      : currentDevice_(cudacore::chooseCUDADevice(streamID)), setDeviceForThisScope_(currentDevice_) {
+      : currentDevice_(cudacore::chooseCUDADevice(streamID)) {
+    cudaCheck(cudaSetDevice(currentDevice_));
     stream_ = cudautils::getCUDAStreamCache().getCUDAStream();
   }
 
   CUDAScopedContextBase::CUDAScopedContextBase(const CUDAProductBase& data)
-      : currentDevice_(data.device()), setDeviceForThisScope_(currentDevice_) {
+      : currentDevice_(data.device()) {
+    cudaCheck(cudaSetDevice(currentDevice_));
     if (data.mayReuseStream()) {
       stream_ = data.streamPtr();
     } else {
@@ -53,7 +55,9 @@ namespace impl {
   }
 
   CUDAScopedContextBase::CUDAScopedContextBase(int device, cudautils::SharedStreamPtr stream)
-      : currentDevice_(device), setDeviceForThisScope_(device), stream_(std::move(stream)) {}
+      : currentDevice_(device), stream_(std::move(stream)) {
+    cudaCheck(cudaSetDevice(currentDevice_));
+  }
 
   ////////////////////
 

diff --git a/HeterogeneousCore/CUDACore/test/test_CUDAScopedContext.cc b/HeterogeneousCore/CUDACore/test/test_CUDAScopedContext.cc
@@ -11,6 +11,7 @@
 #include "HeterogeneousCore/CUDAUtilities/interface/exitSansCUDADevices.h"
 #include "HeterogeneousCore/CUDAUtilities/interface/CUDAStreamCache.h"
 #include "HeterogeneousCore/CUDAUtilities/interface/CUDAEventCache.h"
+#include "HeterogeneousCore/CUDAUtilities/interface/ScopedSetDevice.h"
 
 #include "test_CUDAScopedContextKernels.h"
 
@@ -84,7 +85,7 @@ TEST_CASE("Use of CUDAScopedContext", "[CUDACore]") {
     }
 
     SECTION("Joining multiple CUDA streams") {
-      cuda::device::current::scoped_override_t<> setDeviceForThisScope(defaultDevice);
+      cudautils::ScopedSetDevice setDeviceForThisScope(defaultDevice);
       auto current_device = cuda::device::current::get();
 
       // Mimick a producer on the first CUDA stream

diff --git a/HeterogeneousCore/CUDAUtilities/src/allocate_device.cc b/HeterogeneousCore/CUDAUtilities/src/allocate_device.cc
@@ -1,5 +1,6 @@
 #include "HeterogeneousCore/CUDAUtilities/interface/allocate_device.h"
 #include "HeterogeneousCore/CUDAUtilities/interface/cudaCheck.h"
+#include "HeterogeneousCore/CUDAUtilities/interface/ScopedSetDevice.h"
 #include "FWCore/Utilities/interface/Likely.h"
 
 #include "getCachingDeviceAllocator.h"
@@ -23,7 +24,7 @@ namespace cudautils {
       }
       cuda::throw_if_error(cudautils::allocator::getCachingDeviceAllocator().DeviceAllocate(dev, &ptr, nbytes, stream));
     } else {
-      cuda::device::current::scoped_override_t<> setDeviceForThisScope(dev);
+      ScopedSetDevice setDeviceForThisScope(dev);
       cuda::throw_if_error(cudaMalloc(&ptr, nbytes));
     }
     return ptr;
@@ -33,7 +34,7 @@ namespace cudautils {
     if constexpr (cudautils::allocator::useCaching) {
       cuda::throw_if_error(cudautils::allocator::getCachingDeviceAllocator().DeviceFree(device, ptr));
     } else {
-      cuda::device::current::scoped_override_t<> setDeviceForThisScope(device);
+      ScopedSetDevice setDeviceForThisScope(device);
       cuda::throw_if_error(cudaFree(ptr));
     }
   }