[TIR][FEAT] Require DeclBuffer before use in verify_well_formed

[tqchen]

Summary

Enforce the DeclBuffer-before-use invariant from RFC#70 across all TIR.

Changes

• Track Buffer defs for match_buffers and buffer_bind_scope in TIRVisitorWithPath
• Add UndefinedBufferVerifier to verify_well_formed for all TIR (not just schedulable)
• Update FlattenBuffer to preserve DeclBuffer nodes instead of stripping them
• Fix transform_layout to maintain buffer identity after transformation

[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 significantly enhances the well-formedness checks for buffers within TVM's Tensor IR (TIR) by enforcing that all buffers are explicitly declared before their use. This change aligns with RFC#70, improving the robustness and consistency of buffer management throughout the compiler's intermediate representation. It involves updates to core analysis, transformation, and visitation passes, along with comprehensive test case adjustments to reflect the new invariant.

Highlights

• Buffer Well-Formedness Enforcement: Introduced a new UndefinedBufferVerifier to verify_well_formed that strictly enforces the DeclBuffer-before-use invariant for all TIR buffers, ensuring they are declared within their scope.
• FlattenBuffer Enhancement: Updated FlattenBuffer to correctly emit DeclBuffer nodes for flattened parameter buffer views at the function body level, ensuring proper declaration for these aliased buffers.
• TIRVisitorWithPath Scope Correction: Fixed TIRVisitorWithPath to define SBlock::alloc_buffers before visiting reads/writes, resolving scoping issues where allocated buffers were not in scope for region references.
• Layout Transformation Buffer Identity: Corrected TransformLayoutRewriter to properly update buffer identity in access regions and match_buffer source references after layout transformations, maintaining consistency.
• RenewDefMutator DeclBuffer Support: Added a DeclBufferNode visitor to RenewDefMutator to ensure that buffer data variables within DeclBuffer statements are correctly renewed during mutation.
• Extensive Test Suite Updates: Modified approximately 60 test instances across 21 files, replacing T.Buffer(...) aliases with T.decl_buffer(...), adding check_well_formed=False for expected outputs with buffer views, and adjusting body traversal paths for new DeclBuffer nesting.

🧠 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

• src/s_tir/schedule/primitive/layout_transformation.cc
  • Updated TransformLayoutRewriter to ensure buffer identity is correctly updated in BufferRegion and MatchBufferRegion sources after layout transformations.
• src/s_tir/transform/renew_defs.cc
  • Added a VisitStmt_ method for DeclBufferNode to RenewDefMutator to handle buffer data variable renewal.
• src/tir/analysis/verify_well_formed.cc
  • Introduced UndefinedBufferVerifier to check that buffers are declared before use and are not used outside their declared scope.
  • Integrated UndefinedBufferVerifier into the main VerifyWellFormed function for both PrimFunc and IRModule.
• src/tir/ir/tir_visitor_with_path.cc
  • Expanded AttrStmt context to include DefContext<Buffer> for buffer-related attributes.
  • Added DefContext<Buffer> for buffer_view in AttrStmt visitation.
  • Reordered SBlock visitation to define alloc_buffers before reads and writes to ensure correct scoping.
  • Added DefContext<Buffer> for match_buffers in SBlock visitation.
• src/tir/transform/flatten_buffer.cc
  • Included <unordered_set> header for buffers_used_.
  • Modified Flatten function to emit DeclBuffer nodes for flattened parameter buffer views at the function body level.
  • Updated VisitStmt_(const DeclBufferNode* op) to modify existing DeclBuffer nodes instead of stripping them.
  • Added buffers_used_ set to track accessed buffers for DeclBuffer emission.
• tests/python/s_tir/test_s_tir_renew_defs.py
  • Replaced T.Buffer with T.decl_buffer for buffer declarations.
  • Adjusted the path to retrieve the buffer from BufferStore in _get_buffer_store_buffer.
• tests/python/s_tir/transform/test_s_tir_transform_compact_buffer_region.py
  • Replaced T.Buffer with T.decl_buffer in before and expected prim_func definitions.
• tests/python/s_tir/transform/test_s_tir_transform_inject_double_buffer.py
  • Replaced T.Buffer with T.decl_buffer in various test functions.
• tests/python/s_tir/transform/test_s_tir_transform_inject_ptx_async_copy.py
  • Replaced T.Buffer with T.decl_buffer in test functions.
  • Added A_flattened = T.decl_buffer(...) in Expected prim_func.
  • Added check_well_formed=False to Expected prim_func.
• tests/python/s_tir/transform/test_s_tir_transform_inject_virtual_thread.py
  • Replaced T.Buffer with T.decl_buffer in various test functions.
  • Added check_well_formed=False to expected_func prim_func definitions.
  • Introduced new T.Buffer aliases (B_1) in expected_func to handle buffer views.
• tests/python/s_tir/transform/test_s_tir_transform_loop_partition.py
  • Replaced T.Buffer with T.decl_buffer in numerous test functions.
  • Added T.decl_buffer for function parameters in after and expected_partitioned_concat_single_point.
  • Introduced new T.Buffer aliases (A_1, B_1, etc.) in after and expected functions to handle buffer views.
• tests/python/s_tir/transform/test_s_tir_transform_lower_thread_all_reduce.py
  • Replaced T.Buffer with T.decl_buffer in various test functions.
• tests/python/s_tir/transform/test_s_tir_transform_merge_dynamic_shared_memory_allocations.py
  • Replaced T.Buffer with T.decl_buffer in test functions.
  • Modified assertions to check script properties instead of structural equality due to DeclBuffer nodes.
• tests/python/s_tir/transform/test_s_tir_transform_plan_update_buffer_allocation_location.py
  • Replaced T.Buffer with T.decl_buffer in the before prim_func.
• tests/python/s_tir/transform/test_s_tir_transform_renormalize_split_pattern.py
  • Replaced T.Buffer with T.decl_buffer for flattened buffers in test functions.
• tests/python/s_tir/transform/test_s_tir_transform_rewrite_unsafe_select.py
  • Replaced T.Buffer with T.decl_buffer in test functions.
  • Adjusted body traversal paths (.body.body.body.value) to account for new DeclBuffer nesting.
• tests/python/s_tir/transform/test_s_tir_transform_thread_sync.py
  • Added check_well_formed=False to test_sync_read_thread_id_independent_location.
  • Replaced T.Buffer with T.decl_buffer in various test functions.
• tests/python/tir-analysis/test_tir_analysis_verify_well_formed.py
  • Added test_buffer_in_buffer_map_is_well_formed to verify buffers defined in buffer_map.
  • Added test_decl_buffer_is_well_formed to verify buffers defined via DeclBuffer statement.
  • Added test_alloc_buffer_in_block_is_well_formed to verify buffers defined in SBlock::alloc_buffers.
  • Added test_match_buffer_in_block_is_well_formed to verify buffers defined in SBlock::match_buffers.
  • Added test_error_buffer_used_out_of_decl_scope to test for errors when a buffer is used after its DeclBuffer scope ends.
  • Added test_error_undeclared_buffer_in_schedulable_tir to test for errors when a buffer is used without any prior declaration in a block context.
• tests/python/tir-base/test_tir_stmt_functor_substitute.py
  • Replaced T.Buffer with T.decl_buffer in test functions.
• tests/python/tir-transform/test_tir_transform_convert_ssa.py
  • Replaced T.Buffer with T.decl_buffer in test functions.
• tests/python/tir-transform/test_tir_transform_flatten_buffer.py
  • Replaced T.Buffer with T.decl_buffer in various test functions.
  • Added check_well_formed=False to I.ir_module and T.prim_func definitions where buffer views are expected.
• tests/python/tir-transform/test_tir_transform_lower_tvm_builtin.py
  • Added check_well_formed=False to Expected ir_module and prim_func definitions.
• tests/python/tir-transform/test_tir_transform_pointer_value_type_rewrite.py
  • Replaced T.Buffer with T.decl_buffer in test functions.
  • Added check_well_formed=False to Expected ir_module and prim_func definitions.
  • Introduced new T.Buffer aliases (A_local_1) in Expected functions to handle buffer views.
• tests/python/tir-transform/test_tir_transform_simplify.py
  • Replaced T.Buffer with T.decl_buffer in test functions.
  • Added logic to navigate through DeclBuffer nodes in assertions to reach the inner body.
• tests/python/tir-transform/test_tir_transform_split_host_device.py
  • Replaced T.Buffer with T.decl_buffer for buffer aliases in main function.
• tests/python/tir-transform/test_tir_transform_storage_rewrite.py
  • Replaced T.Buffer with T.decl_buffer in various test functions.
  • Updated assertions to account for the presence of DeclBuffer nodes in the IR structure.
• tests/python/tir-transform/test_tir_transform_unroll_loop.py
  • Replaced T.Buffer with T.decl_buffer in test functions.
• tests/python/tvmscript/test_tvmscript_roundtrip.py
  • Replaced T.Buffer with T.decl_buffer in various test functions.
  • Added check_well_formed=False to T.prim_func definitions where buffer views are used.

Activity

• This pull request supersedes [TIR][FEAT] Require DeclBuffer before use in verify_well_formed #18840, which was closed due to a force push, indicating a refinement of previous work.
• The author has confirmed that all s_tir, tir-transform, tir-base, TVMScript roundtrip, all-platform minimal, set_axis_separator subregion, and transform_layout tests pass, demonstrating thorough validation 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 correctly enforces the DeclBuffer-before-use invariant, a significant step towards better TIR well-formedness as per RFC#70. The changes are comprehensive, including the new UndefinedBufferVerifier, updates to various TIR transforms like FlattenBuffer and RenewDefMutator, and extensive modifications to tests to align with the new requirement. The logic for adding DeclBuffer for flattened parameter buffers and handling buffer scopes in different nodes seems correct and robust. The test updates are thorough, covering both positive cases and expected failures, which ensures the new verification logic is working as intended. I have one minor suggestion for code simplification, but overall, this is a high-quality contribution.

[gemini-code-assist]

This logic for updating the match_buffer is correct, but it can be made more concise and idiomatic by chaining the CopyOnWrite calls. This improves readability without changing the functionality.

Suggested change

	auto new_source = match_buf->source;
	auto* source_n = new_source.CopyOnWrite();
	source_n->buffer = new_buffer_;
	auto new_match = match_buf;
	new_match.CopyOnWrite()->source = new_source;
	return new_match;
	auto new_match = match_buf;
	auto* n = new_match.CopyOnWrite();
	n->source.CopyOnWrite()->buffer = new_buffer_;
	return new_match;