Skip to content

[Fix] Fix frame scope error in T.macro #1308

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
LeiWang1999 merged 4 commits into from
Nov 21, 2025
Merged

[Fix] Fix frame scope error in T.macro #1308

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
7c5a6e9
[Fix] Fix #1307 by adding macro inside function
kurisu6912 Nov 21, 2025
a1e0185
fix lint error
kurisu6912 Nov 21, 2025
c2799bf
add comments and fix lint error
kurisu6912 Nov 21, 2025
94f2c4a
Remove debug print from enter_frame method
LeiWang1999 Nov 21, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 26 additions & 0 deletions testing/python/language/test_tilelang_language_frontend_v2.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -395,5 +395,31 @@ def prim_call_macro():
pass


def frame_inside_macro():

@tilelang.jit
def get_sample_kernel():

@T.macro
def transform(x):
return x + 1

@T.prim_func
def sample_kernel(
num_blocks: T.int32,
idx_out: T.Tensor[(32,), T.int32],
):
with T.Kernel(num_blocks, threads=32) as block_idx: # noqa: F841
fragment = T.alloc_fragment(32, 'int32')
T.copy(idx_out, fragment)

for i in T.Parallel(32):
idx_out[i] = transform(fragment[i])

return sample_kernel

kernel = get_sample_kernel() # noqa: F841


if __name__ == '__main__':
tilelang.testing.main()
22 changes: 20 additions & 2 deletions tilelang/language/v2/builder.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -80,6 +80,10 @@ class MacroFrame(Frame):
...


class ExitedMacroFrame(Frame):
...


class BoolOpFrame(Frame):
...

Expand Down Expand Up @@ -164,8 +168,22 @@ def macro(self, name=None, annotations=None):
save = self.name_inside_frame, self.arg_annotations
self.name_inside_frame = {}
self.arg_annotations = annotations or {}
with self.with_frame(MacroFrame()):
yield
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()
Comment on lines +171 to +186

@coderabbitai coderabbitai Bot Nov 21, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🔴 Critical

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_frame

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

self.name_inside_frame, self.arg_annotations = save

def get(self):
Expand Down
Loading