[Bugfix] Reorganize pass for thread_sync

[silentCoder-dev]

Running the MergeIf pass prior to ThreadSync can lead to a deadlock condition.

Summary by CodeRabbit

• Refactor
  • Reordered compiler optimization passes in the code generation pipeline.
  • Refined conditional logic for memory fence and synchronization operations based on target capabilities.

_{✏️ 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

Modified the optimization pass pipeline in the compilation phase by reordering and conditionally gating fence proxy injection and if-statement merging. Removed early injection from non-warp-specialized flows and relocated the merge-if-statement pass to execute after thread synchronization operations.

Changes

Cohort / File(s)	Summary
Compilation Pass Reordering tilelang/engine/phase.py	Removed MergeIfStmt and InjectFenceProxy from early OptimizeForTarget path; introduced conditional fence proxy injection after MergeSharedMemoryAllocations based on allow_tma_and_warp_specialized or allow_fence_proxy; relocated MergeIfStmt to execute post-ThreadSync operations

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

Poem

🐰 Hops through the passes, reordered with care,
Fences and merges, arranged just right there,
Thread syncs come first, then if-statements align,
The pipeline now flows in a beautiful design! ✨

🚥 Pre-merge checks | ✅ 2 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% 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 clearly and specifically identifies the main change: reorganizing pass execution order for thread_sync to address a deadlock issue, which aligns with the core purpose described in the PR objectives.

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

✨ Finishing touches

• 📝 Generate docstrings

---

🧹 Recent nitpick comments

tilelang/engine/phase.py (1)

274-280: Simplify fence‑proxy injection condition.

The current nested branching injects the fence proxy for any TMA-capable target anyway. You can collapse this to a single allow_fence_proxy check to reduce branching and keep the intent clearer.

♻️ Proposed refactor

-    if allow_tma_and_warp_specialized(pass_ctx=pass_ctx, target=target):
-        mod = tilelang.transform.InjectFenceProxy()(mod)
-    else:
-        if allow_fence_proxy(target=target):
-            # in hopper device, wgmma is an async proxy
-            # so we need to inject a fence proxy before it
-            mod = tilelang.transform.InjectFenceProxy()(mod)
+    if allow_fence_proxy(target=target):
+        # in hopper device, wgmma is an async proxy
+        # so we need to inject a fence proxy before it
+        mod = tilelang.transform.InjectFenceProxy()(mod)

---

📜 Recent review details

Configuration used: defaults

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 5feb225 and 1b76f97.

📒 Files selected for processing (1)

• tilelang/engine/phase.py

🧰 Additional context used

🧬 Code graph analysis (1)

tilelang/engine/phase.py (3)

tilelang/transform/__init__.py (2)

• InjectFenceProxy (253-261)
• MergeIfStmt (194-202)

src/transform/inject_fence_proxy.cc (2)

• InjectFenceProxy (313-321)
• InjectFenceProxy (313-313)

src/transform/merge_if_stmt.cc (2)

• MergeIfStmt (125-130)
• MergeIfStmt (125-125)

⏰ 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). (1)

• GitHub Check: Quick Lint

🔇 Additional comments (1)

tilelang/engine/phase.py (1)

281-284: MergeIf after ThreadSync looks correct.

Placing MergeIfStmt after the shared ThreadSync passes aligns with the deadlock fix intent and the surrounding comments.

_{✏️ Tip: You can disable this entire section by setting review_details to false in your review settings.}

---

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.}

[kurisu6912]

LGTM!