https://github.com/torch/cunn/issues/492

kyoto7250 · Feb 22, 2023 · 1e4d0e5 · 1e4d0e5
1 parent bfb9891
commit 1e4d0e5
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/lib/THCUNN/LookupTable.cu b/lib/THCUNN/LookupTable.cu
@@ -13,6 +13,7 @@
 #include "THCHalf.h"
 #include "THCHalfAutoNumerics.cuh"
 #include "THCTensorSort.cuh"
+#define FULL_MASK 0xffffffff
 
 const int WARP_SIZE = 32;
 
@@ -29,7 +30,7 @@ __device__ __forceinline__ bool warpHasCollision(int val)
   #pragma unroll
   for (int i = 1; i <= 16; i++)
   {
-    dup |= (__shfl(val, (laneId + i) % 32) == val);
+    dup |= (__shfl_sync(FULL_MASK, val, (laneId + i) % 32) == val);
   }
 
 #else
@@ -46,7 +47,7 @@ __device__ __forceinline__ bool warpHasCollision(int val)
 
 #endif
 
-  return __any(dup) != 0;
+  return __any_sync(FULL_MASK, dup) != 0;
 }
 
 template <typename Dtype>