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/extension/group_mm.hpp

@@ -8,38 +8,29 @@
     #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 BlasStream = cudaStream_t;
+    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 BlasStream = hipStream_t;
+    using BlasHandleT = hipblasHandle_t;
 #endif
 
-inline BlasHandle& get_blas_handle() {
-    static BlasHandle handle;
-    return handle;
-}
+// Returns a BLAS handle ready to launch on (device_index, stream). Defined
+// per-build (libtorch_tp_jit.cpp / libtorch_tp_jit_stable.cpp): both return
+// PyTorch's pooled handle (cuBLAS on CUDA, hipBLAS on ROCm), which arrives
+// with the stream, workspace, and math mode already set for the current
+// device and stream. The caller must hold a device guard for device_index.
+BlasHandleT get_op_blas_handle(int32_t device_index, BlasStream stream);
 
 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) {
+        int64_t* ragged_counts, int num_W, int batch_size, int m, int k, int ragged_inner,
+        int32_t device_index, BlasStream stream) {
 
-    auto& blas = get_blas_handle();
+    BlasHandleT handle = get_op_blas_handle(device_index, stream);
     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 +43,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 +57,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 +70,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 +79,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 +88,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 +102,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 +113,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 +124,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,26 @@ 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(int32_t device_index, BlasStream stream) {
+    // The caller's device guard makes device_index current, and `stream` is
+    // that device's current stream, so PyTorch's handle arrives configured
+    // for exactly this (device, stream) with its workspace and math mode
+    // managed by PyTorch. On ROCm the same-named function returns a
+    // hipblasHandle_t.
+    (void)device_index;
+    (void)stream;
+    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,23 @@ 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(int32_t device_index, BlasStream stream) {
+    // The caller's device guard makes device_index current, and `stream` is
+    // that device's current stream, so PyTorch's handle arrives configured
+    // for exactly this (device, stream) with its workspace and math mode
+    // managed by PyTorch. On ROCm builds of libtorch this shim symbol keeps
+    // its name and returns a hipblasHandle_t.
+    (void)device_index;
+    (void)stream;
+    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,51 @@
 #include <cstdint>
+#include <cstdio>
 #include <torch/csrc/inductor/aoti_torch/c/shim.h>
 
+// The stable extension is linked on a machine with CPU-only libtorch, where
+// no libtorch_cuda/libtorch_hip exists. This file builds a stand-in library
+// with the real one's SONAME so the link succeeds; at runtime the DT_NEEDED
+// entry is satisfied by the real library, already loaded via `import torch`,
+// and this file is never opened. It is not shipped in the wheel.
+//
+// These bodies can therefore only execute in a misconfigured process (e.g.
+// the extension loaded without torch). Fail through the shim's error-code
+// contract - TORCH_ERROR_CODE_CHECK at the call site turns this into a
+// catchable exception - rather than returning success with garbage outputs.
+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");
     }
-}
+}