Enable generic functions for 64-bit shuffles (#171)

* enable generic functions for 64-bit shuffles * Update include/cutlass/gpu_generics.h Co-authored-by: Joe Todd <[email protected]> --------- Co-authored-by: Joe Todd <[email protected]>
codeplaysoftware · Jan 9, 2025 · 096165f · 096165f
1 parent 25b7875
commit 096165f
Showing 1 changed file with 12 additions and 8 deletions.
diff --git a/include/cutlass/gpu_generics.h b/include/cutlass/gpu_generics.h
@@ -249,10 +249,11 @@ unsigned int byte_perm(unsigned int x, unsigned int y, unsigned int s) {
 
 // shfl
 
+template<typename T>
 CUTLASS_DEVICE
-unsigned int shfl_up_sync(
+T shfl_up_sync(
   unsigned int const mask,
-  unsigned int const var,
+  T const var,
   int const delta,
   int const width = NumThreadsPerWarp) {
 #if defined(__CUDA_ARCH__)
@@ -264,10 +265,11 @@ unsigned int shfl_up_sync(
 #endif
 }
 
+template<typename T>
 CUTLASS_DEVICE
-unsigned int shfl_down_sync(
+T shfl_down_sync(
   unsigned int const mask,
-  unsigned int const var,
+  T const var,
   int const delta,
   int const width = NumThreadsPerWarp) {
 #if defined(__CUDA_ARCH__)
@@ -279,10 +281,11 @@ unsigned int shfl_down_sync(
 #endif
 }
 
+template<typename T>
 CUTLASS_DEVICE
-unsigned int shfl_sync(
+T shfl_sync(
   unsigned int const mask,
-  unsigned int const var,
+  T const var,
   int const delta,
   int const width = NumThreadsPerWarp) {
 #if defined(__CUDA_ARCH__)
@@ -296,10 +299,11 @@ unsigned int shfl_sync(
 #endif
 }
 
+template<typename T>
 CUTLASS_DEVICE
-unsigned int shfl_xor_sync(
+T shfl_xor_sync(
   unsigned int const mask,
-  unsigned int const var,
+  T const var,
   int const laneMask,
   int const width = NumThreadsPerWarp) {
 #if defined(__CUDA_ARCH__)