Eliminate ATen/cuda/CUDAContext.h and c10/cuda/CUDAGuard.h

Author: pearu

src/libtorchaudio/cuda_utils.h

@@ -1,6 +1,6 @@
 #pragma once
 
-#include <torch/csrc/inductor/aoti_torch/c/shim.h>
+#include <torch/csrc/stable/c/shim.h>
 #include <torch/csrc/stable/device.h>
 
 #include <cuda_runtime_api.h>
@@ -17,9 +17,27 @@ inline cudaStream_t getCurrentCUDAStream(
   return static_cast<cudaStream_t>(stream_ptr);
 }
 
-// A strip-down version of at::cuda::stream_synchronize
-inline void stream_synchronize(cudaStream_t stream) {
-  TA_CUDA_CHECK(cudaStreamSynchronize(stream));
+inline void setCurrentCUDAStream(
+    cudaStream_t stream,
+    torch::stable::DeviceIndex device_index = -1) {
+  TORCH_ERROR_CODE_CHECK(
+      torch_set_current_cuda_stream(static_cast<void*>(stream), device_index));
+}
+
+inline cudaStream_t getStreamFromPool(
+    const bool isHighPriority = false,
+    torch::stable::DeviceIndex device_index = -1) {
+  void* stream_ptr = nullptr;
+  TORCH_ERROR_CODE_CHECK(torch_get_cuda_stream_from_pool(
+      isHighPriority, device_index, &stream_ptr));
+  return static_cast<cudaStream_t>(stream_ptr);
+}
+
+inline void synchronize(
+    cudaStream_t stream,
+    torch::stable::DeviceIndex device_index = -1) {
+  TORCH_ERROR_CODE_CHECK(
+      torch_cuda_stream_synchronize(static_cast<void*>(stream), device_index));
 }
 
 } // namespace libtorchaudio::cuda

src/libtorchaudio/forced_align/gpu/compute.cu

@@ -1,8 +1,9 @@
+#include <libtorchaudio/cuda_utils.h>
 #include <libtorchaudio/utils.h>
 #include <torch/csrc/stable/library.h>
 #include <torch/headeronly/core/Dispatch_v2.h>
 #include <torch/headeronly/core/ScalarType.h>
-#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAException.h>
 
 #include <cub/cub.cuh>
 #include <limits.h>
@@ -120,8 +121,9 @@ void forced_align_impl(
     const Tensor& targets,
     const int64_t blank,
     Tensor& paths) {
-  auto defaultStream = at::cuda::getCurrentCUDAStream();
-  auto cpuDataTranferStream = at::cuda::getStreamFromPool();
+  auto device_index = logProbs.get_device_index();
+  auto defaultStream = libtorchaudio::cuda::getCurrentCUDAStream(device_index);
+  auto cpuDataTranferStream = libtorchaudio::cuda::getStreamFromPool(false, device_index);
   const scalar_t kNegInfinity = -std::numeric_limits<scalar_t>::infinity();
   using target_t = typename std::
       conditional<target_scalar_type == ScalarType::Int, int, int64_t>::type;
@@ -208,12 +210,14 @@ void forced_align_impl(
     C10_CUDA_KERNEL_LAUNCH_CHECK();
     ++backPtrBufferLen;
     if (backPtrBufferLen == kBackPtrBufferSize || t == T - 1) {
-      cpuDataTranferStream.synchronize();
+      //cpuDataTranferStream.synchronize();
+      libtorchaudio::cuda::synchronize(cpuDataTranferStream, device_index);
       // GPU -> GPU copy
       bufferCopy = torch::stable::clone(backPtrBuffer);
       STD_TORCH_CHECK(bufferCopy.is_contiguous(), "unexpected fail, need to implement stable::Tensor::contiguous()")
-      defaultStream.synchronize();
-      at::cuda::setCurrentCUDAStream(cpuDataTranferStream);
+      //defaultStream.synchronize();
+      libtorchaudio::cuda::synchronize(defaultStream, device_index);
+      libtorchaudio::cuda::setCurrentCUDAStream(cpuDataTranferStream, device_index);
       // Copy ASYNC from GPU to CPU
       int64_t offset =
           static_cast<int64_t>(t + 1 - backPtrBufferLen) * S * sizeof(int8_t);
@@ -223,11 +227,12 @@ void forced_align_impl(
           backPtrBufferLen * S * sizeof(int8_t),
           cudaMemcpyDeviceToHost,
           cpuDataTranferStream));
-      at::cuda::setCurrentCUDAStream(defaultStream);
+      libtorchaudio::cuda::setCurrentCUDAStream(defaultStream, device_index);
       backPtrBufferLen = 0;
     }
   }
-  cpuDataTranferStream.synchronize();
+  //cpuDataTranferStream.synchronize();
+  libtorchaudio::cuda::synchronize(cpuDataTranferStream, device_index);
   auto alphasCpu = torchaudio::stable::cpu(alphas);
   auto alphasCpu_a = torchaudio::accessor<scalar_t, 2>(alphasCpu);
   int curIdxOffset = ((T - 1) % 2);

src/libtorchaudio/iir_cuda.cu

@@ -1,8 +1,8 @@
 #include <libtorchaudio/utils.h>
+#include <torch/csrc/stable/accelerator.h>
 #include <torch/headeronly/core/Dispatch_v2.h>
 #include <torch/headeronly/core/ScalarType.h>
-#include <c10/cuda/CUDAGuard.h>
-#include <c10/core/DeviceGuard.h>
+#include <c10/cuda/CUDAException.h>
 
 using torch::headeronly::ScalarType;
 using torch::stable::Tensor;
@@ -64,8 +64,7 @@ Tensor cuda_lfilter_core_loop(
 
   STD_TORCH_CHECK(in.size(2) + a_flipped.size(1) - 1 == padded_out.size(2));
 
-  const at::cuda::OptionalCUDAGuard device_guard(in.get_device_index());
-
+  const torch::stable::accelerator::DeviceGuard device_guard(in.get_device_index());
   const dim3 threads(256);
   const dim3 blocks((N * C + threads.x - 1) / threads.x);