From 6437e1af5e451c78a611fb5fe87c286971a4229c Mon Sep 17 00:00:00 2001 From: asglover <140220574+asglover@users.noreply.github.com> Date: Sat, 12 Sep 2026 20:43:55 -0700 Subject: [PATCH 01/38] Add stable C shim BMM prototype and migration design --- sandbox/cshim_group_gemm/DESIGN.md | 292 ++++++++++++++++++++++++ sandbox/cshim_group_gemm/group_gemm.cpp | 99 ++++++++ sandbox/cshim_group_gemm/smoke_test.py | 71 ++++++ 3 files changed, 462 insertions(+) create mode 100644 sandbox/cshim_group_gemm/DESIGN.md create mode 100644 sandbox/cshim_group_gemm/group_gemm.cpp create mode 100644 sandbox/cshim_group_gemm/smoke_test.py diff --git a/sandbox/cshim_group_gemm/DESIGN.md b/sandbox/cshim_group_gemm/DESIGN.md new file mode 100644 index 00000000..3b919e07 --- /dev/null +++ b/sandbox/cshim_group_gemm/DESIGN.md @@ -0,0 +1,292 @@ +# Move Torch cuBLAS calls to the stable C shim + +Replace OEQ's direct CUDA and ROCm BLAS calls for `libtorch_tp_jit::group_gemm` with PyTorch's `aoti_torch_cuda_bmm_out` C API. Keep the existing grouped operation, tensor layouts, output allocation, and custom autograd rules. PyTorch will own the underlying BLAS handles, streams, library calls, and backend selection. + +Related: [PR #206](https://github.com/PASSIONLab/OpenEquivariance/pull/206). + +## Problem and scope + +OEQ's grouped GEMM is a host loop over nonempty ragged groups, issuing one strided batched GEMM per group. It does not require a vendor's heterogeneous grouped-GEMM API. The only direct BLAS computation calls found are float32/float64 strided batched GEMMs in `extension/group_mm.hpp`, called from the Torch backend in `extension/torch_core.hpp`. + +Using PyTorch's stable ABI to obtain a cuBLAS handle does not make that handle interchangeable with a separately loaded cuBLAS implementation. Earlier H100 experiments passed with matching handle/function owners, while some foreign-owner combinations returned incorrect streams, failed, or crashed. Successful version queries did not establish handle compatibility. Moving the GEMM itself behind PyTorch's C ABI removes this borrowed-handle boundary. + +The current checkout at `dbe854415da7771eba33195534c171adbca5677b` creates its own static BLAS handle; the inspected PR snapshot at `b551d7469db6d7ac688859cec46a3bb2b1c71a5b` borrows Torch's handle. The replacement removes either form of OEQ-side handle management. The final patch must be based on the branch being merged. + +The stable wheels currently target PyTorch 2.10. Use official stable C declarations, with no dependency on ATen's C++ ABI in that build. The C entry point was also exercised successfully on the available PyTorch 2.7 installation; that observation does not establish support for every older Torch release. [PyTorch stable ABI documentation](https://docs.pytorch.org/docs/main/notes/libtorch_stable_abi.html). + +## Public interface + +Include `torch/csrc/inductor/aoti_torch/generated/c_shim_cuda.h`, which declares: + +```cpp +AOTITorchError aoti_torch_cuda_bmm_out( + AtenTensorHandle out, + AtenTensorHandle self, + AtenTensorHandle mat2); +``` + +Use this generated API rather than the deprecated `aoti_torch_bmm_out` spelling. The API takes PyTorch tensor handles, not cuBLAS/hipBLAS handles. OEQ does not call a BLAS version getter or select a vendor library at runtime. [PyTorch 2.10 declaration](https://github.com/pytorch/pytorch/blob/v2.10.0/torch/csrc/inductor/aoti_torch/generated/c_shim_cuda.h#L74). + +Add a Torch-specific helper, such as `group_mm_torch.hpp`, taking the existing input/output pointers, dimensions, CPU ragged counts, dtype, and explicit device index. Both Torch extension variants should use it where their supported Torch versions expose the required C symbols. Keep this helper outside the framework-independent CUDA/HIP kernel backend. + +Create temporary, non-owning tensor views with `aoti_torch_create_tensor_from_blob_v2`. Specify sizes and strides in elements, storage offset zero, and the actual input device. Release each temporary tensor handle with `aoti_torch_delete_tensor_object` through RAII, including error paths. Releasing metadata must not release the caller's input or output storage. [PyTorch's blob implementation](https://github.com/pytorch/pytorch/blob/v2.10.0/torch/csrc/inductor/aoti_torch/shim_common.cpp#L525). + +## Layout mapping + +Let `b = batch_size`, `n = ragged_counts[i]`, and `o` be the sum of preceding counts. Pointer offsets and strides below are measured in elements. + +| Mode | View | Base pointer | Shape | Strides | +| --- | --- | --- | --- | --- | +| `ragged_inner == 0` | Left / input | `B + o*b*k` | `[b,n,k]` | `[k,b*k,1]` | +| `ragged_inner == 0` | Right / weights | `A + i*b*m*k` | `[b,k,m]` | `[m*k,1,k]` | +| `ragged_inner == 0` | Output | `C + o*b*m` | `[b,n,m]` | `[m,b*m,1]` | +| `ragged_inner == 1` | Left | `A + o*b*m` | `[b,m,n]` | `[m,1,b*m]` | +| `ragged_inner == 1` | Right | `B + o*b*k` | `[b,n,k]` | `[k,b*k,1]` | +| `ragged_inner == 1` | Output | `C + i*b*m*k` | `[b,m,k]` | `[m*k,k,1]` | + +Call BMM-out once for each nonempty group. Preserve the existing zero-initialized output and skip empty groups, including the all-empty case. This matters for the weight-gradient output blocks in mode 1. Use 64-bit dimensions and offsets throughout instead of the current narrowing casts to `int`. + +These views preserve OEQ's interleaved batch layout without explicitly transposing or copying the underlying buffers. PyTorch can handle BLAS-compatible strides directly, but some layouts can trigger its internal copy path; avoiding all copies is not an unconditional API guarantee. [BMM layout handling](https://github.com/pytorch/pytorch/blob/v2.10.0/aten/src/ATen/native/cuda/Blas.cpp#L536). + +## Device, stream, and autograd behavior + +Guard the inputs' actual GPU device before making contiguous copies, allocating the output, constructing views, or invoking BMM. Restore the caller's device on return. The 2.10 stable implementation can use `aoti_torch_create_device_guard` / `aoti_torch_delete_device_guard`, whose implementation selects Torch's active accelerator; this avoids requiring a CUDA-specific guard implementation in the shared helper. The existing standalone prototype uses the older CUDA-named guard because it was tested on Torch 2.7. [Generic guard implementation](https://github.com/pytorch/pytorch/blob/v2.10.0/torch/csrc/inductor/aoti_torch/shim_common.cpp#L1468). + +Use Torch's current stream for that device. Do not set a BLAS stream, create a stream, or synchronize inside the operator. Original tensors must remain valid through launch; asynchronous storage lifetime follows the usual PyTorch current-stream contract. No global or cached raw pointers or tensor views are needed. + +Validate the raw-pointer helper's preconditions at the tensor boundary: same GPU device and dtype, supported float32/float64 types, contiguous CPU int64 counts, valid counts length and nonnegative values, compatible input shapes, and a valid mode. Check that counts describe the available rows before constructing views. Keep the existing custom operator schema, fake implementation, and backward formulas. BMM-out is an internal implementation detail; the custom autograd registration remains responsible for derivatives. + +PyTorch's BMM implementation uses alpha 1 and beta 0, matching the existing GEMMs. It may choose a different BLAS backend or kernel and apply Torch's precision/determinism settings, so require numerical agreement rather than bitwise equivalence. [BMM implementation](https://github.com/pytorch/pytorch/blob/v2.10.0/aten/src/ATen/native/cuda/Blas.cpp#L680). + +## Why this should work on ROCm + +Source inspection of PyTorch **v2.10.0** supports using the same generated function on both GPU platforms: + +1. PyTorch's build defines `GENERATED_CXX_TORCH_CUDA` to contain `c_shim_cuda.cpp`, then explicitly adds that source to `torch_hip` under `USE_ROCM`. The generated shim is therefore part of the ROCm library too. [Generated source definition](https://github.com/pytorch/pytorch/blob/v2.10.0/caffe2/CMakeLists.txt#L331), [ROCm library construction](https://github.com/pytorch/pytorch/blob/v2.10.0/caffe2/CMakeLists.txt#L941). +2. The shim generator names the API using the `CUDA` dispatch key and generates a call into the corresponding ATen backend. The exported name remains `aoti_torch_cuda_bmm_out`; do not invent an `aoti_torch_hip_bmm_out` symbol. [Shim generator](https://github.com/pytorch/pytorch/blob/v2.10.0/torchgen/gen_aoti_c_shim.py#L492). +3. BMM reaches Torch's GEMM/batched GEMM implementation. Its HIP conversion maps `cublasSgemmStridedBatched` and `cublasDgemmStridedBatched` to their `hipblas` equivalents. Torch also contains ROCm-specific backend selection, including a double-precision fallback when hipBLASLt cannot handle the operation. [HIP mappings](https://github.com/pytorch/pytorch/blob/v2.10.0/torch/utils/hipify/cuda_to_hip_mappings.py#L6826), [float/double backend selection](https://github.com/pytorch/pytorch/blob/v2.10.0/aten/src/ATen/cuda/CUDABlas.cpp#L779). + +Torch's ROCm tensors use the `cuda` device interface. The helper should use Torch's CUDA device-type value for both these Torch builds, rather than confusing that convention with a distinct external HIP tensor type. [ROCm semantics](https://docs.pytorch.org/docs/main/notes/hip.html). + +This is **source-level evidence**, not an AMD hardware test or a binary export audit of a ROCm wheel. Verify the exports in the supported wheel and run one AMD smoke case before marking ROCm complete. + +## Build and packaging changes + +- Remove direct cuBLAS/rocBLAS includes, calls, and handle management from the Torch grouped-GEMM path. On the PR branch, also remove the now-unused borrowed-handle adapters. +- Remove `CUDA::cublas`, the grouped-GEMM-only `find_package(rocblas)` / `${HIP_BLAS_LIB}`, and the JIT loader's explicit `-lcublas` link flag after confirming there are no remaining consumers. Torch still brings its own BLAS dependencies. +- Resolve the generated BMM symbol through the installed Torch GPU library: `libtorch_cuda` for CUDA and `libtorch_hip` for ROCm. Common tensor/guard C APIs are supplied by Torch's common library. +- OEQ's wheel build currently downloads CPU LibTorch headers/libraries and creates GPU link stubs from `extension/stubs/stream.cpp`. Extend those build-only stubs with the exact generated BMM declaration/definition, sharing the header so signature drift causes a build error. If retaining the prototype's GPU-specific guards, add their GPU exports too; the proposed 2.10 generic guard avoids that need. Never package or execute the fake implementations. Check that installed extensions resolve against the real Torch libraries. +- Apply changes to both the Python extension and the AOTI shared library targets, and to source/JIT builds. The existing HIP CMake target is named `torch_stable_hip` while import/install expectations use `oeq_stable_hip`; resolve that naming mismatch as part of making the HIP artifact load correctly. +- Keep the existing CUDA/HIP runtime and runtime-compiler dependencies needed by OEQ-generated kernels. Removing the BLAS dependency does not make the entire extension independent of the GPU platform. + +The stable-wheel baseline remains 2.10. Confirm the supported source/JIT baseline before sharing every C helper with that build: the prototype's helpers were exercised on 2.7, whereas the proposed generic guard needs its own baseline check. If older JIT support must be preserved, a JIT-only adapter using that installed Torch's ordinary BMM-out/device guard is acceptable; it still eliminates direct vendor BLAS calls and does not enter the stable wheel. + +## JAX scope + +Neither the inspected PR snapshot nor the current JAX extension calls `group_gemm_blas`; its CMake target links the runtime, driver, and NVRTC, without BLAS. This replacement therefore remains entirely within Torch. + +JAX's public FFI has buffers and a GPU stream, but no equivalent BMM shim or public BLAS-handle getter. If a future JAX operation needs this GEMM, express it in the compiled JAX graph where shapes permit, or provide a separate FFI implementation whose BLAS handles and calls come from the same library. [JAX GPU FFI](https://docs.jax.dev/en/latest/ffi.html#ffi-calls-on-a-gpu), [public XLA FFI API](https://github.com/openxla/xla/blob/main/xla/ffi/api/c_api.h#L757). + +## Existing prototype and validation + +A standalone C++ prototype already implements both layout branches using only Torch C APIs. Its single float32 forward smoke case passed on an H100 PCIe with Torch 2.7.0 / CUDA 12.8: counts `[2,0,5,1]`, batch 3, m 4, k 5, a nondefault stream, and maximum absolute error `4.37e-7` against an independent CPU float64 reference. + +Earlier ctypes calls to the same C shim passed 40 cases spanning both dtypes and modes. Four selected profiles showed one GEMM kernel with no observed copy or GPU allocation inside the BMM call. These measurements do not cover the CPU cost of temporary tensor metadata or establish end-to-end performance. + +The integrated extension, its autograd registration, the final wheel link setup, and ROCm execution remain to be validated. No additional GPU experiments were run for this design. + +## Acceptance criteria + +- Both grouped-GEMM modes use Torch BMM, with no direct vendor BLAS symbols or handles in OEQ's Torch extension. +- Existing layouts, empty groups, float32/float64 behavior, operator schema, fake implementation, and gradients are preserved. +- Device guarding and a nondefault current stream work; the caller's current device is restored. +- Stable CUDA and HIP artifacts resolve the official C symbols in the real installed Torch GPU library; link stubs do not ship. Audit direct dynamic dependencies to confirm BLAS is now Torch's responsibility. +- Run focused correctness checks for both modes/dtypes and one relevant autograd case during integration, plus one ROCm smoke case when AMD hardware is available. Keep this bounded; no exhaustive vendor-version matrix is needed to establish the prototype. +- Check representative overhead before claiming performance parity. Graph capture and broader multi-device coverage should be checked where required by the supported operator contract. + +## Prototype code + +The following is the existing standalone CUDA prototype, not the final production registration or generic-device-guard adaptation. The caller supplies validated buffers and zero-initializes outputs where empty groups need to remain zero. + +```cpp +// Prototype for the PyTorch CUDA backend. Uses only PyTorch's stable C ABI. +#include + +#include +#include +#include +#include +#include +#include + +namespace { + +void check(AOTITorchError status) { + if (status != AOTI_TORCH_SUCCESS) + throw std::runtime_error("PyTorch C shim failed"); +} + +using Tensor = std::unique_ptr< + std::remove_pointer_t, + decltype(&aoti_torch_delete_tensor_object)>; +using DeviceGuard = std::unique_ptr< + std::remove_pointer_t, + decltype(&aoti_torch_delete_cuda_guard)>; + +// Only tensor metadata is created; the caller retains ownership of the buffer. +Tensor view(void* data, std::array sizes, + std::array strides, int32_t dtype, int32_t device) { + AtenTensorHandle tensor = nullptr; + check(aoti_torch_create_tensor_from_blob_v2( + data, 3, sizes.data(), strides.data(), 0, dtype, + aoti_torch_device_type_cuda(), device, &tensor, + aoti_torch_layout_strided(), nullptr, 0)); + return Tensor(tensor, aoti_torch_delete_tensor_object); +} + +} // namespace + +template +void group_gemm_cshim( + T* A, T* B, T* C, const int64_t* ragged_counts, int num_groups, + int64_t batch, int64_t m, int64_t k, int ragged_inner, int32_t device) { + static_assert(std::is_same_v || std::is_same_v); + const int32_t dtype = std::is_same_v + ? aoti_torch_dtype_float32() : aoti_torch_dtype_float64(); + + CUDAGuardHandle raw_guard = nullptr; + check(aoti_torch_create_cuda_guard(device, &raw_guard)); + DeviceGuard guard(raw_guard, aoti_torch_delete_cuda_guard); + + int64_t offset = 0; + for (int i = 0; i < num_groups; ++i) { + const int64_t n = ragged_counts[i]; + if (n == 0) continue; // Preserve the original empty-group behavior. + + if (ragged_inner == 0) { + // [batch, n, k] @ [batch, k, m] -> [batch, n, m] + auto input = view(B + offset * batch * k, + {batch, n, k}, {k, batch * k, 1}, dtype, device); + auto weight = view(A + i * batch * m * k, + {batch, k, m}, {m * k, 1, k}, dtype, device); + auto output = view(C + offset * batch * m, + {batch, n, m}, {m, batch * m, 1}, dtype, device); + check(aoti_torch_cuda_bmm_out(output.get(), input.get(), weight.get())); + } else { + // [batch, m, n] @ [batch, n, k] -> [batch, m, k] + auto left = view(A + offset * batch * m, + {batch, m, n}, {m, 1, batch * m}, dtype, device); + auto right = view(B + offset * batch * k, + {batch, n, k}, {k, batch * k, 1}, dtype, device); + auto output = view(C + i * batch * m * k, + {batch, m, k}, {m * k, k, 1}, dtype, device); + check(aoti_torch_cuda_bmm_out(output.get(), left.get(), right.get())); + } + offset += n; + } +} + +// Small ctypes entry point for the smoke test, not production registration code. +extern "C" int oeq_group_gemm_cshim( + int dtype, void* A, void* B, void* C, const int64_t* counts, + int groups, int64_t batch, int64_t m, int64_t k, int inner, int32_t device) { + try { + if (dtype == 0) { + group_gemm_cshim(static_cast(A), static_cast(B), + static_cast(C), counts, groups, + batch, m, k, inner, device); + } else if (dtype == 1) { + group_gemm_cshim(static_cast(A), static_cast(B), + static_cast(C), counts, groups, + batch, m, k, inner, device); + } else { + throw std::runtime_error("Expected float32 or float64"); + } + return 0; + } catch (const std::exception& error) { + std::fprintf(stderr, "%s\n", error.what()); + return 1; + } +} +``` + +
+Reproduce the single CUDA smoke case + +Save the C++ above as `group_gemm.cpp` and this script beside it as `smoke_test.py`, then run `python3 smoke_test.py` with a CUDA-enabled Torch installation and a C++ compiler. The script builds only the standalone prototype and executes one configuration. + +```python +"""Compile the C++ prototype and run one float32 CUDA correctness case.""" +import ctypes +from pathlib import Path +import subprocess + +import torch + + +def main(): + root = Path(__file__).resolve().parent + torch_root = Path(torch.__file__).resolve().parent + library = root / "group_gemm.so" + subprocess.run( + [ + "g++", "-std=c++17", "-O2", "-shared", "-fPIC", "-DUSE_CUDA", + f"-I{torch_root / 'include'}", str(root / "group_gemm.cpp"), + f"-L{torch_root / 'lib'}", f"-Wl,-rpath,{torch_root / 'lib'}", + "-Wl,--no-undefined", "-ltorch_cuda", "-ltorch_cpu", + "-o", str(library), + ], + check=True, + ) + lib = ctypes.CDLL(str(library)) + call = lib.oeq_group_gemm_cshim + i64 = ctypes.c_int64 + call.argtypes = [ + ctypes.c_int, ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p, + ctypes.POINTER(i64), ctypes.c_int, i64, i64, i64, + ctypes.c_int, ctypes.c_int32, + ] + call.restype = ctypes.c_int + + # One case includes differently sized groups, an empty group, interleaved + # batches, and a nondefault stream. The reference runs on the CPU. + counts = [2, 0, 5, 1] + batch, m, k = 3, 4, 5 + torch.manual_seed(123) + weights_cpu = torch.randn(len(counts), batch, m, k) + input_cpu = torch.randn(sum(counts), batch, k) + expected = torch.empty(sum(counts), batch, m, dtype=torch.float64) + offset = 0 + for i, n in enumerate(counts): + expected[offset:offset + n] = torch.einsum( + "bmk,nbk->nbm", weights_cpu[i].double(), + input_cpu[offset:offset + n].double(), + ) + offset += n + + torch.cuda.set_device(0) + stream = torch.cuda.Stream() + with torch.cuda.stream(stream): + weights = weights_cpu.cuda() + inputs = input_cpu.cuda() + output = torch.full(expected.shape, float("nan"), device="cuda") + status = call( + 0, weights.data_ptr(), inputs.data_ptr(), output.data_ptr(), + (i64 * len(counts))(*counts), len(counts), batch, m, k, 0, 0, + ) + assert status == 0, f"C shim prototype returned {status}" + stream.synchronize() + actual = output.cpu().double() + torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-5) + print( + f"PASS: one float32 case, counts={counts}, batch={batch}, m={m}, k={k}; " + f"max absolute error={(actual - expected).abs().max().item():.3g}; " + f"GPU={torch.cuda.get_device_name(0)}, torch={torch.__version__}" + ) + + +if __name__ == "__main__": + main() +``` + +
diff --git a/sandbox/cshim_group_gemm/group_gemm.cpp b/sandbox/cshim_group_gemm/group_gemm.cpp new file mode 100644 index 00000000..1f7462a7 --- /dev/null +++ b/sandbox/cshim_group_gemm/group_gemm.cpp @@ -0,0 +1,99 @@ +// Prototype for the PyTorch CUDA backend. Uses only PyTorch's stable C ABI. +#include + +#include +#include +#include +#include +#include +#include + +namespace { + +void check(AOTITorchError status) { + if (status != AOTI_TORCH_SUCCESS) + throw std::runtime_error("PyTorch C shim failed"); +} + +using Tensor = std::unique_ptr< + std::remove_pointer_t, + decltype(&aoti_torch_delete_tensor_object)>; +using DeviceGuard = std::unique_ptr< + std::remove_pointer_t, + decltype(&aoti_torch_delete_cuda_guard)>; + +// Only tensor metadata is created; the caller retains ownership of the buffer. +Tensor view(void* data, std::array sizes, + std::array strides, int32_t dtype, int32_t device) { + AtenTensorHandle tensor = nullptr; + check(aoti_torch_create_tensor_from_blob_v2( + data, 3, sizes.data(), strides.data(), 0, dtype, + aoti_torch_device_type_cuda(), device, &tensor, + aoti_torch_layout_strided(), nullptr, 0)); + return Tensor(tensor, aoti_torch_delete_tensor_object); +} + +} // namespace + +template +void group_gemm_cshim( + T* A, T* B, T* C, const int64_t* ragged_counts, int num_groups, + int64_t batch, int64_t m, int64_t k, int ragged_inner, int32_t device) { + static_assert(std::is_same_v || std::is_same_v); + const int32_t dtype = std::is_same_v + ? aoti_torch_dtype_float32() : aoti_torch_dtype_float64(); + + CUDAGuardHandle raw_guard = nullptr; + check(aoti_torch_create_cuda_guard(device, &raw_guard)); + DeviceGuard guard(raw_guard, aoti_torch_delete_cuda_guard); + + int64_t offset = 0; + for (int i = 0; i < num_groups; ++i) { + const int64_t n = ragged_counts[i]; + if (n == 0) continue; // Preserve the original empty-group behavior. + + if (ragged_inner == 0) { + // [batch, n, k] @ [batch, k, m] -> [batch, n, m] + auto input = view(B + offset * batch * k, + {batch, n, k}, {k, batch * k, 1}, dtype, device); + auto weight = view(A + i * batch * m * k, + {batch, k, m}, {m * k, 1, k}, dtype, device); + auto output = view(C + offset * batch * m, + {batch, n, m}, {m, batch * m, 1}, dtype, device); + check(aoti_torch_cuda_bmm_out(output.get(), input.get(), weight.get())); + } else { + // [batch, m, n] @ [batch, n, k] -> [batch, m, k] + auto left = view(A + offset * batch * m, + {batch, m, n}, {m, 1, batch * m}, dtype, device); + auto right = view(B + offset * batch * k, + {batch, n, k}, {k, batch * k, 1}, dtype, device); + auto output = view(C + i * batch * m * k, + {batch, m, k}, {m * k, k, 1}, dtype, device); + check(aoti_torch_cuda_bmm_out(output.get(), left.get(), right.get())); + } + offset += n; + } +} + +// Small ctypes entry point for the smoke test, not production registration code. +extern "C" int oeq_group_gemm_cshim( + int dtype, void* A, void* B, void* C, const int64_t* counts, + int groups, int64_t batch, int64_t m, int64_t k, int inner, int32_t device) { + try { + if (dtype == 0) { + group_gemm_cshim(static_cast(A), static_cast(B), + static_cast(C), counts, groups, + batch, m, k, inner, device); + } else if (dtype == 1) { + group_gemm_cshim(static_cast(A), static_cast(B), + static_cast(C), counts, groups, + batch, m, k, inner, device); + } else { + throw std::runtime_error("Expected float32 or float64"); + } + return 0; + } catch (const std::exception& error) { + std::fprintf(stderr, "%s\n", error.what()); + return 1; + } +} diff --git a/sandbox/cshim_group_gemm/smoke_test.py b/sandbox/cshim_group_gemm/smoke_test.py new file mode 100644 index 00000000..931d9698 --- /dev/null +++ b/sandbox/cshim_group_gemm/smoke_test.py @@ -0,0 +1,71 @@ +"""Compile the C++ prototype and run one float32 CUDA correctness case.""" +import ctypes +from pathlib import Path +import subprocess + +import torch + + +def main(): + root = Path(__file__).resolve().parent + torch_root = Path(torch.__file__).resolve().parent + library = root / "group_gemm.so" + subprocess.run( + [ + "g++", "-std=c++17", "-O2", "-shared", "-fPIC", "-DUSE_CUDA", + f"-I{torch_root / 'include'}", str(root / "group_gemm.cpp"), + f"-L{torch_root / 'lib'}", f"-Wl,-rpath,{torch_root / 'lib'}", + "-Wl,--no-undefined", "-ltorch_cuda", "-ltorch_cpu", + "-o", str(library), + ], + check=True, + ) + lib = ctypes.CDLL(str(library)) + call = lib.oeq_group_gemm_cshim + i64 = ctypes.c_int64 + call.argtypes = [ + ctypes.c_int, ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p, + ctypes.POINTER(i64), ctypes.c_int, i64, i64, i64, + ctypes.c_int, ctypes.c_int32, + ] + call.restype = ctypes.c_int + + # One case includes differently sized groups, an empty group, interleaved + # batches, and a nondefault stream. The reference runs on the CPU. + counts = [2, 0, 5, 1] + batch, m, k = 3, 4, 5 + torch.manual_seed(123) + weights_cpu = torch.randn(len(counts), batch, m, k) + input_cpu = torch.randn(sum(counts), batch, k) + expected = torch.empty(sum(counts), batch, m, dtype=torch.float64) + offset = 0 + for i, n in enumerate(counts): + expected[offset:offset + n] = torch.einsum( + "bmk,nbk->nbm", weights_cpu[i].double(), + input_cpu[offset:offset + n].double(), + ) + offset += n + + torch.cuda.set_device(0) + stream = torch.cuda.Stream() + with torch.cuda.stream(stream): + weights = weights_cpu.cuda() + inputs = input_cpu.cuda() + output = torch.full(expected.shape, float("nan"), device="cuda") + status = call( + 0, weights.data_ptr(), inputs.data_ptr(), output.data_ptr(), + (i64 * len(counts))(*counts), len(counts), batch, m, k, 0, 0, + ) + assert status == 0, f"C shim prototype returned {status}" + stream.synchronize() + actual = output.cpu().double() + torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-5) + print( + f"PASS: one float32 case, counts={counts}, batch={batch}, m={m}, k={k}; " + f"max absolute error={(actual - expected).abs().max().item():.3g}; " + f"GPU={torch.cuda.get_device_name(0)}, torch={torch.__version__}" + ) + + +if __name__ == "__main__": + main() From 28a6fb6f77dd13f21f5399a3650a4bdd2ed8eccc Mon Sep 17 00:00:00 2001 From: asglover <140220574+asglover@users.noreply.github.com> Date: Sun, 13 Sep 2026 00:02:12 -0700 Subject: [PATCH 02/38] Route Torch grouped GEMM through the stable BMM C shim --- CHANGELOG.md | 3 + openequivariance/CMakeLists.txt | 15 +- .../_torch/extlib/__init__.py | 8 +- .../openequivariance/extension/group_mm.hpp | 178 ++++------ .../extension/libtorch_tp_jit.cpp | 14 +- .../extension/libtorch_tp_jit_stable.cpp | 10 + .../extension/stubs/stream.cpp | 11 +- .../openequivariance/extension/torch_core.hpp | 42 ++- sandbox/cshim_group_gemm/DESIGN.md | 311 ++++-------------- tests/group_gemm_test.py | 173 ++++++++++ tests/import_test.py | 28 ++ 11 files changed, 389 insertions(+), 404 deletions(-) create mode 100644 tests/group_gemm_test.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 3a6f9d04..39a817d9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,9 @@ **Added**: - Beta support for Intel XPUs. +- Route Torch grouped GEMM through the stable BMM C shim on CUDA and ROCm, + using Torch's current device stream and precision settings. Remove OEQ's + direct cuBLAS and rocBLAS dependencies. ### v0.7.0 (2026-09-10) **Added**: diff --git a/openequivariance/CMakeLists.txt b/openequivariance/CMakeLists.txt index 3f3ad3dc..7ed586db 100644 --- a/openequivariance/CMakeLists.txt +++ b/openequivariance/CMakeLists.txt @@ -67,6 +67,7 @@ function(add_stable_extension target_name backend_define link_libraries) # Enforce CXX11 ABI to match LibTorch target_compile_definitions(${target_name} PRIVATE ${backend_define}=1 + TORCH_TARGET_VERSION=0x020a000000000000ULL _GLIBCXX_USE_CXX11_ABI=1 INCLUDE_NB_EXTENSION ) @@ -97,6 +98,7 @@ function(add_stable_extension target_name backend_define link_libraries) target_compile_definitions(${aoti_target_name} PRIVATE ${backend_define}=1 + TORCH_TARGET_VERSION=0x020a000000000000ULL _GLIBCXX_USE_CXX11_ABI=1 ) @@ -117,7 +119,6 @@ endfunction() find_package(CUDAToolkit QUIET) find_package(hip QUIET) -find_package(rocblas QUIET) if(CUDAToolkit_FOUND) message(STATUS "Building stable extension with CUDA backend.") @@ -138,7 +139,6 @@ if(CUDAToolkit_FOUND) CUDA::cudart CUDA::cuda_driver CUDA::nvrtc - CUDA::cublas cuda_stub_lib ) add_stable_extension(oeq_stable_cuda CUDA_BACKEND "${CUDA_LINK_LIBS}") @@ -146,6 +146,7 @@ endif() if(hip_FOUND) message(STATUS "Building stable extension with HIP backend.") + find_package(hiprtc REQUIRED) add_library(hip_stub_lib SHARED ${EXT_DIR}/stubs/stream.cpp) @@ -159,18 +160,12 @@ if(hip_FOUND) CXX_STANDARD 17 ) - if(TARGET roc::rocblas) - set(HIP_BLAS_LIB roc::rocblas) - else() - set(HIP_BLAS_LIB rocblas) - endif() - set(HIP_LINK_LIBS hip_stub_lib hip::host - ${HIP_BLAS_LIB} + hiprtc::hiprtc ) - add_stable_extension(torch_stable_hip HIP_BACKEND "${HIP_LINK_LIBS}") + add_stable_extension(oeq_stable_hip HIP_BACKEND "${HIP_LINK_LIBS}") endif() set(OEQ_SYCL_FOUND FALSE) diff --git a/openequivariance/openequivariance/_torch/extlib/__init__.py b/openequivariance/openequivariance/_torch/extlib/__init__.py index c8599df0..0861f2e0 100644 --- a/openequivariance/openequivariance/_torch/extlib/__init__.py +++ b/openequivariance/openequivariance/_torch/extlib/__init__.py @@ -135,7 +135,7 @@ def load_jit_extension(): ], ) if BACKEND == "cuda": - extra_link_args.extend(["-lcuda", "-lcudart", "-lnvrtc", "-lcublas"]) + extra_link_args.extend(["-lcuda", "-lcudart", "-lnvrtc", "-ltorch_cuda"]) try: torch_libs, cuda_libs = library_paths("cuda") @@ -148,8 +148,10 @@ def load_jit_extension(): extra_cflags.append("-DCUDA_BACKEND") elif BACKEND == "hip": - torch_libs = library_paths("cuda")[0] - extra_link_args.append("-Wl,-rpath," + torch_libs) + hip_lib_dirs = library_paths("cuda") + extra_link_args.append("-Wl,-rpath," + hip_lib_dirs[0]) + extra_link_args.extend("-L" + directory for directory in hip_lib_dirs) + extra_link_args.extend(["-ltorch_hip", "-lhiprtc"]) extra_cflags.append("-DHIP_BACKEND") elif BACKEND == "sycl": import shutil diff --git a/openequivariance/openequivariance/extension/group_mm.hpp b/openequivariance/openequivariance/extension/group_mm.hpp index 19249c39..247b6b3d 100644 --- a/openequivariance/openequivariance/extension/group_mm.hpp +++ b/openequivariance/openequivariance/extension/group_mm.hpp @@ -1,141 +1,73 @@ #pragma once +#include #include #include #include +#include -#ifdef CUDA_BACKEND - #include "cublas_v2.h" - #include +#include - struct BlasHandle { - cublasHandle_t handle; - BlasHandle() { - if (cublasCreate(&handle) != CUBLAS_STATUS_SUCCESS) - throw std::logic_error("CUBLAS initialization failed"); - } - ~BlasHandle() { cublasDestroy(handle); } - }; -#elif defined(HIP_BACKEND) - #include "rocblas/rocblas.h" - #include - - 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); } - }; -#endif +namespace oeq { -inline BlasHandle& get_blas_handle() { - static BlasHandle handle; - return handle; +inline void check_group_mm_shim(AOTITorchError status) { + if (status != AOTI_TORCH_SUCCESS) + throw std::runtime_error("group_gemm: PyTorch C shim failed"); } -template -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) { +using GroupMMTensor = std::unique_ptr< + std::remove_pointer_t, + decltype(&aoti_torch_delete_tensor_object)>; + +inline GroupMMTensor group_mm_view( + AtenTensorHandle tensor, std::array sizes, + std::array strides, int64_t offset) { + AtenTensorHandle view = nullptr; + check_group_mm_shim(aoti_torch__reinterpret_tensor( + tensor, 3, sizes.data(), strides.data(), offset, &view)); + return GroupMMTensor(view, aoti_torch_delete_tensor_object); +} - auto& blas = get_blas_handle(); - T alpha = 1.0, beta = 0.0; - T* A_base = reinterpret_cast(A_raw); - T* B_base = reinterpret_cast(B_raw); - T* C_base = reinterpret_cast(C_raw); +inline void group_gemm_torch( + AtenTensorHandle A, AtenTensorHandle B, AtenTensorHandle C, + const int64_t* ragged_counts, int64_t num_W, int64_t batch_size, + int64_t m, int64_t k, int64_t ragged_inner) { + if (batch_size == 0 || m == 0 || k == 0) + return; - int64_t ragged_offset = 0; - for (int i = 0; i < num_W; i++) { - int M, K, N, lda, ldb, ldc, strideA, strideB, strideC; - T *A, *B, *C; -#ifdef CUDA_BACKEND - cublasOperation_t transa, transb; -#elif defined(HIP_BACKEND) - rocblas_operation transa, transb; -#endif + int64_t offset = 0; + for (int64_t i = 0; i < num_W; ++i) { + const int64_t n = ragged_counts[i]; + if (n == 0) + continue; if (ragged_inner == 0) { - M = m; K = k; N = static_cast(ragged_counts[i]); - A = A_base + (m * k * batch_size * i); - lda = k; strideA = M * K; - B = B_base + (k * batch_size * ragged_offset); - ldb = K * batch_size; strideB = K; - C = C_base + (m * batch_size * ragged_offset); - ldc = M * batch_size; strideC = M; -#ifdef CUDA_BACKEND - transa = CUBLAS_OP_T; transb = CUBLAS_OP_N; -#elif defined(HIP_BACKEND) - transa = rocblas_operation_transpose; transb = rocblas_operation_none; -#endif + auto input = group_mm_view(B, + {batch_size, n, k}, {k, batch_size * k, 1}, + offset * batch_size * k); + auto weight = group_mm_view(A, + {batch_size, k, m}, {m * k, 1, k}, + i * batch_size * m * k); + auto output = group_mm_view(C, + {batch_size, n, m}, {m, batch_size * m, 1}, + offset * batch_size * m); + check_group_mm_shim(aoti_torch_cuda_bmm_out( + output.get(), input.get(), weight.get())); } else { - M = k; K = static_cast(ragged_counts[i]); N = m; - A = B_base + (k * batch_size * ragged_offset); - lda = k * batch_size; strideA = M; - B = A_base + (m * batch_size * ragged_offset); - ldb = m * batch_size; strideB = N; - C = C_base + (m * k * batch_size * i); - ldc = k; strideC = M * N; -#ifdef CUDA_BACKEND - transa = CUBLAS_OP_N; transb = CUBLAS_OP_T; -#elif defined(HIP_BACKEND) - transa = rocblas_operation_none; transb = rocblas_operation_transpose; -#endif - } - ragged_offset += ragged_counts[i]; - - if (ragged_counts[i] > 0) { -#ifdef CUDA_BACKEND - cublasStatus_t stat; - if (std::is_same::value) { - stat = cublasSgemmStridedBatched(blas.handle, - transa, transb, M, N, K, - reinterpret_cast(&alpha), - reinterpret_cast(A), lda, strideA, - reinterpret_cast(B), ldb, strideB, - reinterpret_cast(&beta), - reinterpret_cast(C), ldc, strideC, - batch_size); - } else if (std::is_same::value) { - stat = cublasDgemmStridedBatched(blas.handle, - transa, transb, M, N, K, - reinterpret_cast(&alpha), - reinterpret_cast(A), lda, strideA, - reinterpret_cast(B), ldb, strideB, - reinterpret_cast(&beta), - reinterpret_cast(C), ldc, strideC, - batch_size); - } else { - throw std::logic_error("Unsupported datatype for grouped GEMM!"); - } - if (stat != CUBLAS_STATUS_SUCCESS) - throw std::logic_error("Grouped GEMM failed!"); -#elif defined(HIP_BACKEND) - rocblas_status stat; - if (std::is_same::value) { - stat = rocblas_sgemm_strided_batched(blas.handle, - transa, transb, M, N, K, - reinterpret_cast(&alpha), - reinterpret_cast(A), lda, strideA, - reinterpret_cast(B), ldb, strideB, - reinterpret_cast(&beta), - reinterpret_cast(C), ldc, strideC, - batch_size); - } else if (std::is_same::value) { - stat = rocblas_dgemm_strided_batched(blas.handle, - transa, transb, M, N, K, - reinterpret_cast(&alpha), - reinterpret_cast(A), lda, strideA, - reinterpret_cast(B), ldb, strideB, - reinterpret_cast(&beta), - reinterpret_cast(C), ldc, strideC, - batch_size); - } else { - throw std::logic_error("Unsupported datatype for grouped GEMM!"); - } - if (stat != rocblas_status_success) - throw std::logic_error("Grouped GEMM failed!"); -#endif + auto left = group_mm_view(A, + {batch_size, m, n}, {m, 1, batch_size * m}, + offset * batch_size * m); + auto right = group_mm_view(B, + {batch_size, n, k}, {k, batch_size * k, 1}, + offset * batch_size * k); + auto output = group_mm_view(C, + {batch_size, m, k}, {m * k, k, 1}, + i * batch_size * m * k); + check_group_mm_shim(aoti_torch_cuda_bmm_out( + output.get(), left.get(), right.get())); } + offset += n; } } + +} diff --git a/openequivariance/openequivariance/extension/libtorch_tp_jit.cpp b/openequivariance/openequivariance/extension/libtorch_tp_jit.cpp index c7f9bdb6..08f98d79 100644 --- a/openequivariance/openequivariance/extension/libtorch_tp_jit.cpp +++ b/openequivariance/openequivariance/extension/libtorch_tp_jit.cpp @@ -2,7 +2,7 @@ #include #ifdef CUDA_BACKEND - #include + #include #endif #ifdef HIP_BACKEND @@ -14,9 +14,11 @@ #endif #include +#include #include #include #include +#include #include using Tensor = torch::Tensor; @@ -33,6 +35,16 @@ constexpr Dtype kByte = torch::kByte; #define REGISTER_LIBRARY_IMPL TORCH_LIBRARY_IMPL #define REGISTER_LIBRARY TORCH_LIBRARY +class TensorDeviceGuard { + c10::DeviceGuard guard; +public: + explicit TensorDeviceGuard(const Tensor& tensor) : guard(tensor.device()) {} +}; + +AtenTensorHandle tensor_handle(Tensor& tensor) { + return torch::aot_inductor::tensor_pointer_to_tensor_handle(&tensor); +} + #include "torch_core.hpp" Tensor tensor_to_cpu_contiguous(const Tensor &tensor) { diff --git a/openequivariance/openequivariance/extension/libtorch_tp_jit_stable.cpp b/openequivariance/openequivariance/extension/libtorch_tp_jit_stable.cpp index dc8ec85b..492931d3 100644 --- a/openequivariance/openequivariance/extension/libtorch_tp_jit_stable.cpp +++ b/openequivariance/openequivariance/extension/libtorch_tp_jit_stable.cpp @@ -32,6 +32,16 @@ constexpr Dtype kByte = torch::headeronly::ScalarType::Byte; #define REGISTER_LIBRARY_IMPL STABLE_TORCH_LIBRARY_IMPL #define REGISTER_LIBRARY STABLE_TORCH_LIBRARY +class TensorDeviceGuard { + torch::stable::accelerator::DeviceGuard guard; +public: + explicit TensorDeviceGuard(const Tensor& tensor) : guard(tensor.get_device_index()) {} +}; + +AtenTensorHandle tensor_handle(Tensor& tensor) { + return tensor.get(); +} + #include "torch_core.hpp" Tensor tensor_to_cpu_contiguous(const Tensor &tensor) { diff --git a/openequivariance/openequivariance/extension/stubs/stream.cpp b/openequivariance/openequivariance/extension/stubs/stream.cpp index 16280c13..de45dc06 100644 --- a/openequivariance/openequivariance/extension/stubs/stream.cpp +++ b/openequivariance/openequivariance/extension/stubs/stream.cpp @@ -1,5 +1,7 @@ +#define USE_CUDA + #include -#include +#include extern "C" { #ifdef SYCL_BACKEND @@ -9,7 +11,12 @@ extern "C" { } #else AOTITorchError aoti_torch_get_current_cuda_stream(int32_t device_index, void** ret_stream) { - return 0; + return AOTI_TORCH_FAILURE; + } + + AOTITorchError aoti_torch_cuda_bmm_out( + AtenTensorHandle out, AtenTensorHandle self, AtenTensorHandle mat2) { + return AOTI_TORCH_FAILURE; } #endif } diff --git a/openequivariance/openequivariance/extension/torch_core.hpp b/openequivariance/openequivariance/extension/torch_core.hpp index ac1daa54..f744d96a 100644 --- a/openequivariance/openequivariance/extension/torch_core.hpp +++ b/openequivariance/openequivariance/extension/torch_core.hpp @@ -653,13 +653,40 @@ inline tuple jit_conv_double_backward( inline Tensor group_gemm( Tensor A, Tensor B, Tensor ragged_counts, int64_t num_W, int64_t batch_size, int64_t m, int64_t k, int64_t ragged_inner) { + TCHECK(A.is_cuda() && B.is_cuda(), "group_gemm: A and B must be GPU tensors"); + TCHECK(A.get_device() == B.get_device(), "group_gemm: A and B must be on the same device"); TCHECK(A.scalar_type() == B.scalar_type(), "group_gemm: A and B must have the same dtype"); + TCHECK(A.scalar_type() == kFloat || A.scalar_type() == kDouble, + "group_gemm: unsupported dtype, expected float32 or float64"); + TCHECK(ragged_counts.is_cpu(), "group_gemm: ragged_counts must be on the CPU"); TCHECK(ragged_counts.scalar_type() == kLong, "group_gemm: ragged_counts must be int64"); + TCHECK(num_W >= 0 && batch_size >= 0 && m >= 0 && k >= 0, + "group_gemm: dimensions must be nonnegative"); + TCHECK(ragged_inner == 0 || ragged_inner == 1, "group_gemm: ragged_inner must be 0 or 1"); + TCHECK(ragged_counts.dim() == 1 && ragged_counts.size(0) == num_W, + "group_gemm: ragged_counts must contain num_W entries"); + TCHECK(B.dim() == 3, "group_gemm: B must have shape [rows, batch_size, k]"); + const int64_t rows = B.size(0); + check_tensor(B, {rows, batch_size, k}, A.scalar_type(), "group_gemm B"); + if (ragged_inner == 0) + check_tensor(A, {num_W, batch_size, m, k}, A.scalar_type(), "group_gemm A"); + else + check_tensor(A, {rows, batch_size, m}, A.scalar_type(), "group_gemm A"); + + Tensor rc_c = tensor_contiguous(ragged_counts); + const auto* rc_ptr = static_cast(data_ptr(rc_c)); + int64_t row_count = 0; + for (int64_t i = 0; i < num_W; ++i) { + TCHECK(rc_ptr[i] >= 0 && rc_ptr[i] <= rows - row_count, + "group_gemm: ragged_counts must be nonnegative and sum to the number of rows"); + row_count += rc_ptr[i]; + } + TCHECK(row_count == rows, + "group_gemm: ragged_counts must sum to the number of rows"); + TensorDeviceGuard device_guard(A); Tensor A_c = tensor_contiguous(A); Tensor B_c = tensor_contiguous(B); - Tensor rc_c = tensor_contiguous(ragged_counts); - int64_t* rc_ptr = reinterpret_cast(data_ptr(rc_c)); Tensor C; if (ragged_inner == 0) { @@ -669,15 +696,8 @@ inline Tensor group_gemm( C = tensor_zeros_like(A, make_sizes({num_W, batch_size, m, k})); } - if (A.scalar_type() == kFloat) { - group_gemm_blas(data_ptr(A_c), data_ptr(B_c), data_ptr(C), rc_ptr, - (int)num_W, (int)batch_size, (int)m, (int)k, (int)ragged_inner); - } else if (A.scalar_type() == kDouble) { - group_gemm_blas(data_ptr(A_c), data_ptr(B_c), data_ptr(C), rc_ptr, - (int)num_W, (int)batch_size, (int)m, (int)k, (int)ragged_inner); - } else { - throw std::logic_error("group_gemm: unsupported dtype, expected float32 or float64"); - } + oeq::group_gemm_torch(tensor_handle(A_c), tensor_handle(B_c), tensor_handle(C), + rc_ptr, num_W, batch_size, m, k, ragged_inner); return C; } diff --git a/sandbox/cshim_group_gemm/DESIGN.md b/sandbox/cshim_group_gemm/DESIGN.md index 3b919e07..9626260e 100644 --- a/sandbox/cshim_group_gemm/DESIGN.md +++ b/sandbox/cshim_group_gemm/DESIGN.md @@ -1,292 +1,95 @@ # Move Torch cuBLAS calls to the stable C shim -Replace OEQ's direct CUDA and ROCm BLAS calls for `libtorch_tp_jit::group_gemm` with PyTorch's `aoti_torch_cuda_bmm_out` C API. Keep the existing grouped operation, tensor layouts, output allocation, and custom autograd rules. PyTorch will own the underlying BLAS handles, streams, library calls, and backend selection. +The production Torch extension now implements `libtorch_tp_jit::group_gemm` through `aoti_torch_cuda_bmm_out` on CUDA and ROCm. Both the stable extension and the source/JIT extension use this C entry point. The operator schema, fake implementation, output layouts, and custom backward formulas remain compatible. -Related: [PR #206](https://github.com/PASSIONLab/OpenEquivariance/pull/206). +GPU execution of this implementation is pending. GPU access is paused at the user's request; the earlier prototype results below are not results for this implementation. -## Problem and scope +## Implementation -OEQ's grouped GEMM is a host loop over nonempty ragged groups, issuing one strided batched GEMM per group. It does not require a vendor's heterogeneous grouped-GEMM API. The only direct BLAS computation calls found are float32/float64 strided batched GEMMs in `extension/group_mm.hpp`, called from the Torch backend in `extension/torch_core.hpp`. +The implementation lives in [group_mm.hpp](../../openequivariance/openequivariance/extension/group_mm.hpp), with input validation and output allocation in [torch_core.hpp](../../openequivariance/openequivariance/extension/torch_core.hpp). The algorithm is one BMM-out invocation per nonempty ragged group, using the same interleaved layout as the previous strided batched GEMMs. It does not need a vendor's heterogeneous grouped-GEMM API. -Using PyTorch's stable ABI to obtain a cuBLAS handle does not make that handle interchangeable with a separately loaded cuBLAS implementation. Earlier H100 experiments passed with matching handle/function owners, while some foreign-owner combinations returned incorrect streams, failed, or crashed. Successful version queries did not establish handle compatibility. Moving the GEMM itself behind PyTorch's C ABI removes this borrowed-handle boundary. +Inputs are made contiguous by the registered operator, then passed as borrowed `AtenTensorHandle` values to the shared helper. The stable extension gets these handles from `torch::stable::Tensor::get()`. The source/JIT extension uses Torch's `tensor_pointer_to_tensor_handle` utility to borrow a handle to its local `at::Tensor`; that bridge is compiled against the installed Torch, as the existing JIT extension already is. ATen C++ objects do not cross the stable extension's ABI boundary. -The current checkout at `dbe854415da7771eba33195534c171adbca5677b` creates its own static BLAS handle; the inspected PR snapshot at `b551d7469db6d7ac688859cec46a3bb2b1c71a5b` borrows Torch's handle. The replacement removes either form of OEQ-side handle management. The final patch must be based on the branch being merged. +The helper creates temporary views with `aoti_torch__reinterpret_tensor`. These views retain the original storage and add their element offsets to the original tensor's storage offset. Their handles are released through RAII after each call, including when an error is returned. Original input and output handles remain owned by the caller. This replaces the prototype's `from_blob` construction and avoids reconstructing storage or querying pointer devices for every group. [Torch reinterpretation API](https://github.com/pytorch/pytorch/blob/v2.10.0/torch/csrc/inductor/aoti_torch/shim_common.cpp#L435). -The stable wheels currently target PyTorch 2.10. Use official stable C declarations, with no dependency on ATen's C++ ABI in that build. The C entry point was also exercised successfully on the available PyTorch 2.7 installation; that observation does not establish support for every older Torch release. [PyTorch stable ABI documentation](https://docs.pytorch.org/docs/main/notes/libtorch_stable_abi.html). +The operation checks GPU placement, matching devices and float32/float64 dtypes, tensor shapes, nonnegative dimensions, and mode 0 or 1. Ragged counts must be a CPU int64 vector with one entry per group. Counts may be noncontiguous; the operator makes a contiguous CPU copy before reading them. Counts must be nonnegative and sum to the input row count. Incremental bounds checking avoids overflowing the count sum. These checks establish the bounds needed by the storage-view helper. -## Public interface +Output allocation uses Torch and starts at zero. Empty groups are skipped; in mode 1 this leaves the corresponding weight-gradient block zero. Zero batch, output, or contraction dimensions return the zero-initialized output without invoking BMM. Dimensions and offsets remain int64 throughout the grouped-GEMM implementation. -Include `torch/csrc/inductor/aoti_torch/generated/c_shim_cuda.h`, which declares: +## Layout mapping -```cpp -AOTITorchError aoti_torch_cuda_bmm_out( - AtenTensorHandle out, - AtenTensorHandle self, - AtenTensorHandle mat2); -``` +Let `b = batch_size`, `n = ragged_counts[i]`, and `o` be the sum of preceding counts. Offsets are relative to the logical beginning of each original tensor; strides and offsets are in elements. -Use this generated API rather than the deprecated `aoti_torch_bmm_out` spelling. The API takes PyTorch tensor handles, not cuBLAS/hipBLAS handles. OEQ does not call a BLAS version getter or select a vendor library at runtime. [PyTorch 2.10 declaration](https://github.com/pytorch/pytorch/blob/v2.10.0/torch/csrc/inductor/aoti_torch/generated/c_shim_cuda.h#L74). +| Mode | Tensor | Element offset | View shape | View strides | +| --- | --- | --- | --- | --- | +| 0 | B / input | `o*b*k` | `[b,n,k]` | `[k,b*k,1]` | +| 0 | A / weights | `i*b*m*k` | `[b,k,m]` | `[m*k,1,k]` | +| 0 | C / output | `o*b*m` | `[b,n,m]` | `[m,b*m,1]` | +| 1 | A / left | `o*b*m` | `[b,m,n]` | `[m,1,b*m]` | +| 1 | B / right | `o*b*k` | `[b,n,k]` | `[k,b*k,1]` | +| 1 | C / output | `i*b*m*k` | `[b,m,k]` | `[m*k,k,1]` | -Add a Torch-specific helper, such as `group_mm_torch.hpp`, taking the existing input/output pointers, dimensions, CPU ragged counts, dtype, and explicit device index. Both Torch extension variants should use it where their supported Torch versions expose the required C symbols. Keep this helper outside the framework-independent CUDA/HIP kernel backend. +The final operation for each group is `aoti_torch_cuda_bmm_out(output, left, right)`. The API overwrites the output with alpha 1 and beta 0. Torch may choose its own BLAS kernel or backend. The layout code requests views, but makes no unconditional promise about copies inside Torch's BMM implementation. [Torch BMM implementation](https://github.com/pytorch/pytorch/blob/v2.10.0/aten/src/ATen/native/cuda/Blas.cpp#L536). -Create temporary, non-owning tensor views with `aoti_torch_create_tensor_from_blob_v2`. Specify sizes and strides in elements, storage offset zero, and the actual input device. Release each temporary tensor handle with `aoti_torch_delete_tensor_object` through RAII, including error paths. Releasing metadata must not release the caller's input or output storage. [PyTorch's blob implementation](https://github.com/pytorch/pytorch/blob/v2.10.0/torch/csrc/inductor/aoti_torch/shim_common.cpp#L525). +## Device, stream, precision, and autograd -## Layout mapping +The registered operator guards A's device before making GPU contiguous copies or allocating the output. The stable build uses `torch::stable::accelerator::DeviceGuard`; the JIT build uses `c10::DeviceGuard`. Each restores the caller's current device on return. -Let `b = batch_size`, `n = ragged_counts[i]`, and `o` be the sum of preceding counts. Pointer offsets and strides below are measured in elements. +BMM uses Torch's current stream on that device. OEQ neither creates a BLAS handle nor changes a BLAS stream. Calls remain asynchronous and follow PyTorch's normal storage/stream lifetime contract. -| Mode | View | Base pointer | Shape | Strides | -| --- | --- | --- | --- | --- | -| `ragged_inner == 0` | Left / input | `B + o*b*k` | `[b,n,k]` | `[k,b*k,1]` | -| `ragged_inner == 0` | Right / weights | `A + i*b*m*k` | `[b,k,m]` | `[m*k,1,k]` | -| `ragged_inner == 0` | Output | `C + o*b*m` | `[b,n,m]` | `[m,b*m,1]` | -| `ragged_inner == 1` | Left | `A + o*b*m` | `[b,m,n]` | `[m,1,b*m]` | -| `ragged_inner == 1` | Right | `B + o*b*k` | `[b,n,k]` | `[k,b*k,1]` | -| `ragged_inner == 1` | Output | `C + i*b*m*k` | `[b,m,k]` | `[m*k,k,1]` | +Following Torch's precision, determinism, and preferred-BLAS settings is an intentional behavior change relative to main's independently created handle. There is no OEQ-specific precision override. Custom backward registration remains responsible for differentiation through the grouped operator; the internal BMM-out invocation does not replace that registration. -Call BMM-out once for each nonempty group. Preserve the existing zero-initialized output and skip empty groups, including the all-empty case. This matters for the weight-gradient output blocks in mode 1. Use 64-bit dimensions and offsets throughout instead of the current narrowing casts to `int`. +## ABI and build paths -These views preserve OEQ's interleaved batch layout without explicitly transposing or copying the underlying buffers. PyTorch can handle BLAS-compatible strides directly, but some layouts can trigger its internal copy path; avoiding all copies is not an unconditional API guarantee. [BMM layout handling](https://github.com/pytorch/pytorch/blob/v2.10.0/aten/src/ATen/native/cuda/Blas.cpp#L536). +The official declaration is in `torch/csrc/inductor/aoti_torch/generated/c_shim_cuda.h`. The backend-specific spelling is used instead of the deprecated `aoti_torch_bmm_out`. [Declaration](https://github.com/pytorch/pytorch/blob/v2.10.0/torch/csrc/inductor/aoti_torch/generated/c_shim_cuda.h#L74), [stable ABI policy](https://docs.pytorch.org/docs/main/notes/libtorch_stable_abi.html). -## Device, stream, and autograd behavior +Both `aoti_torch_cuda_bmm_out` and `aoti_torch__reinterpret_tensor` are present in the inspected PyTorch 2.4 headers, matching OEQ's documented source/JIT baseline. The stable build explicitly targets the 2.10 ABI and continues to use the pinned LibTorch headers/libraries. This uses Torch's stated ABI guarantees; it does not try to select cuBLAS versions or infer handle ownership from version queries. -Guard the inputs' actual GPU device before making contiguous copies, allocating the output, constructing views, or invoking BMM. Restore the caller's device on return. The 2.10 stable implementation can use `aoti_torch_create_device_guard` / `aoti_torch_delete_device_guard`, whose implementation selects Torch's active accelerator; this avoids requiring a CUDA-specific guard implementation in the shared helper. The existing standalone prototype uses the older CUDA-named guard because it was tested on Torch 2.7. [Generic guard implementation](https://github.com/pytorch/pytorch/blob/v2.10.0/torch/csrc/inductor/aoti_torch/shim_common.cpp#L1468). +The source/JIT loader explicitly links the installed `torch_cuda` or `torch_hip` library supplying the BMM shim. CUDA also retains its driver/runtime/NVRTC dependencies; ROCm links hipRTC. All direct cuBLAS/rocBLAS calls, handle management, and corresponding OEQ link dependencies have been removed. -Use Torch's current stream for that device. Do not set a BLAS stream, create a stream, or synchronize inside the operator. Original tensors must remain valid through launch; asynchronous storage lifetime follows the usual PyTorch current-stream contract. No global or cached raw pointers or tensor views are needed. +The stable wheel build uses CPU LibTorch plus small GPU link stubs. [The stub](../../openequivariance/openequivariance/extension/stubs/stream.cpp) now declares and defines BMM-out using the official generated header, alongside the existing stream symbol. Stub bodies return failure if accidentally invoked. Only the real Torch GPU libraries are intended at runtime; CMake installs the OEQ targets, not the stubs. Both the Python extension and AOTI targets use this arrangement. -Validate the raw-pointer helper's preconditions at the tensor boundary: same GPU device and dtype, supported float32/float64 types, contiguous CPU int64 counts, valid counts length and nonnegative values, compatible input shapes, and a valid mode. Check that counts describe the available rows before constructing views. Keep the existing custom operator schema, fake implementation, and backward formulas. BMM-out is an internal implementation detail; the custom autograd registration remains responsible for derivatives. +The HIP CMake target is named `oeq_stable_hip`, matching its module entry point and expected artifact name. It explicitly links `hiprtc::hiprtc`. The existing Python loader still selects JIT compilation for HIP; enabling precompiled HIP loading is outside this change. The ROCm JIT path uses the same BMM helper as the stable CUDA build. -PyTorch's BMM implementation uses alpha 1 and beta 0, matching the existing GEMMs. It may choose a different BLAS backend or kernel and apply Torch's precision/determinism settings, so require numerical agreement rather than bitwise equivalence. [BMM implementation](https://github.com/pytorch/pytorch/blob/v2.10.0/aten/src/ATen/native/cuda/Blas.cpp#L680). +## ROCm evidence -## Why this should work on ROCm +PyTorch v2.10 adds the generated `c_shim_cuda.cpp` to `torch_hip` under `USE_ROCM`. The C symbol retains the `cuda` spelling on ROCm. The underlying GEMM implementations are converted to HIP BLAS calls, and Torch handles backend selection, including the double-precision fallback from hipBLASLt. [ROCm library construction](https://github.com/pytorch/pytorch/blob/v2.10.0/caffe2/CMakeLists.txt#L941), [HIP BLAS mappings](https://github.com/pytorch/pytorch/blob/v2.10.0/torch/utils/hipify/cuda_to_hip_mappings.py#L6826), [backend selection](https://github.com/pytorch/pytorch/blob/v2.10.0/aten/src/ATen/cuda/CUDABlas.cpp#L779). -Source inspection of PyTorch **v2.10.0** supports using the same generated function on both GPU platforms: +This supports the implementation choice but does not establish AMD hardware correctness. ROCm execution remains pending. -1. PyTorch's build defines `GENERATED_CXX_TORCH_CUDA` to contain `c_shim_cuda.cpp`, then explicitly adds that source to `torch_hip` under `USE_ROCM`. The generated shim is therefore part of the ROCm library too. [Generated source definition](https://github.com/pytorch/pytorch/blob/v2.10.0/caffe2/CMakeLists.txt#L331), [ROCm library construction](https://github.com/pytorch/pytorch/blob/v2.10.0/caffe2/CMakeLists.txt#L941). -2. The shim generator names the API using the `CUDA` dispatch key and generates a call into the corresponding ATen backend. The exported name remains `aoti_torch_cuda_bmm_out`; do not invent an `aoti_torch_hip_bmm_out` symbol. [Shim generator](https://github.com/pytorch/pytorch/blob/v2.10.0/torchgen/gen_aoti_c_shim.py#L492). -3. BMM reaches Torch's GEMM/batched GEMM implementation. Its HIP conversion maps `cublasSgemmStridedBatched` and `cublasDgemmStridedBatched` to their `hipblas` equivalents. Torch also contains ROCm-specific backend selection, including a double-precision fallback when hipBLASLt cannot handle the operation. [HIP mappings](https://github.com/pytorch/pytorch/blob/v2.10.0/torch/utils/hipify/cuda_to_hip_mappings.py#L6826), [float/double backend selection](https://github.com/pytorch/pytorch/blob/v2.10.0/aten/src/ATen/cuda/CUDABlas.cpp#L779). +## JAX scope -Torch's ROCm tensors use the `cuda` device interface. The helper should use Torch's CUDA device-type value for both these Torch builds, rather than confusing that convention with a distinct external HIP tensor type. [ROCm semantics](https://docs.pytorch.org/docs/main/notes/hip.html). +JAX does not call this grouped-GEMM helper and has no BLAS dependency in its extension target. The replacement is entirely within the Torch frontend. The common CUDA/HIP kernel compilation backend gains no Torch dependency. -This is **source-level evidence**, not an AMD hardware test or a binary export audit of a ROCm wheel. Verify the exports in the supported wheel and run one AMD smoke case before marking ROCm complete. +JAX's public FFI provides buffers and a GPU stream, without an equivalent BMM C shim. Future JAX grouped GEMM would require native JAX graph operations where shapes permit, or a separate FFI BLAS integration. [JAX FFI](https://docs.jax.dev/en/latest/ffi.html#ffi-calls-on-a-gpu). -## Build and packaging changes +## Validation and next GPU run -- Remove direct cuBLAS/rocBLAS includes, calls, and handle management from the Torch grouped-GEMM path. On the PR branch, also remove the now-unused borrowed-handle adapters. -- Remove `CUDA::cublas`, the grouped-GEMM-only `find_package(rocblas)` / `${HIP_BLAS_LIB}`, and the JIT loader's explicit `-lcublas` link flag after confirming there are no remaining consumers. Torch still brings its own BLAS dependencies. -- Resolve the generated BMM symbol through the installed Torch GPU library: `libtorch_cuda` for CUDA and `libtorch_hip` for ROCm. Common tensor/guard C APIs are supplied by Torch's common library. -- OEQ's wheel build currently downloads CPU LibTorch headers/libraries and creates GPU link stubs from `extension/stubs/stream.cpp`. Extend those build-only stubs with the exact generated BMM declaration/definition, sharing the header so signature drift causes a build error. If retaining the prototype's GPU-specific guards, add their GPU exports too; the proposed 2.10 generic guard avoids that need. Never package or execute the fake implementations. Check that installed extensions resolve against the real Torch libraries. -- Apply changes to both the Python extension and the AOTI shared library targets, and to source/JIT builds. The existing HIP CMake target is named `torch_stable_hip` while import/install expectations use `oeq_stable_hip`; resolve that naming mismatch as part of making the HIP artifact load correctly. -- Keep the existing CUDA/HIP runtime and runtime-compiler dependencies needed by OEQ-generated kernels. Removing the BLAS dependency does not make the entire extension independent of the GPU platform. +Local checks passed for the new operator and both adapters against Torch 2.10 headers, the link stub, and the shared helper against Torch 2.4 headers. These were host syntax checks, not full CUDA/ROCm extension builds. The unmodified production view helper also passed eight cases using real Torch 2.10 CPU tensors with the GPU BMM entry point redirected to CPU BMM for this check: both modes/dtypes, nonzero storage offsets, empty groups, and output guard values. All 30 GPU integration cases collect successfully; none has run on a GPU yet. -The stable-wheel baseline remains 2.10. Confirm the supported source/JIT baseline before sharing every C helper with that build: the prototype's helpers were exercised on 2.7, whereas the proposed generic guard needs its own baseline check. If older JIT support must be preserved, a JIT-only adapter using that installed Torch's ordinary BMM-out/device guard is acceptable; it still eliminates direct vendor BLAS calls and does not enter the stable wheel. +[The integration tests](../../tests/group_gemm_test.py) call the real registered operator. They cover both modes/dtypes, noncontiguous inputs and counts, nonzero input storage offsets, empty groups/dimensions, backward gradients, the current stream, device guarding, and invalid counts. The device-guard test requires two GPUs. The same test file supports CUDA and ROCm. -## JAX scope +[The import tests](../../tests/import_test.py) also inspect the extension and AOTI library's ELF dependencies and undefined symbols to check that OEQ has no direct vendor BLAS dependency. These checks run in the existing build-verification workflow for precompiled and JIT imports. -Neither the inspected PR snapshot nor the current JAX extension calls `group_gemm_blas`; its CMake target links the runtime, driver, and NVRTC, without BLAS. This replacement therefore remains entirely within Torch. - -JAX's public FFI has buffers and a GPU stream, but no equivalent BMM shim or public BLAS-handle getter. If a future JAX operation needs this GEMM, express it in the compiled JAX graph where shapes permit, or provide a separate FFI implementation whose BLAS handles and calls come from the same library. [JAX GPU FFI](https://docs.jax.dev/en/latest/ffi.html#ffi-calls-on-a-gpu), [public XLA FFI API](https://github.com/openxla/xla/blob/main/xla/ffi/api/c_api.h#L757). - -## Existing prototype and validation - -A standalone C++ prototype already implements both layout branches using only Torch C APIs. Its single float32 forward smoke case passed on an H100 PCIe with Torch 2.7.0 / CUDA 12.8: counts `[2,0,5,1]`, batch 3, m 4, k 5, a nondefault stream, and maximum absolute error `4.37e-7` against an independent CPU float64 reference. - -Earlier ctypes calls to the same C shim passed 40 cases spanning both dtypes and modes. Four selected profiles showed one GEMM kernel with no observed copy or GPU allocation inside the BMM call. These measurements do not cover the CPU cost of temporary tensor metadata or establish end-to-end performance. - -The integrated extension, its autograd registration, the final wheel link setup, and ROCm execution remain to be validated. No additional GPU experiments were run for this design. - -## Acceptance criteria - -- Both grouped-GEMM modes use Torch BMM, with no direct vendor BLAS symbols or handles in OEQ's Torch extension. -- Existing layouts, empty groups, float32/float64 behavior, operator schema, fake implementation, and gradients are preserved. -- Device guarding and a nondefault current stream work; the caller's current device is restored. -- Stable CUDA and HIP artifacts resolve the official C symbols in the real installed Torch GPU library; link stubs do not ship. Audit direct dynamic dependencies to confirm BLAS is now Torch's responsibility. -- Run focused correctness checks for both modes/dtypes and one relevant autograd case during integration, plus one ROCm smoke case when AMD hardware is available. Keep this bounded; no exhaustive vendor-version matrix is needed to establish the prototype. -- Check representative overhead before claiming performance parity. Graph capture and broader multi-device coverage should be checked where required by the supported operator contract. - -## Prototype code - -The following is the existing standalone CUDA prototype, not the final production registration or generic-device-guard adaptation. The caller supplies validated buffers and zero-initializes outputs where empty groups need to remain zero. - -```cpp -// Prototype for the PyTorch CUDA backend. Uses only PyTorch's stable C ABI. -#include - -#include -#include -#include -#include -#include -#include - -namespace { - -void check(AOTITorchError status) { - if (status != AOTI_TORCH_SUCCESS) - throw std::runtime_error("PyTorch C shim failed"); -} - -using Tensor = std::unique_ptr< - std::remove_pointer_t, - decltype(&aoti_torch_delete_tensor_object)>; -using DeviceGuard = std::unique_ptr< - std::remove_pointer_t, - decltype(&aoti_torch_delete_cuda_guard)>; - -// Only tensor metadata is created; the caller retains ownership of the buffer. -Tensor view(void* data, std::array sizes, - std::array strides, int32_t dtype, int32_t device) { - AtenTensorHandle tensor = nullptr; - check(aoti_torch_create_tensor_from_blob_v2( - data, 3, sizes.data(), strides.data(), 0, dtype, - aoti_torch_device_type_cuda(), device, &tensor, - aoti_torch_layout_strided(), nullptr, 0)); - return Tensor(tensor, aoti_torch_delete_tensor_object); -} - -} // namespace - -template -void group_gemm_cshim( - T* A, T* B, T* C, const int64_t* ragged_counts, int num_groups, - int64_t batch, int64_t m, int64_t k, int ragged_inner, int32_t device) { - static_assert(std::is_same_v || std::is_same_v); - const int32_t dtype = std::is_same_v - ? aoti_torch_dtype_float32() : aoti_torch_dtype_float64(); - - CUDAGuardHandle raw_guard = nullptr; - check(aoti_torch_create_cuda_guard(device, &raw_guard)); - DeviceGuard guard(raw_guard, aoti_torch_delete_cuda_guard); - - int64_t offset = 0; - for (int i = 0; i < num_groups; ++i) { - const int64_t n = ragged_counts[i]; - if (n == 0) continue; // Preserve the original empty-group behavior. - - if (ragged_inner == 0) { - // [batch, n, k] @ [batch, k, m] -> [batch, n, m] - auto input = view(B + offset * batch * k, - {batch, n, k}, {k, batch * k, 1}, dtype, device); - auto weight = view(A + i * batch * m * k, - {batch, k, m}, {m * k, 1, k}, dtype, device); - auto output = view(C + offset * batch * m, - {batch, n, m}, {m, batch * m, 1}, dtype, device); - check(aoti_torch_cuda_bmm_out(output.get(), input.get(), weight.get())); - } else { - // [batch, m, n] @ [batch, n, k] -> [batch, m, k] - auto left = view(A + offset * batch * m, - {batch, m, n}, {m, 1, batch * m}, dtype, device); - auto right = view(B + offset * batch * k, - {batch, n, k}, {k, batch * k, 1}, dtype, device); - auto output = view(C + i * batch * m * k, - {batch, m, k}, {m * k, k, 1}, dtype, device); - check(aoti_torch_cuda_bmm_out(output.get(), left.get(), right.get())); - } - offset += n; - } -} - -// Small ctypes entry point for the smoke test, not production registration code. -extern "C" int oeq_group_gemm_cshim( - int dtype, void* A, void* B, void* C, const int64_t* counts, - int groups, int64_t batch, int64_t m, int64_t k, int inner, int32_t device) { - try { - if (dtype == 0) { - group_gemm_cshim(static_cast(A), static_cast(B), - static_cast(C), counts, groups, - batch, m, k, inner, device); - } else if (dtype == 1) { - group_gemm_cshim(static_cast(A), static_cast(B), - static_cast(C), counts, groups, - batch, m, k, inner, device); - } else { - throw std::runtime_error("Expected float32 or float64"); - } - return 0; - } catch (const std::exception& error) { - std::fprintf(stderr, "%s\n", error.what()); - return 1; - } -} +After GPU access resumes, start with one case of the production operator: + +```sh +pytest -q 'tests/group_gemm_test.py::test_group_gemm_matches_reference[contiguous-0-dtype0]' ``` -
-Reproduce the single CUDA smoke case - -Save the C++ above as `group_gemm.cpp` and this script beside it as `smoke_test.py`, then run `python3 smoke_test.py` with a CUDA-enabled Torch installation and a C++ compiler. The script builds only the standalone prototype and executes one configuration. - -```python -"""Compile the C++ prototype and run one float32 CUDA correctness case.""" -import ctypes -from pathlib import Path -import subprocess - -import torch - - -def main(): - root = Path(__file__).resolve().parent - torch_root = Path(torch.__file__).resolve().parent - library = root / "group_gemm.so" - subprocess.run( - [ - "g++", "-std=c++17", "-O2", "-shared", "-fPIC", "-DUSE_CUDA", - f"-I{torch_root / 'include'}", str(root / "group_gemm.cpp"), - f"-L{torch_root / 'lib'}", f"-Wl,-rpath,{torch_root / 'lib'}", - "-Wl,--no-undefined", "-ltorch_cuda", "-ltorch_cpu", - "-o", str(library), - ], - check=True, - ) - lib = ctypes.CDLL(str(library)) - call = lib.oeq_group_gemm_cshim - i64 = ctypes.c_int64 - call.argtypes = [ - ctypes.c_int, ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p, - ctypes.POINTER(i64), ctypes.c_int, i64, i64, i64, - ctypes.c_int, ctypes.c_int32, - ] - call.restype = ctypes.c_int - - # One case includes differently sized groups, an empty group, interleaved - # batches, and a nondefault stream. The reference runs on the CPU. - counts = [2, 0, 5, 1] - batch, m, k = 3, 4, 5 - torch.manual_seed(123) - weights_cpu = torch.randn(len(counts), batch, m, k) - input_cpu = torch.randn(sum(counts), batch, k) - expected = torch.empty(sum(counts), batch, m, dtype=torch.float64) - offset = 0 - for i, n in enumerate(counts): - expected[offset:offset + n] = torch.einsum( - "bmk,nbk->nbm", weights_cpu[i].double(), - input_cpu[offset:offset + n].double(), - ) - offset += n - - torch.cuda.set_device(0) - stream = torch.cuda.Stream() - with torch.cuda.stream(stream): - weights = weights_cpu.cuda() - inputs = input_cpu.cuda() - output = torch.full(expected.shape, float("nan"), device="cuda") - status = call( - 0, weights.data_ptr(), inputs.data_ptr(), output.data_ptr(), - (i64 * len(counts))(*counts), len(counts), batch, m, k, 0, 0, - ) - assert status == 0, f"C shim prototype returned {status}" - stream.synchronize() - actual = output.cpu().double() - torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-5) - print( - f"PASS: one float32 case, counts={counts}, batch={batch}, m={m}, k={k}; " - f"max absolute error={(actual - expected).abs().max().item():.3g}; " - f"GPU={torch.cuda.get_device_name(0)}, torch={torch.__version__}" - ) - - -if __name__ == "__main__": - main() +Then run the focused integration suite against the stable build and the JIT build in separate processes: + +```sh +pytest -q tests/import_test.py tests/group_gemm_test.py +OEQ_JIT_EXTENSION=1 pytest -q tests/import_test.py tests/group_gemm_test.py ``` -
+Existing symmetric-contraction integration tests exercise the surrounding model and its higher-order derivatives when the optional MACE dependency is available. No end-to-end performance claim is made before measurement. + +## Earlier experiments + +The original [standalone prototype](group_gemm.cpp) and its [smoke script](smoke_test.py) remain as historical experiment artifacts. They use blob views rather than the production helper's views of existing storage. + +The prototype passed one float32 forward case on an H100 PCIe with Torch 2.7.0 / CUDA 12.8: counts `[2,0,5,1]`, batch 3, m 4, k 5, a nondefault stream, and maximum absolute error `4.37e-7` against a CPU float64 reference. Earlier ctypes calls to the BMM C shim passed 40 small cases across both dtypes/modes. Four selected BMM-only profiles showed no copy or GPU allocation inside the call; they did not measure view creation overhead. + +The earlier direct-cuBLAS experiments explain the migration: matching handle/function owners worked, whereas some separately loaded foreign-owner combinations failed or crashed despite successful version queries. They do not establish interchangeability of private BLAS handles. Related: [PR #206](https://github.com/PASSIONLab/OpenEquivariance/pull/206). diff --git a/tests/group_gemm_test.py b/tests/group_gemm_test.py new file mode 100644 index 00000000..6e983446 --- /dev/null +++ b/tests/group_gemm_test.py @@ -0,0 +1,173 @@ +import importlib + +import pytest +import torch + + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), reason="A CUDA or ROCm GPU is required" +) + + +@pytest.fixture(scope="module") +def group_gemm(): + importlib.import_module("openequivariance._torch.symmetric_contraction") + return torch.ops.libtorch_tp_jit.group_gemm + + +def reference(A, B, counts, inner): + pieces = [] + offset = 0 + for i, n in enumerate(counts): + if inner == 0: + pieces.append(torch.einsum("bmk,nbk->nbm", A[i], B[offset : offset + n])) + else: + pieces.append( + torch.einsum( + "nbm,nbk->bmk", A[offset : offset + n], B[offset : offset + n] + ) + ) + offset += n + if inner == 0: + return torch.cat(pieces, dim=0) + return torch.stack(pieces, dim=0) + + +def make_input(shape, dtype, layout="contiguous", device="cuda"): + values = torch.arange(1, 1 + torch.Size(shape).numel(), dtype=dtype, device="cpu") + values = (values.remainder(17) - 8).reshape(shape) / 8 + if layout == "offset": + storage = torch.empty(values.numel() + 7, dtype=dtype, device=device) + result = storage[3 : 3 + values.numel()].view(shape) + result.copy_(values) + return result + if layout == "noncontiguous": + storage = torch.empty((*shape, 2), dtype=dtype, device=device) + result = storage[..., 0] + result.copy_(values) + return result + return values.to(device) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float64]) +@pytest.mark.parametrize("inner", [0, 1]) +@pytest.mark.parametrize("layout", ["contiguous", "offset", "noncontiguous"]) +def test_group_gemm_matches_reference(group_gemm, dtype, inner, layout): + counts = [2, 0, 3] + batch, m, k = 2, 3, 4 + A_shape = (len(counts), batch, m, k) if inner == 0 else (sum(counts), batch, m) + A = make_input(A_shape, dtype, layout) + B = make_input((sum(counts), batch, k), dtype, layout) + expected = reference(A.cpu().double(), B.cpu().double(), counts, inner) + counts_tensor = torch.tensor([n for n in counts for _ in range(2)], device="cpu")[ + ::2 + ] + + actual = group_gemm(A, B, counts_tensor, len(counts), batch, m, k, inner) + + torch.testing.assert_close(actual.cpu().double(), expected, rtol=1e-5, atol=1e-5) + assert actual.dtype == dtype + assert actual.device == A.device + + +@pytest.mark.parametrize("inner", [0, 1]) +@pytest.mark.parametrize( + "counts,batch,m,k", + [ + ([0, 0], 2, 3, 4), + ([], 2, 3, 4), + ([2, 0, 3], 0, 3, 4), + ([2, 0, 3], 2, 0, 4), + ([2, 0, 3], 2, 3, 0), + ], +) +def test_group_gemm_empty_dimensions(group_gemm, inner, counts, batch, m, k): + A_shape = (len(counts), batch, m, k) if inner == 0 else (sum(counts), batch, m) + A = torch.empty(A_shape, device="cuda", dtype=torch.float64) + B = torch.empty((sum(counts), batch, k), device="cuda", dtype=torch.float64) + actual = group_gemm( + A, + B, + torch.tensor(counts, dtype=torch.int64, device="cpu"), + len(counts), + batch, + m, + k, + inner, + ) + expected_shape = ( + (sum(counts), batch, m) if inner == 0 else (len(counts), batch, m, k) + ) + assert actual.shape == expected_shape + torch.testing.assert_close(actual, torch.zeros_like(actual)) + + +@pytest.mark.parametrize("inner", [0, 1]) +def test_group_gemm_backward(group_gemm, inner): + counts = [2, 0, 3] + batch, m, k = 2, 3, 4 + A_shape = (len(counts), batch, m, k) if inner == 0 else (sum(counts), batch, m) + A = make_input(A_shape, torch.float64, "noncontiguous").requires_grad_() + B = make_input((sum(counts), batch, k), torch.float64, "offset").requires_grad_() + A_ref = A.detach().cpu().requires_grad_() + B_ref = B.detach().cpu().requires_grad_() + expected = reference(A_ref, B_ref, counts, inner) + actual = group_gemm( + A, B, torch.tensor(counts, device="cpu"), len(counts), batch, m, k, inner + ) + grad = make_input(actual.shape, torch.float64, "noncontiguous") + + actual_grads = torch.autograd.grad(actual, (A, B), grad) + expected_grads = torch.autograd.grad(expected, (A_ref, B_ref), grad.cpu()) + + for actual_grad, expected_grad in zip(actual_grads, expected_grads): + torch.testing.assert_close( + actual_grad.cpu(), expected_grad, rtol=1e-10, atol=1e-10 + ) + + +def test_group_gemm_current_stream(group_gemm): + counts = torch.tensor([2, 0, 3], device="cpu") + A = torch.zeros((3, 2, 3, 4), device="cuda") + B = torch.zeros((5, 2, 4), device="cuda") + group_gemm(A, B, counts, 3, 2, 3, 4, 0) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + if not torch.version.hip: + torch.cuda._sleep(1_000_000) + A.fill_(2) + B.fill_(3) + actual = group_gemm(A, B, counts, 3, 2, 3, 4, 0).clone() + stream.synchronize() + torch.testing.assert_close(actual, torch.full_like(actual, 24)) + + +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="Two GPUs are required") +def test_group_gemm_device_guard(group_gemm): + with torch.cuda.device(0): + A = torch.full((1, 2, 3, 4), 2.0, device="cuda:1") + B = torch.full((2, 2, 4), 3.0, device="cuda:1") + actual = group_gemm(A, B, torch.tensor([2], device="cpu"), 1, 2, 3, 4, 0) + assert torch.cuda.current_device() == 0 + assert actual.device == torch.device("cuda:1") + torch.testing.assert_close(actual, torch.full_like(actual, 24)) + with pytest.raises(RuntimeError, match="same device"): + group_gemm( + A, B.to("cuda:0"), torch.tensor([2], device="cpu"), 1, 2, 3, 4, 0 + ) + + +@pytest.mark.parametrize("counts", [[-1, 0, 6], [2, 0, 2], [2, 0, 4]]) +def test_group_gemm_rejects_invalid_counts(group_gemm, counts): + A = torch.empty((3, 2, 3, 4), device="cuda") + B = torch.empty((5, 2, 4), device="cuda") + with pytest.raises(RuntimeError, match="ragged_counts"): + group_gemm(A, B, torch.tensor(counts, device="cpu"), 3, 2, 3, 4, 0) + + +def test_group_gemm_requires_cpu_counts(group_gemm): + A = torch.empty((3, 2, 3, 4), device="cuda") + B = torch.empty((5, 2, 4), device="cuda") + with pytest.raises(RuntimeError, match="ragged_counts must be on the CPU"): + group_gemm(A, B, torch.tensor([2, 0, 3], device="cuda"), 3, 2, 3, 4, 0) diff --git a/tests/import_test.py b/tests/import_test.py index bf26af31..aa53630e 100644 --- a/tests/import_test.py +++ b/tests/import_test.py @@ -1,4 +1,9 @@ from importlib.metadata import version +import re +import shutil +import subprocess + +import pytest def test_import(): @@ -14,3 +19,26 @@ def test_extension_built(): assert BUILT_EXTENSION_ERROR is None assert BUILT_EXTENSION + + +def test_extension_has_no_direct_blas_dependency(): + if not shutil.which("readelf"): + pytest.skip("readelf is required to inspect ELF dependencies") + import openequivariance + from openequivariance._torch.extlib import extension_module + + paths = {extension_module.__file__, openequivariance.torch_ext_so_path()} + for path in paths: + dynamic = subprocess.check_output(["readelf", "--dynamic", path], text=True) + dependencies = [line for line in dynamic.splitlines() if "(NEEDED)" in line] + assert not any( + re.search(r"lib(cublas|hipblas|rocblas)", line) for line in dependencies + ) + symbols = subprocess.check_output( + ["readelf", "--dyn-syms", "--wide", path], text=True + ) + undefined = [line for line in symbols.splitlines() if " UND " in line] + assert not any( + re.search(r"\b(cublas|hipblas|rocblas)[A-Za-z_]", line) + for line in undefined + ) From f0da629c0eaa1ab8128703887a09b5ee15c57a7c Mon Sep 17 00:00:00 2001 From: asglover <140220574+asglover@users.noreply.github.com> Date: Sun, 13 Sep 2026 01:09:32 -0700 Subject: [PATCH 03/38] Expand grouped GEMM integration coverage --- tests/group_gemm_test.py | 135 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 135 insertions(+) diff --git a/tests/group_gemm_test.py b/tests/group_gemm_test.py index 6e983446..6bf70330 100644 --- a/tests/group_gemm_test.py +++ b/tests/group_gemm_test.py @@ -1,4 +1,6 @@ import importlib +import subprocess +import sys import pytest import torch @@ -171,3 +173,136 @@ def test_group_gemm_requires_cpu_counts(group_gemm): B = torch.empty((5, 2, 4), device="cuda") with pytest.raises(RuntimeError, match="ragged_counts must be on the CPU"): group_gemm(A, B, torch.tensor([2, 0, 3], device="cuda"), 3, 2, 3, 4, 0) + + +@pytest.mark.parametrize("inner", [0, 1]) +@pytest.mark.parametrize( + "counts,batch,m,k", + [ + ([1, 0, 5], 1, 1, 7), + ([1, 0, 5], 3, 7, 1), + ([1, 0, 5], 3, 1, 1), + ([1, 7, 0, 13], 4, 17, 29), + ], +) +def test_group_gemm_varied_shapes(group_gemm, inner, counts, batch, m, k): + A_shape = (len(counts), batch, m, k) if inner == 0 else (sum(counts), batch, m) + A = make_input(A_shape, torch.float64, "offset") + B = make_input((sum(counts), batch, k), torch.float64, "offset") + actual = group_gemm( + A, B, torch.tensor(counts, device="cpu"), len(counts), batch, m, k, inner + ) + expected = reference(A.cpu(), B.cpu(), counts, inner) + torch.testing.assert_close(actual.cpu(), expected, rtol=1e-10, atol=1e-10) + + +@pytest.mark.parametrize("inner", [0, 1]) +def test_group_gemm_double_backward(group_gemm, inner): + counts = torch.tensor([1, 0, 2], device="cpu") + A_shape = (3, 2, 2, 3) if inner == 0 else (3, 2, 2) + A = make_input(A_shape, torch.float64).requires_grad_() + B = make_input((3, 2, 3), torch.float64).requires_grad_() + + def operation(A, B): + return group_gemm(A, B, counts, 3, 2, 2, 3, inner) + + assert torch.autograd.gradgradcheck(operation, (A, B), fast_mode=True) + + +@pytest.mark.parametrize("inner", [0, 1]) +def test_group_gemm_graph_replay(group_gemm, inner): + counts = [2, 0, 3] + counts_tensor = torch.tensor(counts, device="cpu") + A_shape = (3, 2, 3, 4) if inner == 0 else (5, 2, 3) + A = make_input(A_shape, torch.float64) + B = make_input((5, 2, 4), torch.float64) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + group_gemm(A, B, counts_tensor, 3, 2, 3, 4, inner) + torch.cuda.current_stream().wait_stream(stream) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + actual = group_gemm(A, B, counts_tensor, 3, 2, 3, 4, inner) + + for scale in (2, 3): + A.mul_(scale) + B.add_(0.25) + graph.replay() + expected = reference(A.cpu(), B.cpu(), counts, inner) + torch.testing.assert_close(actual.cpu(), expected, rtol=1e-10, atol=1e-10) + + +@pytest.mark.parametrize("inner", [0, 1]) +def test_group_gemm_compile(group_gemm, inner): + counts = [2, 0, 3] + counts_tensor = torch.tensor(counts, device="cpu") + A_shape = (3, 2, 3, 4) if inner == 0 else (5, 2, 3) + A = make_input(A_shape, torch.float64).requires_grad_() + B = make_input((5, 2, 4), torch.float64).requires_grad_() + + def operation(A, B, counts): + return group_gemm(A, B, counts, 3, 2, 3, 4, inner) + + compiled = torch.compile(operation, fullgraph=True) + actual = compiled(A, B, counts_tensor) + A_ref = A.detach().cpu().requires_grad_() + B_ref = B.detach().cpu().requires_grad_() + expected = reference(A_ref, B_ref, counts, inner) + torch.testing.assert_close(actual.cpu(), expected, rtol=1e-10, atol=1e-10) + + actual_grads = torch.autograd.grad(actual.sum(), (A, B)) + expected_grads = torch.autograd.grad(expected.sum(), (A_ref, B_ref)) + for actual_grad, expected_grad in zip(actual_grads, expected_grads): + torch.testing.assert_close( + actual_grad.cpu(), expected_grad, rtol=1e-10, atol=1e-10 + ) + + +@pytest.mark.parametrize("inner", [0, 1]) +def test_group_gemm_aoti(group_gemm, inner, tmp_path): + import openequivariance + + class Model(torch.nn.Module): + def forward(self, A, B, counts): + return group_gemm(A, B, counts, 3, 2, 3, 4, inner) + + counts = [2, 0, 3] + counts_tensor = torch.tensor(counts, device="cpu") + A_shape = (3, 2, 3, 4) if inner == 0 else (5, 2, 3) + A = make_input(A_shape, torch.float64) + B = make_input((5, 2, 4), torch.float64) + exported = torch.export.export(Model(), (A, B, counts_tensor), strict=False) + package_path = torch._inductor.aoti_compile_and_package( + exported, package_path=str(tmp_path / "group_gemm.pt2") + ) + inputs_path = tmp_path / "inputs.pt" + torch.save( + (A.cpu(), B.cpu(), counts_tensor, reference(A.cpu(), B.cpu(), counts, inner)), + inputs_path, + ) + result = subprocess.run( + [ + sys.executable, + "-c", + """ +import sys +import torch + +torch.ops.load_library(sys.argv[1]) +model = torch._inductor.aoti_load_package(sys.argv[2]) +A, B, counts, expected = torch.load(sys.argv[3], weights_only=True) +actual = model(A.cuda(), B.cuda(), counts) +torch.testing.assert_close(actual.cpu(), expected, rtol=1e-10, atol=1e-10) +""", + openequivariance.torch_ext_so_path(), + package_path, + str(inputs_path), + ], + capture_output=True, + text=True, + timeout=180, + ) + assert result.returncode == 0, result.stdout + result.stderr From d8f56063218e964dcbeba895e7ee8d2cccea31b1 Mon Sep 17 00:00:00 2001 From: asglover <140220574+asglover@users.noreply.github.com> Date: Sun, 13 Sep 2026 01:19:50 -0700 Subject: [PATCH 04/38] Package stable libraries in wheels and resolve editable library paths --- .github/workflows/verify_extension_build.yml | 3 ++- openequivariance/CMakeLists.txt | 2 +- .../openequivariance/_torch/extlib/__init__.py | 13 ++++++++++++- sandbox/cshim_group_gemm/DESIGN.md | 6 ++++-- 4 files changed, 19 insertions(+), 5 deletions(-) diff --git a/.github/workflows/verify_extension_build.yml b/.github/workflows/verify_extension_build.yml index c6b7c16a..ecb5929a 100644 --- a/.github/workflows/verify_extension_build.yml +++ b/.github/workflows/verify_extension_build.yml @@ -29,10 +29,11 @@ jobs: sudo apt-get update sudo apt install nvidia-cuda-toolkit pip install -r .github/workflows/requirements_cuda_ci.txt - pip install -e "./openequivariance[jax]" + pip install "./openequivariance[jax]" - name: Test CUDA extension build via import run: | + python -c "import openequivariance as oeq; assert oeq.USE_PRECOMPILED_EXTENSION" pytest tests/import_test.py export OEQ_JIT_EXTENSION=1 diff --git a/openequivariance/CMakeLists.txt b/openequivariance/CMakeLists.txt index 7ed586db..f4eb5884 100644 --- a/openequivariance/CMakeLists.txt +++ b/openequivariance/CMakeLists.txt @@ -52,7 +52,7 @@ set(OEQ_SOURCES ${EXT_JSON_DIR}/json11.cpp ) -set(OEQ_INSTALL_DIR "${CMAKE_CURRENT_SOURCE_DIR}/openequivariance/_torch/extlib") +set(OEQ_INSTALL_DIR "openequivariance/_torch/extlib") function(add_stable_extension target_name backend_define link_libraries) # Create nanobind extension diff --git a/openequivariance/openequivariance/_torch/extlib/__init__.py b/openequivariance/openequivariance/_torch/extlib/__init__.py index 0861f2e0..19720e3b 100644 --- a/openequivariance/openequivariance/_torch/extlib/__init__.py +++ b/openequivariance/openequivariance/_torch/extlib/__init__.py @@ -5,6 +5,7 @@ import warnings import sysconfig import contextlib +import importlib.util from pathlib import Path from packaging.version import Version @@ -224,6 +225,16 @@ def load_precompiled_extension(): ) +def _has_precompiled_extension(): + backend = BACKEND + spec = importlib.util.find_spec(f"{__name__}.oeq_stable_{backend}") + return ( + spec is not None + and spec.origin is not None + and Path(spec.origin).with_name(f"liboeq_stable_{backend}_aoti.so").is_file() + ) + + USE_PRECOMPILED_EXTENSION = True WARNING_MESSAGE = "" @@ -245,7 +256,7 @@ def load_precompiled_extension(): AOTI_SO_NAME = f"liboeq_stable_{BACKEND}_aoti.so" -if not os.path.exists(os.path.join(os.path.dirname(__file__), AOTI_SO_NAME)): +if not _has_precompiled_extension(): WARNING_MESSAGE += "Precompiled extension shared object not found.\n" USE_PRECOMPILED_EXTENSION = False diff --git a/sandbox/cshim_group_gemm/DESIGN.md b/sandbox/cshim_group_gemm/DESIGN.md index 9626260e..c5c5ea2d 100644 --- a/sandbox/cshim_group_gemm/DESIGN.md +++ b/sandbox/cshim_group_gemm/DESIGN.md @@ -49,6 +49,8 @@ The source/JIT loader explicitly links the installed `torch_cuda` or `torch_hip` The stable wheel build uses CPU LibTorch plus small GPU link stubs. [The stub](../../openequivariance/openequivariance/extension/stubs/stream.cpp) now declares and defines BMM-out using the official generated header, alongside the existing stream symbol. Stub bodies return failure if accidentally invoked. Only the real Torch GPU libraries are intended at runtime; CMake installs the OEQ targets, not the stubs. Both the Python extension and AOTI targets use this arrangement. +CMake installs the stable libraries into the wheel's package directory. The previous absolute destination wrote them into the source tree, where the wheel's ignore rules excluded them. The loader now resolves the compiled extension's location to find its companion AOTI library; this also supports editable installations where Python sources and compiled libraries live in different directories. Build CI installs a normal wheel and asserts that the stable extension is selected before running the import checks. + The HIP CMake target is named `oeq_stable_hip`, matching its module entry point and expected artifact name. It explicitly links `hiprtc::hiprtc`. The existing Python loader still selects JIT compilation for HIP; enabling precompiled HIP loading is outside this change. The ROCm JIT path uses the same BMM helper as the stable CUDA build. ## ROCm evidence @@ -65,9 +67,9 @@ JAX's public FFI provides buffers and a GPU stream, without an equivalent BMM C ## Validation and next GPU run -Local checks passed for the new operator and both adapters against Torch 2.10 headers, the link stub, and the shared helper against Torch 2.4 headers. These were host syntax checks, not full CUDA/ROCm extension builds. The unmodified production view helper also passed eight cases using real Torch 2.10 CPU tensors with the GPU BMM entry point redirected to CPU BMM for this check: both modes/dtypes, nonzero storage offsets, empty groups, and output guard values. All 30 GPU integration cases collect successfully; none has run on a GPU yet. +Local checks passed for the new operator and both adapters against Torch 2.10 headers, the link stub, and the shared helper against Torch 2.4 headers. These were host syntax checks, not full CUDA/ROCm extension builds. The unmodified production view helper also passed eight cases using real Torch 2.10 CPU tensors with the GPU BMM entry point redirected to CPU BMM for this check: both modes/dtypes, nonzero storage offsets, empty groups, and output guard values. All 46 GPU integration cases collect successfully; none has run on a GPU yet. -[The integration tests](../../tests/group_gemm_test.py) call the real registered operator. They cover both modes/dtypes, noncontiguous inputs and counts, nonzero input storage offsets, empty groups/dimensions, backward gradients, the current stream, device guarding, and invalid counts. The device-guard test requires two GPUs. The same test file supports CUDA and ROCm. +[The integration tests](../../tests/group_gemm_test.py) call the real registered operator. They cover both modes/dtypes, noncontiguous inputs and counts, nonzero input storage offsets, empty and singleton dimensions, varied group sizes, backward and double-backward gradients, the current stream, device guarding, invalid counts, graph replay, compiled training, and AOTI inference in a fresh process that loads only the exported OEQ library. The device-guard test requires two GPUs. The same test file supports CUDA and ROCm. [The import tests](../../tests/import_test.py) also inspect the extension and AOTI library's ELF dependencies and undefined symbols to check that OEQ has no direct vendor BLAS dependency. These checks run in the existing build-verification workflow for precompiled and JIT imports. From ea35c8f62696cf469f13a58e7e93a5b20b9c527b Mon Sep 17 00:00:00 2001 From: asglover <140220574+asglover@users.noreply.github.com> Date: Sun, 13 Sep 2026 01:30:29 -0700 Subject: [PATCH 05/38] Keep Torch adapter helpers local to each extension --- .../openequivariance/extension/libtorch_tp_jit.cpp | 4 ++++ .../openequivariance/extension/libtorch_tp_jit_stable.cpp | 4 ++++ 2 files changed, 8 insertions(+) diff --git a/openequivariance/openequivariance/extension/libtorch_tp_jit.cpp b/openequivariance/openequivariance/extension/libtorch_tp_jit.cpp index 08f98d79..e4600652 100644 --- a/openequivariance/openequivariance/extension/libtorch_tp_jit.cpp +++ b/openequivariance/openequivariance/extension/libtorch_tp_jit.cpp @@ -35,6 +35,8 @@ constexpr Dtype kByte = torch::kByte; #define REGISTER_LIBRARY_IMPL TORCH_LIBRARY_IMPL #define REGISTER_LIBRARY TORCH_LIBRARY +namespace { + class TensorDeviceGuard { c10::DeviceGuard guard; public: @@ -45,6 +47,8 @@ AtenTensorHandle tensor_handle(Tensor& tensor) { return torch::aot_inductor::tensor_pointer_to_tensor_handle(&tensor); } +} + #include "torch_core.hpp" Tensor tensor_to_cpu_contiguous(const Tensor &tensor) { diff --git a/openequivariance/openequivariance/extension/libtorch_tp_jit_stable.cpp b/openequivariance/openequivariance/extension/libtorch_tp_jit_stable.cpp index 492931d3..7792ee1c 100644 --- a/openequivariance/openequivariance/extension/libtorch_tp_jit_stable.cpp +++ b/openequivariance/openequivariance/extension/libtorch_tp_jit_stable.cpp @@ -32,6 +32,8 @@ constexpr Dtype kByte = torch::headeronly::ScalarType::Byte; #define REGISTER_LIBRARY_IMPL STABLE_TORCH_LIBRARY_IMPL #define REGISTER_LIBRARY STABLE_TORCH_LIBRARY +namespace { + class TensorDeviceGuard { torch::stable::accelerator::DeviceGuard guard; public: @@ -42,6 +44,8 @@ AtenTensorHandle tensor_handle(Tensor& tensor) { return tensor.get(); } +} + #include "torch_core.hpp" Tensor tensor_to_cpu_contiguous(const Tensor &tensor) { From 129a36476a691b391802c58c638b8d75bd9a79eb Mon Sep 17 00:00:00 2001 From: asglover <140220574+asglover@users.noreply.github.com> Date: Sun, 13 Sep 2026 01:36:54 -0700 Subject: [PATCH 06/38] Verify compiled grouped GEMM accepts changing counts --- sandbox/cshim_group_gemm/DESIGN.md | 2 ++ tests/group_gemm_test.py | 5 +++++ 2 files changed, 7 insertions(+) diff --git a/sandbox/cshim_group_gemm/DESIGN.md b/sandbox/cshim_group_gemm/DESIGN.md index c5c5ea2d..e1ca774b 100644 --- a/sandbox/cshim_group_gemm/DESIGN.md +++ b/sandbox/cshim_group_gemm/DESIGN.md @@ -37,6 +37,8 @@ The registered operator guards A's device before making GPU contiguous copies or BMM uses Torch's current stream on that device. OEQ neither creates a BLAS handle nor changes a BLAS stream. Calls remain asynchronous and follow PyTorch's normal storage/stream lifetime contract. +GPU graph capture records the group sizes read by the host loop during capture. Direct graph replay requires those counts to remain unchanged; changing group sizes requires recapture. Ordinary calls and compiled execution without graph replay read the counts on each invocation. + Following Torch's precision, determinism, and preferred-BLAS settings is an intentional behavior change relative to main's independently created handle. There is no OEQ-specific precision override. Custom backward registration remains responsible for differentiation through the grouped operator; the internal BMM-out invocation does not replace that registration. ## ABI and build paths diff --git a/tests/group_gemm_test.py b/tests/group_gemm_test.py index 6bf70330..33f7945d 100644 --- a/tests/group_gemm_test.py +++ b/tests/group_gemm_test.py @@ -260,6 +260,11 @@ def operation(A, B, counts): actual_grad.cpu(), expected_grad, rtol=1e-10, atol=1e-10 ) + next_counts = [0, 3, 2] + next_actual = compiled(A, B, torch.tensor(next_counts, device="cpu")) + next_expected = reference(A_ref, B_ref, next_counts, inner) + torch.testing.assert_close(next_actual.cpu(), next_expected, rtol=1e-10, atol=1e-10) + @pytest.mark.parametrize("inner", [0, 1]) def test_group_gemm_aoti(group_gemm, inner, tmp_path): From 9aedaacc97e3d5ce9d5217c909de2683168e2f92 Mon Sep 17 00:00:00 2001 From: asglover <140220574+asglover@users.noreply.github.com> Date: Sun, 13 Sep 2026 03:04:21 -0700 Subject: [PATCH 07/38] Initialize Inductor codecache for fresh-process AOTI validation --- tests/group_gemm_test.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/group_gemm_test.py b/tests/group_gemm_test.py index 33f7945d..8df3d5bf 100644 --- a/tests/group_gemm_test.py +++ b/tests/group_gemm_test.py @@ -295,6 +295,7 @@ def forward(self, A, B, counts): """ import sys import torch +import torch._inductor.codecache torch.ops.load_library(sys.argv[1]) model = torch._inductor.aoti_load_package(sys.argv[2]) From 1d6054362cee937065e94bae7d188705f77e47ad Mon Sep 17 00:00:00 2001 From: asglover <140220574+asglover@users.noreply.github.com> Date: Sun, 13 Sep 2026 03:15:19 -0700 Subject: [PATCH 08/38] Document production CUDA validation of the stable BMM shim --- sandbox/cshim_group_gemm/DESIGN.md | 25 +++------ sandbox/cshim_group_gemm/VALIDATION.md | 73 ++++++++++++++++++++++++++ 2 files changed, 80 insertions(+), 18 deletions(-) create mode 100644 sandbox/cshim_group_gemm/VALIDATION.md diff --git a/sandbox/cshim_group_gemm/DESIGN.md b/sandbox/cshim_group_gemm/DESIGN.md index e1ca774b..ed88da3c 100644 --- a/sandbox/cshim_group_gemm/DESIGN.md +++ b/sandbox/cshim_group_gemm/DESIGN.md @@ -2,7 +2,7 @@ The production Torch extension now implements `libtorch_tp_jit::group_gemm` through `aoti_torch_cuda_bmm_out` on CUDA and ROCm. Both the stable extension and the source/JIT extension use this C entry point. The operator schema, fake implementation, output layouts, and custom backward formulas remain compatible. -GPU execution of this implementation is pending. GPU access is paused at the user's request; the earlier prototype results below are not results for this implementation. +The production implementation passed CUDA validation on an H100 with Torch 2.10's stable and JIT extensions and Torch 2.7's JIT extension. The complete matrix and its hardware limits are recorded in [the validation report](VALIDATION.md). ## Implementation @@ -59,7 +59,7 @@ The HIP CMake target is named `oeq_stable_hip`, matching its module entry point PyTorch v2.10 adds the generated `c_shim_cuda.cpp` to `torch_hip` under `USE_ROCM`. The C symbol retains the `cuda` spelling on ROCm. The underlying GEMM implementations are converted to HIP BLAS calls, and Torch handles backend selection, including the double-precision fallback from hipBLASLt. [ROCm library construction](https://github.com/pytorch/pytorch/blob/v2.10.0/caffe2/CMakeLists.txt#L941), [HIP BLAS mappings](https://github.com/pytorch/pytorch/blob/v2.10.0/torch/utils/hipify/cuda_to_hip_mappings.py#L6826), [backend selection](https://github.com/pytorch/pytorch/blob/v2.10.0/aten/src/ATen/cuda/CUDABlas.cpp#L779). -This supports the implementation choice but does not establish AMD hardware correctness. ROCm execution remains pending. +This supports the implementation choice but does not establish AMD hardware correctness. ROCm execution has not been validated. ## JAX scope @@ -67,28 +67,17 @@ JAX does not call this grouped-GEMM helper and has no BLAS dependency in its ext JAX's public FFI provides buffers and a GPU stream, without an equivalent BMM C shim. Future JAX grouped GEMM would require native JAX graph operations where shapes permit, or a separate FFI BLAS integration. [JAX FFI](https://docs.jax.dev/en/latest/ffi.html#ffi-calls-on-a-gpu). -## Validation and next GPU run +## Validation -Local checks passed for the new operator and both adapters against Torch 2.10 headers, the link stub, and the shared helper against Torch 2.4 headers. These were host syntax checks, not full CUDA/ROCm extension builds. The unmodified production view helper also passed eight cases using real Torch 2.10 CPU tensors with the GPU BMM entry point redirected to CPU BMM for this check: both modes/dtypes, nonzero storage offsets, empty groups, and output guard values. All 46 GPU integration cases collect successfully; none has run on a GPU yet. +The stable wheel and JIT extensions compiled and loaded successfully. Wheel and editable-install checks verified that the installed stable libraries are found and that the build stub is excluded. GPU validation produced 192 successful checks across the three extension configurations and the surrounding model tests, with three skips for the device-guard test that requires two GPUs. [Results, environment corrections, and reproduction commands](VALIDATION.md). + +Earlier local checks also passed for the new operator and both adapters against Torch 2.10 headers, the link stub, and the shared helper against Torch 2.4 headers. The unmodified production view helper passed eight cases using real Torch 2.10 CPU tensors with the GPU BMM entry point redirected to CPU BMM for that check: both modes/dtypes, nonzero storage offsets, empty groups, and output guard values. [The integration tests](../../tests/group_gemm_test.py) call the real registered operator. They cover both modes/dtypes, noncontiguous inputs and counts, nonzero input storage offsets, empty and singleton dimensions, varied group sizes, backward and double-backward gradients, the current stream, device guarding, invalid counts, graph replay, compiled training, and AOTI inference in a fresh process that loads only the exported OEQ library. The device-guard test requires two GPUs. The same test file supports CUDA and ROCm. [The import tests](../../tests/import_test.py) also inspect the extension and AOTI library's ELF dependencies and undefined symbols to check that OEQ has no direct vendor BLAS dependency. These checks run in the existing build-verification workflow for precompiled and JIT imports. -After GPU access resumes, start with one case of the production operator: - -```sh -pytest -q 'tests/group_gemm_test.py::test_group_gemm_matches_reference[contiguous-0-dtype0]' -``` - -Then run the focused integration suite against the stable build and the JIT build in separate processes: - -```sh -pytest -q tests/import_test.py tests/group_gemm_test.py -OEQ_JIT_EXTENSION=1 pytest -q tests/import_test.py tests/group_gemm_test.py -``` - -Existing symmetric-contraction integration tests exercise the surrounding model and its higher-order derivatives when the optional MACE dependency is available. No end-to-end performance claim is made before measurement. +The existing symmetric-contraction integration tests passed all 24 cases for both the stable and JIT extensions under Torch 2.10, including comparison with MACE and higher-order derivatives. No end-to-end performance claim is made by this correctness validation. ## Earlier experiments diff --git a/sandbox/cshim_group_gemm/VALIDATION.md b/sandbox/cshim_group_gemm/VALIDATION.md new file mode 100644 index 00000000..e7b0d60f --- /dev/null +++ b/sandbox/cshim_group_gemm/VALIDATION.md @@ -0,0 +1,73 @@ +# Production validation + +The production implementation passed CUDA validation on 2026-09-13. The final test matrix contains 192 successful checks and three skips, all for the same device-guard test requiring two GPUs. These results cover the registered production operator and installed artifacts, not the earlier standalone prototype. + +## Source and artifact + +- Branch: `move-bmm-calls-to-stable-cshim`, based on main at `dbe854415da7771eba33195534c171adbca5677b`. +- Production source: `d2eac4e`; integration tests: `e4df36c`. +- Wheel: `openequivariance-0.7.0-cp310-cp310-linux_x86_64.whl`. +- Wheel SHA256: `f536c53f1ef2358b6f562931c30c3f55823c32123424c153c96cb808250ea97a`. +- Remote checkout and environments: `/tmp/oeq-cshim-production/`, separate from the workspace-allocation checkout and user-installed packages. + +## Build and packaging results + +All checks in this table ran with `CUDA_VISIBLE_DEVICES=""`. They establish successful compilation, loading, and packaging; they do not establish GPU numerical correctness. + +| Configuration | Result | +| --- | --- | +| Stable wheel, Torch 2.10.0+cu128, Python 3.10 | Both production libraries built; 3 import/dependency tests passed | +| Editable install, Torch 2.10.0+cu128 | Stable extension selected with Python sources and libraries in separate directories; 3 import/dependency tests passed | +| Source/JIT extension, Torch 2.10.0+cu128 | Compiled against the installed Torch; 3 import/dependency tests passed | +| Source/JIT extension, Torch 2.7.0, CUDA 12.8 | Compiled against the installed Torch; 3 import/dependency tests passed | + +The wheel contains its Python extension, companion AOTI library, and JIT sources. It contains no `libtorch_cuda.so` build stub. ELF checks found no direct vendor BLAS library dependency or undefined vendor BLAS function in either production library or either JIT build. Both installed Torch GPU libraries export `aoti_torch_cuda_bmm_out`. + +Validation found and fixed a packaging defect: the absolute CMake install destination placed libraries in the source tree, where wheel ignore rules excluded them. The relative install destination includes the libraries in the wheel. Resolving the extension's actual import location also fixes stable-library discovery for editable installations. + +## GPU validation + +Hardware: one NVIDIA H100 PCIe, driver 580.105.08. Both Torch environments report CUDA 12.8. The first production-wheel forward case passed before running the larger suites. + +| Configuration | Import and grouped-GEMM tests | Symmetric-contraction tests | +| --- | --- | --- | +| Stable wheel, Torch 2.10.0+cu128 | 48 passed, 1 skipped | 24 passed | +| Source/JIT, Torch 2.10.0+cu128 | 48 passed, 1 skipped | 24 passed | +| Source/JIT, installed Torch 2.7.0 / CUDA 12.8 | 48 passed, 1 skipped | Not run in this environment | + +The stable-wheel result combines 46 initial passes with two successful AOTI reruns after correcting fresh-process loader initialization. The Torch 2.7 result is a complete rerun after correcting test-environment dependencies. The table reports final per-case outcomes; it does not count preliminary or repeated passes twice. + +The 46 grouped-GEMM cases cover layouts, empty groups and dimensions, both modes and dtypes, backward and double backward, the current stream, device guarding, invalid counts, graph capture/replay, compiled training with changing counts, and AOTI inference. Each AOTI case loads the exported OEQ library in a new Python process without importing the OEQ Python package, then compares against a CPU reference. Three additional tests check import, successful extension loading, and binary dependencies. + +The existing symmetric-contraction suite compares with MACE for float32/float64 forward, backward, and double backward across three configurations. It also checks compile and export. Both the stable and JIT extensions passed all 24 cases. + +## Issues resolved during validation + +The wheel and editable-discovery fixes above were production changes. GPU execution required no further changes to the GEMM implementation. + +The fresh-process AOTI test exposed a Torch 2.10 loader initialization issue: its package loader accesses `torch._inductor.codecache` before importing it. The subprocess now explicitly imports that module before loading the package. Both AOTI modes then passed for the stable wheel and both JIT builds. + +The Torch 2.10 model-test environment uses MACE 0.3.16 and its pinned e3nn 0.4.4. That e3nn version loads packaged constants containing Python `slice` objects. The test process allowed that specific built-in type with `torch.serialization.add_safe_globals([slice])`; no production code or global loading policy changed. NumPy was 1.26.4 and pytest was 9.1.1. + +The Torch 2.7 environment initially inherited NetworkX 2.4 from the system; its import failed on NumPy's removed `np.int` alias before compilation. Installing NetworkX 3.4.2, SciPy 1.15.3, and SymPy 1.13.3 inside the isolated environment corrected its compiler dependencies. Its final environment used e3nn 0.6.0 and NumPy 1.26.4. System packages and the other checkout were untouched. + +## Reproduction + +After installing the wheel and test dependencies in a matching GPU Torch environment, run: + +```sh +python -c 'import torch; torch.serialization.add_safe_globals([slice]); import pytest; raise SystemExit(pytest.main())' -q tests/import_test.py tests/group_gemm_test.py tests/symmetric_contraction_test.py +OEQ_JIT_EXTENSION=1 python -c 'import torch; torch.serialization.add_safe_globals([slice]); import pytest; raise SystemExit(pytest.main())' -q tests/import_test.py tests/group_gemm_test.py tests/symmetric_contraction_test.py +``` + +The allowlist is only needed for the older e3nn dependency described above. The actual runs used separate JIT and Inductor cache directories for each Torch environment and ran the focused and model suites separately. The Torch 2.7 environment also set `PYTHONNOUSERSITE=1` to exclude the shared machine's user packages. + +## Scope limits + +- ROCm has source-level support evidence in [the design](DESIGN.md), but no AMD hardware run has been performed. +- Multi-GPU device guarding remains unverified on hardware because only one GPU was available. +- Torch 2.4 received header checks only; GPU execution covered Torch 2.7 and 2.10. CUDA versions other than 12.8 were not exercised by this production test matrix. +- Direct GPU graph replay requires fixed ragged counts, as described in the design. +- No end-to-end performance claim is made by these correctness and packaging checks. + +Detailed logs, environment snapshots, and JUnit XML are retained under `/tmp/oeq-cshim-production/logs/` on the validation machine. At the end of testing, the GPU reported zero memory use and no compute processes; its status reservation was released. From 4e600c627224e4e7af955dfc7b135edd81054b34 Mon Sep 17 00:00:00 2001 From: asglover <140220574+asglover@users.noreply.github.com> Date: Sun, 13 Sep 2026 15:32:03 -0700 Subject: [PATCH 09/38] Document AMD hardware validation of the stable BMM shim --- sandbox/cshim_group_gemm/AMD_VALIDATION.md | 75 ++++++++++++++++++++++ sandbox/cshim_group_gemm/DESIGN.md | 8 +-- sandbox/cshim_group_gemm/VALIDATION.md | 3 +- 3 files changed, 81 insertions(+), 5 deletions(-) create mode 100644 sandbox/cshim_group_gemm/AMD_VALIDATION.md diff --git a/sandbox/cshim_group_gemm/AMD_VALIDATION.md b/sandbox/cshim_group_gemm/AMD_VALIDATION.md new file mode 100644 index 00000000..a928d44c --- /dev/null +++ b/sandbox/cshim_group_gemm/AMD_VALIDATION.md @@ -0,0 +1,75 @@ +# AMD production validation + +ROCm validation passed on 2026-09-13 without changing the GEMM implementation. The final matrix contains 192 successful checks and three skips for the device-guard test requiring two GPUs. The normal ROCm JIT path passed with Torch 2.10 and 2.7. The stable HIP artifacts also passed with the isolated loader override described below. + +## Source and environment + +- Branch: `move-bmm-calls-to-stable-cshim`, tested at `4e3dcc8`. +- GPU: one AMD Instinct MI300X VF, 192 GB, `gfx942`. +- Driver reported by `rocm-smi`: `6.19.14.31400000`; Linux kernel `6.8.0-138-generic`. +- Python: 3.12.3. +- Build SDK: `/opt/rocm/core-10.0`, reporting HIP `7.15.26333`; host compiler GCC 13.3.0. +- Test environments: official Torch `2.10.0+rocm7.1` and `2.7.0+rocm6.3` wheels in separate virtual environments. +- Checkout, environments, caches, and logs: `/tmp/oeq-cshim-amd/` on the AMD machine. + +The NVIDIA machine was not used during this validation. + +## Results + +The first production grouped-GEMM case passed before running the focused suites. + +| Configuration | Import and grouped-GEMM tests | Symmetric-contraction tests | +| --- | --- | --- | +| Normal JIT path, Torch 2.10 / ROCm 7.1 | 48 passed, 1 skipped | 24 passed | +| Stable HIP artifacts, Torch 2.10 / ROCm 7.1, explicit test loader override | 48 passed, 1 skipped | 24 passed | +| Normal JIT path, Torch 2.7 / ROCm 6.3 | 48 passed, 1 skipped | Not run in this environment | + +The focused tests cover both GEMM modes, float32/float64, noncontiguous inputs and counts, storage offsets, empty groups and dimensions, backward and double backward, current-stream behavior, graph capture/replay, compiled training with changing counts, and AOTI inference in a fresh process. The model suite compares with MACE for forward, backward, and double backward across three configurations and both dtypes, and checks compile and export. + +Additional checks passed outside this matrix: the initial single GPU case, Torch 2.7 JIT compilation/import with GPU visibility disabled, and three editable-install import/dependency tests with GPU visibility disabled. They are not counted again in the 192-check total. + +## Build, packaging, and BLAS dependency checks + +The HIP wheel built successfully using the branch's CMake configuration: + +- Artifact: `openequivariance-0.7.0-cp312-cp312-linux_x86_64.whl`. +- SHA256: `8ac6dc8cb3313298e613e7425e255fc369f89c797fbdefa1126778932fe76fb8`. +- Packaged libraries: `oeq_stable_hip.cpython-312-x86_64-linux-gnu.so` and `liboeq_stable_hip_aoti.so`. +- The wheel includes the JIT sources and excludes the `libtorch_hip.so` build stub. + +ELF inspection of both stable libraries and both JIT libraries found no direct cuBLAS, hipBLAS, or rocBLAS library dependency and no undefined vendor BLAS functions. All four libraries reference `aoti_torch_cuda_bmm_out`. Both installed ROCm Torch libraries export that symbol and `aoti_torch_get_current_cuda_stream`, retaining the CUDA spelling on AMD. + +The stable libraries depend on `libtorch_hip.so`, `libhiprtc.so.7`, and `libamdhip64.so.7`, along with CPU Torch and standard system libraries. The JIT builds link the Torch GPU library and the HIPRTC library found for their environments: `.so.7` for Torch 2.10 and `.so.6` for Torch 2.7. BLAS selection and handle ownership remain inside Torch. + +The actual editable installation also passed. Its Python source resolved to the checkout, its compiled HIP extension resolved to the virtual environment, and `_has_precompiled_extension()` found the compiled artifacts. The existing ROCm policy selected JIT as intended. The normal wheel was restored afterward. + +## Stable HIP loader policy + +The production loader still unconditionally selects JIT on ROCm. This validation does not enable automatic stable HIP loading. + +To exercise the compiled stable implementation, the wheel was extracted into `/tmp/oeq-cshim-amd/stable210-package`. Only the unconditional HIP precompiled-disable block was removed from that copy of the Python loader. Both shared libraries were verified byte-for-byte against the wheel. The tests asserted that the stable extension was selected and exercised the real registered operator, its existing autograd registration, and the companion AOTI library. No numerical operations were mocked. + +This establishes hardware correctness for the tested stable artifacts while keeping the production loader policy unchanged. + +## Setup and reproduction + +The new machine lacked Python packaging bootstrap support and Python development headers. Pip was bootstrapped inside the isolated environments, and `python3.12-dev` was installed for extension compilation. The initial CMake attempt stopped at the missing Python headers; compilation succeeded after installing them. No GPU test failed, and no production source changes were needed. + +The model-test environment used MACE 0.3.16, its pinned e3nn 0.4.4, NumPy 1.26.4, and pytest 9.1.1. As in the CUDA run, the test process allowed the built-in `slice` type for e3nn's packaged constants. The Torch 2.7 environment used e3nn 0.6.0 and NumPy 1.26.4. + +With the wheel and dependencies installed, the normal ROCm path is exercised by: + +```sh +ROCM_HOME=/opt/rocm/core-10.0 python -c 'import torch; torch.serialization.add_safe_globals([slice]); import pytest; raise SystemExit(pytest.main())' -q tests/import_test.py tests/group_gemm_test.py tests/symmetric_contraction_test.py +``` + +The actual runs used separate JIT and Inductor caches for each Torch version, and ran the focused and model suites separately. For stable-artifact validation, `PYTHONPATH` selected the isolated wheel copy described above. The older Torch environment ran the focused suite without the optional MACE dependency. Full environment snapshots, build logs, binary inspection output, and JUnit XML are retained in `/tmp/oeq-cshim-amd/logs/`. + +## Scope limits + +- Only one AMD GPU was available, so multi-GPU device guarding remains unverified. +- Automatic stable HIP loading remains disabled by existing policy. +- These are correctness and packaging results, with no performance claim. +- Direct GPU graph replay requires fixed ragged counts. + +At the end of GPU testing, `rocm-smi` reported zero GPU utilization, zero allocated VRAM percentage, and no KFD processes. The GPU reservation was released; remaining installation checks used disabled GPU visibility. diff --git a/sandbox/cshim_group_gemm/DESIGN.md b/sandbox/cshim_group_gemm/DESIGN.md index ed88da3c..19be961f 100644 --- a/sandbox/cshim_group_gemm/DESIGN.md +++ b/sandbox/cshim_group_gemm/DESIGN.md @@ -2,7 +2,7 @@ The production Torch extension now implements `libtorch_tp_jit::group_gemm` through `aoti_torch_cuda_bmm_out` on CUDA and ROCm. Both the stable extension and the source/JIT extension use this C entry point. The operator schema, fake implementation, output layouts, and custom backward formulas remain compatible. -The production implementation passed CUDA validation on an H100 with Torch 2.10's stable and JIT extensions and Torch 2.7's JIT extension. The complete matrix and its hardware limits are recorded in [the validation report](VALIDATION.md). +The production implementation passed CUDA validation on an H100 and ROCm validation on an MI300X. Both runs covered Torch 2.10's stable and JIT implementations and Torch 2.7's JIT implementation; stable HIP artifact testing used an isolated loader override because production still selects JIT on ROCm. The complete matrices and hardware limits are recorded in the [CUDA report](VALIDATION.md) and [AMD report](AMD_VALIDATION.md). ## Implementation @@ -53,13 +53,13 @@ The stable wheel build uses CPU LibTorch plus small GPU link stubs. [The stub](. CMake installs the stable libraries into the wheel's package directory. The previous absolute destination wrote them into the source tree, where the wheel's ignore rules excluded them. The loader now resolves the compiled extension's location to find its companion AOTI library; this also supports editable installations where Python sources and compiled libraries live in different directories. Build CI installs a normal wheel and asserts that the stable extension is selected before running the import checks. -The HIP CMake target is named `oeq_stable_hip`, matching its module entry point and expected artifact name. It explicitly links `hiprtc::hiprtc`. The existing Python loader still selects JIT compilation for HIP; enabling precompiled HIP loading is outside this change. The ROCm JIT path uses the same BMM helper as the stable CUDA build. +The HIP CMake target is named `oeq_stable_hip`, matching its module entry point and expected artifact name. It explicitly links `hiprtc::hiprtc`. The existing Python loader still selects JIT compilation for HIP; enabling precompiled HIP loading is outside this change. The ROCm JIT path uses the same BMM helper as the stable CUDA build. The HIP wheel built successfully, and its stable artifacts passed hardware tests with the existing loader restriction removed only in an isolated test copy. ## ROCm evidence PyTorch v2.10 adds the generated `c_shim_cuda.cpp` to `torch_hip` under `USE_ROCM`. The C symbol retains the `cuda` spelling on ROCm. The underlying GEMM implementations are converted to HIP BLAS calls, and Torch handles backend selection, including the double-precision fallback from hipBLASLt. [ROCm library construction](https://github.com/pytorch/pytorch/blob/v2.10.0/caffe2/CMakeLists.txt#L941), [HIP BLAS mappings](https://github.com/pytorch/pytorch/blob/v2.10.0/torch/utils/hipify/cuda_to_hip_mappings.py#L6826), [backend selection](https://github.com/pytorch/pytorch/blob/v2.10.0/aten/src/ATen/cuda/CUDABlas.cpp#L779). -This supports the implementation choice but does not establish AMD hardware correctness. ROCm execution has not been validated. +Hardware validation confirmed these entry points in Torch 2.10 / ROCm 7.1 and Torch 2.7 / ROCm 6.3. On an MI300X, both normal JIT configurations and the stable HIP artifacts passed the focused integration suite; both Torch 2.10 implementations also passed all model tests. The AMD matrix contains 192 successful checks and three skips requiring a second GPU. Binary inspection confirmed no direct vendor BLAS dependency in either stable library or either JIT library. [AMD validation and loader-policy details](AMD_VALIDATION.md). ## JAX scope @@ -69,7 +69,7 @@ JAX's public FFI provides buffers and a GPU stream, without an equivalent BMM C ## Validation -The stable wheel and JIT extensions compiled and loaded successfully. Wheel and editable-install checks verified that the installed stable libraries are found and that the build stub is excluded. GPU validation produced 192 successful checks across the three extension configurations and the surrounding model tests, with three skips for the device-guard test that requires two GPUs. [Results, environment corrections, and reproduction commands](VALIDATION.md). +The stable wheels and JIT extensions compiled and loaded successfully on CUDA and ROCm. Wheel and editable-install checks verified that the installed stable libraries are found and that the build stubs are excluded. Each platform's validation matrix produced 192 successful checks across the three extension configurations and the surrounding model tests, with three skips for the device-guard test that requires two GPUs. [CUDA results](VALIDATION.md), [AMD results](AMD_VALIDATION.md). Earlier local checks also passed for the new operator and both adapters against Torch 2.10 headers, the link stub, and the shared helper against Torch 2.4 headers. The unmodified production view helper passed eight cases using real Torch 2.10 CPU tensors with the GPU BMM entry point redirected to CPU BMM for that check: both modes/dtypes, nonzero storage offsets, empty groups, and output guard values. diff --git a/sandbox/cshim_group_gemm/VALIDATION.md b/sandbox/cshim_group_gemm/VALIDATION.md index e7b0d60f..c45a0565 100644 --- a/sandbox/cshim_group_gemm/VALIDATION.md +++ b/sandbox/cshim_group_gemm/VALIDATION.md @@ -2,6 +2,8 @@ The production implementation passed CUDA validation on 2026-09-13. The final test matrix contains 192 successful checks and three skips, all for the same device-guard test requiring two GPUs. These results cover the registered production operator and installed artifacts, not the earlier standalone prototype. +Subsequent ROCm hardware validation is recorded in [the AMD report](AMD_VALIDATION.md). + ## Source and artifact - Branch: `move-bmm-calls-to-stable-cshim`, based on main at `dbe854415da7771eba33195534c171adbca5677b`. @@ -64,7 +66,6 @@ The allowlist is only needed for the older e3nn dependency described above. The ## Scope limits -- ROCm has source-level support evidence in [the design](DESIGN.md), but no AMD hardware run has been performed. - Multi-GPU device guarding remains unverified on hardware because only one GPU was available. - Torch 2.4 received header checks only; GPU execution covered Torch 2.7 and 2.10. CUDA versions other than 12.8 were not exercised by this production test matrix. - Direct GPU graph replay requires fixed ragged counts, as described in the design. From 81b44510e24ddf02a894a7c9465c167f6a3aed7f Mon Sep 17 00:00:00 2001 From: asglover <140220574+asglover@users.noreply.github.com> Date: Sun, 13 Sep 2026 16:13:33 -0700 Subject: [PATCH 10/38] human changelog --- CHANGELOG.md | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 39a817d9..6e63ef12 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,9 +2,7 @@ **Added**: - Beta support for Intel XPUs. -- Route Torch grouped GEMM through the stable BMM C shim on CUDA and ROCm, - using Torch's current device stream and precision settings. Remove OEQ's - direct cuBLAS and rocBLAS dependencies. +- Removed OEQ's direct cuBLAS and rocBLAS dependencies. Use torch's c shim to instead. ### v0.7.0 (2026-09-10) **Added**: From 7ad9995493bf156bb9d9cd23bc9f363eabda668e Mon Sep 17 00:00:00 2001 From: asglover <140220574+asglover@users.noreply.github.com> Date: Sun, 13 Sep 2026 16:20:30 -0700 Subject: [PATCH 11/38] minimize diff, remove files that should not be tracked --- .github/workflows/verify_extension_build.yml | 3 +- sandbox/cshim_group_gemm/AMD_VALIDATION.md | 75 --------------- sandbox/cshim_group_gemm/DESIGN.md | 88 ----------------- sandbox/cshim_group_gemm/VALIDATION.md | 74 --------------- sandbox/cshim_group_gemm/group_gemm.cpp | 99 -------------------- sandbox/cshim_group_gemm/smoke_test.py | 71 -------------- tests/extension_linkage_test.py | 28 ++++++ tests/import_test.py | 28 ------ 8 files changed, 29 insertions(+), 437 deletions(-) delete mode 100644 sandbox/cshim_group_gemm/AMD_VALIDATION.md delete mode 100644 sandbox/cshim_group_gemm/DESIGN.md delete mode 100644 sandbox/cshim_group_gemm/VALIDATION.md delete mode 100644 sandbox/cshim_group_gemm/group_gemm.cpp delete mode 100644 sandbox/cshim_group_gemm/smoke_test.py create mode 100644 tests/extension_linkage_test.py diff --git a/.github/workflows/verify_extension_build.yml b/.github/workflows/verify_extension_build.yml index ecb5929a..c6b7c16a 100644 --- a/.github/workflows/verify_extension_build.yml +++ b/.github/workflows/verify_extension_build.yml @@ -29,11 +29,10 @@ jobs: sudo apt-get update sudo apt install nvidia-cuda-toolkit pip install -r .github/workflows/requirements_cuda_ci.txt - pip install "./openequivariance[jax]" + pip install -e "./openequivariance[jax]" - name: Test CUDA extension build via import run: | - python -c "import openequivariance as oeq; assert oeq.USE_PRECOMPILED_EXTENSION" pytest tests/import_test.py export OEQ_JIT_EXTENSION=1 diff --git a/sandbox/cshim_group_gemm/AMD_VALIDATION.md b/sandbox/cshim_group_gemm/AMD_VALIDATION.md deleted file mode 100644 index a928d44c..00000000 --- a/sandbox/cshim_group_gemm/AMD_VALIDATION.md +++ /dev/null @@ -1,75 +0,0 @@ -# AMD production validation - -ROCm validation passed on 2026-09-13 without changing the GEMM implementation. The final matrix contains 192 successful checks and three skips for the device-guard test requiring two GPUs. The normal ROCm JIT path passed with Torch 2.10 and 2.7. The stable HIP artifacts also passed with the isolated loader override described below. - -## Source and environment - -- Branch: `move-bmm-calls-to-stable-cshim`, tested at `4e3dcc8`. -- GPU: one AMD Instinct MI300X VF, 192 GB, `gfx942`. -- Driver reported by `rocm-smi`: `6.19.14.31400000`; Linux kernel `6.8.0-138-generic`. -- Python: 3.12.3. -- Build SDK: `/opt/rocm/core-10.0`, reporting HIP `7.15.26333`; host compiler GCC 13.3.0. -- Test environments: official Torch `2.10.0+rocm7.1` and `2.7.0+rocm6.3` wheels in separate virtual environments. -- Checkout, environments, caches, and logs: `/tmp/oeq-cshim-amd/` on the AMD machine. - -The NVIDIA machine was not used during this validation. - -## Results - -The first production grouped-GEMM case passed before running the focused suites. - -| Configuration | Import and grouped-GEMM tests | Symmetric-contraction tests | -| --- | --- | --- | -| Normal JIT path, Torch 2.10 / ROCm 7.1 | 48 passed, 1 skipped | 24 passed | -| Stable HIP artifacts, Torch 2.10 / ROCm 7.1, explicit test loader override | 48 passed, 1 skipped | 24 passed | -| Normal JIT path, Torch 2.7 / ROCm 6.3 | 48 passed, 1 skipped | Not run in this environment | - -The focused tests cover both GEMM modes, float32/float64, noncontiguous inputs and counts, storage offsets, empty groups and dimensions, backward and double backward, current-stream behavior, graph capture/replay, compiled training with changing counts, and AOTI inference in a fresh process. The model suite compares with MACE for forward, backward, and double backward across three configurations and both dtypes, and checks compile and export. - -Additional checks passed outside this matrix: the initial single GPU case, Torch 2.7 JIT compilation/import with GPU visibility disabled, and three editable-install import/dependency tests with GPU visibility disabled. They are not counted again in the 192-check total. - -## Build, packaging, and BLAS dependency checks - -The HIP wheel built successfully using the branch's CMake configuration: - -- Artifact: `openequivariance-0.7.0-cp312-cp312-linux_x86_64.whl`. -- SHA256: `8ac6dc8cb3313298e613e7425e255fc369f89c797fbdefa1126778932fe76fb8`. -- Packaged libraries: `oeq_stable_hip.cpython-312-x86_64-linux-gnu.so` and `liboeq_stable_hip_aoti.so`. -- The wheel includes the JIT sources and excludes the `libtorch_hip.so` build stub. - -ELF inspection of both stable libraries and both JIT libraries found no direct cuBLAS, hipBLAS, or rocBLAS library dependency and no undefined vendor BLAS functions. All four libraries reference `aoti_torch_cuda_bmm_out`. Both installed ROCm Torch libraries export that symbol and `aoti_torch_get_current_cuda_stream`, retaining the CUDA spelling on AMD. - -The stable libraries depend on `libtorch_hip.so`, `libhiprtc.so.7`, and `libamdhip64.so.7`, along with CPU Torch and standard system libraries. The JIT builds link the Torch GPU library and the HIPRTC library found for their environments: `.so.7` for Torch 2.10 and `.so.6` for Torch 2.7. BLAS selection and handle ownership remain inside Torch. - -The actual editable installation also passed. Its Python source resolved to the checkout, its compiled HIP extension resolved to the virtual environment, and `_has_precompiled_extension()` found the compiled artifacts. The existing ROCm policy selected JIT as intended. The normal wheel was restored afterward. - -## Stable HIP loader policy - -The production loader still unconditionally selects JIT on ROCm. This validation does not enable automatic stable HIP loading. - -To exercise the compiled stable implementation, the wheel was extracted into `/tmp/oeq-cshim-amd/stable210-package`. Only the unconditional HIP precompiled-disable block was removed from that copy of the Python loader. Both shared libraries were verified byte-for-byte against the wheel. The tests asserted that the stable extension was selected and exercised the real registered operator, its existing autograd registration, and the companion AOTI library. No numerical operations were mocked. - -This establishes hardware correctness for the tested stable artifacts while keeping the production loader policy unchanged. - -## Setup and reproduction - -The new machine lacked Python packaging bootstrap support and Python development headers. Pip was bootstrapped inside the isolated environments, and `python3.12-dev` was installed for extension compilation. The initial CMake attempt stopped at the missing Python headers; compilation succeeded after installing them. No GPU test failed, and no production source changes were needed. - -The model-test environment used MACE 0.3.16, its pinned e3nn 0.4.4, NumPy 1.26.4, and pytest 9.1.1. As in the CUDA run, the test process allowed the built-in `slice` type for e3nn's packaged constants. The Torch 2.7 environment used e3nn 0.6.0 and NumPy 1.26.4. - -With the wheel and dependencies installed, the normal ROCm path is exercised by: - -```sh -ROCM_HOME=/opt/rocm/core-10.0 python -c 'import torch; torch.serialization.add_safe_globals([slice]); import pytest; raise SystemExit(pytest.main())' -q tests/import_test.py tests/group_gemm_test.py tests/symmetric_contraction_test.py -``` - -The actual runs used separate JIT and Inductor caches for each Torch version, and ran the focused and model suites separately. For stable-artifact validation, `PYTHONPATH` selected the isolated wheel copy described above. The older Torch environment ran the focused suite without the optional MACE dependency. Full environment snapshots, build logs, binary inspection output, and JUnit XML are retained in `/tmp/oeq-cshim-amd/logs/`. - -## Scope limits - -- Only one AMD GPU was available, so multi-GPU device guarding remains unverified. -- Automatic stable HIP loading remains disabled by existing policy. -- These are correctness and packaging results, with no performance claim. -- Direct GPU graph replay requires fixed ragged counts. - -At the end of GPU testing, `rocm-smi` reported zero GPU utilization, zero allocated VRAM percentage, and no KFD processes. The GPU reservation was released; remaining installation checks used disabled GPU visibility. diff --git a/sandbox/cshim_group_gemm/DESIGN.md b/sandbox/cshim_group_gemm/DESIGN.md deleted file mode 100644 index 19be961f..00000000 --- a/sandbox/cshim_group_gemm/DESIGN.md +++ /dev/null @@ -1,88 +0,0 @@ -# Move Torch cuBLAS calls to the stable C shim - -The production Torch extension now implements `libtorch_tp_jit::group_gemm` through `aoti_torch_cuda_bmm_out` on CUDA and ROCm. Both the stable extension and the source/JIT extension use this C entry point. The operator schema, fake implementation, output layouts, and custom backward formulas remain compatible. - -The production implementation passed CUDA validation on an H100 and ROCm validation on an MI300X. Both runs covered Torch 2.10's stable and JIT implementations and Torch 2.7's JIT implementation; stable HIP artifact testing used an isolated loader override because production still selects JIT on ROCm. The complete matrices and hardware limits are recorded in the [CUDA report](VALIDATION.md) and [AMD report](AMD_VALIDATION.md). - -## Implementation - -The implementation lives in [group_mm.hpp](../../openequivariance/openequivariance/extension/group_mm.hpp), with input validation and output allocation in [torch_core.hpp](../../openequivariance/openequivariance/extension/torch_core.hpp). The algorithm is one BMM-out invocation per nonempty ragged group, using the same interleaved layout as the previous strided batched GEMMs. It does not need a vendor's heterogeneous grouped-GEMM API. - -Inputs are made contiguous by the registered operator, then passed as borrowed `AtenTensorHandle` values to the shared helper. The stable extension gets these handles from `torch::stable::Tensor::get()`. The source/JIT extension uses Torch's `tensor_pointer_to_tensor_handle` utility to borrow a handle to its local `at::Tensor`; that bridge is compiled against the installed Torch, as the existing JIT extension already is. ATen C++ objects do not cross the stable extension's ABI boundary. - -The helper creates temporary views with `aoti_torch__reinterpret_tensor`. These views retain the original storage and add their element offsets to the original tensor's storage offset. Their handles are released through RAII after each call, including when an error is returned. Original input and output handles remain owned by the caller. This replaces the prototype's `from_blob` construction and avoids reconstructing storage or querying pointer devices for every group. [Torch reinterpretation API](https://github.com/pytorch/pytorch/blob/v2.10.0/torch/csrc/inductor/aoti_torch/shim_common.cpp#L435). - -The operation checks GPU placement, matching devices and float32/float64 dtypes, tensor shapes, nonnegative dimensions, and mode 0 or 1. Ragged counts must be a CPU int64 vector with one entry per group. Counts may be noncontiguous; the operator makes a contiguous CPU copy before reading them. Counts must be nonnegative and sum to the input row count. Incremental bounds checking avoids overflowing the count sum. These checks establish the bounds needed by the storage-view helper. - -Output allocation uses Torch and starts at zero. Empty groups are skipped; in mode 1 this leaves the corresponding weight-gradient block zero. Zero batch, output, or contraction dimensions return the zero-initialized output without invoking BMM. Dimensions and offsets remain int64 throughout the grouped-GEMM implementation. - -## Layout mapping - -Let `b = batch_size`, `n = ragged_counts[i]`, and `o` be the sum of preceding counts. Offsets are relative to the logical beginning of each original tensor; strides and offsets are in elements. - -| Mode | Tensor | Element offset | View shape | View strides | -| --- | --- | --- | --- | --- | -| 0 | B / input | `o*b*k` | `[b,n,k]` | `[k,b*k,1]` | -| 0 | A / weights | `i*b*m*k` | `[b,k,m]` | `[m*k,1,k]` | -| 0 | C / output | `o*b*m` | `[b,n,m]` | `[m,b*m,1]` | -| 1 | A / left | `o*b*m` | `[b,m,n]` | `[m,1,b*m]` | -| 1 | B / right | `o*b*k` | `[b,n,k]` | `[k,b*k,1]` | -| 1 | C / output | `i*b*m*k` | `[b,m,k]` | `[m*k,k,1]` | - -The final operation for each group is `aoti_torch_cuda_bmm_out(output, left, right)`. The API overwrites the output with alpha 1 and beta 0. Torch may choose its own BLAS kernel or backend. The layout code requests views, but makes no unconditional promise about copies inside Torch's BMM implementation. [Torch BMM implementation](https://github.com/pytorch/pytorch/blob/v2.10.0/aten/src/ATen/native/cuda/Blas.cpp#L536). - -## Device, stream, precision, and autograd - -The registered operator guards A's device before making GPU contiguous copies or allocating the output. The stable build uses `torch::stable::accelerator::DeviceGuard`; the JIT build uses `c10::DeviceGuard`. Each restores the caller's current device on return. - -BMM uses Torch's current stream on that device. OEQ neither creates a BLAS handle nor changes a BLAS stream. Calls remain asynchronous and follow PyTorch's normal storage/stream lifetime contract. - -GPU graph capture records the group sizes read by the host loop during capture. Direct graph replay requires those counts to remain unchanged; changing group sizes requires recapture. Ordinary calls and compiled execution without graph replay read the counts on each invocation. - -Following Torch's precision, determinism, and preferred-BLAS settings is an intentional behavior change relative to main's independently created handle. There is no OEQ-specific precision override. Custom backward registration remains responsible for differentiation through the grouped operator; the internal BMM-out invocation does not replace that registration. - -## ABI and build paths - -The official declaration is in `torch/csrc/inductor/aoti_torch/generated/c_shim_cuda.h`. The backend-specific spelling is used instead of the deprecated `aoti_torch_bmm_out`. [Declaration](https://github.com/pytorch/pytorch/blob/v2.10.0/torch/csrc/inductor/aoti_torch/generated/c_shim_cuda.h#L74), [stable ABI policy](https://docs.pytorch.org/docs/main/notes/libtorch_stable_abi.html). - -Both `aoti_torch_cuda_bmm_out` and `aoti_torch__reinterpret_tensor` are present in the inspected PyTorch 2.4 headers, matching OEQ's documented source/JIT baseline. The stable build explicitly targets the 2.10 ABI and continues to use the pinned LibTorch headers/libraries. This uses Torch's stated ABI guarantees; it does not try to select cuBLAS versions or infer handle ownership from version queries. - -The source/JIT loader explicitly links the installed `torch_cuda` or `torch_hip` library supplying the BMM shim. CUDA also retains its driver/runtime/NVRTC dependencies; ROCm links hipRTC. All direct cuBLAS/rocBLAS calls, handle management, and corresponding OEQ link dependencies have been removed. - -The stable wheel build uses CPU LibTorch plus small GPU link stubs. [The stub](../../openequivariance/openequivariance/extension/stubs/stream.cpp) now declares and defines BMM-out using the official generated header, alongside the existing stream symbol. Stub bodies return failure if accidentally invoked. Only the real Torch GPU libraries are intended at runtime; CMake installs the OEQ targets, not the stubs. Both the Python extension and AOTI targets use this arrangement. - -CMake installs the stable libraries into the wheel's package directory. The previous absolute destination wrote them into the source tree, where the wheel's ignore rules excluded them. The loader now resolves the compiled extension's location to find its companion AOTI library; this also supports editable installations where Python sources and compiled libraries live in different directories. Build CI installs a normal wheel and asserts that the stable extension is selected before running the import checks. - -The HIP CMake target is named `oeq_stable_hip`, matching its module entry point and expected artifact name. It explicitly links `hiprtc::hiprtc`. The existing Python loader still selects JIT compilation for HIP; enabling precompiled HIP loading is outside this change. The ROCm JIT path uses the same BMM helper as the stable CUDA build. The HIP wheel built successfully, and its stable artifacts passed hardware tests with the existing loader restriction removed only in an isolated test copy. - -## ROCm evidence - -PyTorch v2.10 adds the generated `c_shim_cuda.cpp` to `torch_hip` under `USE_ROCM`. The C symbol retains the `cuda` spelling on ROCm. The underlying GEMM implementations are converted to HIP BLAS calls, and Torch handles backend selection, including the double-precision fallback from hipBLASLt. [ROCm library construction](https://github.com/pytorch/pytorch/blob/v2.10.0/caffe2/CMakeLists.txt#L941), [HIP BLAS mappings](https://github.com/pytorch/pytorch/blob/v2.10.0/torch/utils/hipify/cuda_to_hip_mappings.py#L6826), [backend selection](https://github.com/pytorch/pytorch/blob/v2.10.0/aten/src/ATen/cuda/CUDABlas.cpp#L779). - -Hardware validation confirmed these entry points in Torch 2.10 / ROCm 7.1 and Torch 2.7 / ROCm 6.3. On an MI300X, both normal JIT configurations and the stable HIP artifacts passed the focused integration suite; both Torch 2.10 implementations also passed all model tests. The AMD matrix contains 192 successful checks and three skips requiring a second GPU. Binary inspection confirmed no direct vendor BLAS dependency in either stable library or either JIT library. [AMD validation and loader-policy details](AMD_VALIDATION.md). - -## JAX scope - -JAX does not call this grouped-GEMM helper and has no BLAS dependency in its extension target. The replacement is entirely within the Torch frontend. The common CUDA/HIP kernel compilation backend gains no Torch dependency. - -JAX's public FFI provides buffers and a GPU stream, without an equivalent BMM C shim. Future JAX grouped GEMM would require native JAX graph operations where shapes permit, or a separate FFI BLAS integration. [JAX FFI](https://docs.jax.dev/en/latest/ffi.html#ffi-calls-on-a-gpu). - -## Validation - -The stable wheels and JIT extensions compiled and loaded successfully on CUDA and ROCm. Wheel and editable-install checks verified that the installed stable libraries are found and that the build stubs are excluded. Each platform's validation matrix produced 192 successful checks across the three extension configurations and the surrounding model tests, with three skips for the device-guard test that requires two GPUs. [CUDA results](VALIDATION.md), [AMD results](AMD_VALIDATION.md). - -Earlier local checks also passed for the new operator and both adapters against Torch 2.10 headers, the link stub, and the shared helper against Torch 2.4 headers. The unmodified production view helper passed eight cases using real Torch 2.10 CPU tensors with the GPU BMM entry point redirected to CPU BMM for that check: both modes/dtypes, nonzero storage offsets, empty groups, and output guard values. - -[The integration tests](../../tests/group_gemm_test.py) call the real registered operator. They cover both modes/dtypes, noncontiguous inputs and counts, nonzero input storage offsets, empty and singleton dimensions, varied group sizes, backward and double-backward gradients, the current stream, device guarding, invalid counts, graph replay, compiled training, and AOTI inference in a fresh process that loads only the exported OEQ library. The device-guard test requires two GPUs. The same test file supports CUDA and ROCm. - -[The import tests](../../tests/import_test.py) also inspect the extension and AOTI library's ELF dependencies and undefined symbols to check that OEQ has no direct vendor BLAS dependency. These checks run in the existing build-verification workflow for precompiled and JIT imports. - -The existing symmetric-contraction integration tests passed all 24 cases for both the stable and JIT extensions under Torch 2.10, including comparison with MACE and higher-order derivatives. No end-to-end performance claim is made by this correctness validation. - -## Earlier experiments - -The original [standalone prototype](group_gemm.cpp) and its [smoke script](smoke_test.py) remain as historical experiment artifacts. They use blob views rather than the production helper's views of existing storage. - -The prototype passed one float32 forward case on an H100 PCIe with Torch 2.7.0 / CUDA 12.8: counts `[2,0,5,1]`, batch 3, m 4, k 5, a nondefault stream, and maximum absolute error `4.37e-7` against a CPU float64 reference. Earlier ctypes calls to the BMM C shim passed 40 small cases across both dtypes/modes. Four selected BMM-only profiles showed no copy or GPU allocation inside the call; they did not measure view creation overhead. - -The earlier direct-cuBLAS experiments explain the migration: matching handle/function owners worked, whereas some separately loaded foreign-owner combinations failed or crashed despite successful version queries. They do not establish interchangeability of private BLAS handles. Related: [PR #206](https://github.com/PASSIONLab/OpenEquivariance/pull/206). diff --git a/sandbox/cshim_group_gemm/VALIDATION.md b/sandbox/cshim_group_gemm/VALIDATION.md deleted file mode 100644 index c45a0565..00000000 --- a/sandbox/cshim_group_gemm/VALIDATION.md +++ /dev/null @@ -1,74 +0,0 @@ -# Production validation - -The production implementation passed CUDA validation on 2026-09-13. The final test matrix contains 192 successful checks and three skips, all for the same device-guard test requiring two GPUs. These results cover the registered production operator and installed artifacts, not the earlier standalone prototype. - -Subsequent ROCm hardware validation is recorded in [the AMD report](AMD_VALIDATION.md). - -## Source and artifact - -- Branch: `move-bmm-calls-to-stable-cshim`, based on main at `dbe854415da7771eba33195534c171adbca5677b`. -- Production source: `d2eac4e`; integration tests: `e4df36c`. -- Wheel: `openequivariance-0.7.0-cp310-cp310-linux_x86_64.whl`. -- Wheel SHA256: `f536c53f1ef2358b6f562931c30c3f55823c32123424c153c96cb808250ea97a`. -- Remote checkout and environments: `/tmp/oeq-cshim-production/`, separate from the workspace-allocation checkout and user-installed packages. - -## Build and packaging results - -All checks in this table ran with `CUDA_VISIBLE_DEVICES=""`. They establish successful compilation, loading, and packaging; they do not establish GPU numerical correctness. - -| Configuration | Result | -| --- | --- | -| Stable wheel, Torch 2.10.0+cu128, Python 3.10 | Both production libraries built; 3 import/dependency tests passed | -| Editable install, Torch 2.10.0+cu128 | Stable extension selected with Python sources and libraries in separate directories; 3 import/dependency tests passed | -| Source/JIT extension, Torch 2.10.0+cu128 | Compiled against the installed Torch; 3 import/dependency tests passed | -| Source/JIT extension, Torch 2.7.0, CUDA 12.8 | Compiled against the installed Torch; 3 import/dependency tests passed | - -The wheel contains its Python extension, companion AOTI library, and JIT sources. It contains no `libtorch_cuda.so` build stub. ELF checks found no direct vendor BLAS library dependency or undefined vendor BLAS function in either production library or either JIT build. Both installed Torch GPU libraries export `aoti_torch_cuda_bmm_out`. - -Validation found and fixed a packaging defect: the absolute CMake install destination placed libraries in the source tree, where wheel ignore rules excluded them. The relative install destination includes the libraries in the wheel. Resolving the extension's actual import location also fixes stable-library discovery for editable installations. - -## GPU validation - -Hardware: one NVIDIA H100 PCIe, driver 580.105.08. Both Torch environments report CUDA 12.8. The first production-wheel forward case passed before running the larger suites. - -| Configuration | Import and grouped-GEMM tests | Symmetric-contraction tests | -| --- | --- | --- | -| Stable wheel, Torch 2.10.0+cu128 | 48 passed, 1 skipped | 24 passed | -| Source/JIT, Torch 2.10.0+cu128 | 48 passed, 1 skipped | 24 passed | -| Source/JIT, installed Torch 2.7.0 / CUDA 12.8 | 48 passed, 1 skipped | Not run in this environment | - -The stable-wheel result combines 46 initial passes with two successful AOTI reruns after correcting fresh-process loader initialization. The Torch 2.7 result is a complete rerun after correcting test-environment dependencies. The table reports final per-case outcomes; it does not count preliminary or repeated passes twice. - -The 46 grouped-GEMM cases cover layouts, empty groups and dimensions, both modes and dtypes, backward and double backward, the current stream, device guarding, invalid counts, graph capture/replay, compiled training with changing counts, and AOTI inference. Each AOTI case loads the exported OEQ library in a new Python process without importing the OEQ Python package, then compares against a CPU reference. Three additional tests check import, successful extension loading, and binary dependencies. - -The existing symmetric-contraction suite compares with MACE for float32/float64 forward, backward, and double backward across three configurations. It also checks compile and export. Both the stable and JIT extensions passed all 24 cases. - -## Issues resolved during validation - -The wheel and editable-discovery fixes above were production changes. GPU execution required no further changes to the GEMM implementation. - -The fresh-process AOTI test exposed a Torch 2.10 loader initialization issue: its package loader accesses `torch._inductor.codecache` before importing it. The subprocess now explicitly imports that module before loading the package. Both AOTI modes then passed for the stable wheel and both JIT builds. - -The Torch 2.10 model-test environment uses MACE 0.3.16 and its pinned e3nn 0.4.4. That e3nn version loads packaged constants containing Python `slice` objects. The test process allowed that specific built-in type with `torch.serialization.add_safe_globals([slice])`; no production code or global loading policy changed. NumPy was 1.26.4 and pytest was 9.1.1. - -The Torch 2.7 environment initially inherited NetworkX 2.4 from the system; its import failed on NumPy's removed `np.int` alias before compilation. Installing NetworkX 3.4.2, SciPy 1.15.3, and SymPy 1.13.3 inside the isolated environment corrected its compiler dependencies. Its final environment used e3nn 0.6.0 and NumPy 1.26.4. System packages and the other checkout were untouched. - -## Reproduction - -After installing the wheel and test dependencies in a matching GPU Torch environment, run: - -```sh -python -c 'import torch; torch.serialization.add_safe_globals([slice]); import pytest; raise SystemExit(pytest.main())' -q tests/import_test.py tests/group_gemm_test.py tests/symmetric_contraction_test.py -OEQ_JIT_EXTENSION=1 python -c 'import torch; torch.serialization.add_safe_globals([slice]); import pytest; raise SystemExit(pytest.main())' -q tests/import_test.py tests/group_gemm_test.py tests/symmetric_contraction_test.py -``` - -The allowlist is only needed for the older e3nn dependency described above. The actual runs used separate JIT and Inductor cache directories for each Torch environment and ran the focused and model suites separately. The Torch 2.7 environment also set `PYTHONNOUSERSITE=1` to exclude the shared machine's user packages. - -## Scope limits - -- Multi-GPU device guarding remains unverified on hardware because only one GPU was available. -- Torch 2.4 received header checks only; GPU execution covered Torch 2.7 and 2.10. CUDA versions other than 12.8 were not exercised by this production test matrix. -- Direct GPU graph replay requires fixed ragged counts, as described in the design. -- No end-to-end performance claim is made by these correctness and packaging checks. - -Detailed logs, environment snapshots, and JUnit XML are retained under `/tmp/oeq-cshim-production/logs/` on the validation machine. At the end of testing, the GPU reported zero memory use and no compute processes; its status reservation was released. diff --git a/sandbox/cshim_group_gemm/group_gemm.cpp b/sandbox/cshim_group_gemm/group_gemm.cpp deleted file mode 100644 index 1f7462a7..00000000 --- a/sandbox/cshim_group_gemm/group_gemm.cpp +++ /dev/null @@ -1,99 +0,0 @@ -// Prototype for the PyTorch CUDA backend. Uses only PyTorch's stable C ABI. -#include - -#include -#include -#include -#include -#include -#include - -namespace { - -void check(AOTITorchError status) { - if (status != AOTI_TORCH_SUCCESS) - throw std::runtime_error("PyTorch C shim failed"); -} - -using Tensor = std::unique_ptr< - std::remove_pointer_t, - decltype(&aoti_torch_delete_tensor_object)>; -using DeviceGuard = std::unique_ptr< - std::remove_pointer_t, - decltype(&aoti_torch_delete_cuda_guard)>; - -// Only tensor metadata is created; the caller retains ownership of the buffer. -Tensor view(void* data, std::array sizes, - std::array strides, int32_t dtype, int32_t device) { - AtenTensorHandle tensor = nullptr; - check(aoti_torch_create_tensor_from_blob_v2( - data, 3, sizes.data(), strides.data(), 0, dtype, - aoti_torch_device_type_cuda(), device, &tensor, - aoti_torch_layout_strided(), nullptr, 0)); - return Tensor(tensor, aoti_torch_delete_tensor_object); -} - -} // namespace - -template -void group_gemm_cshim( - T* A, T* B, T* C, const int64_t* ragged_counts, int num_groups, - int64_t batch, int64_t m, int64_t k, int ragged_inner, int32_t device) { - static_assert(std::is_same_v || std::is_same_v); - const int32_t dtype = std::is_same_v - ? aoti_torch_dtype_float32() : aoti_torch_dtype_float64(); - - CUDAGuardHandle raw_guard = nullptr; - check(aoti_torch_create_cuda_guard(device, &raw_guard)); - DeviceGuard guard(raw_guard, aoti_torch_delete_cuda_guard); - - int64_t offset = 0; - for (int i = 0; i < num_groups; ++i) { - const int64_t n = ragged_counts[i]; - if (n == 0) continue; // Preserve the original empty-group behavior. - - if (ragged_inner == 0) { - // [batch, n, k] @ [batch, k, m] -> [batch, n, m] - auto input = view(B + offset * batch * k, - {batch, n, k}, {k, batch * k, 1}, dtype, device); - auto weight = view(A + i * batch * m * k, - {batch, k, m}, {m * k, 1, k}, dtype, device); - auto output = view(C + offset * batch * m, - {batch, n, m}, {m, batch * m, 1}, dtype, device); - check(aoti_torch_cuda_bmm_out(output.get(), input.get(), weight.get())); - } else { - // [batch, m, n] @ [batch, n, k] -> [batch, m, k] - auto left = view(A + offset * batch * m, - {batch, m, n}, {m, 1, batch * m}, dtype, device); - auto right = view(B + offset * batch * k, - {batch, n, k}, {k, batch * k, 1}, dtype, device); - auto output = view(C + i * batch * m * k, - {batch, m, k}, {m * k, k, 1}, dtype, device); - check(aoti_torch_cuda_bmm_out(output.get(), left.get(), right.get())); - } - offset += n; - } -} - -// Small ctypes entry point for the smoke test, not production registration code. -extern "C" int oeq_group_gemm_cshim( - int dtype, void* A, void* B, void* C, const int64_t* counts, - int groups, int64_t batch, int64_t m, int64_t k, int inner, int32_t device) { - try { - if (dtype == 0) { - group_gemm_cshim(static_cast(A), static_cast(B), - static_cast(C), counts, groups, - batch, m, k, inner, device); - } else if (dtype == 1) { - group_gemm_cshim(static_cast(A), static_cast(B), - static_cast(C), counts, groups, - batch, m, k, inner, device); - } else { - throw std::runtime_error("Expected float32 or float64"); - } - return 0; - } catch (const std::exception& error) { - std::fprintf(stderr, "%s\n", error.what()); - return 1; - } -} diff --git a/sandbox/cshim_group_gemm/smoke_test.py b/sandbox/cshim_group_gemm/smoke_test.py deleted file mode 100644 index 931d9698..00000000 --- a/sandbox/cshim_group_gemm/smoke_test.py +++ /dev/null @@ -1,71 +0,0 @@ -"""Compile the C++ prototype and run one float32 CUDA correctness case.""" -import ctypes -from pathlib import Path -import subprocess - -import torch - - -def main(): - root = Path(__file__).resolve().parent - torch_root = Path(torch.__file__).resolve().parent - library = root / "group_gemm.so" - subprocess.run( - [ - "g++", "-std=c++17", "-O2", "-shared", "-fPIC", "-DUSE_CUDA", - f"-I{torch_root / 'include'}", str(root / "group_gemm.cpp"), - f"-L{torch_root / 'lib'}", f"-Wl,-rpath,{torch_root / 'lib'}", - "-Wl,--no-undefined", "-ltorch_cuda", "-ltorch_cpu", - "-o", str(library), - ], - check=True, - ) - lib = ctypes.CDLL(str(library)) - call = lib.oeq_group_gemm_cshim - i64 = ctypes.c_int64 - call.argtypes = [ - ctypes.c_int, ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p, - ctypes.POINTER(i64), ctypes.c_int, i64, i64, i64, - ctypes.c_int, ctypes.c_int32, - ] - call.restype = ctypes.c_int - - # One case includes differently sized groups, an empty group, interleaved - # batches, and a nondefault stream. The reference runs on the CPU. - counts = [2, 0, 5, 1] - batch, m, k = 3, 4, 5 - torch.manual_seed(123) - weights_cpu = torch.randn(len(counts), batch, m, k) - input_cpu = torch.randn(sum(counts), batch, k) - expected = torch.empty(sum(counts), batch, m, dtype=torch.float64) - offset = 0 - for i, n in enumerate(counts): - expected[offset:offset + n] = torch.einsum( - "bmk,nbk->nbm", weights_cpu[i].double(), - input_cpu[offset:offset + n].double(), - ) - offset += n - - torch.cuda.set_device(0) - stream = torch.cuda.Stream() - with torch.cuda.stream(stream): - weights = weights_cpu.cuda() - inputs = input_cpu.cuda() - output = torch.full(expected.shape, float("nan"), device="cuda") - status = call( - 0, weights.data_ptr(), inputs.data_ptr(), output.data_ptr(), - (i64 * len(counts))(*counts), len(counts), batch, m, k, 0, 0, - ) - assert status == 0, f"C shim prototype returned {status}" - stream.synchronize() - actual = output.cpu().double() - torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-5) - print( - f"PASS: one float32 case, counts={counts}, batch={batch}, m={m}, k={k}; " - f"max absolute error={(actual - expected).abs().max().item():.3g}; " - f"GPU={torch.cuda.get_device_name(0)}, torch={torch.__version__}" - ) - - -if __name__ == "__main__": - main() diff --git a/tests/extension_linkage_test.py b/tests/extension_linkage_test.py new file mode 100644 index 00000000..77fbb8bc --- /dev/null +++ b/tests/extension_linkage_test.py @@ -0,0 +1,28 @@ +import re +import shutil +import subprocess + +import pytest + + +def test_extension_has_no_direct_blas_dependency(): + if not shutil.which("readelf"): + pytest.skip("readelf is required to inspect ELF dependencies") + import openequivariance + from openequivariance._torch.extlib import extension_module + + paths = {extension_module.__file__, openequivariance.torch_ext_so_path()} + for path in paths: + dynamic = subprocess.check_output(["readelf", "--dynamic", path], text=True) + dependencies = [line for line in dynamic.splitlines() if "(NEEDED)" in line] + assert not any( + re.search(r"lib(cublas|hipblas|rocblas)", line) for line in dependencies + ) + symbols = subprocess.check_output( + ["readelf", "--dyn-syms", "--wide", path], text=True + ) + undefined = [line for line in symbols.splitlines() if " UND " in line] + assert not any( + re.search(r"\b(cublas|hipblas|rocblas)[A-Za-z_]", line) + for line in undefined + ) diff --git a/tests/import_test.py b/tests/import_test.py index aa53630e..bf26af31 100644 --- a/tests/import_test.py +++ b/tests/import_test.py @@ -1,9 +1,4 @@ from importlib.metadata import version -import re -import shutil -import subprocess - -import pytest def test_import(): @@ -19,26 +14,3 @@ def test_extension_built(): assert BUILT_EXTENSION_ERROR is None assert BUILT_EXTENSION - - -def test_extension_has_no_direct_blas_dependency(): - if not shutil.which("readelf"): - pytest.skip("readelf is required to inspect ELF dependencies") - import openequivariance - from openequivariance._torch.extlib import extension_module - - paths = {extension_module.__file__, openequivariance.torch_ext_so_path()} - for path in paths: - dynamic = subprocess.check_output(["readelf", "--dynamic", path], text=True) - dependencies = [line for line in dynamic.splitlines() if "(NEEDED)" in line] - assert not any( - re.search(r"lib(cublas|hipblas|rocblas)", line) for line in dependencies - ) - symbols = subprocess.check_output( - ["readelf", "--dyn-syms", "--wide", path], text=True - ) - undefined = [line for line in symbols.splitlines() if " UND " in line] - assert not any( - re.search(r"\b(cublas|hipblas|rocblas)[A-Za-z_]", line) - for line in undefined - ) From c4e14d59c064f9b858844863de72208eebd8d423 Mon Sep 17 00:00:00 2001 From: asglover <140220574+asglover@users.noreply.github.com> Date: Sun, 13 Sep 2026 16:23:11 -0700 Subject: [PATCH 12/38] specify cuda c shim --- CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 6e63ef12..fe2ed09b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,7 +2,7 @@ **Added**: - Beta support for Intel XPUs. -- Removed OEQ's direct cuBLAS and rocBLAS dependencies. Use torch's c shim to instead. +- Removed OEQ's direct cuBLAS and rocBLAS dependencies. Use torch's cuda c shim to instead. ### v0.7.0 (2026-09-10) **Added**: From bb655b24e6f2ab736fc9c14abddc7b13e12601c9 Mon Sep 17 00:00:00 2001 From: asglover <140220574+asglover@users.noreply.github.com> Date: Sun, 13 Sep 2026 17:01:51 -0700 Subject: [PATCH 13/38] remove target version to minimize PR, this will go on a separate PR --- openequivariance/CMakeLists.txt | 2 -- 1 file changed, 2 deletions(-) diff --git a/openequivariance/CMakeLists.txt b/openequivariance/CMakeLists.txt index f4eb5884..7c6d52ec 100644 --- a/openequivariance/CMakeLists.txt +++ b/openequivariance/CMakeLists.txt @@ -67,7 +67,6 @@ function(add_stable_extension target_name backend_define link_libraries) # Enforce CXX11 ABI to match LibTorch target_compile_definitions(${target_name} PRIVATE ${backend_define}=1 - TORCH_TARGET_VERSION=0x020a000000000000ULL _GLIBCXX_USE_CXX11_ABI=1 INCLUDE_NB_EXTENSION ) @@ -98,7 +97,6 @@ function(add_stable_extension target_name backend_define link_libraries) target_compile_definitions(${aoti_target_name} PRIVATE ${backend_define}=1 - TORCH_TARGET_VERSION=0x020a000000000000ULL _GLIBCXX_USE_CXX11_ABI=1 ) From 165aa90a242e48b0ffe591bc5cc744b60106029e Mon Sep 17 00:00:00 2001 From: asglover <140220574+asglover@users.noreply.github.com> Date: Sun, 13 Sep 2026 17:19:58 -0700 Subject: [PATCH 14/38] simply PR my moving build issues to a separate pr --- openequivariance/CMakeLists.txt | 2 +- .../openequivariance/_torch/extlib/__init__.py | 13 +------------ 2 files changed, 2 insertions(+), 13 deletions(-) diff --git a/openequivariance/CMakeLists.txt b/openequivariance/CMakeLists.txt index 7c6d52ec..b7c6570f 100644 --- a/openequivariance/CMakeLists.txt +++ b/openequivariance/CMakeLists.txt @@ -52,7 +52,7 @@ set(OEQ_SOURCES ${EXT_JSON_DIR}/json11.cpp ) -set(OEQ_INSTALL_DIR "openequivariance/_torch/extlib") +set(OEQ_INSTALL_DIR "${CMAKE_CURRENT_SOURCE_DIR}/openequivariance/_torch/extlib") function(add_stable_extension target_name backend_define link_libraries) # Create nanobind extension diff --git a/openequivariance/openequivariance/_torch/extlib/__init__.py b/openequivariance/openequivariance/_torch/extlib/__init__.py index 19720e3b..0861f2e0 100644 --- a/openequivariance/openequivariance/_torch/extlib/__init__.py +++ b/openequivariance/openequivariance/_torch/extlib/__init__.py @@ -5,7 +5,6 @@ import warnings import sysconfig import contextlib -import importlib.util from pathlib import Path from packaging.version import Version @@ -225,16 +224,6 @@ def load_precompiled_extension(): ) -def _has_precompiled_extension(): - backend = BACKEND - spec = importlib.util.find_spec(f"{__name__}.oeq_stable_{backend}") - return ( - spec is not None - and spec.origin is not None - and Path(spec.origin).with_name(f"liboeq_stable_{backend}_aoti.so").is_file() - ) - - USE_PRECOMPILED_EXTENSION = True WARNING_MESSAGE = "" @@ -256,7 +245,7 @@ def _has_precompiled_extension(): AOTI_SO_NAME = f"liboeq_stable_{BACKEND}_aoti.so" -if not _has_precompiled_extension(): +if not os.path.exists(os.path.join(os.path.dirname(__file__), AOTI_SO_NAME)): WARNING_MESSAGE += "Precompiled extension shared object not found.\n" USE_PRECOMPILED_EXTENSION = False From 5512ddc108f2b6cc85535447da0bfd3369372061 Mon Sep 17 00:00:00 2001 From: asglover <140220574+asglover@users.noreply.github.com> Date: Sun, 13 Sep 2026 17:31:45 -0700 Subject: [PATCH 15/38] simplify pr and move unrelated fixes to separate prs --- openequivariance/CMakeLists.txt | 2 +- openequivariance/openequivariance/_torch/extlib/__init__.py | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/openequivariance/CMakeLists.txt b/openequivariance/CMakeLists.txt index b7c6570f..7f517ab6 100644 --- a/openequivariance/CMakeLists.txt +++ b/openequivariance/CMakeLists.txt @@ -163,7 +163,7 @@ if(hip_FOUND) hip::host hiprtc::hiprtc ) - add_stable_extension(oeq_stable_hip HIP_BACKEND "${HIP_LINK_LIBS}") + add_stable_extension(torch_stable_hip HIP_BACKEND "${HIP_LINK_LIBS}") endif() set(OEQ_SYCL_FOUND FALSE) diff --git a/openequivariance/openequivariance/_torch/extlib/__init__.py b/openequivariance/openequivariance/_torch/extlib/__init__.py index 0861f2e0..4ec88480 100644 --- a/openequivariance/openequivariance/_torch/extlib/__init__.py +++ b/openequivariance/openequivariance/_torch/extlib/__init__.py @@ -148,9 +148,9 @@ def load_jit_extension(): extra_cflags.append("-DCUDA_BACKEND") elif BACKEND == "hip": - hip_lib_dirs = library_paths("cuda") - extra_link_args.append("-Wl,-rpath," + hip_lib_dirs[0]) - extra_link_args.extend("-L" + directory for directory in hip_lib_dirs) + torch_libs = library_paths("cuda")[0] + extra_link_args.append("-Wl,-rpath," + torch_libs) + extra_link_args.extend("-L" + path for path in library_paths("cuda")) extra_link_args.extend(["-ltorch_hip", "-lhiprtc"]) extra_cflags.append("-DHIP_BACKEND") elif BACKEND == "sycl": From 3c723c6ca4d3636cc0e43d13c62a2651f8230e0d Mon Sep 17 00:00:00 2001 From: asglover <140220574+asglover@users.noreply.github.com> Date: Sun, 13 Sep 2026 17:53:57 -0700 Subject: [PATCH 16/38] simplify pr. We're not exposing gemm as an external interface, just test it through symmetric tensor product. The test that we don't rely on blas is over testing. --- tests/extension_linkage_test.py | 28 --- tests/group_gemm_test.py | 314 -------------------------------- 2 files changed, 342 deletions(-) delete mode 100644 tests/extension_linkage_test.py delete mode 100644 tests/group_gemm_test.py diff --git a/tests/extension_linkage_test.py b/tests/extension_linkage_test.py deleted file mode 100644 index 77fbb8bc..00000000 --- a/tests/extension_linkage_test.py +++ /dev/null @@ -1,28 +0,0 @@ -import re -import shutil -import subprocess - -import pytest - - -def test_extension_has_no_direct_blas_dependency(): - if not shutil.which("readelf"): - pytest.skip("readelf is required to inspect ELF dependencies") - import openequivariance - from openequivariance._torch.extlib import extension_module - - paths = {extension_module.__file__, openequivariance.torch_ext_so_path()} - for path in paths: - dynamic = subprocess.check_output(["readelf", "--dynamic", path], text=True) - dependencies = [line for line in dynamic.splitlines() if "(NEEDED)" in line] - assert not any( - re.search(r"lib(cublas|hipblas|rocblas)", line) for line in dependencies - ) - symbols = subprocess.check_output( - ["readelf", "--dyn-syms", "--wide", path], text=True - ) - undefined = [line for line in symbols.splitlines() if " UND " in line] - assert not any( - re.search(r"\b(cublas|hipblas|rocblas)[A-Za-z_]", line) - for line in undefined - ) diff --git a/tests/group_gemm_test.py b/tests/group_gemm_test.py deleted file mode 100644 index 8df3d5bf..00000000 --- a/tests/group_gemm_test.py +++ /dev/null @@ -1,314 +0,0 @@ -import importlib -import subprocess -import sys - -import pytest -import torch - - -pytestmark = pytest.mark.skipif( - not torch.cuda.is_available(), reason="A CUDA or ROCm GPU is required" -) - - -@pytest.fixture(scope="module") -def group_gemm(): - importlib.import_module("openequivariance._torch.symmetric_contraction") - return torch.ops.libtorch_tp_jit.group_gemm - - -def reference(A, B, counts, inner): - pieces = [] - offset = 0 - for i, n in enumerate(counts): - if inner == 0: - pieces.append(torch.einsum("bmk,nbk->nbm", A[i], B[offset : offset + n])) - else: - pieces.append( - torch.einsum( - "nbm,nbk->bmk", A[offset : offset + n], B[offset : offset + n] - ) - ) - offset += n - if inner == 0: - return torch.cat(pieces, dim=0) - return torch.stack(pieces, dim=0) - - -def make_input(shape, dtype, layout="contiguous", device="cuda"): - values = torch.arange(1, 1 + torch.Size(shape).numel(), dtype=dtype, device="cpu") - values = (values.remainder(17) - 8).reshape(shape) / 8 - if layout == "offset": - storage = torch.empty(values.numel() + 7, dtype=dtype, device=device) - result = storage[3 : 3 + values.numel()].view(shape) - result.copy_(values) - return result - if layout == "noncontiguous": - storage = torch.empty((*shape, 2), dtype=dtype, device=device) - result = storage[..., 0] - result.copy_(values) - return result - return values.to(device) - - -@pytest.mark.parametrize("dtype", [torch.float32, torch.float64]) -@pytest.mark.parametrize("inner", [0, 1]) -@pytest.mark.parametrize("layout", ["contiguous", "offset", "noncontiguous"]) -def test_group_gemm_matches_reference(group_gemm, dtype, inner, layout): - counts = [2, 0, 3] - batch, m, k = 2, 3, 4 - A_shape = (len(counts), batch, m, k) if inner == 0 else (sum(counts), batch, m) - A = make_input(A_shape, dtype, layout) - B = make_input((sum(counts), batch, k), dtype, layout) - expected = reference(A.cpu().double(), B.cpu().double(), counts, inner) - counts_tensor = torch.tensor([n for n in counts for _ in range(2)], device="cpu")[ - ::2 - ] - - actual = group_gemm(A, B, counts_tensor, len(counts), batch, m, k, inner) - - torch.testing.assert_close(actual.cpu().double(), expected, rtol=1e-5, atol=1e-5) - assert actual.dtype == dtype - assert actual.device == A.device - - -@pytest.mark.parametrize("inner", [0, 1]) -@pytest.mark.parametrize( - "counts,batch,m,k", - [ - ([0, 0], 2, 3, 4), - ([], 2, 3, 4), - ([2, 0, 3], 0, 3, 4), - ([2, 0, 3], 2, 0, 4), - ([2, 0, 3], 2, 3, 0), - ], -) -def test_group_gemm_empty_dimensions(group_gemm, inner, counts, batch, m, k): - A_shape = (len(counts), batch, m, k) if inner == 0 else (sum(counts), batch, m) - A = torch.empty(A_shape, device="cuda", dtype=torch.float64) - B = torch.empty((sum(counts), batch, k), device="cuda", dtype=torch.float64) - actual = group_gemm( - A, - B, - torch.tensor(counts, dtype=torch.int64, device="cpu"), - len(counts), - batch, - m, - k, - inner, - ) - expected_shape = ( - (sum(counts), batch, m) if inner == 0 else (len(counts), batch, m, k) - ) - assert actual.shape == expected_shape - torch.testing.assert_close(actual, torch.zeros_like(actual)) - - -@pytest.mark.parametrize("inner", [0, 1]) -def test_group_gemm_backward(group_gemm, inner): - counts = [2, 0, 3] - batch, m, k = 2, 3, 4 - A_shape = (len(counts), batch, m, k) if inner == 0 else (sum(counts), batch, m) - A = make_input(A_shape, torch.float64, "noncontiguous").requires_grad_() - B = make_input((sum(counts), batch, k), torch.float64, "offset").requires_grad_() - A_ref = A.detach().cpu().requires_grad_() - B_ref = B.detach().cpu().requires_grad_() - expected = reference(A_ref, B_ref, counts, inner) - actual = group_gemm( - A, B, torch.tensor(counts, device="cpu"), len(counts), batch, m, k, inner - ) - grad = make_input(actual.shape, torch.float64, "noncontiguous") - - actual_grads = torch.autograd.grad(actual, (A, B), grad) - expected_grads = torch.autograd.grad(expected, (A_ref, B_ref), grad.cpu()) - - for actual_grad, expected_grad in zip(actual_grads, expected_grads): - torch.testing.assert_close( - actual_grad.cpu(), expected_grad, rtol=1e-10, atol=1e-10 - ) - - -def test_group_gemm_current_stream(group_gemm): - counts = torch.tensor([2, 0, 3], device="cpu") - A = torch.zeros((3, 2, 3, 4), device="cuda") - B = torch.zeros((5, 2, 4), device="cuda") - group_gemm(A, B, counts, 3, 2, 3, 4, 0) - stream = torch.cuda.Stream() - stream.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(stream): - if not torch.version.hip: - torch.cuda._sleep(1_000_000) - A.fill_(2) - B.fill_(3) - actual = group_gemm(A, B, counts, 3, 2, 3, 4, 0).clone() - stream.synchronize() - torch.testing.assert_close(actual, torch.full_like(actual, 24)) - - -@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="Two GPUs are required") -def test_group_gemm_device_guard(group_gemm): - with torch.cuda.device(0): - A = torch.full((1, 2, 3, 4), 2.0, device="cuda:1") - B = torch.full((2, 2, 4), 3.0, device="cuda:1") - actual = group_gemm(A, B, torch.tensor([2], device="cpu"), 1, 2, 3, 4, 0) - assert torch.cuda.current_device() == 0 - assert actual.device == torch.device("cuda:1") - torch.testing.assert_close(actual, torch.full_like(actual, 24)) - with pytest.raises(RuntimeError, match="same device"): - group_gemm( - A, B.to("cuda:0"), torch.tensor([2], device="cpu"), 1, 2, 3, 4, 0 - ) - - -@pytest.mark.parametrize("counts", [[-1, 0, 6], [2, 0, 2], [2, 0, 4]]) -def test_group_gemm_rejects_invalid_counts(group_gemm, counts): - A = torch.empty((3, 2, 3, 4), device="cuda") - B = torch.empty((5, 2, 4), device="cuda") - with pytest.raises(RuntimeError, match="ragged_counts"): - group_gemm(A, B, torch.tensor(counts, device="cpu"), 3, 2, 3, 4, 0) - - -def test_group_gemm_requires_cpu_counts(group_gemm): - A = torch.empty((3, 2, 3, 4), device="cuda") - B = torch.empty((5, 2, 4), device="cuda") - with pytest.raises(RuntimeError, match="ragged_counts must be on the CPU"): - group_gemm(A, B, torch.tensor([2, 0, 3], device="cuda"), 3, 2, 3, 4, 0) - - -@pytest.mark.parametrize("inner", [0, 1]) -@pytest.mark.parametrize( - "counts,batch,m,k", - [ - ([1, 0, 5], 1, 1, 7), - ([1, 0, 5], 3, 7, 1), - ([1, 0, 5], 3, 1, 1), - ([1, 7, 0, 13], 4, 17, 29), - ], -) -def test_group_gemm_varied_shapes(group_gemm, inner, counts, batch, m, k): - A_shape = (len(counts), batch, m, k) if inner == 0 else (sum(counts), batch, m) - A = make_input(A_shape, torch.float64, "offset") - B = make_input((sum(counts), batch, k), torch.float64, "offset") - actual = group_gemm( - A, B, torch.tensor(counts, device="cpu"), len(counts), batch, m, k, inner - ) - expected = reference(A.cpu(), B.cpu(), counts, inner) - torch.testing.assert_close(actual.cpu(), expected, rtol=1e-10, atol=1e-10) - - -@pytest.mark.parametrize("inner", [0, 1]) -def test_group_gemm_double_backward(group_gemm, inner): - counts = torch.tensor([1, 0, 2], device="cpu") - A_shape = (3, 2, 2, 3) if inner == 0 else (3, 2, 2) - A = make_input(A_shape, torch.float64).requires_grad_() - B = make_input((3, 2, 3), torch.float64).requires_grad_() - - def operation(A, B): - return group_gemm(A, B, counts, 3, 2, 2, 3, inner) - - assert torch.autograd.gradgradcheck(operation, (A, B), fast_mode=True) - - -@pytest.mark.parametrize("inner", [0, 1]) -def test_group_gemm_graph_replay(group_gemm, inner): - counts = [2, 0, 3] - counts_tensor = torch.tensor(counts, device="cpu") - A_shape = (3, 2, 3, 4) if inner == 0 else (5, 2, 3) - A = make_input(A_shape, torch.float64) - B = make_input((5, 2, 4), torch.float64) - stream = torch.cuda.Stream() - stream.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(stream): - for _ in range(3): - group_gemm(A, B, counts_tensor, 3, 2, 3, 4, inner) - torch.cuda.current_stream().wait_stream(stream) - - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph, stream=stream): - actual = group_gemm(A, B, counts_tensor, 3, 2, 3, 4, inner) - - for scale in (2, 3): - A.mul_(scale) - B.add_(0.25) - graph.replay() - expected = reference(A.cpu(), B.cpu(), counts, inner) - torch.testing.assert_close(actual.cpu(), expected, rtol=1e-10, atol=1e-10) - - -@pytest.mark.parametrize("inner", [0, 1]) -def test_group_gemm_compile(group_gemm, inner): - counts = [2, 0, 3] - counts_tensor = torch.tensor(counts, device="cpu") - A_shape = (3, 2, 3, 4) if inner == 0 else (5, 2, 3) - A = make_input(A_shape, torch.float64).requires_grad_() - B = make_input((5, 2, 4), torch.float64).requires_grad_() - - def operation(A, B, counts): - return group_gemm(A, B, counts, 3, 2, 3, 4, inner) - - compiled = torch.compile(operation, fullgraph=True) - actual = compiled(A, B, counts_tensor) - A_ref = A.detach().cpu().requires_grad_() - B_ref = B.detach().cpu().requires_grad_() - expected = reference(A_ref, B_ref, counts, inner) - torch.testing.assert_close(actual.cpu(), expected, rtol=1e-10, atol=1e-10) - - actual_grads = torch.autograd.grad(actual.sum(), (A, B)) - expected_grads = torch.autograd.grad(expected.sum(), (A_ref, B_ref)) - for actual_grad, expected_grad in zip(actual_grads, expected_grads): - torch.testing.assert_close( - actual_grad.cpu(), expected_grad, rtol=1e-10, atol=1e-10 - ) - - next_counts = [0, 3, 2] - next_actual = compiled(A, B, torch.tensor(next_counts, device="cpu")) - next_expected = reference(A_ref, B_ref, next_counts, inner) - torch.testing.assert_close(next_actual.cpu(), next_expected, rtol=1e-10, atol=1e-10) - - -@pytest.mark.parametrize("inner", [0, 1]) -def test_group_gemm_aoti(group_gemm, inner, tmp_path): - import openequivariance - - class Model(torch.nn.Module): - def forward(self, A, B, counts): - return group_gemm(A, B, counts, 3, 2, 3, 4, inner) - - counts = [2, 0, 3] - counts_tensor = torch.tensor(counts, device="cpu") - A_shape = (3, 2, 3, 4) if inner == 0 else (5, 2, 3) - A = make_input(A_shape, torch.float64) - B = make_input((5, 2, 4), torch.float64) - exported = torch.export.export(Model(), (A, B, counts_tensor), strict=False) - package_path = torch._inductor.aoti_compile_and_package( - exported, package_path=str(tmp_path / "group_gemm.pt2") - ) - inputs_path = tmp_path / "inputs.pt" - torch.save( - (A.cpu(), B.cpu(), counts_tensor, reference(A.cpu(), B.cpu(), counts, inner)), - inputs_path, - ) - result = subprocess.run( - [ - sys.executable, - "-c", - """ -import sys -import torch -import torch._inductor.codecache - -torch.ops.load_library(sys.argv[1]) -model = torch._inductor.aoti_load_package(sys.argv[2]) -A, B, counts, expected = torch.load(sys.argv[3], weights_only=True) -actual = model(A.cuda(), B.cuda(), counts) -torch.testing.assert_close(actual.cpu(), expected, rtol=1e-10, atol=1e-10) -""", - openequivariance.torch_ext_so_path(), - package_path, - str(inputs_path), - ], - capture_output=True, - text=True, - timeout=180, - ) - assert result.returncode == 0, result.stdout + result.stderr From f661560b9f16b4f3d263f89182847aa92653391b Mon Sep 17 00:00:00 2001 From: asglover <140220574+asglover@users.noreply.github.com> Date: Sun, 13 Sep 2026 19:12:02 -0700 Subject: [PATCH 17/38] typo --- CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index fe2ed09b..c799a7aa 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,7 +2,7 @@ **Added**: - Beta support for Intel XPUs. -- Removed OEQ's direct cuBLAS and rocBLAS dependencies. Use torch's cuda c shim to instead. +- Removed OEQ's direct cuBLAS and rocBLAS dependencies. Use torch's cuda c shim instead. ### v0.7.0 (2026-09-10) **Added**: From e8c3a2f44a896152f2ef8a42d34013a63d3eb84a Mon Sep 17 00:00:00 2001 From: asglover <140220574+asglover@users.noreply.github.com> Date: Sun, 13 Sep 2026 19:14:33 -0700 Subject: [PATCH 18/38] fmt --- CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index c799a7aa..b35fae8c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,7 +2,7 @@ **Added**: - Beta support for Intel XPUs. -- Removed OEQ's direct cuBLAS and rocBLAS dependencies. Use torch's cuda c shim instead. +- Removed OEQ's direct cuBLAS and rocBLAS dependencies. Use torch's cuda c shim instead. ### v0.7.0 (2026-09-10) **Added**: From b0cb732115a2392692d94a05250c902ecdfb6399 Mon Sep 17 00:00:00 2001 From: asglover <140220574+asglover@users.noreply.github.com> Date: Sat, 24 Oct 2026 23:04:58 -0700 Subject: [PATCH 19/38] SYCL support WIP --- CHANGELOG.md | 3 +- .../SymmetricContraction.py | 8 +- .../extension/backend/backend_sycl.hpp | 6 +- .../openequivariance/extension/group_mm.hpp | 33 +---- .../extension/libtorch_tp_jit.cpp | 41 +++--- .../extension/libtorch_tp_jit_stable.cpp | 76 ++++++++-- .../extension/stubs/stream.cpp | 7 +- .../openequivariance/extension/torch_core.hpp | 110 ++++++++------ tests/symmetric_contraction_test.py | 139 +++++++++++++----- 9 files changed, 271 insertions(+), 152 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index b35fae8c..e4f8e6df 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,7 +2,8 @@ **Added**: - Beta support for Intel XPUs. -- Removed OEQ's direct cuBLAS and rocBLAS dependencies. Use torch's cuda c shim instead. +- Route symmetric contraction's grouped matrix multiplications through PyTorch + on CUDA, HIP, and SYCL, removing direct cuBLAS and rocBLAS dependencies. ### v0.7.0 (2026-09-10) **Added**: diff --git a/openequivariance/openequivariance/_torch/symmetric_contraction/SymmetricContraction.py b/openequivariance/openequivariance/_torch/symmetric_contraction/SymmetricContraction.py index 7f3c38fc..df631551 100644 --- a/openequivariance/openequivariance/_torch/symmetric_contraction/SymmetricContraction.py +++ b/openequivariance/openequivariance/_torch/symmetric_contraction/SymmetricContraction.py @@ -123,10 +123,6 @@ def U_matrix_real( class GroupMM: def __init__(self, dtype, num_elements, batch_size): - if extlib.DEVICE_TYPE == "xpu": - raise NotImplementedError( - "Symmetric contraction is not supported on the SYCL backend." - ) self.num_elements = num_elements self.batch_size = batch_size @@ -323,6 +319,8 @@ def __init__( def forward(self, x: torch.Tensor, y: torch.Tensor): indices = torch.argmax(y, dim=1) + # Recomputed once per forward call and shared by all output contractions. + # Autograd saves these CPU counts for backward and double backward. bincount = torch.bincount(indices, minlength=self.num_elements).to("cpu") permutation = torch.argsort(indices) inverse_perm = torch.argsort(permutation) @@ -410,6 +408,6 @@ def backward(ctx, grad_output): ) -if extlib.BUILT_EXTENSION and extlib.DEVICE_TYPE != "xpu": +if extlib.BUILT_EXTENSION: register_torch_fakes() register_autograd() diff --git a/openequivariance/openequivariance/extension/backend/backend_sycl.hpp b/openequivariance/openequivariance/extension/backend/backend_sycl.hpp index 9a8c225c..4c77fdd8 100644 --- a/openequivariance/openequivariance/extension/backend/backend_sycl.hpp +++ b/openequivariance/openequivariance/extension/backend/backend_sycl.hpp @@ -18,13 +18,15 @@ namespace syclex = sycl::ext::oneapi::experimental; using Stream = sycl::queue *; -Stream get_current_stream(); +namespace oeq { + Stream get_current_stream(); +} inline sycl::queue &resolve_queue(Stream stream) { if (stream != nullptr) { return *stream; } - if (Stream current = get_current_stream()) { + if (Stream current = oeq::get_current_stream()) { return *current; } static sycl::queue default_queue{sycl::gpu_selector_v}; diff --git a/openequivariance/openequivariance/extension/group_mm.hpp b/openequivariance/openequivariance/extension/group_mm.hpp index 247b6b3d..8f04376e 100644 --- a/openequivariance/openequivariance/extension/group_mm.hpp +++ b/openequivariance/openequivariance/extension/group_mm.hpp @@ -2,34 +2,17 @@ #include #include -#include -#include -#include - -#include namespace oeq { -inline void check_group_mm_shim(AOTITorchError status) { - if (status != AOTI_TORCH_SUCCESS) - throw std::runtime_error("group_gemm: PyTorch C shim failed"); -} - -using GroupMMTensor = std::unique_ptr< - std::remove_pointer_t, - decltype(&aoti_torch_delete_tensor_object)>; - -inline GroupMMTensor group_mm_view( - AtenTensorHandle tensor, std::array sizes, +inline Tensor group_mm_view( + const Tensor& tensor, std::array sizes, std::array strides, int64_t offset) { - AtenTensorHandle view = nullptr; - check_group_mm_shim(aoti_torch__reinterpret_tensor( - tensor, 3, sizes.data(), strides.data(), offset, &view)); - return GroupMMTensor(view, aoti_torch_delete_tensor_object); + return oeq::as_strided(tensor, sizes, strides, tensor.storage_offset() + offset); } -inline void group_gemm_torch( - AtenTensorHandle A, AtenTensorHandle B, AtenTensorHandle C, +inline void grouped_gemm_via_torch_bmm( + const Tensor& A, const Tensor& B, Tensor& C, const int64_t* ragged_counts, int64_t num_W, int64_t batch_size, int64_t m, int64_t k, int64_t ragged_inner) { if (batch_size == 0 || m == 0 || k == 0) @@ -51,8 +34,7 @@ inline void group_gemm_torch( auto output = group_mm_view(C, {batch_size, n, m}, {m, batch_size * m, 1}, offset * batch_size * m); - check_group_mm_shim(aoti_torch_cuda_bmm_out( - output.get(), input.get(), weight.get())); + oeq::bmm_out(output, input, weight); } else { auto left = group_mm_view(A, {batch_size, m, n}, {m, 1, batch_size * m}, @@ -63,8 +45,7 @@ inline void group_gemm_torch( auto output = group_mm_view(C, {batch_size, m, k}, {m * k, k, 1}, i * batch_size * m * k); - check_group_mm_shim(aoti_torch_cuda_bmm_out( - output.get(), left.get(), right.get())); + oeq::bmm_out(output, left, right); } offset += n; } diff --git a/openequivariance/openequivariance/extension/libtorch_tp_jit.cpp b/openequivariance/openequivariance/extension/libtorch_tp_jit.cpp index e4600652..17fe4b35 100644 --- a/openequivariance/openequivariance/extension/libtorch_tp_jit.cpp +++ b/openequivariance/openequivariance/extension/libtorch_tp_jit.cpp @@ -1,26 +1,37 @@ +#include #include +#include +#include #include #ifdef CUDA_BACKEND #include + #include "backend_cuda.hpp" #endif #ifdef HIP_BACKEND #include + #include "backend_hip.hpp" #endif #ifdef SYCL_BACKEND #include + #include "backend_sycl.hpp" #endif #include -#include #include #include #include -#include #include +#define TCHECK TORCH_CHECK +#define BOX(x) x +#define REGISTER_LIBRARY_IMPL TORCH_LIBRARY_IMPL +#define REGISTER_LIBRARY TORCH_LIBRARY + +namespace oeq { + using Tensor = torch::Tensor; using Dtype = torch::Dtype; @@ -30,27 +41,21 @@ constexpr Dtype kInt = torch::kInt; constexpr Dtype kLong = torch::kLong; constexpr Dtype kByte = torch::kByte; -#define TCHECK TORCH_CHECK -#define BOX(x) x -#define REGISTER_LIBRARY_IMPL TORCH_LIBRARY_IMPL -#define REGISTER_LIBRARY TORCH_LIBRARY +} // namespace oeq -namespace { +#include "torch_core.hpp" -class TensorDeviceGuard { - c10::DeviceGuard guard; -public: - explicit TensorDeviceGuard(const Tensor& tensor) : guard(tensor.device()) {} -}; +namespace oeq { -AtenTensorHandle tensor_handle(Tensor& tensor) { - return torch::aot_inductor::tensor_pointer_to_tensor_handle(&tensor); +Tensor as_strided(const Tensor& tensor, std::array sizes, + std::array strides, int64_t storage_offset) { + return at::as_strided(tensor, sizes, strides, storage_offset); } +Tensor& bmm_out(Tensor& out, const Tensor& self, const Tensor& mat2) { + return at::bmm_out(out, self, mat2); } -#include "torch_core.hpp" - Tensor tensor_to_cpu_contiguous(const Tensor &tensor) { return tensor.to(torch::kCPU).contiguous(); } @@ -114,7 +119,9 @@ bool tensor_is_on_gpu(const Tensor &tensor) { #endif } -namespace py=pybind11; +} // namespace oeq + +namespace py = pybind11; PYBIND11_MODULE(libtorch_tp_jit, m) { py::class_(m, "DeviceProp") .def(py::init()) diff --git a/openequivariance/openequivariance/extension/libtorch_tp_jit_stable.cpp b/openequivariance/openequivariance/extension/libtorch_tp_jit_stable.cpp index 7792ee1c..96a84600 100644 --- a/openequivariance/openequivariance/extension/libtorch_tp_jit_stable.cpp +++ b/openequivariance/openequivariance/extension/libtorch_tp_jit_stable.cpp @@ -1,14 +1,21 @@ #define USE_CUDA +#include #include +#include +#include #include +#include #include +#include #include +#include #include #include #include #include #include +#include #include #include #ifdef SYCL_BACKEND @@ -17,6 +24,28 @@ extern "C" AOTITorchError aoti_torch_get_current_sycl_queue(void** ret_queue); #endif +#ifdef CUDA_BACKEND + #include "backend_cuda.hpp" +#endif + +#ifdef HIP_BACKEND + #include "backend_hip.hpp" +#endif + +#ifdef SYCL_BACKEND + #include "backend_sycl.hpp" +#endif + +#define TCHECK STD_TORCH_CHECK +#define BOX(x) TORCH_BOX(x) +#define REGISTER_LIBRARY_IMPL STABLE_TORCH_LIBRARY_IMPL +#define REGISTER_LIBRARY STABLE_TORCH_LIBRARY + +#ifndef STABLE_TORCH_ERROR_CODE_CHECK +#define STABLE_TORCH_ERROR_CODE_CHECK(call) TORCH_ERROR_CODE_CHECK(call) +#endif + +namespace oeq { using Tensor = torch::stable::Tensor; using Dtype = torch::headeronly::ScalarType; @@ -27,27 +56,40 @@ constexpr Dtype kInt = torch::headeronly::ScalarType::Int; constexpr Dtype kLong = torch::headeronly::ScalarType::Long; constexpr Dtype kByte = torch::headeronly::ScalarType::Byte; -#define TCHECK STD_TORCH_CHECK -#define BOX(x) TORCH_BOX(x) -#define REGISTER_LIBRARY_IMPL STABLE_TORCH_LIBRARY_IMPL -#define REGISTER_LIBRARY STABLE_TORCH_LIBRARY +} // namespace oeq -namespace { +#include "torch_core.hpp" -class TensorDeviceGuard { - torch::stable::accelerator::DeviceGuard guard; -public: - explicit TensorDeviceGuard(const Tensor& tensor) : guard(tensor.get_device_index()) {} -}; +namespace oeq { -AtenTensorHandle tensor_handle(Tensor& tensor) { - return tensor.get(); +Tensor as_strided(const Tensor& tensor, std::array sizes, + std::array strides, int64_t storage_offset) { + auto sizes_ref = torch::headeronly::IntHeaderOnlyArrayRef(sizes.data(), sizes.size()); + auto strides_ref = torch::headeronly::IntHeaderOnlyArrayRef(strides.data(), strides.size()); + std::array stack{ + torch::stable::detail::from(tensor), + torch::stable::detail::from(sizes_ref), + torch::stable::detail::from(strides_ref), + torch::stable::detail::from(std::optional(storage_offset))}; + STABLE_TORCH_ERROR_CODE_CHECK(torch_call_dispatcher( + "aten::as_strided", "", stack.data(), TORCH_ABI_VERSION)); + return torch::stable::detail::to(stack[0]); } +// Adapted from PyTorch PR #197441, commit +// 646e5d175800d0cce119e9d8a92bd5d79bb926d1, torch/csrc/stable/ops.h: +// https://github.com/pytorch/pytorch/pull/197441 +Tensor& bmm_out(Tensor& out, const Tensor& self, const Tensor& mat2) { + std::array stack{ + torch::stable::detail::from(self), + torch::stable::detail::from(mat2), + torch::stable::detail::from(out)}; + STABLE_TORCH_ERROR_CODE_CHECK(torch_call_dispatcher( + "aten::bmm", "out", stack.data(), TORCH_ABI_VERSION)); + (void)torch::stable::detail::to(stack[0]); + return out; } -#include "torch_core.hpp" - Tensor tensor_to_cpu_contiguous(const Tensor &tensor) { torch::stable::Device device(torch::headeronly::DeviceType::CPU); return torch::stable::contiguous(torch::stable::to(tensor, device)); @@ -89,10 +131,10 @@ Stream get_current_stream() { void* stream_ptr = nullptr; #ifdef SYCL_BACKEND - TORCH_ERROR_CODE_CHECK(aoti_torch_get_current_sycl_queue(&stream_ptr)); + STABLE_TORCH_ERROR_CODE_CHECK(aoti_torch_get_current_sycl_queue(&stream_ptr)); #else auto device_idx = torch::stable::accelerator::getCurrentDeviceIndex(); - TORCH_ERROR_CODE_CHECK(aoti_torch_get_current_cuda_stream(device_idx, &stream_ptr)); + STABLE_TORCH_ERROR_CODE_CHECK(aoti_torch_get_current_cuda_stream(device_idx, &stream_ptr)); #endif return static_cast(stream_ptr); @@ -110,6 +152,8 @@ bool tensor_is_on_gpu(const Tensor &tensor) { #endif } +} // namespace oeq + #ifdef CUDA_BACKEND #define EXTENSION_NAME oeq_stable_cuda #endif diff --git a/openequivariance/openequivariance/extension/stubs/stream.cpp b/openequivariance/openequivariance/extension/stubs/stream.cpp index de45dc06..8676d001 100644 --- a/openequivariance/openequivariance/extension/stubs/stream.cpp +++ b/openequivariance/openequivariance/extension/stubs/stream.cpp @@ -1,7 +1,7 @@ #define USE_CUDA #include -#include +#include extern "C" { #ifdef SYCL_BACKEND @@ -13,10 +13,5 @@ extern "C" { AOTITorchError aoti_torch_get_current_cuda_stream(int32_t device_index, void** ret_stream) { return AOTI_TORCH_FAILURE; } - - AOTITorchError aoti_torch_cuda_bmm_out( - AtenTensorHandle out, AtenTensorHandle self, AtenTensorHandle mat2) { - return AOTI_TORCH_FAILURE; - } #endif } diff --git a/openequivariance/openequivariance/extension/torch_core.hpp b/openequivariance/openequivariance/extension/torch_core.hpp index f744d96a..6d969b3e 100644 --- a/openequivariance/openequivariance/extension/torch_core.hpp +++ b/openequivariance/openequivariance/extension/torch_core.hpp @@ -1,5 +1,6 @@ #pragma once +#include #include #include #include @@ -9,6 +10,7 @@ #include #include #include +#include #include #include @@ -16,30 +18,36 @@ #ifdef CUDA_BACKEND #include "backend_cuda.hpp" +#endif + +#ifdef HIP_BACKEND + #include "backend_hip.hpp" +#endif + +#ifdef SYCL_BACKEND + #include "backend_sycl.hpp" +#endif + +#include "tensorproducts.hpp" +#include "convolution.hpp" + +namespace oeq { + +#ifdef CUDA_BACKEND using JITKernel = CUJITKernel; using GPU_Allocator = CUDA_Allocator; #endif #ifdef HIP_BACKEND - #include "backend_hip.hpp" using JITKernel = HIPJITKernel; using GPU_Allocator = HIP_Allocator; #endif #ifdef SYCL_BACKEND - #include "backend_sycl.hpp" using JITKernel = SYCLJITKernel; using GPU_Allocator = SYCL_Allocator; #endif -#ifndef SYCL_BACKEND - #include "group_mm.hpp" -#endif - -#include "tensorproducts.hpp" -#include "convolution.hpp" - -using namespace std; using json = json11::Json; Dtype enum_to_torch_dtype(int64_t i); @@ -57,6 +65,16 @@ bool tensor_is_on_gpu(const Tensor &tensor); const uint8_t *tensor_data_ptr_u8(const Tensor &tensor); void *data_ptr(const Tensor &tensor); +Tensor as_strided(const Tensor& tensor, std::array sizes, + std::array strides, int64_t storage_offset); +Tensor& bmm_out(Tensor& out, const Tensor& self, const Tensor& mat2); + +} // namespace oeq + +#include "group_mm.hpp" + +namespace oeq { + inline std::string shape_to_string(std::initializer_list shape) { std::ostringstream oss; oss << "["; @@ -108,13 +126,13 @@ inline Dtype enum_to_torch_dtype(int64_t i) { case 4: return kLong; case 5: return kByte; } - throw logic_error("Unsupported tensor datatype!"); + throw std::logic_error("Unsupported tensor datatype!"); } -inline void check_tensor(const Tensor &tensor, - std::initializer_list expected_shape, - Dtype expected_dtype, - std::string tensor_name) { +inline void check_tensor_metadata(const Tensor &tensor, + std::initializer_list expected_shape, + Dtype expected_dtype, + const std::string &tensor_name) { bool shape_ok = (tensor.dim() == static_cast(expected_shape.size())); if (shape_ok) { int64_t i = 0; @@ -131,13 +149,21 @@ inline void check_tensor(const Tensor &tensor, "Shape mismatch for tensor '", tensor_name, "'. Expected: ", shape_to_string(expected_shape), ". Got: ", tensor_sizes_str(tensor)); - TCHECK(tensor_is_on_gpu(tensor), "Tensor '", tensor_name, "' is not on the GPU."); TCHECK(tensor.scalar_type() == expected_dtype, "Dtype mismatch for tensor '", tensor_name, "'. Expected: ", static_cast(expected_dtype), ". Got: ", static_cast(tensor.scalar_type())); } +inline void check_tensor(const Tensor &tensor, + std::initializer_list expected_shape, + Dtype expected_dtype, + const std::string &tensor_name) { + check_tensor_metadata(tensor, expected_shape, expected_dtype, tensor_name); + TCHECK(tensor_is_on_gpu(tensor), "Tensor '", tensor_name, + "' is not on the GPU backend supported by this extension."); +} + inline std::unordered_map parse_json_config(const json &j_obj) { std::unordered_map result; for (const auto &kv : j_obj.object_items()) { @@ -165,7 +191,7 @@ struct KernelProp { idx_dtype(kInt), workspace_dtype(kByte) {} KernelProp( - std::unordered_map &kernel_dims, bool is_convolution) : + std::unordered_map &kernel_dims, bool is_convolution) : L1_dim(kernel_dims.at("L1_dim")), L2_dim(kernel_dims.at("L2_dim")), L3_dim(kernel_dims.at("L3_dim")), @@ -335,7 +361,7 @@ inline Tensor jit_tp_forward( return L3_out; } -inline tuple jit_tp_backward( +inline std::tuple jit_tp_backward( Tensor json_bytes, int64_t hash, Tensor L1_in, Tensor L2_in, @@ -377,10 +403,10 @@ inline tuple jit_tp_backward( stream ); - return tuple(L1_grad, L2_grad, W_grad); + return std::tuple(L1_grad, L2_grad, W_grad); } -inline tuple jit_tp_double_backward( +inline std::tuple jit_tp_double_backward( Tensor json_bytes, int64_t hash, Tensor L1_in, Tensor L2_in, @@ -439,7 +465,7 @@ inline tuple jit_tp_double_backward( stream ); - return tuple(L1_grad, L2_grad, W_grad, L3_dgrad); + return std::tuple(L1_grad, L2_grad, W_grad, L3_dgrad); } @@ -501,7 +527,7 @@ inline Tensor jit_conv_forward( return L3_out; } -inline tuple jit_conv_backward( +inline std::tuple jit_conv_backward( Tensor json_bytes, int64_t hash, Tensor L1_in, Tensor L2_in, @@ -564,10 +590,10 @@ inline tuple jit_conv_backward( data_ptr(transpose_perm_contig), stream); - return tuple(L1_grad, L2_grad, W_grad); + return std::tuple(L1_grad, L2_grad, W_grad); } -inline tuple jit_conv_double_backward( +inline std::tuple jit_conv_double_backward( Tensor json_bytes, int64_t hash, Tensor L1_in, Tensor L2_in, @@ -644,16 +670,16 @@ inline tuple jit_conv_double_backward( stream ); - return tuple(L1_grad, L2_grad, W_grad, L3_dgrad); + return std::tuple(L1_grad, L2_grad, W_grad, L3_dgrad); } // =========================================================== -#ifndef SYCL_BACKEND inline Tensor group_gemm( Tensor A, Tensor B, Tensor ragged_counts, int64_t num_W, int64_t batch_size, int64_t m, int64_t k, int64_t ragged_inner) { - TCHECK(A.is_cuda() && B.is_cuda(), "group_gemm: A and B must be GPU tensors"); + TCHECK(tensor_is_on_gpu(A) && tensor_is_on_gpu(B), + "group_gemm: A and B must be on the GPU backend supported by this extension"); TCHECK(A.get_device() == B.get_device(), "group_gemm: A and B must be on the same device"); TCHECK(A.scalar_type() == B.scalar_type(), "group_gemm: A and B must have the same dtype"); TCHECK(A.scalar_type() == kFloat || A.scalar_type() == kDouble, @@ -667,11 +693,11 @@ inline Tensor group_gemm( "group_gemm: ragged_counts must contain num_W entries"); TCHECK(B.dim() == 3, "group_gemm: B must have shape [rows, batch_size, k]"); const int64_t rows = B.size(0); - check_tensor(B, {rows, batch_size, k}, A.scalar_type(), "group_gemm B"); + check_tensor_metadata(B, {rows, batch_size, k}, A.scalar_type(), "group_gemm B"); if (ragged_inner == 0) - check_tensor(A, {num_W, batch_size, m, k}, A.scalar_type(), "group_gemm A"); + check_tensor_metadata(A, {num_W, batch_size, m, k}, A.scalar_type(), "group_gemm A"); else - check_tensor(A, {rows, batch_size, m}, A.scalar_type(), "group_gemm A"); + check_tensor_metadata(A, {rows, batch_size, m}, A.scalar_type(), "group_gemm A"); Tensor rc_c = tensor_contiguous(ragged_counts); const auto* rc_ptr = static_cast(data_ptr(rc_c)); @@ -684,7 +710,6 @@ inline Tensor group_gemm( TCHECK(row_count == rows, "group_gemm: ragged_counts must sum to the number of rows"); - TensorDeviceGuard device_guard(A); Tensor A_c = tensor_contiguous(A); Tensor B_c = tensor_contiguous(B); @@ -696,15 +721,16 @@ inline Tensor group_gemm( C = tensor_zeros_like(A, make_sizes({num_W, batch_size, m, k})); } - oeq::group_gemm_torch(tensor_handle(A_c), tensor_handle(B_c), tensor_handle(C), + oeq::grouped_gemm_via_torch_bmm(A_c, B_c, C, rc_ptr, num_W, batch_size, m, k, ragged_inner); return C; } -#endif // =========================================================== +} // namespace oeq + // The dispatch key must match the device the tensors live on: XPU for SYCL, // CUDA for both CUDA and HIP (PyTorch maps HIP tensors onto the CUDA key). #ifdef SYCL_BACKEND @@ -714,17 +740,15 @@ inline Tensor group_gemm( #endif REGISTER_LIBRARY_IMPL(libtorch_tp_jit, OEQ_DISPATCH_KEY, m) { - m.impl("jit_tp_forward", BOX(&jit_tp_forward)); - m.impl("jit_tp_backward", BOX(&jit_tp_backward)); - m.impl("jit_tp_double_backward", BOX(&jit_tp_double_backward)); + m.impl("jit_tp_forward", BOX(&oeq::jit_tp_forward)); + m.impl("jit_tp_backward", BOX(&oeq::jit_tp_backward)); + m.impl("jit_tp_double_backward", BOX(&oeq::jit_tp_double_backward)); - m.impl("jit_conv_forward", BOX(&jit_conv_forward)); - m.impl("jit_conv_backward", BOX(&jit_conv_backward)); - m.impl("jit_conv_double_backward", BOX(&jit_conv_double_backward)); + m.impl("jit_conv_forward", BOX(&oeq::jit_conv_forward)); + m.impl("jit_conv_backward", BOX(&oeq::jit_conv_backward)); + m.impl("jit_conv_double_backward", BOX(&oeq::jit_conv_double_backward)); -#ifndef SYCL_BACKEND - m.impl("group_gemm", BOX(&group_gemm)); -#endif + m.impl("group_gemm", BOX(&oeq::group_gemm)); }; REGISTER_LIBRARY(libtorch_tp_jit, m) { @@ -736,7 +760,5 @@ REGISTER_LIBRARY(libtorch_tp_jit, m) { m.def("jit_conv_backward(Tensor json_bytes, int hash, Tensor L1_in, Tensor L2_in, Tensor W, Tensor L3_grad, Tensor rows, Tensor cols, Tensor workspace, Tensor transpose_perm) -> (Tensor, Tensor, Tensor)"); m.def("jit_conv_double_backward(Tensor json_bytes, int hash, Tensor L1_in, Tensor L2_in, Tensor W, Tensor L3_grad, Tensor L1_dgrad, Tensor L2_dgrad, Tensor W_dgrad, Tensor rows, Tensor cols, Tensor workspace, Tensor transpose_perm) -> (Tensor, Tensor, Tensor, Tensor)"); -#ifndef SYCL_BACKEND m.def("group_gemm(Tensor A, Tensor B, Tensor ragged_counts, int num_W, int batch_size, int m, int k, int ragged_inner) -> Tensor"); -#endif }; diff --git a/tests/symmetric_contraction_test.py b/tests/symmetric_contraction_test.py index c94728e0..a20f92b8 100644 --- a/tests/symmetric_contraction_test.py +++ b/tests/symmetric_contraction_test.py @@ -9,13 +9,7 @@ from openequivariance._torch.symmetric_contraction import SymmetricContraction -from conftest import device_type - -if device_type() == "xpu": - pytest.skip( - "Symmetric contraction is not supported on the SYCL backend.", - allow_module_level=True, - ) +from conftest import device_type, torch_accelerator mace_symmetric_contraction = pytest.importorskip("mace.modules.symmetric_contraction") MaceSymmetricContraction = mace_symmetric_contraction.SymmetricContraction @@ -71,10 +65,20 @@ def symmetric_contraction_config(request): @pytest.fixture(params=[torch.float32, torch.float64], ids=["F32", "F64"]) -def dtype(request): +def dtype(request, symmetric_contraction_config): + require_dtype(symmetric_contraction_config.device, request.param) return request.param +def require_dtype(device, dtype): + if ( + device.type == "xpu" + and dtype == torch.float64 + and not torch.xpu.get_device_properties(device).has_fp64 + ): + pytest.skip("This XPU does not support float64") + + @pytest.fixture def labels(symmetric_contraction_config): cfg = symmetric_contraction_config @@ -104,6 +108,20 @@ def node_feats(dtype, symmetric_contraction_config): ) +@pytest.fixture(params=["contiguous", "transposed", "offset"]) +def node_feats_with_layout(request, node_feats): + match request.param: + case "contiguous": + assert node_feats.is_contiguous() + case "transposed": + node_feats = node_feats.transpose(0, 2).contiguous().transpose(0, 2) + assert not node_feats.is_contiguous() + case "offset": + node_feats = torch.cat((torch.zeros_like(node_feats[:1]), node_feats))[1:] + assert node_feats.is_contiguous() and node_feats.storage_offset() > 0 + return node_feats.detach().requires_grad_() + + @pytest.fixture def modules(dtype, symmetric_contraction_config): cfg = symmetric_contraction_config @@ -171,32 +189,48 @@ def random_like(tensor, seed): ) -class TestSymmetricContraction: - def test_matches_mace_forward_backward( - self, modules, node_feats, node_attrs, dtype, symmetric_contraction_config - ): - cfg = symmetric_contraction_config - oeq_module, mace_module = modules - mace_node_feats = node_feats.detach().clone().requires_grad_() +def assert_matches_mace_forward_backward( + modules, node_feats, node_attrs, dtype, oeq_forward=None +): + oeq_module, mace_module = modules + mace_node_feats = node_feats.detach().clone().requires_grad_() + if oeq_forward is None: + oeq_forward = oeq_module - oeq_output = oeq_module(node_feats, node_attrs) - mace_output = mace_module(mace_node_feats, node_attrs) + oeq_output = oeq_forward(node_feats, node_attrs) + mace_output = mace_module(mace_node_feats, node_attrs) - assert oeq_output.shape == (len(cfg.label_values), cfg.irreps_out.dim) - torch.testing.assert_close(oeq_output, mace_output, **tolerance(dtype)) + torch.testing.assert_close(oeq_output, mace_output, **tolerance(dtype)) - output_grad = random_like(oeq_output, seed=4321) - oeq_params, mace_params = matching_trainable_parameters(oeq_module, mace_module) + output_grad = random_like(oeq_output, seed=4321) + oeq_params, mace_params = matching_trainable_parameters(oeq_module, mace_module) - oeq_grads = torch.autograd.grad( - oeq_output, (node_feats, *oeq_params), grad_outputs=output_grad - ) - mace_grads = torch.autograd.grad( - mace_output, (mace_node_feats, *mace_params), grad_outputs=output_grad - ) + oeq_grads = torch.autograd.grad( + oeq_output, (node_feats, *oeq_params), grad_outputs=output_grad + ) + mace_grads = torch.autograd.grad( + mace_output, (mace_node_feats, *mace_params), grad_outputs=output_grad + ) - for oeq_grad, mace_grad in zip(oeq_grads, mace_grads): - torch.testing.assert_close(oeq_grad, mace_grad, **tolerance(dtype)) + for oeq_grad, mace_grad in zip(oeq_grads, mace_grads): + torch.testing.assert_close(oeq_grad, mace_grad, **tolerance(dtype)) + return oeq_output + + +class TestSymmetricContraction: + def test_matches_mace_forward_backward( + self, + modules, + node_feats_with_layout, + node_attrs, + dtype, + symmetric_contraction_config, + ): + output = assert_matches_mace_forward_backward( + modules, node_feats_with_layout, node_attrs, dtype + ) + cfg = symmetric_contraction_config + assert output.shape == (len(cfg.label_values), cfg.irreps_out.dim) def test_matches_mace_double_backward( self, modules, node_feats, node_attrs, dtype, symmetric_contraction_config @@ -250,13 +284,48 @@ def test_matches_mace_double_backward( for oeq_grad, mace_grad in zip(oeq_second_grads, mace_second_grads): torch.testing.assert_close(oeq_grad, mace_grad, **tolerance(dtype)) - def test_compile(self, modules, node_feats, node_attrs): + def test_compile(self, modules, node_feats, node_attrs, dtype): sc, _ = modules - ref = sc(node_feats, node_attrs) - assert torch.allclose(ref, torch.compile(sc)(node_feats, node_attrs), atol=1e-5) - - def test_export(self, modules, node_feats, node_attrs): + compiled = torch.compile(sc) + single_element_attrs = torch.zeros_like(node_attrs) + single_element_attrs[:, -1] = 1 + for attrs in (node_attrs, single_element_attrs): + assert_matches_mace_forward_backward( + modules, node_feats, attrs, dtype, oeq_forward=compiled + ) + + def test_export(self, modules, node_feats, node_attrs, dtype): sc, _ = modules ref = sc(node_feats, node_attrs) exported = torch.export.export(sc, args=(node_feats, node_attrs), strict=False) - assert torch.allclose(ref, exported.module()(node_feats, node_attrs), atol=1e-5) + torch.testing.assert_close( + ref, exported.module()(node_feats, node_attrs), **tolerance(dtype) + ) + + def test_nondefault_stream(self, modules, node_feats, node_attrs, dtype): + accel = torch_accelerator() + current = accel.current_stream(node_feats.device) + stream = accel.Stream(device=node_feats.device) + stream.wait_stream(current) + try: + with accel.stream(stream): + stream_feats = (node_feats.detach() * 1.5).requires_grad_() + assert_matches_mace_forward_backward( + modules, stream_feats, node_attrs, dtype + ) + finally: + stream.synchronize() + + def test_noncurrent_device(self, modules, node_feats, node_attrs, dtype): + accel = torch_accelerator() + if accel.device_count() < 2: + pytest.skip("Requires two GPUs") + current = accel.current_device() + target = torch.device(device_type(), (current + 1) % accel.device_count()) + require_dtype(target, dtype) + modules = tuple(module.to(target) for module in modules) + node_feats = node_feats.detach().to(target).requires_grad_() + node_attrs = node_attrs.to(target) + with accel.device(current): + assert_matches_mace_forward_backward(modules, node_feats, node_attrs, dtype) + assert accel.current_device() == current From fe75c96a1be9906b2f2e85c2fc3a773cd75bf972 Mon Sep 17 00:00:00 2001 From: asglover <140220574+asglover@users.noreply.github.com> Date: Sun, 25 Oct 2026 00:23:41 -0700 Subject: [PATCH 20/38] Use reshape for node features in symmetric contraction x[permutation] keeps the stride order of non-contiguous node features, so x.view fails when there is more than one channel. reshape is still a view for contiguous inputs, and the non-contiguous tensor still reaches group_gemm. --- .../_torch/symmetric_contraction/SymmetricContraction.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/openequivariance/openequivariance/_torch/symmetric_contraction/SymmetricContraction.py b/openequivariance/openequivariance/_torch/symmetric_contraction/SymmetricContraction.py index df631551..5a69a7f2 100644 --- a/openequivariance/openequivariance/_torch/symmetric_contraction/SymmetricContraction.py +++ b/openequivariance/openequivariance/_torch/symmetric_contraction/SymmetricContraction.py @@ -241,7 +241,8 @@ def forward( s = c_tensor.shape out = torch.sum( - c_tensor.view(s[0] * s[1], -1, s[-1]) * x.view(s[0] * s[1], 1, s[-1]), + c_tensor.view(s[0] * s[1], -1, s[-1]) + * x.reshape(s[0] * s[1], 1, s[-1]), dim=2, ).view(s[:-1]) # out = torch.bmm(c_tensor.view(s[0] * s[1], -1, s[-1]), x.view(s[0] * s[1], s[-1], 1)).view(s[:-1]) From 0893733a1af874084f56964b1f2f28df1f473fc0 Mon Sep 17 00:00:00 2001 From: asglover <140220574+asglover@users.noreply.github.com> Date: Sun, 1 Nov 2026 23:03:51 -0800 Subject: [PATCH 21/38] remove comment --- .../_torch/symmetric_contraction/SymmetricContraction.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/openequivariance/openequivariance/_torch/symmetric_contraction/SymmetricContraction.py b/openequivariance/openequivariance/_torch/symmetric_contraction/SymmetricContraction.py index 5a69a7f2..424c2eaf 100644 --- a/openequivariance/openequivariance/_torch/symmetric_contraction/SymmetricContraction.py +++ b/openequivariance/openequivariance/_torch/symmetric_contraction/SymmetricContraction.py @@ -320,8 +320,6 @@ def __init__( def forward(self, x: torch.Tensor, y: torch.Tensor): indices = torch.argmax(y, dim=1) - # Recomputed once per forward call and shared by all output contractions. - # Autograd saves these CPU counts for backward and double backward. bincount = torch.bincount(indices, minlength=self.num_elements).to("cpu") permutation = torch.argsort(indices) inverse_perm = torch.argsort(permutation) From a9aa1e8eeb7ccc308efdc7e300143ca01a042738 Mon Sep 17 00:00:00 2001 From: asglover <140220574+asglover@users.noreply.github.com> Date: Sun, 1 Nov 2026 23:04:35 -0800 Subject: [PATCH 22/38] update changelog --- CHANGELOG.md | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index e4f8e6df..99f3887c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,10 +1,14 @@ ## Latest Changes **Added**: -- Beta support for Intel XPUs. - Route symmetric contraction's grouped matrix multiplications through PyTorch on CUDA, HIP, and SYCL, removing direct cuBLAS and rocBLAS dependencies. +### v0.8.0 (2026-11-01) + +**Added**: +- Beta support for Intel XPUs. + ### v0.7.0 (2026-09-10) **Added**: - Public XLA FFI registration provider From 80b9f22b133f95faf31de4d8d9968cb66b1820f1 Mon Sep 17 00:00:00 2001 From: Vivek Bharadwaj Date: Sun, 1 Nov 2026 21:44:37 -0800 Subject: [PATCH 23/38] Wheeldist. --- .github/cibw_before_all.sh | 32 ++++++++++ .github/workflows/release.yaml | 62 ++++++++++++++----- .github/workflows/requirements_cuda_ci.txt | 2 +- .github/workflows/requirements_release_ci.txt | 1 + .github/workflows/requirements_sycl_ci.txt | 2 +- openequivariance/CMakeLists.txt | 33 ++++++++-- openequivariance/pyproject.toml | 34 ++++++++-- 7 files changed, 140 insertions(+), 26 deletions(-) create mode 100644 .github/cibw_before_all.sh diff --git a/.github/cibw_before_all.sh b/.github/cibw_before_all.sh new file mode 100644 index 00000000..8fbbeaec --- /dev/null +++ b/.github/cibw_before_all.sh @@ -0,0 +1,32 @@ +#!/bin/bash +# Runs inside the manylinux container before cibuildwheel builds +# openequivariance. Installs only the headers and link-time libraries +# for CUDA and HIP; at runtime, the user's PyTorch install provides them. +set -euo pipefail + +CUDA_VERSION=13-0 +ARCH=$(uname -m) + +case "$ARCH" in + x86_64) CUDA_REPO_ARCH=x86_64 ;; + aarch64) CUDA_REPO_ARCH=sbsa ;; +esac + +dnf config-manager --add-repo \ + "https://developer.download.nvidia.com/compute/cuda/repos/rhel8/${CUDA_REPO_ARCH}/cuda-rhel8.repo" +dnf install -y \ + cuda-nvcc-${CUDA_VERSION} \ + cuda-cudart-devel-${CUDA_VERSION} \ + cuda-nvrtc-devel-${CUDA_VERSION} \ + cuda-driver-devel-${CUDA_VERSION} + +# rocm-sdk-core (~415 MB) has the HIP headers, libamdhip64 and libhiprtc, +# without the rest of the ROCm toolkit. AMD publishes no aarch64 builds. +if [ "$ARCH" = x86_64 ]; then + echo "rocm-sdk-core==10.0.0 --hash=sha256:3bf4a72d11aa2a4ee1e90572c73630f937d38c1b7af4dda45b483b54655138a7" \ + > /tmp/rocm-requirements.txt + /opt/python/cp312-cp312/bin/python -m pip install \ + --no-deps --require-hashes --target /opt/rocm-sdk-core \ + --index-url https://download.pytorch.org/whl/rocm10.0 \ + -r /tmp/rocm-requirements.txt +fi diff --git a/.github/workflows/release.yaml b/.github/workflows/release.yaml index b5fe3451..346b0600 100644 --- a/.github/workflows/release.yaml +++ b/.github/workflows/release.yaml @@ -4,8 +4,39 @@ on: workflow_dispatch: jobs: - build-oeq: - name: Build oeq dist + build-oeq-wheels: + name: Build oeq wheels (${{ matrix.arch }}) + runs-on: ${{ matrix.runner }} + timeout-minutes: 60 + strategy: + matrix: + include: + - runner: ubuntu-24.04 + arch: x86_64 + - runner: ubuntu-24.04-arm + arch: aarch64 + steps: + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + with: + fetch-tags: true + - name: set up Python + uses: actions/setup-python@6d4fd22b20465ccc0f7354476588e8df7c9ca0e0 # v7.1.0 + with: + python-version: '3.12' + - name: set version from tag + run: | + python3 -m pip install -r .github/workflows/requirements_release_ci.txt + python3 .github/autoversion.py openequivariance/pyproject.toml + - name: build wheels + run: python3 -m cibuildwheel --output-dir wheelhouse openequivariance + - name: store the wheels + uses: actions/upload-artifact@cf430e030ddbb5b0abf93d22962f4752f3646cd9 # v7.0.2 + with: + name: oeq-wheels-${{ matrix.arch }} + path: wheelhouse/*.whl + + build-oeq-sdist: + name: Build oeq sdist runs-on: ubuntu-latest steps: - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 @@ -14,7 +45,7 @@ jobs: - name: set up Python uses: actions/setup-python@6d4fd22b20465ccc0f7354476588e8df7c9ca0e0 # v7.1.0 with: - python-version: '3.10' + python-version: '3.12' - name: set version from tag run: | python3 -m pip install -r .github/workflows/requirements_release_ci.txt @@ -22,19 +53,19 @@ jobs: - name: install dependencies, then build source tarball run: | cd openequivariance - python3 -m pip install build + python3 -m pip install build python3 -m build --sdist - - name: store the distribution packages + - name: store the distribution packages uses: actions/upload-artifact@cf430e030ddbb5b0abf93d22962f4752f3646cd9 # v7.0.2 with: - name: oeq-package-distributions + name: oeq-sdist path: openequivariance/dist/ pypi-publish: name: Upload oeq to PyPI runs-on: ubuntu-latest - needs: build-oeq - environment: + needs: [build-oeq-wheels, build-oeq-sdist] + environment: name: pypi url: https://pypi.org/p/openequivariance permissions: @@ -43,8 +74,9 @@ jobs: - name: download the distributions uses: actions/download-artifact@9000827ccba6bdab643e8b6fd33ac0654aef8333 # v8.0.2 with: - name: oeq-package-distributions - path: dist/ + pattern: oeq-* + merge-multiple: true + path: dist/ - name: publish package distributions to PyPI uses: pypa/gh-action-pypi-publish@7245906fb61fa33a13554f9d5fb856fe61758086 # v1.14.3 @@ -68,9 +100,9 @@ jobs: - name: install dependencies, then build source tarball run: | cd openequivariance_extjax - python3 -m pip install build - python3 -m build --sdist - - name: store the distribution packages + python3 -m pip install build + python3 -m build --sdist + - name: store the distribution packages uses: actions/upload-artifact@cf430e030ddbb5b0abf93d22962f4752f3646cd9 # v7.0.2 with: name: oeqjax-package-distributions @@ -80,7 +112,7 @@ jobs: name: Upload extjax to PyPI runs-on: ubuntu-latest needs: build-oeq-extjax - environment: + environment: name: pypi url: https://pypi.org/p/openequivariance_extjax permissions: @@ -90,6 +122,6 @@ jobs: uses: actions/download-artifact@9000827ccba6bdab643e8b6fd33ac0654aef8333 # v8.0.2 with: name: oeqjax-package-distributions - path: dist/ + path: dist/ - name: publish package distributions to PyPI uses: pypa/gh-action-pypi-publish@7245906fb61fa33a13554f9d5fb856fe61758086 # v1.14.3 diff --git a/.github/workflows/requirements_cuda_ci.txt b/.github/workflows/requirements_cuda_ci.txt index 4872ba8d..2be0fc59 100644 --- a/.github/workflows/requirements_cuda_ci.txt +++ b/.github/workflows/requirements_cuda_ci.txt @@ -2,5 +2,5 @@ numpy==2.2.5 torch==2.13.0 --index-url https://download.pytorch.org/whl/cu129 pytest==9.0.3 ninja==1.11.1.4 -nanobind==2.10.2 +nanobind==3.1.0 scikit-build-core==0.11.6 \ No newline at end of file diff --git a/.github/workflows/requirements_release_ci.txt b/.github/workflows/requirements_release_ci.txt index 7b69eff7..dcb25b15 100644 --- a/.github/workflows/requirements_release_ci.txt +++ b/.github/workflows/requirements_release_ci.txt @@ -1 +1,2 @@ tomlkit==0.15.2 +cibuildwheel==4.3.5 diff --git a/.github/workflows/requirements_sycl_ci.txt b/.github/workflows/requirements_sycl_ci.txt index d2d92ac3..a4b9ca6e 100644 --- a/.github/workflows/requirements_sycl_ci.txt +++ b/.github/workflows/requirements_sycl_ci.txt @@ -3,5 +3,5 @@ numpy==2.2.5 torch==2.13.0+xpu pytest==9.0.3 ninja==1.11.1.4 -nanobind==2.10.2 +nanobind==3.1.0 scikit-build-core==0.11.6 diff --git a/openequivariance/CMakeLists.txt b/openequivariance/CMakeLists.txt index 7f517ab6..c5bf115f 100644 --- a/openequivariance/CMakeLists.txt +++ b/openequivariance/CMakeLists.txt @@ -1,7 +1,7 @@ cmake_minimum_required(VERSION 3.15...3.30) project(openequivariance_stable_ext) -find_package(Python 3.10 REQUIRED COMPONENTS Interpreter Development.Module) +find_package(Python 3.11 REQUIRED COMPONENTS Interpreter Development.Module ${SKBUILD_SABI_COMPONENT}) # Download LibTorch include(FetchContent) @@ -55,8 +55,9 @@ set(OEQ_SOURCES set(OEQ_INSTALL_DIR "${CMAKE_CURRENT_SOURCE_DIR}/openequivariance/_torch/extlib") function(add_stable_extension target_name backend_define link_libraries) - # Create nanobind extension - nanobind_add_module(${target_name} NB_STATIC ${OEQ_SOURCES}) + # Create nanobind extension in split mode: stable-ABI frontend, with + # libnanobind provided at runtime by the nanobind-backend package. + nanobind_add_module(${target_name} BACKEND_MODULE nanobind_backend ${OEQ_SOURCES}) set_target_properties(${target_name} PROPERTIES CXX_STANDARD 17 @@ -118,6 +119,26 @@ endfunction() find_package(CUDAToolkit QUIET) find_package(hip QUIET) +# The rocm-sdk-core pip package (used for wheel builds) ships HIP headers and +# versioned libraries but no CMake configs, so locate them under ROCM_PATH. +if(NOT hip_FOUND AND DEFINED ENV{ROCM_PATH}) + find_path(HIP_INCLUDE_DIR hip/hip_runtime.h PATHS $ENV{ROCM_PATH}/include NO_DEFAULT_PATH) + find_library(AMDHIP64_LIB NAMES amdhip64 libamdhip64.so.7 PATHS $ENV{ROCM_PATH}/lib NO_DEFAULT_PATH) + find_library(HIPRTC_LIB NAMES hiprtc libhiprtc.so.7 PATHS $ENV{ROCM_PATH}/lib NO_DEFAULT_PATH) + + if(HIP_INCLUDE_DIR AND AMDHIP64_LIB AND HIPRTC_LIB) + add_library(hip::host INTERFACE IMPORTED) + target_include_directories(hip::host INTERFACE ${HIP_INCLUDE_DIR}) + target_compile_definitions(hip::host INTERFACE __HIP_PLATFORM_AMD__=1) + target_link_libraries(hip::host INTERFACE ${AMDHIP64_LIB}) + + add_library(hiprtc::hiprtc INTERFACE IMPORTED) + target_link_libraries(hiprtc::hiprtc INTERFACE ${HIPRTC_LIB}) + + set(hip_FOUND TRUE) + endif() +endif() + if(CUDAToolkit_FOUND) message(STATUS "Building stable extension with CUDA backend.") @@ -144,7 +165,9 @@ endif() if(hip_FOUND) message(STATUS "Building stable extension with HIP backend.") - find_package(hiprtc REQUIRED) + if(NOT TARGET hiprtc::hiprtc) + find_package(hiprtc REQUIRED) + endif() add_library(hip_stub_lib SHARED ${EXT_DIR}/stubs/stream.cpp) @@ -163,7 +186,7 @@ if(hip_FOUND) hip::host hiprtc::hiprtc ) - add_stable_extension(torch_stable_hip HIP_BACKEND "${HIP_LINK_LIBS}") + add_stable_extension(oeq_stable_hip HIP_BACKEND "${HIP_LINK_LIBS}") endif() set(OEQ_SYCL_FOUND FALSE) diff --git a/openequivariance/pyproject.toml b/openequivariance/pyproject.toml index 4e573d2e..3b965157 100644 --- a/openequivariance/pyproject.toml +++ b/openequivariance/pyproject.toml @@ -1,5 +1,5 @@ [build-system] -requires = ["setuptools", "setuptools-scm", "scikit-build-core", "nanobind"] +requires = ["setuptools", "setuptools-scm", "scikit-build-core", "nanobind>=3.0"] build-backend = "scikit_build_core.build" [project] @@ -12,13 +12,14 @@ authors = [ { name="James Demmel" } ] description = "A fast GPU JIT kernel generator for the Clebsch-Gordon Tensor Product" -requires-python = ">=3.10" +requires-python = ">=3.11" dependencies = [ "setuptools", "ninja", "jinja2", "numpy", - "pybind11" + "pybind11", + "nanobind-backend>=1.0" ] readme = "README.md" @@ -27,7 +28,6 @@ license-files = ["LICENSE"] classifiers = [ "Programming Language :: Python :: 3", - "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", "Programming Language :: Python :: 3.13", @@ -77,5 +77,31 @@ addopts = [ [tool.ruff] lint.ignore = ["E741"] +[tool.scikit-build] +# Stable-ABI (abi3) wheel for CPython 3.11+, built with nanobind split mode. +wheel.py-api = "cp311" + [tool.scikit-build.cmake.define] AARCH64 = {env="AARCH64", default="0"} + +[tool.cibuildwheel] +# The abi3 wheel built for cp311 covers all later CPython versions. +build = "cp311-*" +skip = "*-musllinux*" +manylinux-x86_64-image = "manylinux_2_28" +manylinux-aarch64-image = "manylinux_2_28" +before-all = "bash {project}/.github/cibw_before_all.sh" +build-verbosity = 1 + +[tool.cibuildwheel.linux.environment] +PATH = "/usr/local/cuda-13.0/bin:$PATH" +CUDAToolkit_ROOT = "/usr/local/cuda-13.0" +ROCM_PATH = "/opt/rocm-sdk-core/_rocm_sdk_core" + +[tool.cibuildwheel.linux] +# GPU runtimes and libtorch come from the user's PyTorch install; never bundle them. +repair-wheel-command = """auditwheel repair -w {dest_dir} {wheel} \ + --exclude libtorch_cpu.so --exclude libc10.so \ + --exclude libtorch_cuda.so --exclude libtorch_hip.so \ + --exclude libcuda.so.1 --exclude 'libcudart.so.*' --exclude 'libnvrtc.so.*' \ + --exclude 'libamdhip64.so.*' --exclude 'libhiprtc.so.*'""" From 010dc36584b22918252083abd4154ac2667f2267 Mon Sep 17 00:00:00 2001 From: Vivek Bharadwaj Date: Sun, 1 Nov 2026 22:45:44 -0800 Subject: [PATCH 24/38] More work on CIBuildwheel. --- .github/cibw_before_all.sh | 26 +++++---- .github/cibw_rocm_requirements.txt | 2 + README.md | 5 ++ openequivariance/CMakeLists.txt | 57 ++++++++++++++++--- .../_torch/extlib/__init__.py | 20 ++++--- .../extension/libtorch_tp_jit_stable.cpp | 19 ++++--- openequivariance/pyproject.toml | 5 +- 7 files changed, 98 insertions(+), 36 deletions(-) create mode 100644 .github/cibw_rocm_requirements.txt diff --git a/.github/cibw_before_all.sh b/.github/cibw_before_all.sh index 8fbbeaec..08f01618 100644 --- a/.github/cibw_before_all.sh +++ b/.github/cibw_before_all.sh @@ -2,9 +2,13 @@ # Runs inside the manylinux container before cibuildwheel builds # openequivariance. Installs only the headers and link-time libraries # for CUDA and HIP; at runtime, the user's PyTorch install provides them. +# +# Configured from [tool.cibuildwheel.linux.environment] in pyproject.toml: +# OEQ_CUDA_ROOTS space-separated CUDA toolkits to install, e.g. +# "/usr/local/cuda-12.6 /usr/local/cuda-13.0" +# and the pinned rocm-sdk-core in .github/cibw_rocm_requirements.txt. set -euo pipefail -CUDA_VERSION=13-0 ARCH=$(uname -m) case "$ARCH" in @@ -14,19 +18,21 @@ esac dnf config-manager --add-repo \ "https://developer.download.nvidia.com/compute/cuda/repos/rhel8/${CUDA_REPO_ARCH}/cuda-rhel8.repo" -dnf install -y \ - cuda-nvcc-${CUDA_VERSION} \ - cuda-cudart-devel-${CUDA_VERSION} \ - cuda-nvrtc-devel-${CUDA_VERSION} \ - cuda-driver-devel-${CUDA_VERSION} + +for root in $OEQ_CUDA_ROOTS; do + version=${root##*/cuda-} + pkg_version=${version//./-} + dnf install -y \ + cuda-nvcc-${pkg_version} \ + cuda-cudart-devel-${pkg_version} \ + cuda-nvrtc-devel-${pkg_version} \ + cuda-driver-devel-${pkg_version} +done # rocm-sdk-core (~415 MB) has the HIP headers, libamdhip64 and libhiprtc, # without the rest of the ROCm toolkit. AMD publishes no aarch64 builds. if [ "$ARCH" = x86_64 ]; then - echo "rocm-sdk-core==10.0.0 --hash=sha256:3bf4a72d11aa2a4ee1e90572c73630f937d38c1b7af4dda45b483b54655138a7" \ - > /tmp/rocm-requirements.txt /opt/python/cp312-cp312/bin/python -m pip install \ --no-deps --require-hashes --target /opt/rocm-sdk-core \ - --index-url https://download.pytorch.org/whl/rocm10.0 \ - -r /tmp/rocm-requirements.txt + -r "$(dirname "$0")/cibw_rocm_requirements.txt" fi diff --git a/.github/cibw_rocm_requirements.txt b/.github/cibw_rocm_requirements.txt new file mode 100644 index 00000000..5448bf02 --- /dev/null +++ b/.github/cibw_rocm_requirements.txt @@ -0,0 +1,2 @@ +--index-url https://download.pytorch.org/whl/rocm10.0 +rocm-sdk-core==10.0.0 --hash=sha256:3bf4a72d11aa2a4ee1e90572c73630f937d38c1b7af4dda45b483b54655138a7 diff --git a/README.md b/README.md index 34b88cee..bbf31055 100644 --- a/README.md +++ b/README.md @@ -20,6 +20,11 @@ and GCC 9+ available before installing our package via pip install openequivariance ``` +Linux wheels include a precompiled extension for PyTorch built +with CUDA 12 or 13; other setups compile the extension on first import. +To build from source against your local CUDA toolkit instead, run +`pip install --no-binary openequivariance openequivariance`. + We provide up to an order of magnitude acceleration over e3nn perform on par with the latest version of [NVIDIA cuEquivariance](https://github.com/NVIDIA/cuEquivariance), which has a closed-source kernel package. diff --git a/openequivariance/CMakeLists.txt b/openequivariance/CMakeLists.txt index c5bf115f..2c4d2f34 100644 --- a/openequivariance/CMakeLists.txt +++ b/openequivariance/CMakeLists.txt @@ -66,10 +66,11 @@ function(add_stable_extension target_name backend_define link_libraries) ) # Enforce CXX11 ABI to match LibTorch - target_compile_definitions(${target_name} PRIVATE + target_compile_definitions(${target_name} PRIVATE ${backend_define}=1 _GLIBCXX_USE_CXX11_ABI=1 INCLUDE_NB_EXTENSION + EXTENSION_NAME=${target_name} ) target_include_directories(${target_name} PRIVATE @@ -116,7 +117,6 @@ function(add_stable_extension target_name backend_define link_libraries) install(TARGETS ${aoti_target_name} LIBRARY DESTINATION "${OEQ_INSTALL_DIR}") endfunction() -find_package(CUDAToolkit QUIET) find_package(hip QUIET) # The rocm-sdk-core pip package (used for wheel builds) ships HIP headers and @@ -139,30 +139,71 @@ if(NOT hip_FOUND AND DEFINED ENV{ROCM_PATH}) endif() endif() -if(CUDAToolkit_FOUND) - message(STATUS "Building stable extension with CUDA backend.") +# CUDA: one extension per CUDA major version, named oeq_stable_cuda, so +# that it links the same libcudart/libnvrtc major as the user's PyTorch. +# +# OEQ_CUDA_ROOTS (CMake or environment variable, ;- or space-separated) lists +# the toolkit roots to build against, e.g. "/usr/local/cuda-12.6 /usr/local/cuda-13.0". +# If unset, a single extension is built against the toolkit find_package finds. +set(OEQ_CUDA_ROOTS "$ENV{OEQ_CUDA_ROOTS}" CACHE STRING "CUDA toolkit roots to build against") +string(REPLACE " " ";" OEQ_CUDA_ROOTS "${OEQ_CUDA_ROOTS}") + +set(OEQ_CUDA_FOUND FALSE) +if(OEQ_CUDA_ROOTS) + set(OEQ_CUDA_FOUND TRUE) +else() + find_package(CUDAToolkit QUIET) + set(OEQ_CUDA_FOUND ${CUDAToolkit_FOUND}) +endif() +if(OEQ_CUDA_FOUND) add_library(cuda_stub_lib SHARED ${EXT_DIR}/stubs/stream.cpp) target_include_directories(cuda_stub_lib PRIVATE ${LIBTORCH_INCLUDE_DIR} ) - set_target_properties(cuda_stub_lib PROPERTIES - OUTPUT_NAME "torch_cuda" + set_target_properties(cuda_stub_lib PROPERTIES + OUTPUT_NAME "torch_cuda" POSITION_INDEPENDENT_CODE ON CXX_STANDARD 17 ) +endif() +if(OEQ_CUDA_FOUND AND NOT OEQ_CUDA_ROOTS) + message(STATUS "Building stable extension with CUDA ${CUDAToolkit_VERSION} backend.") set(CUDA_LINK_LIBS CUDA::cudart CUDA::cuda_driver CUDA::nvrtc cuda_stub_lib ) - add_stable_extension(oeq_stable_cuda CUDA_BACKEND "${CUDA_LINK_LIBS}") + add_stable_extension(oeq_stable_cuda${CUDAToolkit_VERSION_MAJOR} CUDA_BACKEND "${CUDA_LINK_LIBS}") endif() +foreach(cuda_root ${OEQ_CUDA_ROOTS}) + file(STRINGS "${cuda_root}/include/cuda.h" cuda_version_line REGEX "^#define CUDA_VERSION [0-9]+") + if(NOT cuda_version_line) + message(FATAL_ERROR "No CUDA toolkit found at ${cuda_root}") + endif() + string(REGEX REPLACE ".* ([0-9]+)$" "\\1" cuda_version "${cuda_version_line}") + math(EXPR cuda_major "${cuda_version} / 1000") + math(EXPR cuda_minor "(${cuda_version} % 1000) / 10") + message(STATUS "Building stable extension with CUDA ${cuda_major}.${cuda_minor} backend from ${cuda_root}.") + + set(cuda_lib_dir "${cuda_root}/lib64") + add_library(oeq_cuda${cuda_major}_toolkit INTERFACE) + target_include_directories(oeq_cuda${cuda_major}_toolkit INTERFACE "${cuda_root}/include") + target_link_libraries(oeq_cuda${cuda_major}_toolkit INTERFACE + "${cuda_lib_dir}/libcudart.so" + "${cuda_lib_dir}/libnvrtc.so" + "${cuda_lib_dir}/stubs/libcuda.so" + ) + + set(CUDA_LINK_LIBS oeq_cuda${cuda_major}_toolkit cuda_stub_lib) + add_stable_extension(oeq_stable_cuda${cuda_major} CUDA_BACKEND "${CUDA_LINK_LIBS}") +endforeach() + if(hip_FOUND) message(STATUS "Building stable extension with HIP backend.") if(NOT TARGET hiprtc::hiprtc) @@ -223,6 +264,6 @@ if(OEQ_SYCL_FOUND) endforeach() endif() -if(NOT CUDAToolkit_FOUND AND NOT hip_FOUND AND NOT OEQ_SYCL_FOUND) +if(NOT OEQ_CUDA_FOUND AND NOT hip_FOUND AND NOT OEQ_SYCL_FOUND) message(WARNING "None of CUDAToolkit, HIP or SYCL was found. The stable extension will not be built.") endif() diff --git a/openequivariance/openequivariance/_torch/extlib/__init__.py b/openequivariance/openequivariance/_torch/extlib/__init__.py index 4ec88480..646efd5b 100644 --- a/openequivariance/openequivariance/_torch/extlib/__init__.py +++ b/openequivariance/openequivariance/_torch/extlib/__init__.py @@ -5,6 +5,7 @@ import warnings import sysconfig import contextlib +import importlib from pathlib import Path from packaging.version import Version @@ -51,6 +52,12 @@ def _detect_backend(): DEVICE_TYPE = "xpu" if IS_SYCL else "cuda" +# CUDA extensions are built per CUDA major version, matching torch's libcudart/libnvrtc. +if BACKEND == "cuda": + EXTENSION_NAME = f"oeq_stable_cuda{torch.version.cuda.split('.')[0]}" +else: + EXTENSION_NAME = f"oeq_stable_{BACKEND}" + @contextlib.contextmanager def _maybe_set_max_jobs_env(limit=8): @@ -209,12 +216,9 @@ def load_precompiled_extension(): True # Doesn't actually use libpython, just set this as true anyway ) try: - if BACKEND == "cuda": - import openequivariance._torch.extlib.oeq_stable_cuda as extension_module - elif BACKEND == "hip": - import openequivariance._torch.extlib.oeq_stable_hip as extension_module - elif BACKEND == "sycl": - import openequivariance._torch.extlib.oeq_stable_sycl as extension_module + extension_module = importlib.import_module( + f"openequivariance._torch.extlib.{EXTENSION_NAME}" + ) torch.ops.load_library(extension_module.__file__) BUILT_EXTENSION = True @@ -243,10 +247,10 @@ def load_precompiled_extension(): WARNING_MESSAGE += "SYCL does not support precompiled extension yet.\n" USE_PRECOMPILED_EXTENSION = False -AOTI_SO_NAME = f"liboeq_stable_{BACKEND}_aoti.so" +AOTI_SO_NAME = f"lib{EXTENSION_NAME}_aoti.so" if not os.path.exists(os.path.join(os.path.dirname(__file__), AOTI_SO_NAME)): - WARNING_MESSAGE += "Precompiled extension shared object not found.\n" + WARNING_MESSAGE += f"Precompiled extension {AOTI_SO_NAME} not found.\n" USE_PRECOMPILED_EXTENSION = False diff --git a/openequivariance/openequivariance/extension/libtorch_tp_jit_stable.cpp b/openequivariance/openequivariance/extension/libtorch_tp_jit_stable.cpp index 96a84600..0563c57c 100644 --- a/openequivariance/openequivariance/extension/libtorch_tp_jit_stable.cpp +++ b/openequivariance/openequivariance/extension/libtorch_tp_jit_stable.cpp @@ -154,15 +154,18 @@ bool tensor_is_on_gpu(const Tensor &tensor) { } // namespace oeq -#ifdef CUDA_BACKEND - #define EXTENSION_NAME oeq_stable_cuda -#endif -#ifdef HIP_BACKEND - #define EXTENSION_NAME oeq_stable_hip +// CMake passes EXTENSION_NAME to match the target name (e.g. oeq_stable_cuda13) +#ifndef EXTENSION_NAME + #ifdef CUDA_BACKEND + #define EXTENSION_NAME oeq_stable_cuda + #endif + #ifdef HIP_BACKEND + #define EXTENSION_NAME oeq_stable_hip + #endif + #ifdef SYCL_BACKEND + #define EXTENSION_NAME oeq_stable_sycl + #endif #endif -#ifdef SYCL_BACKEND - #define EXTENSION_NAME oeq_stable_sycl -#endif #ifdef INCLUDE_NB_EXTENSION #include "nanobind/nanobind.h" diff --git a/openequivariance/pyproject.toml b/openequivariance/pyproject.toml index 3b965157..7f3d4c60 100644 --- a/openequivariance/pyproject.toml +++ b/openequivariance/pyproject.toml @@ -94,8 +94,9 @@ before-all = "bash {project}/.github/cibw_before_all.sh" build-verbosity = 1 [tool.cibuildwheel.linux.environment] -PATH = "/usr/local/cuda-13.0/bin:$PATH" -CUDAToolkit_ROOT = "/usr/local/cuda-13.0" +# CUDA toolkits to build against: one extension per major version, chosen at +# runtime from torch.version.cuda. Use the oldest minor PyTorch ships per major. +OEQ_CUDA_ROOTS = "/usr/local/cuda-12.6 /usr/local/cuda-13.0" ROCM_PATH = "/opt/rocm-sdk-core/_rocm_sdk_core" [tool.cibuildwheel.linux] From e5bdea45910948504b75880c82584994a5c75c57 Mon Sep 17 00:00:00 2001 From: Vivek Bharadwaj Date: Sun, 8 Nov 2026 18:38:22 -0800 Subject: [PATCH 25/38] Symlink and other modifications. --- .github/autoversion.py | 44 ---- .github/release.py | 78 +++++++ .github/workflows/release.yaml | 12 +- .github/workflows/verify_extension_build.yml | 12 +- README.md | 7 +- docs/installation.rst | 38 ++- openequivariance/README.md | 221 +----------------- .../_torch/extlib/__init__.py | 104 ++++----- 8 files changed, 149 insertions(+), 367 deletions(-) delete mode 100644 .github/autoversion.py create mode 100644 .github/release.py mode change 100644 => 120000 openequivariance/README.md diff --git a/.github/autoversion.py b/.github/autoversion.py deleted file mode 100644 index dbe884e6..00000000 --- a/.github/autoversion.py +++ /dev/null @@ -1,44 +0,0 @@ -""" -Sets the version in a pyproject.toml from git tags of the form -vX.Y.Z on HEAD. - -Usage: python3 .github/autoversion.py [...] -""" - -import re -import subprocess -import sys -from pathlib import Path - -import tomlkit - -TAG_PATTERN = re.compile(r"v([0-9]+\.[0-9]+\.[0-9]+)") - - -def get_version_from_tag(): - tags = subprocess.run( - ["git", "tag", "--points-at", "HEAD"], - check=True, - capture_output=True, - text=True, - ).stdout.splitlines() - - versions = [m.group(1) for m in map(TAG_PATTERN.fullmatch, tags) if m] - assert len(versions) == 1 - return versions[0] - - -def set_version(pyproject, version): - doc = tomlkit.parse(pyproject.read_text()) - assert "version" in doc["project"] - doc["project"]["version"] = version - pyproject.write_text(tomlkit.dumps(doc)) - - -if __name__ == "__main__": - if len(sys.argv) < 2: - sys.exit(__doc__) - - version = get_version_from_tag() - for path in sys.argv[1:]: - set_version(Path(path), version) diff --git a/.github/release.py b/.github/release.py new file mode 100644 index 00000000..e221c4b0 --- /dev/null +++ b/.github/release.py @@ -0,0 +1,78 @@ +""" +Release pipeline steps for .github/workflows/release.yaml. + + python3 .github/release.py resolve-symlinks + python3 .github/release.py set-version [...] +""" + +import argparse +import re +import shutil +import subprocess +from pathlib import Path + +TAG_PATTERN = re.compile(r"v([0-9]+\.[0-9]+\.[0-9]+)") +GIT_SYMLINK_MODE = "120000" + + +def resolve_symlinks(args): + """Replaces every symlink tracked by git with a copy of its target.""" + entries = subprocess.run( + ["git", "ls-files", "-s", "-z"], + check=True, + capture_output=True, + text=True, + ).stdout.split("\0") + + for entry in filter(None, entries): + info, path = entry.split("\t", 1) + if info.split()[0] == GIT_SYMLINK_MODE: + link = Path(path) + target = link.resolve() + link.unlink() + if target.is_dir(): + shutil.copytree(target, link) + else: + shutil.copy2(target, link) + print(f"Resolved symlink {link} -> {target}") + + +def get_version_from_tag(): + tags = subprocess.run( + ["git", "tag", "--points-at", "HEAD"], + check=True, + capture_output=True, + text=True, + ).stdout.splitlines() + + versions = [m.group(1) for m in map(TAG_PATTERN.fullmatch, tags) if m] + assert len(versions) == 1 + return versions[0] + + +def set_version(args): + """Sets the version in each pyproject.toml from the git tag of the form + vX.Y.Z on HEAD.""" + import tomlkit + + version = get_version_from_tag() + for pyproject in args.pyprojects: + doc = tomlkit.parse(pyproject.read_text()) + assert "version" in doc["project"] + doc["project"]["version"] = version + pyproject.write_text(tomlkit.dumps(doc)) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description=__doc__.split("\n\n")[0].strip()) + steps = parser.add_subparsers(dest="step", required=True) + + step = steps.add_parser("resolve-symlinks", help=resolve_symlinks.__doc__) + step.set_defaults(func=resolve_symlinks) + + step = steps.add_parser("set-version", help=set_version.__doc__) + step.add_argument("pyprojects", nargs="+", type=Path) + step.set_defaults(func=set_version) + + args = parser.parse_args() + args.func(args) diff --git a/.github/workflows/release.yaml b/.github/workflows/release.yaml index 346b0600..7f5afd7d 100644 --- a/.github/workflows/release.yaml +++ b/.github/workflows/release.yaml @@ -23,10 +23,12 @@ jobs: uses: actions/setup-python@6d4fd22b20465ccc0f7354476588e8df7c9ca0e0 # v7.1.0 with: python-version: '3.12' + - name: replace symlinks with copies + run: python3 .github/release.py resolve-symlinks - name: set version from tag run: | python3 -m pip install -r .github/workflows/requirements_release_ci.txt - python3 .github/autoversion.py openequivariance/pyproject.toml + python3 .github/release.py set-version openequivariance/pyproject.toml - name: build wheels run: python3 -m cibuildwheel --output-dir wheelhouse openequivariance - name: store the wheels @@ -46,10 +48,12 @@ jobs: uses: actions/setup-python@6d4fd22b20465ccc0f7354476588e8df7c9ca0e0 # v7.1.0 with: python-version: '3.12' + - name: replace symlinks with copies + run: python3 .github/release.py resolve-symlinks - name: set version from tag run: | python3 -m pip install -r .github/workflows/requirements_release_ci.txt - python3 .github/autoversion.py openequivariance/pyproject.toml + python3 .github/release.py set-version openequivariance/pyproject.toml - name: install dependencies, then build source tarball run: | cd openequivariance @@ -93,10 +97,12 @@ jobs: uses: actions/setup-python@6d4fd22b20465ccc0f7354476588e8df7c9ca0e0 # v7.1.0 with: python-version: '3.10' + - name: replace symlinks with copies + run: python3 .github/release.py resolve-symlinks - name: set version from tag run: | python3 -m pip install -r .github/workflows/requirements_release_ci.txt - python3 .github/autoversion.py openequivariance_extjax/pyproject.toml + python3 .github/release.py set-version openequivariance_extjax/pyproject.toml - name: install dependencies, then build source tarball run: | cd openequivariance_extjax diff --git a/.github/workflows/verify_extension_build.yml b/.github/workflows/verify_extension_build.yml index c6b7c16a..44afeeee 100644 --- a/.github/workflows/verify_extension_build.yml +++ b/.github/workflows/verify_extension_build.yml @@ -20,7 +20,7 @@ jobs: - name: Set up Python uses: actions/setup-python@6d4fd22b20465ccc0f7354476588e8df7c9ca0e0 # v7.1.0 with: - python-version: "3.10" + python-version: "3.12" cache: 'pip' cache-dependency-path: '**/requirements_cuda_ci.txt' - name: Install dependencies @@ -35,10 +35,6 @@ jobs: run: | pytest tests/import_test.py - export OEQ_JIT_EXTENSION=1 - - pytest tests/import_test.py - - name: Test JAX extension build run: | XLA_DIRECT_DOWNLOAD=1 pip install -e "./openequivariance_extjax" --no-build-isolation @@ -52,7 +48,7 @@ jobs: - name: Set up Python uses: actions/setup-python@6d4fd22b20465ccc0f7354476588e8df7c9ca0e0 # v7.1.0 with: - python-version: "3.10" + python-version: "3.12" cache: 'pip' cache-dependency-path: '**/requirements_sycl_ci.txt' @@ -78,7 +74,3 @@ jobs: export CXX=icpx pytest tests/import_test.py - - export OEQ_JIT_EXTENSION=1 - - pytest tests/import_test.py \ No newline at end of file diff --git a/README.md b/README.md index bbf31055..34de8a61 100644 --- a/README.md +++ b/README.md @@ -13,16 +13,15 @@ that [e3nn](https://e3nn.org/) supports commonly found in graph neural networks (e.g. [Nequip](https://github.com/mir-group/nequip) or [MACE](https://github.com/ACEsuit/mace)). To get -started with PyTorch, ensure that you have PyTorch -and GCC 9+ available before installing our package via +started with PyTorch (2.10 or newer), install our package via ```bash pip install openequivariance ``` Linux wheels include a precompiled extension for PyTorch built -with CUDA 12 or 13; other setups compile the extension on first import. -To build from source against your local CUDA toolkit instead, run +with CUDA 12, CUDA 13, or ROCm 7. For other setups, build the +extension from source (requires GCC 9+ and your CUDA / HIP toolkit) with `pip install --no-binary openequivariance openequivariance`. We provide up to an order of magnitude acceleration over e3nn perform on par with the latest diff --git a/docs/installation.rst b/docs/installation.rst index 17feab14..144583bf 100644 --- a/docs/installation.rst +++ b/docs/installation.rst @@ -8,11 +8,11 @@ Installation You need the following to install OpenEquivariance: - A Linux system equipped with an NVIDIA / AMD / Intel graphics card. -- Either PyTorch >= 2.4 (>= 2.10 for AOTI, export, and Intel GPUs), or - JAX>0.5.0 with CUDA, RocM, or XPU support. -- GCC 9+ and the CUDA / HIP toolkit, or the oneAPI DPC++ compiler - (``icpx``) for Intel GPUs. The command - ``c++ --version`` should return >= 9.0; see below for details on +- Either PyTorch >= 2.10, or + JAX>0.5.0 with CUDA, RocM, or XPU support. +- For source builds only: GCC 9+ and the CUDA / HIP toolkit, or the + oneAPI DPC++ compiler (``icpx``) for Intel GPUs. The command + ``c++ --version`` should return >= 9.0; see below for details on setting an alternate compiler. .. note:: @@ -32,32 +32,26 @@ You need the following to install OpenEquivariance: pip install openequivariance python -c "import openequivariance" - The second line triggers a build of the C++ extension we use to compile - kernels, which can take a couple of minutes. Subsequent imports are - much faster since this extension is cached. + On Linux, the wheel includes a precompiled C++ extension for PyTorch + built with CUDA 12, CUDA 13, or ROCm 7. If you see a warning that no + precompiled extension matches your PyTorch build, build it from + source instead (this requires GCC 9+ and the CUDA / HIP toolkit): - To support ``torch.compile``, ``torch.export``, and - JITScript, OpenEquivariance needs to compile a C++ extension - tightly integrated with PyTorch. If you see a warning that - this extension could not be compiled, first check: + .. code-block:: bash - .. code-block:: bash + pip install --no-binary openequivariance openequivariance - c++ --version - - To build the extension with an alternate compiler, set the - ``CC`` and ``CXX`` - environment variable and retry the import: + To build the extension with an alternate compiler, set the + ``CC`` and ``CXX`` environment variables before running ``pip``: .. code-block:: bash export CC=/path/to/your/gcc export CXX=/path/to/your/g++ - python -c "import openequivariance" - - These configuration steps are required only ONCE after - installation (or upgrade) with pip. + On Intel GPUs, the extension is still compiled on first import, + which can take a couple of minutes; subsequent imports use the + cached build. .. tab:: JAX NVIDIA GPUs diff --git a/openequivariance/README.md b/openequivariance/README.md deleted file mode 100644 index 3efe59a7..00000000 --- a/openequivariance/README.md +++ /dev/null @@ -1,220 +0,0 @@ -# OpenEquivariance -[![OEQ C++ Extension Build Verification](https://github.com/PASSIONLab/OpenEquivariance/actions/workflows/verify_extension_build.yml/badge.svg?event=push)](https://github.com/PASSIONLab/OpenEquivariance/actions/workflows/verify_extension_build.yml) -[![License](https://img.shields.io/badge/License-BSD_3--Clause-blue.svg)](https://opensource.org/licenses/BSD-3-Clause) - -[[PyTorch Examples]](#pytorch-examples) -[[JAX Examples]](#jax-examples) -[[Citation and Acknowledgements]](#citation-and-acknowledgements) - -OpenEquivariance is a CUDA and HIP kernel generator for the Clebsch-Gordon tensor product, -a key kernel in rotation-equivariant deep neural networks. -It implements some of the tensor products -that [e3nn](https://e3nn.org/) supports -commonly found in graph neural networks -(e.g. [Nequip](https://github.com/mir-group/nequip) or -[MACE](https://github.com/ACEsuit/mace)). To get -started with PyTorch, ensure that you have PyTorch -and GCC 9+ available before installing our package via - -```bash -pip install openequivariance -``` - -We provide up to an order of magnitude acceleration over e3nn perform on par with the latest -version of [NVIDIA cuEquivariance](https://github.com/NVIDIA/cuEquivariance), -which has a closed-source kernel package. -We also offer fused equivariant graph -convolutions that can reduce -computation and memory consumption significantly. - -For detailed instructions on tests, benchmarks, MACE / Nequip, and our API, -check out the [documentation](https://passionlab.github.io/OpenEquivariance). - -⭐️ **JAX**: Our latest update brings -support for JAX. For NVIDIA GPUs, -install it (after installing JAX) -with the following two commands strictly in order: - -``` bash -pip install openequivariance[jax] -pip install openequivariance_extjax --no-build-isolation -``` - -For AMD GPUs: -``` bash -pip install openequivariance[jax] -JAX_HIP=1 pip install openequivariance_extjax --no-build-isolation -``` - -See the section below for example usage and -our [API page](https://passionlab.github.io/OpenEquivariance/api/) for more details. - -## PyTorch Examples -Here's a CG tensor product implemented by e3nn: - -```python -import torch -import e3nn.o3 as o3 - -gen = torch.Generator(device='cuda') - -batch_size = 1000 -X_ir, Y_ir, Z_ir = o3.Irreps("1x2e"), o3.Irreps("1x3e"), o3.Irreps("1x2e") -X = torch.rand(batch_size, X_ir.dim, device='cuda', generator=gen) -Y = torch.rand(batch_size, Y_ir.dim, device='cuda', generator=gen) - -instructions=[(0, 0, 0, "uvu", True)] - -tp_e3nn = o3.TensorProduct(X_ir, Y_ir, Z_ir, instructions, - shared_weights=False, internal_weights=False).to('cuda') -W = torch.rand(batch_size, tp_e3nn.weight_numel, device='cuda', generator=gen) - -Z = tp_e3nn(X, Y, W) -print(torch.norm(Z)) -``` - -And here's the same tensor product using openequivariance. We require that your -tensors are stored on a CUDA device for this to work: - -```python -import openequivariance as oeq - -problem = oeq.TPProblem(X_ir, Y_ir, Z_ir, instructions, shared_weights=False, internal_weights=False) -tp_fast = oeq.TensorProduct(problem, torch_op=True) - -Z = tp_fast(X, Y, W) # Reuse X, Y, W from earlier -print(torch.norm(Z)) -``` - -Our interface for `oeq.TPProblem` is almost a strict superset of -`o3.TensorProduct` (two key differences: we -impose `internal_weights=False` and add support for multiple datatypes). -You can pass e3nn `Irreps` instances directly or -use `oeq.Irreps`, which is identical. - -We recommend reading the [e3nn documentation and API reference](https://docs.e3nn.org/en/latest/) first, then using our kernels -as drop-in replacements. We support most "uvu" and "uvw" tensor products; -see [this section](#tensor-products-we-accelerate) for an up-to-date list of supported configurations. - -**Important**: For many configurations, our kernels return results identical to -e3nn up to floating point roundoff (this includes all "uvu" problems with -multiplicity 1 for all irreps in the second input). For other configurations -(e.g. any "uvw" connection modes), we return identical -results up to a well-defined reordering of the weights relative to e3nn. - -If you're executing tensor products as part of a message passing graph -neural network, we offer fused kernels that save both memory and compute time: - -```python -from torch_geometric import EdgeIndex - -node_ct, nonzero_ct = 3, 4 - -# Receiver, sender indices for message passing GNN -edge_index = EdgeIndex( - [[0, 1, 1, 2], # Receiver - [1, 0, 2, 1]], # Sender - device='cuda', - dtype=torch.long) - -X = torch.rand(node_ct, X_ir.dim, device='cuda', generator=gen) -Y = torch.rand(nonzero_ct, Y_ir.dim, device='cuda', generator=gen) -W = torch.rand(nonzero_ct, problem.weight_numel, device='cuda', generator=gen) - -tp_conv = oeq.TensorProductConv(problem, torch_op=True, deterministic=False) # Reuse problem from earlier -Z = tp_conv.forward(X, Y, W, edge_index[0], edge_index[1]) # Z has shape [node_ct, z_ir.dim] -print(torch.norm(Z)) -``` - -If you can guarantee `EdgeIndex` is sorted by receiver index and supply the transpose -permutation, we can provide even greater speedup (and deterministic results) -by avoiding atomics: - -```python -_, sender_perm = edge_index.sort_by("col") # Sort by sender index -edge_index, receiver_perm = edge_index.sort_by("row") # Sort by receiver index - -# Now we can use the faster deterministic algorithm -tp_conv = oeq.TensorProductConv(problem, torch_op=True, deterministic=True) -Z = tp_conv.forward(X, Y[receiver_perm], W[receiver_perm], edge_index[0], edge_index[1], sender_perm) -print(torch.norm(Z)) -``` -**Note**: you don't need Pytorch geometric to use our kernels. When -`deterministic=False`, the `sender` and `receiver` indices can have -arbitrary order. - -## JAX Examples -After installation, use the library -as follows. Set `OEQ_NOTORCH=1` -in your environment to avoid the PyTorch import in -the regular `openequivariance` package. -```python -import jax -import os - -os.environ["OEQ_NOTORCH"] = "1" -import openequivariance as oeq - -seed = 42 -key = jax.random.PRNGKey(seed) - -batch_size = 1000 -X_ir, Y_ir, Z_ir = oeq.Irreps("1x2e"), oeq.Irreps("1x3e"), oeq.Irreps("1x2e") -problem = oeq.TPProblem(X_ir, Y_ir, Z_ir, [(0, 0, 0, "uvu", True)], shared_weights=False, internal_weights=False) - - -node_ct, nonzero_ct = 3, 4 -edge_index = jax.numpy.array( - [ - [0, 1, 1, 2], - [1, 0, 2, 1], - ], - dtype=jax.numpy.int32, # NOTE: This int32, not int64 -) - -X = jax.random.uniform(key, shape=(node_ct, X_ir.dim), minval=0.0, maxval=1.0, dtype=jax.numpy.float32) -Y = jax.random.uniform(key, shape=(nonzero_ct, Y_ir.dim), - minval=0.0, maxval=1.0, dtype=jax.numpy.float32) -W = jax.random.uniform(key, shape=(nonzero_ct, problem.weight_numel), - minval=0.0, maxval=1.0, dtype=jax.numpy.float32) - -tp_conv = oeq.jax.TensorProductConv(problem, deterministic=False) -Z = tp_conv.forward( - X, Y, W, edge_index[0], edge_index[1] -) -print(jax.numpy.linalg.norm(Z)) - -# Test JAX JIT -jitted = jax.jit(lambda X, Y, W, e1, e2: tp_conv.forward(X, Y, W, e1, e2)) -print(jax.numpy.linalg.norm(jitted(X, Y, W, edge_index[0], edge_index[1]))) -``` - -## Citation and Acknowledgements -If you find this code useful, please cite our paper: - -```bibtex -@inbook{openequivariance, -author={Vivek Bharadwaj and Austin Glover and Aydin Buluc and James Demmel}, -title={An Efficient Sparse Kernel Generator for O(3)-Equivariant Deep Networks}, -booktitle = {SIAM Conference on Applied and Computational Discrete Algorithms (ACDA25)}, -chapter = {}, -url={https://arxiv.org/abs/2501.13986}, -publisher={Society for Industrial and Applied Mathematics}, -year={2025} -} -``` - -Our codebase includes a lightweight clone of -[e3nn](https://e3nn.org/)'s frontend interface (in particular, the -`TensorProduct` and `Irreps` classes). We removed references to Pytorch -and separated the implementation from the problem description (for future -frontend support outside of torch). We also extracted the Wigner 3j tensor generating code from QuTiP. Thank you to the current -developers and maintainers! - -## Copyright - -Copyright (c) 2025, The Regents of the University of California, through Lawrence Berkeley National Laboratory (subject to receipt of any required approvals from the U.S. Dept. of Energy). All rights reserved. - -If you have questions about your rights to use or distribute this software, please contact Berkeley Lab's Intellectual Property Office at IPO@lbl.gov. - -NOTICE. This Software was developed under funding from the U.S. Department of Energy and the U.S. Government consequently retains certain rights. As such, the U.S. Government has been granted for itself and others acting on its behalf a paid-up, nonexclusive, irrevocable, worldwide license in the Software to reproduce, distribute copies to the public, prepare derivative works, and perform publicly and display publicly, and to permit others to do so. \ No newline at end of file diff --git a/openequivariance/README.md b/openequivariance/README.md new file mode 120000 index 00000000..32d46ee8 --- /dev/null +++ b/openequivariance/README.md @@ -0,0 +1 @@ +../README.md \ No newline at end of file diff --git a/openequivariance/openequivariance/_torch/extlib/__init__.py b/openequivariance/openequivariance/_torch/extlib/__init__.py index 646efd5b..679f408c 100644 --- a/openequivariance/openequivariance/_torch/extlib/__init__.py +++ b/openequivariance/openequivariance/_torch/extlib/__init__.py @@ -11,7 +11,6 @@ import torch -from openequivariance.core.logging import getLogger oeq_root = str(Path(__file__).parent.parent.parent) @@ -88,6 +87,8 @@ def _wait_for_torch_build_lock(name, timeout=300): def load_jit_extension(): + """Compile the SYCL extension at first import. CUDA and HIP use only the + precompiled stable extension.""" global \ BUILT_EXTENSION, \ BUILT_EXTENSION_ERROR, \ @@ -141,50 +142,30 @@ def load_jit_extension(): f"-l{python_lib_name}", ], ) - if BACKEND == "cuda": - extra_link_args.extend(["-lcuda", "-lcudart", "-lnvrtc", "-ltorch_cuda"]) + import shutil - try: - torch_libs, cuda_libs = library_paths("cuda") - extra_link_args.append("-Wl,-rpath," + torch_libs) - extra_link_args.append("-L" + cuda_libs) - if os.path.exists(cuda_libs + "/stubs"): - extra_link_args.append("-L" + cuda_libs + "/stubs") - except Exception as e: - getLogger().info(str(e)) - - extra_cflags.append("-DCUDA_BACKEND") - elif BACKEND == "hip": - torch_libs = library_paths("cuda")[0] - extra_link_args.append("-Wl,-rpath," + torch_libs) - extra_link_args.extend("-L" + path for path in library_paths("cuda")) - extra_link_args.extend(["-ltorch_hip", "-lhiprtc"]) - extra_cflags.append("-DHIP_BACKEND") - elif BACKEND == "sycl": - import shutil - - cxx = os.environ.get("CXX", "") - if "icpx" not in os.path.basename(cxx): - if shutil.which("icpx") is None: - BUILT_EXTENSION_ERROR = ( - "The SYCL backend requires the oneAPI DPC++ compiler. " - "Put 'icpx' on your PATH or set CXX to it." - ) - return - os.environ["CXX"] = "icpx" + cxx = os.environ.get("CXX", "") + if "icpx" not in os.path.basename(cxx): + if shutil.which("icpx") is None: + BUILT_EXTENSION_ERROR = ( + "The SYCL backend requires the oneAPI DPC++ compiler. " + "Put 'icpx' on your PATH or set CXX to it." + ) + return + os.environ["CXX"] = "icpx" - extra_cflags.extend(["-fsycl", "-DSYCL_BACKEND"]) - extra_link_args.extend(["-fsycl", "-lsycl", "-ltorch_xpu", "-lc10_xpu"]) + extra_cflags.extend(["-fsycl", "-DSYCL_BACKEND"]) + extra_link_args.extend(["-fsycl", "-lsycl", "-ltorch_xpu", "-lc10_xpu"]) - for lib_dir in library_paths("xpu"): - extra_link_args.append("-Wl,-rpath," + lib_dir) - extra_link_args.append("-L" + lib_dir) + for lib_dir in library_paths("xpu"): + extra_link_args.append("-Wl,-rpath," + lib_dir) + extra_link_args.append("-L" + lib_dir) torch_sources = [oeq_root + "/extension/" + src for src in torch_sources] include_dirs = ( [oeq_root + "/extension/" + d for d in include_dirs] + extra_include_dirs - + include_paths("xpu" if BACKEND == "sycl" else "cuda") + + include_paths("xpu") ) with warnings.catch_warnings(): @@ -228,38 +209,33 @@ def load_precompiled_extension(): ) -USE_PRECOMPILED_EXTENSION = True -WARNING_MESSAGE = "" - -if os.getenv("OEQ_JIT_EXTENSION", "0") == "1": - WARNING_MESSAGE += "Environment variable OEQ_JIT_EXTENSION=1 is set.\n" - USE_PRECOMPILED_EXTENSION = False - -if Version(torch.__version__) <= Version("2.9.9"): - WARNING_MESSAGE += f"PyTorch version {torch.__version__} is < 2.10, minimum required for precompiled extension. Please upgrade to 2.10.\n" - USE_PRECOMPILED_EXTENSION = False - -if BACKEND == "hip": - WARNING_MESSAGE += "HIP does not support precompiled extension yet.\n" - USE_PRECOMPILED_EXTENSION = False - -if BACKEND == "sycl": - WARNING_MESSAGE += "SYCL does not support precompiled extension yet.\n" - USE_PRECOMPILED_EXTENSION = False - AOTI_SO_NAME = f"lib{EXTENSION_NAME}_aoti.so" -if not os.path.exists(os.path.join(os.path.dirname(__file__), AOTI_SO_NAME)): - WARNING_MESSAGE += f"Precompiled extension {AOTI_SO_NAME} not found.\n" - USE_PRECOMPILED_EXTENSION = False +# CUDA and HIP load only the precompiled stable extension, from a wheel or a +# source build; SYCL is still compiled at first import. +USE_PRECOMPILED_EXTENSION = BACKEND != "sycl" +_SOURCE_BUILD_HINT = ( + "To build the extension for your setup from source, run:\n" + " pip install --no-binary openequivariance openequivariance" +) -if USE_PRECOMPILED_EXTENSION: - load_precompiled_extension() -else: - WARNING_MESSAGE += "For these reasons, falling back to JIT compilation of OpenEquivariance extension. If another process holds the build lock, this waits up to 5 minutes, then raises an error naming the lock file to delete.\n" - warnings.warn(WARNING_MESSAGE, stacklevel=3) +if not USE_PRECOMPILED_EXTENSION: load_jit_extension() +elif Version(torch.__version__) < Version("2.10"): + BUILT_EXTENSION_ERROR = ( + f"OpenEquivariance requires PyTorch >= 2.10, found {torch.__version__}." + ) +elif not os.path.exists(os.path.join(os.path.dirname(__file__), AOTI_SO_NAME)): + BUILT_EXTENSION_ERROR = ( + f"No precompiled OpenEquivariance extension {EXTENSION_NAME} for PyTorch " + f"{torch.__version__}. {_SOURCE_BUILD_HINT}" + ) +else: + load_precompiled_extension() + +if BUILT_EXTENSION_ERROR is not None: + warnings.warn(str(BUILT_EXTENSION_ERROR), stacklevel=3) def torch_ext_so_path(): From 88e6b4a111169c943e7cc13381077fb87f421189 Mon Sep 17 00:00:00 2001 From: Vivek Bharadwaj Date: Mon, 9 Nov 2026 21:01:11 -0800 Subject: [PATCH 26/38] Changes to extlib initialization. --- openequivariance/openequivariance/__init__.py | 10 - .../_torch/extlib/__init__.py | 370 ++++++++---------- tests/export_test.py | 1 - 3 files changed, 162 insertions(+), 219 deletions(-) diff --git a/openequivariance/openequivariance/__init__.py b/openequivariance/openequivariance/__init__.py index 1a3681f5..7c36b4de 100644 --- a/openequivariance/openequivariance/__init__.py +++ b/openequivariance/openequivariance/__init__.py @@ -59,8 +59,6 @@ def _check_package_editable(): ) from openequivariance._torch.extlib import ( - LINKED_LIBPYTHON, - LINKED_LIBPYTHON_ERROR, BUILT_EXTENSION, BUILT_EXTENSION_ERROR, USE_PRECOMPILED_EXTENSION, @@ -83,14 +81,6 @@ def torch_ext_so_path(): import openequivariance_extjax import openequivariance.jax as jax - # TODO-someday: enable - # extjax_version = version("openequivariance_extjax") - # if extjax_version != __version__: - # warnings.warn( - # f"openequivariance_extjax version {extjax_version} does not match " - # f"openequivariance version {__version__}. Ensure both versions match." - # ) - except Exception as e: error = e diff --git a/openequivariance/openequivariance/_torch/extlib/__init__.py b/openequivariance/openequivariance/_torch/extlib/__init__.py index 679f408c..785260b0 100644 --- a/openequivariance/openequivariance/_torch/extlib/__init__.py +++ b/openequivariance/openequivariance/_torch/extlib/__init__.py @@ -11,250 +11,201 @@ import torch - +assert Version(torch.__version__) >= Version("2.10"): oeq_root = str(Path(__file__).parent.parent.parent) -BUILT_EXTENSION = False -BUILT_EXTENSION_ERROR = None - -LINKED_LIBPYTHON = False -LINKED_LIBPYTHON_ERROR = None - -extension_module = None - - -def _detect_backend(): - """HIP builds report a ``torch.version.cuda`` of None, so HIP is tested first.""" - if torch.version.hip: - return "hip" - if torch.version.cuda: - return "cuda" - if getattr(torch.version, "xpu", None): - return "sycl" - return None - - -BACKEND = _detect_backend() - -assert BACKEND is not None, ( - "Only the CUDA, HIP and XPU (SYCL) backends are supported. " - "No supported accelerator was detected in this PyTorch build." +SOURCE_BUILD_HINT = ( + "To build the extension for your setup from source, run:\n" + " pip install --no-binary openequivariance openequivariance" ) -IS_HIP = BACKEND == "hip" -IS_SYCL = BACKEND == "sycl" - -if IS_SYCL and Version(torch.__version__) < Version("2.10"): - raise RuntimeError( - f"The SYCL backend requires PyTorch >= 2.10, found {torch.__version__}." - ) - -DEVICE_TYPE = "xpu" if IS_SYCL else "cuda" - -# CUDA extensions are built per CUDA major version, matching torch's libcudart/libnvrtc. -if BACKEND == "cuda": - EXTENSION_NAME = f"oeq_stable_cuda{torch.version.cuda.split('.')[0]}" -else: - EXTENSION_NAME = f"oeq_stable_{BACKEND}" - - -@contextlib.contextmanager -def _maybe_set_max_jobs_env(limit=8): - if "MAX_JOBS" in os.environ: - yield - return - os.environ["MAX_JOBS"] = str(min(limit, os.cpu_count() or 1)) - try: - yield - finally: - os.environ.pop("MAX_JOBS", None) - - -def _wait_for_torch_build_lock(name, timeout=300): - from torch.utils.cpp_extension import _get_build_directory - - lock_path = os.path.join(_get_build_directory(name, False), "lock") - start = time.time() - while os.path.exists(lock_path): - if time.time() - start > timeout: - raise RuntimeError( - f"Timed out after {timeout} seconds waiting for another process " - f"to finish building the OpenEquivariance extension '{name}'. " - f"If no other build is running, a previous build was likely " - f"killed partway; delete {lock_path} and retry." - ) - time.sleep(1) - - -def load_jit_extension(): - """Compile the SYCL extension at first import. CUDA and HIP use only the - precompiled stable extension.""" - global \ - BUILT_EXTENSION, \ - BUILT_EXTENSION_ERROR, \ - LINKED_LIBPYTHON, \ - LINKED_LIBPYTHON_ERROR, \ - extension_module - - # Locate libpython (required for AOTI) - try: - python_lib_dir = sysconfig.get_config_var("LIBDIR") - major, minor = sys.version_info.major, sys.version_info.minor - python_lib_name = f"python{major}.{minor}" - - libpython_so = os.path.join(python_lib_dir, f"lib{python_lib_name}.so") - libpython_a = os.path.join(python_lib_dir, f"lib{python_lib_name}.a") - if not (os.path.exists(libpython_so) or os.path.exists(libpython_a)): - raise FileNotFoundError( - f"libpython not found, tried {libpython_so} and {libpython_a}" - ) +class Extension: + def __init__(): + self.built_extension = False + self.built_extension_error = None + self.ext_module_python = None + self.aoti_module_path = None + +class JITExtension(Extension): + """ + Only SYCL uses PyTorch's JIT compilation flow; CUDA and HIP now use + precompiled extensions shipped with the wheel. + """ + @contextlib.contextmanager + def maybe_set_max_jobs_env(limit=8): + if "MAX_JOBS" in os.environ: + yield + return + os.environ["MAX_JOBS"] = str(min(limit, os.cpu_count() or 1)) + try: + yield + finally: + os.environ.pop("MAX_JOBS", None) + + def wait_for_torch_build_lock(name, timeout=300): + from torch.utils.cpp_extension import _get_build_directory + + lock_path = os.path.join(_get_build_directory(name, False), "lock") + start = time.time() + while os.path.exists(lock_path): + if time.time() - start > timeout: + raise RuntimeError( + f"Timed out after {timeout} seconds waiting for another process " + f"to finish building the OpenEquivariance extension '{name}'. " + f"If no other build is running, a previous build was likely " + f"killed partway; delete {lock_path} and retry." + ) + time.sleep(1) - LINKED_LIBPYTHON = True - except Exception as e: - LINKED_LIBPYTHON_ERROR = f"Error linking libpython:\n{e}\nSysconfig variables:\n{sysconfig.get_config_vars()}" - try: - from torch.utils.cpp_extension import library_paths, include_paths + def __init__(): + # The JIT extension is a pybind11 module, so it links libpython (when one + # exists) to stay loadable from AOTI C++ programs without an interpreter. + python_lib_dir = sysconfig.get_config_var("LIBDIR") or "" + python_lib_name = f"python{sys.version_info.major}.{sys.version_info.minor}" + found_libpython = any( + os.path.exists(os.path.join(python_lib_dir, f"lib{python_lib_name}{ext}")) + for ext in (".so", ".a") + ) - extra_cflags = ["-O3"] - torch_sources = ["libtorch_tp_jit.cpp", "json11/json11.cpp"] + try: + from torch.utils.cpp_extension import library_paths, include_paths - include_dirs, extra_link_args = (["backend"], ["-Wl,--no-as-needed"]) - extra_include_dirs = [] + extra_cflags = ["-O3"] + torch_sources = ["libtorch_tp_jit.cpp", "json11/json11.cpp"] - try: - import pybind11 + include_dirs, extra_link_args = (["backend"], ["-Wl,--no-as-needed"]) + extra_include_dirs = [] - extra_include_dirs.append(pybind11.get_include()) - except Exception as e: - BUILT_EXTENSION_ERROR = ( - "Could not locate pybind11 include path required for JIT " - f"OpenEquivariance extension compilation: {e}" - ) - return + try: + import pybind11 - if LINKED_LIBPYTHON: - extra_link_args.pop() - extra_link_args.extend( - [ - f"-Wl,--no-as-needed,-rpath,{python_lib_dir}", - f"-L{python_lib_dir}", - f"-l{python_lib_name}", - ], - ) - import shutil - - cxx = os.environ.get("CXX", "") - if "icpx" not in os.path.basename(cxx): - if shutil.which("icpx") is None: - BUILT_EXTENSION_ERROR = ( - "The SYCL backend requires the oneAPI DPC++ compiler. " - "Put 'icpx' on your PATH or set CXX to it." + extra_include_dirs.append(pybind11.get_include()) + except Exception as e: + self.built_extension_error = ( + "Could not locate pybind11 include path required for JIT " + f"OpenEquivariance extension compilation: {e}" ) return - os.environ["CXX"] = "icpx" - extra_cflags.extend(["-fsycl", "-DSYCL_BACKEND"]) - extra_link_args.extend(["-fsycl", "-lsycl", "-ltorch_xpu", "-lc10_xpu"]) - - for lib_dir in library_paths("xpu"): - extra_link_args.append("-Wl,-rpath," + lib_dir) - extra_link_args.append("-L" + lib_dir) - - torch_sources = [oeq_root + "/extension/" + src for src in torch_sources] - include_dirs = ( - [oeq_root + "/extension/" + d for d in include_dirs] - + extra_include_dirs - + include_paths("xpu") - ) + if found_libpython: + extra_link_args.pop() + extra_link_args.extend( + [ + f"-Wl,--no-as-needed,-rpath,{python_lib_dir}", + f"-L{python_lib_dir}", + f"-l{python_lib_name}", + ], + ) + import shutil + + cxx = os.environ.get("CXX", "") + if "icpx" not in os.path.basename(cxx): + if shutil.which("icpx") is None: + self.built_extension_error = ( + "The SYCL backend requires the oneAPI DPC++ compiler. " + "Put 'icpx' on your PATH or set CXX to it." + ) + return + os.environ["CXX"] = "icpx" - with warnings.catch_warnings(): - warnings.simplefilter("ignore") + extra_cflags.extend(["-fsycl", "-DSYCL_BACKEND"]) + extra_link_args.extend(["-fsycl", "-lsycl", "-ltorch_xpu", "-lc10_xpu"]) - try: - _wait_for_torch_build_lock("libtorch_tp_jit") - with _maybe_set_max_jobs_env(): - extension_module = torch.utils.cpp_extension.load( - "libtorch_tp_jit", - torch_sources, - extra_cflags=extra_cflags, - extra_include_paths=include_dirs, - extra_ldflags=extra_link_args, - ) - torch.ops.load_library(extension_module.__file__) - BUILT_EXTENSION = True - except Exception as e: - # If compiling torch fails (e.g. low gcc version), we should fall back to the - # version that takes integer pointers as args (but is untraceable to PyTorch JIT / export). - BUILT_EXTENSION_ERROR = e - except Exception as e: - BUILT_EXTENSION_ERROR = f"Error JIT-compiling OpenEquivariance Extension: {e}" + for lib_dir in library_paths("xpu"): + extra_link_args.append("-Wl,-rpath," + lib_dir) + extra_link_args.append("-L" + lib_dir) + torch_sources = [oeq_root + "/extension/" + src for src in torch_sources] + include_dirs = ( + [oeq_root + "/extension/" + d for d in include_dirs] + + extra_include_dirs + + include_paths("xpu") + ) -def load_precompiled_extension(): - global BUILT_EXTENSION, BUILT_EXTENSION_ERROR, LINKED_LIBPYTHON, extension_module - LINKED_LIBPYTHON = ( - True # Doesn't actually use libpython, just set this as true anyway - ) - try: - extension_module = importlib.import_module( - f"openequivariance._torch.extlib.{EXTENSION_NAME}" - ) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + + try: + _wait_for_torch_build_lock("libtorch_tp_jit") + with _maybe_set_max_jobs_env(): + self.ext_module_python = torch.utils.cpp_extension.load( + "libtorch_tp_jit", + torch_sources, + extra_cflags=extra_cflags, + extra_include_paths=include_dirs, + extra_ldflags=extra_link_args, + ) + torch.ops.load_library(self.ext_module_python.__file__) + self.built_extension = True + self.aoti_module_path =self.ext_module_python.__file__ + + except Exception as e: + # If compiling torch fails (e.g. low gcc version), we should fall back to the + # version that takes integer pointers as args (but is untraceable to PyTorch JIT / export). + self.built_extension_error = e + except Exception as e: + self.built_extension_error = f"Error JIT-compiling OpenEquivariance Extension: {e}" - torch.ops.load_library(extension_module.__file__) - BUILT_EXTENSION = True - except Exception as e: - BUILT_EXTENSION_ERROR = ( - f"Error loading precompiled OpenEquivariance Extension: {e}" - ) +class PrecompiledExtension(Extension): + def __init__(): + try: + ext_module_py_name = f"oeq_stable_cuda{torch.version.cuda.split('.')[0]}" -AOTI_SO_NAME = f"lib{EXTENSION_NAME}_aoti.so" + self.ext_module_python = importlib.import_module( + f"openequivariance._torch.extlib.{ext_module_py_name}" + ) -# CUDA and HIP load only the precompiled stable extension, from a wheel or a -# source build; SYCL is still compiled at first import. -USE_PRECOMPILED_EXTENSION = BACKEND != "sycl" + torch.ops.load_library(extension_module.__file__) + self.built_extension = True -_SOURCE_BUILD_HINT = ( - "To build the extension for your setup from source, run:\n" - " pip install --no-binary openequivariance openequivariance" -) + dirname = os.path.dirname(extension_module.__file__) + self.aoti_module_path = os.path.join(dirname, f"lib{ext_module_py_name}_aoti.so") -if not USE_PRECOMPILED_EXTENSION: - load_jit_extension() -elif Version(torch.__version__) < Version("2.10"): - BUILT_EXTENSION_ERROR = ( - f"OpenEquivariance requires PyTorch >= 2.10, found {torch.__version__}." - ) -elif not os.path.exists(os.path.join(os.path.dirname(__file__), AOTI_SO_NAME)): - BUILT_EXTENSION_ERROR = ( - f"No precompiled OpenEquivariance extension {EXTENSION_NAME} for PyTorch " - f"{torch.__version__}. {_SOURCE_BUILD_HINT}" - ) -else: - load_precompiled_extension() + except Exception as e: + self.built_extension_error = ( + f"Error loading precompiled OpenEquivariance Extension: {e}" + ) -if BUILT_EXTENSION_ERROR is not None: - warnings.warn(str(BUILT_EXTENSION_ERROR), stacklevel=3) +def _detect_backend(): + """HIP builds report a ``torch.version.cuda`` of None, so HIP is tested first.""" + if torch.version.hip: + return "hip" + if torch.version.cuda: + return "cuda" + if getattr(torch.version, "xpu", None): + return "sycl" + assert False, "Invalid Backend" -def torch_ext_so_path(): - if not USE_PRECOMPILED_EXTENSION: - return extension_module.__file__ - else: - dirname = os.path.dirname(extension_module.__file__) - return os.path.join(dirname, AOTI_SO_NAME) +BACKEND = _detect_backend() +IS_HIP = BACKEND == "hip" +IS_SYCL = BACKEND == "sycl" +DEVICE_TYPE = "xpu" if IS_SYCL else "cuda" +extension = None +if BACKEND == "sycl": + extension = JITExtension() +else: + extension = PrecompiledExtension() +sys.modules["oeq_utilities"] = extension.ext_module_python +BUILT_EXTENSION = extension.BUILT_EXTENSION +BUILT_EXTENSION_ERROR = extension.BUILT_EXTENSION_ERROR -sys.modules["oeq_utilities"] = extension_module if BUILT_EXTENSION: from oeq_utilities import ( DeviceProp, GPUTimer, ) -else: + def torch_ext_so_path(): + return extension.aoti_module_path() + +else: + if extension.built_extension_error is not None: + warnings.warn(str(BUILT_EXTENSION_ERRROR), stacklevel=3) def _raise_import_error_helper(import_target: str): if not BUILT_EXTENSION: raise ImportError( @@ -266,3 +217,6 @@ def DeviceProp(*args, **kwargs): def GPUTimer(*args, **kwargs): _raise_import_error_helper("GPUTimer") + + def torch_ext_so_path(): + _raise_import_error_helper("torch_ext_so_path") diff --git a/tests/export_test.py b/tests/export_test.py index 1b77df7a..f84a5310 100644 --- a/tests/export_test.py +++ b/tests/export_test.py @@ -144,7 +144,6 @@ def test_aoti(tp_and_inputs): def test_aoti_cpp_inference(problem_and_irreps): - assert oeq.LINKED_LIBPYTHON, oeq.LINKED_LIBPYTHON_ERROR problem, X_ir, Y_ir, _ = problem_and_irreps cmake_prefix_path = torch.utils.cmake_prefix_path torch_ext_so_path = oeq.torch_ext_so_path() From 4d3fb94693bf25cfd36a92f324b4cd5736a626ad Mon Sep 17 00:00:00 2001 From: Vivek Bharadwaj Date: Mon, 9 Nov 2026 21:26:05 -0800 Subject: [PATCH 27/38] More progress on wheeldist. --- openequivariance/openequivariance/__init__.py | 10 +-- .../_torch/extlib/__init__.py | 78 +++++++++++++------ 2 files changed, 57 insertions(+), 31 deletions(-) diff --git a/openequivariance/openequivariance/__init__.py b/openequivariance/openequivariance/__init__.py index 7c36b4de..b8e02d74 100644 --- a/openequivariance/openequivariance/__init__.py +++ b/openequivariance/openequivariance/__init__.py @@ -58,17 +58,15 @@ def _check_package_editable(): ] ) - from openequivariance._torch.extlib import ( - BUILT_EXTENSION, - BUILT_EXTENSION_ERROR, - USE_PRECOMPILED_EXTENSION, - ) + from openequivariance._torch.extlib import BUILT_EXTENSION, BUILT_EXTENSION_ERROR def torch_ext_so_path(): """ :returns: Path to a ``.so`` file that must be linked to use OpenEquivariance - from the PyTorch C++ Interface. + from the PyTorch C++ Interface, or ``None`` if PyTorch support is + disabled (``OEQ_NOTORCH=1``). + :raises ImportError: If the OpenEquivariance extension failed to load. """ try: return torch_ext_so_path_internal() diff --git a/openequivariance/openequivariance/_torch/extlib/__init__.py b/openequivariance/openequivariance/_torch/extlib/__init__.py index 785260b0..94555480 100644 --- a/openequivariance/openequivariance/_torch/extlib/__init__.py +++ b/openequivariance/openequivariance/_torch/extlib/__init__.py @@ -11,7 +11,10 @@ import torch -assert Version(torch.__version__) >= Version("2.10"): +if Version(torch.__version__) < Version("2.10"): + raise RuntimeError( + f"OpenEquivariance requires PyTorch >= 2.10, found {torch.__version__}." + ) oeq_root = str(Path(__file__).parent.parent.parent) SOURCE_BUILD_HINT = ( @@ -19,18 +22,22 @@ " pip install --no-binary openequivariance openequivariance" ) + class Extension: - def __init__(): + def __init__(self): self.built_extension = False self.built_extension_error = None self.ext_module_python = None self.aoti_module_path = None + class JITExtension(Extension): """ Only SYCL uses PyTorch's JIT compilation flow; CUDA and HIP now use precompiled extensions shipped with the wheel. """ + + @staticmethod @contextlib.contextmanager def maybe_set_max_jobs_env(limit=8): if "MAX_JOBS" in os.environ: @@ -42,6 +49,7 @@ def maybe_set_max_jobs_env(limit=8): finally: os.environ.pop("MAX_JOBS", None) + @staticmethod def wait_for_torch_build_lock(name, timeout=300): from torch.utils.cpp_extension import _get_build_directory @@ -57,8 +65,8 @@ def wait_for_torch_build_lock(name, timeout=300): ) time.sleep(1) - - def __init__(): + def __init__(self): + super().__init__() # The JIT extension is a pybind11 module, so it links libpython (when one # exists) to stay loadable from AOTI C++ programs without an interpreter. python_lib_dir = sysconfig.get_config_var("LIBDIR") or "" @@ -127,8 +135,8 @@ def __init__(): warnings.simplefilter("ignore") try: - _wait_for_torch_build_lock("libtorch_tp_jit") - with _maybe_set_max_jobs_env(): + self.wait_for_torch_build_lock("libtorch_tp_jit") + with self.maybe_set_max_jobs_env(): self.ext_module_python = torch.utils.cpp_extension.load( "libtorch_tp_jit", torch_sources, @@ -138,34 +146,49 @@ def __init__(): ) torch.ops.load_library(self.ext_module_python.__file__) self.built_extension = True - self.aoti_module_path =self.ext_module_python.__file__ + self.aoti_module_path = self.ext_module_python.__file__ except Exception as e: # If compiling torch fails (e.g. low gcc version), we should fall back to the # version that takes integer pointers as args (but is untraceable to PyTorch JIT / export). self.built_extension_error = e except Exception as e: - self.built_extension_error = f"Error JIT-compiling OpenEquivariance Extension: {e}" + self.built_extension_error = ( + f"Error JIT-compiling OpenEquivariance Extension: {e}" + ) class PrecompiledExtension(Extension): - def __init__(): - try: + def __init__(self, backend): + super().__init__() + # CUDA extensions are built per CUDA major version, matching torch's + # libcudart / libnvrtc. + if backend == "cuda": ext_module_py_name = f"oeq_stable_cuda{torch.version.cuda.split('.')[0]}" + else: + ext_module_py_name = f"oeq_stable_{backend}" + + aoti_module_path = os.path.join( + os.path.dirname(__file__), f"lib{ext_module_py_name}_aoti.so" + ) + if not os.path.exists(aoti_module_path): + self.built_extension_error = ( + f"No precompiled OpenEquivariance extension {ext_module_py_name} " + f"for PyTorch {torch.__version__}. {SOURCE_BUILD_HINT}" + ) + return + try: self.ext_module_python = importlib.import_module( f"openequivariance._torch.extlib.{ext_module_py_name}" ) - - torch.ops.load_library(extension_module.__file__) + torch.ops.load_library(self.ext_module_python.__file__) self.built_extension = True - - dirname = os.path.dirname(extension_module.__file__) - self.aoti_module_path = os.path.join(dirname, f"lib{ext_module_py_name}_aoti.so") - + self.aoti_module_path = aoti_module_path except Exception as e: self.built_extension_error = ( - f"Error loading precompiled OpenEquivariance Extension: {e}" + f"Error loading precompiled OpenEquivariance extension " + f"{ext_module_py_name}: {e}\n{SOURCE_BUILD_HINT}" ) @@ -177,7 +200,11 @@ def _detect_backend(): return "cuda" if getattr(torch.version, "xpu", None): return "sycl" - assert False, "Invalid Backend" + raise RuntimeError( + "Only the CUDA, HIP and XPU (SYCL) backends are supported. " + "No supported accelerator was detected in this PyTorch build." + ) + BACKEND = _detect_backend() IS_HIP = BACKEND == "hip" @@ -188,24 +215,25 @@ def _detect_backend(): if BACKEND == "sycl": extension = JITExtension() else: - extension = PrecompiledExtension() -sys.modules["oeq_utilities"] = extension.ext_module_python -BUILT_EXTENSION = extension.BUILT_EXTENSION -BUILT_EXTENSION_ERROR = extension.BUILT_EXTENSION_ERROR + extension = PrecompiledExtension(BACKEND) +BUILT_EXTENSION = extension.built_extension +BUILT_EXTENSION_ERROR = extension.built_extension_error if BUILT_EXTENSION: + sys.modules["oeq_utilities"] = extension.ext_module_python from oeq_utilities import ( DeviceProp, GPUTimer, ) def torch_ext_so_path(): - return extension.aoti_module_path() + return extension.aoti_module_path else: - if extension.built_extension_error is not None: - warnings.warn(str(BUILT_EXTENSION_ERRROR), stacklevel=3) + if BUILT_EXTENSION_ERROR is not None: + warnings.warn(str(BUILT_EXTENSION_ERROR), stacklevel=3) + def _raise_import_error_helper(import_target: str): if not BUILT_EXTENSION: raise ImportError( From 924585f75a109e311aa9c7cf46b0267d61aa10e8 Mon Sep 17 00:00:00 2001 From: Vivek Bharadwaj Date: Mon, 9 Nov 2026 21:54:21 -0800 Subject: [PATCH 28/38] Eliminated a file. --- .github/cibw_before_all.sh | 38 ------------ .github/release.py | 61 ++++++++++++++++++- .../_torch/extlib/__init__.py | 4 -- .../extension/libtorch_tp_jit_stable.cpp | 19 +----- .../extension/nb_utilities.hpp | 35 +++++++++++ openequivariance/pyproject.toml | 6 +- .../src/jax_ffi_registration.cpp | 17 +----- 7 files changed, 101 insertions(+), 79 deletions(-) delete mode 100644 .github/cibw_before_all.sh create mode 100644 openequivariance/openequivariance/extension/nb_utilities.hpp diff --git a/.github/cibw_before_all.sh b/.github/cibw_before_all.sh deleted file mode 100644 index 08f01618..00000000 --- a/.github/cibw_before_all.sh +++ /dev/null @@ -1,38 +0,0 @@ -#!/bin/bash -# Runs inside the manylinux container before cibuildwheel builds -# openequivariance. Installs only the headers and link-time libraries -# for CUDA and HIP; at runtime, the user's PyTorch install provides them. -# -# Configured from [tool.cibuildwheel.linux.environment] in pyproject.toml: -# OEQ_CUDA_ROOTS space-separated CUDA toolkits to install, e.g. -# "/usr/local/cuda-12.6 /usr/local/cuda-13.0" -# and the pinned rocm-sdk-core in .github/cibw_rocm_requirements.txt. -set -euo pipefail - -ARCH=$(uname -m) - -case "$ARCH" in - x86_64) CUDA_REPO_ARCH=x86_64 ;; - aarch64) CUDA_REPO_ARCH=sbsa ;; -esac - -dnf config-manager --add-repo \ - "https://developer.download.nvidia.com/compute/cuda/repos/rhel8/${CUDA_REPO_ARCH}/cuda-rhel8.repo" - -for root in $OEQ_CUDA_ROOTS; do - version=${root##*/cuda-} - pkg_version=${version//./-} - dnf install -y \ - cuda-nvcc-${pkg_version} \ - cuda-cudart-devel-${pkg_version} \ - cuda-nvrtc-devel-${pkg_version} \ - cuda-driver-devel-${pkg_version} -done - -# rocm-sdk-core (~415 MB) has the HIP headers, libamdhip64 and libhiprtc, -# without the rest of the ROCm toolkit. AMD publishes no aarch64 builds. -if [ "$ARCH" = x86_64 ]; then - /opt/python/cp312-cp312/bin/python -m pip install \ - --no-deps --require-hashes --target /opt/rocm-sdk-core \ - -r "$(dirname "$0")/cibw_rocm_requirements.txt" -fi diff --git a/.github/release.py b/.github/release.py index e221c4b0..e3df285a 100644 --- a/.github/release.py +++ b/.github/release.py @@ -1,18 +1,24 @@ """ -Release pipeline steps for .github/workflows/release.yaml. +Release pipeline steps for .github/workflows/release.yaml. python3 .github/release.py resolve-symlinks python3 .github/release.py set-version [...] + python3 .github/release.py install-gpu-toolkits [...] """ import argparse +import platform import re import shutil import subprocess +import sys from pathlib import Path TAG_PATTERN = re.compile(r"v([0-9]+\.[0-9]+\.[0-9]+)") GIT_SYMLINK_MODE = "120000" +CUDA_ROOT_PATTERN = re.compile(r"/usr/local/cuda-([0-9]+)\.([0-9]+)") +CUDA_REPO_ARCH = {"x86_64": "x86_64", "aarch64": "sbsa"} +ROCM_REQUIREMENTS = Path(__file__).parent / "cibw_rocm_requirements.txt" def resolve_symlinks(args): @@ -63,6 +69,55 @@ def set_version(args): pyproject.write_text(tomlkit.dumps(doc)) +def install_gpu_toolkits(args): + """Installs the CUDA and HIP headers and link-time libraries inside the + manylinux container (cibuildwheel before-all); at runtime, the user's + PyTorch install provides them.""" + arch = platform.machine() + subprocess.run( + [ + "dnf", + "config-manager", + "--add-repo", + "https://developer.download.nvidia.com/compute/cuda/repos/rhel8/" + f"{CUDA_REPO_ARCH[arch]}/cuda-rhel8.repo", + ], + check=True, + ) + + for root in args.cuda_roots: + match = CUDA_ROOT_PATTERN.fullmatch(root) + assert match, f"expected a CUDA root like /usr/local/cuda-12.6, got {root!r}" + pkg_version = "-".join(match.groups()) + packages = [ + "cuda-nvcc", + "cuda-cudart-devel", + "cuda-nvrtc-devel", + "cuda-driver-devel", + ] + subprocess.run( + ["dnf", "install", "-y", *(f"{pkg}-{pkg_version}" for pkg in packages)], + check=True, + ) + + if arch == "x86_64": + subprocess.run( + [ + sys.executable, + "-m", + "pip", + "install", + "--no-deps", + "--require-hashes", + "--target", + "/opt/rocm-sdk-core", + "-r", + str(ROCM_REQUIREMENTS), + ], + check=True, + ) + + if __name__ == "__main__": parser = argparse.ArgumentParser(description=__doc__.split("\n\n")[0].strip()) steps = parser.add_subparsers(dest="step", required=True) @@ -74,5 +129,9 @@ def set_version(args): step.add_argument("pyprojects", nargs="+", type=Path) step.set_defaults(func=set_version) + step = steps.add_parser("install-gpu-toolkits", help=install_gpu_toolkits.__doc__) + step.add_argument("cuda_roots", nargs="+") + step.set_defaults(func=install_gpu_toolkits) + args = parser.parse_args() args.func(args) diff --git a/openequivariance/openequivariance/_torch/extlib/__init__.py b/openequivariance/openequivariance/_torch/extlib/__init__.py index 94555480..c76bc0cf 100644 --- a/openequivariance/openequivariance/_torch/extlib/__init__.py +++ b/openequivariance/openequivariance/_torch/extlib/__init__.py @@ -67,8 +67,6 @@ def wait_for_torch_build_lock(name, timeout=300): def __init__(self): super().__init__() - # The JIT extension is a pybind11 module, so it links libpython (when one - # exists) to stay loadable from AOTI C++ programs without an interpreter. python_lib_dir = sysconfig.get_config_var("LIBDIR") or "" python_lib_name = f"python{sys.version_info.major}.{sys.version_info.minor}" found_libpython = any( @@ -161,8 +159,6 @@ def __init__(self): class PrecompiledExtension(Extension): def __init__(self, backend): super().__init__() - # CUDA extensions are built per CUDA major version, matching torch's - # libcudart / libnvrtc. if backend == "cuda": ext_module_py_name = f"oeq_stable_cuda{torch.version.cuda.split('.')[0]}" else: diff --git a/openequivariance/openequivariance/extension/libtorch_tp_jit_stable.cpp b/openequivariance/openequivariance/extension/libtorch_tp_jit_stable.cpp index 0563c57c..56c57a9e 100644 --- a/openequivariance/openequivariance/extension/libtorch_tp_jit_stable.cpp +++ b/openequivariance/openequivariance/extension/libtorch_tp_jit_stable.cpp @@ -168,23 +168,8 @@ bool tensor_is_on_gpu(const Tensor &tensor) { #endif #ifdef INCLUDE_NB_EXTENSION - #include "nanobind/nanobind.h" - #include "nanobind/stl/string.h" - namespace nb = nanobind; + #include "nb_utilities.hpp" NB_MODULE(EXTENSION_NAME, m) { - nb::class_(m, "DeviceProp") - .def(nb::init()) - .def_ro("name", &DeviceProp::name) - .def_ro("warpsize", &DeviceProp::warpsize) - .def_ro("major", &DeviceProp::major) - .def_ro("minor", &DeviceProp::minor) - .def_ro("multiprocessorCount", &DeviceProp::multiprocessorCount) - .def_ro("maxSharedMemPerBlock", &DeviceProp::maxSharedMemPerBlock); - - nb::class_(m, "GPUTimer") - .def(nb::init<>()) - .def("start", &GPUTimer::start) - .def("stop_clock_get_elapsed", &GPUTimer::stop_clock_get_elapsed) - .def("clear_L2_cache", &GPUTimer::clear_L2_cache); + oeq::bind_utilities(m); } #endif diff --git a/openequivariance/openequivariance/extension/nb_utilities.hpp b/openequivariance/openequivariance/extension/nb_utilities.hpp new file mode 100644 index 00000000..e5c3d5cf --- /dev/null +++ b/openequivariance/openequivariance/extension/nb_utilities.hpp @@ -0,0 +1,35 @@ +#pragma once + +#include "nanobind/nanobind.h" +#include "nanobind/stl/string.h" + +namespace oeq { + +inline void bind_utilities(nanobind::module_ &m) { + namespace nb = nanobind; + + if (nb::handle existing = nb::type()) { + m.attr("DeviceProp") = existing; + } else { + nb::class_(m, "DeviceProp") + .def(nb::init()) + .def_ro("name", &DeviceProp::name) + .def_ro("warpsize", &DeviceProp::warpsize) + .def_ro("major", &DeviceProp::major) + .def_ro("minor", &DeviceProp::minor) + .def_ro("multiprocessorCount", &DeviceProp::multiprocessorCount) + .def_ro("maxSharedMemPerBlock", &DeviceProp::maxSharedMemPerBlock); + } + + if (nb::handle existing = nb::type()) { + m.attr("GPUTimer") = existing; + } else { + nb::class_(m, "GPUTimer") + .def(nb::init<>()) + .def("start", &GPUTimer::start) + .def("stop_clock_get_elapsed", &GPUTimer::stop_clock_get_elapsed) + .def("clear_L2_cache", &GPUTimer::clear_L2_cache); + } +} + +} // namespace oeq diff --git a/openequivariance/pyproject.toml b/openequivariance/pyproject.toml index 7f3d4c60..d695ac03 100644 --- a/openequivariance/pyproject.toml +++ b/openequivariance/pyproject.toml @@ -85,22 +85,18 @@ wheel.py-api = "cp311" AARCH64 = {env="AARCH64", default="0"} [tool.cibuildwheel] -# The abi3 wheel built for cp311 covers all later CPython versions. build = "cp311-*" skip = "*-musllinux*" manylinux-x86_64-image = "manylinux_2_28" manylinux-aarch64-image = "manylinux_2_28" -before-all = "bash {project}/.github/cibw_before_all.sh" +before-all = "/opt/python/cp312-cp312/bin/python {project}/.github/release.py install-gpu-toolkits $OEQ_CUDA_ROOTS" build-verbosity = 1 [tool.cibuildwheel.linux.environment] -# CUDA toolkits to build against: one extension per major version, chosen at -# runtime from torch.version.cuda. Use the oldest minor PyTorch ships per major. OEQ_CUDA_ROOTS = "/usr/local/cuda-12.6 /usr/local/cuda-13.0" ROCM_PATH = "/opt/rocm-sdk-core/_rocm_sdk_core" [tool.cibuildwheel.linux] -# GPU runtimes and libtorch come from the user's PyTorch install; never bundle them. repair-wheel-command = """auditwheel repair -w {dest_dir} {wheel} \ --exclude libtorch_cpu.so --exclude libc10.so \ --exclude libtorch_cuda.so --exclude libtorch_hip.so \ diff --git a/openequivariance_extjax/src/jax_ffi_registration.cpp b/openequivariance_extjax/src/jax_ffi_registration.cpp index 86307f95..83d8d12f 100644 --- a/openequivariance_extjax/src/jax_ffi_registration.cpp +++ b/openequivariance_extjax/src/jax_ffi_registration.cpp @@ -10,6 +10,8 @@ #include "backend/backend_hip.hpp" #endif +#include "nb_utilities.hpp" + namespace nb = nanobind; namespace { @@ -47,18 +49,5 @@ NB_MODULE(openequivariance_extjax, m) { }); m.def("is_hip", &IsHip); - nb::class_(m, "DeviceProp") - .def(nb::init()) - .def_ro("name", &DeviceProp::name) - .def_ro("warpsize", &DeviceProp::warpsize) - .def_ro("major", &DeviceProp::major) - .def_ro("minor", &DeviceProp::minor) - .def_ro("multiprocessorCount", &DeviceProp::multiprocessorCount) - .def_ro("maxSharedMemPerBlock", &DeviceProp::maxSharedMemPerBlock); - - nb::class_(m, "GPUTimer") - .def(nb::init<>()) - .def("start", &GPUTimer::start) - .def("stop_clock_get_elapsed", &GPUTimer::stop_clock_get_elapsed) - .def("clear_L2_cache", &GPUTimer::clear_L2_cache); + oeq::bind_utilities(m); } From 9e497449e987534e89608c8a840ae81dc1f53709 Mon Sep 17 00:00:00 2001 From: Vivek Bharadwaj Date: Mon, 9 Nov 2026 22:04:15 -0800 Subject: [PATCH 29/38] More modifications to release script. --- .github/cibw_rocm_requirements.txt | 2 -- .github/release.py | 44 +++++++++++++++++------------- 2 files changed, 25 insertions(+), 21 deletions(-) delete mode 100644 .github/cibw_rocm_requirements.txt diff --git a/.github/cibw_rocm_requirements.txt b/.github/cibw_rocm_requirements.txt deleted file mode 100644 index 5448bf02..00000000 --- a/.github/cibw_rocm_requirements.txt +++ /dev/null @@ -1,2 +0,0 @@ ---index-url https://download.pytorch.org/whl/rocm10.0 -rocm-sdk-core==10.0.0 --hash=sha256:3bf4a72d11aa2a4ee1e90572c73630f937d38c1b7af4dda45b483b54655138a7 diff --git a/.github/release.py b/.github/release.py index e3df285a..a9d79dfb 100644 --- a/.github/release.py +++ b/.github/release.py @@ -12,13 +12,17 @@ import shutil import subprocess import sys +import tempfile from pathlib import Path TAG_PATTERN = re.compile(r"v([0-9]+\.[0-9]+\.[0-9]+)") GIT_SYMLINK_MODE = "120000" CUDA_ROOT_PATTERN = re.compile(r"/usr/local/cuda-([0-9]+)\.([0-9]+)") CUDA_REPO_ARCH = {"x86_64": "x86_64", "aarch64": "sbsa"} -ROCM_REQUIREMENTS = Path(__file__).parent / "cibw_rocm_requirements.txt" +ROCM_REQUIREMENTS = """\ +--index-url https://download.pytorch.org/whl/rocm10.0 +rocm-sdk-core==10.0.0 --hash=sha256:3bf4a72d11aa2a4ee1e90572c73630f937d38c1b7af4dda45b483b54655138a7 +""" def resolve_symlinks(args): @@ -70,9 +74,8 @@ def set_version(args): def install_gpu_toolkits(args): - """Installs the CUDA and HIP headers and link-time libraries inside the - manylinux container (cibuildwheel before-all); at runtime, the user's - PyTorch install provides them.""" + """Installs the CUDA and HIP headers and link-time libraries inside a + manylinux container.""" arch = platform.machine() subprocess.run( [ @@ -101,21 +104,24 @@ def install_gpu_toolkits(args): ) if arch == "x86_64": - subprocess.run( - [ - sys.executable, - "-m", - "pip", - "install", - "--no-deps", - "--require-hashes", - "--target", - "/opt/rocm-sdk-core", - "-r", - str(ROCM_REQUIREMENTS), - ], - check=True, - ) + with tempfile.TemporaryDirectory() as tmp: + requirements = Path(tmp) / "rocm_requirements.txt" + requirements.write_text(ROCM_REQUIREMENTS) + subprocess.run( + [ + sys.executable, + "-m", + "pip", + "install", + "--no-deps", + "--require-hashes", + "--target", + "/opt/rocm-sdk-core", + "-r", + str(requirements), + ], + check=True, + ) if __name__ == "__main__": From d6f6ee170f8624a7e1be9a6f3af5dea00b906877 Mon Sep 17 00:00:00 2001 From: Vivek Bharadwaj Date: Mon, 9 Nov 2026 22:12:22 -0800 Subject: [PATCH 30/38] Changes to get rid of pybind11 dependency. --- .github/workflows/verify_extension_build.yml | 2 +- docs/installation.rst | 6 +++- .../_torch/extlib/__init__.py | 34 +++++++++---------- openequivariance/pyproject.toml | 8 +++-- 4 files changed, 28 insertions(+), 22 deletions(-) diff --git a/.github/workflows/verify_extension_build.yml b/.github/workflows/verify_extension_build.yml index 44afeeee..35e3275d 100644 --- a/.github/workflows/verify_extension_build.yml +++ b/.github/workflows/verify_extension_build.yml @@ -66,7 +66,7 @@ jobs: run: | python -m pip install --upgrade pip pip install -r .github/workflows/requirements_sycl_ci.txt - pip install -e "./openequivariance" + pip install -e "./openequivariance[sycl]" - name: Test SYCL extension build via import run: | diff --git a/docs/installation.rst b/docs/installation.rst index 144583bf..58f7ff0c 100644 --- a/docs/installation.rst +++ b/docs/installation.rst @@ -51,7 +51,11 @@ You need the following to install OpenEquivariance: On Intel GPUs, the extension is still compiled on first import, which can take a couple of minutes; subsequent imports use the - cached build. + cached build. Install the extra build dependencies with: + + .. code-block:: bash + + pip install openequivariance[sycl] .. tab:: JAX NVIDIA GPUs diff --git a/openequivariance/openequivariance/_torch/extlib/__init__.py b/openequivariance/openequivariance/_torch/extlib/__init__.py index c76bc0cf..7feb0b78 100644 --- a/openequivariance/openequivariance/_torch/extlib/__init__.py +++ b/openequivariance/openequivariance/_torch/extlib/__init__.py @@ -6,6 +6,7 @@ import sysconfig import contextlib import importlib +import importlib.util from pathlib import Path from packaging.version import Version @@ -74,6 +75,19 @@ def __init__(self): for ext in (".so", ".a") ) + missing = [ + name + for name in ("setuptools", "ninja") + if importlib.util.find_spec(name) is None + ] + if missing: + self.built_extension_error = ( + f"The SYCL backend compiles its extension at first import, which " + f"needs {', '.join(missing)}. Install them with:\n" + " pip install openequivariance[sycl]" + ) + return + try: from torch.utils.cpp_extension import library_paths, include_paths @@ -81,18 +95,6 @@ def __init__(self): torch_sources = ["libtorch_tp_jit.cpp", "json11/json11.cpp"] include_dirs, extra_link_args = (["backend"], ["-Wl,--no-as-needed"]) - extra_include_dirs = [] - - try: - import pybind11 - - extra_include_dirs.append(pybind11.get_include()) - except Exception as e: - self.built_extension_error = ( - "Could not locate pybind11 include path required for JIT " - f"OpenEquivariance extension compilation: {e}" - ) - return if found_libpython: extra_link_args.pop() @@ -123,11 +125,9 @@ def __init__(self): extra_link_args.append("-L" + lib_dir) torch_sources = [oeq_root + "/extension/" + src for src in torch_sources] - include_dirs = ( - [oeq_root + "/extension/" + d for d in include_dirs] - + extra_include_dirs - + include_paths("xpu") - ) + include_dirs = [ + oeq_root + "/extension/" + d for d in include_dirs + ] + include_paths("xpu") with warnings.catch_warnings(): warnings.simplefilter("ignore") diff --git a/openequivariance/pyproject.toml b/openequivariance/pyproject.toml index d695ac03..c29865ab 100644 --- a/openequivariance/pyproject.toml +++ b/openequivariance/pyproject.toml @@ -14,11 +14,8 @@ authors = [ description = "A fast GPU JIT kernel generator for the Clebsch-Gordon Tensor Product" requires-python = ">=3.11" dependencies = [ - "setuptools", - "ninja", "jinja2", "numpy", - "pybind11", "nanobind-backend>=1.0" ] readme = "README.md" @@ -60,6 +57,11 @@ dev = [ "cmake" ] +sycl = [ + "setuptools", + "ninja" +] + jax = [ "scikit-build-core", "setuptools-scm", From c9f36041f6c3ececdaf86cde0bf8acce7d96460a Mon Sep 17 00:00:00 2001 From: Vivek Bharadwaj Date: Tue, 10 Nov 2026 19:50:49 -0800 Subject: [PATCH 31/38] Removed some comments. --- openequivariance/CMakeLists.txt | 14 ++------------ .../openequivariance/_torch/extlib/__init__.py | 12 ++++++------ openequivariance/pyproject.toml | 2 ++ 3 files changed, 10 insertions(+), 18 deletions(-) diff --git a/openequivariance/CMakeLists.txt b/openequivariance/CMakeLists.txt index 2c4d2f34..4a193984 100644 --- a/openequivariance/CMakeLists.txt +++ b/openequivariance/CMakeLists.txt @@ -52,11 +52,9 @@ set(OEQ_SOURCES ${EXT_JSON_DIR}/json11.cpp ) -set(OEQ_INSTALL_DIR "${CMAKE_CURRENT_SOURCE_DIR}/openequivariance/_torch/extlib") +set(OEQ_INSTALL_DIR "openequivariance/_torch/extlib") function(add_stable_extension target_name backend_define link_libraries) - # Create nanobind extension in split mode: stable-ABI frontend, with - # libnanobind provided at runtime by the nanobind-backend package. nanobind_add_module(${target_name} BACKEND_MODULE nanobind_backend ${OEQ_SOURCES}) set_target_properties(${target_name} PROPERTIES @@ -119,8 +117,6 @@ endfunction() find_package(hip QUIET) -# The rocm-sdk-core pip package (used for wheel builds) ships HIP headers and -# versioned libraries but no CMake configs, so locate them under ROCM_PATH. if(NOT hip_FOUND AND DEFINED ENV{ROCM_PATH}) find_path(HIP_INCLUDE_DIR hip/hip_runtime.h PATHS $ENV{ROCM_PATH}/include NO_DEFAULT_PATH) find_library(AMDHIP64_LIB NAMES amdhip64 libamdhip64.so.7 PATHS $ENV{ROCM_PATH}/lib NO_DEFAULT_PATH) @@ -139,12 +135,6 @@ if(NOT hip_FOUND AND DEFINED ENV{ROCM_PATH}) endif() endif() -# CUDA: one extension per CUDA major version, named oeq_stable_cuda, so -# that it links the same libcudart/libnvrtc major as the user's PyTorch. -# -# OEQ_CUDA_ROOTS (CMake or environment variable, ;- or space-separated) lists -# the toolkit roots to build against, e.g. "/usr/local/cuda-12.6 /usr/local/cuda-13.0". -# If unset, a single extension is built against the toolkit find_package finds. set(OEQ_CUDA_ROOTS "$ENV{OEQ_CUDA_ROOTS}" CACHE STRING "CUDA toolkit roots to build against") string(REPLACE " " ";" OEQ_CUDA_ROOTS "${OEQ_CUDA_ROOTS}") @@ -175,7 +165,7 @@ if(OEQ_CUDA_FOUND AND NOT OEQ_CUDA_ROOTS) set(CUDA_LINK_LIBS CUDA::cudart CUDA::cuda_driver - CUDA::nvrtc + ${CUDA_nvrtc_LIBRARY} cuda_stub_lib ) add_stable_extension(oeq_stable_cuda${CUDAToolkit_VERSION_MAJOR} CUDA_BACKEND "${CUDA_LINK_LIBS}") diff --git a/openequivariance/openequivariance/_torch/extlib/__init__.py b/openequivariance/openequivariance/_torch/extlib/__init__.py index 7feb0b78..54f31155 100644 --- a/openequivariance/openequivariance/_torch/extlib/__init__.py +++ b/openequivariance/openequivariance/_torch/extlib/__init__.py @@ -164,10 +164,12 @@ def __init__(self, backend): else: ext_module_py_name = f"oeq_stable_{backend}" - aoti_module_path = os.path.join( - os.path.dirname(__file__), f"lib{ext_module_py_name}_aoti.so" + full_name = f"openequivariance._torch.extlib.{ext_module_py_name}" + spec = importlib.util.find_spec(full_name) + aoti_module_path = spec and os.path.join( + os.path.dirname(spec.origin), f"lib{ext_module_py_name}_aoti.so" ) - if not os.path.exists(aoti_module_path): + if spec is None or not os.path.exists(aoti_module_path): self.built_extension_error = ( f"No precompiled OpenEquivariance extension {ext_module_py_name} " f"for PyTorch {torch.__version__}. {SOURCE_BUILD_HINT}" @@ -175,9 +177,7 @@ def __init__(self, backend): return try: - self.ext_module_python = importlib.import_module( - f"openequivariance._torch.extlib.{ext_module_py_name}" - ) + self.ext_module_python = importlib.import_module(full_name) torch.ops.load_library(self.ext_module_python.__file__) self.built_extension = True self.aoti_module_path = aoti_module_path diff --git a/openequivariance/pyproject.toml b/openequivariance/pyproject.toml index c29865ab..d8663f1a 100644 --- a/openequivariance/pyproject.toml +++ b/openequivariance/pyproject.toml @@ -82,6 +82,8 @@ lint.ignore = ["E741"] [tool.scikit-build] # Stable-ABI (abi3) wheel for CPython 3.11+, built with nanobind split mode. wheel.py-api = "cp311" +build-dir = "build/{wheel_tag}" +editable.rebuild = true [tool.scikit-build.cmake.define] AARCH64 = {env="AARCH64", default="0"} From 17f9c9e26ed205922d2da73c6605fae63e53b807 Mon Sep 17 00:00:00 2001 From: Vivek Bharadwaj Date: Tue, 10 Nov 2026 19:53:43 -0800 Subject: [PATCH 32/38] Edits. --- openequivariance/pyproject.toml | 1 - 1 file changed, 1 deletion(-) diff --git a/openequivariance/pyproject.toml b/openequivariance/pyproject.toml index d8663f1a..7e4e9212 100644 --- a/openequivariance/pyproject.toml +++ b/openequivariance/pyproject.toml @@ -80,7 +80,6 @@ addopts = [ lint.ignore = ["E741"] [tool.scikit-build] -# Stable-ABI (abi3) wheel for CPython 3.11+, built with nanobind split mode. wheel.py-api = "cp311" build-dir = "build/{wheel_tag}" editable.rebuild = true From 9bf26019640d7619daaabcd110a8e16104774d1f Mon Sep 17 00:00:00 2001 From: Vivek Bharadwaj Date: Tue, 10 Nov 2026 20:42:22 -0800 Subject: [PATCH 33/38] Adding ROCM information to pyproject.toml. --- .github/release.py | 25 ++++++++++++++++++------- openequivariance/CMakeLists.txt | 11 +++++++++-- openequivariance/pyproject.toml | 13 ++++++++++++- 3 files changed, 39 insertions(+), 10 deletions(-) diff --git a/.github/release.py b/.github/release.py index a9d79dfb..3e482dd5 100644 --- a/.github/release.py +++ b/.github/release.py @@ -19,10 +19,11 @@ GIT_SYMLINK_MODE = "120000" CUDA_ROOT_PATTERN = re.compile(r"/usr/local/cuda-([0-9]+)\.([0-9]+)") CUDA_REPO_ARCH = {"x86_64": "x86_64", "aarch64": "sbsa"} -ROCM_REQUIREMENTS = """\ ---index-url https://download.pytorch.org/whl/rocm10.0 -rocm-sdk-core==10.0.0 --hash=sha256:3bf4a72d11aa2a4ee1e90572c73630f937d38c1b7af4dda45b483b54655138a7 -""" +ROCM_INDEX_PATTERN = re.compile( + r"https://download\.pytorch\.org/whl/rocm[0-9]+\.[0-9]+" +) +ROCM_VERSION_PATTERN = re.compile(r"[0-9]+\.[0-9]+\.[0-9]+") +SHA256_PATTERN = re.compile(r"[0-9a-f]{64}") def resolve_symlinks(args): @@ -74,8 +75,8 @@ def set_version(args): def install_gpu_toolkits(args): - """Installs the CUDA and HIP headers and link-time libraries inside a - manylinux container.""" + """Installs the CUDA and HIP headers and link-time libraries inside a + manylinux container.""" arch = platform.machine() subprocess.run( [ @@ -104,9 +105,16 @@ def install_gpu_toolkits(args): ) if arch == "x86_64": + assert ROCM_INDEX_PATTERN.fullmatch(args.rocm_index_url) + assert ROCM_VERSION_PATTERN.fullmatch(args.rocm_sdk_version) + assert SHA256_PATTERN.fullmatch(args.rocm_sdk_sha256) with tempfile.TemporaryDirectory() as tmp: requirements = Path(tmp) / "rocm_requirements.txt" - requirements.write_text(ROCM_REQUIREMENTS) + requirements.write_text( + f"--index-url {args.rocm_index_url}\n" + f"rocm-sdk-core=={args.rocm_sdk_version} " + f"--hash=sha256:{args.rocm_sdk_sha256}\n" + ) subprocess.run( [ sys.executable, @@ -136,6 +144,9 @@ def install_gpu_toolkits(args): step.set_defaults(func=set_version) step = steps.add_parser("install-gpu-toolkits", help=install_gpu_toolkits.__doc__) + step.add_argument("--rocm-index-url", required=True) + step.add_argument("--rocm-sdk-version", required=True) + step.add_argument("--rocm-sdk-sha256", required=True) step.add_argument("cuda_roots", nargs="+") step.set_defaults(func=install_gpu_toolkits) diff --git a/openequivariance/CMakeLists.txt b/openequivariance/CMakeLists.txt index 4a193984..2cdf63e4 100644 --- a/openequivariance/CMakeLists.txt +++ b/openequivariance/CMakeLists.txt @@ -8,15 +8,22 @@ include(FetchContent) if (AARCH64 OR CMAKE_SYSTEM_PROCESSOR MATCHES "(arm)|(ARM)|(aarch64)") message(STATUS "Targeting ARM architecture (AARCH64 override or processor match)") - set(LIBTORCH_URL "https://github.com/PASSIONLab/libtorch_aarch64/raw/refs/heads/main/pkg-libtorch.tar.zst") + set(LIBTORCH_ARCH AARCH64) else() message(STATUS "Did not match arm/aarch64, performing x86 install") - set(LIBTORCH_URL "https://download.pytorch.org/libtorch/cpu/libtorch-shared-with-deps-2.10.0%2Bcpu.zip") + set(LIBTORCH_ARCH X86_64) +endif() + +set(LIBTORCH_URL "${LIBTORCH_URL_${LIBTORCH_ARCH}}") +set(LIBTORCH_SHA256 "${LIBTORCH_SHA256_${LIBTORCH_ARCH}}") +if(NOT LIBTORCH_URL OR NOT LIBTORCH_SHA256) + message(FATAL_ERROR "LIBTORCH_URL / LIBTORCH_SHA256 not set") endif() FetchContent_Declare( libtorch URL "${LIBTORCH_URL}" + URL_HASH SHA256=${LIBTORCH_SHA256} ) message(STATUS "Downloading LibTorch...") diff --git a/openequivariance/pyproject.toml b/openequivariance/pyproject.toml index 7e4e9212..65d425d1 100644 --- a/openequivariance/pyproject.toml +++ b/openequivariance/pyproject.toml @@ -86,18 +86,29 @@ editable.rebuild = true [tool.scikit-build.cmake.define] AARCH64 = {env="AARCH64", default="0"} +LIBTORCH_URL_X86_64 = "https://download.pytorch.org/libtorch/cpu/libtorch-shared-with-deps-2.10.0%2Bcpu.zip" +LIBTORCH_SHA256_X86_64 = "c5bf8efda9224a2d971b19d1ef6cf3ba6fee8ab53e69c49427db003d1d300496" +LIBTORCH_URL_AARCH64 = "https://github.com/PASSIONLab/libtorch_aarch64/raw/742ffd07f47a31a94f6b6aae2cbe129e6b667d94/pkg-libtorch.tar.zst" +LIBTORCH_SHA256_AARCH64 = "9fdebdb56766317c71eb4c9f1b750baef4ebc4e6ba85b427f701fdc20844f909" [tool.cibuildwheel] build = "cp311-*" skip = "*-musllinux*" manylinux-x86_64-image = "manylinux_2_28" manylinux-aarch64-image = "manylinux_2_28" -before-all = "/opt/python/cp312-cp312/bin/python {project}/.github/release.py install-gpu-toolkits $OEQ_CUDA_ROOTS" +before-all = """/opt/python/cp312-cp312/bin/python {project}/.github/release.py install-gpu-toolkits \ + --rocm-index-url $OEQ_ROCM_INDEX_URL \ + --rocm-sdk-version $OEQ_ROCM_SDK_VERSION \ + --rocm-sdk-sha256 $OEQ_ROCM_SDK_SHA256 \ + $OEQ_CUDA_ROOTS""" build-verbosity = 1 [tool.cibuildwheel.linux.environment] OEQ_CUDA_ROOTS = "/usr/local/cuda-12.6 /usr/local/cuda-13.0" ROCM_PATH = "/opt/rocm-sdk-core/_rocm_sdk_core" +OEQ_ROCM_INDEX_URL = "https://download.pytorch.org/whl/rocm10.0" +OEQ_ROCM_SDK_VERSION = "10.0.0" +OEQ_ROCM_SDK_SHA256 = "3bf4a72d11aa2a4ee1e90572c73630f937d38c1b7af4dda45b483b54655138a7" [tool.cibuildwheel.linux] repair-wheel-command = """auditwheel repair -w {dest_dir} {wheel} \ From 8776ddf77ff53ff71a9b9eed41b8c39a3dea189a Mon Sep 17 00:00:00 2001 From: Vivek Bharadwaj Date: Tue, 10 Nov 2026 21:16:08 -0800 Subject: [PATCH 34/38] Updated documentation --- README.md | 7 +------ docs/installation.rst | 14 +++++++------- 2 files changed, 8 insertions(+), 13 deletions(-) diff --git a/README.md b/README.md index 34de8a61..24fa7a0f 100644 --- a/README.md +++ b/README.md @@ -13,17 +13,12 @@ that [e3nn](https://e3nn.org/) supports commonly found in graph neural networks (e.g. [Nequip](https://github.com/mir-group/nequip) or [MACE](https://github.com/ACEsuit/mace)). To get -started with PyTorch (2.10 or newer), install our package via +started, make sure that you have PyTorch installed and run: ```bash pip install openequivariance ``` -Linux wheels include a precompiled extension for PyTorch built -with CUDA 12, CUDA 13, or ROCm 7. For other setups, build the -extension from source (requires GCC 9+ and your CUDA / HIP toolkit) with -`pip install --no-binary openequivariance openequivariance`. - We provide up to an order of magnitude acceleration over e3nn perform on par with the latest version of [NVIDIA cuEquivariance](https://github.com/NVIDIA/cuEquivariance), which has a closed-source kernel package. diff --git a/docs/installation.rst b/docs/installation.rst index 58f7ff0c..5979348e 100644 --- a/docs/installation.rst +++ b/docs/installation.rst @@ -10,7 +10,7 @@ You need the following to install OpenEquivariance: - A Linux system equipped with an NVIDIA / AMD / Intel graphics card. - Either PyTorch >= 2.10, or JAX>0.5.0 with CUDA, RocM, or XPU support. -- For source builds only: GCC 9+ and the CUDA / HIP toolkit, or the +- To build from source: GCC 9+ and the CUDA / HIP toolkit, or the oneAPI DPC++ compiler (``icpx``) for Intel GPUs. The command ``c++ --version`` should return >= 9.0; see below for details on setting an alternate compiler. @@ -32,17 +32,17 @@ You need the following to install OpenEquivariance: pip install openequivariance python -c "import openequivariance" - On Linux, the wheel includes a precompiled C++ extension for PyTorch - built with CUDA 12, CUDA 13, or ROCm 7. If you see a warning that no - precompiled extension matches your PyTorch build, build it from - source instead (this requires GCC 9+ and the CUDA / HIP toolkit): + OpenEquivariance bundles a precompiled `.so` files for two recent + CUDA versions and a HIP version. If you're building for an architecture + we don't support or need a different CUDA verison, you'll need to build + the package from source: .. code-block:: bash pip install --no-binary openequivariance openequivariance - To build the extension with an alternate compiler, set the - ``CC`` and ``CXX`` environment variables before running ``pip``: + Set the environment variables ``CC`` and ``CXX`` before the pip install + to build with a different compiler: .. code-block:: bash From 40f6f4f17eba7678f0db0c6a41dff9a538427c3f Mon Sep 17 00:00:00 2001 From: Vivek Bharadwaj Date: Tue, 10 Nov 2026 21:17:47 -0800 Subject: [PATCH 35/38] Further documentation updates. --- docs/installation.rst | 39 --------------------------------------- 1 file changed, 39 deletions(-) diff --git a/docs/installation.rst b/docs/installation.rst index 5979348e..a0436956 100644 --- a/docs/installation.rst +++ b/docs/installation.rst @@ -114,45 +114,6 @@ Here are some tested environment configuration files. If you use OpenEquivarianc on a major cluster, send us a pull request to add your configuration! -.. tab:: NERSC Perlmutter (NVIDIA A100) - - .. code-block:: bash - :caption: env.sh (last updated June 2025) - - module load gcc - module load conda - - # Deactivate any base environments - for i in $(seq ${CONDA_SHLVL}); do - conda deactivate - done - - conda activate - - -.. tab:: OLCF Frontier (AMD MI250x) - - You need to install a HIP-enabled verison of PyTorch to use our package. - Follow the steps `here `_. - - - .. code-block:: bash - :caption: env.sh (last updated June 2025) - - module load PrgEnv-gnu/8.6.0 - module load miniforge3/23.11.0-0 - module load rocm/6.4.0 - module load craype-accel-amd-gfx90a - - for i in $(seq ${CONDA_SHLVL}); do - conda deactivate - done - - conda activate - export CC=cc - export CXX=CC - - .. tab:: ALCF Aurora (Intel Data Center GPU Max) Aurora provides a PyTorch build with XPU support in its ``frameworks`` From 52ab3bf185a92d5e7bd2a66755c60042b6354f1d Mon Sep 17 00:00:00 2001 From: Vivek Bharadwaj Date: Tue, 10 Nov 2026 21:25:52 -0800 Subject: [PATCH 36/38] More fixes. --- docs/installation.rst | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/docs/installation.rst b/docs/installation.rst index a0436956..8405f77f 100644 --- a/docs/installation.rst +++ b/docs/installation.rst @@ -32,9 +32,9 @@ You need the following to install OpenEquivariance: pip install openequivariance python -c "import openequivariance" - OpenEquivariance bundles a precompiled `.so` files for two recent + OpenEquivariance bundles a precompiled library files for two recent CUDA versions and a HIP version. If you're building for an architecture - we don't support or need a different CUDA verison, you'll need to build + we don't support or need a different CUDA version, you'll need to build the package from source: .. code-block:: bash @@ -91,11 +91,11 @@ You need the following to install OpenEquivariance: .. code-block:: bash - pip install "git+https://github.com/PASSIONLab/OpenEquivariance#subdirectory=openequivariance[jax]" - pip install "git+https://github.com/PASSIONLab/OpenEquivariance#subdirectory=openequivariance_extjax --no-build-isolation" + pip install "openequivariance[jax] @ git+https://github.com/PASSIONLab/OpenEquivariance#subdirectory=openequivariance" + pip install "git+https://github.com/PASSIONLab/OpenEquivariance#subdirectory=openequivariance_extjax" --no-build-isolation # Use the command below for JAX+AMD - # JAX_HIP=1 pip install "git+https://github.com/PASSIONLab/OpenEquivariance#subdirectory=openequivariance_extjax --no-build-isolation" + # JAX_HIP=1 pip install "git+https://github.com/PASSIONLab/OpenEquivariance#subdirectory=openequivariance_extjax" --no-build-isolation If you're using JAX, set the environment variable From c89a9adb232e73f7ce0a4d060798bb5a4d138304 Mon Sep 17 00:00:00 2001 From: Vivek Bharadwaj Date: Tue, 10 Nov 2026 22:06:54 -0800 Subject: [PATCH 37/38] Minor changes. --- .github/release.py | 48 ++++++++++++-------- .github/workflows/verify_extension_build.yml | 21 +++++++++ openequivariance/CMakeLists.txt | 2 +- openequivariance/pyproject.toml | 17 +++---- 4 files changed, 60 insertions(+), 28 deletions(-) diff --git a/.github/release.py b/.github/release.py index 3e482dd5..d7a645b8 100644 --- a/.github/release.py +++ b/.github/release.py @@ -3,7 +3,8 @@ python3 .github/release.py resolve-symlinks python3 .github/release.py set-version [...] - python3 .github/release.py install-gpu-toolkits [...] + python3 .github/release.py install-cuda-toolkits [...] + python3 .github/release.py install-rocm-sdk """ import argparse @@ -13,6 +14,7 @@ import subprocess import sys import tempfile +import tomllib from pathlib import Path TAG_PATTERN = re.compile(r"v([0-9]+\.[0-9]+\.[0-9]+)") @@ -74,17 +76,16 @@ def set_version(args): pyproject.write_text(tomlkit.dumps(doc)) -def install_gpu_toolkits(args): - """Installs the CUDA and HIP headers and link-time libraries inside a - manylinux container.""" - arch = platform.machine() +def install_cuda_toolkits(args): + """Installs the CUDA headers and link-time libraries for each toolkit root + inside a manylinux container.""" subprocess.run( [ "dnf", "config-manager", "--add-repo", "https://developer.download.nvidia.com/compute/cuda/repos/rhel8/" - f"{CUDA_REPO_ARCH[arch]}/cuda-rhel8.repo", + f"{CUDA_REPO_ARCH[platform.machine()]}/cuda-rhel8.repo", ], check=True, ) @@ -104,16 +105,23 @@ def install_gpu_toolkits(args): check=True, ) - if arch == "x86_64": - assert ROCM_INDEX_PATTERN.fullmatch(args.rocm_index_url) - assert ROCM_VERSION_PATTERN.fullmatch(args.rocm_sdk_version) - assert SHA256_PATTERN.fullmatch(args.rocm_sdk_sha256) + +def install_rocm_sdk(args): + """Installs rocm-sdk-core (HIP headers, libamdhip64 and libhiprtc), pinned + in [tool.openequivariance.rocm-sdk-core] of pyproject.toml, into a target + directory; skipped on non-x86_64 machines.""" + if platform.machine() == "x86_64": + with open(args.pyproject, "rb") as f: + pin = tomllib.load(f)["tool"]["openequivariance"]["rocm-sdk-core"] + assert ROCM_INDEX_PATTERN.fullmatch(pin["index-url"]) + assert ROCM_VERSION_PATTERN.fullmatch(pin["version"]) + assert SHA256_PATTERN.fullmatch(pin["sha256"]) with tempfile.TemporaryDirectory() as tmp: requirements = Path(tmp) / "rocm_requirements.txt" requirements.write_text( - f"--index-url {args.rocm_index_url}\n" - f"rocm-sdk-core=={args.rocm_sdk_version} " - f"--hash=sha256:{args.rocm_sdk_sha256}\n" + f"--index-url {pin['index-url']}\n" + f"rocm-sdk-core=={pin['version']} " + f"--hash=sha256:{pin['sha256']}\n" ) subprocess.run( [ @@ -124,7 +132,7 @@ def install_gpu_toolkits(args): "--no-deps", "--require-hashes", "--target", - "/opt/rocm-sdk-core", + str(args.target), "-r", str(requirements), ], @@ -143,12 +151,14 @@ def install_gpu_toolkits(args): step.add_argument("pyprojects", nargs="+", type=Path) step.set_defaults(func=set_version) - step = steps.add_parser("install-gpu-toolkits", help=install_gpu_toolkits.__doc__) - step.add_argument("--rocm-index-url", required=True) - step.add_argument("--rocm-sdk-version", required=True) - step.add_argument("--rocm-sdk-sha256", required=True) + step = steps.add_parser("install-cuda-toolkits", help=install_cuda_toolkits.__doc__) step.add_argument("cuda_roots", nargs="+") - step.set_defaults(func=install_gpu_toolkits) + step.set_defaults(func=install_cuda_toolkits) + + step = steps.add_parser("install-rocm-sdk", help=install_rocm_sdk.__doc__) + step.add_argument("pyproject", type=Path) + step.add_argument("target", type=Path) + step.set_defaults(func=install_rocm_sdk) args = parser.parse_args() args.func(args) diff --git a/.github/workflows/verify_extension_build.yml b/.github/workflows/verify_extension_build.yml index 35e3275d..7843cb87 100644 --- a/.github/workflows/verify_extension_build.yml +++ b/.github/workflows/verify_extension_build.yml @@ -39,6 +39,27 @@ jobs: run: | XLA_DIRECT_DOWNLOAD=1 pip install -e "./openequivariance_extjax" --no-build-isolation + verify_hip_extension: + if: ${{ github.event.label.name == 'ci-ready' || github.event_name != 'pull_request' }} + runs-on: ubuntu-latest + + steps: + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + - name: Set up Python + uses: actions/setup-python@6d4fd22b20465ccc0f7354476588e8df7c9ca0e0 # v7.1.0 + with: + python-version: "3.12" + + - name: Install pinned rocm-sdk-core + run: | + python .github/release.py install-rocm-sdk openequivariance/pyproject.toml "$RUNNER_TEMP/rocm-sdk-core" + + - name: Build HIP extension + env: + ROCM_PATH: ${{ runner.temp }}/rocm-sdk-core/_rocm_sdk_core + run: | + pip wheel --no-deps -w dist ./openequivariance + verify_sycl_extension: if: ${{ github.event.label.name == 'ci-ready' || github.event_name != 'pull_request' }} runs-on: ubuntu-latest diff --git a/openequivariance/CMakeLists.txt b/openequivariance/CMakeLists.txt index aa4231c5..1ceeadbd 100644 --- a/openequivariance/CMakeLists.txt +++ b/openequivariance/CMakeLists.txt @@ -260,5 +260,5 @@ if(OEQ_SYCL_FOUND) endif() if(NOT OEQ_CUDA_FOUND AND NOT hip_FOUND AND NOT OEQ_SYCL_FOUND) - message(WARNING "None of CUDAToolkit, HIP or SYCL was found. The stable extension will not be built.") + message(FATAL_ERROR "None of CUDAToolkit, HIP or SYCL was found") endif() diff --git a/openequivariance/pyproject.toml b/openequivariance/pyproject.toml index 65d425d1..4fee2c0a 100644 --- a/openequivariance/pyproject.toml +++ b/openequivariance/pyproject.toml @@ -96,19 +96,15 @@ build = "cp311-*" skip = "*-musllinux*" manylinux-x86_64-image = "manylinux_2_28" manylinux-aarch64-image = "manylinux_2_28" -before-all = """/opt/python/cp312-cp312/bin/python {project}/.github/release.py install-gpu-toolkits \ - --rocm-index-url $OEQ_ROCM_INDEX_URL \ - --rocm-sdk-version $OEQ_ROCM_SDK_VERSION \ - --rocm-sdk-sha256 $OEQ_ROCM_SDK_SHA256 \ - $OEQ_CUDA_ROOTS""" +before-all = [ + "/opt/python/cp312-cp312/bin/python {project}/.github/release.py install-cuda-toolkits $OEQ_CUDA_ROOTS", + "/opt/python/cp312-cp312/bin/python {project}/.github/release.py install-rocm-sdk {package}/pyproject.toml /opt/rocm-sdk-core", +] build-verbosity = 1 [tool.cibuildwheel.linux.environment] OEQ_CUDA_ROOTS = "/usr/local/cuda-12.6 /usr/local/cuda-13.0" ROCM_PATH = "/opt/rocm-sdk-core/_rocm_sdk_core" -OEQ_ROCM_INDEX_URL = "https://download.pytorch.org/whl/rocm10.0" -OEQ_ROCM_SDK_VERSION = "10.0.0" -OEQ_ROCM_SDK_SHA256 = "3bf4a72d11aa2a4ee1e90572c73630f937d38c1b7af4dda45b483b54655138a7" [tool.cibuildwheel.linux] repair-wheel-command = """auditwheel repair -w {dest_dir} {wheel} \ @@ -116,3 +112,8 @@ repair-wheel-command = """auditwheel repair -w {dest_dir} {wheel} \ --exclude libtorch_cuda.so --exclude libtorch_hip.so \ --exclude libcuda.so.1 --exclude 'libcudart.so.*' --exclude 'libnvrtc.so.*' \ --exclude 'libamdhip64.so.*' --exclude 'libhiprtc.so.*'""" + +[tool.openequivariance.rocm-sdk-core] +index-url = "https://download.pytorch.org/whl/rocm10.0" +version = "10.0.0" +sha256 = "3bf4a72d11aa2a4ee1e90572c73630f937d38c1b7af4dda45b483b54655138a7" From 6f8cc6a95a0152f1ef0dd0f664d094e3fdecd059 Mon Sep 17 00:00:00 2001 From: Vivek Bharadwaj Date: Tue, 10 Nov 2026 22:40:07 -0800 Subject: [PATCH 38/38] Fix CI failures. --- .github/workflows/requirements_cuda_ci.txt | 3 ++- .github/workflows/verify_extension_build.yml | 2 ++ 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/.github/workflows/requirements_cuda_ci.txt b/.github/workflows/requirements_cuda_ci.txt index 2be0fc59..09f2376a 100644 --- a/.github/workflows/requirements_cuda_ci.txt +++ b/.github/workflows/requirements_cuda_ci.txt @@ -1,5 +1,6 @@ +--extra-index-url https://download.pytorch.org/whl/cu129 numpy==2.2.5 -torch==2.13.0 --index-url https://download.pytorch.org/whl/cu129 +torch==2.13.0+cu129 pytest==9.0.3 ninja==1.11.1.4 nanobind==3.1.0 diff --git a/.github/workflows/verify_extension_build.yml b/.github/workflows/verify_extension_build.yml index 7843cb87..c33a4da2 100644 --- a/.github/workflows/verify_extension_build.yml +++ b/.github/workflows/verify_extension_build.yml @@ -87,6 +87,8 @@ jobs: run: | python -m pip install --upgrade pip pip install -r .github/workflows/requirements_sycl_ci.txt + source /opt/intel/oneapi/setvars.sh + export CXX=icpx pip install -e "./openequivariance[sycl]" - name: Test SYCL extension build via import