[REFATOR][TIR] Phase out AllocConst

[tqchen]

This PR phases out alloc const node in the TIR.
This node was oroginally introduced to include embedded weights into the allocation. However, the presence of the particular IR couples the data allocation and the weight placement, which is not as desirable especialy when weights get large. A better approach is to have extra annotation on the allocation and store weights separately either outside module or as part of module/function attribute.

As a result, we phases out this node which can help us to simplify code logic in the codebase.

[gemini-code-assist]

Summary of Changes

Hello @tqchen, 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 systematically removes the AllocateConst node from the TensorIR (TIR) framework. This change aims to improve the flexibility and modularity of TVM by separating the concerns of data allocation and constant weight placement, which were previously tightly coupled. The removal simplifies various parts of the compiler infrastructure, from IR definitions and builders to analysis passes and code generation backends, paving the way for more scalable and efficient handling of large constant data.

Highlights

• Phase out AllocateConst TIR node: The AllocateConst node in the TensorIR (TIR) is being completely removed. This node was originally designed to embed weights directly into memory allocations.
• Decoupling Data Allocation and Weight Placement: The primary motivation for this removal is to decouple data allocation from weight placement. The previous approach was found to be undesirable, especially when dealing with large weights, as it tightly coupled these two concerns.
• Simplified Codebase: Removing AllocateConst simplifies the overall code logic within the codebase, leading to a cleaner and more maintainable architecture.
• New Approach for Weights: A new strategy for handling weights will involve using extra annotations on allocations and storing weights separately, either outside the module or as part of module/function attributes.

🧠 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/ir/module.h
  • Removed the kConstants attribute and its associated documentation, which was used to store runtime Tensors from AllocateConst nodes.
• include/tvm/script/ir_builder/tir/frame.h
  • Removed the AllocateConstFrameNode and AllocateConstFrame classes, eliminating the IR builder components for AllocateConst.
• include/tvm/script/ir_builder/tir/ir.h
  • Removed the tvm::runtime::Tensor import and the AllocateConst function definition from the TIR IR builder interface.
• include/tvm/tir/stmt.h
  • Removed the AllocateConstNode and AllocateConst class definitions, effectively deleting the TIR statement for constant allocation.
• include/tvm/tir/stmt_functor.h
  • Removed VisitStmt_ overrides and dispatch macros related to AllocateConstNode from StmtFunctor, StmtVisitor, and StmtMutator.
• include/tvm/tir/transform.h
  • Removed declarations for the BindParams and ExtractPrimFuncConstants passes, as they were tied to AllocateConst.
• python/tvm/script/ir_builder/tir/frame.py
  • Removed the AllocateConstFrame class registration and its __enter__ method.
• python/tvm/script/ir_builder/tir/ir.py
  • Removed imports for numpy and tvm.runtime.tensor, the allocate_const function, and its export from the module's __all__ list.
• python/tvm/tir/init.py
  • Removed AllocateConst from the module's __all__ export list.
• python/tvm/tir/functor.py
  • Removed references to AllocateConst in the initializers and method lists of PyStmtExprVisitor and PyStmtExprMutator.
• python/tvm/tir/stmt.py
  • Removed the Tensor import and the AllocateConst class definition and registration.
• python/tvm/tir/transform/transform.py
  • Removed the ExtractPrimFuncConstants function.
• src/s_tir/schedule/primitive/cache_read_write.cc
  • Removed handling of AllocateConst nodes during cache stage insertion and cache location detection.
• src/s_tir/schedule/transform.cc
  • Removed handling of AllocateConst nodes during the leaf block removal process.
• src/s_tir/transform/plan_update_buffer_allocation_location.cc
  • Updated comments to remove references to AllocateConst nodes.
• src/script/ir_builder/tir/frame.cc
  • Removed reflection registration and the ExitWithScope implementation for AllocateConstFrameNode.
• src/script/ir_builder/tir/ir.cc
  • Removed the AllocateConst function implementation and its FFI registration.
• src/script/printer/tir/stmt.cc
  • Removed the PrintTensor helper function and the IRDocsifier dispatch for tir::AllocateConst.
• src/target/llvm/codegen_llvm.cc
  • Removed the VisitStmt_ implementation for AllocateConstNode.
• src/target/llvm/codegen_llvm.h
  • Removed the VisitStmt_ declaration for AllocateConstNode.
• src/target/source/codegen_c.cc
  • Removed the VisitStmt_ implementation for AllocateConstNode, which handled global static variable allocation.
• src/target/source/codegen_c.h
  • Removed the VisitStmt_ declaration for AllocateConstNode.
• src/target/source/codegen_webgpu.cc
  • Removed the VisitStmt_ implementation for AllocateConstNode.
• src/target/source/codegen_webgpu.h
  • Removed the VisitStmt_ declaration for AllocateConstNode.
• src/te/operation/create_primfunc.cc
  • Removed CreatePrimFuncWithConstants functions and calls to tir::BindParams.
• src/te/operation/create_primfunc.h
  • Removed CreatePrimFuncWithConstants declarations.
• src/tir/analysis/estimate_flops.cc
  • Removed the VisitStmt_ override for AllocateConstNode.
• src/tir/analysis/var_use_def_analysis.cc
  • Removed the VisitStmt_ implementation for AllocateConstNode.
• src/tir/analysis/var_use_def_analysis.h
  • Removed the VisitStmt_ declaration for AllocateConstNode.
• src/tir/ir/py_functor.cc
  • Removed f_visit_allocate_const and related dispatch/default dispatch for PyStmtExprVisitorNode and PyStmtExprMutatorNode.
• src/tir/ir/stmt.cc
  • Removed reflection registration and the constructor implementation for AllocateConst.
• src/tir/ir/stmt_functor.cc
  • Removed VisitStmt_ implementations for AllocateConstNode from StmtVisitor and StmtMutator.
• src/tir/ir/tir_visitor_with_path.cc
  • Removed the VisitStmt_ implementation for AllocateConstNode.
• src/tir/ir/tir_visitor_with_path.h
  • Removed the VisitStmt_ declaration for AllocateConstNode.
• src/tir/transform/bind_params.cc
  • Removed file containing the BindParams pass logic.
• src/tir/transform/extract_constants.cc
  • Removed file containing the ExtractPrimFuncConstants pass logic.
• src/tir/transform/ir_utils.cc
  • Removed handling of AllocateConstNode in the MergeNest utility function.
• src/tir/transform/ir_utils.h
  • Removed the BindParams function declaration.
• src/tir/transform/remove_weight_layout_rewrite_block.cc
  • Removed the AllocateConstRewrite and CollectAllocateConstBufferVars classes, along with related logic in WeightLayoutRewriteBlockRemover.
• src/tir/transform/renew_defs.cc
  • Removed STMT_REGENERATE_VAR_DEF for AllocateConstNode.
• src/tir/transform/storage_rewrite.cc
  • Removed the kAllocateConstNode flag from BufferVarInfo, and removed VisitStmt_ for AllocateConstNode from VectorTypeAccessChecker and VectorTypeRewriter. Also removed the rewrite_allocate_const_node parameter from VectorTypeRewriter constructor and PointerValueTypeRewrite function.
• tests/python/ir/test_node_reflection.py
  • Removed the test_alloc_const function.
• tests/python/relax/test_meta_schedule_relax_integration.py
  • Removed imports for numpy, pytest, tempfile, IRModule, tir, transform. Removed several PrimFunc definitions (conv2d, conv2d0, fused_conv2d_add) that used T.allocate_const. Removed the test_using_anchor_trace function and related module equality expectations.
• tests/python/s_tir/schedule/test_tir_schedule_cache_read_write.py
  • Removed the test_cache_read_allocate_const and test_cache_write_allocate_const functions.
• tests/python/s_tir/schedule/test_tir_schedule_compute_at.py
  • Removed the test_compute_at_allocate_const and test_compute_inline_allocate_const functions.
• tests/python/tir-transform/test_tir_transform_extract_constants.py
  • Removed file containing tests for the ExtractPrimFuncConstants pass.
• tests/python/tvmscript/test_tvmscript_complete.py
  • Removed the test_access_region_for_decl_buffer function.
• tests/python/tvmscript/test_tvmscript_ir_builder_tir.py
  • Removed the test_ir_builder_tir_allocate_const function.
• tests/python/tvmscript/test_tvmscript_roundtrip.py
  • Removed the module_const and constant functions and their references.

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 is a large refactoring that phases out the AllocConst node from TIR. The changes are extensive, touching many parts of the codebase including core IR definitions, statement functors, IR builders, various transformation passes, codegen backends, and tests. The removal of AllocateConst and its related logic (AllocateConstNode, AllocateConstFrame, BindParams pass, ExtractPrimFuncConstants pass) seems to be thorough and consistent across all modified files. The code is cleaner and simplified as a result. The changes appear correct and align with the goal of decoupling data allocation and weight placement as described in the pull request description. I have no specific comments as the refactoring seems well-executed.