From 421941993d169c943f2c364bfe9c48b603339fd1 Mon Sep 17 00:00:00 2001 From: Neil Bickford <57467222+NeilBickford-NV@users.noreply.github.com> Date: Tue, 7 Nov 2023 11:47:18 -0800 Subject: CUDA: Fixes for NVRTC 12.x and warp mask ambiguity; adds CC 8.x warp reduction intrinsics. (#3314) * CUDA: Fixes for NVRTC 12.x, warp mask ambiguity; add reduction partial specializations. * Fixes running NVRTC on CUDA 12 without a specified profile (used in testing, e.g. `slang-test -api cuda -category wave`) * Fixes mask ambiguity between getting the lane index from threadId.x and a full mask of threads. * Adds partial specializations for compute capability 8.x warp reduction intrinsics. * Fix formatting --- prelude/slang-cuda-prelude.h | 41 +++++++++++++++++++++++++++++++++++++++-- 1 file changed, 39 insertions(+), 2 deletions(-) (limited to 'prelude') diff --git a/prelude/slang-cuda-prelude.h b/prelude/slang-cuda-prelude.h index 24e400d2d..38f8a721a 100644 --- a/prelude/slang-cuda-prelude.h +++ b/prelude/slang-cuda-prelude.h @@ -53,7 +53,8 @@ # define SLANG_CUDA_WARP_SIZE 32 #endif -#define SLANG_CUDA_WARP_MASK (SLANG_CUDA_WARP_SIZE - 1) +#define SLANG_CUDA_WARP_MASK (SLANG_CUDA_WARP_SIZE - 1) // Used for masking threadIdx.x to the warp lane index +#define SLANG_CUDA_WARP_BITMASK (~int(0)) // #define SLANG_FORCE_INLINE inline @@ -1418,10 +1419,16 @@ __inline__ __device__ bool _waveIsSingleLane(WarpMask mask) } // Returns the power of 2 size of run of set bits. Returns 0 if not a suitable run. +// Examples: +// 0b00000000'00000000'00000000'11111111 -> 8 +// 0b11111111'11111111'11111111'11111111 -> 32 +// 0b00000000'00000000'00000000'00011111 -> 0 (since 5 is not a power of 2) +// 0b00000000'00000000'00000000'11110000 -> 0 (since the run of bits does not start at the LSB) +// 0b00000000'00000000'00000000'00100111 -> 0 (since it is not a single contiguous run) __inline__ __device__ int _waveCalcPow2Offset(WarpMask mask) { // This should be the most common case, so fast path it - if (mask == SLANG_CUDA_WARP_MASK) + if (mask == SLANG_CUDA_WARP_BITMASK) { return SLANG_CUDA_WARP_SIZE; } @@ -1647,6 +1654,36 @@ __inline__ __device__ T _waveMin(WarpMask mask, T val) { return _waveReduceScala template __inline__ __device__ T _waveMax(WarpMask mask, T val) { return _waveReduceScalar, T>(mask, val); } +// Fast-path specializations when CUDA warp reduce operators are available +#if __CUDA_ARCH__ >= 800 // 8.x or higher +template<> +__inline__ __device__ unsigned _waveOr(WarpMask mask, unsigned val) { return __reduce_or_sync(mask, val); } + +template<> +__inline__ __device__ unsigned _waveAnd(WarpMask mask, unsigned val) { return __reduce_and_sync(mask, val); } + +template<> +__inline__ __device__ unsigned _waveXor(WarpMask mask, unsigned val) { return __reduce_xor_sync(mask, val); } + +template<> +__inline__ __device__ unsigned _waveSum(WarpMask mask, unsigned val) { return __reduce_add_sync(mask, val); } + +template<> +__inline__ __device__ int _waveSum(WarpMask mask, int val) { return __reduce_add_sync(mask, val); } + +template<> +__inline__ __device__ unsigned _waveMin(WarpMask mask, unsigned val) { return __reduce_min_sync(mask, val); } + +template<> +__inline__ __device__ int _waveMin(WarpMask mask, int val) { return __reduce_min_sync(mask, val); } + +template<> +__inline__ __device__ unsigned _waveMax(WarpMask mask, unsigned val) { return __reduce_max_sync(mask, val); } + +template<> +__inline__ __device__ int _waveMax(WarpMask mask, int val) { return __reduce_max_sync(mask, val); } +#endif + // Multiple -- cgit v1.2.3