Skip to content

refactor(ir): switch to pointer identity for variable management in IR transforms#625

Open
lyfne123 wants to merge 1 commit intohw-native-sys:mainfrom
lyfne123:main
Open

refactor(ir): switch to pointer identity for variable management in IR transforms#625
lyfne123 wants to merge 1 commit intohw-native-sys:mainfrom
lyfne123:main

Conversation

@lyfne123
Copy link
Collaborator

  • Updated multiple IR transformation files to utilize pointer identity (const Var*) instead of string-based name_hint_ for variable identification.
  • Enhanced variable substitution and mapping methods to improve type safety and reduce reliance on string comparisons.
  • Refactored relevant functions to ensure consistent handling of variable identities across transformations.
  • Improved overall code clarity and maintainability by centralizing variable identity logic.

This refactor aligns with recent changes aimed at enhancing the robustness of variable management in the IR layer.

@coderabbitai
Copy link

coderabbitai bot commented Mar 19, 2026

Note

Reviews paused

It 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 reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

  • @coderabbitai resume to resume automatic reviews.
  • @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

  • ▶️ Resume reviews
  • 🔍 Trigger review
📝 Walkthrough

Walkthrough

The PR switches multiple IR transform passes from name-based (string) variable maps to pointer-identity (const Var*) maps and updates PTO codegen loop-carried tensor iter_arg init resolution to use GetOrCreateTensorView, with related adjustments in basic memory reuse and new unit test coverage.

Changes

Cohort / File(s) Summary
PTO Codegen Tensor View Resolution
src/codegen/pto/pto_codegen.cpp
For tensor-typed (non-scalar) loop-carried iter_arg init values, SSA names are obtained via GetOrCreateTensorView(init_var) (falls back to alternative name maps and follows IterArg::initValue_) instead of requiring a direct tensor_to_view_ lookup.
Transform Passes — Pointer-identity migration
src/ir/transforms/convert_tensor_to_tile_ops_pass.cpp, src/ir/transforms/flatten_tile_nd_to_2d_pass.cpp, src/ir/transforms/resolve_transpose_layout_pass.cpp
Replaced name-keyed substitution/collector maps with pointer-identity keyed maps (std::unordered_map<const Var*, VarPtr> / std::unordered_set<const Var*>); updated substitution APIs, loop/if/assign handling, and caller/param update flows to use var.get()/orig_iter_arg.get() identity lookups; preserved cross-scope name matching where needed via auxiliary name->Var maps.
Basic memory reuse loop-body fixup
src/ir/transforms/basic_memory_reuse_pass.cpp
When iter-arg memref patching occurs, the pass now re-visits/remutates the loop body (VisitStmt(for_stmt->body_)) and erases old iter_arg entries from var_remap_ to avoid stale mappings; reconstructed ForStmt uses the patched body.
Tests — memory reuse
tests/ut/ir/transforms/test_basic_memory_reuse.py
Added helper to find first ForStmt and new test test_loop_fixup_rewrites_iter_arg_body_refs_and_removes_stale_alloc verifying iter_arg init-value memref fixup inside loop body and that stale allocs are removed after the pass.

Sequence Diagram(s)

(omitted — changes are internal refactors and targeted lookup/fix behavior that do not introduce a new multi-component runtime control flow requiring a sequence diagram)

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~60 minutes

Possibly related issues

Possibly related PRs

Suggested reviewers

  • Hzfengsy

Poem

🐰
I hop through maps both old and new,
From names to pointers, tidy too.
Iter-arg views I gently trace,
Fixing bodies, cleaning space.
A tiny hop — a cleaner place 🥕✨

🚥 Pre-merge checks | ✅ 2 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 53.13% which is insufficient. The required threshold is 80.00%. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (2 passed)
Check name Status Explanation
Title check ✅ Passed The title accurately summarizes the main refactoring effort across all modified files: switching from string-based variable identification to pointer identity-based approach.
Description check ✅ Passed The description is directly related to the changeset, explaining the core refactoring goal of moving from name-based to pointer-identity-based variable management in IR transforms.

✏️ Tip: You can configure your own custom pre-merge checks in the settings.

📝 Coding Plan
  • Generate coding plan for human review comments

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.

❤️ Share

Comment @coderabbitai help to get the list of available commands and usage tips.

@gemini-code-assist
Copy link
Contributor

Summary of Changes

Hello, 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 introduces a fundamental change in how variables are identified and managed within the Intermediate Representation (IR) layer. By transitioning from string-based naming to pointer identity, the system gains improved robustness against naming collisions and enhances type safety. This refactoring impacts several core IR transformation passes, ensuring a more consistent and reliable approach to variable handling throughout the compilation pipeline.

Highlights

  • Variable Identification: Switched from string-based name_hint_ to pointer identity (const Var*) for variable management in IR transformations, enhancing type safety and reducing reliance on string comparisons.
  • Variable Substitution and Mapping: Enhanced variable substitution and mapping methods across multiple IR transformation files to consistently use pointer identity for Var and IterArg objects.
  • Code Refactoring: Refactored relevant functions and data structures (e.g., unordered_map keys) to centralize variable identity logic, improving overall code clarity and maintainability.
  • Dependency Update: Removed the custom VarSubstitutionMutator class in favor of a new utility function SubstituteVars for variable substitution, streamlining the transformation process.
Using Gemini Code Assist

The 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 /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature Command Description
Code Review /gemini review Performs a code review for the current pull request in its current state.
Pull Request Summary /gemini summary Provides a summary of the current pull request in its current state.
Comment @gemini-code-assist Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help /gemini help Displays a list of available commands.

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 .gemini/ folder in the base of the repository. Detailed instructions can be found here.

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

  1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution.

Copy link
Contributor

@gemini-code-assist gemini-code-assist bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request refactors several IR transformation passes to use pointer identity (const Var*) instead of string-based name_hint_ for variable identification and substitution. This is a significant improvement for the robustness and maintainability of the IR layer. The changes are consistently applied across convert_tensor_to_tile_ops_pass.cpp, flatten_tile_nd_to_2d_pass.cpp, and resolve_transpose_layout_pass.cpp. A new centralized SubstituteVars utility has been introduced, replacing local mutator implementations. The handling of complex cross-scope variable matching in flatten_tile_nd_to_2d_pass.cpp with a hybrid pointer/name approach is well-considered. The overall refactoring is thorough and correct.

…R transforms

- Updated multiple IR transformation files to utilize pointer identity (`const Var*`) instead of string-based `name_hint_` for variable identification.
- Enhanced variable substitution and mapping methods to improve type safety and reduce reliance on string comparisons.
- Refactored relevant functions to ensure consistent handling of variable identities across transformations.
- Improved overall code clarity and maintainability by centralizing variable identity logic.

This refactor aligns with recent changes aimed at enhancing the robustness of variable management in the IR layer.
Copy link

@coderabbitai coderabbitai bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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/flatten_tile_nd_to_2d_pass.cpp`:
- Around line 337-345: The current code in flatten_tile_nd_to_2d_pass.cpp
substitutes an IfStmt/ForStmt/WhileStmt return var by directly reusing a
body-local Var (see if_stmt->return_vars_, then_ctx.name_to_new_var and
ctx.Insert(rv, it->second)), which breaks the invariant that control-flow nodes
have fresh defs; instead, for each matched rv create a fresh Var with the same
type/name-hint as the original (e.g., new_rv), push new_rv into new_return_vars,
insert the mapping from the original rv to new_rv into ctx (ctx.Insert(rv,
new_rv)), and then map the body-local var to that new_rv so uses after the
branch/loop refer to the fresh control-flow def (ensure the same fix is applied
in the analogous blocks referenced at lines 383-387 and 424-428).

In `@tests/ut/ir/transforms/test_basic_memory_reuse.py`:
- Around line 1054-1076: The helper _find_first_for_stmt currently skips
ir.ScopeStmt, causing false negatives when loops are wrapped in scopes; update
_find_first_for_stmt to detect isinstance(stmt, ir.ScopeStmt) and recurse into
the inner scope (e.g., call _find_first_for_stmt on stmt.body or the ScopeStmt's
contained statement sequence) just like it does for ir.WhileStmt and ir.IfStmt
so the search descends into scope wrappers.

ℹ️ Review info
⚙️ Run configuration

Configuration used: Repository UI

Review profile: CHILL

Plan: Pro

Run ID: 946fa49b-bc49-48a4-b218-31c7479afdbb

📥 Commits

Reviewing files that changed from the base of the PR and between dfc7cb8 and 7123597.

📒 Files selected for processing (6)
  • src/codegen/pto/pto_codegen.cpp
  • src/ir/transforms/basic_memory_reuse_pass.cpp
  • src/ir/transforms/convert_tensor_to_tile_ops_pass.cpp
  • src/ir/transforms/flatten_tile_nd_to_2d_pass.cpp
  • src/ir/transforms/resolve_transpose_layout_pass.cpp
  • tests/ut/ir/transforms/test_basic_memory_reuse.py
🚧 Files skipped from review as they are similar to previous changes (1)
  • src/codegen/pto/pto_codegen.cpp

Comment on lines +337 to 345
// Substitute return_vars using the branch contexts (name-based cross-scope matching)
std::vector<VarPtr> new_return_vars;
new_return_vars.reserve(if_stmt->return_vars_.size());
for (const auto& rv : if_stmt->return_vars_) {
auto it = then_ctx.var_map.find(rv->name_hint_);
if (it != then_ctx.var_map.end()) {
auto it = then_ctx.name_to_new_var.find(rv->name_hint_);
if (it != then_ctx.name_to_new_var.end()) {
new_return_vars.push_back(it->second);
ctx.var_map[rv->name_hint_] = it->second;
ctx.Insert(rv, it->second);
} else {
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

Keep return_vars_ as fresh control-flow defs.

Here it->second is a var defined inside the branch or loop body. Reusing that pointer as the IfStmt/ForStmt/WhileStmt return var breaks the invariant that return_vars_ are separate defs for the control-flow node, and later passes can substitute post-branch/post-loop uses to a body-local var instead of the statement output. Please rebuild a fresh Var for the rewritten type and map rv -> new_rv, rather than wiring rv directly to the body var.

Minimal fix pattern
-        if (it != body_ctx.name_to_new_var.end()) {
-          new_return_vars.push_back(it->second);
-          ctx.Insert(rv, it->second);
+        if (it != body_ctx.name_to_new_var.end()) {
+          auto new_rv = std::make_shared<Var>(rv->name_hint_, it->second->GetType(), rv->span_);
+          new_return_vars.push_back(new_rv);
+          ctx.Insert(rv, new_rv);
         } else {

Also applies to: 383-387, 424-428

🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@src/ir/transforms/flatten_tile_nd_to_2d_pass.cpp` around lines 337 - 345, The
current code in flatten_tile_nd_to_2d_pass.cpp substitutes an
IfStmt/ForStmt/WhileStmt return var by directly reusing a body-local Var (see
if_stmt->return_vars_, then_ctx.name_to_new_var and ctx.Insert(rv, it->second)),
which breaks the invariant that control-flow nodes have fresh defs; instead, for
each matched rv create a fresh Var with the same type/name-hint as the original
(e.g., new_rv), push new_rv into new_return_vars, insert the mapping from the
original rv to new_rv into ctx (ctx.Insert(rv, new_rv)), and then map the
body-local var to that new_rv so uses after the branch/loop refer to the fresh
control-flow def (ensure the same fix is applied in the analogous blocks
referenced at lines 383-387 and 424-428).

Comment on lines +1054 to +1076
def _find_first_for_stmt(stmt):
"""Return the first ForStmt found in a statement tree."""
if isinstance(stmt, ir.ForStmt):
return stmt
if isinstance(stmt, ir.SeqStmts):
for child in stmt.stmts:
found = _find_first_for_stmt(child)
if found is not None:
return found
if isinstance(stmt, ir.OpStmts):
for child in stmt.stmts:
found = _find_first_for_stmt(child)
if found is not None:
return found
if isinstance(stmt, ir.IfStmt):
found = _find_first_for_stmt(stmt.then_body)
if found is not None:
return found
if stmt.else_body is not None:
return _find_first_for_stmt(stmt.else_body)
if isinstance(stmt, ir.WhileStmt):
return _find_first_for_stmt(stmt.body)
return None
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟡 Minor

Descend into scope wrappers here.

This helper skips ir.ScopeStmt, so the new regression test will false-negative as soon as a normalization pass wraps the loop body in a scope.

🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@tests/ut/ir/transforms/test_basic_memory_reuse.py` around lines 1054 - 1076,
The helper _find_first_for_stmt currently skips ir.ScopeStmt, causing false
negatives when loops are wrapped in scopes; update _find_first_for_stmt to
detect isinstance(stmt, ir.ScopeStmt) and recurse into the inner scope (e.g.,
call _find_first_for_stmt on stmt.body or the ScopeStmt's contained statement
sequence) just like it does for ir.WhileStmt and ir.IfStmt so the search
descends into scope wrappers.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant