gemini-code-assist[bot] commented on PR #18876: URL: https://github.com/apache/tvm/pull/18876#issuecomment-4007004644
## 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 introduces a significant refactoring of the TIR (Tensor IR) by removing the explicit 'body' field from `AllocBuffer` and `DeclBuffer` statements. This change simplifies the IR structure, aligning these buffer operations with other flat statements like `Bind`. The primary impact is a more consistent and streamlined representation of buffer allocation and declaration, requiring extensive updates across the compiler stack to adapt to this new flat semantic model. The refactor aims to improve the clarity and maintainability of the IR, ensuring that buffer lifetimes are managed through sequential statement ordering rather than nested scopes. ### Highlights * **Core IR Refactoring**: The `body` field has been removed from `AllocBufferNode` and `DeclBufferNode`, making these statements 'flat' and consistent with `Bind` semantics. This means buffer scope now implicitly extends to the end of the enclosing `SeqStmt`. * **TVMScript IR Builder Updates**: The Python `T.alloc_buffer` and `T.decl_buffer` functions in TVMScript now directly return `Buffer` objects and emit flat statements, rather than being used as context managers. This simplifies their usage and aligns with the new IR structure. * **Widespread Adaption**: Over 60 files across core IR, codegen backends (LLVM, NVPTX, AMDGPU, C, CUDA, Metal, WebGPU, SPIRV), various transforms (e.g., `compact_buffer_region`, `inject_double_buffer`, `lower_thread_allreduce`), and analysis passes have been updated to correctly handle the new flat `AllocBuffer` and `DeclBuffer` nodes. * **Deferred Remapping for Flat IR**: New mechanisms, such as `pending_dbuffer_allocs_` in `InjectDoubleBuffer` and a `DeferredRemapper` pass in `LowerThreadAllreduce`, have been introduced to correctly handle buffer remappings in the context of flat IR where allocation statements might appear before their remapping logic. 🧠**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/script/ir_builder/tir/ir.h** * Updated `DeclBuffer` and `AllocBuffer` FFI functions to return `Buffer` instead of frame objects. * **include/tvm/tir/stmt.h** * Removed the `body` field from `DeclBufferNode` and `AllocBufferNode`. * Adjusted constructors for `DeclBuffer` and `AllocBuffer` to no longer accept a `body` parameter. * **python/tvm/script/ir_builder/tir/ir.py** * Modified `alloc_buffer` and `decl_buffer` Python functions to return `Buffer` objects directly. * Updated docstrings and examples to reflect the non-context-manager usage. * **python/tvm/tir/stmt.py** * Removed the `body` attribute and corresponding parameter from `AllocBuffer` and `DeclBuffer` Python classes. * **src/relax/op/tensor/inspect.cc** * Adapted `LegalizeTensorShape` to use the flat `DeclBuffer` within a `SeqStmt`. * **src/s_tir/analysis/calculate_allocated_memory.cc** * Modified `AllocBufferCalculator` to visit the `AllocBufferNode` directly instead of its now-removed `body`. * **src/s_tir/analysis/estimate_flops.cc** * Updated `FlopEstimator` to handle `AllocBufferNode` and `DeclBufferNode` without a `body`. * **src/s_tir/backend/adreno/inject_texture_alloc.cc** * Adjusted `TextureAllocInjector` to construct `SeqStmt` with flat `Bind` and `AllocBuffer` statements. * **src/s_tir/schedule/primitive/cache_read_write.cc** * Simplified `InsertCacheStage` and `CacheLocDetector` by removing logic that previously traversed nested `DeclBuffer`/`AllocBuffer` bodies. * **src/s_tir/schedule/transform.cc** * Removed code that peeled off `DeclBuffer` and `AllocBuffer` nodes from SBlock bodies in `LeafBlockRemovalPlan`. * **src/s_tir/transform/compact_buffer_region.cc** * Updated `BufferAccessRegionCollector` to track pending flat `AllocBuffer` nodes and compact them upon scope exit. * Introduced `CompactPendingFlatAllocBuffers` method to process deferred `AllocBuffer` compaction. * **src/s_tir/transform/inject_double_buffer.cc** * Modified `DoubleBufferInjector` to defer processing of `AllocBuffer` nodes, storing them in `pending_dbuffer_allocs_` for later re-emission. * **src/s_tir/transform/inject_ptx_ldg32.cc** * Replaced nested `AllocBuffer` statements with `SeqStmt::Flatten` for `AllocBuffer`. * **src/s_tir/transform/inject_virtual_thread.cc** * Updated `VarTouchedAnalysis` and `VTInjector` to reflect the absence of a `body` in `AllocBuffer` nodes. * **src/s_tir/transform/lower_match_buffer.cc** * Adjusted `MatchBufferLower` to correctly handle buffer remapping with flat IR, ensuring remapped buffers are accessible. * **src/s_tir/transform/lower_opaque_block.cc** * Used `SeqStmt::Flatten` to combine `AllocBuffer` with the block body. * **src/s_tir/transform/lower_thread_allreduce.cc** * Introduced deferred remapping logic for `AllocBuffer` and `DeclBuffer` nodes in `ThreadAllreduceBuilder`. * Added a `DeferredRemapper` post-processing pass to apply these deferred remappings for flat IR. * Modified allocation prepending to use `SeqStmt::Flatten` for new allocations. * **src/s_tir/transform/lower_vtcm_alloc.cc** * Updated `VtcmAllocator` to directly return a `Bind` statement for `AllocBuffer` without a nested body. * **src/s_tir/transform/merge_shared_memory_allocations.cc** * Modified `SharedMemoryRewriter` to use `SeqStmt::Flatten` for new `AllocBuffer` statements. * Changed `AllocBuffer` and `DeclBuffer` visit methods to return `Evaluate(0)` when the buffer is merged or unused. * **src/s_tir/transform/profile_instrumentation.cc** * Updated `LoopAnalyzer` to no longer traverse the body of `AllocBufferNode`. * **src/script/ir_builder/tir/frame.cc** * Modified `AllocateFrameNode` and `DeclBufferFrameNode` to emit flat `AllocBuffer`/`DeclBuffer` statements followed by their body statements, handling single vs. multiple statements for parent frames. * **src/script/ir_builder/tir/ir.cc** * Updated `DeclBuffer` and `AllocBuffer` FFI functions to directly add statements to the parent and return the `Buffer` object. * **src/script/printer/tir/stmt.cc** * Adjusted `DeclBufferDoc` and `AllocBuffer` printing logic to reflect the flat statement structure, removing context manager syntax. * **src/target/llvm/codegen_amdgpu.cc** * Removed `VisitStmt(op->body)` call for `AllocBufferNode`. * **src/target/llvm/codegen_llvm.cc** * Removed `VisitStmt(op->body)` calls for `AllocBufferNode` and `DeclBufferNode`. * **src/target/llvm/codegen_nvptx.cc** * Removed `VisitStmt(op->body)` call for `AllocBufferNode`. * **src/target/source/codegen_c.cc** * Removed `PrintStmt(op->body)` calls for `DeclBufferNode` and `AllocBufferNode`. * **src/target/source/codegen_cuda.cc** * Removed `PrintStmt(op->body)` call for `AllocBufferNode`. * **src/target/source/codegen_metal.cc** * Removed `PrintStmt(op->body)` call for `AllocBufferNode`. * **src/target/source/codegen_webgpu.cc** * Removed `PrintStmt(op->body)` call for `AllocBufferNode`. * **src/target/spirv/codegen_spirv.cc** * Removed `VisitStmt(op->body)` calls for `AllocBufferNode` and `DeclBufferNode`. * **src/tir/ir/script/script_complete.cc** * Simplified `ScriptCompleter` for `AllocBufferNode` and `DeclBufferNode` to only register the buffer, as they no longer have bodies. * **src/tir/ir/stmt.cc** * Updated `DeclBuffer` and `AllocBuffer` constructors to remove the `body` parameter. * **src/tir/ir/stmt_functor.cc** * Removed `VisitStmt(op->body)` calls from `StmtVisitor` and `StmtMutator` for `AllocBufferNode` and `DeclBufferNode`. * **src/tir/ir/tir_visitor_with_path.cc** * Modified `TIRVisitorWithPath` to push `AllocBuffer` and `DeclBuffer` definitions into the current scope as siblings, rather than visiting a nested body. * **src/tir/ir/tir_visitor_with_path.h** * Updated `TIRVisitorWithPath` to allow `DefContext<Buffer>` in `bind_scope_`. * **src/tir/transform/common_subexpr_elim.cc** * Simplified `CommonSubexpressionEliminator` for `AllocBufferNode` and `DeclBufferNode` to reflect their new flat semantics. * **src/tir/transform/flatten_buffer.cc** * Used `SeqStmt::Flatten` when inserting `DeclBuffer` statements. * **src/tir/transform/ir_utils.cc** * Updated `MergeNest` to use `SeqStmt::Flatten` for `AllocBuffer` and `DeclBuffer` nodes. * Simplified `IRConvertSSA` for `DeclBufferNode` and `AllocBufferNode` by removing scope management around their bodies. * **src/tir/transform/lower_tvm_builtin.cc** * Adapted `BuiltinLower` to handle flat `AllocBuffer` nodes, pushing free statements to the enclosing scope and constructing the allocation as a `SeqStmt`. * **src/tir/transform/lower_warp_memory.cc** * Adapted `WarpAccessRewriter` to take a `body` parameter for `Rewrite` and used `SeqStmt::Flatten` for `AllocBuffer`. * Modified `WarpMemoryRewriter` to process `SeqStmt` to find and rewrite warp `AllocBuffer` nodes. * **src/tir/transform/remove_no_op.cc** * Removed logic that checked for `body` in `AllocBufferNode` and `DeclBufferNode`. * **src/tir/transform/split_host_device.cc** * Used `SeqStmt::Flatten` when inserting `DeclBuffer` statements. * **src/tir/transform/storage_rewrite.cc** * Updated `StoragePlanRewriter` to return `Evaluate(0)` for stripped `AllocBuffer`/`DeclBuffer` nodes and used flat `AllocBuffer` in `PrepareNewAlloc`. * **src/tir/transform/tvm_ffi_binder.cc** * Used flat `DeclBuffer` in `DecodeAllParams`. * **src/tir/transform/unsupported_dtype_legalize.cc** * Updated `ComputeLegalizer` and `StorageLegalizer` to use flat `DeclBuffer` and `AllocBuffer`. * **tests/cpp/ir_functor_test.cc** * Updated test cases to reflect `AllocBuffer` and `DeclBuffer` no longer having bodies, using `SeqStmt` to combine them with other statements. * **tests/python/s_tir/transform/test_s_tir_transform_rewrite_unsafe_select.py** * Adjusted access paths for `Evaluate` nodes due to the new flat `SeqStmt` structure. * **tests/python/tir-base/test_tir_constructor.py** * Removed assertions related to the `body` field for `AllocBuffer`. * **tests/python/tir-base/test_tir_specialize.py** * Adjusted access path for `A_buf` due to the new flat `SeqStmt` structure. * **tests/python/tir-transform/test_tir_transform_convert_ssa.py** * Updated test case to use `SeqStmt` for `AllocBuffer` and `Evaluate`. * **tests/python/tir-transform/test_tir_transform_lower_tvm_builtin.py** * Removed structural equality assertion and added checks for specific keywords in the script output to verify lowered builtins and flat `DeclBuffer`. * **tests/python/tir-transform/test_tir_transform_narrow_datatype.py** * Updated `lower_func_body` to correctly navigate `SeqStmt` to find `For` nodes. * **tests/python/tir-transform/test_tir_transform_remove_no_op.py** * Added `xfail` markers for tests related to dead allocation removal, indicating it's not yet implemented for flat `AllocBuffer`. * **tests/python/tir-transform/test_tir_transform_simplify.py** * Updated navigation logic to find `For` loops within `SeqStmt`. * **tests/python/tir-transform/test_tir_transform_storage_rewrite.py** * Adjusted assertions to account for `AllocBuffer` being the first element of a `SeqStmt`. * **tests/python/tir-transform/test_tir_transform_unroll_loop.py** * Updated test cases to use flat `AllocBuffer` and `DeclBuffer` syntax. * **tests/python/tvmscript/test_tvmscript_ir_builder_tir.py** * Updated assertions for `T.allocate` and `T.decl_buffer` to expect a `SeqStmt` containing the buffer statement and an `Evaluate` statement. * **tests/python/tvmscript/test_tvmscript_printer_structural_equal.py** * Adjusted `AccessPath` for structural equality checks due to `SeqStmt` wrapping. * **tests/python/tvmscript/test_tvmscript_printer_tir.py** * Updated expected print output for `T.allocate` and `T.decl_buffer` to reflect flat statements and `SeqStmt` usage. </details> <details> <summary><b>Activity</b></summary> * The pull request introduces a fundamental change to the TIR structure by removing the `body` field from `AllocBuffer` and `DeclBuffer` nodes. * This refactoring required extensive modifications across the codebase, touching core IR definitions, script builders, code generation backends, and various transformation passes. * The author has confirmed that all existing test suites pass, indicating a successful integration of this significant change. </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]
