Skip to content

Conversation

@tqchen
Copy link
Member

@tqchen tqchen commented Feb 9, 2026

This PR moves remaining related data structures to s_tir.

  • Moves sblock_dependency_info and sblock_scope.
  • Moves related analyssis.
  • Hides the data_type_rewriter to private functions.

@gemini-code-assist
Copy link
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 undertakes a significant refactoring effort to enhance the organization and clarity of the Schedulable TensorIR (S-TIR) components within the TVM project. By relocating core S-TIR data structures and analysis functions to a dedicated s_tir namespace and directory, the change aims to create a clearer separation of concerns between general TensorIR and its schedulable counterpart. Additionally, an internal utility, data_type_rewriter, has been moved to an internal source directory, further streamlining the public API.

Highlights

  • S-TIR Data Structure Relocation: Key S-TIR related data structures, including sblock_dependency_info, sblock_scope, and data_layout, have been moved from the general tvm/tir namespace to the dedicated tvm/s_tir namespace and directory structure. This improves modularity and organization.
  • Analysis Function Migration: Analysis utilities specific to S-TIR, such as GetSBlockAccessRegion, GetSBlockReadWriteRegion, DetectBufferAccessLCA, and FindAnchorBlock, have been migrated from tvm/tir/analysis.h to the new tvm/s_tir/analysis.h header and their corresponding Python bindings.
  • Data Type Rewriter Internalization: The data_type_rewriter header has been moved from a public include/tvm/tir/data_type_rewriter.h to an internal src/tir/ir/data_type_rewriter.h, effectively hiding it from the public API and treating it as an internal implementation detail.
  • Extensive Reference Updates: All internal and external references across C++ and Python files have been updated to reflect the new locations and naming conventions of the moved data structures and analysis functions, ensuring consistency and correctness.

🧠 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 header file for S-TIR specific analysis utilities.
    • Includes declarations for GetSBlockAccessRegion, GetSBlockReadWriteRegion, DetectBufferAccessLCA, and FindAnchorBlock.
  • include/tvm/s_tir/schedule/state.h
    • Updated include path from tvm/tir/block_scope.h to tvm/s_tir/sblock_scope.h.
  • include/tvm/tir/analysis.h
    • Removed S-TIR specific analysis function declarations (GetSBlockAccessRegion, GetSBlockReadWriteRegion, DetectBufferAccessLCA, FindAnchorBlock).
    • Added include for the new tvm/s_tir/analysis.h.
  • include/tvm/tir/block_dependence_info.h -> include/tvm/s_tir/sblock_dependence_info.h
    • File renamed and moved to s_tir directory.
    • Updated include guards and class names to reflect SBlockDependenceInfo.
    • Updated internal include from tvm/tir/block_scope.h to tvm/s_tir/sblock_scope.h.
    • FFI object registration names changed from tir.SBlockDependenceInfo to s_tir.SBlockDependenceInfo.
  • include/tvm/tir/block_scope.h -> include/tvm/s_tir/sblock_scope.h
    • File renamed and moved to s_tir directory.
    • Updated include guards from TVM_TIR_BLOCK_SCOPE_H_ to TVM_S_TIR_SBLOCK_SCOPE_H_.
    • FFI object registration names changed from tir.StmtSRef, tir.Dependency, tir.SBlockScope to s_tir.StmtSRef, s_tir.Dependency, s_tir.SBlockScope respectively.
  • include/tvm/tir/data_layout.h -> include/tvm/s_tir/data_layout.h
    • File renamed and moved to s_tir directory.
    • Updated include guards from TVM_TIR_DATA_LAYOUT_H_ to TVM_S_TIR_DATA_LAYOUT_H_.
    • FFI object registration names changed from tir.Layout to s_tir.Layout and tir.BijectiveLayout to s_tir.BijectiveLayout.
  • include/tvm/tir/data_type_rewriter.h -> src/tir/ir/data_type_rewriter.h
    • Header file renamed and moved to src/tir/ir directory.
    • Updated include guard from TVM_TIR_DATA_TYPE_REWRITER_H_ to TVM_TIR_IR_DATA_TYPE_REWRITER_H_.
  • include/tvm/tir/utils.h -> include/tvm/s_tir/utils.h
    • File renamed and moved to s_tir directory.
    • Updated include guards from TVM_TIR_UTILS_H_ to TVM_S_TIR_UTILS_H_.
    • Updated internal include from tvm/tir/block_scope.h to tvm/s_tir/sblock_scope.h.
  • include/tvm/topi/einsum.h
    • Removed include for tvm/tir/data_layout.h.
  • include/tvm/topi/transform.h
    • Updated include path from tvm/tir/data_layout.h to tvm/s_tir/data_layout.h.
  • python/tvm/s_tir/init.py
    • Updated import for SBlockDependenceInfo to reflect its new location.
    • Added imports for Layout, BijectiveLayout, bijective_layout, layout from s_tir.data_layout.
    • Added import for the new analysis module under s_tir.
  • python/tvm/s_tir/_ffi_api.py
    • Changed FFI API initialization from tvm_ffi.init_ffi_api("tir", __name__) to tvm_ffi.init_ffi_api("s_tir", __name__).
  • python/tvm/s_tir/analysis/init.py
    • Added new Python file for S-TIR analysis functions.
    • Provides Python bindings for get_sblock_access_region, get_sblock_read_write_region, detect_buffer_access_lca, and find_anchor_sblock.
  • python/tvm/s_tir/analysis/_ffi_api.py
    • Added new Python FFI API initialization for s_tir.analysis.
  • python/tvm/s_tir/block_dependence_info.py
    • Updated docstring to refer to SBlockDependenceInfoNode.
    • Updated import path for SBlockScope and StmtSRef from .block_scope to .sblock_scope.
  • python/tvm/s_tir/block_scope.py -> python/tvm/s_tir/sblock_scope.py
    • File renamed and moved to s_tir directory.
    • Updated FFI object registration from tir.StmtSRef, tir.Dependency, tir.SBlockScope to s_tir.StmtSRef, s_tir.Dependency, s_tir.SBlockScope respectively.
  • python/tvm/s_tir/sblock_dependence_info.py
    • Added new Python file for SBlockDependenceInfo class.
    • Provides methods get_sref and get_sblock_scope.
  • python/tvm/s_tir/schedule/init.py
    • Updated import path for SBlockScope, Dependency, DepKind, StmtSRef from ..block_scope to ..sblock_scope.
  • python/tvm/s_tir/schedule/state.py
    • Updated import path for SBlockScope and StmtSRef from ..block_scope to ..sblock_scope.
  • python/tvm/tir/init.py
    • Removed imports for Layout, BijectiveLayout, bijective_layout, layout as they are now in tvm.s_tir.
  • python/tvm/tir/analysis/analysis.py
    • Removed Python bindings for S-TIR specific analysis functions (get_sblock_access_region, get_sblock_read_write_region, detect_buffer_access_lca, find_anchor_sblock).
  • python/tvm/tir/data_layout.py -> python/tvm/s_tir/data_layout.py
    • File renamed and moved to s_tir directory.
    • Updated FFI object registration from tir.Layout to s_tir.Layout and tir.BijectiveLayout to s_tir.BijectiveLayout.
  • python/tvm/topi/utils.py
    • Updated imports for bijective_layout and layout to reference tvm.s_tir instead of tvm.tir.
  • src/contrib/msc/core/ir/graph.h
    • Updated include path from tvm/tir/data_layout.h to tvm/s_tir/data_layout.h.
  • src/contrib/msc/core/ir/graph_builder.h
    • Updated include path from tvm/tir/data_layout.h to tvm/s_tir/data_layout.h.
  • src/contrib/msc/core/ir/plugin.h
    • Updated include path from tvm/tir/data_layout.h to tvm/s_tir/data_layout.h.
  • src/relax/op/op_common.h
    • Updated include path from tvm/tir/data_layout.h to tvm/s_tir/data_layout.h.
  • src/relax/transform/infer_amp_utils.h
    • Removed include for tvm/tir/data_layout.h.
  • src/relax/transform/infer_layout_utils.h
    • Updated include path from tvm/tir/data_layout.h to tvm/s_tir/data_layout.h.
  • src/s_tir/analysis/find_anchor_sblock.cc
    • Added new C++ source file implementing FindAnchorBlock for S-TIR.
    • Includes FFI registration for s_tir.analysis.find_anchor_sblock.
  • src/s_tir/schedule/primitive/blockize_tensorize.cc
    • Updated include path for data_type_rewriter.h from <tvm/tir/data_type_rewriter.h> to ../../../tir/ir/data_type_rewriter.h.
  • src/s_tir/schedule/utils.h
    • Updated include path for tvm/s_tir/utils.h and removed tvm/tir/utils.h.
  • src/te/operation/create_primfunc.cc
    • Updated include path for data_type_rewriter.h from <tvm/tir/data_type_rewriter.h> to ../../tir/ir/data_type_rewriter.h.
  • src/tir/analysis/block_access_region_detector.cc -> src/s_tir/analysis/sblock_access_region_detector.cc
    • File renamed and moved to s_tir/analysis directory.
    • Updated file header comment.
    • Updated FFI registration names from tir.analysis.GetSBlockAccessRegion to s_tir.analysis.GetSBlockAccessRegion and tir.analysis.GetSBlockReadWriteRegion to s_tir.analysis.GetSBlockReadWriteRegion.
  • src/tir/analysis/buffer_access_lca_detector.cc -> src/s_tir/analysis/sblock_buffer_access_lca_detector.cc
    • File renamed and moved to s_tir/analysis directory.
    • Updated file header comment.
    • Updated FFI registration name from tir.analysis.detect_buffer_access_lca to s_tir.analysis.detect_buffer_access_lca.
  • src/tir/analysis/stmt_finding.cc
    • Removed implementations of GetEnclosingLoop and FindAnchorBlock as they are now in s_tir/analysis.
  • src/tir/ir/block_dependence_info.cc -> src/s_tir/sblock_dependence_info.cc
    • File renamed and moved to s_tir directory.
    • Updated internal includes.
    • Updated class names from BlockDependenceInfoNode to SBlockDependenceInfoNode and BlockDependenceInfoCollector to SBlockDependenceInfoCollector.
    • Updated FFI registration names from tir.SBlockDependenceInfo to s_tir.SBlockDependenceInfo.
  • src/tir/ir/block_scope.cc
    • Updated includes from tvm/tir/block_scope.h and tvm/tir/utils.h to tvm/s_tir/sblock_scope.h and tvm/s_tir/utils.h.
    • Updated FFI registration names from tir.StmtSRefStmt, tir.StmtSRefParent, etc., to s_tir.StmtSRefStmt, s_tir.StmtSRefParent, etc.
  • src/tir/ir/data_layout.cc -> src/s_tir/data_layout.cc
    • File renamed and moved to s_tir directory.
    • Updated internal include path.
    • Updated FFI registration names from tir.Layout, tir.LayoutIndexOf, etc., to s_tir.Layout, s_tir.LayoutIndexOf, etc.
  • src/tir/ir/data_type_rewriter.cc
    • Updated include path for its own header from <tvm/tir/data_type_rewriter.h> to data_type_rewriter.h.
  • src/tir/ir/stmt_functor.cc
    • Updated include path for data_type_rewriter.h from <tvm/tir/data_type_rewriter.h> to data_type_rewriter.h.
  • src/tir/transform/force_narrow_index_to_i32.cc
    • Updated include path for data_type_rewriter.h from <tvm/tir/data_type_rewriter.h> to ../ir/data_type_rewriter.h.
  • src/tir/transform/narrow_datatype.cc
    • Updated include path for data_type_rewriter.h from <tvm/tir/data_type_rewriter.h> to ../ir/data_type_rewriter.h.
  • tests/cpp/data_type_rewriter_test.cc
    • Removed test file as data_type_rewriter is now an internal utility.
  • tests/python/s_tir/base/test_tir_block_dependence_info.py -> tests/python/s_tir/base/test_sblock_dependence_info.py
    • Test file renamed to reflect sblock_dependence_info.
    • Updated import for DepKind from tvm.s_tir.block_scope to tvm.s_tir.sblock_scope.
  • tests/python/tir-analysis/test_tir_analysis_detect_buffer_access_lca.py -> tests/python/s_tir/analysis/test_sblock_buffer_access_lca.py
    • Test file renamed and moved to s_tir/analysis directory.
    • Updated imports from tvm.tir to tvm.s_tir and tir.analysis to s_tir.analysis.
  • tests/python/tir-analysis/test_tir_analysis_get_block_access_region.py -> tests/python/s_tir/analysis/test_sblock_access_region.py
    • Test file renamed and moved to s_tir/analysis directory.
    • Updated imports from tvm.tir to tvm.s_tir and tir.analysis to s_tir.analysis.
  • tests/python/tir-base/test_tir_data_layout.py -> tests/python/s_tir/base/test_tir_data_layout.py
    • Test file renamed and moved to s_tir/base directory.
    • Updated imports from tvm.tir.layout and tvm.tir.bijective_layout to tvm.s_tir.layout and tvm.s_tir.bijective_layout.
  • tests/python/tir-base/test_tir_te_extern_primfunc.py -> tests/python/s_tir/base/test_tir_te_extern_primfunc.py
    • Test file renamed and moved to s_tir/base directory.
  • tests/scripts/task_python_unittest.sh
    • Added s_tir/analysis to the list of Python test files to be run.
Activity
  • The pull request was created by tqchen.
  • The changes involve a significant refactoring and relocation of S-TIR related components.
  • No specific reviewer comments or discussions are available in the provided context.
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.

Copy link
Contributor

@gemini-code-assist gemini-code-assist bot left a comment

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 refactors the Schedulable TensorIR (S-TIR) analysis and data layout components by moving them from the tvm.tir namespace to a dedicated tvm.s_tir namespace. Key changes include creating new header files (include/tvm/s_tir/analysis.h, include/tvm/s_tir/sblock_dependence_info.h, include/tvm/s_tir/sblock_scope.h, include/tvm/s_tir/data_layout.h, include/tvm/s_tir/utils.h) and corresponding Python modules (python/tvm/s_tir/analysis/__init__.py, python/tvm/s_tir/analysis/_ffi_api.py, python/tvm/s_tir/sblock_dependence_info.py), and updating existing files to reflect these new paths and namespaces. Specifically, functions like GetSBlockAccessRegion, GetSBlockReadWriteRegion, DetectBufferAccessLCA, and FindAnchorBlock are now part of tvm.s_tir.analysis. Additionally, Layout and BijectiveLayout objects, along with StmtSRef, SBlockScope, and SBlockDependenceInfo, have been moved to the s_tir namespace, involving renaming classes, FFI registration names, and updating include paths across various C++ and Python files. The DataTypeLegalizer test file was also removed as part of this reorganization.

@tqchen tqchen force-pushed the block-scope branch 2 times, most recently from 2969f6f to f4071a9 Compare February 9, 2026 19:57
This PR moves remaining related data structures to s_tir.
- Moves sblock_dependency_info and sblock_scope.
- Moves related analyssis.
- Hides the data_type_rewriter to private functions.
@tqchen tqchen force-pushed the block-scope branch 3 times, most recently from fc01ca2 to 5d9e9ae Compare February 9, 2026 21:16
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.

2 participants