[Feat] Allow print macro call stack in device assert#1616
[Feat] Allow print macro call stack in device assert#1616LeiWang1999 merged 6 commits intotile-ai:mainfrom
Conversation
|
👋 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! 🚀 |
📝 WalkthroughWalkthroughAdds per-file/line span tracking and propagation through macros and generated closures (Builder fields and APIs, DSLMutator filename propagation, SpanAttacher injection); modifies device_assert signature to optionally suppress stack info; updates TVM submodule pointer. Changes
Sequence Diagram(s)sequenceDiagram
autonumber
participant Caller as Macro Caller
participant Builder as Builder
participant DSL as DSLMutator / SpanAttacher
participant Closure as Generated Closure
participant Device as tl (CUDA)
rect rgb(245,250,255)
Caller->>Builder: enter macro (file,line,name) / set_fileline
Builder->>Builder: push (file,line,name) onto macro_fileline_stack
Builder->>DSL: mutate AST (pass filename)
DSL->>Closure: inject per-statement __tb.set_fileline(__tb_fl,__tb_fn,...)
Builder->>Caller: return closure with __tb_fl / __tb_fn bound
end
rect rgb(245,255,245)
Closure->>Device: device_assert(condition, msg, no_stack_info)
alt no_stack_info == False
Device->>Device: call device_assert_with_msg(get_stack_str(...))
else no_stack_info == True
Device->>Device: call device_assert or device_assert_with_msg(msg) depending on msg
end
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing touches
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 |
![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
🤖 Fix all issues with AI Agents
In @tilelang/language/print_op.py:
- Around line 127-132: get_stack_str can raise AttributeError when
Builder.current() returns None; modify get_stack_str to first retrieve builder =
Builder.current(), and if builder is None return the original msg (or msg +
newline) without attempting to call get_fileline_stack; otherwise call
builder.get_fileline_stack(stacklevel) and proceed as before—refer to the
get_stack_str function and Builder.current() to locate where to add this None
check and early-return behavior.
🧹 Nitpick comments (3)
tilelang/language/v2/ast.py (1)
582-591: Consider using iterable unpacking for cleaner concatenation.The static analysis tool suggests using iterable unpacking instead of list concatenation. This is a minor stylistic improvement.
🔎 Suggested refactor
def visit(self, node: ast.AST): node = self.generic_visit(node) if isinstance(node, ast.stmt): - return quote(f"__tb.set_fileline({self.filename_var}, {node.lineno}, {self.func_name_var})") + [node] + return [*quote(f"__tb.set_fileline({self.filename_var}, {node.lineno}, {self.func_name_var})"), node] return nodetilelang/language/v2/builder.py (2)
695-697: Consider using iterable unpacking for cleaner concatenation.Per static analysis suggestion, the list concatenation could use iterable unpacking for consistency with Python best practices.
🔎 Suggested refactor
def get_fileline_stack(self, stacklevel=1): - stack = self.macro_fileline_stack + [(self.current_file, self.current_line, self.current_macro_name)] + stack = [*self.macro_fileline_stack, (self.current_file, self.current_line, self.current_macro_name)] return stack[: len(stack) - stacklevel + 1]
699-704: Remove commented-out code.The
get_fileline_stack_msghelper is commented out. If it's not needed, it should be removed to keep the codebase clean. If it's intended for future use, consider implementing it properly or tracking it in an issue.🔎 Suggested removal
- # def get_fileline_stack_msg(self, stacklevel=1): - # stack = self.get_fileline_stack(stacklevel) - # msg = '' - # for filename, lineno, macro_name in stack: - # msg += f'{filename}:{lineno} {macro_name}\n' - # return msg -
📜 Review details
Configuration used: defaults
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (4)
3rdparty/tvmtilelang/language/print_op.pytilelang/language/v2/ast.pytilelang/language/v2/builder.py
🧰 Additional context used
🧬 Code graph analysis (2)
tilelang/language/v2/builder.py (2)
tilelang/language/kernel.py (1)
pop(38-45)tilelang/language/frame.py (1)
pop(36-50)
tilelang/language/print_op.py (1)
tilelang/language/v2/builder.py (5)
current(182-184)get_fileline_stack(695-697)macro(197-224)macro(752-788)get(226-227)
🪛 Ruff (0.14.10)
tilelang/language/v2/builder.py
696-696: Consider iterable unpacking instead of concatenation
Replace with iterable unpacking
(RUF005)
tilelang/language/v2/ast.py
590-590: Consider iterable unpacking instead of concatenation
Replace with iterable unpacking
(RUF005)
⏰ 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). (1)
- GitHub Check: Test for Python 3.12 with Metal (on macos-latest)
🔇 Additional comments (7)
3rdparty/tvm (1)
1-1: This is a submodule version update for TVM; no action required on the code review side. The tilelang implementation files referenced in the AI summary already exist in the repository and were modified in prior commits supporting this feature. This PR only bumps the TVM submodule reference to enable the "allow device assert print macro call stack" functionality.Likely an incorrect or invalid review comment.
tilelang/language/v2/ast.py (2)
474-486: LGTM! File/line context injection for debugging.The changes correctly capture the function name and inject per-statement file/line tracking via
SpanAttacher. The generated closure properly initializes__tb_fland__tb_fnvariables that are referenced by the injectedset_filelinecalls.
648-648: LGTM!Using
Path(filename).nameto extract just the base filename is appropriate for stack traces, keeping error messages concise and avoiding exposure of full filesystem paths.tilelang/language/print_op.py (1)
136-150: LGTM! Theno_stack_infoparameter provides flexibility.The implementation correctly:
- Defaults to including stack info for better debugging
- Allows users to opt-out with
no_stack_info=Truewhen performance is critical- Emits a warning about potential slowdown when using messages without stack info
tilelang/language/v2/builder.py (3)
175-179: LGTM! Proper initialization of file/line tracking state.The new fields are properly initialized with sensible defaults (
"<unknown>"for file/macro,0for line, empty list for stack).
219-223: Stack management looks correct, but consider exception safety.The push/pop of
macro_fileline_stackis balanced within the context manager. However, if an exception is raised during macro execution (betweenyieldat line 221 andpop()at line 223), the stack won't be popped. This might leave stale entries in the stack.Since this is primarily for debugging/error messages and exceptions typically unwind the entire builder context anyway, this may be acceptable. Consider if this needs a
try/finallyfor robustness.
690-693: LGTM!The
set_filelinemethod is a straightforward setter that updates the current file/line/macro context, called before each statement during macro execution.
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Fix all issues with AI Agents
In @tilelang/language/v2/ast.py:
- Around line 479-480: The generated f-strings embedding filename/name in ast.py
(the lines that set __tb_fl and __tb_fn) can break if self.filename or name
contain quotes/backslashes; change the code generation to embed a safely-escaped
Python literal (use repr(self.filename) and repr(name) or an equivalent escaping
helper) instead of raw f"'{...}'" so the produced source is always valid and
safe; update the code that emits "__tb_fl" and "__tb_fn" accordingly
(references: the AST code path that writes __tb_fl = ... and __tb_fn = ...).
- Line 648: The code passes Path(filename).name to DSLMutator which strips
directory context and produces ambiguous filenames in user-visible stack
messages; change the argument so DSLMutator receives a disambiguating path
(e.g., the relative path from the repo root or the absolute path) instead of
just the basename: locate the DSLMutator construction (DSLMutator(nonlocals,
func.__globals__, Path(filename).name)) and replace the third parameter with a
path that preserves directories (for example Path(filename).as_posix() or a
relative path computed from Path.cwd() or the project root) so get_stack_str()
shows a full/relative path like tilelang/language/v2/utils.py rather than just
utils.py.
🧹 Nitpick comments (1)
tilelang/language/v2/ast.py (1)
582-591: Consider using iterable unpacking for cleaner code.The concatenation
quote(...) + [node]works correctly but can be more idiomatically expressed using iterable unpacking.🔎 Proposed refactor using iterable unpacking
class SpanAttacher(ast.NodeTransformer): def __init__(self, filename_var: str, func_name_var: str): self.filename_var = filename_var self.func_name_var = func_name_var def visit(self, node: ast.AST): node = self.generic_visit(node) if isinstance(node, ast.stmt) and hasattr(node, 'lineno'): - return quote(f"__tb.set_fileline({self.filename_var}, {node.lineno}, {self.func_name_var})") + [node] + return [*quote(f"__tb.set_fileline({self.filename_var}, {node.lineno}, {self.func_name_var})"), node] return nodeBased on static analysis hints from Ruff.
📜 Review details
Configuration used: defaults
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
tilelang/language/v2/ast.py
🧰 Additional context used
🪛 Ruff (0.14.10)
tilelang/language/v2/ast.py
590-590: Consider iterable unpacking instead of concatenation
Replace with iterable unpacking
(RUF005)
🔇 Additional comments (3)
tilelang/language/v2/ast.py (3)
4-4: LGTM!The
Pathimport is correctly added to support filename extraction at line 648.
253-253: LGTM!The
filenameparameter is correctly added and stored for later use in generating file-level tracking context.Also applies to: 258-258
475-475: Theset_filelinemethod is already defined in theBuilderclass (line 690 intilelang/language/v2/builder.py) with the correct signature:def set_fileline(self, filename: str, lineno: int, name: str). SinceBuilderextendsBaseBuilder, the method is accessible to the generated code at line 475 of ast.py that calls__tb.set_fileline(...)through theSpanAttacher.Likely an incorrect or invalid review comment.
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (1)
tilelang/language/v2/ast.py (1)
582-592: Consider iterable unpacking for cleaner syntax.The
SpanAttacherclass correctly instruments statements with span tracking. However, line 590 uses list concatenation where iterable unpacking would be more idiomatic.🔎 Proposed refactor using iterable unpacking
def visit(self, node: ast.AST): node = self.generic_visit(node) if isinstance(node, ast.stmt) and hasattr(node, "lineno"): - return quote(f"__tb.set_fileline({self.filename_var}, {node.lineno}, {self.func_name_var})") + [node] + return [*quote(f"__tb.set_fileline({self.filename_var}, {node.lineno}, {self.func_name_var})"), node] return nodeBased on static analysis hints.
📜 Review details
Configuration used: defaults
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
tilelang/language/v2/ast.py
🧰 Additional context used
🪛 Ruff (0.14.10)
tilelang/language/v2/ast.py
590-590: Consider iterable unpacking instead of concatenation
Replace with iterable unpacking
(RUF005)
⏰ 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 Metal (on macos-latest)
- GitHub Check: Test for Python 3.12 with Nightly-ROCm-7.1 (on self-hosted-amd)
- GitHub Check: Test for Python 3.12 with CUDA-12.8 (on self-hosted-nvidia)
🔇 Additional comments (3)
tilelang/language/v2/ast.py (3)
4-4: LGTM!The import additions are appropriate:
Pathis used for filename handling at line 648, and theParamSpeccompatibility layer follows standard patterns for Python 3.9 support.Also applies to: 11-13
253-258: LGTM!The addition of the
filenameparameter toDSLMutator.__init__correctly enables per-file tracking for stack traces in device assertions.
474-475: LGTM!Capturing the function name locally before applying
SpanAttacheris correct, as the AST transformation might modify node properties. TheSpanAttacherapplication properly instruments the function body with per-statement span tracking.
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (2)
tilelang/language/v2/builder.py (2)
690-692: Consider adding parameter validation and documentation.The method works correctly for typical usage, but could be more robust:
- The
stacklevelparameter lacks documentation explaining its behavior (how many frames to trim from the top).- No validation for edge cases (e.g.,
stacklevel <= 0orstacklevel > len(stack)).- Static analysis suggests using iterable unpacking instead of list concatenation for better performance.
🔎 Optional improvements
def get_fileline_stack(self, stacklevel=1): + """Get the file/line call stack up to the specified level. + + Args: + stacklevel: Number of frames to include from bottom (1 = all frames). + + Returns: + List of (filename, lineno, macro_name) tuples representing the call stack. + """ - stack = self.macro_fileline_stack + [(self.current_file, self.current_line, self.current_macro_name)] + stack = [*self.macro_fileline_stack, (self.current_file, self.current_line, self.current_macro_name)] return stack[: len(stack) - stacklevel + 1]
694-699: Consider removing commented-out code.The commented-out
get_fileline_stack_msgmethod adds maintenance overhead. If this functionality is needed, consider either:
- Implementing and documenting it as an active method
- Removing it entirely to keep the codebase clean
If it's temporarily kept for reference during development, consider adding a TODO comment explaining when it will be removed or activated.
📜 Review details
Configuration used: defaults
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
tilelang/language/v2/builder.py
🧰 Additional context used
🧬 Code graph analysis (1)
tilelang/language/v2/builder.py (2)
tilelang/language/kernel.py (1)
pop(38-45)tilelang/language/frame.py (1)
pop(36-50)
🪛 Ruff (0.14.10)
tilelang/language/v2/builder.py
691-691: Consider iterable unpacking instead of concatenation
Replace with iterable unpacking
(RUF005)
⏰ 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). (1)
- GitHub Check: Test for Python 3.12 with Metal (on macos-latest)
🔇 Additional comments (4)
tilelang/language/v2/builder.py (4)
176-180: LGTM! Initialization of file/line tracking fields.The new fields are properly initialized with sensible defaults, and the comment clearly explains the purpose of the stack.
220-220: LGTM! Macro call stack tracking.The push operation correctly records the caller's file/line/macro context before entering a macro frame. This integrates well with the contextmanager pattern to ensure balanced stack operations.
224-224: LGTM! Stack cleanup on macro exit.The pop operation properly restores the call stack when exiting a macro frame.
685-688: LGTM! Simple setter for file/line context.The method correctly updates the current file/line/macro tracking fields.
Removed commented-out method get_fileline_stack_msg from builder.py.
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (2)
tilelang/language/v2/builder.py (2)
690-692: Consider adding bounds validation for thestacklevelparameter.If
stacklevelexceedslen(stack), the slicing could return an empty list or unexpectedly truncated results. Consider adding validation to ensurestacklevelis within reasonable bounds (e.g.,1 <= stacklevel <= len(stack)), or document the expected behavior for out-of-range values.🛡️ Proposed bounds checking
def get_fileline_stack(self, stacklevel=1): stack = self.macro_fileline_stack + [(self.current_file, self.current_line, self.current_macro_name)] + if stacklevel < 1: + stacklevel = 1 + if stacklevel > len(stack): + stacklevel = len(stack) return stack[: len(stack) - stacklevel + 1]
691-691: Consider iterable unpacking for clearer intent.As suggested by static analysis, you could use iterable unpacking instead of list concatenation for slightly better performance and clarity.
♻️ Proposed refactor using iterable unpacking
- stack = self.macro_fileline_stack + [(self.current_file, self.current_line, self.current_macro_name)] + stack = [*self.macro_fileline_stack, (self.current_file, self.current_line, self.current_macro_name)]
📜 Review details
Configuration used: defaults
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
tilelang/language/v2/builder.py
🧰 Additional context used
🧬 Code graph analysis (1)
tilelang/language/v2/builder.py (2)
tilelang/language/kernel.py (1)
pop(38-45)tilelang/language/frame.py (1)
pop(36-50)
🪛 Ruff (0.14.10)
tilelang/language/v2/builder.py
691-691: Consider iterable unpacking instead of concatenation
Replace with iterable unpacking
(RUF005)
⏰ 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). (1)
- GitHub Check: Test for Python 3.12 with Metal (on macos-latest)
🔇 Additional comments (3)
tilelang/language/v2/builder.py (3)
176-180: LGTM! Clear initialization of macro call stack tracking fields.The new fields support tracking file/line/macro context through macro expansion. The default values are appropriate sentinel values, and the comment on line 179 helpfully clarifies that the stack tracks caller context.
220-224: LGTM! Exception-safe push/pop pattern for macro call stack tracking.The
@contextmanagerdecorator ensures line 224 executes even if exceptions occur within the macro body, maintaining stack consistency. The balanced push before entry (line 220) and pop after exit (line 224) correctly tracks the macro invocation context.
685-688: LGTM! Simple and clear setter for tracking current location context.
2 participants
as title.
Summary by CodeRabbit
New Features
Chores
✏️ Tip: You can customize this high-level summary in your review settings.