[Fix] Fix frame scope error in T.macro #1308
Conversation
WalkthroughA new test function Changes
Sequence Diagram(s)sequenceDiagram
participant UserCode as User Code
participant MacroCtx as macro() context
participant FrameStack as Frame Stack
UserCode->>MacroCtx: enter macro context
MacroCtx->>FrameStack: push MacroFrame
Note right of FrameStack: MacroFrame is active during body
rect rgb(220, 240, 230)
MacroCtx->>MacroCtx: execute macro body (bindings, transforms)
MacroCtx->>FrameStack: perform bind/assign operations
end
MacroCtx->>FrameStack: replace MacroFrame with ExitedMacroFrame
Note left of FrameStack: ExitedMacroFrame marks post-macro state
MacroCtx->>UserCode: yield / exit macro context
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes
Poem
Pre-merge checks and finishing touches❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✨ Finishing touches
🧪 Generate unit tests (beta)
Tip 📝 Customizable high-level summaries are now available in beta!You can now customize how CodeRabbit generates the high-level summary in your pull requests — including its content, structure, tone, and formatting.
Example instruction:
Note: This feature is currently in beta for Pro-tier users, and pricing will be announced later. 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. Comment |
|
👋 Hi! Thank you for contributing to the TileLang project. Please remember to run We appreciate you taking this step! Our team will review your contribution, and we look forward to your awesome work! 🚀 |
Removed debug print statement from enter_frame method.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
testing/python/language/test_tilelang_language_frontend_v2.py (1)
398-423: Turnframe_inside_macrointo an actual test and strengthen the checkThe
frame_inside_macrohelper nicely exercises the “macro inside jit’d function” scenario, but as written:
- Its name doesn’t start with
test_, so normal pytest discovery won’t run it.- It only builds
kernel = get_sample_kernel()without invoking the kernel or asserting anything.To make this a proper regression test for the T.macro frame fix, consider:
-def frame_inside_macro(): +def test_frame_inside_macro(): @@ - kernel = get_sample_kernel() # noqa: F841 + kernel = get_sample_kernel() + # A simple smoke run is enough; we just care that it compiles and runs without frame-scope errors. + num_blocks = 1 + idx_out = tilelang.testing.alloc_tensor((32,), "int32") + kernel(num_blocks, idx_out)(Or, if there’s an existing testing utility for tensor allocation/assertion in this module, reuse that instead.)
This will ensure the scenario that previously triggered the frame-scope error is continuously covered in CI.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
testing/python/language/test_tilelang_language_frontend_v2.py(1 hunks)tilelang/language/v2/builder.py(3 hunks)
🧰 Additional context used
🧠 Learnings (1)
📚 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/language/test_tilelang_language_frontend_v2.py
🧬 Code graph analysis (1)
testing/python/language/test_tilelang_language_frontend_v2.py (5)
tilelang/jit/__init__.py (3)
jit(293-294)jit(298-310)jit(313-382)tilelang/language/v2/builder.py (4)
macro(163-187)macro(568-605)prim_func(156-160)prim_func(674-767)tilelang/language/allocate.py (1)
alloc_fragment(59-70)tilelang/language/copy.py (1)
copy(13-101)tilelang/language/loop.py (1)
Parallel(11-31)
🔇 Additional comments (1)
tilelang/language/v2/builder.py (1)
83-86: Macro frame stack handling viaExitedMacroFramelooks soundUsing a lightweight
ExitedMacroFramesentinel and manually appending/relabeling theMacroFrame(instead of wrapping it inwith_frame) is a reasonable way to keep frames created inside macros (e.g.,LetStmtaround temporaries) alive until the enclosing TIR control frame exits, while ensuring:
find_frame_idx(MacroFrame)only sees active macros; and- outer bindings (like
c = foo(1); d = c) still occur within the same frame stack that includes macro‑generated lets.This aligns with the comment in Lines [175]-[184] and should address the previous frame‑scope error without obvious new correctness issues.
Also applies to: 171-186
There was a problem hiding this comment.
Actionable comments posted: 1
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
tilelang/language/v2/builder.py(2 hunks)
⏰ 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). (3)
- GitHub Check: Test for Python 3.12 with CUDA-12.8 (on self-hosted-nvidia)
- GitHub Check: Test for Python 3.12 with Metal (on macos-latest)
- GitHub Check: Test for Python 3.12 with ROCm-6.3 (on self-hosted-amd)
| pos = len(self.frames) | ||
| # here we add a ExitedMacroFrame to preserve the frame stack inside macro | ||
| # because macro may bind some variable, and return it | ||
| # | ||
| # ```py | ||
| # @T.macro | ||
| # def foo(x): | ||
| # y = x + 1 | ||
| # return y | ||
| # @T.prim_func | ||
| # def bar(): | ||
| # c = foo(1) # macro generates let y = x + 1 | ||
| # d = c # d = c should lay inside frame of `let y = x + 1` | ||
| self.frames.append(MacroFrame()) | ||
| yield | ||
| self.frames[pos] = ExitedMacroFrame() |
There was a problem hiding this comment.
Macro bindings still break outside the macro scope
Inside bind() we stash self.frames[frame_idx] in self.name_inside_frame[name]. During the macro body that object is the MacroFrame you just appended. After the yield, you replace the entry in self.frames with a new ExitedMacroFrame() instance, but the values stored in self.name_inside_frame still point at the old MacroFrame. When rval() runs later it checks frame not in self.frames and immediately raises the same runtime error (#1307) because the original MacroFrame is no longer present. So variables bound inside the macro still cannot be read after the macro returns. This is a regression blocker.
Patch idea:
- self.frames.append(MacroFrame())
+ macro_frame = MacroFrame()
+ self.frames.append(macro_frame)
yield
- self.frames[pos] = ExitedMacroFrame()
+ exited_frame = ExitedMacroFrame()
+ self.frames[pos] = exited_frame
+ for key, frame in list(self.name_inside_frame.items()):
+ if frame is macro_frame:
+ self.name_inside_frame[key] = exited_frameThis keeps the frame object referenced by existing bindings inside the active self.frames list, unblocking the scope check. Please address before merging.
🤖 Prompt for AI Agents
In tilelang/language/v2/builder.py around lines 171-186, the code appends a new
MacroFrame then after the yield replaces self.frames[pos] with a fresh
ExitedMacroFrame(), but existing bindings in self.name_inside_frame still
reference the original MacroFrame and later fail the "frame not in self.frames"
check; fix by not creating a new frame object: update/convert the existing
MacroFrame in-place to represent an exited macro (e.g., set an "exited" flag or
mutate its internals to match ExitedMacroFrame semantics) OR if you must replace
the object, iterate over self.name_inside_frame and rebind any entries that
pointed to the old frame to the new ExitedMacroFrame so all stored references
remain valid.
* [Fix] Fix tile-ai#1307 by adding macro inside function * fix lint error * add comments and fix lint error * Remove debug print from enter_frame method Removed debug print statement from enter_frame method. --------- Co-authored-by: Lei Wang <34334180+LeiWang1999@users.noreply.github.com>
2 participants
This pr fixes the error in #1307 by disabling frame collapse in T.macro:
ExitedMacroFrameto replace the macro frameSummary by CodeRabbit
Tests
Improvements
✏️ Tip: You can customize this high-level summary in your review settings.