Support larger lookup in permute (#5086)

Summary:
Pull Request resolved: #5086

X-link: https://github.com/facebookresearch/FBGEMM/pull/2094

For lengths per shard exceeding 2^31, we avoid overflow resulting in undefined behavior.

Reviewed By: spcyppt

Differential Revision: D86209662

fbshipit-source-id: 6d51290f3436629571677091c42b76b6f98e5790

Author: kausv

fbgemm_gpu/src/sparse_ops/sparse_permute_2d.cu

@@ -20,7 +20,7 @@ template <
     typename indices_t,
     typename weights_t>
 __global__ __launch_bounds__(kMaxThreads) void permute_2D_data_kernel(
-    int32_t len,
+    int64_t len,
     int32_t T,
     int32_t B,
     const indices_t* __restrict__ indices,