Remove USE_CUDA and USE_ROCM in engine.cpp

Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/15893

Differential Revision: D13627319

Pulled By: zdevito

fbshipit-source-id: 7c72c1c6cc242143fb66383423c668c9b9810884

diff --git a/aten/src/ATen/Context.h b/aten/src/ATen/Context.h
index 4238b5f..77a8c1f 100644 (file)
--- a/aten/src/ATen/Context.h
+++ b/aten/src/ATen/Context.h
@@ -203,6 +203,16 @@ static inline bool hasHIP() {
   return globalContext().hasHIP();
 }
 
+static inline size_t getNumGPUs() {
+  if (hasCUDA()) {
+    return detail::getCUDAHooks().getNumGPUs();
+  }
+  if (hasHIP()) {
+    return detail::getHIPHooks().getNumGPUs();
+  }
+  return 0;
+}
+
 static inline bool hasMKL() {
   return globalContext().hasMKL();
 }

diff --git a/torch/csrc/autograd/engine.cpp b/torch/csrc/autograd/engine.cpp
index 6063baa..44ab9c5 100644 (file)
--- a/torch/csrc/autograd/engine.cpp
+++ b/torch/csrc/autograd/engine.cpp
@@ -27,16 +27,6 @@
 #include <queue>
 #include <TH/TH.h>
 
-#ifdef USE_CUDA
-#include <cuda.h>
-#include <c10/cuda/CUDAGuard.h>
-#endif  // USE_CUDA
-
-#ifdef USE_ROCM
-#include <hip/hip_runtime.h>
-#include <ATen/hip/impl/HIPGuardImplMasqueradingAsCUDA.h>
-#endif  // USE_ROCM
-
 namespace torch { namespace autograd {
 
 // NB: -1 indicates the CPU worker!
@@ -217,22 +207,20 @@ Engine::~Engine() = default;
 // not CUDA.
 auto Engine::thread_init(int device) -> void {
   THInferNumThreads();
-#if defined(USE_CUDA)
   // NB: We MUST NOT construct the guard for device -1,
-  // as in some settings we compile with USE_CUDA, but
+  // as in some settings we compile with cuda, but
   // have lazy stubs for CUDA functionality (so actually
   // attempting to setup a guard(-1) will cause an
   // error, because it will still query cudaGetDevice).
-  at::cuda::OptionalCUDAGuard guard;
-  if (device != -1) {
-    guard.set_index(device);
-  }
-#elif defined(USE_ROCM)
-  at::cuda::OptionalHIPGuardMasqueradingAsCUDA guard;
+  at::OptionalDeviceGuard guard;
   if (device != -1) {
-    guard.set_index(device);
+    if (at::hasCUDA()) {
+      guard.reset_device(at::Device(at::DeviceType::CUDA, device));
+    }
+    if (at::hasHIP()) {
+      guard.reset_device(at::Device(at::DeviceType::HIP, device));
+    }
   }
-#endif
   worker_device = device;
   thread_main(nullptr);
 }
@@ -644,29 +632,7 @@ auto Engine::ready_queue(int device) -> ReadyQueue& {
 }
 
 auto Engine::start_threads() -> void {
-  int num_devices = 0;
-#ifdef USE_CUDA
-  {
-    int num_cuda_devices = 0;
-    // check for case of compiled with CUDA but no available devices
-    if (cudaGetDeviceCount(&num_cuda_devices) != cudaSuccess) {
-      cudaGetLastError();
-    } else {
-      num_devices += num_cuda_devices;
-    }
-  }
-#endif
-#ifdef USE_ROCM
-  {
-    int num_hip_devices = 0;
-    // check for case of compiled with CUDA but no available devices
-    if (hipGetDeviceCount(&num_hip_devices) != hipSuccess) {
-      hipGetLastError();
-    } else {
-      num_devices += num_hip_devices;
-    }
-  }
-#endif
+  int num_devices = at::getNumGPUs();
   // One for CPU, plus one for every GPU device
   int num_threads = num_devices + 1;
   ready_queues = std::vector<std::shared_ptr<ReadyQueue>>(num_threads);

diff --git a/torch/csrc/jit/ivalue.h b/torch/csrc/jit/ivalue.h
index 87d617c..d3130db 100644 (file)
--- a/torch/csrc/jit/ivalue.h
+++ b/torch/csrc/jit/ivalue.h
@@ -1,3 +1,4 @@
+#pragma once
 #include <ATen/core/ivalue.h>
 
 namespace torch {