From 1ff4c69b407bfc756e4d084bcdcdbda89171e358 Mon Sep 17 00:00:00 2001 From: Jeff Daily Date: Sat, 20 Jun 2026 00:50:17 +0000 Subject: [PATCH] [ROCm] Add AMD GPU support via ROCm/HIP This adds AMD GPU support to dgSPARSE-Lib through ROCm/HIP, alongside the existing CUDA path. A CUDAExtension built against a ROCm PyTorch wheel automatically runs torch.utils.hipify on the extension's CUDA sources, translating the CUDA runtime and cuSPARSE calls to their HIP equivalents at build time. The changes here make the existing kernels and build script compatible with that hipified output, so the library builds and runs on AMD GPUs with no separate code path to maintain. What changed: - setup.py: detect a ROCm build via torch.version.hip, treat it as a valid GPU build (ROCm has CUDA_HOME=None), and link hipsparse instead of cusparse. On Windows with ROCm, a BuildExtension subclass registers .hip as a C++ source extension and routes the host op-wrapper .cpp through hipcc via a generated shim (MSVC's cl.exe cannot parse the HIP runtime headers' GCC __attribute__ syntax), enabling ninja so include paths with spaces are escaped. - The CUDA headers guard their CUDA-only includes (device_atomic_functions.h, device_launch_parameters.h) under USE_ROCM, since HIP provides these through hip_runtime.h; the warp shuffles use a 64-bit full mask under ROCm (the __shfl*_sync intrinsics require it there) while keeping width-32 subgroup semantics that work on wave64 and wave32. - A few cuSPARSE/cuBLAS symbols that hipify does not currently map are aliased under USE_ROCM so the hipified sources compile unchanged; version.cpp returns HIP_VERSION on ROCm. Test Plan: Built and tested on an AMD Instinct MI250X (gfx90a) with ROCm 7.2.1; kernel compilation additionally covers gfx942, gfx950, and gfx1100. On Windows, built and tested on an AMD Radeon RX 9070 XT (gfx1201, RDNA4), where all five SpMM tests (spmm_sum/max/min/mean forward and spmm_sum backward) pass. The CUDA build path is unchanged; it was reconfirmed by building the GPU extension with nvcc (CUDA 12.8) against an upstream CUDA PyTorch wheel. This work was authored with the assistance of Claude, an AI assistant by Anthropic. --- .gitignore | 12 ++++++ README.md | 9 +++++ include/cuda/csr2csc.cuh | 9 +++++ include/cuda/cuda_util.cuh | 50 +++++++++++++++++++---- include/cuda/sddmm_cuda.cuh | 20 ++++++---- include/cuda/spmm_cuda.cuh | 4 ++ setup.py | 70 ++++++++++++++++++++++++++++----- src/cuda/spconv_cuda.cu | 10 +++++ src/sddmm/coosddmm_ebalance.cuh | 10 +++-- src/sddmm/csrsddmm_ebalance.cuh | 10 +++-- src/spmm.cpp | 8 ++++ src/util/cuda_util.cuh | 12 +++++- src/version.cpp | 8 ++++ 13 files changed, 198 insertions(+), 34 deletions(-) diff --git a/.gitignore b/.gitignore index 179c357..3d843df 100644 --- a/.gitignore +++ b/.gitignore @@ -11,3 +11,15 @@ build/ setup.cfg .vscode *.DS_Store + +# Auto-generated hipified files (torch.utils.hipify) +include/hip/ +include/gspmm_hip.h +src/hip/ +*_hip.cpp +*_hip.h +# Windows ROCm build shims (byte copies of .cpp routed through hipcc) and +# their hipified outputs +*_winhip.cu +*_winhip.hip +*.pyd diff --git a/README.md b/README.md index 4d2d196..c764a92 100644 --- a/README.md +++ b/README.md @@ -28,6 +28,15 @@ Or you can build from source pip install -e . ``` +### Building on AMD GPUs (ROCm) + +dgSPARSE also builds on AMD GPUs with ROCm. Install a ROCm build of PyTorch, then build from source the same way -- the CUDA sources are translated to HIP at build time by PyTorch's `torch.utils.hipify`, and the library links against hipSPARSE instead of cuSPARSE: + +```bash +export PYTORCH_ROCM_ARCH=gfx90a # your AMD GPU arch (e.g. gfx90a, gfx942, gfx1100, gfx1201) +pip install -e . +``` + A demo for SpMM inference time compared to other main-stream library. (Tested on RTX 3090 with feature=64). ![image1](benchmark/datasets_comparison.jpg) diff --git a/include/cuda/csr2csc.cuh b/include/cuda/csr2csc.cuh index 9eabaec..e9ec1f4 100644 --- a/include/cuda/csr2csc.cuh +++ b/include/cuda/csr2csc.cuh @@ -5,6 +5,15 @@ #include "cuda_util.cuh" +// torch's hipify lacks mappings for these cusparse Csr2csc symbols; provide +// them under USE_ROCM so the hipified copy compiles without modification. +#ifdef USE_ROCM +#define cusparseCsr2cscEx2_bufferSize hipsparseCsr2cscEx2_bufferSize +#define cusparseCsr2cscEx2 hipsparseCsr2cscEx2 +#define CUSPARSE_ACTION_NUMERIC HIPSPARSE_ACTION_NUMERIC +#define CUSPARSE_CSR2CSC_ALG1 HIPSPARSE_CSR2CSC_ALG1 +#endif + void csr2cscKernel(int m, int n, int nnz, int devid, int *csrRowPtr, int *csrColInd, float *csrVal, int *cscColPtr, int *cscRowInd, float *cscVal) { diff --git a/include/cuda/cuda_util.cuh b/include/cuda/cuda_util.cuh index 1eb14cf..9faf0d0 100644 --- a/include/cuda/cuda_util.cuh +++ b/include/cuda/cuda_util.cuh @@ -1,17 +1,27 @@ #ifndef UTIL_H #define UTIL_H + +#ifdef USE_ROCM +#include +#else +#include "device_atomic_functions.h" +#include "device_launch_parameters.h" #include #include #include +#endif + #include #include -#include "device_atomic_functions.h" -#include "device_launch_parameters.h" - #define CEIL(x, y) (((x) + (y)-1) / (y)) +// ROCm 7.2.1+ requires 64-bit mask for warp sync functions +#ifdef USE_ROCM +#define FULLMASK 0xffffffffffffffffULL +#else #define FULLMASK 0xffffffff +#endif #define MIN(a, b) ((a < b) ? a : b) #define MAX(a, b) ((a < b) ? b : a) @@ -113,6 +123,17 @@ enum gespmmAlg_t { if (tmps == segid && lane_id < 16) \ v += tmpv; +#ifdef USE_ROCM +#define checkCudaError(a) \ + do { \ + if (hipSuccess != (a)) { \ + fprintf(stderr, "Hip runTime error in line %d of file %s \ + : %s \n", \ + __LINE__, __FILE__, hipGetErrorString(hipGetLastError())); \ + exit(EXIT_FAILURE); \ + } \ + } while (0) +#else #define checkCudaError(a) \ do { \ if (cudaSuccess != (a)) { \ @@ -122,7 +143,19 @@ enum gespmmAlg_t { exit(EXIT_FAILURE); \ } \ } while (0) +#endif +#ifdef USE_ROCM +#define checkCuSparseError(a) \ + do { \ + if (HIPSPARSE_STATUS_SUCCESS != (a)) { \ + fprintf(stderr, "HipSparse runTime error in line %d of file %s \ + : %s \n", \ + __LINE__, __FILE__, hipGetErrorString(hipGetLastError())); \ + exit(EXIT_FAILURE); \ + } \ + } while (0) +#else #define checkCuSparseError(a) \ do { \ if (CUSPARSE_STATUS_SUCCESS != (a)) { \ @@ -132,6 +165,7 @@ enum gespmmAlg_t { exit(EXIT_FAILURE); \ } \ } while (0) +#endif __device__ __forceinline__ float sum_reduce(float acc, float x) { return acc + x; } @@ -255,10 +289,10 @@ template __device__ __forceinline__ void AllReduce4(data *multi, int stride, int warpSize) { for (; stride > 0; stride >>= 1) { - multi[0] += __shfl_xor_sync(0xffffffff, multi[0], stride, warpSize); - multi[1] += __shfl_xor_sync(0xffffffff, multi[1], stride, warpSize); - multi[2] += __shfl_xor_sync(0xffffffff, multi[2], stride, warpSize); - multi[3] += __shfl_xor_sync(0xffffffff, multi[3], stride, warpSize); + multi[0] += __shfl_xor_sync(FULLMASK, multi[0], stride, warpSize); + multi[1] += __shfl_xor_sync(FULLMASK, multi[1], stride, warpSize); + multi[2] += __shfl_xor_sync(FULLMASK, multi[2], stride, warpSize); + multi[3] += __shfl_xor_sync(FULLMASK, multi[3], stride, warpSize); } } @@ -266,7 +300,7 @@ template __device__ __forceinline__ void AllReduce(data multi, int stride, int warpSize) { for (; stride > 0; stride >>= 1) { - multi += __shfl_xor_sync(0xffffffff, multi, stride, warpSize); + multi += __shfl_xor_sync(FULLMASK, multi, stride, warpSize); } } diff --git a/include/cuda/sddmm_cuda.cuh b/include/cuda/sddmm_cuda.cuh index bec436c..7e72643 100644 --- a/include/cuda/sddmm_cuda.cuh +++ b/include/cuda/sddmm_cuda.cuh @@ -1,14 +1,18 @@ #ifndef SDDMM_CUDA #define SDDMM_CUDA +#ifdef USE_ROCM +#include +#else +#include "device_atomic_functions.h" +#include "device_launch_parameters.h" #include #include #include +#endif #include "../gspmm.h" #include "cuda_util.cuh" -#include "device_atomic_functions.h" -#include "device_launch_parameters.h" __global__ void sddmmCOO4Scale(int D_kcols, const unsigned long Size, int *S_cooRowInd, int *S_cooColInd, @@ -73,7 +77,7 @@ __global__ void sddmmCOO4Scale(int D_kcols, const unsigned long Size, multi += D1tmp0 * D2tmp0; } for (int stride = 16; stride > 0; stride >>= 1) { - multi += __shfl_xor_sync(0xffffffff, multi, stride, 32); + multi += __shfl_xor_sync(FULLMASK, multi, stride, 32); } if (threadIdx.x == 0 && threadIdx.y == 0) { O_cooVal[eid] = multi; @@ -144,7 +148,7 @@ __global__ void sddmmCOO2Scale(int D_kcols, const unsigned long Size, multi += D1tmp0 * D2tmp0; } for (int stride = 16; stride > 0; stride >>= 1) { - multi += __shfl_xor_sync(0xffffffff, multi, stride, 32); + multi += __shfl_xor_sync(FULLMASK, multi, stride, 32); } if (threadIdx.x == 0 && threadIdx.y == 0) { O_cooVal[eid] = multi; @@ -211,7 +215,7 @@ __global__ void sddmmCOO1Scale(int D_kcols, const unsigned long Size, multi += D1tmp0 * D2tmp0; } for (int stride = 16; stride > 0; stride >>= 1) { - multi += __shfl_xor_sync(0xffffffff, multi, stride, 32); + multi += __shfl_xor_sync(FULLMASK, multi, stride, 32); } if (threadIdx.x == 0 && threadIdx.y == 0) { O_cooVal[eid] = multi; @@ -299,7 +303,7 @@ __global__ void sddmmCSR2Scale(const int S_mrows, int D_kcols, multi += D1tmp0 * D2tmp0; } for (int stride = 16; stride > 0; stride >>= 1) { - multi += __shfl_xor_sync(0xffffffff, multi, stride, 32); + multi += __shfl_xor_sync(FULLMASK, multi, stride, 32); } if (REDUCE::Op == MEAN && length > 0) { multi /= length; @@ -389,7 +393,7 @@ __global__ void sddmmCSR1Scale(const int S_mrows, int D_kcols, multi += D1tmp0 * D2tmp0; } for (int stride = 16; stride > 0; stride >>= 1) { - multi += __shfl_xor_sync(0xffffffff, multi, stride, 32); + multi += __shfl_xor_sync(FULLMASK, multi, stride, 32); } if (REDUCE::Op == MEAN && length > 0) { multi /= length; @@ -498,7 +502,7 @@ __global__ void sddmmCSR1Scale_with_mask(const int S_mrows, int D_kcols, // multi += D1tmp0 * D2tmp0; } for (int stride = 16; stride > 0; stride >>= 1) { - multi += __shfl_xor_sync(0xffffffff, multi, stride, 32); + multi += __shfl_xor_sync(FULLMASK, multi, stride, 32); } if (threadIdx.x == 0 && threadIdx.y == 0) { O_csrVal[eid] = multi; diff --git a/include/cuda/spmm_cuda.cuh b/include/cuda/spmm_cuda.cuh index 8874063..cba5394 100644 --- a/include/cuda/spmm_cuda.cuh +++ b/include/cuda/spmm_cuda.cuh @@ -1,8 +1,12 @@ #ifndef SPMM_CUDA #define SPMM_CUDA +#ifdef USE_ROCM +#include +#else #include #include +#endif #include "../gspmm.h" #include "cuda_util.cuh" diff --git a/setup.py b/setup.py index 6f90d44..c220f07 100644 --- a/setup.py +++ b/setup.py @@ -1,6 +1,8 @@ import glob import os import os.path as osp +import shutil +import sys from itertools import product import torch @@ -11,16 +13,41 @@ CUDAExtension, ) +# Detect ROCm build (HIP backend) +IS_ROCM = hasattr(torch.version, 'hip') and torch.version.hip is not None +IS_WINDOWS = sys.platform == 'win32' + __version__ = '0.1.1' URL = 'https://github.com/dgSPARSE/dgSPARSE-Lib' WITH_CUDA = False if torch.cuda.is_available(): - WITH_CUDA = CUDA_HOME is not None + # ROCm builds have CUDA_HOME=None but are still valid GPU builds + WITH_CUDA = CUDA_HOME is not None or IS_ROCM suffices = ['cuda'] if WITH_CUDA else ['cpu'] if os.getenv('FORCE_CUDA', '0') == '1': suffices = ['cuda'] -print(f'Building with CUDA: {WITH_CUDA}, ', 'CUDA_HOME:', CUDA_HOME) +print(f'Building with CUDA: {WITH_CUDA}, IS_ROCM: {IS_ROCM}, CUDA_HOME:', + CUDA_HOME) + + +# On Windows with ROCm, torch's BuildExtension adds .cu/.cuh to MSVC's +# _cpp_extensions but not .hip. After hipify, .cu sources become .hip and +# MSVC's compile() rejects them before spawn() can route them to hipcc. +# Subclass to also register .hip as a C++ extension on Windows+ROCm. +# use_ninja=True is required on Windows+ROCm: win_wrap_ninja_compile replaces +# spaces in MSVC include paths with backslash-escapes before passing to hipcc +# (hipcc forwards -I paths to clang without quoting, so unescaped spaces cause +# clang to split the path into separate tokens). The non-ninja path lacks this +# fix and fails with "no such file or directory: 'Files'" errors. +class HIPBuildExtension(BuildExtension): + + def build_extensions(self): + if IS_WINDOWS and IS_ROCM and hasattr(self.compiler, + '_cpp_extensions'): + if '.hip' not in self.compiler._cpp_extensions: + self.compiler._cpp_extensions.append('.hip') + super().build_extensions() def get_extensions(): @@ -34,12 +61,15 @@ def get_extensions(): undef_macros = [] libraries = [] extra_compile_args = {'cxx': ['-O2']} - extra_link_args = [ - '-s', - '-lm', - '-ldl', - ] - extra_link_args += ['-lcusparse'] if suffix == 'cuda' else [] + # -s/-lm/-ldl are POSIX-only; skip them on Windows + extra_link_args = [] if IS_WINDOWS else ['-s', '-lm', '-ldl'] + if suffix == 'cuda': + if IS_ROCM: + # On Windows lld-link uses .lib names; on Linux use -l prefix + extra_link_args += ['hipsparse.lib' + ] if IS_WINDOWS else ['-lhipsparse'] + else: + extra_link_args += ['-lcusparse'] if suffix == 'cuda': define_macros += [('WITH_CUDA', None)] @@ -49,7 +79,21 @@ def get_extensions(): extra_compile_args['nvcc'] = nvcc_flags name = main.split(os.sep)[-1][:-4] - sources = [main] + + # On Windows with ROCm, the host .cpp op-wrapper includes + # torch/extension.h, which pulls in c10/cuda/CUDAGuard.h and the hip + # headers (amd_hip_vector_types.h) whose GCC __attribute__ syntax MSVC + # cl.exe cannot parse. Route the host wrapper through the device + # toolchain (hipcc) by presenting it as a .cu file; hipify then renames + # the shim to _hip.cu and hipcc compiles it. + if IS_WINDOWS and IS_ROCM and suffix == 'cuda' and main.endswith( + '.cpp'): + shim = main[:-4] + '_winhip.cu' + shutil.copyfile(main, shim) + main_src = shim + else: + main_src = main + sources = [main_src] path = osp.join(extensions_dir, 'cuda', f'{name}_cuda.cu') if suffix == 'cuda' and osp.exists(path): @@ -117,7 +161,13 @@ def get_extensions(): ext_modules=get_extensions(), cmdclass={ 'build_ext': - BuildExtension.with_options(no_python_abi_suffix=True, use_ninja=False) + HIPBuildExtension.with_options( + no_python_abi_suffix=True, + # On Windows with ROCm, ninja is required: win_wrap_ninja_compile + # escapes spaces in MSVC include paths before forwarding to hipcc + # (the non-ninja single-compile path lacks this workaround). + use_ninja=IS_WINDOWS and IS_ROCM, + ) }, packages=find_packages(), include_package_data=True, diff --git a/src/cuda/spconv_cuda.cu b/src/cuda/spconv_cuda.cu index 91ab444..4ad07a8 100644 --- a/src/cuda/spconv_cuda.cu +++ b/src/cuda/spconv_cuda.cu @@ -11,6 +11,16 @@ #include #include +// ROCm/HIP compatibility: hipify does not map these symbols +#ifdef USE_ROCM +#ifndef CUBLAS_COMPUTE_16F +#define CUBLAS_COMPUTE_16F HIPBLAS_COMPUTE_16F +#endif +#ifndef CUBLAS_TENSOR_OP_MATH +#define CUBLAS_TENSOR_OP_MATH HIPBLAS_DEFAULT_MATH +#endif +#endif + #include "../../include/cuda/cuda_util.cuh" #include "../../include/cuda/spconv.cuh" #include "../../include/cuda/spconv_cuda.h" diff --git a/src/sddmm/coosddmm_ebalance.cuh b/src/sddmm/coosddmm_ebalance.cuh index 531a725..e5ad210 100644 --- a/src/sddmm/coosddmm_ebalance.cuh +++ b/src/sddmm/coosddmm_ebalance.cuh @@ -1,4 +1,8 @@ +#ifdef USE_ROCM +#include +#else #include +#endif #include "../util/cuda_util.cuh" @@ -65,7 +69,7 @@ __global__ void sddmm_coo_ebalance_vec4(int D_kcols, const int Size, multi += D1tmp0 * D2tmp0; } for (int stride = 16; stride > 0; stride >>= 1) { - multi += __shfl_xor_sync(0xffffffff, multi, stride, 32); + multi += __shfl_xor_sync(FULLMASK, multi, stride, 32); } if (threadIdx.x == 0 && threadIdx.y == 0) { O_cooVal[eid] = multi; @@ -136,7 +140,7 @@ __global__ void sddmm_coo_ebalance_vec2(int D_kcols, const int Size, multi += D1tmp0 * D2tmp0; } for (int stride = 16; stride > 0; stride >>= 1) { - multi += __shfl_xor_sync(0xffffffff, multi, stride, 32); + multi += __shfl_xor_sync(FULLMASK, multi, stride, 32); } if (threadIdx.x == 0 && threadIdx.y == 0) { O_cooVal[eid] = multi; @@ -203,7 +207,7 @@ __global__ void sddmm_coo_ebalance_scalar(int D_kcols, const int Size, multi += D1tmp0 * D2tmp0; } for (int stride = 16; stride > 0; stride >>= 1) { - multi += __shfl_xor_sync(0xffffffff, multi, stride, 32); + multi += __shfl_xor_sync(FULLMASK, multi, stride, 32); } if (threadIdx.x == 0 && threadIdx.y == 0) { O_cooVal[eid] = multi; diff --git a/src/sddmm/csrsddmm_ebalance.cuh b/src/sddmm/csrsddmm_ebalance.cuh index 625390b..282d27d 100644 --- a/src/sddmm/csrsddmm_ebalance.cuh +++ b/src/sddmm/csrsddmm_ebalance.cuh @@ -1,4 +1,8 @@ +#ifdef USE_ROCM +#include +#else #include +#endif #include "../util/cuda_util.cuh" @@ -48,7 +52,7 @@ __global__ void sddmm_csr_ebalance_vec4(const int S_mrows, int D_kcols, multi += D1tmp0 * D2tmp0; } for (int stride = 16; stride > 0; stride >>= 1) { - multi += __shfl_xor_sync(0xffffffff, multi, stride, 32); + multi += __shfl_xor_sync(FULLMASK, multi, stride, 32); } if (threadIdx.x == 0 && threadIdx.y == 0) { O_csrVal[eid] = multi; @@ -123,7 +127,7 @@ __global__ void sddmm_csr_ebalance_vec2(const int S_mrows, int D_kcols, multi += D1tmp0 * D2tmp0; } for (int stride = 16; stride > 0; stride >>= 1) { - multi += __shfl_xor_sync(0xffffffff, multi, stride, 32); + multi += __shfl_xor_sync(FULLMASK, multi, stride, 32); } if (threadIdx.x == 0 && threadIdx.y == 0) { O_csrVal[eid] = multi; @@ -197,7 +201,7 @@ __global__ void sddmm_csr_ebalance_scalar(const int S_mrows, int D_kcols, multi += D1tmp0 * D2tmp0; } for (int stride = 16; stride > 0; stride >>= 1) { - multi += __shfl_xor_sync(0xffffffff, multi, stride, 32); + multi += __shfl_xor_sync(FULLMASK, multi, stride, 32); } if (threadIdx.x == 0 && threadIdx.y == 0) { O_csrVal[eid] = multi; diff --git a/src/spmm.cpp b/src/spmm.cpp index 77ddc2b..1ce1fe6 100644 --- a/src/spmm.cpp +++ b/src/spmm.cpp @@ -268,3 +268,11 @@ TORCH_LIBRARY(dgsparse_spmm, m) { m.def("spmm_mean", &spmm_mean); m.def("csr2csc", &csr2csc); } + +// On Windows, torch's BuildExtension exports PyInit_ for every +// CUDAExtension. The spmm extension is loaded via torch.ops.load_library +// (not Python import), so PyInit_ is never called, but MSVC requires it +// to be resolvable at link time. +#ifdef _WIN32 +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {} +#endif diff --git a/src/util/cuda_util.cuh b/src/util/cuda_util.cuh index 6db6478..a816aaa 100644 --- a/src/util/cuda_util.cuh +++ b/src/util/cuda_util.cuh @@ -2,18 +2,26 @@ #pragma once +#ifdef USE_ROCM +#include +#else +#include "device_launch_parameters.h" #include #include #include - -#include "device_launch_parameters.h" +#endif /// heuristic choice of thread-block size const int RefThreadPerBlock = 256; #define CEIL(x, y) (((x) + (y)-1) / (y)) +// ROCm 7.2.1+ requires 64-bit mask for warp sync functions +#ifdef USE_ROCM +#define FULLMASK 0xffffffffffffffffULL +#else #define FULLMASK 0xffffffff +#endif #define DIV_UP(x, y) (((x) + (y)-1) / (y)) #define MIN(a, b) ((a < b) ? a : b) #define MAX(a, b) ((a < b) ? b : a) diff --git a/src/version.cpp b/src/version.cpp index c9a08f7..80222fb 100644 --- a/src/version.cpp +++ b/src/version.cpp @@ -5,12 +5,20 @@ #include #ifdef WITH_CUDA +#ifdef USE_ROCM +#include +#else #include #endif +#endif int64_t cuda_version() noexcept { #ifdef WITH_CUDA +#ifdef USE_ROCM + return HIP_VERSION; +#else return CUDA_VERSION; +#endif #else return -1; #endif