gemini-code-assist[bot] commented on PR #18873: URL: https://github.com/apache/tvm/pull/18873#issuecomment-3999112191
## Summary of Changes Hello, I'm Gemini Code Assist[^1]! 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](https://codeassist.google/code-review/login) to enable Memory in your admin console.** <details> <summary><b>Changelog</b></summary> * **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 `LetStmt`s 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`. </details> <details> <summary><b>Activity</b></summary> * 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. </details> <details> <summary><b>Using Gemini Code Assist</b></summary> <br> The full guide for Gemini Code Assist can be found on our [documentation page](https://developers.google.com/gemini-code-assist/docs/review-github-code), here are some quick tips. <b>Invoking Gemini</b> 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. <b>Customization</b> 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](https://developers.google.com/gemini-code-assist/docs/customize-gemini-behavior-github). <b>Limitations & Feedback</b> Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with :thumbsup: and :thumbsdown: 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](https://google.qualtrics.com/jfe/form/SV_2cyuGuTWsEw84yG). <b>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](https://cloud.google.com/products/gemini/code-assist).</b> </details> [^1]: Review the [Privacy Notices](https://policies.google.com/privacy), [Generative AI Prohibited Use Policy](https://policies.google.com/terms/generative-ai/use-policy), [Terms of Service](https://policies.google.com/terms), and learn how to configure Gemini Code Assist in GitHub [here](https://developers.google.com/gemini-code-assist/docs/customize-gemini-behavior-github). Gemini can make mistakes, so double check it and [use code with caution](https://support.google.com/legal/answer/13505487). -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
