[REFACTOR][S-TIR] Migrate more transform to s_tir

[tqchen]

This PR migrate more transformations into s_tir, to reduce the passes in tir namespace.

[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 substantial refactoring effort within the TVM Tensor IR (TIR) framework. Its primary goal is to streamline the tir namespace by relocating a suite of transformation passes to the more specialized s_tir namespace. This strategic move ensures that the core tir namespace remains focused on essential, deterministic passes vital for low-level code generation, while allowing s_tir to house more advanced or target-specific optimizations. The changes span across C++ implementations, Python bindings, and various integration points, ensuring a clean separation and proper functionality of the migrated passes.

Highlights

• Namespace Migration: Migrated a significant number of TIR transformation passes from the tvm::tir::transform namespace to the tvm::s_tir::transform namespace in C++ source and header files.
• Python Bindings and Usage Updates: Updated Python FFI bindings for the migrated passes and adjusted their usage in Python pipelines (e.g., Adreno pipeline) and test files to reflect the new s_tir.transform module.
• Codebase Reorganization: Renamed numerous C++ source, header, and Python test files to align with the s_tir namespace, ensuring consistent code organization.

🧠 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/transform.h
  • Added declarations for HoistIfThenElse, HoistExpression, RenormalizeSplitPattern, RewriteUnsafeSelect, InstrumentBoundCheckers, InjectPTXLDG32, InstrumentProfileIntrinsics, LowerVtcmAlloc, ThreadSync, InferFragment, LowerThreadAllreduce, LowerAsyncDMA, InjectPTXAsyncCopy, and MergeSharedMemoryAllocations passes.
• python/tvm/s_tir/backend/adreno/pipeline.py
  • Updated calls to LowerAsyncDMA, HoistIfThenElse, RenormalizeSplitPattern, RewriteUnsafeSelect, InstrumentBoundCheckers, InjectPTXLDG32, InstrumentProfileIntrinsics, VerifyVTCMLimit, LowerVtcmAlloc, ThreadSync, InferFragment, LowerThreadAllreduce, InjectPTXAsyncCopy, and MergeSharedMemoryAllocations to use the s_tir.transform namespace.
• python/tvm/s_tir/pipeline.py
  • Updated calls to LowerAsyncDMA, HoistIfThenElse, RenormalizeSplitPattern, RewriteUnsafeSelect, InstrumentBoundCheckers, InjectPTXLDG32, InstrumentProfileIntrinsics, VerifyVTCMLimit, LowerVtcmAlloc, ThreadSync, InferFragment, LowerThreadAllreduce, InjectPTXAsyncCopy, and MergeSharedMemoryAllocations to use the s_tir.transform namespace.
• python/tvm/s_tir/transform/transform.py
  • Added Python FFI bindings for HoistIfThenElse, RenormalizeSplitPattern, RewriteUnsafeSelect, InstrumentBoundCheckers, InjectPTXLDG32, InstrumentProfileIntrinsics, VerifyVTCMLimit, LowerVtcmAlloc, ThreadSync, InferFragment, LowerThreadAllreduce, LowerAsyncDMA, InjectPTXAsyncCopy, and MergeSharedMemoryAllocations.
• src/s_tir/analysis/calculate_allocated_memory.cc
  • Renamed from src/tir/analysis/calculate_allocated_memory.cc.
  • Updated the pass name registration to s_tir.VerifyVTCMLimit.
  • Added FFI registration for s_tir.transform.VerifyVTCMLimit.
• src/s_tir/meta_schedule/postproc/verify_gpu_code.cc
  • Updated the call to MergeSharedMemoryAllocations from tir::transform to s_tir::transform.
• src/s_tir/transform/bound_checker.cc
  • Renamed from src/tir/transform/bound_checker.cc.
  • Updated include path to tvm/s_tir/transform.h.
  • Changed namespace from tvm::tir to tvm::s_tir.
  • Updated pass name and FFI registration to s_tir.InstrumentBoundCheckers.
• src/s_tir/transform/hoist_expression.cc
  • Renamed from src/tir/transform/hoist_expression.cc.
  • Updated include path to tvm/s_tir/transform.h.
  • Changed namespace from tvm::tir to tvm::s_tir.
  • Updated pass config options and FFI registrations for HoistExpression and HoistIfThenElse to use s_tir prefix.
• src/s_tir/transform/inject_ptx_async_copy.cc
  • Renamed from src/tir/transform/inject_ptx_async_copy.cc.
  • Updated include path to tvm/s_tir/transform.h.
  • Changed namespace from tvm::tir to tvm::s_tir.
  • Updated pass name and FFI registration to s_tir.InjectPTXAsyncCopy.
• src/s_tir/transform/inject_ptx_ldg32.cc
  • Renamed from src/tir/transform/inject_ptx_ldg32.cc.
  • Updated include path to tvm/s_tir/transform.h.
  • Changed namespace from tvm::tir to tvm::s_tir.
  • Updated pass name and FFI registration to s_tir.InjectPTXLDG32.
• src/s_tir/transform/lower_async_dma.cc
  • Renamed from src/tir/transform/lower_async_dma.cc.
  • Updated include path to tvm/s_tir/transform.h.
  • Changed namespace from tvm::tir to tvm::s_tir.
  • Updated pass name and FFI registration to s_tir.LowerAsyncDMA.
• src/s_tir/transform/lower_thread_allreduce.cc
  • Renamed from src/tir/transform/lower_thread_allreduce.cc.
  • Updated include path to tvm/s_tir/transform.h.
  • Changed namespace from tvm::tir to tvm::s_tir.
  • Updated pass name and FFI registration to s_tir.LowerThreadAllreduce.
• src/s_tir/transform/lower_vtcm_alloc.cc
  • Renamed from src/tir/transform/lower_vtcm_alloc.cc.
  • Updated include path to tvm/s_tir/transform.h.
  • Changed namespace from tvm::tir to tvm::s_tir.
  • Updated pass name and FFI registration to s_tir.LowerVtcmAlloc.
• src/s_tir/transform/merge_shared_memory_allocations.cc
  • Renamed from src/tir/transform/merge_shared_memory_allocations.cc.
  • Updated include path to tvm/s_tir/transform.h.
  • Changed namespace from tvm::tir to tvm::s_tir.
  • Updated pass name and FFI registration to s_tir.MergeSharedMemoryAllocations.
• src/s_tir/transform/profile_instrumentation.cc
  • Renamed from src/tir/transform/profile_instrumentation.cc.
  • Updated include path to tvm/s_tir/transform.h.
  • Changed namespace from tvm::tir to tvm::s_tir.
  • Updated pass name and FFI registration to s_tir.InstrumentProfileIntrinsics.
• src/s_tir/transform/renormalize_split_pattern.cc
  • Renamed from src/tir/transform/renormalize_split_pattern.cc.
  • Updated include path to tvm/s_tir/transform.h.
  • Changed namespace from tvm::tir to tvm::s_tir.
  • Updated pass name and FFI registration to s_tir.RenormalizeSplitPattern.
• src/s_tir/transform/rewrite_unsafe_select.cc
  • Renamed from src/tir/transform/rewrite_unsafe_select.cc.
  • Updated include path to tvm/s_tir/transform.h.
  • Changed namespace from tvm::tir to tvm::s_tir.
  • Updated pass name and FFI registration to s_tir.RewriteUnsafeSelect.
• src/s_tir/transform/storage_access.cc
  • Renamed from src/tir/transform/storage_access.cc.
  • Changed namespace from tvm::tir to tvm::s_tir.
• src/s_tir/transform/storage_access.h
  • Renamed from src/tir/transform/storage_access.h.
  • Updated include guards from TVM_TIR_TRANSFORM_STORAGE_ACCESS_H_ to TVM_S_TIR_TRANSFORM_STORAGE_ACCESS_H_.
  • Changed namespace from tvm::tir to tvm::s_tir.
• src/s_tir/transform/tensorcore_infer_fragment.cc
  • Renamed from src/tir/transform/tensorcore_infer_fragment.cc.
  • Updated include path to tvm/s_tir/transform.h.
  • Changed namespace from tvm::tir to tvm::s_tir.
  • Updated pass name and FFI registration to s_tir.InferFragment.
• src/s_tir/transform/thread_storage_sync.cc
  • Renamed from src/tir/transform/thread_storage_sync.cc.
  • Updated include path to tvm/s_tir/transform.h.
  • Changed namespace from tvm::tir to tvm::s_tir.
  • Updated pass name and FFI registration to s_tir.ThreadSync.
• tests/python/s_tir/transform/test_s_tir_transform_hoist_expression.py
  • Renamed from tests/python/tir-transform/test_tir_transform_hoist_expression.py.
  • Updated imports from tvm.tir to tvm.s_tir.
  • Updated calls to tvm.tir.transform.HoistExpression to tvm.s_tir.transform.HoistExpression.
• tests/python/s_tir/transform/test_s_tir_transform_hoist_if.py
  • Renamed from tests/python/tir-transform/test_tir_transform_hoist_if.py.
  • Updated imports from tvm.te to tvm.te, s_tir.
  • Updated calls to tvm.tir.transform.HoistIfThenElse to tvm.s_tir.transform.HoistIfThenElse.
• tests/python/s_tir/transform/test_s_tir_transform_inject_ptx_async_copy.py
  • Renamed from tests/python/tir-transform/test_tir_transform_inject_ptx_async_copy.py.
  • Updated imports to include tvm.s_tir.
  • Updated calls to tvm.tir.transform.InjectPTXAsyncCopy and tvm.tir.transform.MergeSharedMemoryAllocations to use s_tir.transform.
• tests/python/s_tir/transform/test_s_tir_transform_inject_ptx_ldg32.py
  • Renamed from tests/python/tir-transform/test_tir_transform_inject_ptx_ldg32.py.
  • Updated imports to include tvm.s_tir.
  • Updated calls to tvm.tir.transform.InjectPTXLDG32 to tvm.s_tir.transform.InjectPTXLDG32.
• tests/python/s_tir/transform/test_s_tir_transform_lower_thread_all_reduce.py
  • Renamed from tests/python/tir-transform/test_tir_transform_lower_thread_all_reduce.py.
  • Updated imports to include tvm.s_tir.
  • Updated calls to tvm.tir.transform.LowerThreadAllreduce to tvm.s_tir.transform.LowerThreadAllreduce.
• tests/python/s_tir/transform/test_s_tir_transform_merge_dynamic_shared_memory_allocations.py
  • Renamed from tests/python/tir-transform/test_tir_transform_merge_dynamic_shared_memory_allocations.py.
  • Updated imports to include tvm.s_tir.
  • Updated calls to tvm.tir.transform.MergeSharedMemoryAllocations to tvm.s_tir.transform.MergeSharedMemoryAllocations.
• tests/python/s_tir/transform/test_s_tir_transform_profiling_instr.py
  • Renamed from tests/python/tir-transform/test_tir_transform_profiling_instr.py.
  • Updated imports to include tvm.s_tir.
  • Updated calls to tvm.tir.transform.InstrumentProfileIntrinsics to tvm.s_tir.transform.InstrumentProfileIntrinsics.
• tests/python/s_tir/transform/test_s_tir_transform_renormalize_split_pattern.py
  • Renamed from tests/python/tir-transform/test_tir_transform_renormalize_split_pattern.py.
  • Updated imports to include tvm.s_tir.
  • Updated calls to tvm.tir.transform.RenormalizeSplitPattern to tvm.s_tir.transform.RenormalizeSplitPattern.
• tests/python/s_tir/transform/test_s_tir_transform_rewrite_unsafe_select.py
  • Renamed from tests/python/tir-transform/test_tir_transform_rewrite_unsafe_select.py.
  • Updated imports to include tvm.s_tir.
  • Updated calls to tvm.tir.transform.RewriteUnsafeSelect to tvm.s_tir.transform.RewriteUnsafeSelect.
• tests/python/s_tir/transform/test_s_tir_transform_thread_sync.py
  • Renamed from tests/python/tir-transform/test_tir_transform_thread_sync.py.
  • Updated imports to include tvm.s_tir.
  • Updated calls to tvm.tir.transform.ThreadSync to tvm.s_tir.transform.ThreadSync.

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 significant refactoring that migrates many TIR transformation passes from the tir namespace to the new s_tir namespace. The changes primarily consist of moving files, updating namespaces and include paths, and adjusting Python bindings and pipelines accordingly. The refactoring is well-executed, but I've found a few inconsistencies and typos that should be addressed to ensure correctness and maintain code quality. My main concerns are around pass configuration keys that were not updated during the migration, which could lead to tests not running with the intended configurations.