Gather/Scatter. (pytorch#3743)

Summary:

X-link: facebookresearch/FBGEMM#824

Minor improvements.

Differential Revision: D70226730

Author: levendlee

fbgemm_gpu/experimental/gen_ai/src/gather_scatter/gather_scatter.cu

@@ -54,10 +54,12 @@ __global__ static void gather_or_scatter_along_first_dim_kernel(
   int M = cute::get<0>(problem_shape);
   int N = cute::get<1>(problem_shape);
   int K = cute::get<2>(problem_shape);
+  int L = cute::get<3>(problem_shape);
 
   static_assert(cute::is_static<TileShape>::value);
   constexpr int kBlkNOrM = cute::size<0>(tile_shape);
   constexpr int kBlkK = cute::size<1>(tile_shape);
+  constexpr int kBlkL = cute::size<2>(tile_shape);
 
   using SmemT = SharedStorage<kBlkNOrM, DataType, IndexType, SmemLayout>;
   constexpr int kPipeMax = SmemT::kPipeMax;
@@ -79,15 +81,17 @@ __global__ static void gather_or_scatter_along_first_dim_kernel(
 
   if (threadIdx.x == 0) {
     // Tensors on HBM.
-    cute::Tensor gA = tma_load_input.get_tma_tensor(cute::make_shape(M, K));
-    cute::Tensor gB = tma_store_output.get_tma_tensor(cute::make_shape(N, K));
+    cute::Tensor gA = tma_load_input.get_tma_tensor(cute::make_shape(M, K, L));
+    cute::Tensor gB =
+        tma_store_output.get_tma_tensor(cute::make_shape(N, K, L));
     // Tensors on SMEM.
     cute::Tensor sA = cute::make_tensor(
         cute::make_smem_ptr(smem.data.data()), cute::group<0, 2>(SmemLayout{}));
 
-    constexpr int kTmaTransactionBytes = kBlkK * sizeof(DataType);
+    constexpr int kTmaTransactionBytes = kBlkK * kBlkL * sizeof(DataType);
     const int kNumKTiles = ((K + kBlkK - 1) / kBlkK);
-    const int kNumNOrMKTiles = kBlkNOrM * kNumKTiles;
+    const int kNumNOrMKTiles =
+        std::min(kBlkNOrM, indexing_dim - n_or_m_offset) * kNumKTiles;
     const int kNumIterations = kNumNOrMKTiles + kPipeMax - 1;
 
     for (int iteration = 0; iteration < kNumIterations; ++iteration) {
@@ -110,12 +114,12 @@ __global__ static void gather_or_scatter_along_first_dim_kernel(
 
         cute::Tensor tAgA = cute::local_tile(
             gA,
-            cute::Tile<cute::_1, cute::Int<kBlkK>>{},
-            cute::make_coord(m, k));
+            cute::Tile<cute::_1, cute::Int<kBlkK>, cute::Int<kBlkL>>{},
+            cute::make_coord(m, k, 0));
         cute::Tensor tAsA = cute::local_tile(
             sA,
-            cute::Tile<cute::_1, cute::Int<kBlkK>>{},
-            cute::make_coord(load_pipe, 0));
+            cute::Tile<cute::_1, cute::Int<kBlkK>, cute::Int<kBlkL>>{},
+            cute::make_coord(load_pipe, 0, 0));
 
         auto& tma_load_mbar = smem.tma_load_barrier[load_pipe];
         cute::initialize_barrier(smem.tma_load_barrier[load_pipe], 1);
@@ -149,12 +153,12 @@ __global__ static void gather_or_scatter_along_first_dim_kernel(
 
         cute::Tensor tAgB = cute::local_tile(
             gB,
-            cute::Tile<cute::_1, cute::Int<kBlkK>>{},
-            cute::make_coord(n, k));
+            cute::Tile<cute::_1, cute::Int<kBlkK>, cute::Int<kBlkL>>{},
+            cute::make_coord(n, k, 0));
         cute::Tensor tAsA = cute::local_tile(
             sA,
-            cute::Tile<cute::_1, cute::Int<kBlkK>>{},
-            cute::make_coord(store_pipe, 0));
+            cute::Tile<cute::_1, cute::Int<kBlkK>, cute::Int<kBlkL>>{},
+            cute::make_coord(store_pipe, 0, 0));
 
         auto tma_store_per_cta = tma_store_output.get_slice(0);
         cute::copy(
@@ -205,36 +209,51 @@ void gather_or_scatter_along_first_dim(
   assert(dst.dtype() == TorchDTypeTrait<DataType>::dtype());
   assert(index.dtype() == TorchDTypeTrait<IndexType>::dtype());
 
+  constexpr int L = 256;
   const int M = src.size(0);
-  const int K = src.size(1);
   const int N = dst.size(0);
+  const int K = src.size(1) / L;
 
-  auto src_gmem_layout =
-      cute::make_layout(cute::make_shape(M, K), cute::make_stride(K, 1));
+  auto src_gmem_layout = cute::make_layout(
+      cute::make_shape(M, K, L), cute::make_stride(K * L, L, 1));
   auto src_gmem_tensor = cute::make_tensor(
       cute::make_gmem_ptr(reinterpret_cast<DataType*>(src.data_ptr())),
       src_gmem_layout);
 
-  auto dst_gmem_layout =
-      cute::make_layout(cute::make_shape(N, K), cute::make_stride(K, 1));
+  auto dst_gmem_layout = cute::make_layout(
+      cute::make_shape(N, K, L), cute::make_stride(K * L, L, 1));
   auto dst_gmem_tensor = cute::make_tensor(
       cute::make_gmem_ptr(reinterpret_cast<DataType*>(dst.data_ptr())),
       dst_gmem_layout);
 
   constexpr int kBlkNOrM = 1;
-  constexpr int kBlkK = 256;
-  constexpr int kPipeMax = 4;
+  constexpr int kBlkK = 4;
+  constexpr int kBlkL = L;
+  constexpr int kPipeMax = 3;
 
   auto smem_layout = cute::make_layout(
-      cute::make_shape(cute::Int<kPipeMax>{}, cute::_1{}, cute::Int<kBlkK>{}),
-      cute::make_stride(cute::Int<kBlkK>{}, cute::Int<kBlkK>{}, cute::_1{}));
+      cute::make_shape(
+          cute::Int<kPipeMax>{},
+          cute::_1{},
+          cute::Int<kBlkK>{},
+          cute::Int<kBlkL>{}),
+      cute::make_stride(
+          cute::Int<kBlkK * kBlkL>{},
+          cute::Int<kBlkK * kBlkL>{},
+          cute::Int<kBlkL>{},
+          cute::_1{}));
   auto tma_load = cute::make_tma_copy(
-      cute::SM90_TMA_LOAD{}, src_gmem_tensor, smem_layout(0, cute::_, cute::_));
+      cute::SM90_TMA_LOAD{},
+      src_gmem_tensor,
+      smem_layout(0, cute::_, cute::_, cute::_));
   auto tma_store = cute::make_tma_copy(
-      TMAStoreInst{}, dst_gmem_tensor, smem_layout(0, cute::_, cute::_));
+      TMAStoreInst{},
+      dst_gmem_tensor,
+      smem_layout(0, cute::_, cute::_, cute::_));
 
-  auto problem_shape = cute::make_shape(M, N, K);
-  auto tile_shape = cute::make_shape(cute::Int<kBlkNOrM>{}, cute::Int<kBlkK>{});
+  auto problem_shape = cute::make_shape(M, N, K, L);
+  auto tile_shape = cute::make_shape(
+      cute::Int<kBlkNOrM>{}, cute::Int<kBlkK>{}, cute::Int<kBlkL>{});
 
   using SmemT =
       SharedStorage<kBlkNOrM, DataType, IndexType, decltype(smem_layout)>;

fbgemm_gpu/experimental/gen_ai/test/gather_scatter/gather_scatter_test.py

@@ -78,6 +78,12 @@ def ref_fn():
         _test_gather_along_first_dim(1024, 1024, 1024)
         _test_gather_along_first_dim(1024, 1024, 1024, compile=True)
 
+        _test_gather_along_first_dim(1, 1, 5120)
+        _test_gather_along_first_dim(128, 128, 5120)
+        _test_gather_along_first_dim(2048, 2048, 5120)
+        _test_gather_along_first_dim(4096, 4096, 5120)
+        _test_gather_along_first_dim(8192, 8192, 5120)
+
     def test_scatter_add_along_first_dim(self) -> None:
         def _test_scatter_add_along_first_dim(
             M: int, N: int, K: int, compile: bool = False
@@ -101,7 +107,7 @@ def _test_scatter_add_along_first_dim(
             logger.info("Running PyTorch")
             ref_dst.scatter_add_(0, indices_2d, src)
 
-            torch.testing.assert_close(test_dst, ref_dst, atol=1e-3, rtol=2e-2)
+            torch.testing.assert_close(test_dst, ref_dst, atol=1e-3, rtol=2.1e-2)
 
             def fn():
                 op = torch.ops.fbgemm.scatter_add_along_first_dim
@@ -137,6 +143,12 @@ def ref_fn():
         _test_scatter_add_along_first_dim(1024, 1024, 1024)
         _test_scatter_add_along_first_dim(1024, 1024, 1024, compile=True)
 
+        _test_scatter_add_along_first_dim(1, 1, 5120)
+        _test_scatter_add_along_first_dim(128, 128, 5120)
+        _test_scatter_add_along_first_dim(2048, 2048, 5120)
+        _test_scatter_add_along_first_dim(4096, 4096, 5120)
+        _test_scatter_add_along_first_dim(8192, 8192, 5120)
+
 
 if __name__ == "__main__":
     unittest.main()