[Bugfix] Fix variable scoping issue in InjectSoftwarePipeline for transitive LetStmt dependencies

[LeiWang1999]

Summary

• Fix variable scoping bug in InjectSoftwarePipeline pass where variables with transitive dependencies on loop variables were incorrectly scoped
• Handle chained LetStmt dependencies (e.g., id = ids[i], id2 = ids2[id]) correctly

Problem

When using pipelined loops with chained variable definitions like:

for i in T.Pipelined(4, num_stages=1):
    id = ids[i]        # depends on loop var i
    id2 = ids2[id]     # depends on id (transitively on loop var)
    T.copy(KV[id2, :], A)

After the InjectSoftwarePipeline transformation, the variable id would become "floating" (undefined in scope) because:

1. id2's definition was placed at an outer scope
2. id's definition (via LetStmt) was inside an inner scope
3. This caused id2 to reference an undefined variable id

Root Cause

Two issues in src/transform/inject_pipeline.cc:

1. Only direct loop variable dependencies were tracked: The code only checked if a LetStmt's value directly used the loop variable, ignoring transitive dependencies through other LetStmts.

2. Wrong LetStmt wrapping order: When building the LetStmt chain, forward iteration produced incorrect nesting where inner variables were referenced by outer definitions.

Solution

1. Track transitive dependencies: Now collects all variables from previously recorded loop-var-dependent LetStmts and checks if current LetStmt depends on any of them.

2. Reverse iteration for LetStmt wrapping: Iterate in reverse order so that earlier definitions scope over later ones, producing correct nesting:

   LetStmt(id, ids[...],
     LetStmt(id2, ids2[id],
       body))
   

Test plan

• Added test case test_tilelang_issue_1263.py that reproduces the issue
• Verified the IR transformation produces correct variable scoping

🤖 Generated with Claude Code

Summary by CodeRabbit

Release Notes

• Bug Fixes

  • Fixed variable dependency handling in pipeline injection for improved correctness.

• Refactor

  • Simplified elementwise example by removing autotuning capability.

• Chores

  • Updated submodule 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

This PR updates a TVM submodule reference, removes autotuning infrastructure from an elementwise example, and refines pipeline injection logic to correctly handle transitive LetStmt dependencies by iterating wrappers in reverse order for proper scoping.

Changes

Cohort / File(s)	Summary
Submodule Update 3rdparty/tvm	Updated TVM submodule pointer to new commit da7f19b6908045a1f9bf94cb7e044beaa32421b6.
Autotuning Removal examples/elementwise/example_elementwise_add.py, examples/elementwise/test_example_elementwise.py	Removed @tilelang.autotune decorator, eliminated get_configs() function, dropped use_autotune parameter from main(), and removed test_example_elementwise_add_autotune() test.
Pipeline Injection Logic src/transform/inject_pipeline.cc	Reversed iteration order when wrapping rewritten blocks with LetStmt substitutions; enhanced LetStmt dependency detection to account for transitive variable dependencies beyond direct loop variable usage.
Test Configuration Update testing/python/issue/test_tilelang_issue_1263.py	Replaced JIT decorator with explicit tilelang.compile() calls using and without specific pass configurations.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

Possibly related PRs

• tile-ai/tilelang#1399: Modifies pipeline injection logic in src/transform/inject_pipeline.cc with LetStmt wrapping and dependency handling similar to this PR.
• tile-ai/tilelang#1525: Refactors PipelineInjector/PipelineBodyRewriter behavior in src/transform/inject_pipeline.cc alongside buffer allocation changes.
• tile-ai/tilelang#1445: Directly modifies examples/elementwise/example_elementwise_add.py autotuning infrastructure and decorator handling.

Suggested reviewers

• kurisu6912
• SiriusNEO

Poem

🐰 In pipelines deep where Let statements flow,
We chase dependencies, watching them grow,
Reverse the order, let scoping be right,
Autotuning fades—pure simplicity's light! ✨

🚥 Pre-merge checks | ✅ 3

✅ Passed checks (3 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title directly addresses the main change: fixing a variable scoping bug in InjectSoftwarePipeline for transitive LetStmt dependencies, which aligns with the primary modification in src/transform/inject_pipeline.cc and the new test case.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.

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

[Copilot]

Pull request overview

This PR fixes a variable scoping bug in the InjectSoftwarePipeline transformation pass that occurred when LetStmt variables had transitive dependencies on loop variables through other LetStmt definitions.

Changes:

• Track transitive dependencies by maintaining a set of all variables dependent on the loop variable, not just direct dependencies
• Reverse the iteration order when wrapping LetStmts to ensure proper variable scoping (earlier definitions scope over later ones)
• Add comprehensive test case with multiple compilation configurations to verify the fix

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.

File	Description
testing/python/issue/test_tilelang_issue_1263.py	Updated test to use tilelang.compile() with multiple configurations instead of @tilelang.jit() decorator
src/transform/inject_pipeline.cc	Fixed variable scoping by tracking transitive LetStmt dependencies and reversing LetStmt wrapping order

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

The test case validates compilation succeeds but doesn't verify the correctness of the variable scoping fix. Consider adding assertions to check that the IR transformation produces the expected LetStmt nesting structure, or add a runtime test that would fail with incorrect scoping.

[coderabbitai]

Actionable comments posted: 1

🤖 Fix all issues with AI agents

In @3rdparty/tvm:
- Line 1: The PR references the wrong commit for the InjectSoftwarePipeline fix;
find the actual commit that modifies inject_pipeline.cc and addresses transitive
LetStmt dependencies (the change should touch
InjectSoftwarePipeline/inject_pipeline.cc), then update the referenced commit
hash and message wherever it is documented or cited in the PR/changelog to point
to that correct commit and ensure the description matches the actual diff (i.e.,
confirm the commit modifies inject_pipeline.cc and fixes transitive LetStmt
dependency handling in InjectSoftwarePipeline).

🧹 Nitpick comments (1)

examples/elementwise/example_elementwise_add.py (1)

31-31: Remove unused use_autotune parameter.

The use_autotune parameter is no longer used in the function body after removing the autotuning functionality. This dead parameter should be removed for clarity.

♻️ Suggested fix

-def main(M=1024, N=1024, use_autotune=False):
+def main(M=1024, N=1024):

📜 Review details

Configuration used: defaults

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between fd260e3 and 08cc5bf.

📒 Files selected for processing (5)

• 3rdparty/tvm
• examples/elementwise/example_elementwise_add.py
• examples/elementwise/test_example_elementwise.py
• src/transform/inject_pipeline.cc
• testing/python/issue/test_tilelang_issue_1263.py

💤 Files with no reviewable changes (1)

• examples/elementwise/test_example_elementwise.py

🧰 Additional context used

🧠 Learnings (3)

📚 Learning: 2025-11-14T07:56:11.098Z

Learnt from: lucifer1004
Repo: tile-ai/tilelang PR: 1256
File: testing/python/jit/test_tilelang_jit_gemm_nvrtc.py:55-115
Timestamp: 2025-11-14T07:56:11.098Z
Learning: In `testing/python/jit/test_tilelang_jit_gemm_nvrtc.py`, the global function `tilelang_callback_cuda_postproc` registered via `tvm.register_global_func(..., override=True)` is intentionally not restored after the test completes, as the persistent behavior is expected.

Applied to files:

• testing/python/issue/test_tilelang_issue_1263.py

📚 Learning: 2026-01-06T05:20:51.649Z

Learnt from: silentCoder-dev
Repo: tile-ai/tilelang PR: 1606
File: testing/python/transform/test_tilelang_transform_hoist_broadcast_values.py:30-30
Timestamp: 2026-01-06T05:20:51.649Z
Learning: In `testing/python/transform/test_tilelang_transform_hoist_broadcast_values.py`, the test validates that the `hoist_broadcast_values` transformation pass correctly identifies and hoists broadcast operations by checking for patterns in the generated kernel source code. The specific literal values used (e.g., 430) are not important for the test's purpose, as it does not validate numerical precision or actual stored tensor values.

Applied to files:

• testing/python/issue/test_tilelang_issue_1263.py

📚 Learning: 2025-12-18T04:50:00.512Z

Learnt from: silentCoder-dev
Repo: tile-ai/tilelang PR: 1464
File: testing/python/language/test_tilelang_language_rand.py:14-14
Timestamp: 2025-12-18T04:50:00.512Z
Learning: In `testing/python/language/test_tilelang_language_rand.py`, the TileLang kernel uses `blk_M = M` (single block) and calls `rng_rand()` four times per element to align results with the Triton implementation, which uses `blk_M = 128` (multiple blocks) and calls the RNG once per element. These differences compensate for internal RNG behavior differences between TileLang and Triton.

Applied to files:

• testing/python/issue/test_tilelang_issue_1263.py

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

• 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 (4)

testing/python/issue/test_tilelang_issue_1263.py (1)

27-34: LGTM! Good test coverage for the transitive dependency fix.

The test correctly exercises the transitive LetStmt dependency pattern (id = ids[i] → id2 = ids2[id]) that this PR fixes. Testing with both default settings and with warp specialization/TMA lowering disabled provides good coverage of different code paths.

src/transform/inject_pipeline.cc (2)

841-861: LGTM! Correct reverse iteration for proper LetStmt scoping.

The reverse iteration ensures that earlier LetStmt definitions (e.g., id = ids[i]) correctly scope over later ones that depend on them (e.g., id2 = ids2[id]). The explanatory comments clearly document the expected nesting behavior.

---

1089-1118: LGTM! Proper transitive dependency tracking for LetStmts.

The implementation correctly handles transitive dependencies by building a dependent_vars set that includes both the loop variable and all previously recorded loop-var-dependent variables. This ensures chains like id = ids[i] → id2 = ids2[id] are both captured and later substituted correctly.

examples/elementwise/example_elementwise_add.py (1)

35-38: LGTM!

The simplified kernel invocation with fixed configuration values is appropriate for an example file.