[TIR] Add VisitBufferDef/VisitBufferUse to base StmtVisitor/StmtMutator

[tqchen]

Summary

Add virtual VisitBufferDef/VisitBufferUse hooks to StmtVisitor, StmtMutator,
StmtExprVisitor, StmtExprMutator, and TIRVisitorWithPath.

These centralize buffer field visiting (shape, strides, elem_offset) at definition
sites only, and distinguish definition vs use sites so analyses like UndefinedVars
correctly track variable scoping.

Key Design

• VisitBufferDef(buffer, alloc_data) — called at definition sites (AllocBuffer, DeclBuffer, SBlock alloc_buffers). Visits buffer shape, strides, elem_offset. alloc_data=true means data var is a new allocation (DEF); false means data references existing allocation (USE).

• VisitBufferUse(buffer) — called at use sites (BufferLoad, BufferStore, SBlock reads/writes). Default is empty (no field re-visiting).

• StmtMutator::VisitBufferDef maintains a buffer_remap_ map so mutated buffers propagate to use sites automatically.

• Removes duplicated buffer field visiting from ~10 subclasses by centralizing it in the base classes.

Test plan

• test_tir_analysis_undefined_vars.py — new tests for def/use distinction
• test_tir_transform_convert_ssa.py — 13 tests
• test_tir_transform_vectorize.py — 72 tests
• test_tir_transform_simplify.py — new test for buffer simplification
• test_minimal_target_codegen_llvm.py — LLVM codegen smoke test
• ir_functor_test.cc — C++ functor tests

[gemini-code-assist]

Summary of Changes

Hello, I'm Gemini Code Assist¹! 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 refactors how buffers are visited and mutated within the TIR (Tensor IR) framework by introducing explicit virtual methods for buffer definition and use sites. This change centralizes the logic for handling buffer fields, such as shape and strides, and ensures that variable scoping is correctly maintained during analysis and transformation passes. The primary impact is a more robust and maintainable IR infrastructure, particularly for analyses like UndefinedVars and various mutation passes, by clearly separating how buffers are declared versus how they are referenced.

Highlights

• New Virtual Methods for Buffer Visiting: Introduced VisitBufferDef(buffer, alloc_data) and VisitBufferUse(buffer) as new virtual methods in StmtVisitor, StmtMutator, StmtExprVisitor, StmtExprMutator, and TIRVisitorWithPath to centralize buffer field visiting logic.
• Definition vs. Use Site Distinction: The new methods explicitly distinguish between buffer definition sites (e.g., AllocBuffer, DeclBuffer, SBlock alloc_buffers) and use sites (e.g., BufferLoad, BufferStore, SBlock reads/writes), allowing for more accurate variable scoping analysis.
• Centralized Buffer Field Visiting: Buffer shape, strides, and elem_offset fields are now visited primarily at definition sites via VisitBufferDef, preventing redundant re-visiting at use sites where associated variables might be out of scope.
• Improved StmtMutator Buffer Remapping: The StmtMutator now includes a buffer_remap_ map to automatically propagate mutated buffers to their use sites, ensuring consistency across transformations.
• Enhanced Variable Use/Definition Analysis: The VarUseDefAnalyzer has been updated to correctly track buffer data variables as definitions or uses based on AllocBuffer (definition) vs. DeclBuffer (use), resolving issues with UndefinedVars analysis.
• Refactored Buffer Handling in Passes: Numerous passes and visitors, including RenewDefMutator, IndexDataTypeRewriter, PrimFuncSpecializer, IRSubstitute, IRConvertSSA, and StmtSimplifier, have been refactored to leverage the new VisitBufferDef/VisitBufferUse hooks, simplifying their implementations and improving correctness.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Changelog

• include/tvm/tir/stmt_functor.h
  • Added virtual VisitBufferDef and VisitBufferUse methods to StmtVisitor and StmtMutator.
  • Introduced ffi::Map<Buffer, Buffer> buffer_remap_ to StmtMutator for tracking buffer remappings.
  • Reordered inheritance for StmtExprVisitor and StmtExprMutator to ExprVisitor then StmtVisitor/StmtMutator.
  • Overrode VisitExpr_ for BufferLoadNode in StmtExprVisitor and StmtExprMutator to call VisitBufferUse.
• src/s_tir/transform/renew_defs.cc
  • Overrode VisitBufferDef and VisitBufferUse to use DefineBuffer and UseOrRemapBuffer respectively.
  • Removed direct VisitStmt_ overrides for AllocBufferNode, DeclBufferNode, BufferStoreNode, and BufferLoadNode, delegating to base class methods.
• src/tir/analysis/var_use_def_analysis.cc
  • Implemented VisitBufferDef to handle data variable definitions for AllocBuffer and uses for DeclBuffer.
  • Implemented VisitBufferUse to track buffer and its data pointer as uses.
  • Removed specific VisitStmt_ overrides for DeclBufferNode, AllocBufferNode, and BufferStoreNode, and VisitExpr_ for BufferLoadNode.
• src/tir/analysis/var_use_def_analysis.h
  • Updated VarUseDefAnalyzer to declare overrides for VisitBufferDef and VisitBufferUse.
  • Removed declarations for VisitStmt_ of DeclBufferNode, BufferStoreNode, and VisitExpr_ of BufferLoadNode.
• src/tir/analysis/verify_well_formed.cc
  • Changed UndefinedBufferVerifier::Visit(const Buffer&, AccessPath) to VisitBufferUse(const Buffer&, AccessPath).
• src/tir/ir/data_type_rewriter.cc
  • Overrode VisitBufferDef and VisitBufferUse to manage is_enabled_ flag and call base class methods.
  • Updated VisitStmt_ for SBlockNode to use VisitBufferDef for alloc_buffers and match_buffers.
  • Updated VisitBlockAnnotations and VisitBufferRegion to use VisitBufferUse.
  • Removed custom VisitBuffer and GetRemappedBuffer methods and the buffer_remap_ member.
• src/tir/ir/data_type_rewriter.h
  • Updated IndexDataTypeRewriter to declare overrides for VisitBufferDef and VisitBufferUse.
  • Removed declarations for VisitStmt_ of DeclBufferNode, AllocBufferNode, and VisitBuffer/GetRemappedBuffer methods.
• src/tir/ir/specialize.cc
  • Overrode VisitBufferUse to use GetNewBuffer for buffer remapping.
  • Removed specific VisitStmt_ for BufferStoreNode and VisitExpr_ for BufferLoadNode.
• src/tir/ir/stmt_functor.cc
  • Implemented default VisitBufferDef and VisitBufferUse for StmtVisitor and StmtMutator.
  • Updated StmtVisitor::VisitStmt_ for AllocBufferNode, DeclBufferNode, BufferStoreNode, and SBlockNode to call VisitBufferDef/VisitBufferUse.
  • Updated StmtMutator::VisitStmt_ for AllocBufferNode, DeclBufferNode, BufferStoreNode, and SBlockNode to use VisitBufferDef/VisitBufferUse and buffer_remap_.
  • Modified StmtMutator::Internal::Mutate for BufferRegion and MatchBufferRegion to use VisitBufferUse and VisitBufferDef.
• src/tir/ir/tir_visitor_with_path.cc
  • Implemented VisitBufferDef and VisitBufferUse methods.
  • Updated EnterDef to call VisitBufferDef.
  • Updated Visit(BufferRegion), VisitStmt_(BufferStoreNode), and VisitExpr_(BufferLoadNode) to call VisitBufferUse.
• src/tir/ir/tir_visitor_with_path.h
  • Added declarations for VisitBufferUse and VisitBufferDef.
  • Removed the generic virtual void Visit(const Buffer& obj, ffi::reflection::AccessPath path) declaration.
• src/tir/transform/ir_utils.cc
  • Overrode VisitBufferDef in IRConvertSSA to return the original buffer, preventing conflicting remapping with its internal SSA conversion logic.
• src/tir/transform/simplify.cc
  • Overrode VisitBufferDef in StmtSimplifier to return the original buffer, preventing simplification of buffer definition fields that could lead to identity divergence.
• tests/cpp/ir_functor_test.cc
  • Updated StmtVisitor test assertions to reflect increased visit counts due to VisitBufferDef visiting buffer shape fields.
  • Updated StmtMutator test assertions to reflect that buffers are now mutated when their shape fields change via VisitBufferDef.
• tests/python/tir-analysis/test_tir_analysis_undefined_vars.py
  • Added new tests to verify that DeclBuffer's data and elem_offset variables are correctly identified as uses (undefined) by UndefinedVars.
  • Added new tests to verify that AllocBuffer's data variable is correctly identified as a definition (not undefined) by UndefinedVars.
• tests/python/tir-transform/test_tir_transform_simplify.py
  • Added new tests to confirm that LetStmts used in buffer definition fields (e.g., elem_offset, shape, strides) are retained after simplification.
  • Added a regression test (test_simplify_buffer_identity_well_formed) to ensure Simplify does not cause buffer identity divergence between DeclBuffer and BufferLoad.

Activity

• The pull request introduces new virtual methods and refactors existing visitors/mutators to centralize buffer handling logic.
• New tests have been added to test_tir_analysis_undefined_vars.py to validate the correct distinction between buffer definition and use sites.
• Existing tests in test_tir_transform_convert_ssa.py, test_tir_transform_vectorize.py, test_tir_transform_simplify.py, test_minimal_target_codegen_llvm.py, and ir_functor_test.cc have been updated or passed, indicating compatibility and correctness of the changes.

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.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

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. ↩

[gemini-code-assist]

Code Review

This pull request introduces VisitBufferDef and VisitBufferUse hooks to the base statement visitors and mutators (StmtVisitor, StmtMutator, StmtExprVisitor, StmtExprMutator, TIRVisitorWithPath). This is a significant and beneficial refactoring that centralizes buffer visiting logic, distinguishing between buffer definition sites and use sites. This improves correctness for analyses like UndefinedVars and simplifies many downstream visitors/mutators by removing duplicated boilerplate code. The changes are well-implemented and supported by a comprehensive set of new and updated tests. My only suggestion is to correct a comment in stmt_functor.h to accurately reflect the default behavior of VisitBufferUse.