Add AMD GPU support via ROCm/HIP

.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

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)
 

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) {

include/cuda/cuda_util.cuh

@@ -1,17 +1,27 @@
 #ifndef UTIL_H
 #define UTIL_H
+
+#ifdef USE_ROCM
+#include <hip/hip_runtime.h>
+#else
+#include "device_atomic_functions.h"
+#include "device_launch_parameters.h"
 #include <cuda.h>
 #include <cuda_runtime.h>
 #include <cuda_runtime_api.h>
+#endif
+
 #include <stdio.h>
 #include <string.h>
 
-#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,18 +289,18 @@ template <typename data>
 __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);
   }
 }
 
 template <typename data>
 __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);
   }
 }
 

include/cuda/sddmm_cuda.cuh

@@ -1,14 +1,18 @@
 #ifndef SDDMM_CUDA
 #define SDDMM_CUDA
 
+#ifdef USE_ROCM
+#include <hip/hip_runtime.h>
+#else
+#include "device_atomic_functions.h"
+#include "device_launch_parameters.h"
 #include <cuda.h>
 #include <cuda_runtime.h>
 #include <cuda_runtime_api.h>
+#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;

include/cuda/spmm_cuda.cuh

@@ -1,8 +1,12 @@
 #ifndef SPMM_CUDA
 #define SPMM_CUDA
 
+#ifdef USE_ROCM
+#include <hip/hip_runtime.h>
+#else
 #include <cuda.h>
 #include <cuda_runtime_api.h>
+#endif
 
 #include "../gspmm.h"
 #include "cuda_util.cuh"

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,

src/cuda/spconv_cuda.cu

@@ -11,6 +11,16 @@
 #include <tuple>
 #include <vector>
 
+// 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"