[REFACTOR][S-TIR] More migrations to s-tir

[tqchen]

This PR migrates more passes into s-tir namespace
and keep the tir minimal for low-level passes.

Also cleans up unused passes

[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 undertakes a significant refactoring effort to migrate a set of analysis functions and transform passes from the general tvm::tir namespace to the more specialized tvm::s_tir namespace. The primary goal is to consolidate S-TIR specific functionalities, making the tvm::tir namespace more minimal and focused on low-level operations. This change impacts both C++ headers and source files, as well as Python bindings and usage examples, ensuring that the codebase adheres to the new organizational structure for schedulable TensorIR components.

Highlights

• Analysis Function Migration: Migrated several TIR analysis functions, including EstimateTIRFlops, IsPureFunction, VerifyGPUCode, IdentifyMemCpy, CalculateAllocatedBytes, GetVTCMCompactionPasses, VerifyVTCMLimit, and OOBChecker, from the tvm::tir namespace to the tvm::s_tir namespace. This centralizes S-TIR specific analysis utilities.
• Transform Pass Migration: Moved various transform passes, such as RenewDefs, DefaultGPUSchedule, RemoveWeightLayoutRewriteBlock, RemoveStoreUndef, DecorateDeviceScope, and UseAssumeToReduceBranches, from the tvm::tir::transform namespace to tvm::s_tir::transform. This ensures that these passes are correctly associated with the S-TIR framework.
• Adreno Backend Refactoring: Renamed and relocated Adreno backend transform headers and source files from tvm/tir/backend/adreno to tvm/s_tir/backend/adreno, updating their namespaces and FFI registrations accordingly. This aligns the Adreno backend with the S-TIR structure.
• Codebase Updates: Updated numerous C++ and Python files across the repository to reflect the new s_tir namespace for the migrated analysis functions and transform passes. This includes changes in relax passes, meta_schedule components, and various test files, ensuring consistent usage of the S-TIR API.

🧠 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/s_tir/analysis.h
  • Added new s_tir namespace with various analysis functions and passes, including EstimateTIRFlops, IsPureFunction, VerifyGPUCode, IdentifyMemCpy, CalculateAllocatedBytes, GetVTCMCompactionPasses, VerifyVTCMLimit, and OOBChecker.
• include/tvm/s_tir/backend/adreno/transform.h
  • Renamed from include/tvm/tir/backend/adreno/transform.h.
  • Updated include guards and namespace to TVM_S_TIR_BACKEND_ADRENO_TRANSFORM_H_ and tvm::s_tir::backend::adreno::transform respectively.
  • Removed redundant using declarations for PassContextNode, PassInfo, PassInfoNode, PassNode, and Sequential.
• include/tvm/s_tir/transform.h
  • Added RenewDefs function.
  • Added transform passes: DefaultGPUSchedule, RemoveWeightLayoutRewriteBlock, RemoveStoreUndef, DecorateDeviceScope, and UseAssumeToReduceBranches.
• include/tvm/tir/analysis.h
  • Removed EstimateTIRFlops functions.
  • Removed IsPureFunction function.
  • Removed VerifyGPUCode function.
  • Removed GetVTCMCompactionPasses function.
  • Removed VerifyVTCMLimit functions.
  • Removed IdentifyMemCpy struct and function.
  • Removed CalculateAllocatedBytes functions.
  • Removed OOBChecker pass from transform namespace.
• include/tvm/tir/stmt_functor.h
  • Removed RenewDefs function declaration.
• include/tvm/tir/transform.h
  • Removed RewriteUnsafeSelect pass.
  • Removed InstrumentBoundCheckers pass.
  • Removed DecorateDeviceScope pass.
  • Removed ThreadSync pass.
  • Removed LowerThreadAllreduce pass.
  • Removed InferFragment pass.
  • Removed HoistIfThenElse pass.
  • Removed HoistExpression pass.
  • Removed LowerVtcmAlloc pass.
  • Removed LowerAsyncDMA pass.
  • Removed MergeSharedMemoryAllocations pass.
  • Removed RenormalizeSplitPattern pass.
  • Removed InjectPTXAsyncCopy pass.
  • Removed InjectPTXLDG32 pass.
  • Removed RemoveWeightLayoutRewriteBlock pass.
  • Removed InstrumentProfileIntrinsics pass.
  • Removed DefaultGPUSchedule pass.
  • Removed UseAssumeToReduceBranches pass.
  • Removed RemoveStoreUndef pass.
  • Removed VerifyVTCMLimit pass.
• python/tvm/contrib/hexagon/meta_schedule.py
  • Updated import of RemoveWeightLayoutRewriteBlock to tvm.s_tir.transform.
• python/tvm/contrib/msc/core/runtime/runner.py
  • Updated call to DefaultGPUSchedule to use tvm.s_tir.transform.
• python/tvm/contrib/msc/framework/tvm/runtime/runner.py
  • Updated calls to DefaultGPUSchedule to use tvm.s_tir.transform.
• python/tvm/s_tir/init.py
  • Added import for _ffi_api.
  • Added renew_defs function, wrapping the FFI API call.
• python/tvm/s_tir/analysis/init.py
  • Added import for tvm and Union.
  • Added Python wrappers for verify_gpu_code, calculate_allocated_bytes, estimate_tir_flops, OOBChecker, get_vtcm_compaction_passes, is_pure_function, and assert_pure_function.
• python/tvm/s_tir/backend/adreno/init.py
  • Added import for transform.
• python/tvm/s_tir/backend/adreno/pipeline.py
  • Updated calls to TextureFlatten and InjectTextureAlloc to use s_tir.backend.adreno.transform.
• python/tvm/s_tir/backend/adreno/transform/init.py
  • Renamed from python/tvm/tir/backend/adreno/transform/__init__.py.
• python/tvm/s_tir/backend/adreno/transform/_ffi_api.py
  • Renamed from python/tvm/tir/backend/adreno/transform/_ffi_api.py.
  • Updated FFI API initialization to s_tir.backend.adreno.transform.
• python/tvm/s_tir/backend/adreno/transform/transform.py
  • Renamed from python/tvm/tir/backend/adreno/transform/transform.py.
• python/tvm/s_tir/meta_schedule/builder/local_builder.py
  • Updated import of RemoveWeightLayoutRewriteBlock to tvm.s_tir.transform.
• python/tvm/s_tir/transform/transform.py
  • Added Python wrappers for DefaultGPUSchedule, RemoveWeightLayoutRewriteBlock, RemoveStoreUndef, DecorateDeviceScope, and UseAssumeToReduceBranches.
• python/tvm/tir/analysis/analysis.py
  • Removed imports for Dict, Union, and tvm.
  • Removed Python wrappers for verify_gpu_code, calculate_allocated_bytes, estimate_tir_flops, OOBChecker, has_if_then_else, get_vtcm_compaction_passes, is_pure_function, and assert_pure_function.
• python/tvm/tir/backend/adreno/init.py
  • Removed import for transform.
• python/tvm/tir/stmt_functor.py
  • Removed renew_defs function.
• python/tvm/tir/transform/transform.py
  • Removed Optional import.
  • Removed Python wrappers for RemoveStoreUndef, RewriteUnsafeSelect, InstrumentBoundCheckers, DecorateDeviceScope, ThreadSync, LowerThreadAllreduce, InferFragment, HoistIfThenElse, HoistExpression, MergeSharedMemoryAllocations, RenormalizeSplitPattern, InjectPTXAsyncCopy, RemoveWeightLayoutRewriteBlock, InstrumentProfileIntrinsics, DefaultGPUSchedule, UseAssumeToReduceBranches, LowerAsyncDMA, InjectPTXLDG32, and LowerVtcmAlloc.
• src/relax/distributed/transform/lower_global_view_to_local_view.cc
  • Added include for tvm/s_tir/transform.h.
  • Updated call to RenewDefs to use s_tir::RenewDefs.
• src/relax/transform/attach_attr_layout_free_buffers.cc
  • Added include for tvm/s_tir/transform.h.
  • Updated call to RenewDefs to use s_tir::RenewDefs.
• src/relax/transform/compute_prim_value.cc
  • Added include for tvm/s_tir/transform.h.
  • Updated call to RenewDefs to use s_tir::RenewDefs.
• src/relax/transform/fuse_tir.cc
  • Added include for tvm/s_tir/transform.h.
  • Updated calls to RenewDefs to use s_tir::RenewDefs.
• src/relax/transform/split_call_tir_by_pattern.cc
  • Added include for tvm/s_tir/transform.h.
  • Updated calls to RenewDefs to use s_tir::RenewDefs.
• src/relax/transform/split_layout_rewrite_preproc.cc
  • Added include for tvm/s_tir/transform.h.
  • Updated calls to RenewDefs to use s_tir::RenewDefs.
• src/s_tir/analysis/calculate_allocated_memory.cc
  • Updated namespace from tvm::tir to tvm::s_tir.
  • Added using namespace tvm::tir;.
  • Updated FFI API registration from tir.analysis.calculate_allocated_bytes to s_tir.analysis.calculate_allocated_bytes.
  • Updated GetVTCMCapacity to use tvm::transform::PassContext.
  • Updated FFI API registration from tir.analysis.get_vtcm_compaction_passes to s_tir.analysis.get_vtcm_compaction_passes.
  • Removed FFI API registration for tir.transform.VerifyVTCMLimit.
• src/s_tir/analysis/estimate_flops.cc
  • Renamed from src/tir/analysis/estimate_flops.cc.
  • Updated namespace from tvm::tir to tvm::s_tir.
  • Added using namespace tvm::tir;.
  • Updated FFI API registration from tir.analysis.EstimateTIRFlops to s_tir.analysis.EstimateTIRFlops.
• src/s_tir/analysis/find_anchor_sblock.cc
  • Added include for tvm/s_tir/analysis.h.
  • Updated call to EstimateTIRFlops to use s_tir::EstimateTIRFlops.
• src/s_tir/analysis/identify_memcpy.cc
  • Renamed from src/tir/analysis/identify_memcpy.cc.
  • Updated namespace from tvm::tir to tvm::s_tir.
  • Added using namespace tvm::tir;.
  • Updated FFI API registration from tir.analysis._identify_memcpy to s_tir.analysis._identify_memcpy.
• src/s_tir/analysis/is_pure_function.cc
  • Renamed from src/tir/analysis/is_pure_function.cc.
  • Updated include path for tir_visitor_with_path.h.
  • Updated namespace from tvm::tir to tvm::s_tir.
  • Added using namespace tvm::tir;.
  • Updated FFI API registration from tir.analysis.is_pure_function to s_tir.analysis.is_pure_function.
• src/s_tir/analysis/oob_checker.cc
  • Renamed from src/tir/analysis/oob_checker.cc.
  • Updated include path for error.h.
  • Updated namespace from tvm::tir to tvm::s_tir.
  • Added using namespace tvm::tir;.
  • Updated pass creation and FFI API registration to use s_tir.analysis.OOBChecker.
• src/s_tir/analysis/verify_gpu_code.cc
  • Renamed from src/tir/analysis/verify_gpu_code.cc.
  • Updated include path for ir_utils.h.
  • Updated namespace from tvm::tir to tvm::s_tir.
  • Added using namespace tvm::tir;.
  • Updated attribute key check to tir::attr::thread_extent.
  • Updated FFI API registration and pass creation to use s_tir.analysis.verify_gpu_code and s_tir.VerifyGPUCode.
• src/s_tir/backend/adreno/inject_texture_alloc.cc
  • Renamed from src/tir/backend/adreno/inject_texture_alloc.cc.
  • Added include for tvm/s_tir/backend/adreno/transform.h.
  • Updated include path for ir_utils.h.
  • Updated namespace from tvm::tir to tvm::s_tir.
  • Added using namespace tvm::tir;.
  • Updated pass creation and FFI API registration to use s_tir.backend.adreno.InjectTextureAlloc.
• src/s_tir/backend/adreno/texture_flatten.cc
  • Renamed from src/tir/backend/adreno/texture_flatten.cc.
  • Added include for tvm/s_tir/backend/adreno/transform.h.
  • Updated namespace from tvm::tir to tvm::s_tir.
  • Added using namespace tvm::tir;.
  • Updated pass creation and FFI API registration to use s_tir.backend.adreno.TextureFlatten.
• src/s_tir/meta_schedule/arg_info.cc
  • Added include for tvm/s_tir/transform.h.
  • Updated call to RemoveWeightLayoutRewriteBlock to use transform::RemoveWeightLayoutRewriteBlock.
• src/s_tir/meta_schedule/feature_extractor/per_store_feature.cc
  • Updated call to RemoveWeightLayoutRewriteBlock to use s_tir::transform.
• src/s_tir/meta_schedule/postproc/verify_gpu_code.cc
  • Updated call to VerifyGPUCode to use s_tir::VerifyGPUCode.
• src/s_tir/meta_schedule/postproc/verify_vtcm_limit.cc
  • Added include for tvm/s_tir/analysis.h.
  • Updated calls to VerifyVTCMLimit and GetVTCMCompactionPasses to use s_tir namespace.
• src/s_tir/meta_schedule/task_scheduler/task_scheduler.cc
  • Added include for tvm/s_tir/analysis.h.
  • Updated call to EstimateTIRFlops to use s_tir::EstimateTIRFlops.
• src/s_tir/transform/decorate_device_scope.cc
  • Renamed from src/tir/transform/decorate_device_scope.cc.
  • Added include for tvm/s_tir/transform.h.
  • Updated namespace from tvm::tir to tvm::s_tir.
  • Added using namespace tvm::tir;.
  • Renamed internal function DecorateDeviceScope to DecorateDeviceScopeImpl.
  • Updated pass creation and FFI API registration to use s_tir.DecorateDeviceScope.
• src/s_tir/transform/default_gpu_schedule.cc
  • Renamed from src/tir/transform/default_gpu_schedule.cc.
  • Updated include path for utils.h.
  • Updated namespace from tvm::tir to tvm::s_tir.
  • Added using namespace tvm::tir;.
  • Updated attribute check to tir::attr::kIsScheduled.
  • Updated pass creation and FFI API registration to use s_tir.transform.DefaultGPUSchedule.
• src/s_tir/transform/lower_async_dma.cc
  • Updated include for tvm/s_tir/analysis.h.
  • Updated call to IdentifyMemCpy to use s_tir::IdentifyMemCpy.
• src/s_tir/transform/remove_store_undef.cc
  • Renamed from src/tir/transform/remove_store_undef.cc.
  • Added include for tvm/s_tir/transform.h.
  • Updated namespace from tvm::tir to tvm::s_tir.
  • Added using namespace tvm::tir;.
  • Updated pass names and FFI API registration to use s_tir.RemoveStoreUndefInternal, s_tir.ValidateAllUndefRemoved, and s_tir.RemoveStoreUndef.
• src/s_tir/transform/remove_weight_layout_rewrite_block.cc
  • Renamed from src/tir/transform/remove_weight_layout_rewrite_block.cc.
  • Added include for tvm/s_tir/transform.h.
  • Updated namespace from tvm::tir to tvm::s_tir.
  • Added using namespace tvm::tir;.
  • Updated attribute check to tir::attr::meta_schedule_layout_rewrite_preproc.
  • Updated pass creation and FFI API registration to use s_tir.RemoveWeightLayoutRewriteBlock.
• src/s_tir/transform/renew_defs.cc
  • Renamed from src/tir/transform/renew_defs.cc.
  • Added include for tvm/s_tir/transform.h.
  • Updated include path for functor_common.h.
  • Updated namespace from tvm::tir to tvm::s_tir.
  • Added using namespace tvm::tir;.
  • Updated FFI API registration from tir.RenewDefs to s_tir.RenewDefs.
• src/s_tir/transform/using_assume_to_reduce_branches.cc
  • Renamed from src/tir/transform/using_assume_to_reduce_branches.cc.
  • Added include for tvm/s_tir/transform.h.
  • Updated namespace from tvm::tir to tvm::s_tir.
  • Added using namespace tvm::tir;.
  • Updated pass creation and FFI API registration to use s_tir.UseAssumeToReduceBranches.
• src/tir/ir/function.cc
  • Updated include for tvm/s_tir/analysis.h.
  • Updated call to IsPureFunction to use s_tir::IsPureFunction.
• src/tir/transform/bind_target.cc
  • Added include for tvm/s_tir/transform.h.
  • Updated call to RenewDefs to use s_tir::RenewDefs.
• tests/python/contrib/test_msc/test_plugin.py
  • Updated call to DefaultGPUSchedule to use tvm.s_tir.transform.
• tests/python/contrib/test_msc/test_translate_tensorrt.py
  • Updated call to DefaultGPUSchedule to use tvm.s_tir.transform.
• tests/python/relax/test_codegen_cublas.py
  • Updated call to DefaultGPUSchedule to use tvm.s_tir.transform.
• tests/python/relax/test_codegen_cutlass.py
  • Updated calls to DefaultGPUSchedule to use tvm.s_tir.transform.
• tests/python/relax/test_contrib_vllm.py
  • Updated call to DefaultGPUSchedule to use tvm.s_tir.transform.
• tests/python/relax/test_eliminate_pad_branch_using_buffer_assumption.py
  • Updated calls to UseAssumeToReduceBranches to use tvm.s_tir.transform.
• tests/python/relax/test_frontend_nn_op.py
  • Updated import to include s_tir.
  • Updated calls to DefaultGPUSchedule to use s_tir.transform.
• tests/python/relax/test_transform_codegen_pass.py
  • Updated import to include s_tir.
  • Updated call to DefaultGPUSchedule to use s_tir.transform.
• tests/python/relax/test_transform_meta_schedule_tuning.py
  • Updated call to DefaultGPUSchedule to use tvm.s_tir.transform.
• tests/python/relax/test_vm_cuda_graph.py
  • Updated call to DefaultGPUSchedule to use tvm.s_tir.transform.
• tests/python/relax/test_vm_multi_device.py
  • Updated call to DefaultGPUSchedule to use tvm.s_tir.transform.
• tests/python/s_tir/analysis/test_s_tir_analysis_calculate_allocated_memory.py
  • Renamed from tests/python/tir-analysis/test_tir_analysis_calculate_allocated_memory.py.
  • Updated calls to calculate_allocated_bytes to use tvm.s_tir.analysis.
• tests/python/s_tir/analysis/test_s_tir_analysis_estimate_tir_flops.py
  • Renamed from tests/python/tir-analysis/test_tir_analysis_estimate_tir_flops.py.
  • Updated import and call to estimate_tir_flops to use tvm.s_tir.analysis.
• tests/python/s_tir/analysis/test_s_tir_analysis_identify_memcpy.py
  • Renamed from tests/python/tir-analysis/test_tir_analysis_identify_memcpy.py.
  • Updated call to _identify_memcpy to use tvm.s_tir.analysis.
• tests/python/s_tir/analysis/test_s_tir_analysis_is_pure_function.py
  • Renamed from tests/python/tir-analysis/test_tir_analysis_is_pure_function.py.
  • Updated import of is_pure_function and assert_pure_function to use tvm.s_tir.analysis.
• tests/python/s_tir/analysis/test_s_tir_analysis_oob.py
  • Renamed from tests/python/tir-analysis/test_tir_analysis_oob.py.
  • Updated calls to OOBChecker to use tvm.s_tir.analysis.
• tests/python/s_tir/test_s_tir_renew_defs.py
  • Renamed from tests/python/tir-base/test_tir_renew_defs.py.
  • Updated calls to renew_defs to use tvm.s_tir.
• tests/python/s_tir/transform/test_s_tir_transform_decorate_device_scope.py
  • Renamed from tests/python/tir-transform/test_tir_transform_decorate_device_scope.py.
  • Updated call to DecorateDeviceScope to use tvm.s_tir.transform.
• tests/python/s_tir/transform/test_s_tir_transform_default_gpu_schedule.py
  • Renamed from tests/python/tir-transform/test_transform_default_gpu_schedule.py.
  • Updated import of DefaultGPUSchedule to use tvm.s_tir.transform.
• tests/python/s_tir/transform/test_s_tir_transform_inject_double_buffer.py
  • Updated call to ThreadSync to use tvm.s_tir.transform.
• tests/python/s_tir/transform/test_s_tir_transform_remove_undef.py
  • Renamed from tests/python/tir-transform/test_tir_transform_remove_undef.py.
  • Updated calls to RemoveStoreUndef to use tvm.s_tir.transform.
• tests/python/s_tir/transform/test_s_tir_transform_remove_weight_layout_rewrite_block.py
  • Renamed from tests/python/tir-transform/test_tir_transform_remove_weight_layout_rewrite_block.py.
  • Updated call to RemoveWeightLayoutRewriteBlock to use tvm.s_tir.transform.

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 migrates many passes and analysis functions from the tir namespace to the s_tir namespace. The changes are extensive but appear to be consistent and correct across both C++ and Python code. I've identified a couple of minor issues with Python type hints that should be addressed for better code clarity and correctness.