Repository: NVIDIA/Fuser
Pull request #4231: Multi-tile Blackwell MMA unit tests
Author: login: zasdfgbnm
Created: 2025-04-10T21:03:40Z | Merged: 2025-04-15T17:31:05Z

This PR implements indexing of TMem for MMA result, and adds the multiple-tile test.

Remaining things to fixes before we can really consider the MMA being supported:
- We are still using `nanosleep` for sync, we need to replace it with mbarrier.
- The predicate of MMA is incorrect. The correct way is to let one thread submit the job, but currently all threads are submitting jobs, which means that we are duplicating the compute, which will be wrong result if the problem size is > than the `K` of the instruction because some numbers will be added multiple times.


Diff:
diff --git a/csrc/device_lower/pass/index.cpp b/csrc/device_lower/pass/index.cpp
index 6217e925dbf..25605ddd0d4 100644
--- a/csrc/device_lower/pass/index.cpp
+++ b/csrc/device_lower/pass/index.cpp
@@ -2578,6 +2578,40 @@ Val* getOuterStrideBytes(TensorView* tv, const MmaOp* mma) {
   return SimplifyingIrBuilder::mulExpr(stride, dataTypeSize(tv->dtype()));
 }
 
+namespace {
+
+Val* indexBlackwellMmaOutput(
+    const MmaOp* mma,
+    const std::vector<ForLoop*>& for_loops) {
+  TensorView* tmem_tv = mma->out()->as<TensorView>();
+  NVF_ERROR(tmem_tv->getMemoryType() == MemoryType::Tensor, "Invalid tmem_tv");
+  const auto& tmem_info = GpuLower::current()->tmemInfo();
+  const auto& tensor_indexer = GpuLower::current()->tensorIndexer();
+
+  const std::vector<IterDomain*>& column_allocation_domain =
+      tmem_info.allocation.getTVInfo(tmem_tv).column_allocation;
+
+  std::vector<Val*> indices = tensor_indexer.getIndexFor(
+      mma, true, column_allocation_domain, for_loops);
+  Val* stride = tmem_tv->fusion()->oneVal();
+  Val* column_index = tmem_tv->fusion()->zeroVal();
+  for (const auto& [id, idx] :
+       std::ranges::views::reverse(zip(column_allocation_domain, indices))) {
+    column_index = SimplifyingIrBuilder::addExpr(
+        column_index, SimplifyingIrBuilder::mulExpr(idx, stride));
+    stride = SimplifyingIrBuilder::mulExpr(stride, id->extent());
+  }
+
+  column_index =
+      SimplifyingIrBuilder::maybeCastExpr(DataType::UInt16, column_index);
+
+  Val* index = SimplifyingIrBuilder::arrayExpr(std::vector<Val*>{
+      mma->fusion()->zeroVal(DataType::UInt16), column_index});
+  return GpuLower::current()->commonScalarMap().hoistScalar(index, for_loops);
+}
+
+} // namespace
+
 // Reference for smem strides:
 // https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#strides
 void IndexLowering::handle(const MmaOp* mma) {
@@ -2653,10 +2687,7 @@ void IndexLowering::handle(const MmaOp* mma) {
   }
   Val* out = nullptr;
   if (mma->out()->as<TensorView>()->getMemoryType() == MemoryType::Tensor) {
-    // TODO: hardcoded zero index for now
-    Val* index = IrBuilder::create<Val>(
-        std::vector<int64_t>{0, 0},
-        ArrayType(std::make_shared<DataType>(DataType::UInt16), 2));
+    auto index = indexBlackwellMmaOutput(mma, for_loops_);
     out = IrBuilder::create<kir::TensorIndex>(
         mma->out()->as<TensorView>(), index, DataType::TMemAddress);
   } else {
diff --git a/csrc/device_lower/pass/inline_ptx.cpp b/csrc/device_lower/pass/inline_ptx.cpp
index 610da496732..fdf463a013d 100644
--- a/csrc/device_lower/pass/inline_ptx.cpp
+++ b/csrc/device_lower/pass/inline_ptx.cpp
@@ -402,7 +402,8 @@ class LowerToInlinePtx : public kir::ExprMutator {
         IrBuilder::create<kir::Asm>(
             "nanosleep.u32",
             std::vector<Val*>{},
-            std::vector<Val*>{IrBuilder::create<Val>(100000, DataType::UInt32)},
+            std::vector<Val*>{
+                IrBuilder::create<Val>(100000000, DataType::UInt32)},
             kir::Asm::Options{/*volatile=*/true}));
   }
 
diff --git a/tests/cpp/test_mma.cpp b/tests/cpp/test_mma.cpp
index 36b05b02dea..44f4f69023c 100644
--- a/tests/cpp/test_mma.cpp
+++ b/tests/cpp/test_mma.cpp
@@ -19,6 +19,7 @@
 #include <scheduler/mma_utils.h>
 #include <scheduler/tools/inlining.h>
 #include <algorithm>
+#include <bit>
 #include <iterator>
 #include <unordered_map>
 
@@ -908,18 +909,13 @@ TEST_P(HopperSS, SingleTileTransposed) {
   NVF_CHECK(at::allclose(cg_outputs[0].as<at::Tensor>(), tref, 1e-5, 1e-5));
 }
 
-TEST_P(HopperSS, MultipleTile) {
-  Fusion fusion;
-  FusionGuard fg(&fusion);
-
-  constexpr int64_t num_tiles = 2;
-
-  auto shapes = matmulAtInputShape3DSS(
-      num_tiles * getM(macro),
-      num_tiles * getN(macro),
-      num_tiles * getK(macro),
-      layout);
-
+const char* skipIfSwizzleNotCompatibleWithTiling(
+    int64_t num_tiles,
+    MmaLayout layout,
+    MmaMacro macro,
+    MmaInputSmemSwizzle swizzle_a,
+    MmaInputSmemSwizzle swizzle_b) {
+  constexpr int64_t dtype_size = 2;
   const char* skip_reason =
       "This test stores smem inputs on the inner dimension densely, "
       "which is not compatible with this macro and swizzle mode "
@@ -931,13 +927,13 @@ TEST_P(HopperSS, MultipleTile) {
         ? getK(macro)
         : getM(macro);
     int64_t inner_size = num_tiles * inner_tile_size;
-    int64_t swizzle_size = getBytesFromSwizzle(swizzle_a) / dataTypeSize(dtype);
+    int64_t swizzle_size = getBytesFromSwizzle(swizzle_a) / dtype_size;
     bool instruction_tile_span_multiple_swizzle = inner_size > swizzle_size;
     bool span_uneven_swizzle = inner_tile_size % swizzle_size != 0 &&
         swizzle_size % inner_tile_size != 0;
 
     if (instruction_tile_span_multiple_swizzle && span_uneven_swizzle) {
-      GTEST_SKIP() << skip_reason;
+      return skip_reason;
     }
   }
 
@@ -947,15 +943,35 @@ TEST_P(HopperSS, MultipleTile) {
         ? getN(macro)
         : getK(macro);
     int64_t inner_size = num_tiles * inner_tile_size;
-    int64_t swizzle_size = getBytesFromSwizzle(swizzle_b) / dataTypeSize(dtype);
+    int64_t swizzle_size = getBytesFromSwizzle(swizzle_b) / dtype_size;
     bool instruction_tile_span_multiple_swizzle = inner_size > swizzle_size;
     bool span_uneven_swizzle = inner_tile_size % swizzle_size != 0 &&
         swizzle_size % inner_tile_size != 0;
 
     if (instruction_tile_span_multiple_swizzle && span_uneven_swizzle) {
-      GTEST_SKIP() << skip_reason;
+      return skip_reason;
     }
   }
+  return nullptr;
+}
+
+TEST_P(HopperSS, MultipleTile) {
+  Fusion fusion;
+  FusionGuard fg(&fusion);
+
+  constexpr int64_t num_tiles = 2;
+
+  const char* skip_reason = skipIfSwizzleNotCompatibleWithTiling(
+      num_tiles, layout, macro, swizzle_a, swizzle_b);
+  if (skip_reason) {
+    GTEST_SKIP() << skip_reason;
+  }
+
+  auto shapes = matmulAtInputShape3DSS(
+      num_tiles * getM(macro),
+      num_tiles * getN(macro),
+      num_tiles * getK(macro),
+      layout);
 
   auto tv0 = makeConcreteTensor(shapes.first, dtype);
   auto tv1 = makeConcreteTensor(shapes.second, dtype);
@@ -1262,6 +1278,179 @@ TEST_P(Blackwell1CTAM128SS, SingleTile) {
   NVF_CHECK(at::allclose(cg_outputs[0].as<at::Tensor>(), tref, 1e-5, 1e-5));
 }
 
+TEST_P(Blackwell1CTAM128SS, MultipleTile) {
+  Fusion fusion;
+  FusionGuard fg(&fusion);
+
+  constexpr int64_t num_tiles = 2;
+
+  const int64_t num_cols = getN(macro) * num_tiles * num_tiles;
+
+  if (num_cols > 512) {
+    GTEST_SKIP() << "Skipping test due to excessive number of columns";
+  }
+
+  const char* skip_reason = skipIfSwizzleNotCompatibleWithTiling(
+      num_tiles, layout, macro, swizzle_a, swizzle_b);
+  if (skip_reason) {
+    GTEST_SKIP() << skip_reason;
+  }
+
+  // TODO: currently we are not doing tiling on K dimension because we are not
+  // generating the correct predicate correctly. The correct way to do MMA is to
+  // let one thread submit the instruction, but currently all threads are
+  // submitting the instruction, so we are doing duplicate computation. If there
+  // are multiple tiles on K dimension, the result will be incorrect because we
+  // are adding the same number multiple times.
+  auto shapes = matmulAtInputShape3DSS(
+      num_tiles * getM(macro),
+      num_tiles * getN(macro),
+      getK(macro), // TODO: num_tiles * getK(macro),
+      layout);
+
+  auto tv0 = makeConcreteTensor(shapes.first, dtype);
+  auto tv1 = makeConcreteTensor(shapes.second, dtype);
+  fusion.addInput(tv0);
+  fusion.addInput(tv1);
+
+  // Just doing a gmem->smem copy
+  tv0 = set(tv0);
+  tv0->setMemoryType(MemoryType::Shared);
+  tv1 = set(tv1);
+  tv1->setMemoryType(MemoryType::Shared);
+
+  int axes = 0;
+  switch (layout) {
+    case MmaLayout::NT:
+      axes = 0;
+      break;
+    case MmaLayout::TT:
+    case MmaLayout::NN:
+      axes = 1;
+      break;
+    case MmaLayout::TN:
+      axes = 2;
+      break;
+    default:
+      NVF_ERROR("Invalid layout");
+  }
+  auto tv2 = fusedMultiplySum(tv0, tv1, {axes});
+
+  // Reorder the accumulator as [M, N, K]
+  switch (layout) {
+    case MmaLayout::TT:
+      // [M, K, N] -> [M, N, K]
+      tv2->reorder({{-2, -1}});
+      break;
+    case MmaLayout::TN:
+      // [M, N, K]
+      break;
+    case MmaLayout::NT:
+      // [K, M, N] -> [M, N, K]
+      tv2->reorder({{-3, -1}});
+      break;
+    case MmaLayout::NN:
+      // [N, K, M] -> [M, N, K]
+      tv2->reorder({{-1, -3}});
+      break;
+    default:
+      NVF_ERROR("Invalid layout");
+  }
+  tv2->commitLeafToLogical();
+
+  fusion.addOutput(tv2);
+
+  auto mma_ops = ir_utils::getOpsOfType<MmaOp>(&fusion);
+  NVF_CHECK(
+      1 == mma_ops.size(),
+      "Invalid number of MmaOp instances in fusion definition, expected 1, got ",
+      mma_ops.size());
+  mma_ops.front()->setMacro(macro);
+
+  auto mma_result = tv2->cacheBefore();
+  mma_result->setMemoryType(MemoryType::Tensor);
+
+  auto register_result = tv2->cacheBefore();
+  register_result->definition()->as<LoadStoreOp>()->setOpType(
+      LoadStoreOpType::LdTMem);
+
+  // Bring related dims to innermost, that is:
+  // - Reorder tv0 as [1, M, K] or [1, K, M]
+  // - Reorder tv1 as [1, N, K] or [1, K, N]
+  matmul_utils::moveInnerBroadcastLeft(tv0);
+  matmul_utils::moveInnerBroadcastLeft(tv1);
+
+  // Note that we do not split tv0 and tv1 by tile. We just directly swizzle the
+  // entire CTA. This means, the location of core matrices of each instruction
+  // will be discontiguous. For example, something like this:
+  //  it0 it0 it1 it1
+  //  it0 it0 it1 it1
+  //  it2 it2 it3 it3
+  //  it2 it2 it3 it3
+  // where itX refers to "instruction tile X".
+  //
+  // Being discontiguous is not a problem, as long as different core matrices
+  // are stored in a strided manner, and we will be able to infer the correct
+  // stride.
+  tv0->applyMmaSwizzle(swizzle_a);
+  tv1->applyMmaSwizzle(swizzle_b);
+
+  naivelyParallelize(tv0);
+  naivelyParallelize(tv1);
+
+  // Schedule the MMA result
+  {
+    // Split by tile
+    mma_result->split(-3, getM(macro));
+    mma_result->split(-2, getN(macro));
+    mma_result->split(-1, getK(macro));
+    // [Mo, Mi, No, Ni, Ko, Ki] -> [Mo, No, Ko, Mi, Ni, Ki]
+    mma_result->reorder({{-5, -3}, {-3, -2}});
+    mma_result->axis(-1)->parallelize(ParallelType::Mma);
+    mma_result->axis(-2)->parallelize(ParallelType::Mma);
+    mma_result->axis(-3)->parallelize(ParallelType::Mma);
+    // Set the allocation domain
+    AbstractTensor s(mma_result->getLoopDomain());
+    // [Mo, No, Ko, Mi, Ni, Ki] -> [Mi, |, Mo, No, Ko, Ni, Ki]
+    s.reorder({{3, 0}});
+    mma_result->setAllocationDomain(s.as<IterDomain*>(), true);
+    mma_result->setTMemDimSepPos(1);
+  }
+
+  // Schedule TMem load and gmem store
+  for (auto tv : {register_result, tv2}) {
+    // Split by tile
+    tv->split(-2, getM(macro));
+    tv->split(-1, getN(macro));
+    // [Mo, Mi, No, Ni] -> [Mo, No, Mi, Ni]
+    tv->reorder({{-3, -2}});
+    tv->axis(-2)->parallelize(ParallelType::TIDx);
+  }
+  if (std::has_single_bit(static_cast<uint64_t>(getN(macro)))) {
+    // We can only vectorize if N is a power of 2
+    register_result->axis(-1)->parallelize(ParallelType::Vectorize);
+  }
+
+  inlineMost();
+
+  auto inputs = matmulAtInput3DSS(
+      num_tiles * getM(macro),
+      num_tiles * getN(macro),
+      getK(macro), // TODO: num_tiles * getK(macro),
+      layout,
+      data_type_to_aten(dtype));
+
+  KernelExecutor ke;
+  ke.compile(
+      &fusion, {inputs.first, inputs.second}, LaunchParams(), matmul_cparams);
+  auto cg_outputs = ke.run({inputs.first, inputs.second});
+  auto tref = atMatmul(
+      inputs.first.squeeze().to(at::kFloat),
+      inputs.second.squeeze().to(at::kFloat),
+      layout);
+  NVF_CHECK(at::allclose(cg_outputs[0].as<at::Tensor>(), tref, 1e-5, 1e-5));
+}
+
 INSTANTIATE_TEST_SUITE_P(
     MmaTest,
     Blackwell1CTAM128SS,


Reviews:
user: login: zasdfgbnm
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-15T06:55:30Z
user: login: zasdfgbnm
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-15T06:56:30Z
user: login: jacobhinkle
body: LGTM
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-04-15T15:16:30Z
user: login: zasdfgbnm
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-15T17:16:54Z
user: login: zasdfgbnm
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-15T17:17:37Z

Review comments:
diff_hunk: @@ -908,18 +909,13 @@ TEST_P(HopperSS, SingleTileTransposed) {
   NVF_CHECK(at::allclose(cg_outputs[0].as<at::Tensor>(), tref, 1e-5, 1e-5));
 }
 
-TEST_P(HopperSS, MultipleTile) {
-  Fusion fusion;
-  FusionGuard fg(&fusion);
-
-  constexpr int64_t num_tiles = 2;
-
-  auto shapes = matmulAtInputShape3DSS(
-      num_tiles * getM(macro),
-      num_tiles * getN(macro),
-      num_tiles * getK(macro),
-      layout);
-
+const char* skipIfSwizzleNotCompatibleWithTiling(
path: tests/cpp/test_mma.cpp
user: login: zasdfgbnm
body: Check pulled from `HopperSS.MultipleTile`
created_at: 2025-04-15T06:55:30Z
author_association: COLLABORATOR
diff_hunk: @@ -1262,6 +1278,179 @@ TEST_P(Blackwell1CTAM128SS, SingleTile) {
   NVF_CHECK(at::allclose(cg_outputs[0].as<at::Tensor>(), tref, 1e-5, 1e-5));
 }
 
+TEST_P(Blackwell1CTAM128SS, MultipleTile) {
path: tests/cpp/test_mma.cpp
user: login: zasdfgbnm
body: This is pretty similar to `HopperSS.MultipleTile`
created_at: 2025-04-15T06:56:30Z
author_association: COLLABORATOR
diff_hunk: @@ -2578,6 +2578,41 @@ Val* getOuterStrideBytes(TensorView* tv, const MmaOp* mma) {
   return SimplifyingIrBuilder::mulExpr(stride, dataTypeSize(tv->dtype()));
 }
 
+namespace {
+
+Val* indexBlackwellMmaOutput(
+    const MmaOp* mma,
+    const std::vector<ForLoop*>& for_loops) {
+  TensorView* tmem_tv = mma->out()->as<TensorView>();
+  NVF_ERROR(tmem_tv->getMemoryType() == MemoryType::Tensor, "Invalid tmem_tv");
+  const auto& tmem_info = GpuLower::current()->tmemInfo();
+  const auto& tensor_indexer = GpuLower::current()->tensorIndexer();
+
+  const std::vector<IterDomain*>& column_allocation_domain =
+      tmem_info.allocation.getTVInfo(tmem_tv).column_allocation;
+
+  std::vector<Val*> indices = tensor_indexer.getIndexFor(
+      mma, true, column_allocation_domain, for_loops);
+  Val* stride = tmem_tv->fusion()->oneVal();
+  Val* column_index = tmem_tv->fusion()->zeroVal();
+  for (const auto& [id, idx] :
+       zip(std::ranges::views::reverse(column_allocation_domain),
+           std::ranges::views::reverse(indices))) {
path: csrc/device_lower/pass/index.cpp
user: login: jacobhinkle
body: Just curious, does this work?
```suggestion
  for (const auto& [id, idx] :
       std::ranges::views::reverse(zip(column_allocation_domain,indices))) {
```
created_at: 2025-04-15T15:04:16Z
author_association: COLLABORATOR
diff_hunk: @@ -1262,6 +1278,179 @@ TEST_P(Blackwell1CTAM128SS, SingleTile) {
   NVF_CHECK(at::allclose(cg_outputs[0].as<at::Tensor>(), tref, 1e-5, 1e-5));
 }
 
+TEST_P(Blackwell1CTAM128SS, MultipleTile) {
+  Fusion fusion;
+  FusionGuard fg(&fusion);
+
+  constexpr int64_t num_tiles = 2;
+
+  const int64_t num_cols = getN(macro) * num_tiles * num_tiles;
+
+  if (num_cols > 512) {
+    GTEST_SKIP() << "Skipping test due to excessive number of columns";
+  }
+
+  const char* skip_reason = skipIfSwizzleNotCompatibleWithTiling(
+      num_tiles, layout, macro, swizzle_a, swizzle_b);
+  if (skip_reason) {
+    GTEST_SKIP() << skip_reason;
+  }
+
+  // TODO: currently we are not doing tiling on K dimension because we are not
+  // generating the correct predicate correctly. The correct way to do MMA is to
+  // let one thread submit the instruction, but currently all threads are
+  // submitting the instruction, so we are doing duplicate computation. If there
+  // are multiple tiles on K dimension, the result will be incorrect because we
+  // are adding the same number multiple times.
+  auto shapes = matmulAtInputShape3DSS(
+      num_tiles * getM(macro),
+      num_tiles * getN(macro),
+      getK(macro), // TODO: num_tiles * getK(macro),
+      layout);
+
+  auto tv0 = makeConcreteTensor(shapes.first, dtype);
+  auto tv1 = makeConcreteTensor(shapes.second, dtype);
+  fusion.addInput(tv0);
+  fusion.addInput(tv1);
+
+  // Just doing a gmem->smem copy
+  tv0 = set(tv0);
+  tv0->setMemoryType(MemoryType::Shared);
+  tv1 = set(tv1);
+  tv1->setMemoryType(MemoryType::Shared);
+
+  int axes = 0;
+  switch (layout) {
+    case MmaLayout::NT:
+      axes = 0;
+      break;
+    case MmaLayout::TT:
+    case MmaLayout::NN:
+      axes = 1;
+      break;
+    case MmaLayout::TN:
+      axes = 2;
+      break;
+    default:
+      NVF_ERROR("Invalid layout");
+  }
+  auto tv2 = fusedMultiplySum(tv0, tv1, {axes});
+
+  // Reorder the accumulator as [M, N, K]
path: tests/cpp/test_mma.cpp
user: login: jacobhinkle
body: You may want to keep `commitLeafToLogical` here but if not, you could permute the inputs before the fusedMultiplySum and call `scheduler_utils::scheduleInputToSkipIntermediates` on each of them. That's what we do in the Hopper scheduler these days and likely what we'd want to do with the Blackwell scheduler so that we can standardize on MNK order. Anyway, no need to change it here since it's not the subject of this test, just an FYI.
created_at: 2025-04-15T15:11:31Z
author_association: COLLABORATOR
diff_hunk: @@ -1262,6 +1278,179 @@ TEST_P(Blackwell1CTAM128SS, SingleTile) {
   NVF_CHECK(at::allclose(cg_outputs[0].as<at::Tensor>(), tref, 1e-5, 1e-5));
 }
 
+TEST_P(Blackwell1CTAM128SS, MultipleTile) {
+  Fusion fusion;
+  FusionGuard fg(&fusion);
+
+  constexpr int64_t num_tiles = 2;
+
+  const int64_t num_cols = getN(macro) * num_tiles * num_tiles;
+
+  if (num_cols > 512) {
+    GTEST_SKIP() << "Skipping test due to excessive number of columns";
+  }
+
+  const char* skip_reason = skipIfSwizzleNotCompatibleWithTiling(
+      num_tiles, layout, macro, swizzle_a, swizzle_b);
+  if (skip_reason) {
+    GTEST_SKIP() << skip_reason;
+  }
+
+  // TODO: currently we are not doing tiling on K dimension because we are not
+  // generating the correct predicate correctly. The correct way to do MMA is to
+  // let one thread submit the instruction, but currently all threads are
+  // submitting the instruction, so we are doing duplicate computation. If there
+  // are multiple tiles on K dimension, the result will be incorrect because we
+  // are adding the same number multiple times.
+  auto shapes = matmulAtInputShape3DSS(
+      num_tiles * getM(macro),
+      num_tiles * getN(macro),
+      getK(macro), // TODO: num_tiles * getK(macro),
+      layout);
+
+  auto tv0 = makeConcreteTensor(shapes.first, dtype);
+  auto tv1 = makeConcreteTensor(shapes.second, dtype);
+  fusion.addInput(tv0);
+  fusion.addInput(tv1);
+
+  // Just doing a gmem->smem copy
+  tv0 = set(tv0);
+  tv0->setMemoryType(MemoryType::Shared);
+  tv1 = set(tv1);
+  tv1->setMemoryType(MemoryType::Shared);
+
+  int axes = 0;
+  switch (layout) {
+    case MmaLayout::NT:
+      axes = 0;
+      break;
+    case MmaLayout::TT:
+    case MmaLayout::NN:
+      axes = 1;
+      break;
+    case MmaLayout::TN:
+      axes = 2;
+      break;
+    default:
+      NVF_ERROR("Invalid layout");
+  }
+  auto tv2 = fusedMultiplySum(tv0, tv1, {axes});
+
+  // Reorder the accumulator as [M, N, K]
+  switch (layout) {
+    case MmaLayout::TT:
+      // [M, K, N] -> [M, N, K]
+      tv2->reorder({{-2, -1}});
+      break;
+    case MmaLayout::TN:
+      // [M, N, K]
+      break;
+    case MmaLayout::NT:
+      // [K, M, N] -> [M, N, K]
+      tv2->reorder({{-3, -1}});
+      break;
+    case MmaLayout::NN:
+      // [N, K, M] -> [M, N, K]
+      tv2->reorder({{-1, -3}});
+      break;
+    default:
+      NVF_ERROR("Invalid layout");
+  }
+  tv2->commitLeafToLogical();
+
+  fusion.addOutput(tv2);
+
+  auto mma_ops = ir_utils::getOpsOfType<MmaOp>(&fusion);
+  NVF_CHECK(
+      1 == mma_ops.size(),
+      "Invalid number of MmaOp instances in fusion definition, expected 1, got ",
+      mma_ops.size());
+  mma_ops.front()->setMacro(macro);
+
+  auto mma_result = tv2->cacheBefore();
+  mma_result->setMemoryType(MemoryType::Tensor);
+
+  auto register_result = tv2->cacheBefore();
+  register_result->definition()->as<LoadStoreOp>()->setOpType(
+      LoadStoreOpType::LdTMem);
+
+  // Bring related dims to innermost, that is:
+  // - Reorder tv0 as [1, M, K] or [1, K, M]
+  // - Reorder tv1 as [1, N, K] or [1, K, N]
+  matmul_utils::moveInnerBroadcastLeft(tv0);
+  matmul_utils::moveInnerBroadcastLeft(tv1);
+
+  // Note that we do not split tv0 and tv1 by tile. We just directly swizzle the
+  // entire CTA. This means, the location of core matrices of each instruction
+  // will be discontiguous. For example, something like this:
+  //  it0 it0 it1 it1
+  //  it0 it0 it1 it1
+  //  it2 it2 it3 it3
+  //  it2 it2 it3 it3
+  // where itX refers to "instruction tile X".
+  //
+  // Being discontiguous is not a problem, as long as different core matrices
+  // are stored in a strided manner, and we will be able to infer the correct
+  // stride.
+  tv0->applyMmaSwizzle(swizzle_a);
+  tv1->applyMmaSwizzle(swizzle_b);
+
+  naivelyParallelize(tv0);
+  naivelyParallelize(tv1);
+
+  // Schedule the MMA result
+  {
+    // Split by tile
+    mma_result->split(-3, getM(macro));
+    mma_result->split(-2, getN(macro));
+    mma_result->split(-1, getK(macro));
+    // [Mo, Mi, No, Ni, Ko, Ki] -> [Mo, No, Ko, Mi, Ni, Ki]
+    mma_result->reorder({{-5, -3}, {-3, -2}});
+    mma_result->axis(-1)->parallelize(ParallelType::Mma);
+    mma_result->axis(-2)->parallelize(ParallelType::Mma);
+    mma_result->axis(-3)->parallelize(ParallelType::Mma);
+    // Set the allocation domain
+    AbstractTensor s(mma_result->getLoopDomain());
+    // [Mo, No, Ko, Mi, Ni, Ki] -> [Mi, |, Mo, No, Ko, Ni, Ki]
+    s.reorder({{3, 0}});
+    mma_result->setAllocationDomain(s.as<IterDomain*>(), true);
+    mma_result->setTMemDimSepPos(1);
+  }
+
+  // Schedule TMem load and gmem store
+  for (auto tv : {register_result, tv2}) {
+    // Split by tile
+    tv->split(-2, getM(macro));
+    tv->split(-1, getN(macro));
+    // [Mo, Mi, No, Ni] -> [Mo, No, Mi, Ni]
+    tv->reorder({{-3, -2}});
+    tv->axis(-2)->parallelize(ParallelType::TIDx);
+  }
+  if (std::has_single_bit(static_cast<uint64_t>(getN(macro)))) {
path: tests/cpp/test_mma.cpp
user: login: jacobhinkle
body: TIL about `std::has_single_bit` :+1: 
created_at: 2025-04-15T15:16:10Z
author_association: COLLABORATOR
diff_hunk: @@ -2578,6 +2578,41 @@ Val* getOuterStrideBytes(TensorView* tv, const MmaOp* mma) {
   return SimplifyingIrBuilder::mulExpr(stride, dataTypeSize(tv->dtype()));
 }
 
+namespace {
+
+Val* indexBlackwellMmaOutput(
+    const MmaOp* mma,
+    const std::vector<ForLoop*>& for_loops) {
+  TensorView* tmem_tv = mma->out()->as<TensorView>();
+  NVF_ERROR(tmem_tv->getMemoryType() == MemoryType::Tensor, "Invalid tmem_tv");
+  const auto& tmem_info = GpuLower::current()->tmemInfo();
+  const auto& tensor_indexer = GpuLower::current()->tensorIndexer();
+
+  const std::vector<IterDomain*>& column_allocation_domain =
+      tmem_info.allocation.getTVInfo(tmem_tv).column_allocation;
+
+  std::vector<Val*> indices = tensor_indexer.getIndexFor(
+      mma, true, column_allocation_domain, for_loops);
+  Val* stride = tmem_tv->fusion()->oneVal();
+  Val* column_index = tmem_tv->fusion()->zeroVal();
+  for (const auto& [id, idx] :
+       zip(std::ranges::views::reverse(column_allocation_domain),
+           std::ranges::views::reverse(indices))) {
path: csrc/device_lower/pass/index.cpp
user: login: zasdfgbnm
body: It works. There used to be a bug in zip making this not working, but that bug has been fixed.
created_at: 2025-04-15T17:16:54Z
author_association: COLLABORATOR
diff_hunk: @@ -1262,6 +1278,179 @@ TEST_P(Blackwell1CTAM128SS, SingleTile) {
   NVF_CHECK(at::allclose(cg_outputs[0].as<at::Tensor>(), tref, 1e-5, 1e-5));
 }
 
+TEST_P(Blackwell1CTAM128SS, MultipleTile) {
+  Fusion fusion;
+  FusionGuard fg(&fusion);
+
+  constexpr int64_t num_tiles = 2;
+
+  const int64_t num_cols = getN(macro) * num_tiles * num_tiles;
+
+  if (num_cols > 512) {
+    GTEST_SKIP() << "Skipping test due to excessive number of columns";
+  }
+
+  const char* skip_reason = skipIfSwizzleNotCompatibleWithTiling(
+      num_tiles, layout, macro, swizzle_a, swizzle_b);
+  if (skip_reason) {
+    GTEST_SKIP() << skip_reason;
+  }
+
+  // TODO: currently we are not doing tiling on K dimension because we are not
+  // generating the correct predicate correctly. The correct way to do MMA is to
+  // let one thread submit the instruction, but currently all threads are
+  // submitting the instruction, so we are doing duplicate computation. If there
+  // are multiple tiles on K dimension, the result will be incorrect because we
+  // are adding the same number multiple times.
+  auto shapes = matmulAtInputShape3DSS(
+      num_tiles * getM(macro),
+      num_tiles * getN(macro),
+      getK(macro), // TODO: num_tiles * getK(macro),
+      layout);
+
+  auto tv0 = makeConcreteTensor(shapes.first, dtype);
+  auto tv1 = makeConcreteTensor(shapes.second, dtype);
+  fusion.addInput(tv0);
+  fusion.addInput(tv1);
+
+  // Just doing a gmem->smem copy
+  tv0 = set(tv0);
+  tv0->setMemoryType(MemoryType::Shared);
+  tv1 = set(tv1);
+  tv1->setMemoryType(MemoryType::Shared);
+
+  int axes = 0;
+  switch (layout) {
+    case MmaLayout::NT:
+      axes = 0;
+      break;
+    case MmaLayout::TT:
+    case MmaLayout::NN:
+      axes = 1;
+      break;
+    case MmaLayout::TN:
+      axes = 2;
+      break;
+    default:
+      NVF_ERROR("Invalid layout");
+  }
+  auto tv2 = fusedMultiplySum(tv0, tv1, {axes});
+
+  // Reorder the accumulator as [M, N, K]
path: tests/cpp/test_mma.cpp
user: login: zasdfgbnm
body: Thanks for the information. I will look into it when I started working on scheduler.
created_at: 2025-04-15T17:17:37Z
author_association: COLLABORATOR

PR comments:
user: login: github-actions[bot]
created_at: 2025-04-10T21:04:31Z
body: <!-- llm action combined persistent comment, dont modify this line -->

*Review updated until commit https://github.com/NVIDIA/Fuser/commit/d2a936f8723a6922ccde2cd4648a29b8d005fb26*


### **Description**
- Added indexing for MMA result in `index.cpp`

- Replaced `nanosleep` with a larger value in `inline_ptx.cpp`

- Implemented multiple-tile MMA unit tests in `test_mma.cpp`

- Added helper function `skipIfSwizzleNotCompatibleWithTiling` in `test_mma.cpp`


___



### **Changes walkthrough** 📝
<table><thead><tr><th></th><th align="left">Relevant files</th></tr></thead><tbody><tr><td><strong>Enhancement</strong></td><td><table>
<tr>
  <td>
    <details>
      <summary><strong>index.cpp</strong><dd><code>Implement MMA result indexing</code>&nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; </dd></summary>
<hr>

csrc/device_lower/pass/index.cpp

<li>Added namespace with <code>indexBlackwellMmaOutput</code> function for indexing MMA <br>result<br> <li> Updated <code>IndexLowering::handle</code> to use <code>indexBlackwellMmaOutput</code> for <br>Tensor memory type


</details>


  </td>
  <td><a href="https://github.com/NVIDIA/Fuser/pull/4231/files#diff-60e927bdf6f0ab7d9206208a2515d01c1a0b2a355253224d8ac63a4a32a4aa69">+35/-4</a>&nbsp; &nbsp; </td>

</tr>

<tr>
  <td>
    <details>
      <summary><strong>inline_ptx.cpp</strong><dd><code>Increase nanosleep value</code>&nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; </dd></summary>
<hr>

csrc/device_lower/pass/inline_ptx.cpp

- Increased the value in `nanosleep` from 100000 to 100000000


</details>


  </td>
  <td><a href="https://github.com/NVIDIA/Fuser/pull/4231/files#diff-e792cd4cdec50993f43539170c6810414339141e65ca216487582e423259a86b">+2/-1</a>&nbsp; &nbsp; &nbsp; </td>

</tr>
</table></td></tr><tr><td><strong>Tests</strong></td><td><table>
<tr>
  <td>
    <details>
      <summary><strong>test_mma.cpp</strong><dd><code>Implement multiple-tile MMA unit tests</code>&nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; </dd></summary>
<hr>

tests/cpp/test_mma.cpp

<li>Added <code>skipIfSwizzleNotCompatibleWithTiling</code> function to check swizzle <br>compatibility<br> <li> Implemented <code>MultipleTile</code> test for Blackwell1CTAM128SS<br> <li> Updated <code>MultipleTile</code> test for HopperSS to use <br><code>skipIfSwizzleNotCompatibleWithTiling</code>


</details>


  </td>
  <td><a href="https://github.com/NVIDIA/Fuser/pull/4231/files#diff-25d36694d3c86c362f1f90c453f9c821f4c989ea3c8bddf4670201c62c82d9e2">+205/-16</a></td>

</tr>
</table></td></tr></tr></tbody></table>

---

### PR Reviewer Guide 🔍

Here are some key observations to aid the review process:

<table>
<tr><td>🧪&nbsp;<strong>PR contains tests</strong></td></tr>
<tr><td>⚡&nbsp;<strong>Recommended focus areas for review</strong><br><br>

<details><summary><a href='https://github.com/NVIDIA/Fuser/pull/4231/files#diff-60e927bdf6f0ab7d9206208a2515d01c1a0b2a355253224d8ac63a4a32a4aa69R2583-R2611'><strong>Predicate Issue</strong></a>

The predicate of MMA is incorrect. The current implementation allows all threads to submit jobs, which can lead to incorrect results if the problem size exceeds the `K` of the instruction.</summary>

```c++
Val* indexBlackwellMmaOutput(
    const MmaOp* mma,
    const std::vector<ForLoop*>& for_loops) {
  TensorView* tmem_tv = mma->out()->as<TensorView>();
  NVF_ERROR(tmem_tv->getMemoryType() == MemoryType::Tensor, "Invalid tmem_tv");
  const auto& tmem_info = GpuLower::current()->tmemInfo();
  const auto& tensor_indexer = GpuLower::current()->tensorIndexer();

  const std::vector<IterDomain*>& column_allocation_domain =
      tmem_info.allocation.getTVInfo(tmem_tv).column_allocation;

  std::vector<Val*> indices = tensor_indexer.getIndexFor(
      mma, true, column_allocation_domain, for_loops);
  Val* stride = tmem_tv->fusion()->oneVal();
  Val* column_index = tmem_tv->fusion()->zeroVal();
  for (const auto& [id, idx] :
       std::ranges::views::reverse(zip(column_allocation_domain, indices))) {
    column_index = SimplifyingIrBuilder::addExpr(
        column_index, SimplifyingIrBuilder::mulExpr(idx, stride));
    stride = SimplifyingIrBuilder::mulExpr(stride, id->extent());
  }

  column_index =
      SimplifyingIrBuilder::maybeCastExpr(DataType::UInt16, column_index);

  Val* index = SimplifyingIrBuilder::arrayExpr(std::vector<Val*>{
      mma->fusion()->zeroVal(DataType::UInt16), column_index});
  return GpuLower::current()->commonScalarMap().hoistScalar(index, for_loops);
}
```

</details>

<details><summary><a href='https://github.com/NVIDIA/Fuser/pull/4231/files#diff-e792cd4cdec50993f43539170c6810414339141e65ca216487582e423259a86bR400-R407'><strong>Sync Method</strong></a>

The PR uses `nanosleep` for synchronization, which is not a reliable method. It should be replaced with `mbarrier` for proper synchronization.</summary>

```c++
registerReplace(
    mma,
    IrBuilder::create<kir::Asm>(
        "nanosleep.u32",
        std::vector<Val*>{},
        std::vector<Val*>{
            IrBuilder::create<Val>(100000000, DataType::UInt32)},
        kir::Asm::Options{/*volatile=*/true}));
```

</details>

<details><summary><a href='https://github.com/NVIDIA/Fuser/pull/4231/files#diff-25d36694d3c86c362f1f90c453f9c821f4c989ea3c8bddf4670201c62c82d9e2R1281-R1452'><strong>Tiling on K Dimension</strong></a>

The PR does not handle tiling on the K dimension correctly. The predicate for submitting MMA instructions is incorrect, leading to duplicate computations if there are multiple tiles on the K dimension.</summary>

```c++
TEST_P(Blackwell1CTAM128SS, MultipleTile) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  constexpr int64_t num_tiles = 2;

  const int64_t num_cols = getN(macro) * num_tiles * num_tiles;

  if (num_cols > 512) {
    GTEST_SKIP() << "Skipping test due to excessive number of columns";
  }

  const char* skip_reason = skipIfSwizzleNotCompatibleWithTiling(
      num_tiles, layout, macro, swizzle_a, swizzle_b);
  if (skip_reason) {
    GTEST_SKIP() << skip_reason;
  }

  // TODO: currently we are not doing tiling on K dimension because we are not
  // generating the correct predicate correctly. The correct way to do MMA is to
  // let one thread submit the instruction, but currently all threads are
  // submitting the instruction, so we are doing duplicate computation. If there
  // are multiple tiles on K dimension, the result will be incorrect because we
  // are adding the same number multiple times.
  auto shapes = matmulAtInputShape3DSS(
      num_tiles * getM(macro),
      num_tiles * getN(macro),
      getK(macro), // TODO: num_tiles * getK(macro),
      layout);

  auto tv0 = makeConcreteTensor(shapes.first, dtype);
  auto tv1 = makeConcreteTensor(shapes.second, dtype);
  fusion.addInput(tv0);
  fusion.addInput(tv1);

  // Just doing a gmem->smem copy
  tv0 = set(tv0);
  tv0->setMemoryType(MemoryType::Shared);
  tv1 = set(tv1);
  tv1->setMemoryType(MemoryType::Shared);

  int axes = 0;
  switch (layout) {
    case MmaLayout::NT:
      axes = 0;
      break;
    case MmaLayout::TT:
    case MmaLayout::NN:
      axes = 1;
      break;
    case MmaLayout::TN:
      axes = 2;
      break;
    default:
      NVF_ERROR("Invalid layout");
  }
  auto tv2 = fusedMultiplySum(tv0, tv1, {axes});

  // Reorder the accumulator as [M, N, K]
  switch (layout) {
    case MmaLayout::TT:
      // [M, K, N] -> [M, N, K]
      tv2->reorder({{-2, -1}});
      break;
    case MmaLayout::TN:
      // [M, N, K]
      break;
    case MmaLayout::NT:
      // [K, M, N] -> [M, N, K]
      tv2->reorder({{-3, -1}});
      break;
    case MmaLayout::NN:
      // [N, K, M] -> [M, N, K]
      tv2->reorder({{-1, -3}});
      break;
    default:
      NVF_ERROR("Invalid layout");
  }
  tv2->commitLeafToLogical();

  fusion.addOutput(tv2);

  auto mma_ops = ir_utils::getOpsOfType<MmaOp>(&fusion);
  NVF_CHECK(
      1 == mma_ops.size(),
      "Invalid number of MmaOp instances in fusion definition, expected 1, got ",
      mma_ops.size());
  mma_ops.front()->setMacro(macro);

  auto mma_result = tv2->cacheBefore();
  mma_result->setMemoryType(MemoryType::Tensor);

  auto register_result = tv2->cacheBefore();
  register_result->definition()->as<LoadStoreOp>()->setOpType(
      LoadStoreOpType::LdTMem);

  // Bring related dims to innermost, that is:
  // - Reorder tv0 as [1, M, K] or [1, K, M]
  // - Reorder tv1 as [1, N, K] or [1, K, N]
  matmul_utils::moveInnerBroadcastLeft(tv0);
  matmul_utils::moveInnerBroadcastLeft(tv1);

  // Note that we do not split tv0 and tv1 by tile. We just directly swizzle the
  // entire CTA. This means, the location of core matrices of each instruction
  // will be discontiguous. For example, something like this:
  //  it0 it0 it1 it1
  //  it0 it0 it1 it1
  //  it2 it2 it3 it3
  //  it2 it2 it3 it3
  // where itX refers to "instruction tile X".
  //
  // Being discontiguous is not a problem, as long as different core matrices
  // are stored in a strided manner, and we will be able to infer the correct
  // stride.
  tv0->applyMmaSwizzle(swizzle_a);
  tv1->applyMmaSwizzle(swizzle_b);

  naivelyParallelize(tv0);
  naivelyParallelize(tv1);

  // Schedule the MMA result
  {
    // Split by tile
    mma_result->split(-3, getM(macro));
    mma_result->split(-2, getN(macro));
    mma_result->split(-1, getK(macro));
    // [Mo, Mi, No, Ni, Ko, Ki] -> [Mo, No, Ko, Mi, Ni, Ki]
    mma_result->reorder({{-5, -3}, {-3, -2}});
    mma_result->axis(-1)->parallelize(ParallelType::Mma);
    mma_result->axis(-2)->parallelize(ParallelType::Mma);
    mma_result->axis(-3)->parallelize(ParallelType::Mma);
    // Set the allocation domain
    AbstractTensor s(mma_result->getLoopDomain());
    // [Mo, No, Ko, Mi, Ni, Ki] -> [Mi, |, Mo, No, Ko, Ni, Ki]
    s.reorder({{3, 0}});
    mma_result->setAllocationDomain(s.as<IterDomain*>(), true);
    mma_result->setTMemDimSepPos(1);
  }

  // Schedule TMem load and gmem store
  for (auto tv : {register_result, tv2}) {
    // Split by tile
    tv->split(-2, getM(macro));
    tv->split(-1, getN(macro));
    // [Mo, Mi, No, Ni] -> [Mo, No, Mi, Ni]
    tv->reorder({{-3, -2}});
    tv->axis(-2)->parallelize(ParallelType::TIDx);
  }
  if (std::has_single_bit(static_cast<uint64_t>(getN(macro)))) {
    // We can only vectorize if N is a power of 2
    register_result->axis(-1)->parallelize(ParallelType::Vectorize);
  }

  inlineMost();

  auto inputs = matmulAtInput3DSS(
      num_tiles * getM(macro),
      num_tiles * getN(macro),
      getK(macro), // TODO: num_tiles * getK(macro),
      layout,
      data_type_to_aten(dtype));

  KernelExecutor ke;
  ke.compile(
      &fusion, {inputs.first, inputs.second}, LaunchParams(), matmul_cparams);
  auto cg_outputs = ke.run({inputs.first, inputs.second});
  auto tref = atMatmul(
      inputs.first.squeeze().to(at::kFloat),
      inputs.second.squeeze().to(at::kFloat),
      layout);
  NVF_CHECK(at::allclose(cg_outputs[0].as<at::Tensor>(), tref, 1e-5, 1e-5));
}
```

</details>

</td></tr>
</table>

author_association: NONE
user: login: zasdfgbnm
created_at: 2025-04-10T22:14:52Z
body: !test
author_association: COLLABORATOR
user: login: zasdfgbnm
created_at: 2025-04-15T00:00:57Z
body: !test
author_association: COLLABORATOR
user: login: zasdfgbnm
created_at: 2025-04-15T17:10:06Z
body: !build
author_association: COLLABORATOR