Fix ROCm builds on 6.4.3+ and on GPU-less builders

csrc/cuda/utils.cuh

@@ -6,19 +6,26 @@
   AT_ASSERTM(x.device().is_cuda(), #x " must be CUDA tensor")
 #define CHECK_INPUT(x) AT_ASSERTM(x, "Input mismatch")
 
-__device__ __inline__ at::Half __shfl_up_sync(const unsigned mask,
+// On ROCm, __shfl_*_sync requires a 64-bit mask; on CUDA it's 32-bit.
+#ifdef USE_ROCM
+  using warp_mask_t = unsigned long long;
+#else
+  using warp_mask_t = unsigned int;
+#endif
+
+__device__ __inline__ at::Half __shfl_up_sync(const warp_mask_t mask,
                                               const at::Half var,
                                               const unsigned int delta) {
   return __shfl_up_sync(mask, var.operator __half(), delta);
 }
 
-__device__ __inline__ at::Half __shfl_down_sync(const unsigned mask,
+__device__ __inline__ at::Half __shfl_down_sync(const warp_mask_t mask,
                                                 const at::Half var,
                                                 const unsigned int delta) {
   return __shfl_down_sync(mask, var.operator __half(), delta);
 }
 
-__device__ __inline__ at::Half __shfl_sync(const unsigned mask,
+__device__ __inline__ at::Half __shfl_sync(const warp_mask_t mask,
                                            const at::Half var,
                                            const int delta) {
   return __shfl_sync(mask, var.operator __half(), delta);
@@ -43,9 +50,30 @@ __shfl(const at::Half var, const int delta) {
 __device__ __inline__ at::Half __ldg(const at::Half* ptr) {
   return __ldg(reinterpret_cast<const __half*>(ptr));
 }
-#define SHFL_UP_SYNC(mask, var, delta) __shfl_up(var, delta)
-#define SHFL_DOWN_SYNC(mask, var, delta) __shfl_down(var, delta)
-#define SHFL_SYNC(mask, var, delta) __shfl(var, delta)
+
+__device__ __inline__ at::Half __shfl_up(const at::Half var,
+                                         const unsigned int delta,
+                                         const int width) {
+  return __shfl_up(var.operator __half(), delta, width);
+}
+
+__device__ __inline__ at::Half __shfl_down(const at::Half var,
+                                           const unsigned int delta,
+                                           const int width) {
+  return __shfl_down(var.operator __half(), delta, width);
+}
+
+__device__ __inline__ at::Half __shfl(const at::Half var, const int delta,
+                                      const int width) {
+  return __shfl(var.operator __half(), delta, width);
+}
+
+// CUDA's `__shfl_*_sync` default to a width of `warpSize`, i.e. 32. HIP
+// defaults to the wavefront size, which is 64 on CDNA, so the width has to be
+// passed explicitly to preserve the 32-lane semantics these kernels assume.
+#define SHFL_UP_SYNC(mask, var, delta) __shfl_up(var, delta, 32)
+#define SHFL_DOWN_SYNC(mask, var, delta) __shfl_down(var, delta, 32)
+#define SHFL_SYNC(mask, var, delta) __shfl(var, delta, 32)
 #else
 #define SHFL_UP_SYNC __shfl_up_sync
 #define SHFL_DOWN_SYNC __shfl_down_sync

setup.py

@@ -19,8 +19,12 @@
 URL = 'https://github.com/rusty1s/pytorch_sparse'
 
 WITH_CUDA = False
-if torch.cuda.is_available():
-    WITH_CUDA = CUDA_HOME is not None or torch.version.hip
+if torch.version.hip is not None:
+    # A ROCm build of PyTorch always targets HIP, and no GPU needs to be
+    # visible at build time (e.g. when building wheels on a CI runner):
+    WITH_CUDA = True
+elif torch.cuda.is_available():
+    WITH_CUDA = CUDA_HOME is not None
 suffices = ['cpu', 'cuda'] if WITH_CUDA else ['cpu']
 if os.getenv('FORCE_CUDA', '0') == '1':
     suffices = ['cuda', 'cpu']
@@ -144,7 +148,7 @@ def get_extensions():
 
 # work-around hipify abs paths
 include_package_data = True
-if torch.cuda.is_available() and torch.version.hip:
+if torch.version.hip is not None:
     include_package_data = False
 
 setup(