Add temporary shim functions.

Author: pearu

src/libtorchaudio/cuda_utils.h

@@ -1,5 +1,6 @@
 #pragma once
 
+#include <libtorchaudio/shim_temporary.h>
 #include <torch/csrc/stable/c/shim.h>
 #include <torch/csrc/stable/device.h>
 
@@ -20,24 +21,24 @@ inline cudaStream_t getCurrentCUDAStream(
 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));
+  TORCH_ERROR_CODE_CHECK(tmp_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(
+  TORCH_ERROR_CODE_CHECK(tmp_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));
+  TORCH_ERROR_CODE_CHECK(tmp_torch_cuda_stream_synchronize(
+      static_cast<void*>(stream), device_index));
 }
 
 } // namespace libtorchaudio::cuda

src/libtorchaudio/shim_temporary.h

@@ -0,0 +1,36 @@
+#pragma once
+// TODO: remove this file once https://github.com/pytorch/pytorch/pull/169376
+// has landed.
+
+#include <c10/cuda/CUDAStream.h>
+#include <torch/csrc/inductor/aoti_torch/utils.h>
+#include <torch/csrc/stable/c/shim.h>
+
+inline AOTITorchError tmp_torch_set_current_cuda_stream(
+    void* stream,
+    int32_t device_index) {
+  AOTI_TORCH_CONVERT_EXCEPTION_TO_ERROR_CODE({
+    at::cuda::setCurrentCUDAStream(at::cuda::getStreamFromExternal(
+        static_cast<cudaStream_t>(stream), device_index));
+  });
+}
+
+inline AOTITorchError tmp_torch_get_cuda_stream_from_pool(
+    const bool isHighPriority,
+    int32_t device_index,
+    void** ret_stream) {
+  AOTI_TORCH_CONVERT_EXCEPTION_TO_ERROR_CODE({
+    *(cudaStream_t*)(ret_stream) =
+        at::cuda::getStreamFromPool(isHighPriority, device_index);
+  });
+}
+
+inline AOTITorchError tmp_torch_cuda_stream_synchronize(
+    void* stream,
+    int32_t device_index) {
+  AOTI_TORCH_CONVERT_EXCEPTION_TO_ERROR_CODE({
+    at::cuda::getStreamFromExternal(
+        static_cast<cudaStream_t>(stream), device_index)
+        .synchronize();
+  });
+}