fix(codegen): propagate tensor view across plain tensor aliases#1606
Conversation
|
Important Review skippedAuto incremental reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
📝 WalkthroughWalkthroughThis PR extends PTOCodegen's assignment visitor to propagate tensor views through plain alias assignments, allowing downstream tile.store operations to resolve and lower without explicit tensor-view creation. Tests validate view propagation for variable and loop iter_arg sources. ChangesTensor Alias View Propagation
🎯 2 (Simple) | ⏱️ ~12 minutes Possibly related PRs
Suggested labels
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. 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 |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request implements support for plain tensor aliases in the PTO codegen. When a loop-result ForStmt is folded into a plain AssignStmt (e.g., lhs_tensor = rhs_var), the AssignStmt visitor now propagates the tensor view, SSA name, and base pointer from the RHS to the LHS. This prevents downstream failures during subsequent tile.store operations. Two regression tests have been added to verify this behavior for both plain variable and loop IterArg RHS aliases. There are no review comments, so no feedback is provided.
There was a problem hiding this comment.
Pull request overview
This PR fixes a PTO codegen gap where AssignStmt did not propagate tensor-view metadata across plain tensor aliases (lhs_tensor = rhs_var). That gap shows up when an empty loop-result ForStmt is simplified into a plain alias (notably after certain constant-trip pl.pipeline folds), causing later tile.store codegen to fail because the alias has no registered tensor view.
Changes:
- Extend
PTOCodegen::VisitStmt_(AssignStmtPtr)to detect plain tensor aliases and bind tensor view / SSA name / base ptr from RHS to LHS (using non-fatalTryGetTensorView). - Add two regression tests covering the plain-
VarRHS case and theIterArgRHS case (loop-carried tensor alias). - Update the
TryGetTensorViewheader comment to include the new alias-propagation use case.
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 1 comment.
| File | Description |
|---|---|
src/codegen/pto/pto_codegen.cpp |
Adds alias-detection logic in AssignStmt visitor to propagate tensor-view bindings for later consumers (e.g., tile.store). |
include/pypto/codegen/pto/pto_codegen.h |
Expands TryGetTensorView doc comment to cover alias propagation as a non-fatal lookup use case. |
tests/ut/codegen/test_pto_codegen.py |
Adds regression tests ensuring tile.store resolves tensor views when its output tensor is a plain alias (Var RHS and IterArg RHS). |
PTO codegen registers loop-result tensor views only inside the ForStmt
visitor. When a constant-trip stage-2 pl.pipeline's statically-empty
main loop is folded by Simplify into a plain alias `t__rv = t__iter`,
that registration is bypassed and the AssignStmt visitor propagated
nothing for tensor aliases. A later tile.store into the alias then
tripped GetOrCreateTensorView's INTERNAL_CHECK ("Tensor view not found
for parameter: <name>"), an opaque failure pointing at a synthetic var.
Propagate the tensor view / SSA name / base ptr from RHS to LHS for a
plain `lhs_tensor = rhs_var` alias, mirroring the ForStmt loop-result
registration. Resolution uses the non-fatal TryGetTensorView and falls
through on a view miss, so the change is strictly additive (never
introduces a new throw).
Add two regression tests covering the plain-Var and IterArg RHS legs;
both fail without the fix with the exact "Tensor view not found" error.
Hzfengsy
force-pushed
the
fix/pto-codegen-tensor-alias-view
branch
from
a7583f3 to
e7cf6fe
Compare
VisitStmt_(ForStmtPtr) registers the loop-result tensor view / SSA name for its return var but does not call RegisterBasePtr; the prior comment overstated it. Reword to match, and note that this branch additionally propagates the base-ptr mapping for element-wise alias consumers (pl.read / pl.write / store_scalar), as VisitStmt_(IfStmtPtr) does for merged tensors. Comment-only; no behavior change. Addresses review feedback on hw-native-sys#1606.
Summary
PTO codegen registers loop-result tensor views only inside the
ForStmtvisitor. When a constant-trip stage-2pl.pipeline's statically-empty main loop is folded bySimplifyinto a plain aliast__rv = t__iter, that registration is bypassed and theAssignStmtvisitor propagated nothing for tensor aliases. A latertile.storeinto the alias then trippedGetOrCreateTensorView'sINTERNAL_CHECKwith an opaqueTensor view not found for parameter: <synthetic-var>, giving no hint at the real cause.This adds a branch to
PTOCodegen::VisitStmt_(AssignStmtPtr)that, for a plainlhs_tensor = rhs_varalias, propagates the tensor view / SSA name / base ptr from RHS to LHS — mirroring theForStmtloop-result registration the fold bypasses. Resolution uses the non-fatalTryGetTensorViewand falls through on a view miss, so the change is strictly additive (never introduces a new throw).Testing
tests/ut/codegen/test_pto_codegen.pycovering the plain-VarandIterArgRHS legs (the latter is the exact post-foldt__rv = t__itershape). Both fail without the fix with the documented "Tensor view not found" error and pass with it.tests/ut/suite: 5382 passed, 28 skipped.Notes
TryGetTensorViewhelper already exists upstream; this change only adds a new caller.