feat(ir): fix ForStmt loop-carry MemRef handling in InitMemRef and BasicMemoryReuse#622
Conversation
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request addresses a critical aspect of memory management within loop structures by refining how memory references ( Highlights
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here. Footnotes
|
|
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:
📝 WalkthroughWalkthroughInitMemRef now derives ForStmt return_var MemRefs from the loop body Yield values and exposes a shared FindYieldStmt helper; BasicMemoryReuse adds a ForStmtYieldFixupMutator that inserts Changes
Sequence DiagramsequenceDiagram
participant ForStmt
participant InitMemRef
participant BasicMemoryReuse
participant YieldFixup as ForStmtYieldFixupMutator
ForStmt->>InitMemRef: Visit ForStmt
Note over InitMemRef: Locate YieldStmt in body
InitMemRef->>InitMemRef: Set initValue & iter_arg MemRef (group A)
InitMemRef->>InitMemRef: Set yield & return_var MemRef (group B)
InitMemRef-->>ForStmt: Emit mutated ForStmt
ForStmt->>BasicMemoryReuse: Process ForStmt
BasicMemoryReuse->>BasicMemoryReuse: ApplyMemRefSharing
BasicMemoryReuse->>YieldFixup: Check iter_arg vs yield MemRef
alt Divergent MemRefs
YieldFixup->>YieldFixup: Create moved Var with initValue MemRef
YieldFixup->>YieldFixup: Insert tile.move (AssignStmt) before YieldStmt
YieldFixup->>YieldFixup: Replace YieldStmt values with moved Vars
YieldFixup->>YieldFixup: Patch iter_args & return_vars to unified MemRef
else Aligned or no moves needed
YieldFixup->>YieldFixup: Patch iter_args & return_vars types as needed
end
YieldFixup-->>ForStmt: Return ForStmt with unified loop-carry MemRefs
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related issues
Possibly related PRs
Suggested reviewers
🚥 Pre-merge checks | ✅ 3✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. 📝 Coding Plan
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 Tip You can customize the high-level summary generated by CodeRabbit.Configure the |
![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 significantly improves the handling of MemRef assignments for ForStmt loop-carry variables within the InitMemRef and BasicMemoryReuse passes. The refactoring in InitMemRefMutator provides more precise control over MemRef inheritance, ensuring iter_args inherit from initValue and return_vars from yield values. The introduction of ForStmtYieldFixupMutator in BasicMemoryReuse is a critical addition, correctly resolving potential MemRef mismatches between iter_arg and yield value by inserting tile.move operations. The extraction of FindYieldStmt into a shared utility header promotes code reuse. The updated documentation clearly explains these complex interactions, and the new test cases provide excellent validation for the revised MemRef relationships. Overall, these changes enhance the correctness and robustness of memory management within loop constructs.
![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
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
src/ir/transforms/basic_memory_reuse_pass.cpp (1)
915-924:⚠️ Potential issue | 🔴 CriticalRun the loop-carry fixup even when generic reuse finds nothing.
Lines 915-917 return before
ForStmtYieldFixupMutatorand alloc cleanup run. That skips the new loop-only repair path for cases whereyieldstill differs frominitValue/iter_argbut no ordinary reuse candidate exists.💡 Minimal fix
- if (reuse_map.empty()) { - return func; - } - - // Step 4: Apply MemRef sharing - StmtPtr new_body = ApplyMemRefSharing(func->body_, reuse_map, analysis_result.var_sharing_groups); + StmtPtr new_body = func->body_; + if (!reuse_map.empty()) { + new_body = ApplyMemRefSharing(func->body_, reuse_map, analysis_result.var_sharing_groups); + }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/ir/transforms/basic_memory_reuse_pass.cpp` around lines 915 - 924, The early return when reuse_map.empty() skips loop-carry repairs and alloc cleanup; change control flow so that even when reuse_map is empty you still run ForStmtYieldFixupMutator and the alloc cleanup path on func->body_. Specifically, remove or bypass the return that exits when reuse_map.empty(), then always set StmtPtr new_body = (reuse_map.empty() ? func->body_ : ApplyMemRefSharing(func->body_, reuse_map, analysis_result.var_sharing_groups)); then invoke ForStmtYieldFixupMutator yield_fixup; new_body = yield_fixup.VisitStmt(new_body); and proceed with the existing alloc cleanup logic so loop-only yield/iter_arg mismatches are fixed.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@docs/en/dev/passes/13-basic_memory_reuse.md`:
- Around line 47-48: Update the docs to reflect the actual pass order: move the
"ForStmt yield fixup" bullet so ForStmtYieldFixupMutator is listed before the
"Remove redundant allocs"/RemoveUnusedAllocStatements step (the C++
implementation in src/ir/transforms/basic_memory_reuse_pass.cpp runs
ForStmtYieldFixupMutator first), and make the same reordering in the
corresponding zh-CN page (also adjust the duplicate occurrence around lines
141-142 in the English doc to match); ensure the bullets reference the same step
names (ForStmtYieldFixupMutator and RemoveUnusedAllocStatements/Remove redundant
allocs) so documentation matches the code.
---
Outside diff comments:
In `@src/ir/transforms/basic_memory_reuse_pass.cpp`:
- Around line 915-924: The early return when reuse_map.empty() skips loop-carry
repairs and alloc cleanup; change control flow so that even when reuse_map is
empty you still run ForStmtYieldFixupMutator and the alloc cleanup path on
func->body_. Specifically, remove or bypass the return that exits when
reuse_map.empty(), then always set StmtPtr new_body = (reuse_map.empty() ?
func->body_ : ApplyMemRefSharing(func->body_, reuse_map,
analysis_result.var_sharing_groups)); then invoke ForStmtYieldFixupMutator
yield_fixup; new_body = yield_fixup.VisitStmt(new_body); and proceed with the
existing alloc cleanup logic so loop-only yield/iter_arg mismatches are fixed.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: dc947204-7489-4e28-a8f2-9dd293282c49
📒 Files selected for processing (8)
docs/en/dev/passes/12-init_memref.mddocs/en/dev/passes/13-basic_memory_reuse.mddocs/zh-cn/dev/passes/12-init_memref.mddocs/zh-cn/dev/passes/13-basic_memory_reuse.mdinclude/pypto/ir/transforms/utils/memref_utils.hsrc/ir/transforms/basic_memory_reuse_pass.cppsrc/ir/transforms/init_memref.cpptests/ut/ir/transforms/test_init_memref.py
There was a problem hiding this comment.
🧹 Nitpick comments (1)
tests/ut/ir/transforms/test_init_memref.py (1)
500-503: Consider renaming to avoid variable shadowing.
init_tileis first defined as anir.Var(line 427), then shadowed here asir.TileType. This variable reuse across different types can be confusing when reading the test.♻️ Suggested rename for clarity
- init_tile = cast(ir.TileType, init_value.type) - iter_arg_tile = cast(ir.TileType, loop_iter_arg.type) - yield_tile = cast(ir.TileType, yield_var.type) - return_var_tile = cast(ir.TileType, loop_return_var.type) + init_tile_type = cast(ir.TileType, init_value.type) + iter_arg_tile_type = cast(ir.TileType, loop_iter_arg.type) + yield_tile_type = cast(ir.TileType, yield_var.type) + return_var_tile_type = cast(ir.TileType, loop_return_var.type) - assert init_tile.shares_memref_with(iter_arg_tile), "initValue and iter_arg must share MemRef" + assert init_tile_type.shares_memref_with(iter_arg_tile_type), "initValue and iter_arg must share MemRef" # Group B: yield value and return_var share the same MemRef - assert yield_tile.shares_memref_with(return_var_tile), "yield value and return_var must share MemRef" + assert yield_tile_type.shares_memref_with(return_var_tile_type), "yield value and return_var must share MemRef" # Group A and B have different MemRefs (yield gets independent MemRef from tile.add) - assert not yield_tile.shares_memref_with(iter_arg_tile), ( + assert not yield_tile_type.shares_memref_with(iter_arg_tile_type), ( "yield value should get its own MemRef from tile.add, not share with iter_arg" )🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/ut/ir/transforms/test_init_memref.py` around lines 500 - 503, The variable init_tile is shadowed: earlier it held an ir.Var and here it's reassigned to an ir.TileType, which is confusing; rename the TileType instance (e.g., init_tile_type or init_tile_ty) and update its usages (the casts of init_value.type, loop_iter_arg.type, yield_var.type, loop_return_var.type) so the original ir.Var named init_tile remains distinct from the newly named TileType variables.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Nitpick comments:
In `@tests/ut/ir/transforms/test_init_memref.py`:
- Around line 500-503: The variable init_tile is shadowed: earlier it held an
ir.Var and here it's reassigned to an ir.TileType, which is confusing; rename
the TileType instance (e.g., init_tile_type or init_tile_ty) and update its
usages (the casts of init_value.type, loop_iter_arg.type, yield_var.type,
loop_return_var.type) so the original ir.Var named init_tile remains distinct
from the newly named TileType variables.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 51414509-f94b-4c1a-910f-157c239a4afc
📒 Files selected for processing (8)
docs/en/dev/passes/12-init_memref.mddocs/en/dev/passes/13-basic_memory_reuse.mddocs/zh-cn/dev/passes/12-init_memref.mddocs/zh-cn/dev/passes/13-basic_memory_reuse.mdinclude/pypto/ir/transforms/utils/memref_utils.hsrc/ir/transforms/basic_memory_reuse_pass.cppsrc/ir/transforms/init_memref.cpptests/ut/ir/transforms/test_init_memref.py
✅ Files skipped from review due to trivial changes (2)
- docs/en/dev/passes/12-init_memref.md
- docs/zh-cn/dev/passes/12-init_memref.md
🚧 Files skipped from review as they are similar to previous changes (4)
- docs/en/dev/passes/13-basic_memory_reuse.md
- include/pypto/ir/transforms/utils/memref_utils.h
- docs/zh-cn/dev/passes/13-basic_memory_reuse.md
- src/ir/transforms/basic_memory_reuse_pass.cpp
…sicMemoryReuse InitMemRef now correctly assigns MemRefs to ForStmt loop-carry variables in two groups: initValue+iter_arg share one MemRef, yield+return_var share another. BasicMemoryReuse adds ForStmtYieldFixupMutator to reconcile the two groups by inserting tile.move when needed. Also extracts shared FindYieldStmt utility to memref_utils.h.
2 participants
Summary
Fix ForStmt loop-carry variable MemRef assignment in InitMemRef and BasicMemoryReuse passes.
InitMemRef changes
VisitStmt_(ForStmtPtr)to manually traverse ForStmt fields, ensuring correct MemRef inheritance order:initValue+iter_argshare the same MemRef (iter_arg inherits from initValue)yield value+return_varshare the same MemRef (return_var inherits from yield value)BasicMemoryReuse changes
ForStmtYieldFixupMutator(Step 4.5) to reconcile Group A and Group B after memory reuse:tile.movebefore yield to copy into initValue's bufferShared utilities
FindYieldStmtfrom both passes intomemref_utils.hto eliminate duplicationTesting
test_init_memref_for_stmt_loop_carry_memref_relationships)Files Changed
include/pypto/ir/transforms/utils/memref_utils.h— Added sharedFindYieldStmtsrc/ir/transforms/init_memref.cpp— Refactored ForStmt MemRef assignmentsrc/ir/transforms/basic_memory_reuse_pass.cpp— AddedForStmtYieldFixupMutatortests/ut/ir/transforms/test_init_memref.py— Updated test for two-group MemRef modeldocs/en/dev/passes/12-init_memref.md— Added ForStmt loop-carry sectiondocs/en/dev/passes/13-basic_memory_reuse.md— Added yield fixup stepdocs/zh-cn/dev/passes/12-init_memref.md— Chinese translationdocs/zh-cn/dev/passes/13-basic_memory_reuse.md— Chinese translation