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