From 593007c04ab37b64c41ef4cb46ab3ee963f7d2c6 Mon Sep 17 00:00:00 2001 From: Harkirat Gill Date: Mon, 22 Jun 2026 14:12:27 -0400 Subject: [PATCH] fix: use 64-bit warp mask for ROCm HIP compatibility --- csrc/cuda/utils.cuh | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/csrc/cuda/utils.cuh b/csrc/cuda/utils.cuh index 747a8e2c..8c85368b 100644 --- a/csrc/cuda/utils.cuh +++ b/csrc/cuda/utils.cuh @@ -6,19 +6,25 @@ 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, +#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);