[Feature] Support TMA store in T.tma_copy() - #1981

testing/python/language/test_tilelang_language_tma_store.py (1)

83-87: Tighten the source assertion so this test catches accidental auto-waits.

Because Line 50 already adds an explicit T.tma_store_wait(), this test still passes if lowering regresses and inserts a second implicit wait. Checking only tma_store_arrive doesn't pin the new contract; please also assert that the emitted source contains only the single user-authored store-wait invocation.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@testing/python/language/test_tilelang_language_tma_store.py` around lines 83
- 87, The current test only asserts that "tma_store_arrive" appears but doesn't
prevent an extra implicit wait from being generated; update the assertions on
kernel_source (the string returned by kernel.get_kernel_source()) to also verify
that "tma_store_wait" appears exactly once (e.g., assert
kernel_source.count("tma_store_wait") == 1) so the emitted source contains only
the single user-authored T.tma_store_wait invocation and no additional implicit
waits.

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

src/target/codegen_cutedsl.cc (1)

641-643: Add the same tma_store_wait arg guards in CuTeDSL codegen.

This code path also assumes op->args[0] is present and IntImm. Guarding count/type keeps diagnostics consistent across backends.

Suggested hardening

  } else if (op->op.same_as(tl::tma_store_wait())) {
-    int count = Downcast<IntImm>(op->args[0])->value;
+    ICHECK_EQ(op->args.size(), 1) << "tma_store_wait expects 1 argument";
+    const auto* count_imm = op->args[0].as<IntImmNode>();
+    ICHECK(count_imm) << "tma_store_wait expects IntImm count";
+    int count = count_imm->value;
     PrintIndent();
     stream << "tl.tma_store_wait(" << count << ")\n";

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@src/target/codegen_cutedsl.cc` around lines 641 - 643, The tma_store_wait
emission assumes op->args[0] exists and is an IntImm; update the code around the
Downcast<IntImm>(op->args[0])->value usage in codegen_cutedsl.cc to first guard
that op->args is non-empty and that op->args[0] is an IntImm (e.g., using a safe
type check instead of direct Downcast), and if the guard fails emit a clear
diagnostic or fall back to a safe behavior before calling PrintIndent() and
writing "tl.tma_store_wait(...)" so diagnostics and behavior match other
backends.

src/target/codegen_cuda.cc (1)

2037-2039: Validate tma_store_wait arg count/type before use.

This path reads op->args[0] unguarded. Adding explicit checks makes failures deterministic and easier to debug when malformed IR reaches codegen.

Suggested hardening

  } else if (op->op.same_as(tl::tma_store_wait())) {
-    int count = Downcast<IntImm>(op->args[0])->value;
+    ICHECK_EQ(op->args.size(), 1) << "tma_store_wait expects 1 argument";
+    const auto* count_imm = op->args[0].as<IntImmNode>();
+    ICHECK(count_imm) << "tma_store_wait expects IntImm count";
+    int count = count_imm->value;
     this->PrintIndent();
     this->stream << "tl::tma_store_wait<" << count << ">();\n";

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@src/target/codegen_cuda.cc` around lines 2037 - 2039, The code is accessing
op->args[0] without validation before Downcast<IntImm>, so add explicit guards
that op->args has at least one element and that op->args[0] is an IntImm before
using it in the emitted tl::tma_store_wait template; if the checks fail, emit a
clear diagnostic (e.g., an ICHECK/LOG + throw or return an error) mentioning the
operator and that tma_store_wait's count must be a constant IntImm. Update the
block around the Downcast<IntImm>(op->args[0])->value and its emission of
"tl::tma_store_wait<...>()" to perform these validations and produce a
deterministic, informative failure path instead of undefined behavior.