Back to home page

Project CMSSW displayed by LXR

 
 

    


File indexing completed on 2024-04-06 12:15:45

0001 #ifndef HeterogeneousCore_CUDAUtilities_ScopedSetDevice_h
0002 #define HeterogeneousCore_CUDAUtilities_ScopedSetDevice_h
0003 
0004 #include "HeterogeneousCore/CUDAUtilities/interface/cudaCheck.h"
0005 
0006 #include <cuda_runtime.h>
0007 
0008 namespace cms {
0009   namespace cuda {
0010     class ScopedSetDevice {
0011     public:
0012       // Store the original device, without setting a new one
0013       ScopedSetDevice() {
0014         // Store the original device
0015         cudaCheck(cudaGetDevice(&originalDevice_));
0016       }
0017 
0018       // Store the original device, and set a new current device
0019       explicit ScopedSetDevice(int device) : ScopedSetDevice() {
0020         // Change the current device
0021         set(device);
0022       }
0023 
0024       // Restore the original device
0025       ~ScopedSetDevice() {
0026         // Intentionally don't check the return value to avoid
0027         // exceptions to be thrown. If this call fails, the process is
0028         // doomed anyway.
0029         cudaSetDevice(originalDevice_);
0030       }
0031 
0032       // Set a new current device, without changing the original device
0033       // that will be restored when this object is destroyed
0034       void set(int device) {
0035         // Change the current device
0036         cudaCheck(cudaSetDevice(device));
0037       }
0038 
0039     private:
0040       int originalDevice_;
0041     };
0042   }  // namespace cuda
0043 }  // namespace cms
0044 
0045 #endif