[Language] Add missing while statement #1254
Conversation
WalkthroughThis PR adds while loop support to TileLang by introducing a comprehensive TIR IR type-stub file with 80+ function signatures, modifying the v2 AST visitor to recursively process While nodes, updating the builder to dynamically evaluate while loop conditions, and adding a test case. Changes
Sequence DiagramsequenceDiagram
participant User
participant AST as v2/ast.py<br/>(visit_While)
participant Builder as v2/builder.py<br/>(while handling)
participant TIR as TIR IR
User->>AST: Visit While node
AST->>AST: generic_visit(node)
Note over AST: Recursively transform<br/>While node children
AST->>Builder: Process While statement
Builder->>Builder: Evaluate condition
alt Condition → PrimExpr
Builder->>TIR: Emit while frame<br/>(valid primitive)
TIR-->>Builder: IR generated
else Non-PrimExpr truthy
Builder->>Builder: Raise infinite-loop error
else Non-PrimExpr falsy
Builder->>Builder: Log warning<br/>(skipped loop)
end
Builder-->>User: While loop handled
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes
Areas requiring extra attention:
Possibly related PRs
Suggested reviewers
Poem
Pre-merge checks and finishing touches❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✨ Finishing touches
🧪 Generate unit tests (beta)
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! 🚀 |
![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 (2)
tilelang/language/v2/ast.py (1)
471-477: While visitor now matches other control‑flow visitorsCalling
self.generic_visit(node)before emitting thectx_whiletransformation is correct and consistent with theIf/Forhandlers; it ensures nested expressions and bodies are mutated before IR is generated. You might optionally annotate the parameter asnode: ast.Whilefor consistency with other visit methods, but the behavior looks good as‑is.testing/python/language/test_tilelang_language_frontend_v2.py (1)
345-361: Strengthentest_while_loopby initializing the accumulator explicitlyThe new test covers the
whilesemantics nicely, but two small tweaks would make it more robust and readable:
sum = T.alloc_var(T.int32)does not document the initial value; depending onalloc_var’s default here is a bit implicit for a reference test. It’s safer and clearer to initialize explicitly:- sum = T.alloc_var(T.int32) + sum = T.alloc_var(T.int32, 0)
- (Optional) Using the same name
test_while_loopfor both the test function and the innerprim_funcworks but is a bit confusing at a glance; consider renaming the inner function to something likewhile_sum_kernelfor clarity.The rest of the test (JIT compilation, kernel invocation, and the
sum(range(10))check) looks solid.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (4)
testing/python/language/test_tilelang_language_frontend_v2.py(1 hunks)tilelang/language/tir/ir.pyi(1 hunks)tilelang/language/v2/ast.py(1 hunks)tilelang/language/v2/builder.py(1 hunks)
🧰 Additional context used
🧬 Code graph analysis (2)
tilelang/language/v2/builder.py (1)
tilelang/language/ast/ir.py (1)
While(1078-1093)
testing/python/language/test_tilelang_language_frontend_v2.py (2)
tilelang/jit/__init__.py (3)
jit(279-280)jit(284-295)jit(298-365)tilelang/language/v2/builder.py (2)
prim_func(151-155)prim_func(643-736)
🪛 Ruff (0.14.4)
tilelang/language/v2/builder.py
299-302: Avoid specifying long messages outside the exception class
(TRY003)
⏰ 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: Quick Lint
🔇 Additional comments (1)
tilelang/language/tir/ir.pyi (1)
1-106: New TIR IR typing stub looks reasonable; please verify a few signatures against TVMThe overall structure of this new
ir.pyistub (generic_T,Span | Nonespans, and the large set of math/TIR intrinsics) looks consistent and should give much better editor/type‑checker support.Given how large and API‑surface‑sensitive this stub is, it would be good to sanity‑check a handful of representative functions against the actual
tvm.tirAPI for the TVM version you depend on (e.g., return types ofvscale,tvm_tuple, and the varioustvm_*/ptx_*intrinsics, and the import location ofSpan). That will help avoid subtle type‑mismatch noise later on.If you’d like, I can sketch a small script to introspect
tvm.tirand compare callable signatures to this stub.
| def ctx_while(self, cond): | ||
| self.check_continue_break() | ||
| raise RuntimeError("while loops are not supported in TileLang builder") | ||
| cond_v = cond() | ||
| cond_v_unwrap = unwrap_cond(cond_v) | ||
| if not isinstance(cond_v_unwrap, PrimExpr): | ||
| if cond_v_unwrap: | ||
| raise RuntimeError( | ||
| f'Infinite while loop detected in TileLang\n' | ||
| f'Condition: {cond_v}({type(cond_v)}) => {cond_v_unwrap}({type(cond_v_unwrap)})\n' | ||
| ) | ||
| else: | ||
| logger.warning( | ||
| 'While loop with constant false condition detected in Tilelang, the loop body will never be executed.\n', | ||
| f'Condition: {cond_v}({type(cond_v)}) => {cond_v_unwrap}({type(cond_v_unwrap)})\n', | ||
| stack_info=True, | ||
| stacklevel=2) | ||
| with self.with_frame(tir.While(cond_v_unwrap)): | ||
| yield None | ||
|
|
There was a problem hiding this comment.
🧩 Analysis chain
Fix logging call in ctx_while to avoid runtime errors
The new ctx_while logic overall makes sense (evaluate cond() once to obtain a PrimExpr, detect obviously infinite loops, and then build a tir.While frame), but the warning branch for constant‑false conditions has a subtle bug:
logger.warning(
'While loop with constant false condition detected in Tilelang, the loop body will never be executed.\n',
f'Condition: {cond_v}({type(cond_v)}) => {cond_v_unwrap}({type(cond_v_unwrap)})\n',
stack_info=True,
stacklevel=2)logging treats the second positional argument as *args for msg % args formatting. Since the format string contains no % placeholders, this will raise TypeError: not all arguments converted during string formatting when this branch is hit.
You can fix this by folding the condition detail into the message string and passing only keyword arguments:
- else:
- logger.warning(
- 'While loop with constant false condition detected in Tilelang, the loop body will never be executed.\n',
- f'Condition: {cond_v}({type(cond_v)}) => {cond_v_unwrap}({type(cond_v_unwrap)})\n',
- stack_info=True,
- stacklevel=2)
+ else:
+ logger.warning(
+ 'While loop with constant false condition detected in TileLang; '
+ f'the loop body will never be executed.\n'
+ f'Condition: {cond_v}({type(cond_v)}) => {cond_v_unwrap}({type(cond_v_unwrap)})\n',
+ stack_info=True,
+ stacklevel=2,
+ )Everything else in ctx_while looks consistent with the existing builder design (use of unwrap_cond, check_continue_break, and tir.While).
I'll verify the review comment by examining the actual code and confirming the logging bug.
# Verify Python logging behavior with extra positional args
import logging
# Test the exact scenario from the review
test_logger = logging.getLogger("test")
handler = logging.StreamHandler()
test_logger.addHandler(handler)
test_logger.setLevel(logging.WARNING)
# This should demonstrate the bug
try:
test_logger.warning(
'First message\n',
f'Second message\n',
stack_info=True,
stacklevel=2
)
print("No error occurred - bug description may be incorrect")
except TypeError as e:
print(f"TypeError caught as expected: {e}")🏁 Script executed:
#!/bin/bash
# Extract the ctx_while method to verify the exact code
fd -t f "builder.py" | head -1 | xargs -I {} sed -n '293,311p' {}Length of output: 870
🏁 Script executed:
#!/bin/bash
# Search for ctx_while in the file to understand full context
fd -t f "builder.py" | head -1 | xargs -I {} grep -n "def ctx_while" {}Length of output: 34
🏁 Script executed:
#!/bin/bash
# Get a broader view of the method and surrounding code
fd -t f "builder.py" | head -1 | xargs -I {} sed -n '290,315p' {}Length of output: 1129
I need to verify the Python logging behavior. Let me run the test:
#!/usr/bin/env python3
# Verify Python logging behavior with extra positional args
import logging
import sys
# Set up logging to see the actual error
logging.basicConfig(level=logging.WARNING, stream=sys.stdout)
test_logger = logging.getLogger("test")
# Simulate the exact scenario from the code
cond_v = "some_condition"
cond_v_unwrap = False
try:
test_logger.warning(
'While loop with constant false condition detected in Tilelang, the loop body will never be executed.\n',
f'Condition: {cond_v}({type(cond_v)}) => {cond_v_unwrap}({type(cond_v_unwrap)})\n',
stack_info=True,
stacklevel=2
)
print("SUCCESS: No error occurred - the logging call worked")
except TypeError as e:
print(f"CONFIRMED BUG: TypeError caught: {e}")
except Exception as e:
print(f"OTHER ERROR: {type(e).__name__}: {e}")Fix logging call in ctx_while to avoid runtime errors
The code has been verified and the issue is confirmed. The logger.warning() call passes two positional string arguments:
logger.warning(
'While loop with constant false condition detected in Tilelang, the loop body will never be executed.\n',
f'Condition: {cond_v}({type(cond_v)}) => {cond_v_unwrap}({type(cond_v_unwrap)})\n',
stack_info=True,
stacklevel=2)Python's logging module treats extra positional arguments as formatting arguments (via msg % args). Since the first message string contains no % placeholders, this will raise TypeError: not all arguments converted during string formatting when the constant-false condition branch is executed.
The suggested fix combines the messages into a single string and uses only keyword arguments:
else:
- logger.warning(
- 'While loop with constant false condition detected in Tilelang, the loop body will never be executed.\n',
- f'Condition: {cond_v}({type(cond_v)}) => {cond_v_unwrap}({type(cond_v_unwrap)})\n',
- stack_info=True,
- stacklevel=2)
+ logger.warning(
+ 'While loop with constant false condition detected in TileLang; '
+ f'the loop body will never be executed.\n'
+ f'Condition: {cond_v}({type(cond_v)}) => {cond_v_unwrap}({type(cond_v_unwrap)})\n',
+ stack_info=True,
+ stacklevel=2,
+ )🧰 Tools
🪛 Ruff (0.14.4)
299-302: Avoid specifying long messages outside the exception class
(TRY003)
🤖 Prompt for AI Agents
In tilelang/language/v2/builder.py around lines 293 to 311, the logger.warning
call passes two positional string arguments which the logging module interprets
as format args and can raise a TypeError; combine the two message parts into a
single formatted string (e.g. one f-string that includes the condition and
types) and call logger.warning with that single message and the existing keyword
arguments stack_info=True and stacklevel=2 so no extra positional args are
passed.
2 participants
This pr allow
while condin tilelang functionSummary by CodeRabbit
New Features
Tests