Create a macro for OptionalCUDAGuard (#2265)

Summary: Pull Request resolved: #2265 As title Reviewed By: jspark1105 Differential Revision: D52770256 fbshipit-source-id: 94ca1c549da6de99e0cc894dbb0b3f9838a99cf3
pytorch · Jan 15, 2024 · 5782e7e · 5782e7e
1 parent 7e50cd9
commit 5782e7e
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/fbgemm_gpu/include/fbgemm_gpu/fbgemm_cuda_utils.cuh b/fbgemm_gpu/include/fbgemm_gpu/fbgemm_cuda_utils.cuh
@@ -63,6 +63,10 @@ namespace cub = hipcub;
 
 #define DEVICE_INLINE __device__ inline __attribute__((always_inline))
 
+#define CUDA_DEVICE_GUARD(TENSOR)           \
+  at::cuda::OptionalCUDAGuard device_guard; \
+  device_guard.set_index(TENSOR.get_device())
+
 // Warp size
 #ifdef USE_ROCM
 static constexpr int32_t kWarpSize = 64;