[Bugfix] Introduce a flag to avoid unnecessary broadcast hoist and enable for let stmt

[LeiWang1999]

as title, examples pls checkout updated test.

Summary by CodeRabbit

• Bug Fixes

  • Improved handling of broadcast value hoisting in let statements and nested buffer scenarios, ensuring correct preservation of broadcasts in complex transformations.

• Tests

  • Added test cases validating broadcast hoisting behavior in let statements and nested buffer store contexts.

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

[Copilot]

Pull request overview

This pull request introduces a bugfix to avoid unnecessary broadcast hoisting and extends support for hoisting broadcasts in LetStmt expressions. The change adds a hoist_enabled flag that controls when broadcast values should be hoisted, and implements a new visitor method for LetStmt nodes.

• Adds hoist_enabled flag to control when broadcast hoisting occurs
• Implements visit_let_stmt_ method to handle broadcasts in LetStmt value expressions
• Updates test suite with a new test case demonstrating LetStmt broadcast hoisting

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 4 comments.

File	Description
tilelang/transform/hoist_broadcast_values.py	Adds hoist_enabled flag and visit_let_stmt_ method to support broadcast hoisting in LetStmt expressions; modifies visit_buffer_store_ to enable/disable the flag appropriately
testing/python/transform/test_tilelang_transform_hoist_broadcast_values.py	Adds test_transform_hoist_let_stmt test case to verify LetStmt broadcast hoisting behavior

---

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

[coderabbitai]

📝 Walkthrough

Walkthrough

The PR enhances the HoistBroadcastValuesMutator to support LetStmt broadcast hoisting through a conditional hoist flag and new visitor method. It adds per-statement context preservation and includes two test cases validating hoisting behavior in LetStmt and nested BufferStore scenarios.

Changes

Cohort / File(s)	Summary
Core transformation logic tilelang/transform/hoist_broadcast_values.py	Added hoist_enabled flag to conditionally gate broadcast hoisting. Introduced visit_let_stmt_() method for LetStmt value expression hoisting. Enhanced visit_buffer_store_() to preserve and restore hoist state and pending definitions across statement contexts, enabling nested scope handling.
Test coverage testing/python/transform/test_tilelang_transform_hoist_broadcast_values.py	Added test_transform_hoist_let_stmt() to verify hoisting of LetStmt broadcast values. Added test_transform_hoist_let_stmt_with_nested_bufferstore_broadcasts() to validate preservation of broadcasts from both LetStmt and nested BufferStore contexts.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

Possibly related PRs

• [Feature] Add hoist_broadcast_values pass #1606: Directly modifies the same hoist_broadcast_values transformation logic and mutator class implementation.

Poem

🐰 A rabbit hops through broadcast trees,
Hoisting LetStmts with such ease,
Nested scopes now dance with grace,
Each context finds its perfect place! ✨

🚥 Pre-merge checks | ✅ 2 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 9.09% 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 change: introducing a flag to control broadcast hoisting and enabling it for let statements, which aligns with the core modifications in both the transform logic and test cases.

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

✨ Finishing touches

• 📝 Generate docstrings

---

📜 Recent review details

Configuration used: defaults

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between d5503cd and 341c1fb.

📒 Files selected for processing (2)

• testing/python/transform/test_tilelang_transform_hoist_broadcast_values.py
• tilelang/transform/hoist_broadcast_values.py

🧰 Additional context used

🧠 Learnings (4)

📓 Common learnings

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:45.325Z
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.

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

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:45.325Z
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:

• tilelang/transform/hoist_broadcast_values.py

📚 Learning: 2025-09-12T09:47:46.474Z

Learnt from: kurisu6912
Repo: tile-ai/tilelang PR: 794
File: tilelang/transform/add_bufstore_wrapper.py:30-33
Timestamp: 2025-09-12T09:47:46.474Z
Learning: In TVM's PyStmtExprMutator, visit_block_ methods typically call super().visit_block_(op) to process child nodes and update internal state, but return the original op when the block itself doesn't need transformation. The pattern `return op` is correct for blocks that serve as containers where mutations happen at deeper levels.

Applied to files:

• tilelang/transform/hoist_broadcast_values.py

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

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:45.325Z
Learning: 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. Do not rely on particular values (like 430); the test should validate behavior, not numerical precision or actual stored values. This guideline can be applied to all tests in the testing/python/transform directory.

Applied to files:

• testing/python/transform/test_tilelang_transform_hoist_broadcast_values.py

🧬 Code graph analysis (2)

tilelang/transform/hoist_broadcast_values.py (2)

tilelang/language/ast/ir.py (1)

• LetStmt (874-902)

tilelang/analysis/ast_printer.py (1)

• visit_stmt (32-76)

testing/python/transform/test_tilelang_transform_hoist_broadcast_values.py (2)

tilelang/language/ast/ir.py (1)

• decl_buffer (1134-1202)

tilelang/language/v2/dtypes.py (2)

• float8_e4m3fn (341-341)
• float8_e4m3fnx8 (344-344)

⏰ 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 Metal (on macos-latest)
• GitHub Check: Test for Python 3.12 with Nightly-ROCm-7.1 (on self-hosted-amd)

🔇 Additional comments (6)

tilelang/transform/hoist_broadcast_values.py (4)

19-20: LGTM: Flag enables selective hoisting.

The hoist_enabled flag correctly gates hoisting to specific statement contexts (BufferStore, LetStmt), preventing unnecessary hoisting at the top level. This aligns well with the PR objective.

---

23-23: LGTM: Conditional hoisting properly enforced.

The addition of self.hoist_enabled check ensures broadcasts are only hoisted in intended contexts, working in tandem with the flag set by parent statement visitors.

---

40-66: Excellent state management for nested statements.

The save/restore pattern correctly isolates hoisting contexts:

1. Saves outer state to prevent contamination
2. Enables hoisting locally and captures pending defs
3. Wraps the statement with LetStmts for hoisted values
4. Restores outer state to prevent leakage
5. Returns the transformed statement (critical for mutations to take effect)

This design elegantly handles nested BufferStore/LetStmt cases where inner hoisting shouldn't affect outer scopes.

---

67-102: Well-designed LetStmt hoisting with proper scope isolation.

The implementation correctly:

1. Hoists broadcasts from the LetStmt's value expression (lines 73-80)
2. Disables hoisting before visiting the body (lines 83-84), allowing nested statements like BufferStore to manage their own hoisting context
3. Wraps the resulting LetStmt with hoisted definitions from the value (lines 93-96)
4. Preserves outer scope state (lines 99-100)

The intentional disabling of hoisting before visiting the body ensures that each statement type controls its own hoisting context. This produces the correct nesting structure in cases like:

Let var1 = <expr with broadcasts> In
  BufferStore(...broadcasts...)

Result:

Let var1 = <hoisted> In
  Let val1 = <value hoisted> In
    Let var2 = <bufferstore hoisted> In
      BufferStore(...)

testing/python/transform/test_tilelang_transform_hoist_broadcast_values.py (2)

85-102: LGTM: Clear test for basic LetStmt hoisting.

This test correctly validates that broadcasts in a LetStmt's value expression are hoisted and defined as separate variables before the LetStmt. The before/after structure clearly shows the expected transformation.

---

105-133: Excellent regression test for nested hoisting.

This test validates the critical nested case where a LetStmt has broadcasts in its value expression and its body (a BufferStore) also contains broadcasts. The expected output correctly shows:

1. Broadcasts from LetStmt value hoisted first (broadcast_var, broadcast_var_1)
2. LetStmt definition using hoisted variables
3. Broadcasts from BufferStore hoisted in its own scope (broadcast_var_2)
4. BufferStore using hoisted variable

The docstring clearly explains the scenario being tested. This provides strong coverage of the state management logic.

---

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

[LeiWang1999]

LTCP, Merged.