Repository: NVIDIA/Fuser
Pull request #4233: Use elect-sync predicate for Blackwell MMA
Author: login: zasdfgbnm
Created: 2025-04-10T23:16:47Z | Merged: 2025-04-15T22:12:09Z

Blackwell MMA should be launched by one thread, instead of every thread.

Before:

```C++
mma
nanosleep
epilogue
```

After:

```C++
__syncthreads();
if (elect-sync) {
  mma
}
__syncthreads();
nanosleep
epilogue
```


Diff:
diff --git a/csrc/device_lower/pass/inline_ptx.cpp b/csrc/device_lower/pass/inline_ptx.cpp
index fdf463a013d..704475a0152 100644
--- a/csrc/device_lower/pass/inline_ptx.cpp
+++ b/csrc/device_lower/pass/inline_ptx.cpp
@@ -394,17 +394,6 @@ class LowerToInlinePtx : public kir::ExprMutator {
                 /*volatile=*/true,
                 /*memory=*/false,
                 /*readable_outputs=*/{0}}));
-    // TODO: This is clearly a wrong way to sync, but as an intermediate step to
-    // enable incremental development, we use nanosleep to sync the mma. We
-    // should replace this with a correct sync method.
-    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}));
   }
 
   void handle(MmaOp* mma) final {
diff --git a/csrc/device_lower/pass/insert_syncs.cpp b/csrc/device_lower/pass/insert_syncs.cpp
index c5f9d79e91c..892eb59db9b 100644
--- a/csrc/device_lower/pass/insert_syncs.cpp
+++ b/csrc/device_lower/pass/insert_syncs.cpp
@@ -462,6 +462,20 @@ class ReadAfterWriteSyncs : public kir::ExprMutator {
         fill_async_mma_pipeline_ = true;
         // async mma pipeline has not been flushed yet.
         flush_async_mma_pipeline_ = false;
+      } else if (mma->isBlackwell()) {
+        // TODO: This is clearly a wrong way to sync, but as an intermediate
+        // step to enable incremental development, we use nanosleep to sync the
+        // mma. We should replace this with a correct sync method.
+        registerInsertBefore(expr, IrBuilder::create<kir::BlockSync>());
+        registerInsertAfter(
+            expr,
+            IrBuilder::create<kir::Asm>(
+                "nanosleep.u32",
+                std::vector<Val*>{},
+                std::vector<Val*>{
+                    IrBuilder::create<Val>(4000000000, DataType::UInt32)},
+                kir::Asm::Options{/*volatile=*/true}));
+        registerInsertAfter(expr, IrBuilder::create<kir::BlockSync>());
       }
     } else if (ir_utils::isCpAsyncBulkStore(expr)) {
       // Add a fence before TMA store so that writes in the generic proxy is
diff --git a/csrc/device_lower/pass/unroll.cpp b/csrc/device_lower/pass/unroll.cpp
index 2851e5a1a67..a68bd2869ff 100644
--- a/csrc/device_lower/pass/unroll.cpp
+++ b/csrc/device_lower/pass/unroll.cpp
@@ -157,8 +157,10 @@ void UnrollPass::dispatch(Expr* expr) {
       pred = IrBuilder::create<kir::Predicate>(PredicateType::Vectorize);
     }
 
-    // short-circuit: wrap tma expressions with elect sync predicate
-    if (ir_utils::isCpAsyncBulk(expr)) {
+    // short-circuit: wrap tma and blackwell mma expressions with elect sync
+    // predicate
+    if (ir_utils::isCpAsyncBulk(expr) ||
+        (expr->isA<MmaOp>() && expr->as<MmaOp>()->isBlackwell())) {
       // If we need a predicate, put expr inside an if then else
       auto elect_sync_pred = IrBuilder::create<kir::Predicate>(
           PredicateType::ElectSync, expr, thread_pred);
diff --git a/csrc/predicate_compute.cpp b/csrc/predicate_compute.cpp
index 9661ae403ed..49b48fbea1e 100644
--- a/csrc/predicate_compute.cpp
+++ b/csrc/predicate_compute.cpp
@@ -552,7 +552,10 @@ Val* createSingleExpressionElectSync(
     const std::vector<ForLoop*>& loops) {
   NVF_ERROR(pred->expr() != nullptr);
   NVF_ERROR(
-      ir_utils::isCpAsyncBulk(pred->expr()), "Limited to TMA expressions");
+      ir_utils::isCpAsyncBulk(pred->expr()) ||
+          (pred->expr()->isA<MmaOp>() &&
+           pred->expr()->as<MmaOp>()->isBlackwell()),
+      "Limited to TMA/Blackwell MMA expressions");
 
   TensorView* out_tv = ir_utils::getTvOutput(pred->expr());
   Val* zero = IrBuilder::create<Val>(0L, PrimDataType::UInt64);
@@ -566,7 +569,7 @@ Val* createSingleExpressionElectSync(
       continue;
     }
 
-    // Case 1: TMA expression uses ParallelDim to launch multiple
+    // Case 1: TMA/Blackwell MMA expression uses ParallelDim to launch multiple
     // operations simultaneously. Use parallel domain predicate if it
     // exists.
     auto pred_info_it = pred_map.find(pt);
@@ -577,8 +580,8 @@ Val* createSingleExpressionElectSync(
           SimplifyingIrBuilder::logicalAndExpr(parallel_dom_pred, tid_pred);
     }
 
-    // Case 2: ParallelDim is used by CTA but not the TMA expression.
-    // Select a single thread along ParallelDim.
+    // Case 2: ParallelDim is used by CTA but not the TMA/Blackwell MMA
+    // expression. Select a single thread along ParallelDim.
     bool is_tv_tid_parallelized = std::any_of(
         out_tv->domain()->loop().begin(),
         out_tv->domain()->loop().end(),
diff --git a/tests/cpp/test_mma.cpp b/tests/cpp/test_mma.cpp
index 44f4f69023c..52abfdb1cea 100644
--- a/tests/cpp/test_mma.cpp
+++ b/tests/cpp/test_mma.cpp
@@ -1305,7 +1305,7 @@ TEST_P(Blackwell1CTAM128SS, MultipleTile) {
   auto shapes = matmulAtInputShape3DSS(
       num_tiles * getM(macro),
       num_tiles * getN(macro),
-      getK(macro), // TODO: num_tiles * getK(macro),
+      num_tiles * getK(macro),
       layout);
 
   auto tv0 = makeConcreteTensor(shapes.first, dtype);
@@ -1410,10 +1410,14 @@ TEST_P(Blackwell1CTAM128SS, MultipleTile) {
     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);
+    // Note that we intentionally reorder both the loop domain and the
+    // allocation domain. The reason for reordering the loop domain is to
+    // prevent `inlineMost` from inlining the MmaOp. The reason for not
+    // wanting to inline is because we want to test the case where the
+    // TMem has multiple tiles, so that we know our indexing is correct.
+    mma_result->reorder({{3, 0}});
+    mma_result->setAllocationDomain(mma_result->getLoopDomain(), true);
     mma_result->setTMemDimSepPos(1);
   }
 
@@ -1436,7 +1440,7 @@ TEST_P(Blackwell1CTAM128SS, MultipleTile) {
   auto inputs = matmulAtInput3DSS(
       num_tiles * getM(macro),
       num_tiles * getN(macro),
-      getK(macro), // TODO: num_tiles * getK(macro),
+      num_tiles * getK(macro),
       layout,
       data_type_to_aten(dtype));
 


Reviews:
user: login: zasdfgbnm
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-15T18:03:15Z
user: login: zasdfgbnm
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-15T18:04:54Z
user: login: zasdfgbnm
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-15T18:10:47Z
user: login: copilot-pull-request-reviewer[bot]
body: ## Pull Request Overview

This PR refactors the launch logic for Blackwell MMA so that the MMA operation is executed by a single thread using an elect-sync predicate. Key changes include:
- Updating test cases in tests/cpp/test_mma.cpp to align with the new execution logic.
- Modifying predicate_compute.cpp and device lower passes (unroll.cpp, insert_syncs.cpp, inline_ptx.cpp) to support Blackwell MMA in addition to TMA expressions.

### Reviewed Changes

Copilot reviewed 5 out of 5 changed files in this pull request and generated 1 comment.

<details>
<summary>Show a summary per file</summary>

| File                                       | Description                                                           |
| ------------------------------------------ | --------------------------------------------------------------------- |
| tests/cpp/test_mma.cpp                     | Removed outdated TODO comments and updated getK multiplications.      |
| csrc/predicate_compute.cpp                 | Expanded error conditions and messages to include Blackwell MMA.      |
| csrc/device_lower/pass/unroll.cpp          | Adjusted predicate wrapping to incorporate Blackwell MMA expressions. |
| csrc/device_lower/pass/insert_syncs.cpp    | Added a temporary nanosleep-based sync for Blackwell MMA operations.    |
| csrc/device_lower/pass/inline_ptx.cpp        | Removed the nanosleep workaround previously used for synchronization. |
</details>





state: COMMENTED
author_association: NONE
submitted_at: 2025-04-15T18:12:44Z
user: login: rdspring1
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-15T21:14:02Z
user: login: rdspring1
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-04-15T21:15:17Z

Review comments:
diff_hunk: @@ -394,17 +394,6 @@ class LowerToInlinePtx : public kir::ExprMutator {
                 /*volatile=*/true,
                 /*memory=*/false,
                 /*readable_outputs=*/{0}}));
-    // TODO: This is clearly a wrong way to sync, but as an intermediate step to
path: csrc/device_lower/pass/inline_ptx.cpp
user: login: zasdfgbnm
body: Moved to `insert_syncs.cpp`
created_at: 2025-04-15T18:03:15Z
author_association: COLLABORATOR
diff_hunk: @@ -1410,10 +1410,14 @@ TEST_P(Blackwell1CTAM128SS, MultipleTile) {
     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);
+    // Note that we intentionally reorder both the loop domain and the
path: tests/cpp/test_mma.cpp
user: login: zasdfgbnm
body: I just realized that in my previous PR, we were fully inlining the `MmaOp`, as a result, the TMem only has one tile, and the computed index is always zero. This is not what we want to test. So I slightly modified here to prevent the inlining of `MmaOp`.
created_at: 2025-04-15T18:04:54Z
author_association: COLLABORATOR
diff_hunk: @@ -1305,7 +1305,7 @@ TEST_P(Blackwell1CTAM128SS, MultipleTile) {
   auto shapes = matmulAtInputShape3DSS(
       num_tiles * getM(macro),
       num_tiles * getN(macro),
-      getK(macro), // TODO: num_tiles * getK(macro),
+      num_tiles * getK(macro),
path: tests/cpp/test_mma.cpp
user: login: zasdfgbnm
body: After fixing predicate, we can finally tile on K.
created_at: 2025-04-15T18:10:47Z
author_association: COLLABORATOR
diff_hunk: @@ -462,6 +462,20 @@ class ReadAfterWriteSyncs : public kir::ExprMutator {
         fill_async_mma_pipeline_ = true;
         // async mma pipeline has not been flushed yet.
         flush_async_mma_pipeline_ = false;
+      } else if (mma->isBlackwell()) {
+        // TODO: This is clearly a wrong way to sync, but as an intermediate
+        // step to enable incremental development, we use nanosleep to sync the
+        // mma. We should replace this with a correct sync method.
+        registerInsertBefore(expr, IrBuilder::create<kir::BlockSync>());
+        registerInsertAfter(
+            expr,
+            IrBuilder::create<kir::Asm>(
+                "nanosleep.u32",
+                std::vector<Val*>{},
+                std::vector<Val*>{
+                    IrBuilder::create<Val>(4000000000, DataType::UInt32)},
path: csrc/device_lower/pass/insert_syncs.cpp
user: login: Copilot
body: Using a nanosleep-based synchronization for Blackwell MMA may lead to race conditions under heavy load. Consider replacing it with a proper synchronization method as soon as possible.
```suggestion
        // Proper synchronization for Blackwell MMA operations.
        // Use __syncwarp to synchronize threads within a warp.
        registerInsertBefore(expr, IrBuilder::create<kir::BlockSync>());
        registerInsertAfter(
            expr,
            IrBuilder::create<kir::Asm>(
                "__syncwarp();",
                std::vector<Val*>{},
                std::vector<Val*>(),
```
created_at: 2025-04-15T18:12:43Z
author_association: NONE
diff_hunk: @@ -462,6 +462,20 @@ class ReadAfterWriteSyncs : public kir::ExprMutator {
         fill_async_mma_pipeline_ = true;
         // async mma pipeline has not been flushed yet.
         flush_async_mma_pipeline_ = false;
+      } else if (mma->isBlackwell()) {
+        // TODO: This is clearly a wrong way to sync, but as an intermediate
+        // step to enable incremental development, we use nanosleep to sync the
+        // mma. We should replace this with a correct sync method.
+        registerInsertBefore(expr, IrBuilder::create<kir::BlockSync>());
+        registerInsertAfter(
+            expr,
+            IrBuilder::create<kir::Asm>(
+                "nanosleep.u32",
+                std::vector<Val*>{},
+                std::vector<Val*>{
+                    IrBuilder::create<Val>(4000000000, DataType::UInt32)},
path: csrc/device_lower/pass/insert_syncs.cpp
user: login: rdspring1
body: I wonder if you add `This is clearly a wrong way`, copilot will complain.
created_at: 2025-04-15T21:14:02Z
author_association: COLLABORATOR

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

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


### **Description**
- Use elect-sync predicate for Blackwell MMA operations

- Insert syncs before and after Blackwell MMA operations

- Update predicate checks to include Blackwell MMA

- Adjust test cases for Blackwell MMA with multiple tiles


___



### **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>inline_ptx.cpp</strong><dd><code>Remove incorrect MMA sync</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; </dd></summary>
<hr>

csrc/device_lower/pass/inline_ptx.cpp

- Removed nanosleep sync for MMA operations


</details>


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

</tr>

<tr>
  <td>
    <details>
      <summary><strong>insert_syncs.cpp</strong><dd><code>Insert syncs for Blackwell MMA</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; </dd></summary>
<hr>

csrc/device_lower/pass/insert_syncs.cpp

<li>Insert block sync before and after Blackwell MMA operations<br> <li> Add nanosleep sync for Blackwell MMA operations


</details>


  </td>
  <td><a href="https://github.com/NVIDIA/Fuser/pull/4233/files#diff-49bea61a8cde014ec0396c89d0654813136e0bd3ffe8b3a5974ee9ccf3a5fbb8">+14/-0</a>&nbsp; &nbsp; </td>

</tr>

<tr>
  <td>
    <details>
      <summary><strong>unroll.cpp</strong><dd><code>Update predicate for Blackwell MMA</code>&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/unroll.cpp

- Update predicate check to include Blackwell MMA


</details>


  </td>
  <td><a href="https://github.com/NVIDIA/Fuser/pull/4233/files#diff-3cf8b0f1d2f6ec1070cf8066665544a4fcdb524fd4f64be7764e24705fa30b24">+4/-2</a>&nbsp; &nbsp; &nbsp; </td>

</tr>

<tr>
  <td>
    <details>
      <summary><strong>predicate_compute.cpp</strong><dd><code>Update predicate checks for Blackwell MMA</code>&nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; </dd></summary>
<hr>

csrc/predicate_compute.cpp

<li>Update predicate checks to include Blackwell MMA<br> <li> Adjust comments to reflect changes


</details>


  </td>
  <td><a href="https://github.com/NVIDIA/Fuser/pull/4233/files#diff-8bc95e21925afa42b6322b589a82564d783e91d484a3fd0d4ec4e7f7bfb8f9ed">+7/-4</a>&nbsp; &nbsp; &nbsp; </td>

</tr>

<tr>
  <td>
    <details>
      <summary><strong>test_mma.cpp</strong><dd><code>Update Blackwell MMA test cases</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; </dd></summary>
<hr>

tests/cpp/test_mma.cpp

<li>Correct K dimension in test shapes<br> <li> Update comments and reordering logic for Blackwell MMA tests


</details>


  </td>
  <td><a href="https://github.com/NVIDIA/Fuser/pull/4233/files#diff-25d36694d3c86c362f1f90c453f9c821f4c989ea3c8bddf4670201c62c82d9e2">+9/-5</a>&nbsp; &nbsp; &nbsp; </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/4233/files#diff-49bea61a8cde014ec0396c89d0654813136e0bd3ffe8b3a5974ee9ccf3a5fbb8R469-R478'><strong>Incorrect Sync Method</strong></a>

The PR uses `nanosleep.u32` as a sync method for Blackwell MMA, which is not a correct approach. This should be replaced with a proper synchronization method.</summary>

```c++
registerInsertBefore(expr, IrBuilder::create<kir::BlockSync>());
registerInsertAfter(
    expr,
    IrBuilder::create<kir::Asm>(
        "nanosleep.u32",
        std::vector<Val*>{},
        std::vector<Val*>{
            IrBuilder::create<Val>(4000000000, DataType::UInt32)},
        kir::Asm::Options{/*volatile=*/true}));
registerInsertAfter(expr, IrBuilder::create<kir::BlockSync>());
```

</details>

<details><summary><a href='https://github.com/NVIDIA/Fuser/pull/4233/files#diff-e792cd4cdec50993f43539170c6810414339141e65ca216487582e423259a86bR394-R400'><strong>Unnecessary Code</strong></a>

The commented-out code for `nanosleep.u32` should be removed if it is no longer needed.</summary>

```c++
              /*volatile=*/true,
              /*memory=*/false,
              /*readable_outputs=*/{0}}));
}

void handle(MmaOp* mma) final {
  if (mma->isTuring() || mma->isAmpere()) {
```

</details>

<details><summary><a href='https://github.com/NVIDIA/Fuser/pull/4233/files#diff-25d36694d3c86c362f1f90c453f9c821f4c989ea3c8bddf4670201c62c82d9e2R1308-R1308'><strong>Test Code Update</strong></a>

The test code has been updated to use `num_tiles * getK(macro)`, but it is unclear if this change is correct and if it affects the test's validity.</summary>

```c++
num_tiles * getK(macro),
```

</details>

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

author_association: NONE
user: login: zasdfgbnm
created_at: 2025-04-10T23:58:20Z
body: !test
author_association: COLLABORATOR
user: login: zasdfgbnm
created_at: 2025-04-14T15:07:42Z
body: !test
author_association: COLLABORATOR
user: login: zasdfgbnm
created_at: 2025-04-15T18:02:37Z
body: !test
author_association: COLLABORATOR
user: login: zasdfgbnm
created_at: 2025-04-15T18:05:15Z
body: !test
author_association: COLLABORATOR