fix(pass): store iter-arg returns to InOut params in ConvertTensorToTileOps#783
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdded utilities to detect/resolve tile-alias chains and locate yields, introduced an analysis to derive per-InCore iter-argument mappings from orchestration loops, extended TransformIncoreFunction to accept these mappings and track merged return indices, and added a store-sinking phase that inserts tile.store into IfStmt branches and updates returns accordingly. Changes
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
Suggested reviewers
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. 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 |
There was a problem hiding this comment.
Code Review
This pull request introduces an optimization to the ConvertTensorToTileOps pass by sinking tile.store operations into IfStmt branches when return values correspond to loop iteration arguments. It adds a new analysis phase to map InCore function returns to call arguments and includes utilities for resolving tile alias chains and manipulating statement lists. Feedback was provided regarding an inconsistency in the documentation for FindYieldStmt, which should be updated to reflect its recursive implementation.
f49d451 to
b11bf7c
Compare
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@src/ir/transforms/convert_tensor_to_tile_ops_pass.cpp`:
- Around line 1634-1642: The code currently picks last_if_stmt and
unconditionally replaces its result with fresh return_vars but only updates the
final ReturnStmt via tensor_to_tile, leaving any subsequent statements
referencing the old IfStmt result broken; fix by only sinking (modifying
last_if_stmt, last_if_index, and generating new_if_stmt) when that IfStmt is the
terminal non-return statement (i.e., last_if_index == new_stmts.size()-1 or all
following statements are safe returns), otherwise after creating new_if_stmt
perform a suffix rewrite that substitutes old_rv -> new_rv in every statement
after last_if_index (use the same tensor_to_tile substitution logic you use for
the ReturnStmt) so no later statement keeps referencing the dead/typed-out old
IfStmt; refer to last_if_stmt, last_if_index, new_stmts, new_if_stmt,
tensor_to_tile, ReturnStmt, and old_rv/new_rv when implementing the checks and
suffix substitution.
- Around line 466-468: The code currently stores a single mapping into
result[gvar->name_] (when discovering mapping for an InCore call inside a
ForStmt), so the last loop wins and can mis-map loop-carried args across
different call sites; change the assignment logic to detect conflicting mappings
for the same InCore function: when you find a mapping for gvar->name_, compare
it to any existing mapping in result[gvar->name_]; if they match do nothing, if
they differ mark that function as ambiguous (e.g., store a special sentinel or a
skip flag) and do not apply sinking for that callee in TransformIncoreFunction;
apply the same compare-or-mark-ambiguous change to the other occurrence handling
mappings (the block around the second occurrence noted in the review) so all
discovered mappings for a given InCore function must match before enabling store
sinking.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 5c4d4389-5ae4-46b8-a3f5-9237f9d1177e
📒 Files selected for processing (1)
src/ir/transforms/convert_tensor_to_tile_ops_pass.cpp
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (2)
src/ir/transforms/convert_tensor_to_tile_ops_pass.cpp (2)
1632-1640:⚠️ Potential issue | 🔴 CriticalOnly sink when the rewritten
IfStmtis terminal, or rewrite suffix uses.The transform rewrites
last_if_stmtin place, but statements after Line 1750 are not re-substituted from old return vars to new ones. If the chosen IfStmt is non-terminal, downstream users can reference stale vars/types.Minimal safe guard
- if (!sink_candidates.empty()) { + if (!sink_candidates.empty() && last_if_index + 1 == new_stmts.size()) {Also applies to: 1748-1751
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/ir/transforms/convert_tensor_to_tile_ops_pass.cpp` around lines 1632 - 1640, The current sink picks the last IfStmt (last_if_stmt) from new_stmts without ensuring it is terminal; only sink (rewrite) when that IfStmt is terminal (no statements after its index) or else update the trailing statements to substitute old return variables/types to the newly rewritten ones; locate the selection logic that builds last_if_stmt/last_if_index from new_stmts and either (a) add a guard that skips sinking unless last_if_index == new_stmts.size()-1, or (b) after rewriting IfStmt, walk new_stmts from last_if_index+1 onward and apply the same return-var/type remapping used for the IfStmt so downstream uses are updated.
466-468:⚠️ Potential issue | 🟠 MajorHandle conflicting iter-arg mappings per callee instead of last-writer-wins.
At Line 467, a later loop can overwrite an earlier mapping for the same InCore callee. If call-site mappings differ, sinking can target the wrong parameter for one caller.
Suggested fix
std::unordered_map<std::string, IterArgMapping> AnalyzeIterArgMappings( const std::vector<FunctionPtr>& functions) { std::unordered_map<std::string, IterArgMapping> result; + std::unordered_set<std::string> ambiguous_callees; @@ - if (!mapping.empty()) { - result[gvar->name_] = std::move(mapping); + if (!mapping.empty()) { + if (ambiguous_callees.count(gvar->name_) > 0) { + break; + } + auto it_existing = result.find(gvar->name_); + if (it_existing == result.end()) { + result.emplace(gvar->name_, mapping); + } else if (it_existing->second != mapping) { + result.erase(it_existing); + ambiguous_callees.insert(gvar->name_); + } break; // Found the right InCore call for this ForStmt }
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@src/ir/transforms/convert_tensor_to_tile_ops_pass.cpp`:
- Around line 368-374: The loop currently skips traversing ForStmt bodies when
for_stmt->iter_args_ is empty, causing nested loops/calls to be missed; change
the control flow so that if stmt is a ForStmt you always push its body (e.g.,
worklist.push_back(FlattenToStmts(for_stmt->body_))) into the worklist even when
for_stmt->iter_args_.empty(), and only skip recursion/continue for non-ForStmt
cases (still handle WhileStmt as before); update the conditional around
As<ForStmt>(stmt) and the use of iter_args_ to ensure ForStmt bodies are always
traversed.
---
Duplicate comments:
In `@src/ir/transforms/convert_tensor_to_tile_ops_pass.cpp`:
- Around line 1632-1640: The current sink picks the last IfStmt (last_if_stmt)
from new_stmts without ensuring it is terminal; only sink (rewrite) when that
IfStmt is terminal (no statements after its index) or else update the trailing
statements to substitute old return variables/types to the newly rewritten ones;
locate the selection logic that builds last_if_stmt/last_if_index from new_stmts
and either (a) add a guard that skips sinking unless last_if_index ==
new_stmts.size()-1, or (b) after rewriting IfStmt, walk new_stmts from
last_if_index+1 onward and apply the same return-var/type remapping used for the
IfStmt so downstream uses are updated.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: f8c4b7f5-3cdc-4c4f-a829-6181d665b3ad
📒 Files selected for processing (1)
src/ir/transforms/convert_tensor_to_tile_ops_pass.cpp
b11bf7c to
46048ff
Compare
When an InCore function returns tile values that feed back as iter-args in a ForStmt loop, the pass now stores to the existing In param (auto-promoted to InOut by UpgradeWrittenTensorParamDirections) instead of adding new Out params. The tile.store is placed outside the IfStmt, referencing the phi variable directly. Supporting fixes in downstream passes: - InitMemRef: tile alias assignments (a = b) now share the source's MemRef instead of allocating a fresh one, preventing empty-memref aliases in IfStmt branches - MemoryReuse YieldFixupMutator: when IfStmt branches yield tiles with different MemRefs, insert tile.move in the else branch to unify to the then-branch's canonical MemRef (mirroring the existing ForStmt tile.move pattern) Key additions: - AnalyzeIterArgMappings: scans orchestration functions to map InCore return indices to their corresponding call arg positions via yield/iter-arg tracing - TransformIncoreFunction Phase 3: for returns with iter-arg mappings, tile.store targets the In param directly (no new Out param, no tensor.create at call site)
46048ff to
871edd6
Compare
Summary
tile.storenow targets the existing In param (auto-promoted to InOut) instead of adding new Out params. The store is placed outside the IfStmt, referencing the phi variable directly — no store sinking into branches needed.a = b) now share the source's MemRef instead of allocating a fresh one, preventing empty-memref aliases in IfStmt branches.tile.moveis inserted in the else branch to unify to the then-branch's canonical MemRef (mirroring the existing ForStmt pattern).