[Feature] Add read option to TMA store wait (#2300)

Add read option to TMA store wait

docs/programming_guides/instructions.md

@@ -235,7 +235,10 @@ Barriers, TMA, warp‑group
 - Parity ops: `T.mbarrier_wait_parity(barrier, parity)`, `T.mbarrier_arrive(barrier)`.
 - Expect tx: `T.mbarrier_expect_tx(...)`; sugar: `T.barrier_wait(id, parity=None)`.
 - TMA: `T.create_tma_descriptor(...)`, `T.tma_load(...)`,
-  `T.tma_store_arrive(...)`, `T.tma_store_wait(...)`.
+  `T.tma_store_arrive(...)`, `T.tma_store_wait(count=0, read=True)`.
+  `read=True` emits `cp.async.bulk.wait_group.read`; use `read=False` for
+  `cp.async.bulk.wait_group` when the wait must include destination writes
+  becoming visible.
 - Proxy/fences: `T.fence_proxy_async(...)`, `T.warpgroup_fence_operand(...)`.
 - Warp‑group: `T.warpgroup_arrive()`, `T.warpgroup_commit_batch()`,
   `T.warpgroup_wait(num_mma)`, `T.wait_wgmma(id)`.

src/backend/cuda/codegen/codegen_cuda.cc

@@ -2482,8 +2482,11 @@ void CodeGenTileLangCUDA::VisitExpr_(const CallNode *op, std::ostream &os) {
     print_extern_call_stmt("tl::tma_store_arrive");
   } else if (op->op.same_as(tl::tma_store_wait())) {
     int count = Downcast<IntImm>(op->args[0])->value;
+    bool read =
+        op->args.size() == 1 || Downcast<IntImm>(op->args[1])->value != 0;
     this->PrintIndent();
-    this->stream << "tl::tma_store_wait<" << count << ">();\n";
+    this->stream << "tl::tma_store_wait<" << count << ", "
+                 << (read ? "true" : "false") << ">();\n";
   } else if (op->op.same_as(tl::warpgroup_arrive())) {
     print_extern_call_stmt("tl::warpgroup_arrive");
   } else if (op->op.same_as(tl::warpgroup_commit_batch())) {

src/backend/cuda/codegen/codegen_cutedsl.cc

@@ -731,8 +731,11 @@ void CodeGenTileLangCuTeDSL::VisitExpr_(const CallNode *op,
     print_extern_call_stmt("tl.tma_store_arrive");
   } else if (op->op.same_as(tl::tma_store_wait())) {
     int count = Downcast<IntImm>(op->args[0])->value;
+    bool read =
+        op->args.size() == 1 || Downcast<IntImm>(op->args[1])->value != 0;
     PrintIndent();
-    stream << "tl.tma_store_wait(" << count << ")\n";
+    stream << "tl.tma_store_wait(" << count
+           << ", read=" << (read ? "True" : "False") << ")\n";
   } else if (op->op.same_as(tl::warpgroup_arrive())) {
     PrintIndent();
     stream << "tl.warpgroup_arrive()\n";

src/backend/cuda/op/atomic_add.cc

@@ -491,7 +491,7 @@ struct AtomicAdd {
     seq.push_back(tma_reduce);
     seq.push_back(Evaluate(Call(DataType::Handle(), tma_store_arrive(), {})));
     seq.push_back(Evaluate(Call(DataType::Handle(), tma_store_wait(),
-                                {IntImm(DataType::Int(32), 0)})));
+                                {IntImm(DataType::Int(32), 0), Bool(true)})));
     return IfThenElse(EQ(T.thread_var, T.thread_bounds->min),
                       SeqStmt(std::move(seq)));
   }

src/backend/cuda/op/copy.cc

@@ -1699,7 +1699,7 @@ Stmt Copy::LowerBulk(const CopyNode &op, const LowerArgs &T,
     seq.push_back(Evaluate(Call(DataType::Handle(), tma_store_arrive(), {})));
     if (!GetIsTmaCopy(op)) {
       seq.push_back(Evaluate(Call(DataType::Handle(), tma_store_wait(),
-                                  {IntImm(DataType::Int(32), 0)})));
+                                  {IntImm(DataType::Int(32), 0), Bool(true)})));
     }
     tma_copy = SeqStmt(std::move(seq));
   }
@@ -2034,7 +2034,7 @@ Stmt Copy::LowerBulk1D(const CopyNode &op, const LowerArgs &T,
     seq.push_back(Evaluate(Call(DataType::Handle(), tma_store_arrive(), {})));
     if (!GetIsTmaCopy(op)) {
       seq.push_back(Evaluate(Call(DataType::Handle(), tma_store_wait(),
-                                  {IntImm(DataType::Int(32), 0)})));
+                                  {IntImm(DataType::Int(32), 0), Bool(true)})));
     }
     tma_copy = SeqStmt(std::move(seq));
   }

src/op/builtin.cc

@@ -327,7 +327,7 @@ TIR_DEFINE_TL_BUILTIN(tma_store_arrive)
                                Integer(CallEffectKind::kOpaque));
 
 TIR_DEFINE_TL_BUILTIN(tma_store_wait)
-    .set_num_inputs(1)
+    .set_num_inputs(2)
     .set_attr<TCallEffectKind>("TCallEffectKind",
                                Integer(CallEffectKind::kOpaque));
 TIR_DEFINE_TL_BUILTIN(set_max_nreg)

src/tl_templates/cuda/barrier.h

@@ -142,8 +142,12 @@ TL_DEVICE void tma_store_arrive() {
   asm volatile("cp.async.bulk.commit_group;");
 }
 
-template <int Count> TL_DEVICE void tma_store_wait() {
-  asm volatile("cp.async.bulk.wait_group.read %0;" : : "n"(Count) : "memory");
+template <int Count, bool Read = true> TL_DEVICE void tma_store_wait() {
+  if constexpr (Read) {
+    asm volatile("cp.async.bulk.wait_group.read %0;" : : "n"(Count) : "memory");
+  } else {
+    asm volatile("cp.async.bulk.wait_group %0;" : : "n"(Count) : "memory");
+  }
 }
 
 TL_DEVICE void syncthreads_partial(uint64_t &smem_barrier) {

testing/python/language/test_tilelang_language_tma_copy.py

@@ -217,7 +217,7 @@ def matmul_tma_copy_store(
             # Store result: fragment -> shared -> global via T.tma_copy (no barrier needed)
             T.copy(C_local, C_shared)
             T.tma_copy(C_shared, C[by * block_M, bx * block_N])
-            T.tma_store_wait()
+            T.tma_store_wait(read=False)
 
     return main
 
@@ -273,7 +273,7 @@ def fp4_tma_copy_roundtrip(M=128, N=256, block_M=64, block_N=128):
             T.barrier_arrive(mbar)
             T.mbarrier_wait_parity(mbar, 0)
             T.tma_copy(A_shared, B[by * block_M, bx * block_N])
-            T.tma_store_wait()
+            T.tma_store_wait(read=False)
 
     return main
 

testing/python/language/test_tilelang_language_tma_gather_scatter.py

@@ -51,7 +51,7 @@ def gather_scatter_program(N: int, K: int, K_box: int, in_dtype: str = "float16"
             if T.shuffle_elect(128):
                 T.tma_scatter4(smem, Dst, 0, [r0, r1, r2, r3])
                 T.tma_store_arrive()
-            T.tma_store_wait(0)
+            T.tma_store_wait(0, read=False)
 
     return main
 
@@ -140,7 +140,7 @@ def gather_scatter_swizzled_program(N: int, K: int, K_box: int, swizzle_kind: st
             if T.shuffle_elect(128):
                 T.tma_scatter4(smem, Dst, 0, [r0, r1, r2, r3])
                 T.tma_store_arrive()
-            T.tma_store_wait(0)
+            T.tma_store_wait(0, read=False)
 
     return main
 

testing/python/language/test_tilelang_language_tma_store.py

@@ -126,6 +126,15 @@ def full_shape_tma_store_1d(dtype, threads):
     return main
 
 
+def tma_store_wait_read_option(read):
+    @T.prim_func
+    def main():
+        with T.Kernel(threads=128):
+            T.tma_store_wait(read=read)
+
+    return main
+
+
 def run_auto_tma_store_copy():
     M = N = 256
     block_M = block_N = 128
@@ -141,6 +150,7 @@ def run_auto_tma_store_copy():
     kernel_source = kernel.get_kernel_source()
     assert "tma_store_arrive" in kernel_source, "Expected auto tma_store_arrive in kernel source"
     assert "tma_store_wait" in kernel_source, "Expected auto tma_store_wait in kernel source"
+    assert "tl::tma_store_wait<0, true>();" in kernel_source, "Expected auto tma_store_wait to use default read wait"
 
     profiler = kernel.get_profiler()
 
@@ -161,6 +171,20 @@ def run_full_shape_tma_store_1d_codegen():
     assert "CUtensorMap" not in kernel_source, "Expected pointer-based 1D TMA store"
 
 
+def run_tma_store_wait_read_option_codegen():
+    default_kernel = tilelang.compile(
+        tma_store_wait_read_option(True),
+        pass_configs={tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True},
+    )
+    no_read_kernel = tilelang.compile(
+        tma_store_wait_read_option(False),
+        pass_configs={tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True},
+    )
+
+    assert "tl::tma_store_wait<0, true>();" in default_kernel.get_kernel_source()
+    assert "tl::tma_store_wait<0, false>();" in no_read_kernel.get_kernel_source()
+
+
 @tilelang.testing.requires_cuda
 @tilelang.testing.requires_cuda_compute_version_ge(9, 0)
 def test_tma_store_2_stages():
@@ -185,6 +209,12 @@ def test_plain_copy_full_shape_tma_store_uses_1d():
     run_full_shape_tma_store_1d_codegen()
 
 
+@tilelang.testing.requires_cuda
+@tilelang.testing.requires_cuda_compute_version_ge(9, 0)
+def test_tma_store_wait_read_option_codegen():
+    run_tma_store_wait_read_option_codegen()
+
+
 if __name__ == "__main__":
     # tilelang.testing.main()
     test_tma_store_2_stages()

tilelang/contrib/cutedsl/cpasync.py

@@ -238,10 +238,12 @@ def tma_store_arrive(*, loc=None, ip=None) -> None:
 def tma_store_wait(count: int, *, read=None, loc=None, ip=None) -> None:
     """
     Wait for TMA_STORE operations to complete.
-    Corresponds to PTX instruction: cp.async.bulk.wait_group.read <count>;
+    Corresponds to PTX instruction: cp.async.bulk.wait_group{.read} <count>;
 
     :param count: The number of outstanding bulk async groups to wait for
     :type count: Int
+    :param read: Whether to use the PTX .read modifier
+    :type read: Optional[bool]
     """
     nvvm.cp_async_bulk_wait_group(group=count, read=read, loc=loc, ip=ip)
 

tilelang/language/builtin.py

@@ -358,20 +358,23 @@ def tma_store_arrive():
     return tirx.call_intrin("handle", tirx.op.Op.get("tl.tma_store_arrive"))
 
 
-def tma_store_wait(count: int = 0):
+def tma_store_wait(count: int = 0, read: bool = True):
     """Wait for completion of TMA store operations.
 
     Waits until the number of outstanding TMA store groups is at most ``count``.
-    Maps to the PTX instruction ``cp.async.bulk.wait_group.read <count>``.
+    Maps to ``cp.async.bulk.wait_group.read <count>`` by default, or
+    ``cp.async.bulk.wait_group <count>`` when ``read`` is false.
 
     Args:
         count (int): The maximum number of outstanding store groups allowed
             to remain in flight. Defaults to 0 (wait for all stores to complete).
+        read (bool): Whether to use the PTX ``.read`` modifier, which only
+            waits for the source reads to complete. Defaults to True.
 
     Returns:
         tirx.Call: A handle to the store wait operation
     """
-    return tirx.call_intrin("handle", tirx.op.Op.get("tl.tma_store_wait"), count)
+    return tirx.call_intrin("handle", tirx.op.Op.get("tl.tma_store_wait"), count, read)
 
 
 def set_max_nreg(reg_count: int, is_inc: int):