[BugFix] Missing Recursive Loop Var Checking in Loop Unswitching

[kurisu6912]

UsesLoopVarThroughLetBindings does not check the variable recursively in the let binding stack. This pr fixes it.

Summary by CodeRabbit

• Bug Fixes

  • Improved detection of loop-variable usage through nested bindings, enhancing correctness of loop-unswitching and related optimizations.

• Tests

  • Added a test covering kernels with multiple local bindings to ensure correct behavior when hoisting/unswitching is applied.

[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

Replaced a direct UsesVar check with a recursive UsesLoopVarThroughLetBindings call to detect loop-variable usage through nested Let-bindings in loop-unswitching logic, and added a test exercising behavior with multiple Let-bindings during hoisting/unswitching.

Changes

Cohort / File(s)	Summary
Loop Unswitching Logic src/transform/loop_unswitching.cc	Replaced direct UsesVar call on a Let-binding's bound expression with a recursive UsesLoopVarThroughLetBindings(...) invocation to analyze loop-variable usage through nested Let-bindings during condition checks.
Test Coverage testing/python/transform/test_tilelang_transform_loop_unswitching.py	Added test_no_hoist_multiple_let to exercise a fused mapping kernel with multiple Let-bindings and verify behavior during compilation/unswitching.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

Possibly related PRs

• [Feature] Enhance Loop Unswitching with Let Binding and Condition Handling #1766: Related work improving Let-binding handling and hoisting/replacement logic in loop unswitching.

Suggested reviewers

• LeiWang1999

Poem

🐰
Through nested Lets I hop and peep,
I sniff the loop-var, shallow then deep.
A recursive trail I joyfully chart,
Uncovering uses tucked in apart.
Hooray—unswitching gets a clever start!

🚥 Pre-merge checks | ✅ 2 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 50.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 accurately describes the main bug fix: adding recursive loop variable checking in loop unswitching, which directly matches the core change in the codebase.

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

🤖 Fix all issues with AI agents

In `@src/transform/loop_unswitching.cc`:
- Line 357: Replace the noisy production LOG call that prints loop-invariant IR
with a debug-only log and fix its spacing: change the LOG(INFO) << "Loop
Invariant" << ffi::GetRef<IfThenElse>(op); occurrence in loop_unswitching.cc to
use DLOG(INFO) or VLOG(1) (depending on desired verbosity) and ensure the
message includes a space or separator (e.g., "Loop Invariant: ") before the
ffi::GetRef<IfThenElse>(op) output so it only emits in debug/verbose builds and
is properly formatted.

🧹 Nitpick comments (1)

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

389-400: Test only checks that compilation doesn't crash — consider adding a structural assertion.

All other tests in this file use _check(before, expected) to verify the transformed IR structurally. This test only calls .compile(), so it validates that the bug no longer causes a crash or invalid IR, but doesn't assert that the condition is actually not hoisted.

If it's difficult to express the expected output via @T.prim_func for this kernel style, a comment explaining why this test is compile-only and what invariant it guards would be helpful for future maintainers.

Based on learnings: "For Python tests of the tilelang transform passes, focus assertions on structural patterns in the generated kernel source (e.g., hoisting behavior) rather than specific numeric literals."