[Refactor] Improve CallNode handling to include annotations in various operations

[LeiWang1999]

as title.

Summary by CodeRabbit

• Bug Fixes

  • Fixed annotation handling in atomic operations to prevent unwanted metadata propagation.
  • Improved annotation preservation through CUDA dispatch and code transformation passes.
  • Added null-check for call validation in CUDA shuffle operations.

• Chores

  • Updated third-party dependencies.

_{✏️ Tip: You can customize this high-level summary in your review settings.}

[github-actions]

👋 Hi! Thank you for contributing to the TileLang project.

Please remember to run pre-commit run --all-files in the root directory of the project to ensure your changes are properly linted and formatted. This will help ensure your contribution passes the format check.

We appreciate you taking this step! Our team will review your contribution, and we look forward to your awesome work! 🚀

[coderabbitai]

📝 Walkthrough

Walkthrough

Updates annotation handling across TVM IR transformations to ensure Call node annotations are properly propagated through CUDA intrinsic dispatchers and barrier injection passes. Filters "use_tma" annotation from atomic operations while preserving it in node metadata. Includes TVM submodule reference update.

Changes

Cohort / File(s)	Summary
TVM Submodule Update 3rdparty/tvm	Updated submodule commit reference from da7f19b6 to 0794c13a; no observable changes to code logic.
Call Annotation Propagation src/target/intrin_rule_cuda.cc	DispatchCUDAWarpActiveMask and DispatchCUDAShuffle now propagate source call annotations to reconstructed Call nodes. Added null-check for CallNode in DispatchCUDAShuffle.
TMA Barrier Annotation Propagation src/transform/inject_tma_barrier.cc	All Call node reconstructions now include op->annotations as fourth parameter across TmaExpectTxRewriter, BarrierCreationRewriter, and TmaBarrierRewriter. Preserves metadata through rewrites.
Atomic Operation Annotation Filtering src/op/atomic_add.cc	AtomicAddNode::MakeSIMTLoop and AtomicAddNode::Lower now copy annotations, remove "use_tma" entry, and pass filtered set to atomic/tma_store calls, preventing "use_tma" propagation to atomic operations.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~12 minutes

Possibly related PRs

• [Bugfix] Fix missing annotations for default CallNode Visitor #1659: Addresses TVM fix for default CallNode visitor behavior; both PRs handle Call-node annotation propagation and coordinate with TVM submodule updates.
• [Feature]: Add test for atomicadd auto vectorize and remove useless code #1019: Modifies AtomicAddNode implementation internals; shares changes to how atomic-add extern calls are constructed.
• [Language] Support atomic add with ret #870: Modifies src/op/atomic_add.cc; overlaps in atomic add call argument and annotation handling.

Poem

🐰 Annotations hop through IR trees,
From Call to Call, with practiced ease,
Filtering out what shouldn't be,
While metadata runs wild and free,
The CUDA code now knows its place! ✨

🚥 Pre-merge checks | ✅ 2 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 28.57% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (2 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title accurately describes the main refactoring objective: improving CallNode handling to include annotations across various operations.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing touches

• 📝 Generate docstrings

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Actionable comments posted: 1

🤖 Fix all issues with AI agents

In @src/target/intrin_rule_cuda.cc:
- Around line 119-123: DispatchCUDAWarpActiveMask dereferences call without
checking e->as<CallNode()]; add the same null-check used in DispatchCUDAShuffle
by inserting ICHECK(call != nullptr) after obtaining const CallNode *call =
e.as<CallNode>(); so you bail early with a clear invariant if e is not a
CallNode before using call->dtype, call->args, or call->annotations; keep the
rest of the function constructing the Call(...) intact.

📜 Review details

Configuration used: defaults

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 5e90edc and d7e935c.

📒 Files selected for processing (4)

• 3rdparty/tvm
• src/op/atomic_add.cc
• src/target/intrin_rule_cuda.cc
• src/transform/inject_tma_barrier.cc

🧰 Additional context used

🧠 Learnings (1)

📚 Learning: 2025-09-12T09:47:46.474Z

Learnt from: kurisu6912
Repo: tile-ai/tilelang PR: 794
File: tilelang/transform/add_bufstore_wrapper.py:30-33
Timestamp: 2025-09-12T09:47:46.474Z
Learning: In TVM's PyStmtExprMutator, visit_block_ methods typically call super().visit_block_(op) to process child nodes and update internal state, but return the original op when the block itself doesn't need transformation. The pattern `return op` is correct for blocks that serve as containers where mutations happen at deeper levels.

Applied to files:

• src/transform/inject_tma_barrier.cc

⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (3)

• GitHub Check: Test for Python 3.12 with CUDA-12.8 (on self-hosted-nvidia)
• GitHub Check: Test for Python 3.12 with Nightly-ROCm-7.1 (on self-hosted-amd)
• GitHub Check: Test for Python 3.12 with Metal (on macos-latest)

🔇 Additional comments (7)

3rdparty/tvm (1)

1-1: Verify TVM API compatibility—internal code review passed, external verification required.

Internal review confirms consistent annotation handling across all modified files (src/op/atomic_add.cc, src/target/intrin_rule_cuda.cc, src/transform/inject_tma_barrier.cc). All Call constructors properly use the 4-argument form with annotations parameter. No security advisories found for the TVM commit.

However, external verification of TVM API changes between the old and new commits is still required to ensure the CallNode construction patterns remain compatible with the updated TVM version.

src/op/atomic_add.cc (2)

283-287: LGTM! Annotation filtering is correctly implemented.

The pattern of copying the annotations map, erasing "use_tma", and passing the filtered annotations to the atomicadd_elem_op call is correct. This ensures the use_tma annotation remains on the AtomicAddNode for decision-making (e.g., in GetUseTMA()) while not being forwarded to the underlying atomic operation.

---

396-403: LGTM! Consistent annotation handling in the TMA store path.

The same filtering pattern is correctly applied here: copying annotations, erasing "use_tma", and passing the filtered map to the tma_store call. This maintains consistency with MakeSIMTLoop and ensures proper annotation propagation throughout the lowering pipeline.

src/target/intrin_rule_cuda.cc (1)

125-133: LGTM! Proper null-check and annotation propagation.

The added ICHECK(call != nullptr) provides a clear diagnostic if the input is malformed, and the annotations are correctly forwarded to the new Call node. This aligns with the PR's goal of consistent annotation handling across CUDA intrinsic dispatchers.

src/transform/inject_tma_barrier.cc (3)

165-179: LGTM! Annotations preserved in TMA load rewrite.

The op->annotations are correctly propagated when reconstructing the Call node with the updated barrier argument. This ensures any metadata attached to the original TMA load call survives the rewrite.

---

350-389: LGTM! Annotations preserved in barrier creation rewrite.

The op->annotations are correctly forwarded when constructing the new Call node with the modified barrier argument list. This maintains any metadata attached to the original create_list_of_mbarrier call.

---

507-568: LGTM! Consistent annotation propagation across all TMA barrier rewrite paths.

All Call node reconstructions in TmaBarrierRewriter::VisitExpr_ correctly include op->annotations, ensuring metadata is preserved through:

• 1D TMA load promotion (line 524)
• TMA load barrier ID assignment (line 540)
• mbarrier_expect_tx → ptx_arrive_barrier_expect_tx conversion (line 555)
• mbarrier_expect_tx argument update (line 557)
• ptx_arrive_barrier passthrough (line 565)