[Feature] Enhance Loop Unswitching with Let Binding and Condition Handling

[LeiWang1999]

• Updated the loop unswitching pass to support hoisting Let-bound variables alongside their conditions, improving optimization for loops with variable dependencies.
• Modified the CallCheckerExcludingIf class to exclude matching If nodes when checking for CallNode presence, ensuring safety during unswitching.
• Enhanced IfBranchReplacer to replace all instances of If nodes with matching conditions, streamlining the transformation process.
• Added comprehensive tests to validate the new behavior, including scenarios with multiple Let-bound variables and identical conditions across If statements.

Summary by CodeRabbit

Release Notes

• Enhancements
  • Loop unswitching optimization now robustly consolidates multiple identical conditions within loops.
  • Improved handling of variables bound in conditional expressions, ensuring proper hoisting and dependency ordering.
  • Enhanced support for complex nested conditional structures while maintaining code correctness.

[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

The pull request enhances the loop unswitching transformation to robustly handle multiple identical conditions and let-bound variables. The implementation broadens condition exclusion logic, introduces a let-binding hoisting mechanism, updates the IfBranchReplacer to consolidate matching conditions, and integrates hoisted let-bindings around transformed loop structures while maintaining SSA properties.

Changes

Cohort / File(s)	Summary
Loop Unswitching Enhancement src/transform/loop_unswitching.cc	Reworks exclusion logic to match any structurally equal condition; introduces LetVarCollector and hoisted_let_bindings tracking; updates IfBranchReplacer to consolidate all matching If nodes and remove associated LetStmts; integrates let-binding propagation into LoopUnswitcher with reverse-order hoisting to preserve dependencies and SSA properties.
Test Coverage testing/python/transform/test_tilelang_transform_loop_unswitching.py	Adds five new test cases: test_hoist_let_bound_variable, test_hoist_multiple_let_bound_variables, test_multiple_identical_conditions, test_multiple_identical_conditions_with_else, and test_no_hoist_let_bound_loop_variant, validating let-binding hoisting, condition consolidation, and loop-variant exclusion.

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~45 minutes

Possibly related PRs

• tile-ai/tilelang#1747: Earlier modification to LoopUnswitching implementation that this PR builds upon with enhanced let-binding and condition consolidation features.

Poem

🐰 Through nested loops, conditions hop and bound,
Multiple paths now merged, let-bindings found,
Variables hoisted high with care and grace,
SSA preserved in their rightful place! ✨

🚥 Pre-merge checks | ✅ 2 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 26.92% 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 summarizes the main changes: enhancing loop unswitching with Let binding and condition handling improvements. It matches the core modifications to src/transform/loop_unswitching.cc and the comprehensive test coverage added.

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

---

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: 0

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

src/transform/loop_unswitching.cc (1)

185-215: ⚠️ Potential issue | 🟠 Major

Handle transitive let-binding dependencies when checking loop variance.

The current check only inspects one level of let bindings. If a condition uses x, x depends on y, and y depends on the loop var, the condition is incorrectly treated as invariant and can be hoisted. That breaks correctness for chained lets.

🐛 Proposed fix: recursively follow let bindings when checking loop variance

 bool UsesLoopVarThroughLetBindings(
     const PrimExpr &cond, const Var &loop_var,
     const std::unordered_map<const VarNode *, PrimExpr> *let_bindings) {
-  // Check if condition directly uses loop variable
-  if (UsesVar(cond, [&](const VarNode *v) { return v == loop_var.get(); })) {
-    return true;
-  }
-
-  // Check if any Let-bound variable used in condition has a binding that uses
-  // the loop variable
-  if (let_bindings) {
-    bool uses_loop_var = false;
-    PostOrderVisit(cond, [&](const ObjectRef &obj) {
-      if (uses_loop_var)
-        return;
-      if (const auto *var_node = obj.as<VarNode>()) {
-        auto it = let_bindings->find(var_node);
-        if (it != let_bindings->end()) {
-          // Check if the bound expression uses the loop variable
-          if (UsesVar(it->second,
-                      [&](const VarNode *v) { return v == loop_var.get(); })) {
-            uses_loop_var = true;
-          }
-        }
-      }
-    });
-    if (uses_loop_var) {
-      return true;
-    }
-  }
-  return false;
+  class LoopVarChecker : public ExprVisitor {
+   public:
+    LoopVarChecker(const Var &loop_var,
+                   const std::unordered_map<const VarNode *, PrimExpr> *bindings)
+        : loop_var_(loop_var.get()), let_bindings_(bindings) {}
+
+    void VisitExpr_(const VarNode *op) final {
+      if (uses_loop_var_) return;
+      if (op == loop_var_) {
+        uses_loop_var_ = true;
+        return;
+      }
+      if (!let_bindings_ || visiting_.count(op)) return;
+      auto it = let_bindings_->find(op);
+      if (it != let_bindings_->end()) {
+        visiting_.insert(op);
+        VisitExpr(it->second);
+      }
+    }
+
+    bool uses_loop_var_ = false;
+
+   private:
+    const VarNode *loop_var_;
+    const std::unordered_map<const VarNode *, PrimExpr> *let_bindings_;
+    std::unordered_set<const VarNode *> visiting_;
+  };
+
+  LoopVarChecker checker(loop_var, let_bindings);
+  checker(cond);
+  return checker.uses_loop_var_;
 }

🧹 Nitpick comments (1)

testing/python/transform/test_tilelang_transform_loop_unswitching.py (1)

222-387: Consider adding a chained let-binding test case.

The new cases cover direct let-bound invariance well. A test where the condition uses x, x depends on y, and y depends on the loop var would guard against transitive dependency regressions.