Skip to content

[REFACTOR][S-TIR] Initialize the s_tir module#18712

Merged
tqchen merged 1 commit into
apache:mainfrom
tqchen:s-tir
Feb 5, 2026
Merged

[REFACTOR][S-TIR] Initialize the s_tir module#18712
tqchen merged 1 commit into
apache:mainfrom
tqchen:s-tir

Conversation

@tqchen

@tqchen tqchen commented Feb 4, 2026

Copy link
Copy Markdown
Member

This PR initalizes the s_tir for scheduable TensorIR. The change mainly starts from python side, the we will gradually move towards the c++ side in followup PRs. The python main change:

tir.Schedule => s_tir.Schedule

@gemini-code-assist

Copy link
Copy Markdown
Contributor

Summary of Changes

Hello @tqchen, I'm Gemini Code Assist1! 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 establishes the foundational s_tir module within the TVM framework, specifically for schedulable TensorIR. The changes are concentrated on the Python API, transitioning existing tvm.tir.Schedule references to the new tvm.s_tir.Schedule namespace. This strategic refactoring aims to modularize and enhance the TensorIR scheduling capabilities, setting the stage for subsequent C++-side implementations.

Highlights

  • Introduction of S-TIR Module: This pull request initializes the s_tir module, which is designed for schedulable TensorIR, laying the groundwork for future development.
  • Python API Refactoring: The primary change involves renaming the Python-side API from tvm.tir.Schedule to tvm.s_tir.Schedule across numerous files.
  • File Renames and Relocations: Several core tir schedule-related files, including block_dependence_info.py, block_scope.py, pipeline.py, and all files within tir/schedule and tir/tensor_intrin, have been renamed and moved to the new s_tir directory.
  • Documentation and Configuration Updates: Documentation files (.rst) and mypy configuration (mypy.ini) have been updated to reflect the new s_tir module and its associated schedule classes.
  • Pipeline Default Change: The tvm.tir.pipeline now defaults to using the s_tir pipeline, indicating a shift in the core scheduling mechanism.

🧠 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
  • apps/android_rpc/tests/android_rpc_test.py
    • Updated tvm.tir.Schedule to tvm.s_tir.Schedule.
  • apps/ios_rpc/tests/ios_rpc_test.py
    • Updated tvm.tir.Schedule to tvm.s_tir.Schedule.
  • docs/deep_dive/tensor_ir/tutorials/tir_transformation.py
    • Updated tvm.tir.Schedule to tvm.s_tir.Schedule.
  • docs/how_to/tutorials/cross_compilation_and_rpc.py
    • Updated tvm.tir.Schedule to tvm.s_tir.Schedule.
  • docs/reference/api/python/index.rst
    • Removed tir/schedule from tvm.tir toctree.
    • Added s_tir/schedule to a new tvm.s_tir toctree.
  • docs/reference/api/python/tir/schedule.rst
    • Renamed to docs/reference/api/python/s_tir/schedule.rst.
    • Updated module reference from tvm.tir.schedule to tvm.s_tir.schedule.
  • jvm/core/src/test/scripts/prepare_test_libs.py
    • Updated tvm.tir.Schedule to tvm.s_tir.Schedule.
  • mypy.ini
    • Updated mypy configuration to refer to python.tvm.s_tir.schedule.* instead of python.tvm.tir.schedule.*.
  • python/tvm/init.py
    • Added import s_tir to the tvm package initialization.
  • python/tvm/dlight/adreno/convolution.py
    • Imported s_tir module.
    • Updated type hints and instantiations from tir.Schedule to s_tir.Schedule.
  • python/tvm/dlight/adreno/fallback.py
    • Imported s_tir module.
    • Updated type hints and instantiations from tir.Schedule to s_tir.Schedule.
  • python/tvm/dlight/adreno/layout_transform.py
    • Imported s_tir module.
    • Updated type hints and instantiations from tir.Schedule to s_tir.Schedule.
  • python/tvm/dlight/adreno/pool.py
    • Imported s_tir module.
    • Updated type hints and instantiations from tir.Schedule to s_tir.Schedule.
  • python/tvm/dlight/adreno/utils.py
    • Imported s_tir module.
    • Updated type hints from tir.Schedule to s_tir.Schedule.
  • python/tvm/dlight/analysis/common_analysis.py
    • Imported s_tir module.
    • Updated type hints and instantiations from tir.Schedule to s_tir.Schedule.
  • python/tvm/dlight/analysis/gemv.py
    • Imported s_tir module.
    • Updated type hints from tir.Schedule to s_tir.Schedule.
  • python/tvm/dlight/base/common_schedules.py
    • Imported s_tir module.
    • Updated type hints from tir.Schedule to s_tir.Schedule.
  • python/tvm/dlight/base/schedule_rule.py
    • Imported s_tir module.
    • Updated type hints from tir.Schedule to s_tir.Schedule.
  • python/tvm/dlight/base/transform.py
    • Imported s_tir module.
    • Updated type hints from tir.Schedule to s_tir.Schedule.
  • python/tvm/dlight/base/utils.py
    • Imported s_tir module.
    • Updated type hints from tir.Schedule to s_tir.Schedule.
  • python/tvm/dlight/cpu/gemv.py
    • Imported s_tir module.
    • Updated type hints and instantiations from tir.Schedule to s_tir.Schedule.
  • python/tvm/dlight/gpu/fallback.py
    • Imported s_tir module.
    • Updated type hints and instantiations from tir.Schedule to s_tir.Schedule.
  • python/tvm/dlight/gpu/gemv.py
    • Imported s_tir module.
    • Updated type hints and instantiations from tir.Schedule to s_tir.Schedule.
  • python/tvm/dlight/gpu/general_reduction.py
    • Imported s_tir module.
    • Updated type hints and instantiations from tir.Schedule to s_tir.Schedule.
  • python/tvm/dlight/gpu/low_batch_gemv.py
    • Imported s_tir module.
    • Updated type hints and instantiations from tir.Schedule to s_tir.Schedule.
  • python/tvm/dlight/gpu/matmul.py
    • Imported s_tir module.
    • Updated type hints and instantiations from tir.Schedule to s_tir.Schedule.
    • Updated tvm.tir.tensor_intrin.metal and tvm.tir.tensor_intrin.cuda imports to tvm.s_tir.tensor_intrin.
  • python/tvm/dlight/gpu/reduction.py
    • Imported s_tir module.
    • Updated type hints and instantiations from tir.Schedule to s_tir.Schedule.
  • python/tvm/dlight/gpu/rmsnorm.py
    • Updated instantiations from tir.Schedule to tvm.s_tir.Schedule.
  • python/tvm/dlight/gpu/transpose.py
    • Imported s_tir module.
    • Updated type hints and instantiations from tir.Schedule to s_tir.Schedule.
  • python/tvm/exec/gpu_memory_bandwidth.py
    • Removed tir import.
    • Updated type hints and instantiations from tir.Schedule to tvm.s_tir.Schedule.
  • python/tvm/meta_schedule/builder/local_builder.py
    • Updated tvm.tir.tensor_intrin import to tvm.s_tir.tensor_intrin.
  • python/tvm/meta_schedule/database/database.py
    • Updated tvm.tir.schedule import to tvm.s_tir.schedule.
    • Updated type hints for schedule.
  • python/tvm/meta_schedule/database/schedule_fn_database.py
    • Updated tvm.tir import to tvm.s_tir.
  • python/tvm/meta_schedule/mutator/mutator.py
    • Updated tvm.tir.schedule import to tvm.s_tir.schedule.
  • python/tvm/meta_schedule/post_optimization/space.py
    • Updated docstring.
    • Updated tvm.tir import to tvm.s_tir.
  • python/tvm/meta_schedule/postproc/postproc.py
    • Updated tvm.tir.schedule import to tvm.s_tir.schedule.
    • Updated type hints for sch parameter.
  • python/tvm/meta_schedule/schedule/cuda/layout_transform.py
    • Updated tvm.tir.schedule import to tvm.s_tir.schedule.
    • Updated type hints for sch parameter.
  • python/tvm/meta_schedule/schedule_rule/multi_level_tiling.py
    • Updated tvm.tir.schedule import to tvm.s_tir.schedule.
  • python/tvm/meta_schedule/schedule_rule/schedule_rule.py
    • Updated tvm.tir.schedule import to tvm.s_tir.schedule.
    • Updated type hints for sch parameter.
  • python/tvm/meta_schedule/search_strategy/search_strategy.py
    • Updated tvm.tir.schedule import to tvm.s_tir.schedule.
    • Updated type hints for sch parameter.
  • python/tvm/meta_schedule/space_generator/space_generator.py
    • Updated tvm.tir.schedule import to tvm.s_tir.schedule.
    • Updated type hints for design_spaces return value.
  • python/tvm/meta_schedule/testing/dummy_object.py
    • Updated tvm.tir.schedule import to tvm.s_tir.schedule.
  • python/tvm/meta_schedule/testing/space_generation.py
    • Updated tvm.tir imports to tvm.s_tir.
  • python/tvm/meta_schedule/testing/tune_te.py
    • Removed tir import.
    • Updated type hints from tir.Schedule to s_tir.Schedule.
  • python/tvm/meta_schedule/testing/validate_database.py
    • Updated tvm.tir imports to tvm.s_tir.
  • python/tvm/meta_schedule/tir_integration.py
    • Added import tvm.
    • Updated type hints from tir.Schedule to tvm.s_tir.Schedule.
  • python/tvm/meta_schedule/trace_apply.py
    • Updated tvm.tir.schedule import to tvm.s_tir.
  • python/tvm/meta_schedule/tune_context.py
    • Updated tvm.tir imports to tvm.s_tir.
  • python/tvm/relax/frontend/nn/llm/kv_cache.py
    • Imported s_tir module.
    • Updated type hints and instantiations from tir.Schedule to tvm.s_tir.Schedule.
  • python/tvm/relax/frontend/nn/llm/tree_attn.py
    • Imported s_tir module.
    • Updated type hints and instantiations from tir.Schedule to s_tir.Schedule.
  • python/tvm/relax/transform/legalize_ops/manipulate.py
    • Imported s_tir module.
    • Updated type hints and instantiations from tir.Schedule to tvm.s_tir.Schedule.
  • python/tvm/s_tir/init.py
    • Added new file to initialize the s_tir namespace.
  • python/tvm/s_tir/_ffi_api.py
    • Added new file for FFI APIs of tvm.s_tir.
  • python/tvm/s_tir/backend/init.py
    • Added new file to initialize s_tir.backend namespace.
  • python/tvm/s_tir/backend/adreno/init.py
    • Added new file to initialize s_tir.backend.adreno namespace.
  • python/tvm/s_tir/backend/adreno/pipeline.py
    • Added new file defining the default TIR pipeline for Adreno under s_tir.
  • python/tvm/s_tir/block_dependence_info.py
    • Renamed from python/tvm/tir/block_dependence_info.py.
  • python/tvm/s_tir/block_scope.py
    • Renamed from python/tvm/tir/block_scope.py.
  • python/tvm/s_tir/pipeline.py
    • Renamed from python/tvm/tir/backend/adreno/pipeline.py.
    • Modified to define default_s_tir_pipeline and update tir_pipeline.PIPELINE_MAP.
  • python/tvm/s_tir/schedule/init.py
    • Renamed from python/tvm/tir/schedule/__init__.py.
  • python/tvm/s_tir/schedule/_ffi_api.py
    • Renamed from python/tvm/tir/schedule/_ffi_api.py.
    • Updated docstring to reflect tvm.s_tir.schedule.
  • python/tvm/s_tir/schedule/_type_checker.py
    • Renamed from python/tvm/tir/schedule/_type_checker.py.
  • python/tvm/s_tir/schedule/analysis.py
    • Renamed from python/tvm/tir/schedule/analysis.py.
    • Updated imports to use tvm.tir and tvm.s_tir.schedule.
  • python/tvm/s_tir/schedule/instruction.py
    • Renamed from python/tvm/tir/schedule/instruction.py.
  • python/tvm/s_tir/schedule/schedule.py
    • Renamed from python/tvm/tir/schedule/schedule.py.
    • Updated imports and example code to use tvm.s_tir.Schedule.
  • python/tvm/s_tir/schedule/state.py
    • Renamed from python/tvm/tir/schedule/state.py.
  • python/tvm/s_tir/schedule/testing.py
    • Renamed from python/tvm/tir/schedule/testing.py.
    • Updated imports and type hints to use tvm.s_tir.schedule.
  • python/tvm/s_tir/schedule/trace.py
    • Renamed from python/tvm/tir/schedule/trace.py.
    • Updated imports and Python representation to use tvm.s_tir.Schedule.
  • python/tvm/s_tir/schedule/transform.py
    • Renamed from python/tvm/tir/schedule/transform.py.
    • Updated imports to use tvm.s_tir.schedule.
  • python/tvm/s_tir/tensor_intrin/init.py
    • Renamed from python/tvm/tir/tensor_intrin/__init__.py.
  • python/tvm/s_tir/tensor_intrin/arm_cpu.py
    • Renamed from python/tvm/tir/tensor_intrin/arm_cpu.py.
  • python/tvm/s_tir/tensor_intrin/cuda.py
    • Renamed from python/tvm/tir/tensor_intrin/cuda.py.
  • python/tvm/s_tir/tensor_intrin/dot_product_common.py
    • Renamed from python/tvm/tir/tensor_intrin/dot_product_common.py.
  • python/tvm/s_tir/tensor_intrin/hexagon.py
    • Renamed from python/tvm/tir/tensor_intrin/hexagon.py.
  • python/tvm/s_tir/tensor_intrin/metal.py
    • Renamed from python/tvm/tir/tensor_intrin/metal.py.
  • python/tvm/s_tir/tensor_intrin/riscv_cpu.py
    • Renamed from python/tvm/tir/tensor_intrin/riscv_cpu.py.
  • python/tvm/s_tir/tensor_intrin/rocm.py
    • Renamed from python/tvm/tir/tensor_intrin/rocm.py.
  • python/tvm/s_tir/tensor_intrin/x86.py
    • Renamed from python/tvm/tir/tensor_intrin/x86.py.
  • python/tvm/testing/tir.py
    • Updated instantiations from tvm.tir.Schedule to tvm.s_tir.Schedule.
  • python/tvm/tir/init.py
    • Removed imports related to schedule, block_dependence_info, and backend as they are now part of s_tir.
  • python/tvm/tir/backend/init.py
    • Updated docstring.
    • Removed pipeline import.
  • python/tvm/tir/backend/adreno/init.py
    • Updated docstring.
    • Removed pipeline import.
  • python/tvm/tir/pipeline.py
    • Removed default_tir_pipeline function.
    • Updated get_tir_pipeline to default to the s_tir pipeline.
  • src/tir/schedule/trace.cc
    • Updated Python import statement for Schedule from tvm.tir to tvm.s_tir.
  • tests/python/all-platform-minimal-test/test_minimal_target_codegen_llvm.py
    • Updated tvm.tir.Schedule to tvm.s_tir.Schedule.
  • tests/python/codegen/test_gpu_codegen_allreduce.py
    • Updated tvm.tir.Schedule to tvm.s_tir.Schedule.
  • tests/python/codegen/test_target_codegen_arm.py
    • Updated tvm.tir.Schedule to tvm.s_tir.Schedule.
  • tests/python/codegen/test_target_codegen_bool.py
    • Updated tvm.tir.Schedule to tvm.s_tir.Schedule.
  • tests/python/codegen/test_target_codegen_cross_llvm.py
    • Updated tvm.tir.Schedule to tvm.s_tir.Schedule.
  • tests/python/codegen/test_target_codegen_cuda.py
    • Updated tvm.tir.Schedule to tvm.s_tir.Schedule.
  • tests/python/codegen/test_target_codegen_cuda_fp4.py
    • Updated tvm.tir.Schedule to tvm.s_tir.Schedule.
  • tests/python/codegen/test_target_codegen_cuda_fp8.py
    • Updated tvm.tir.Schedule to tvm.s_tir.Schedule.
  • tests/python/codegen/test_target_codegen_device.py
    • Updated tir.Schedule to tvm.s_tir.Schedule.
  • tests/python/codegen/test_target_codegen_gpu_common.py
    • Updated tvm.tir.Schedule to tvm.s_tir.Schedule.
  • tests/python/codegen/test_target_codegen_llvm.py
    • Updated tir.Schedule to tvm.s_tir.Schedule.
  • tests/python/codegen/test_target_codegen_metal.py
    • Updated tvm.tir.Schedule to tvm.s_tir.Schedule.
  • tests/python/codegen/test_target_codegen_opencl.py
    • Updated tvm.tir.Schedule to tvm.s_tir.Schedule.
  • tests/python/codegen/test_target_codegen_rocm.py
    • Updated tir.Schedule to tvm.s_tir.Schedule.
  • tests/python/codegen/test_target_codegen_vulkan.py
    • Updated tvm.tir imports to tvm.s_tir.
    • Updated tir.Schedule to tvm.s_tir.Schedule.
  • tests/python/contrib/test_hexagon/test_2d_physical_buffers.py
    • Updated tvm.tir.Schedule to tvm.s_tir.Schedule.
  • tests/python/contrib/test_hexagon/test_async_dma_pipeline.py
    • Updated tvm.tir.Schedule to tvm.s_tir.Schedule.
  • tests/python/contrib/test_hexagon/test_benchmark_maxpool2d.py
    • Updated tir.Schedule to tvm.s_tir.Schedule.
  • tests/python/contrib/test_hexagon/test_meta_schedule.py
    • Updated tvm.tir.tensor_intrin.hexagon import to tvm.s_tir.tensor_intrin.hexagon.
    • Updated tvm.tir.Schedule to tvm.s_tir.Schedule.
  • tests/python/contrib/test_hexagon/test_parallel_hvx.py
    • Updated tvm.tir.Schedule to tvm.s_tir.Schedule.
  • tests/python/contrib/test_hexagon/test_parallel_hvx_load_vtcm.py
    • Updated tvm.tir.Schedule to tvm.s_tir.Schedule.
  • tests/python/contrib/test_hexagon/test_parallel_scalar.py
    • Updated tvm.tir.Schedule to tvm.s_tir.Schedule.
  • tests/python/contrib/test_hexagon/test_sigmoid.py
    • Updated tir.Schedule to tvm.s_tir.Schedule.
  • tests/python/contrib/test_hexagon/test_software_pipeline_async.py
    • Updated tir.Schedule to tvm.s_tir.Schedule.
  • tests/python/contrib/test_hexagon/test_vtcm.py
    • Updated tir.Schedule to tvm.s_tir.Schedule.
  • tests/python/contrib/test_hexagon/test_vtcm_bandwidth.py
    • Updated tvm.tir.tensor_intrin.hexagon import to tvm.s_tir.tensor_intrin.hexagon.
    • Updated tvm.tir.Schedule to tvm.s_tir.Schedule.
  • tests/python/dlight/test_primitives.py
    • Updated tvm.tir.Schedule to tvm.s_tir.Schedule.
  • tests/python/meta_schedule/test_meta_schedule_cost_model.py
    • Updated tvm.tir.schedule.schedule import to tvm.s_tir.schedule.schedule.
  • tests/python/meta_schedule/test_meta_schedule_database.py
    • Updated tvm.tir import to tvm.s_tir.
    • Updated tir.Schedule to tvm.s_tir.Schedule.
  • tests/python/meta_schedule/test_meta_schedule_feature_extractor_per_store_feature.py
    • Imported s_tir module.
    • Updated type hints and instantiations from tir.Schedule to s_tir.Schedule.
  • tests/python/meta_schedule/test_meta_schedule_measure_callback.py
    • Updated tvm.tir.schedule import to tvm.s_tir.schedule.
  • tests/python/meta_schedule/test_meta_schedule_mma_tensorize.py
    • Updated tvm.tir.schedule import to tvm.s_tir.schedule.
    • Updated tvm.tir.tensor_intrin import to tvm.s_tir.tensor_intrin.
  • tests/python/meta_schedule/test_meta_schedule_mutator_mutate_compute_location.py
    • Updated tvm.tir import to tvm.s_tir.
  • tests/python/meta_schedule/test_meta_schedule_mutator_mutate_parallel.py
    • Updated tvm.tir import to tvm.s_tir.
  • tests/python/meta_schedule/test_meta_schedule_mutator_mutate_thread_binding.py
    • Updated tvm.tir import to tvm.s_tir.
  • tests/python/meta_schedule/test_meta_schedule_mutator_mutate_tile_size.py
    • Updated tvm.tir import to tvm.s_tir.
  • tests/python/meta_schedule/test_meta_schedule_mutator_mutate_unroll.py
    • Updated tvm.tir import to tvm.s_tir.
  • tests/python/meta_schedule/test_meta_schedule_post_order_apply.py
    • Updated tvm.tir.schedule import to tvm.s_tir.schedule.
    • Updated tvm.tir.schedule.schedule.ScheduleError to tvm.s_tir.schedule.schedule.ScheduleError.
  • tests/python/meta_schedule/test_meta_schedule_postproc_disallow_async_strided_mem_copy.py
    • Updated tir.Schedule to tvm.s_tir.Schedule.
  • tests/python/meta_schedule/test_meta_schedule_postproc_disallow_dynamic_loop.py
    • Updated tir.Schedule to tvm.s_tir.Schedule.
  • tests/python/meta_schedule/test_meta_schedule_postproc_rewrite_cooperative_fetch.py
    • Updated tir.Schedule to tvm.s_tir.Schedule.
  • tests/python/meta_schedule/test_meta_schedule_postproc_rewrite_layout.py
    • Updated tvm.tir.schedule.testing import to tvm.s_tir.schedule.testing.
    • Updated tvm.tir.Schedule to tvm.s_tir.Schedule.
  • tests/python/meta_schedule/test_meta_schedule_postproc_rewrite_parallel_vectorize_unroll.py
    • Updated tvm.tir.schedule imports to tvm.s_tir.schedule.
  • tests/python/meta_schedule/test_meta_schedule_postproc_rewrite_reduction_block.py
    • Updated tir.Schedule to tvm.s_tir.Schedule.
  • tests/python/meta_schedule/test_meta_schedule_postproc_rewrite_tensorize.py
    • Updated tvm.tir.tensor_intrin imports to tvm.s_tir.tensor_intrin.
    • Updated tvm.tir.Schedule to tvm.s_tir.Schedule.
  • tests/python/meta_schedule/test_meta_schedule_postproc_rewrite_unbound_block.py
    • Updated tir.Schedule to tvm.s_tir.Schedule.
  • tests/python/meta_schedule/test_meta_schedule_postproc_verify_gpu_code.py
    • Updated tir.Schedule to tvm.s_tir.Schedule.
  • tests/python/meta_schedule/test_meta_schedule_postproc_verify_vtcm_limit.py
    • Updated tir.Schedule to tvm.s_tir.Schedule.
  • tests/python/meta_schedule/test_meta_schedule_schedule_rule_apply_custom_rule.py
    • Updated type hints from tvm.tir.Schedule to tvm.s_tir.Schedule.
  • tests/python/meta_schedule/test_meta_schedule_schedule_rule_auto_inline.py
    • Updated tvm.tir import to tvm.s_tir.
    • Updated tvm.tir.ScheduleError to tvm.s_tir.ScheduleError.
  • tests/python/meta_schedule/test_meta_schedule_schedule_rule_mlt_intrin.py
    • Updated tvm.tir.tensor_intrin imports to tvm.s_tir.tensor_intrin.
  • tests/python/meta_schedule/test_meta_schedule_schedule_rule_mlt_tc.py
    • Updated tvm.tir.tensor_intrin.cuda import to tvm.s_tir.tensor_intrin.cuda.
  • tests/python/meta_schedule/test_meta_schedule_search_strategy.py
    • Updated tvm.tir.schedule import to tvm.s_tir.schedule.
  • tests/python/meta_schedule/test_meta_schedule_space_cuda_async_multiple_initialization.py
    • Updated tvm.tir import to tvm.s_tir.
  • tests/python/meta_schedule/test_meta_schedule_space_generator.py
    • Updated tvm.tir.schedule import to tvm.s_tir.schedule.
  • tests/python/meta_schedule/test_meta_schedule_task_scheduler.py
    • Updated tvm.tir import to tvm.s_tir.
  • tests/python/meta_schedule/test_meta_schedule_trace_apply.py
    • Updated tvm.tir imports to tvm.s_tir.
  • tests/python/meta_schedule/test_meta_schedule_tune_tir.py
    • Updated tvm.tir.schedule import to tvm.s_tir.schedule.
  • tests/python/relax/test_analysis_suggest_layout_transforms.py
    • Updated tir.Schedule to tvm.s_tir.Schedule.
  • tests/python/relax/test_frontend_dynamo.py
    • Updated tir.Schedule to tvm.s_tir.Schedule.
  • tests/python/relax/test_transform_few_shot_tuning.py
    • Updated tvm.tir.tensor_intrin imports to tvm.s_tir.tensor_intrin.
  • tests/python/relax/test_transform_meta_schedule_apply_database.py
    • Updated tir.Schedule to tvm.s_tir.Schedule.
  • tests/python/relax/test_vm_build.py
    • Updated tvm.tir.Schedule to tvm.s_tir.Schedule.
  • tests/python/relax/texture/test_texture_nd.py
    • Updated type hints and instantiations from tir.Schedule to tvm.s_tir.Schedule.
  • tests/python/runtime/test_evaluator_with_preproc.py
    • Updated tvm.tir.Schedule to tvm.s_tir.Schedule.
  • tests/python/runtime/test_runtime_module_load.py
    • Updated tvm.tir.Schedule to tvm.s_tir.Schedule.
  • tests/python/runtime/test_runtime_rpc.py
    • Updated tvm.tir.Schedule to tvm.s_tir.Schedule.
  • tests/python/s_tir/base/test_tir_block_dependence_info.py
    • Renamed from tests/python/tir-base/test_tir_block_dependence_info.py.
    • Updated imports to use tvm.s_tir.
  • tests/python/s_tir/schedule/test_tir_schedule_analysis.py
    • Renamed from tests/python/tir-schedule/test_tir_schedule_analysis.py.
    • Updated imports to use tvm.s_tir and tvm.s_tir.schedule.analysis.
    • Updated instantiations from tvm.tir.Schedule to tvm.s_tir.Schedule.
  • tests/python/s_tir/schedule/test_tir_schedule_annotate_buffer_access.py
    • Renamed from tests/python/tir-schedule/test_tir_schedule_annotate_buffer_access.py.
    • Updated imports to use tvm.s_tir.schedule.testing.
    • Updated instantiations from tir.Schedule to tvm.s_tir.Schedule.
  • tests/python/s_tir/schedule/test_tir_schedule_block_scope.py
    • Renamed from tests/python/tir-schedule/test_tir_schedule_block_scope.py.
    • Updated imports to use tvm.s_tir and tvm.s_tir.schedule.
    • Updated instantiations from tir.ScheduleState to tvm.s_tir.ScheduleState.
  • tests/python/s_tir/schedule/test_tir_schedule_blockize.py
    • Renamed from tests/python/tir-schedule/test_tir_schedule_blockize.py.
    • Updated imports to use tvm.s_tir.schedule.testing.
    • Updated instantiations from tir.Schedule to tvm.s_tir.Schedule.
  • tests/python/s_tir/schedule/test_tir_schedule_cache_index.py
    • Renamed from tests/python/tir-schedule/test_tir_schedule_cache_index.py.
    • Updated imports to use tvm.s_tir.schedule.testing.
    • Updated instantiations from tvm.tir.Schedule to tvm.s_tir.Schedule.
  • tests/python/s_tir/schedule/test_tir_schedule_cache_read_write.py
    • Renamed from tests/python/tir-schedule/test_tir_schedule_cache_read_write.py.
    • Updated imports to use tvm.s_tir.schedule.testing.
    • Updated instantiations from tir.Schedule to tvm.s_tir.Schedule and tvm.tir.ScheduleError to tvm.s_tir.ScheduleError.
  • tests/python/s_tir/schedule/test_tir_schedule_compute_at.py
    • Renamed from tests/python/tir-schedule/test_tir_schedule_compute_at.py.
    • Updated imports to use tvm.s_tir.schedule.testing.
    • Updated instantiations from tir.Schedule to tvm.s_tir.Schedule and tvm.tir.ScheduleError to tvm.s_tir.ScheduleError.
  • tests/python/s_tir/schedule/test_tir_schedule_compute_inline.py
    • Renamed from tests/python/tir-schedule/test_tir_schedule_compute_inline.py.
    • Updated imports to use tvm.s_tir.tensor_intrin and tvm.s_tir.schedule.testing.
    • Updated instantiations from tir.Schedule to tvm.s_tir.Schedule and tvm.tir.ScheduleError to tvm.s_tir.ScheduleError.
  • tests/python/s_tir/schedule/test_tir_schedule_decompose_padding.py
    • Renamed from tests/python/tir-schedule/test_tir_schedule_decompose_padding.py.
    • Updated imports to use tvm.s_tir.schedule.testing.
    • Updated instantiations from tir.Schedule to tvm.s_tir.Schedule.
  • tests/python/s_tir/schedule/test_tir_schedule_error.py
    • Renamed from tests/python/tir-schedule/test_tir_schedule_error.py.
    • Updated imports to use tvm.s_tir.
    • Updated instantiations from tir.Schedule to tvm.s_tir.Schedule and tir.ScheduleError to tvm.s_tir.ScheduleError.
  • tests/python/s_tir/schedule/test_tir_schedule_for_kind.py
    • Renamed from tests/python/tir-schedule/test_tir_schedule_for_kind.py.
    • Updated imports to use tvm.s_tir.schedule.testing.
    • Updated instantiations from tir.Schedule to tvm.s_tir.Schedule and tvm.tir.ScheduleError to tvm.s_tir.ScheduleError.
  • tests/python/s_tir/schedule/test_tir_schedule_fuse_reduction_epilogue.py
    • Renamed from tests/python/tir-schedule/test_tir_schedule_fuse_reduction_epilogue.py.
    • Updated imports to use tvm.s_tir.schedule.testing.
    • Updated instantiations from tir.Schedule to tvm.s_tir.Schedule and tvm.tir.ScheduleError to tvm.s_tir.ScheduleError.
  • tests/python/s_tir/schedule/test_tir_schedule_fuse_reduction_epilogue_clipping.py
    • Renamed from tests/python/tir-schedule/test_tir_schedule_fuse_reduction_epilogue_clipping.py.
    • Updated imports to use tvm.s_tir.schedule.testing.
    • Updated instantiations from tir.Schedule to tvm.s_tir.Schedule.
  • tests/python/s_tir/schedule/test_tir_schedule_fuse_reduction_epilogue_relu.py
    • Renamed from tests/python/tir-schedule/test_tir_schedule_fuse_reduction_epilogue_relu.py.
    • Updated imports to use tvm.s_tir.schedule.testing.
    • Updated instantiations from tir.Schedule to tvm.s_tir.Schedule.
  • tests/python/s_tir/schedule/test_tir_schedule_instruction.py
    • Renamed from tests/python/tir-schedule/test_tir_schedule_instruction.py.
    • Updated imports to use tvm.s_tir.schedule.
  • tests/python/s_tir/schedule/test_tir_schedule_merge.py
    • Renamed from tests/python/tir-schedule/test_tir_schedule_merge.py.
    • Updated imports to use tvm.s_tir.schedule.testing.
    • Updated instantiations from tir.Schedule to tvm.s_tir.Schedule and tvm.tir.ScheduleError to tvm.s_tir.ScheduleError.
  • tests/python/s_tir/schedule/test_tir_schedule_pad_einsum.py
    • Renamed from tests/python/tir-schedule/test_tir_schedule_pad_einsum.py.
    • Updated imports to use tvm.s_tir.schedule.testing.
    • Updated instantiations from tir.Schedule to tvm.s_tir.Schedule.
  • tests/python/s_tir/schedule/test_tir_schedule_partition.py
    • Renamed from tests/python/tir-schedule/test_tir_schedule_partition.py.
    • Updated imports to use tvm.s_tir.schedule.testing.
    • Updated instantiations from tir.Schedule to tvm.s_tir.Schedule and tvm.tir.ScheduleError to tvm.s_tir.ScheduleError.
  • tests/python/s_tir/schedule/test_tir_schedule_read_write_at.py
    • Renamed from tests/python/tir-schedule/test_tir_schedule_read_write_at.py.
    • Updated imports to use tvm.s_tir.schedule.testing.
    • Updated instantiations from tir.Schedule to tvm.s_tir.Schedule.
  • tests/python/s_tir/schedule/test_tir_schedule_reduction.py
    • Renamed from tests/python/tir-schedule/test_tir_schedule_reduction.py.
    • Updated imports to use tvm.s_tir.schedule.testing.
    • Updated instantiations from tir.Schedule to tvm.s_tir.Schedule.
  • tests/python/s_tir/schedule/test_tir_schedule_reindex.py
    • Renamed from tests/python/tir-schedule/test_tir_schedule_reindex.py.
    • Updated imports to use tvm.s_tir.schedule.
    • Updated instantiations from tir.Schedule to tvm.s_tir.Schedule.
  • tests/python/s_tir/schedule/test_tir_schedule_reorder.py
    • Renamed from tests/python/tir-schedule/test_tir_schedule_reorder.py.
    • Updated imports to use tvm.s_tir.schedule.testing.
    • Updated instantiations from tir.Schedule to tvm.s_tir.Schedule and tvm.tir.ScheduleError to tvm.s_tir.ScheduleError.
  • tests/python/s_tir/schedule/test_tir_schedule_reorder_block_iter_var.py
    • Renamed from tests/python/tir-schedule/test_tir_schedule_reorder_block_iter_var.py.
    • Updated imports to use tvm.s_tir.schedule.testing.
    • Updated instantiations from tir.Schedule to tvm.s_tir.Schedule and tvm.tir.ScheduleError to tvm.s_tir.ScheduleError.
  • tests/python/s_tir/schedule/test_tir_schedule_rfactor.py
    • Renamed from tests/python/tir-schedule/test_tir_schedule_rfactor.py.
    • Updated imports to use tvm.s_tir.schedule.testing.
    • Updated instantiations from tir.Schedule to tvm.s_tir.Schedule and tvm.tir.ScheduleError to tvm.s_tir.ScheduleError.
  • tests/python/s_tir/schedule/test_tir_schedule_rolling_buffer.py
    • Renamed from tests/python/tir-schedule/test_tir_schedule_rolling_buffer.py.
    • Updated imports to use tvm.s_tir and tvm.s_tir.schedule.testing.
    • Updated instantiations from tir.Schedule to tvm.s_tir.Schedule and tvm.tir.ScheduleError to tvm.s_tir.ScheduleError.
  • tests/python/s_tir/schedule/test_tir_schedule_sampling.py
    • Renamed from tests/python/tir-schedule/test_tir_schedule_sampling.py.
    • Updated imports to use tvm.s_tir.schedule.testing.
    • Updated instantiations from tir.Schedule to tvm.s_tir.Schedule and tvm.tir.schedule.ScheduleError to tvm.s_tir.schedule.ScheduleError.
  • tests/python/s_tir/schedule/test_tir_schedule_set_axis_separator.py
    • Renamed from tests/python/tir-schedule/test_tir_schedule_set_axis_separator.py.
    • Updated imports to use tvm.s_tir.schedule.testing.
    • Updated instantiations from tir.Schedule to tvm.s_tir.Schedule.
  • tests/python/s_tir/schedule/test_tir_schedule_set_dtype.py
    • Renamed from tests/python/tir-schedule/test_tir_schedule_set_dtype.py.
    • Updated imports to use tvm.s_tir.schedule.testing.
    • Updated instantiations from tir.Schedule to tvm.s_tir.Schedule and tvm.tir.ScheduleError to tvm.s_tir.ScheduleError.
  • tests/python/s_tir/schedule/test_tir_schedule_set_scope.py
    • Renamed from tests/python/tir-schedule/test_tir_schedule_set_scope.py.
    • Updated imports to use tvm.s_tir.schedule.testing.
    • Updated instantiations from tir.Schedule to tvm.s_tir.Schedule and tvm.tir.ScheduleError to tvm.s_tir.ScheduleError.
  • tests/python/s_tir/schedule/test_tir_schedule_split_fuse.py
    • Renamed from tests/python/tir-schedule/test_tir_schedule_split_fuse.py.
    • Updated imports to use tvm.s_tir.schedule.testing.
    • Updated instantiations from tir.Schedule to tvm.s_tir.Schedule and tvm.tir.ScheduleError to tvm.s_tir.ScheduleError.
  • tests/python/s_tir/schedule/test_tir_schedule_state.py
    • Renamed from tests/python/tir-schedule/test_tir_schedule_state.py.
    • Updated instantiations from tir.ScheduleState to tvm.s_tir.ScheduleState.
  • tests/python/s_tir/schedule/test_tir_schedule_state_cached_flags.py
    • Renamed from tests/python/tir-schedule/test_tir_schedule_state_cached_flags.py.
    • Updated imports to use tvm.s_tir and tvm.s_tir.schedule.state.
    • Updated instantiations from tir.ScheduleState to tvm.s_tir.ScheduleState.
  • tests/python/s_tir/schedule/test_tir_schedule_storage_align.py
    • Renamed from tests/python/tir-schedule/test_tir_schedule_storage_align.py.
    • Updated imports to use tvm.s_tir.schedule.testing.
    • Updated instantiations from tir.Schedule to tvm.s_tir.Schedule and tir.ScheduleError to tvm.s_tir.ScheduleError.
  • tests/python/s_tir/schedule/test_tir_schedule_tensorize.py
    • Renamed from tests/python/tir-schedule/test_tir_schedule_tensorize.py.
    • Updated imports to use tvm.s_tir.schedule.testing and tvm.s_tir.tensor_intrin.
    • Updated instantiations from tir.Schedule to tvm.s_tir.Schedule.
  • tests/python/s_tir/schedule/test_tir_schedule_tensorize_ldmatrix_mma_numeric.py
    • Renamed from tests/python/tir-schedule/test_tir_schedule_tensorize_ldmatrix_mma_numeric.py.
    • Updated imports to use tvm.s_tir.tensor_intrin.cuda.
  • tests/python/s_tir/schedule/test_tir_schedule_tensorize_mfma_numeric.py
    • Renamed from tests/python/tir-schedule/test_tir_schedule_tensorize_mfma_numeric.py.
    • Updated imports to use tvm.s_tir.tensor_intrin.rocm.
  • tests/python/s_tir/schedule/test_tir_schedule_trace.py
    • Renamed from tests/python/tir-schedule/test_tir_schedule_trace.py.
    • Updated imports to use tvm.s_tir and tvm.s_tir.schedule.
  • tests/python/s_tir/schedule/test_tir_schedule_transform.py
    • Renamed from tests/python/tir-schedule/test_tir_schedule_transform.py.
    • Updated imports to use tvm.s_tir and tvm.s_tir.schedule.transform.
  • tests/python/s_tir/schedule/test_tir_schedule_transform_layout.py
    • Renamed from tests/python/tir-schedule/test_tir_schedule_transform_layout.py.
    • Updated imports to use tvm.s_tir.schedule.testing.
    • Updated instantiations from tir.Schedule to tvm.s_tir.Schedule and tvm.tir.schedule.schedule.ScheduleError to tvm.s_tir.schedule.schedule.ScheduleError.
  • tests/python/s_tir/schedule/test_tir_schedule_utilities.py
    • Renamed from tests/python/tir-schedule/test_tir_schedule_utilities.py.
    • Updated imports to use tvm.s_tir.schedule.testing.
    • Updated instantiations from tir.Schedule to tvm.s_tir.Schedule.
  • tests/python/te/test_te_create_primfunc.py
    • Imported s_tir module.
    • Updated instantiations from tir.Schedule to s_tir.Schedule.
  • tests/python/testing/test_type_annotation_checker.py
    • Updated tvm.tir.schedule._type_checker import to tvm.s_tir.schedule._type_checker.
  • tests/python/tir-analysis/test_tir_analysis_calculate_allocated_memory.py
    • Updated tir.Schedule to tvm.s_tir.Schedule.
  • tests/python/tir-analysis/test_tir_analysis_get_block_access_region.py
    • Updated tvm.tir.schedule.Schedule to tvm.s_tir.schedule.Schedule.
  • tests/python/tir-analysis/test_tir_analysis_oob.py
    • Updated tvm.tir.ScheduleError to tvm.s_tir.ScheduleError.
  • tests/python/tir-base/test_tir_intrin.py
    • Updated tir.Schedule to tvm.s_tir.Schedule.
  • tests/python/tir-base/test_tir_ptx_mma.py
    • Updated tvm.tir.Schedule to tvm.s_tir.Schedule.
  • tests/python/tir-base/test_tir_ptx_mma_sp.py
    • Updated tvm.tir.Schedule to tvm.s_tir.Schedule.
  • tests/python/tir-base/test_tir_specialize.py
    • Updated tvm.tir.schedule.testing import to tvm.s_tir.schedule.testing.
  • tests/python/tir-base/test_tir_texture_scope.py
    • Updated tir.Schedule to tvm.s_tir.Schedule.
  • tests/python/tir-base/test_tir_unsafe_hide_buffer_access.py
    • Updated tvm.tir.schedule.testing import to tvm.s_tir.schedule.testing.
    • Updated tir.Schedule to tvm.s_tir.Schedule.
  • tests/python/tir-transform/test_tir_transform_inject_software_pipeline.py
    • Updated tvm.tir.tensor_intrin.cuda import to tvm.s_tir.tensor_intrin.cuda.
    • Updated tvm.tir.Schedule to tvm.s_tir.Schedule.
  • tests/python/tir-transform/test_tir_transform_lower_tvm_builtin.py
    • Updated tvm.tir.schedule.testing import to tvm.s_tir.schedule.testing.
  • tests/python/tir-transform/test_tir_transform_plan_update_buffer_allocation_location.py
    • Updated tvm.tir.tensor_intrin.hexagon import to tvm.s_tir.tensor_intrin.hexagon.
  • tests/python/tvmscript/test_tvmscript_error_report.py
    • Updated tir.Schedule to tvm.s_tir.Schedule.
    • Updated tvm.tir.ScheduleError to tvm.s_tir.ScheduleError.
  • tests/python/tvmscript/test_tvmscript_syntax_sugar.py
    • Updated tvm.tir.schedule.testing import to tvm.s_tir.schedule.testing.
  • tests/scripts/task_python_unittest.sh
    • Updated test file paths to reflect the s_tir renaming.
  • web/tests/python/relax_rpc_test.py
    • Updated tvm.tir.Schedule to tvm.s_tir.Schedule.
  • web/tests/python/webgpu_rpc_test.py
    • Updated tvm.tir.Schedule to tvm.s_tir.Schedule.
Activity
  • The pull request was created by tqchen.
  • The author's description indicates that this PR is an initial step to introduce the s_tir module, focusing on Python-side changes first, with C++ changes to follow in subsequent PRs.
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 gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request introduces a major refactoring by initializing the s_tir module for schedulable TensorIR, primarily by replacing tvm.tir.Schedule with tvm.s_tir.Schedule across the codebase. The changes are extensive, touching tests, documentation, and various parts of the dlight and meta_schedule components. Overall, the refactoring appears to be thorough and consistent. I've identified one issue in a test file that would cause a NameError due to a missing import qualification. Addressing this will ensure the tests pass and the refactoring is complete.

@tqchen

tqchen commented Feb 4, 2026

Copy link
Copy Markdown
Member Author

cc @Hzfengsy

@tqchen tqchen force-pushed the s-tir branch 2 times, most recently from c0a58e6 to 43a2c26 Compare February 5, 2026 01:22
This PR initalizes the s_tir for scheduable TensorIR.
The change mainly starts from python side, the we will gradually
move towards the c++ side in followup PRs. The python main change:

tir.Schedule => s_tir.Schedule
@tqchen tqchen merged commit d76c729 into apache:main Feb 5, 2026
12 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants