Faiss GPU large transpose fix (#2075)

Summary:
Pull Request resolved: #2075

This is a fix for #1996

namely, that large transposition jobs (where one of the dimensions is > 65535) will still work by performing a loop over the gridDim.y

Reviewed By: mdouze

Differential Revision: D31481494

fbshipit-source-id: af35af36cce27ce7d44128d95cc229dd5c4b4b56

faiss/gpu/utils/Transpose.cuh

@@ -101,14 +101,14 @@ __global__ void transposeOuter(
         IndexT t1,
         IndexT t2,
         IndexT i1) {
-    IndexT gt1 = blockIdx.y;
     IndexT gt2 = blockIdx.x;
+    for (IndexT gt1 = blockIdx.y; gt1 < t1; gt1 += gridDim.y) {
+        auto curIn = in + i1 * (gt1 * t2 + gt2);
+        auto curOut = out + i1 * (gt2 * t1 + gt1);
 
-    in += i1 * (gt1 * t2 + gt2);
-    out += i1 * (gt2 * t1 + gt1);
-
-    for (IndexT i = threadIdx.x; i < i1; i += blockDim.x) {
-        out[i] = in[i];
+        for (IndexT i = threadIdx.x; i < i1; i += blockDim.x) {
+            curOut[i] = curIn[i];
+        }
     }
 }
 
@@ -166,7 +166,10 @@ void runTransposeAny(
             innerSize *= in.getSize(i);
         }
 
-        auto grid = dim3(in.getSize(1), in.getSize(0));
+        // The grid y dimension is more limited; we do a grid loop if necessary
+        int maxGridY = getCurrentDeviceProperties().maxGridSize[1];
+        auto grid = dim3(in.getSize(1), std::min(in.getSize(0), maxGridY));
+
         int block = (innerSize < maxThreads) ? innerSize : maxThreads;
 
         if (totalSize <= (size_t)std::numeric_limits<int>::max()) {