Patch LookupTable

See torch/issues/492 for reference
ashwin2802 · Mar 6, 2021 · 4ea8700 · 4ea8700
1 parent 1ae6aa0
commit 4ea8700
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 3 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>

diff --git a/rocks/cunn-scm-1.rockspec b/rocks/cunn-scm-1.rockspec
@@ -2,7 +2,8 @@ package = "cunn"
 version = "scm-1"
 
 source = {
-   url = "git://github.com/torch/cunn.git",
+   url = "git://github.com/ashwin2802/cunn.git",
+   branch = "cuda11"
 }
 
 description = {