[Language] Support Consequential assignments like 'a = b = c = 1' #992
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! 🚀 |
WalkthroughAdds a new CUDA-backed TileLang test exercising chained equal assignments, and extends the TileLang parser override to handle Assign nodes with chained writes and local.var buffers, including subscript handling and read-then-write semantics. No existing public APIs are modified beyond introducing new test helper functions and a new Assign visitor. Changes
Sequence Diagram(s)sequenceDiagram
autonumber
actor User
participant Parser as Parser Override
participant VarTbl as Variable Table
participant Eval as RHS Evaluator
participant Buf as Buffer Ops
User->>Parser: Visit Assign (targets, value)
Parser->>Parser: Validate targets / normalize slices
Parser->>Eval: Evaluate RHS once
loop For each LHS target
alt LHS is Subscript
Parser->>Buf: buffer_store(base, indices, RHS)
else LHS is Name and in VarTbl
Parser->>Buf: buffer_load(var) Note over Parser,Buf: Read-then-write semantics
alt Loaded is local.var BufferLoad with idx == [0]
Parser->>Buf: buffer_store(local.var, [0], RHS)
else
Parser->>Parser: General in-place update path
end
else
Parser->>Parser: Fallback eval_assign(LHS, RHS)
end
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
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 |
![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 (3)
tilelang/language/overrides/parser.py (1)
49-55: Align indices handling with AugAssign (wrap non-tuple slice in list).AugAssign wraps the single index in a list; Assign currently passes a scalar. While
buffer_storecoerces scalars, keeping both consistent improves readability and avoids surprises.- else: - indices = self.eval_expr(lhs.slice) + else: + indices = [self.eval_expr(lhs.slice)]testing/python/language/test_tilelang_laguange_chain_equal.py (2)
30-32: Honor dtype parameter when allocating tensors.
dtypearg is ignored; tensors are hardcoded totorch.float32. Use the provided dtype for correctness and future coverage.- A = torch.zeros((N,), dtype=torch.float32, device="cuda") - B = torch.zeros((N,), dtype=torch.float32, device="cuda") - C = torch.zeros((N,), dtype=torch.float32, device="cuda") + torch_dtype = getattr(torch, dtype) # e.g., "float32" -> torch.float32 + A = torch.zeros((N,), dtype=torch_dtype, device="cuda") + B = torch.zeros((N,), dtype=torch_dtype, device="cuda") + C = torch.zeros((N,), dtype=torch_dtype, device="cuda")
1-46: Nit: fix filename typo for test discoverability and consistency.Consider renaming
test_tilelang_laguange_chain_equal.py->test_tilelang_language_chain_equal.py.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
testing/python/language/test_tilelang_laguange_chain_equal.py(1 hunks)tilelang/language/overrides/parser.py(1 hunks)
🧰 Additional context used
🧬 Code graph analysis (2)
testing/python/language/test_tilelang_laguange_chain_equal.py (2)
tilelang/jit/__init__.py (1)
jit(244-317)tilelang/transform/pass_config.py (1)
PassConfigKey(6-104)
tilelang/language/overrides/parser.py (2)
tilelang/language/ast/ir.py (1)
buffer_store(1263-1300)tilelang/language/parser/parser.py (1)
bind_assign_value(114-160)
🪛 Ruff (0.13.3)
tilelang/language/overrides/parser.py
20-20: Possible hardcoded password assigned to argument: "token"
(S106)
| if isinstance(node.value, doc.Subscript): | ||
| check_slices = [] | ||
| if isinstance(node.value.slice, doc.Slice): | ||
| check_slices = [node.value.slice] | ||
| elif isinstance(node.value.slice, doc.Tuple): | ||
| for part in node.value.slice.elts: | ||
| if isinstance(part, doc.Slice): | ||
| check_slices.append(part) | ||
| for slice_node in check_slices: | ||
| if not slice_node.step and slice_node.upper and slice_node.lower: | ||
| slice_node.step = doc.Constant( | ||
| 1, | ||
| None, | ||
| 1, | ||
| 1, | ||
| slice_node.upper.lineno, | ||
| slice_node.upper.end_col_offset + 1, | ||
| slice_node.upper.lineno, | ||
| slice_node.upper.end_col_offset + 2, | ||
| ) | ||
|
|
There was a problem hiding this comment.
doc.Constant constructed with invalid arg list (likely runtime error).
The doc.Constant(...) call passes 9 positional args. TVM’s doc nodes expect Constant(value, kind, lineno, col, end_lineno, end_col) (value + kind + 4 positions). The extra three 1 literals make the signature invalid and will crash when this path is hit.
Fix by removing the extraneous args and keeping only the positional span:
- slice_node.step = doc.Constant(
- 1,
- None,
- 1,
- 1,
- 1,
- slice_node.upper.lineno,
- slice_node.upper.end_col_offset + 1,
- slice_node.upper.lineno,
- slice_node.upper.end_col_offset + 2,
- )
+ slice_node.step = doc.Constant(
+ 1,
+ None,
+ slice_node.upper.lineno,
+ slice_node.upper.end_col_offset + 1,
+ slice_node.upper.lineno,
+ slice_node.upper.end_col_offset + 2,
+ )🤖 Prompt for AI Agents
In tilelang/language/overrides/parser.py around lines 26 to 46, the
doc.Constant(...) call is constructed with nine positional arguments (including
three extra numeric 1s) which does not match the expected signature and will
raise at runtime; replace the call so it only passes value, kind, lineno, col,
end_lineno, end_col (e.g. value=1, kind=None, lineno=slice_node.upper.lineno,
col=slice_node.upper.end_col_offset+1, end_lineno=slice_node.upper.lineno,
end_col=slice_node.upper.end_col_offset+2) — remove the extraneous 1 literals
and supply the six correct arguments in that order.
…le-ai#992) * chained assignments * test update * [Lint]: [pre-commit.ci] auto fixes [...] --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2 participants
This pull request introduces a new override for handling assignment statements in the
tilelangparser, specifically to support advanced assignment features such as chained writes andlocal.varbuffers. The main change is the addition of a customtilelang_visit_assignmethod for theAssignnode type, which enhances assignment handling to align with the requirements of the tilelang language.Assignment node handling improvements:
tilelang_visit_assignmethod to override assignment behavior, supporting chained writes and assignments tolocal.varbuffers. This includes logic for slice handling, buffer stores, and context-aware assignment to variables and buffer indices.Summary by CodeRabbit
New Features
Tests