Multi device bugfix

openequivariance/CMakeLists.txt

@@ -117,7 +117,7 @@ endfunction()
 
 find_package(CUDAToolkit QUIET)
 find_package(hip QUIET)
-find_package(rocblas QUIET)
+find_package(hipblas QUIET)
 
 if(CUDAToolkit_FOUND)
     message(STATUS "Building stable extension with CUDA backend.")
@@ -159,10 +159,10 @@ if(hip_FOUND)
         CXX_STANDARD 17
     )
 
-    if(TARGET roc::rocblas)
-        set(HIP_BLAS_LIB roc::rocblas)
+    if(TARGET roc::hipblas)
+        set(HIP_BLAS_LIB roc::hipblas)
     else()
-        set(HIP_BLAS_LIB rocblas)
+        set(HIP_BLAS_LIB hipblas)
     endif()
 
     set(HIP_LINK_LIBS

openequivariance/openequivariance/_torch/TensorProductConv.py

@@ -87,7 +87,11 @@ def _init_class(self):
 
         self.allocate_workspace(self.workspace_size)
 
-        self.dummy_transpose_perm = torch.zeros(1, dtype=torch.int64, device="cuda")
+        self.register_buffer(
+            "dummy_transpose_perm",
+            torch.zeros(1, dtype=torch.int64, device=self.workspace_buffer.device),
+            persistent=False,
+        )
         self.weight_numel = self.config.weight_numel
         self.kernel = string_to_tensor(self.kernel_string)
         self.L3_dim = self.kernel_prop["L3_dim"]
@@ -129,7 +133,9 @@ def _apply(self, fn, recurse=True):
             self.to(result.dtype)
             self._applying = False
 
-        return super()._apply(fn, recurse)
+        out = super()._apply(fn, recurse)
+        self.workspace_ptr = self.workspace_buffer.data_ptr()
+        return out
 
     def __getstate__(self):
         return self.input_args
@@ -184,10 +190,15 @@ def forward(
             sender_perm,
         )
 
-    def allocate_workspace(self, size_bytes):
+    def allocate_workspace(self, size_bytes, device=None):
         self.workspace_size = size_bytes
-        self.workspace_buffer = torch.zeros(
-            size_bytes, dtype=torch.uint8, device="cuda"
+        if device is None:
+            device = getattr(self, "workspace_buffer", None)
+            device = device.device if device is not None else "cuda"
+        self.register_buffer(
+            "workspace_buffer",
+            torch.zeros(size_bytes, dtype=torch.uint8, device=device),
+            persistent=False,
         )
         self.workspace_ptr = self.workspace_buffer.data_ptr()
         logger.info(f"Convolution requires {size_bytes // 1000000}MB of workspace.")

openequivariance/openequivariance/extension/group_mm.hpp

@@ -8,38 +8,21 @@
     #include "cublas_v2.h"
     #include <cuda_runtime.h>
 
-    struct BlasHandle {
-        cublasHandle_t handle;
-        BlasHandle() {
-            if (cublasCreate(&handle) != CUBLAS_STATUS_SUCCESS)
-                throw std::logic_error("CUBLAS initialization failed");
-        }
-        ~BlasHandle() { cublasDestroy(handle); }
-    };
+    using BlasHandleT = cublasHandle_t;
 #elif defined(HIP_BACKEND)
-    #include "rocblas/rocblas.h"
+    #include <hipblas/hipblas.h>
     #include <hip/hip_runtime.h>
 
-    struct BlasHandle {
-        rocblas_handle handle;
-        BlasHandle() {
-            if (rocblas_create_handle(&handle) != rocblas_status_success)
-                throw std::logic_error("rocBLAS initialization failed");
-        }
-        ~BlasHandle() { rocblas_destroy_handle(handle); }
-    };
+    using BlasHandleT = hipblasHandle_t;
 #endif
 
-inline BlasHandle& get_blas_handle() {
-    static BlasHandle handle;
-    return handle;
-}
+BlasHandleT get_op_blas_handle();
 
 template<typename T>
 void group_gemm_blas(void* A_raw, void* B_raw, void* C_raw,
         int64_t* ragged_counts, int num_W, int batch_size, int m, int k, int ragged_inner) {
 
-    auto& blas = get_blas_handle();
+    BlasHandleT handle = get_op_blas_handle();
     T alpha = 1.0, beta = 0.0;
     T* A_base = reinterpret_cast<T*>(A_raw);
     T* B_base = reinterpret_cast<T*>(B_raw);
@@ -52,7 +35,7 @@ void group_gemm_blas(void* A_raw, void* B_raw, void* C_raw,
 #ifdef CUDA_BACKEND
         cublasOperation_t transa, transb;
 #elif defined(HIP_BACKEND)
-        rocblas_operation transa, transb;
+        hipblasOperation_t transa, transb;
 #endif
 
         if (ragged_inner == 0) {
@@ -66,7 +49,7 @@ void group_gemm_blas(void* A_raw, void* B_raw, void* C_raw,
 #ifdef CUDA_BACKEND
             transa = CUBLAS_OP_T; transb = CUBLAS_OP_N;
 #elif defined(HIP_BACKEND)
-            transa = rocblas_operation_transpose; transb = rocblas_operation_none;
+            transa = HIPBLAS_OP_T; transb = HIPBLAS_OP_N;
 #endif
         } else {
             M = k; K = static_cast<int>(ragged_counts[i]); N = m;
@@ -79,7 +62,7 @@ void group_gemm_blas(void* A_raw, void* B_raw, void* C_raw,
 #ifdef CUDA_BACKEND
             transa = CUBLAS_OP_N; transb = CUBLAS_OP_T;
 #elif defined(HIP_BACKEND)
-            transa = rocblas_operation_none; transb = rocblas_operation_transpose;
+            transa = HIPBLAS_OP_N; transb = HIPBLAS_OP_T;
 #endif
         }
         ragged_offset += ragged_counts[i];
@@ -88,7 +71,7 @@ void group_gemm_blas(void* A_raw, void* B_raw, void* C_raw,
 #ifdef CUDA_BACKEND
             cublasStatus_t stat;
             if (std::is_same<T, float>::value) {
-                stat = cublasSgemmStridedBatched(blas.handle,
+                stat = cublasSgemmStridedBatched(handle,
                     transa, transb, M, N, K,
                     reinterpret_cast<float*>(&alpha),
                     reinterpret_cast<float*>(A), lda, strideA,
@@ -97,7 +80,7 @@ void group_gemm_blas(void* A_raw, void* B_raw, void* C_raw,
                     reinterpret_cast<float*>(C), ldc, strideC,
                     batch_size);
             } else if (std::is_same<T, double>::value) {
-                stat = cublasDgemmStridedBatched(blas.handle,
+                stat = cublasDgemmStridedBatched(handle,
                     transa, transb, M, N, K,
                     reinterpret_cast<double*>(&alpha),
                     reinterpret_cast<double*>(A), lda, strideA,
@@ -111,9 +94,9 @@ void group_gemm_blas(void* A_raw, void* B_raw, void* C_raw,
             if (stat != CUBLAS_STATUS_SUCCESS)
                 throw std::logic_error("Grouped GEMM failed!");
 #elif defined(HIP_BACKEND)
-            rocblas_status stat;
+            hipblasStatus_t stat;
             if (std::is_same<T, float>::value) {
-                stat = rocblas_sgemm_strided_batched(blas.handle,
+                stat = hipblasSgemmStridedBatched(handle,
                     transa, transb, M, N, K,
                     reinterpret_cast<float*>(&alpha),
                     reinterpret_cast<float*>(A), lda, strideA,
@@ -122,7 +105,7 @@ void group_gemm_blas(void* A_raw, void* B_raw, void* C_raw,
                     reinterpret_cast<float*>(C), ldc, strideC,
                     batch_size);
             } else if (std::is_same<T, double>::value) {
-                stat = rocblas_dgemm_strided_batched(blas.handle,
+                stat = hipblasDgemmStridedBatched(handle,
                     transa, transb, M, N, K,
                     reinterpret_cast<double*>(&alpha),
                     reinterpret_cast<double*>(A), lda, strideA,
@@ -133,7 +116,7 @@ void group_gemm_blas(void* A_raw, void* B_raw, void* C_raw,
             } else {
                 throw std::logic_error("Unsupported datatype for grouped GEMM!");
             }
-            if (stat != rocblas_status_success)
+            if (stat != HIPBLAS_STATUS_SUCCESS)
                 throw std::logic_error("Grouped GEMM failed!");
 #endif
         }

openequivariance/openequivariance/extension/libtorch_tp_jit.cpp

@@ -3,10 +3,13 @@
 
 #ifdef CUDA_BACKEND
     #include <ATen/cuda/CUDAContext.h>
+    #include <c10/cuda/CUDAGuard.h>
 #endif
 
 #ifdef HIP_BACKEND
+    #include <ATen/hip/HIPContext.h>
     #include <c10/hip/HIPStream.h>
+    #include <c10/hip/HIPGuard.h>
 #endif
 
 #include <ATen/Operators.h>
@@ -29,6 +32,13 @@ constexpr Dtype kByte = torch::kByte;
 #define REGISTER_LIBRARY_IMPL TORCH_LIBRARY_IMPL
 #define REGISTER_LIBRARY TORCH_LIBRARY
 
+#ifdef CUDA_BACKEND
+    using DeviceGuard = c10::cuda::CUDAGuard;
+#endif
+#ifdef HIP_BACKEND
+    using DeviceGuard = c10::hip::HIPGuard;
+#endif
+
 #include "torch_core.hpp"
 
 Tensor tensor_to_cpu_contiguous(const Tensor &tensor) {
@@ -74,15 +84,19 @@ void *data_ptr(const Tensor &tensor) {
         throw std::logic_error("Unsupported tensor datatype!");
 }
 
-Stream get_current_stream() {
+Stream get_current_stream(int32_t device_index) {
 #ifdef CUDA_BACKEND
-    return c10::cuda::getCurrentCUDAStream();
+    return c10::cuda::getCurrentCUDAStream(device_index);
 #endif
 #ifdef HIP_BACKEND
-    return c10::hip::getCurrentHIPStream();
+    return c10::hip::getCurrentHIPStream(device_index);
 #endif
 }
 
+BlasHandleT get_op_blas_handle() {
+    return at::cuda::getCurrentCUDABlasHandle();
+}
+
 namespace py=pybind11;
 PYBIND11_MODULE(libtorch_tp_jit, m) {
     py::class_<DeviceProp>(m, "DeviceProp")

openequivariance/openequivariance/extension/libtorch_tp_jit_stable.cpp

@@ -11,6 +11,7 @@
 #include <torch/headeronly/util/Exception.h>
 #include <torch/headeronly/util/shim_utils.h>
 #include <torch/csrc/inductor/aoti_torch/c/shim.h>
+#include <torch/csrc/stable/c/shim.h>
 
 
 using Tensor = torch::stable::Tensor;
@@ -27,6 +28,8 @@ constexpr Dtype kByte = torch::headeronly::ScalarType::Byte;
 #define REGISTER_LIBRARY_IMPL STABLE_TORCH_LIBRARY_IMPL
 #define REGISTER_LIBRARY STABLE_TORCH_LIBRARY
 
+using DeviceGuard = torch::stable::accelerator::DeviceGuard;
+
 #include "torch_core.hpp"
 
 Tensor tensor_to_cpu_contiguous(const Tensor &tensor) {
@@ -66,16 +69,16 @@ void *data_ptr(const Tensor &tensor) {
     return tensor.data_ptr();
 }
 
-Stream get_current_stream() {
-    auto device_idx = torch::stable::accelerator::getCurrentDeviceIndex();
+Stream get_current_stream(int32_t device_index) {
     void* stream_ptr = nullptr;
-    TORCH_ERROR_CODE_CHECK(aoti_torch_get_current_cuda_stream(device_idx, &stream_ptr));
+    TORCH_ERROR_CODE_CHECK(aoti_torch_get_current_cuda_stream(device_index, &stream_ptr));
+    return static_cast<Stream>(stream_ptr);
+}
 
-    #ifdef CUDA_BACKEND
-        return static_cast<Stream>(stream_ptr); 
-    #elif defined(HIP_BACKEND)
-        return static_cast<Stream>(stream_ptr);
-    #endif
+BlasHandleT get_op_blas_handle() {
+    void* handle = nullptr;
+    TORCH_ERROR_CODE_CHECK(torch_get_current_cuda_blas_handle(&handle));
+    return static_cast<BlasHandleT>(handle);
 }
 
 #ifdef CUDA_BACKEND

openequivariance/openequivariance/extension/stubs/stream.cpp

@@ -1,8 +1,41 @@
 #include <cstdint>
+#include <cstdio>
 #include <torch/csrc/inductor/aoti_torch/c/shim.h>
 
+namespace {
+
+AOTITorchError oeq_stub_called(const char *name) {
+    fprintf(stderr,
+            "OpenEquivariance: link-time stub '%s' executed; the real "
+            "libtorch_cuda/libtorch_hip is not loaded (import torch before "
+            "loading the OpenEquivariance extension).\n",
+            name);
+    return AOTI_TORCH_FAILURE;
+}
+
+} // namespace
+
 extern "C" {
     AOTITorchError aoti_torch_get_current_cuda_stream(int32_t device_index, void** ret_stream) {
-        return 0;
+        if (ret_stream) *ret_stream = nullptr;
+        return oeq_stub_called("aoti_torch_get_current_cuda_stream");
+    }
+
+    AOTITorchError aoti_torch_create_device_guard(int32_t device_index, DeviceGuardHandle* ret_guard) {
+        if (ret_guard) *ret_guard = nullptr;
+        return oeq_stub_called("aoti_torch_create_device_guard");
+    }
+
+    AOTITorchError aoti_torch_delete_device_guard(DeviceGuardHandle guard) {
+        return oeq_stub_called("aoti_torch_delete_device_guard");
+    }
+
+    AOTITorchError aoti_torch_device_guard_set_index(DeviceGuardHandle guard, int32_t device_index) {
+        return oeq_stub_called("aoti_torch_device_guard_set_index");
+    }
+
+    AOTITorchError torch_get_current_cuda_blas_handle(void** ret_handle) {
+        if (ret_handle) *ret_handle = nullptr;
+        return oeq_stub_called("torch_get_current_cuda_blas_handle");
     }
-}
+}