[BugFix] LoopUnswitching: gate non-trivial else behind PassConfig

[LeiWang1999]

This PR makes LoopUnswitching more conservative by default, and adds an opt-in PassConfig to re-enable more aggressive hoisting.

Motivation:

• In kernels with later synchronization injection (e.g. partial thread sync / fence proxy), unswitching a per-thread predicate or duplicating a non-trivial else branch can introduce control-flow that breaks correctness.

Changes:

• Default: only unswitch when the else-version of the loop body is side-effect free (no-op / Evaluate-only).
• Add PassConfig (Python: ) to allow unswitching even when the else-version is non-trivial.
• Keep the existing safety guard that blocks unswitching predicates that depend on threadIdx.*.
• Add unit tests covering default vs enabled behavior.

This keeps the canonical optimization -> while avoiding code-size blowups and unsafe control-flow splits by default.

Summary by CodeRabbit

• New Features

  • Added configuration option to enable loop unswitching even when else-branches contain side effects.

• Improvements

  • Enhanced loop unswitching to better detect and handle per-thread predicates, with conservative default behavior.

• Tests

  • Added comprehensive tests for configuration-driven loop unswitching behavior and per-thread predicate handling.

[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

This PR enhances the loop unswitching pass with configuration-driven constraints to prevent unsafe optimizations. It adds a configuration option to control non-trivial else-branch hoisting, introduces per-thread variable tracking (threadIdx) to prevent incorrect unswitching, and adds utility functions for analyzing variable usage through Let bindings and statement side-effect detection.

Changes

Cohort / File(s)	Summary
Configuration & Headers src/op/builtin.h, src/op/builtin.cc, tilelang/transform/pass_config.py	Added new public configuration key TL_LOOP_UNSWITCHING_ALLOW_NON_TRIVIAL_ELSE (default: false) to control whether loop unswitching hoists branches with side effects.
Loop Unswitching Pass src/transform/loop_unswitching.cc	Extended loop unswitching with: new helper functions UsesVarsThroughLetBindings() and IsSideEffectFreeStmt() for analysis; per-thread variable tracking (thread_idx_vars_in_scope_) to prevent unswitching when threadIdx-dependent predicates exist; enhanced LoopUnswitcher constructor accepting allow_non_trivial_else flag; updated IsLoopInvariant() and HoistableIfFinder to check disallowed variables; public API change to ApplyLoopUnswitching() now accepts allow_non_trivial_else parameter.
Test Suite testing/python/transform/test_tilelang_transform_loop_unswitching.py	Added helper _check_with_config() to run tests under configured PassContext; introduced four new test cases covering thread-local predicate rejection, non-trivial else hoisting when enabled, and side-effect handling under various configurations.

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~60 minutes

Possibly related issues

• [BUG] Loop unswitching pass causes buggy codegen #1804 — Addresses safety concerns in loop unswitching by preventing hoisting of thread-local predicates and non-trivial else branches, which directly aligns with the constraint-checking logic introduced here.

Possibly related PRs

• [Feature] Enhance Loop Unswitching with Let Binding and Condition Handling #1766 — Modifies loop unswitching pass to enhance Let-binding and If-condition handling, providing complementary variable analysis infrastructure used in this PR.
• [Feature] Implement LoopUnswitching Pass #1747 — Introduces the foundational LoopUnswitching pass and configuration knobs; this PR extends it with per-thread safety checks and configurable else-branch handling.
• [BugFix] Missing Recursive Loop Var Checking in Loop Unswitching #1801 — Adds recursive variable-usage checking through Let bindings (UsesVarsThroughLetBindings), which is a core utility introduced here for predicate analysis.

Poem

🐰 A rabbit hops through loops so neat,
But stops when threadIdx you meet!
With configs wise and side-checks keen,
Safe unswitching—the best we've seen! 🌿

🚥 Pre-merge checks | ✅ 2 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 36.84% 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 describes the main change: adding a PassConfig gate to control when non-trivial else branches can be unswitched in loop unswitching, which is the core objective of the PR.

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

✨ Finishing touches

• 📝 Generate docstrings

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment

---

No actionable comments were generated in the recent review. 🎉

🧹 Recent nitpick comments

src/transform/loop_unswitching.cc (1)

521-527: thread_idx_vars_in_scope_ should be private.

This member is internal mutable state used only by VisitStmt_ and follows the trailing-underscore naming convention that signals a private field. Move it to the private: section alongside allow_non_trivial_else_.

♻️ Suggested diff

 class LoopUnswitcher : public StmtExprMutator {
 public:
   explicit LoopUnswitcher(bool allow_non_trivial_else)
       : allow_non_trivial_else_(allow_non_trivial_else) {}
 
-  std::unordered_set<const VarNode *> thread_idx_vars_in_scope_;
-
   Stmt VisitStmt_(const ForNode *op) final {
     // ...
   }
 
 private:
   bool allow_non_trivial_else_{false};
+  std::unordered_set<const VarNode *> thread_idx_vars_in_scope_;
 };

---

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