You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by tq...@apache.org on 2026/09/17 19:31:39 UTC
(tvm) branch main updated: Move S-TIR specific node and TensorIntrin out of TIRX (#20378)
This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new d4d1ad2d10 Move S-TIR specific node and TensorIntrin out of TIRX (#20378)
d4d1ad2d10 is described below
commit d4d1ad2d107387da227335f210593781e9240ae1
Author: Tianqi Chen <tq...@users.noreply.github.com>
AuthorDate: Thu Sep 17 15:31:32 2026 -0400
Move S-TIR specific node and TensorIntrin out of TIRX (#20378)
Move SBlock, SBlockRealize, MatchBufferRegion and TensorIntrin into
S-TIR ownership, including native definitions, reflection, Python APIs
and serialized identities. Preserve shared TensorRegion compatibility
for both typed and older object-form buffer-region graphs.
Reuse inherited traversal through explicit S-TIR analyzer, SSA,
verification and simplification classes. Generic TIRX traversal uses
structural definitions and remapping; storage alignment and schedulable
dtype normalization stay in S-TIR. The shared statement simplifier and
its configuration implementation belong to TIRX, retaining their
existing runtime identities.
Adapt scheduling, TE, Relax and script consumers while preserving
current script syntax and pipeline order. Define matched buffers before
read/write regions, and limit private-function inlining to natively
dispatchable TIRX roots. Type-preserving substitution remains structural
before lowering; dtype-changing statement substitution targets core TIRX
after block lowering.
Apply common variable, buffer and type well-formedness checks across
parsed modules, then apply TIRX-specific execution checks to functions
without the S-TIR attribute. This also checks ordinary TIRX functions in
mixed modules; the legacy parse flag remains accepted without selecting
verification. Separate script ownership remains a follow-up.
---
include/tvm/relax/analysis.h | 3 +-
include/tvm/relax/distributed/axis_group_graph.h | 15 +-
include/tvm/s_tir/analysis.h | 12 +-
include/tvm/s_tir/sblock_scope.h | 17 +-
include/tvm/s_tir/schedule/schedule.h | 1 +
include/tvm/s_tir/schedule/state.h | 1 +
include/tvm/s_tir/stmt.h | 177 ++-
include/tvm/s_tir/stmt_functor.h | 128 +++
include/tvm/s_tir/tensor_intrin.h | 90 ++
include/tvm/s_tir/transform.h | 8 +
include/tvm/s_tir/utils.h | 13 +-
include/tvm/script/ir_builder/base.h | 2 +-
include/tvm/tirx/analysis.h | 8 +-
include/tvm/tirx/function.h | 56 -
include/tvm/tirx/script/builder/frame.h | 5 +-
include/tvm/tirx/script/builder/ir.h | 2 +-
include/tvm/tirx/stmt.h | 171 +--
include/tvm/tirx/stmt_functor.h | 35 +-
.../backend/trn/tile_primitive/binary/default.py | 3 +-
.../trn/tile_primitive/compose_op/binary_chain.py | 3 +-
.../trn/tile_primitive/compose_op/binary_reduce.py | 6 +-
.../trn/tile_primitive/compose_op/unary_reduce.py | 6 +-
.../tvm/backend/trn/tile_primitive/copy/default.py | 13 +-
.../tvm/backend/trn/tile_primitive/gemm/default.py | 6 +-
.../backend/trn/tile_primitive/private_alloc.py | 6 +-
.../backend/trn/tile_primitive/reduction/utils.py | 6 +-
.../backend/trn/tile_primitive/select/default.py | 3 +-
.../tvm/backend/trn/tile_primitive/unary/utils.py | 6 +-
python/tvm/ir/json_compact.py | 23 +-
python/tvm/relax/analysis/analysis.py | 3 +-
python/tvm/s_tir/__init__.py | 3 +-
python/tvm/s_tir/_tensor_intrin.py | 80 ++
python/tvm/s_tir/analysis/__init__.py | 16 +-
python/tvm/s_tir/backend/adreno/pipeline.py | 2 +-
.../tvm/s_tir/dlight/analysis/common_analysis.py | 4 +-
python/tvm/s_tir/dlight/analysis/gemv.py | 6 +-
python/tvm/s_tir/dlight/gpu/low_batch_gemv.py | 8 +-
python/tvm/s_tir/dlight/gpu/matmul.py | 10 +-
python/tvm/s_tir/dlight/gpu/reduction.py | 2 +-
python/tvm/s_tir/dlight/gpu/rmsnorm.py | 3 +-
python/tvm/s_tir/pipeline.py | 2 +-
python/tvm/s_tir/sblock_dependence_info.py | 3 +-
python/tvm/s_tir/sblock_scope.py | 3 +-
python/tvm/s_tir/schedule/schedule.py | 5 +-
python/tvm/s_tir/schedule/state.py | 3 +-
python/tvm/s_tir/stmt.py | 179 +++
python/tvm/s_tir/tensor_intrin/arm_cpu.py | 2 +-
python/tvm/s_tir/tensor_intrin/cuda.py | 3 +-
python/tvm/s_tir/tensor_intrin/metal.py | 3 +-
python/tvm/s_tir/transform/transform.py | 10 +
python/tvm/script/parser/core/entry.py | 15 +-
python/tvm/script/parser/ir/entry.py | 2 +-
python/tvm/tirx/__init__.py | 4 +-
python/tvm/tirx/analysis/analysis.py | 5 +-
python/tvm/tirx/function.py | 68 +-
python/tvm/tirx/script/parser/entry.py | 3 +-
python/tvm/tirx/stmt.py | 152 +--
src/relax/analysis/layout_transformation.cc | 35 +-
src/relax/analysis/tir_op_pattern_kind.cc | 32 +-
src/relax/backend/task_extraction.cc | 11 +-
.../transform/lower_global_view_to_local_view.cc | 50 +-
src/relax/transform/fuse_tir.cc | 37 +-
src/relax/transform/legalize_ops.cc | 3 +-
src/relax/transform/split_call_tir_by_pattern.cc | 54 +-
.../transform/split_layout_rewrite_preproc.cc | 44 +-
src/s_tir/analysis/calculate_allocated_memory.cc | 3 +-
src/s_tir/analysis/conditional_bounds.cc | 1 +
src/s_tir/analysis/domain_touched.cc | 10 +-
src/s_tir/analysis/estimate_flops.cc | 4 +-
src/s_tir/analysis/find_anchor_sblock.cc | 43 +-
src/s_tir/analysis/identify_memcpy.cc | 6 +-
src/s_tir/analysis/is_pure_function.cc | 5 +-
src/s_tir/analysis/oob_checker.cc | 6 +-
.../analysis/sblock_access_region_detector.cc | 47 +-
.../analysis/sblock_buffer_access_lca_detector.cc | 28 +-
src/s_tir/analysis/verify_gpu_code.cc | 3 +-
src/s_tir/analysis/verify_well_formed.cc | 168 +++
src/s_tir/backend/adreno/inject_texture_alloc.cc | 11 +-
src/s_tir/backend/adreno/texture_flatten.cc | 2 +-
src/s_tir/data_layout.cc | 3 +-
src/s_tir/ir/data_type_rewriter.cc | 241 ++++
src/s_tir/ir/data_type_rewriter.h | 69 ++
.../ir/ir_mutator_with_analyzer.cc} | 31 +-
src/s_tir/ir/ir_mutator_with_analyzer.h | 60 +
.../ir/ir_visitor_with_analyzer.cc} | 30 +-
src/s_tir/ir/ir_visitor_with_analyzer.h | 49 +
src/s_tir/ir/tir_visitor_with_path.cc | 81 ++
src/s_tir/ir/tir_visitor_with_path.h | 54 +
.../feature_extractor/per_store_feature.cc | 5 +-
src/s_tir/meta_schedule/module_equality.cc | 8 +-
.../meta_schedule/postproc/rewrite_tensorize.cc | 4 +-
.../postproc/rewrite_unbound_block.cc | 1 +
.../meta_schedule/postproc/verify_gpu_code.cc | 2 +-
.../schedule_rule/cross_thread_reduction.cc | 3 +-
.../schedule_rule/multi_level_tiling.cc | 2 +-
.../multi_level_tiling_tensor_core.cc | 24 +-
.../multi_level_tiling_wide_vector.cc | 5 +-
.../multi_level_tiling_with_intrin.cc | 5 +-
.../meta_schedule/schedule_rule/schedule_rule.cc | 3 +-
src/s_tir/meta_schedule/trace_apply.cc | 3 +-
src/s_tir/meta_schedule/utils.h | 15 +-
src/s_tir/sblock_dependence_info.cc | 15 +-
src/s_tir/sblock_scope.cc | 9 +-
src/s_tir/schedule/analysis.h | 1 +
src/s_tir/schedule/analysis/analysis.cc | 4 +-
src/s_tir/schedule/analysis/reducer.cc | 1 +
src/s_tir/schedule/analysis/verify.cc | 1 +
src/s_tir/schedule/concrete_schedule.cc | 6 +-
src/s_tir/schedule/concrete_schedule.h | 1 +
src/s_tir/schedule/error.h | 1 +
src/s_tir/schedule/ir_comparator.cc | 1 +
src/s_tir/schedule/ir_comparator.h | 1 +
src/s_tir/schedule/primitive.h | 2 +
src/s_tir/schedule/primitive/annotate.cc | 1 +
src/s_tir/schedule/primitive/block_annotate.cc | 2 +-
src/s_tir/schedule/primitive/blockize_tensorize.cc | 8 +-
src/s_tir/schedule/primitive/cache_index.cc | 1 +
.../schedule/primitive/cache_index_helpers.cc | 15 +-
src/s_tir/schedule/primitive/cache_index_helpers.h | 10 +-
src/s_tir/schedule/primitive/cache_read_write.cc | 1 +
src/s_tir/schedule/primitive/compute_at.cc | 1 +
src/s_tir/schedule/primitive/decompose_padding.cc | 1 +
src/s_tir/schedule/primitive/for_kind.cc | 1 +
src/s_tir/schedule/primitive/get_block_loop.cc | 1 +
src/s_tir/schedule/primitive/hide_buffer_access.cc | 1 +
.../schedule/primitive/layout_transformation.cc | 19 +-
.../schedule/primitive/loop_transformation.cc | 1 +
src/s_tir/schedule/primitive/pad_einsum.cc | 1 +
src/s_tir/schedule/primitive/reduction.cc | 1 +
.../schedule/primitive/reorder_block_iter_var.cc | 1 +
src/s_tir/schedule/primitive/rolling_buffer.cc | 1 +
src/s_tir/schedule/schedule.cc | 1 +
src/s_tir/schedule/state.cc | 3 +-
src/s_tir/schedule/traced_schedule.cc | 2 +
src/s_tir/schedule/traced_schedule.h | 1 +
src/s_tir/schedule/transform.cc | 10 +-
src/s_tir/schedule/transform.h | 11 +-
src/s_tir/schedule/utils.h | 12 +-
src/s_tir/stmt.cc | 426 +++++++
src/s_tir/stmt_functor.cc | 201 ++++
src/s_tir/tensor_intrin.cc | 100 ++
src/s_tir/transform/annotate_irregular_loop.cc | 2 +-
src/s_tir/transform/bound_checker.cc | 2 +-
src/s_tir/transform/canonicalize_loop.cc | 2 +-
src/s_tir/transform/compact_buffer_region.cc | 4 +-
src/s_tir/transform/convert_blocks_to_opaque.cc | 3 +-
src/s_tir/transform/default_gpu_schedule.cc | 21 +-
src/s_tir/transform/hoist_expression.cc | 23 +-
src/s_tir/transform/inject_double_buffer.cc | 6 +-
src/s_tir/transform/inject_permuted_layout.cc | 5 +-
src/s_tir/transform/inject_ptx_async_copy.cc | 3 +-
src/s_tir/transform/inject_ptx_ldg32.cc | 3 +-
src/s_tir/transform/inject_software_pipeline.cc | 4 +-
src/s_tir/transform/inject_virtual_thread.cc | 20 +-
src/s_tir/transform/ir_utils.cc | 201 ++++
src/s_tir/transform/ir_utils.h | 66 ++
src/s_tir/transform/lift_thread_binding.cc | 2 +-
src/s_tir/transform/loop_partition.cc | 7 +-
src/s_tir/transform/lower_async_dma.cc | 28 +-
.../transform/lower_cross_thread_reduction.cc | 3 +-
src/s_tir/transform/lower_init_block.cc | 3 +-
src/s_tir/transform/lower_match_buffer.cc | 4 +-
src/s_tir/transform/lower_opaque_block.cc | 4 +-
src/s_tir/transform/lower_thread_allreduce.cc | 2 +-
src/s_tir/transform/lower_vtcm_alloc.cc | 2 +-
.../manifest_shared_memory_local_stage.cc | 2 +-
src/s_tir/transform/memhammer_lower_auto_copy.cc | 2 +-
src/s_tir/transform/memhammer_rewrite_rule.h | 2 +-
.../transform/memhammer_tensorcore_rewrite.cc | 1 +
.../transform/merge_shared_memory_allocations.cc | 2 +-
.../plan_update_buffer_allocation_location.cc | 4 +-
src/s_tir/transform/profile_instrumentation.cc | 2 +-
src/s_tir/transform/remove_store_undef.cc | 3 +-
.../remove_weight_layout_rewrite_block.cc | 2 +-
src/s_tir/transform/renew_defs.cc | 3 +-
src/s_tir/transform/renormalize_split_pattern.cc | 5 +-
src/s_tir/transform/rewrite_unsafe_select.cc | 2 +-
src/s_tir/transform/stmt_simplify.cc | 94 ++
src/{tirx => s_tir}/transform/stmt_simplify.h | 24 +-
src/s_tir/transform/storage_access.h | 2 +-
src/s_tir/transform/tensorcore_infer_fragment.cc | 25 +-
src/s_tir/transform/thread_storage_sync.cc | 3 +-
src/s_tir/transform/transform_mma_buffer_layout.cc | 4 +-
src/s_tir/transform/unify_thread_binding.cc | 3 +-
.../transform/using_assume_to_reduce_branches.cc | 4 +-
src/te/operation/create_primfunc.cc | 132 +--
src/tirx/analysis/verify_tirx_well_formed.cc | 37 -
src/tirx/analysis/verify_well_formed.cc | 370 +-----
...{verify_well_formed.cc => verify_well_formed.h} | 246 +---
src/tirx/ir/data_type_rewriter.cc | 212 +---
src/tirx/ir/data_type_rewriter.h | 15 +-
src/tirx/ir/function.cc | 67 --
src/tirx/ir/ir_mutator_with_analyzer.cc | 26 +-
src/tirx/ir/ir_mutator_with_analyzer.h | 11 +-
src/tirx/ir/ir_visitor_with_analyzer.cc | 14 +-
src/tirx/ir/ir_visitor_with_analyzer.h | 5 +-
src/tirx/ir/script/script_complete.cc | 32 +-
src/tirx/ir/specialize.cc | 45 +-
src/tirx/ir/stmt.cc | 379 ------
src/tirx/ir/stmt_functor.cc | 132 +--
src/tirx/ir/tir_visitor_with_path.cc | 63 +-
src/tirx/ir/tir_visitor_with_path.h | 16 +-
src/tirx/script/builder/frame.cc | 15 +-
src/tirx/script/builder/ir.cc | 10 +-
src/tirx/script/printer/block.cc | 22 +-
src/tirx/script/printer/buffer.cc | 7 +-
src/tirx/script/printer/function.cc | 17 +-
src/tirx/script/printer/stmt.cc | 3 +-
src/tirx/script/printer/utils.h | 2 +-
src/tirx/transform/bind_target.cc | 4 +-
src/tirx/transform/flatten_buffer.cc | 56 -
src/tirx/transform/force_narrow_index_to_i32.cc | 27 +-
src/tirx/transform/inline_private_functions.cc | 19 +-
src/tirx/transform/ir_utils.cc | 1207 ++++++++------------
src/tirx/transform/ir_utils.h | 147 ++-
src/tirx/transform/lower_tirx_opaque.cc | 2 +-
src/tirx/transform/lower_warp_memory.cc | 2 +-
src/tirx/transform/narrow_datatype.cc | 11 +-
src/tirx/transform/remove_no_op.cc | 2 +-
src/tirx/transform/stmt_simplify.cc | 320 +++---
src/tirx/transform/stmt_simplify.h | 63 +
src/tirx/transform/storage_rewrite.cc | 8 +-
src/tirx/transform/tvm_ffi_binder.cc | 1 +
tests/cpp/ir_functor_test.cc | 33 +-
.../test_analysis_suggest_layout_transforms.py | 4 +-
tests/python/relax/test_transform_fuse_ops.py | 10 +-
tests/python/relax/test_transform_fuse_tir.py | 4 +-
.../s_tir/analysis/test_sblock_access_region.py | 2 +-
.../s_tir/base/test_sblock_dependence_info.py | 4 +-
...e_postproc_rewrite_parallel_vectorize_unroll.py | 4 +-
...eta_schedule_schedule_rule_apply_custom_rule.py | 2 +-
.../s_tir/schedule/test_tir_schedule_analysis.py | 3 +-
.../schedule/test_tir_schedule_block_scope.py | 4 +-
.../s_tir/schedule/test_tir_schedule_state.py | 8 +-
.../test_tir_schedule_state_cached_flags.py | 4 +-
.../s_tir/schedule/test_tir_schedule_tensorize.py | 18 +-
tests/python/s_tir/test_s_tir_renew_defs.py | 2 +-
.../test_s_tir_transform_compact_buffer_region.py | 4 +-
...est_s_tir_transform_convert_blocks_to_opaque.py | 2 +-
...est_s_tir_transform_inject_software_pipeline.py | 2 +-
.../test_s_tir_transform_lower_match_buffer.py | 2 +-
...st_s_tir_transform_memhammer_lower_auto_copy.py | 12 +-
.../test_s_tir_transform_mma_buffer_layout.py | 4 +-
.../test_s_tir_transform_unify_thread_binding.py | 2 +-
.../test_tir_analysis_verify_well_formed.py | 24 +-
tests/python/tirx-base/test_tir_constructor.py | 2 +-
.../test_tir_inline_private_functions.py | 4 +-
.../test_tir_transform_flatten_buffer.py | 33 +-
...test_tir_transform_force_narrow_index_to_i32.py | 71 +-
.../test_tir_transform_narrow_datatype.py | 63 +-
tests/python/tirx/test_printer_tir_namespaces.py | 10 +-
tests/python/tvmscript/test_tvmscript_complete.py | 10 +-
.../tvmscript/test_tvmscript_error_report.py | 4 +-
.../tvmscript/test_tvmscript_ir_builder_tir.py | 16 +-
.../python/tvmscript/test_tvmscript_printer_tir.py | 4 +-
tests/python/tvmscript/test_tvmscript_roundtrip.py | 40 +-
256 files changed, 4507 insertions(+), 3862 deletions(-)
diff --git a/include/tvm/relax/analysis.h b/include/tvm/relax/analysis.h
index 454e218b94..092745f171 100644
--- a/include/tvm/relax/analysis.h
+++ b/include/tvm/relax/analysis.h
@@ -29,6 +29,7 @@
#include <tvm/relax/expr.h>
#include <tvm/relax/op_attr_types.h>
#include <tvm/relax/type.h>
+#include <tvm/s_tir/stmt.h>
#include <tvm/tirx/function.h>
#include <tvm/tirx/index_map.h>
@@ -646,7 +647,7 @@ TVM_DLL bool CheckWellFormed(ffi::Variant<IRModule, Function> obj, bool check_ty
* from the object (block or buffer) to it's index map transformation.
*/
-TVM_DLL ffi::Map<tirx::SBlock, ffi::Map<ffi::ObjectRef, tirx::IndexMap>> SuggestLayoutTransforms(
+TVM_DLL ffi::Map<s_tir::SBlock, ffi::Map<ffi::ObjectRef, tirx::IndexMap>> SuggestLayoutTransforms(
const Function& fn, ffi::Array<tirx::IndexMap> write_buffer_transformations);
/* \brief Collect variables whose value can be computed at compile-time
diff --git a/include/tvm/relax/distributed/axis_group_graph.h b/include/tvm/relax/distributed/axis_group_graph.h
index 1b94df4abc..9590394297 100644
--- a/include/tvm/relax/distributed/axis_group_graph.h
+++ b/include/tvm/relax/distributed/axis_group_graph.h
@@ -23,8 +23,9 @@
#include <tvm/arith/iter_affine_map.h>
#include <tvm/relax/distributed/type.h>
#include <tvm/relax/expr.h>
+#include <tvm/s_tir/stmt.h>
+#include <tvm/s_tir/stmt_functor.h>
#include <tvm/tirx/function.h>
-#include <tvm/tirx/stmt_functor.h>
#include <algorithm>
#include <limits>
@@ -65,7 +66,7 @@ Var GetShardingVarFromIndex(PrimExpr index, ffi::Map<Var, Range> var_range,
* \brief Construct an axis group graph from a PrimFunc. Two buffer axis are connected if they
* are accessed by the same index.
*/
-class BufferAxisGraphExtractor : public StmtExprVisitor {
+class BufferAxisGraphExtractor : public s_tir::StmtExprVisitor {
public:
static std::vector<std::vector<TIRVarAxis>> GetTIRVarAxisGraph(const PrimFunc& prim_func) {
auto extractor = ffi::make_object<BufferAxisGraphExtractor>();
@@ -119,14 +120,14 @@ class BufferAxisGraphExtractor : public StmtExprVisitor {
private:
ffi::Optional<VisitInterrupt> Visit_(const BufferStoreNode* op) final {
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(op));
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(s_tir::StmtExprVisitor::Visit_(op));
buffer_access_indices_.push_back({op->buffer, op->indices});
return std::nullopt;
}
ffi::Optional<VisitInterrupt> Visit_(const TensorLoadNode* op) final {
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(op));
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(s_tir::StmtExprVisitor::Visit_(op));
buffer_access_indices_.push_back({op->source.as_or_throw<tvm::tirx::BufferVar>(), op->indices});
return std::nullopt;
@@ -158,12 +159,12 @@ class BufferAxisGraphExtractor : public StmtExprVisitor {
return true;
}
- ffi::Optional<VisitInterrupt> Visit_(const SBlockNode* op) final {
+ ffi::Optional<VisitInterrupt> Visit_(const s_tir::SBlockNode* op) final {
if (op->name_hint == "root") {
- return StmtExprVisitor::Visit_(op);
+ return s_tir::StmtExprVisitor::Visit_(op);
}
buffer_access_indices_.clear();
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(op));
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(s_tir::StmtExprVisitor::Visit_(op));
iter_var_range_.clear();
for (const auto& iter_var : op->iter_vars) {
iter_var_range_.Set(iter_var->var, iter_var->dom);
diff --git a/include/tvm/s_tir/analysis.h b/include/tvm/s_tir/analysis.h
index 4a7c72d386..acaf2a4ba6 100644
--- a/include/tvm/s_tir/analysis.h
+++ b/include/tvm/s_tir/analysis.h
@@ -27,6 +27,7 @@
#include <tvm/ir/module.h>
#include <tvm/ir/prim/expr.h>
#include <tvm/ir/transform.h>
+#include <tvm/s_tir/stmt.h>
#include <tvm/target/target.h>
#include <tvm/tirx/function.h>
#include <tvm/tirx/stmt.h>
@@ -48,7 +49,7 @@ namespace tirx {
* - third: opaque regions
*/
TVM_DLL ffi::Array<ffi::Array<TensorRegion>> GetSBlockAccessRegion(
- const SBlock& block, const ffi::Map<Var, BufferVar>& buffer_var_map);
+ const s_tir::SBlock& block, const ffi::Map<Var, BufferVar>& buffer_var_map);
/*!
* \brief Auto detect the block read/write region according to its body stmt. An opaque access will
@@ -59,7 +60,7 @@ TVM_DLL ffi::Array<ffi::Array<TensorRegion>> GetSBlockAccessRegion(
* \return An array only consisting of the read regions and write regions of the input block
*/
TVM_DLL ffi::Array<ffi::Array<TensorRegion>> GetSBlockReadWriteRegion(
- const SBlock& block, const ffi::Map<Var, BufferVar>& buffer_var_map);
+ const s_tir::SBlock& block, const ffi::Map<Var, BufferVar>& buffer_var_map);
/*!
* \brief Detect the lowest common ancestor(LCA) of buffer access, including both high-level
@@ -85,7 +86,7 @@ TVM_DLL ffi::Map<BufferVar, ffi::Optional<Stmt>> DetectBufferAccessLCA(const Pri
* \param mod The input TIR module.
* \return The anchor block if found, nullptr otherwise.
*/
-const tirx::SBlockNode* FindAnchorBlock(const IRModule& mod);
+const s_tir::SBlockNode* FindAnchorBlock(const IRModule& mod);
} // namespace tirx
@@ -97,6 +98,11 @@ class Analyzer;
namespace s_tir {
using namespace tvm::tirx;
+/*! \brief Verify variable/buffer definitions, load types and schedulable block boundaries. */
+TVM_DLL bool VerifyWellFormed(const tirx::PrimFunc& func, bool assert_mode = true);
+/*! \brief Verify S-TIR or mixed modules, including definitions shared across functions. */
+TVM_DLL bool VerifyWellFormed(const IRModule& mod, bool assert_mode = true);
+
/*!
* \brief Estimate the FLOPs of a TIR fragment.
* \param stmt The TIR fragment to be estimated.
diff --git a/include/tvm/s_tir/sblock_scope.h b/include/tvm/s_tir/sblock_scope.h
index 34ff87025f..3509f2a6db 100644
--- a/include/tvm/s_tir/sblock_scope.h
+++ b/include/tvm/s_tir/sblock_scope.h
@@ -26,9 +26,10 @@
#define TVM_S_TIR_SBLOCK_SCOPE_H_
#include <tvm/ir/module.h>
+#include <tvm/s_tir/stmt.h>
+#include <tvm/s_tir/stmt_functor.h>
#include <tvm/tirx/function.h>
#include <tvm/tirx/stmt.h>
-#include <tvm/tirx/stmt_functor.h>
#include <unordered_map>
#include <utility>
@@ -41,7 +42,7 @@ namespace tirx {
* \brief An object that refers to schedulable elements (block/for-loop) in TensorIR, aka "sref".
*
* Glossary
- * - SBlock sref: A StmtSRef that points to a TensorIR SBlock.
+ * - s_tir::SBlock sref: A StmtSRef that points to a TensorIR s_tir::SBlock.
* - Loop sref: A StmtSRef that points to a TensorIR for loop.
* - Parent sref: The parent reference of an sref is the block or loop reference to the closest
schedulable statement. We define closest to be the nearest schedulable statement of an ancestor in
@@ -87,7 +88,7 @@ class StmtSRefNode : public ffi::Object {
* It serves the same purpose as `ffi::ObjectRef::as`, but does not acquire strong reference to
* `stmt`
* \tparam StmtType The type that `this->stmt` to be downcasted to. Presumably
- * tvm::tirx::SBlockNode or tvm::tirx::ForNode
+ * tvm::s_tir::SBlockNode or tvm::tirx::ForNode
* \return nullptr if type check fails, otherwise the casted result for `this->stmt`
*/
template <typename StmtType>
@@ -144,13 +145,13 @@ class StmtSRef : public ffi::ObjectRef {
TVM_DLL static StmtSRef RootMark();
};
-class SRefTreeCreator : public StmtExprVisitor {
+class SRefTreeCreator : public s_tir::StmtExprVisitor {
public:
- using StmtExprVisitor::Visit_;
+ using s_tir::StmtExprVisitor::Visit_;
ffi::Optional<VisitInterrupt> Visit(ffi::AnyView value) override {
if (value.as<ExprNode>()) return std::nullopt;
- return StmtExprVisitor::Visit(value);
+ return s_tir::StmtExprVisitor::Visit(value);
}
/*!
@@ -185,7 +186,7 @@ class SRefTreeCreator : public StmtExprVisitor {
ffi::Optional<VisitInterrupt> Visit_(const ForNode* loop) final;
- ffi::Optional<VisitInterrupt> Visit_(const SBlockRealizeNode* realize) final;
+ ffi::Optional<VisitInterrupt> Visit_(const s_tir::SBlockRealizeNode* realize) final;
ffi::Optional<VisitInterrupt> Visit_(const SeqStmtNode* seq_stmt) final;
@@ -251,7 +252,7 @@ class Dependency : public ffi::ObjectRef {
* For example even leaf nodes have a scope node, even though they have no dependencies.
*
* Glossary:
- * - SBlock scope: A contiguous subtree of the sref tree, rooted at each SBlock sref,
+ * - s_tir::SBlock scope: A contiguous subtree of the sref tree, rooted at each s_tir::SBlock sref,
* whose components are:
* - scope root: a block sref
* - internal srefs: loop srefs
diff --git a/include/tvm/s_tir/schedule/schedule.h b/include/tvm/s_tir/schedule/schedule.h
index b149eb3b2c..7af08a1630 100644
--- a/include/tvm/s_tir/schedule/schedule.h
+++ b/include/tvm/s_tir/schedule/schedule.h
@@ -23,6 +23,7 @@
#include <tvm/s_tir/random_engine.h>
#include <tvm/s_tir/schedule/state.h>
#include <tvm/s_tir/schedule/trace.h>
+#include <tvm/s_tir/stmt.h>
#include <tvm/tirx/index_map.h>
namespace tvm {
diff --git a/include/tvm/s_tir/schedule/state.h b/include/tvm/s_tir/schedule/state.h
index bbaeb30396..feb1650ee2 100644
--- a/include/tvm/s_tir/schedule/state.h
+++ b/include/tvm/s_tir/schedule/state.h
@@ -27,6 +27,7 @@
#include <tvm/ir/module.h>
#include <tvm/ir/prim/expr.h>
#include <tvm/s_tir/sblock_scope.h>
+#include <tvm/s_tir/stmt.h>
#include <tvm/tirx/function.h>
#include <unordered_map>
diff --git a/include/tvm/s_tir/stmt.h b/include/tvm/s_tir/stmt.h
index 57dcd70aab..eefe47d860 100644
--- a/include/tvm/s_tir/stmt.h
+++ b/include/tvm/s_tir/stmt.h
@@ -18,7 +18,7 @@
*/
/*!
* \file tvm/s_tir/stmt.h
- * \brief S-TIR (Schedulable TIR) statement attribute declarations.
+ * \brief S-TIR (Schedulable TIR) statements and attributes.
*
* This file contains attribute keys that are specific to the schedulable TIR
* (S-TIR) layer, including meta_schedule annotations and schedule primitive /
@@ -27,18 +27,185 @@
#ifndef TVM_S_TIR_STMT_H_
#define TVM_S_TIR_STMT_H_
-#include <tvm/ir/prim/expr.h>
+#include <tvm/tirx/stmt.h>
namespace tvm {
namespace s_tir {
+
+/*!
+ * \brief Match introduces a constraint that the source buffer region can be remapped to the data
+ * layout specified by the buffer field. The constraint can be checked in later part of lowering (or
+ * optionally during runtime).
+ *
+ * MatchBufferRegion provides a mechanism to represent data layout and compactness constraints in
+ * low-level hardware primitives in the IR and defer the check after the sequence of
+ * transformations.
+ */
+class MatchBufferRegionNode : public ffi::Object {
+ public:
+ /*! \brief The target buffer. */
+ tirx::BufferVar buffer;
+ /*! \brief The source buffer region. */
+ TensorRegion source;
+
+ static void RegisterReflection() {
+ namespace refl = tvm::ffi::reflection;
+ refl::ObjectDef<MatchBufferRegionNode>()
+ .def_ro("buffer", &MatchBufferRegionNode::buffer, refl::AttachFieldFlag::SEqHashDefSimple())
+ .def_ro("source", &MatchBufferRegionNode::source);
+ }
+
+ static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("s_tir.MatchBufferRegion", MatchBufferRegionNode, ffi::Object);
+};
+
+/*!
+ * \brief Managed reference to MatchBufferRegionNode.
+ * \sa MatchBufferRegionNode
+ */
+class MatchBufferRegion : public ffi::ObjectRef {
+ public:
+ TVM_DLL explicit MatchBufferRegion(tirx::BufferVar buffer, TensorRegion source);
+
+ TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(MatchBufferRegion, ffi::ObjectRef,
+ MatchBufferRegionNode);
+ TVM_DEFINE_OBJECT_REF_COW_METHOD(MatchBufferRegionNode);
+};
+
+/*!
+ * \brief A block is a basic schedule unit in TIR.
+ * \note SBlock's body is parameterized by iter vars.
+ * \code
+ *
+ * with T.sblock(name):
+ * v0 = T.axis.S(domain, value0)
+ * v1 = T.axis.R(domain, value1)
+ * ...
+ * T.reads([buffer0[start:end, ...], ...])
+ * T.writes([buffer1[start:end, ...], ...])
+ * T.where(predicate)
+ * buffer2 = T.alloc_buffer(shape, dtype)
+ * buffer3 = T.match_buffer(source_buffer[start:end, ...])
+ * T.attr({attr_key: attr_value, ...})
+ * with T.init():
+ * // init body
+ * // body
+ *
+ * \endcode
+ */
+class SBlockNode : public tirx::StmtNode {
+ public:
+ /*! \brief The variables of the block. */
+ ffi::Array<tirx::IterVar> iter_vars;
+ /*! \brief The read buffer regions of the block. */
+ ffi::Array<TensorRegion> reads;
+ /*! \brief The write buffer regions of the block. */
+ ffi::Array<TensorRegion> writes;
+ /*! \brief The name_hint of the block. */
+ ffi::String name_hint;
+ /*! \brief The buffer allocated in the block. */
+ ffi::Array<tirx::BufferVar> alloc_buffers;
+ /*! \brief The match buffer regions. */
+ ffi::Array<MatchBufferRegion> match_buffers;
+ /*! \brief The annotation of the block. */
+ ffi::Map<ffi::String, ffi::Any> annotations;
+ /*!
+ * \brief The init statement is executed during the first iteration of reduction loops in a
+ * reduction block. The optional init field allows us to represent initialization and
+ * reduction update in a single block and transform them collectively.
+ * We also provide primitives to decompose the init into a separate block during scheduling.
+ * Init field is `std::nullopt` if there is no reduction iter_vars
+ */
+ ffi::Optional<tirx::Stmt> init;
+ /*! \brief The body of the block. */
+ tirx::Stmt body;
+
+ static void RegisterReflection() {
+ namespace refl = tvm::ffi::reflection;
+ refl::ObjectDef<SBlockNode>()
+ .def_ro("iter_vars", &SBlockNode::iter_vars)
+ .def_ro("reads", &SBlockNode::reads)
+ .def_ro("writes", &SBlockNode::writes)
+ .def_ro("name_hint", &SBlockNode::name_hint, refl::AttachFieldFlag::SEqHashIgnore())
+ .def_ro("alloc_buffers", &SBlockNode::alloc_buffers,
+ refl::AttachFieldFlag::SEqHashDefSimple())
+ .def_ro("match_buffers", &SBlockNode::match_buffers)
+ .def_ro("annotations", &SBlockNode::annotations)
+ .def_ro("init", &SBlockNode::init)
+ .def_ro("body", &SBlockNode::body);
+ }
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("s_tir.SBlock", SBlockNode, tirx::StmtNode);
+};
+
+/*!
+ * \brief Managed reference to SBlockNode.
+ * \sa SBlockNode
+ */
+class SBlock : public tirx::Stmt {
+ public:
+ TVM_DLL explicit SBlock(
+ ffi::Array<tirx::IterVar> iter_vars, ffi::Array<TensorRegion> reads,
+ ffi::Array<TensorRegion> writes, ffi::String name_hint, tirx::Stmt body,
+ ffi::Optional<tirx::Stmt> init = std::nullopt,
+ ffi::Array<tirx::BufferVar> alloc_buffers = ffi::Array<tirx::BufferVar>(),
+ ffi::Array<MatchBufferRegion> match_buffers = ffi::Array<MatchBufferRegion>(),
+ ffi::Map<ffi::String, ffi::Any> annotations = ffi::Map<ffi::String, ffi::Any>(),
+ Span span = Span());
+
+ TVM_DLL explicit SBlock(ffi::String name_hint, tirx::Stmt body,
+ ffi::Array<tirx::BufferVar> alloc_buffers = ffi::Array<tirx::BufferVar>(),
+ Span span = Span());
+
+ TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(SBlock, tirx::Stmt, SBlockNode);
+ TVM_DEFINE_OBJECT_REF_COW_METHOD(SBlockNode);
+};
+
+/*!
+ * \brief A block realization node represents execution of the block at the binding values.
+ */
+class SBlockRealizeNode : public tirx::StmtNode {
+ public:
+ /*! \brief The corresponding values of the iter vars. */
+ ffi::Array<PrimExpr> iter_values;
+ /*!
+ * \brief The predicate of the block realization, the block will only be executed when the
+ * predicate is true.
+ */
+ PrimExpr predicate;
+ /*! \brief The block to be realized. */
+ SBlock block;
+
+ static void RegisterReflection() {
+ namespace refl = tvm::ffi::reflection;
+ refl::ObjectDef<SBlockRealizeNode>()
+ .def_ro("iter_values", &SBlockRealizeNode::iter_values)
+ .def_ro("predicate", &SBlockRealizeNode::predicate)
+ .def_ro("block", &SBlockRealizeNode::block);
+ }
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("s_tir.SBlockRealize", SBlockRealizeNode, tirx::StmtNode);
+};
+
+/*!
+ * \brief Managed reference to BlockRealizeNode
+ * \sa BlockRealizeNode
+ */
+class SBlockRealize : public tirx::Stmt {
+ public:
+ TVM_DLL explicit SBlockRealize(ffi::Array<PrimExpr> iter_values, PrimExpr predicate, SBlock block,
+ Span span = Span());
+
+ TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(SBlockRealize, tirx::Stmt, SBlockRealizeNode);
+ TVM_DEFINE_OBJECT_REF_COW_METHOD(SBlockRealizeNode);
+};
+
namespace attr {
/*!
* \brief Annotations for invoking and synchronizing asynchronous operations.
*/
constexpr const char* async_commit_queue_scope = "async_commit_queue_scope";
-constexpr const char* async_wait_queue_scope = "async_wait_queue_scope";
-constexpr const char* async_wait_inflight_count = "async_wait_inflight_count";
+constexpr const char* async_wait_queue_scope = tirx::attr::async_wait_queue_scope;
+constexpr const char* async_wait_inflight_count = tirx::attr::async_wait_inflight_count;
/*!
* \brief Mark that the attached statement runs asynchronously.
@@ -77,7 +244,7 @@ constexpr const char* pragma_loop_partition_hint = "pragma_loop_partition_hint";
constexpr const char* reduce_scope = "reduce_scope";
/*! \brief Mark launching of a virtual thread. */
-constexpr const char* virtual_thread = "virtual_thread";
+constexpr const char* virtual_thread = tirx::attr::virtual_thread;
// -----------------------------------------------------------------------
// meta_schedule annotations
diff --git a/include/tvm/s_tir/stmt_functor.h b/include/tvm/s_tir/stmt_functor.h
new file mode 100644
index 0000000000..8d4734c07c
--- /dev/null
+++ b/include/tvm/s_tir/stmt_functor.h
@@ -0,0 +1,128 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+/*!
+ * \file tvm/s_tir/stmt_functor.h
+ * \brief Native traversal and statement dispatch for schedulable TIR.
+ */
+#ifndef TVM_S_TIR_STMT_FUNCTOR_H_
+#define TVM_S_TIR_STMT_FUNCTOR_H_
+
+#include <tvm/s_tir/stmt.h>
+#include <tvm/tirx/stmt_functor.h>
+
+namespace tvm {
+namespace s_tir {
+
+/*!
+ * \brief Extend TIRX statement dispatch with schedulable blocks.
+ * \tparam FType The statement signature, using the native Dispatch API.
+ */
+template <typename FType>
+class StmtFunctor;
+
+template <typename R, typename... Args>
+class StmtFunctor<R(const tirx::Stmt&, Args...)>
+ : public tirx::StmtFunctor<R(const tirx::Stmt&, Args...)> {
+ using Parent = tirx::StmtFunctor<R(const tirx::Stmt&, Args...)>;
+
+ public:
+ TVM_DEFINE_OBJECT_FUNCTOR_DEFAULT_CONSTRUCTOR(StmtFunctor, Parent)
+ using Parent::Dispatch_;
+
+ virtual R Dispatch_(const SBlockNode* op, Args... args) {
+ return this->DispatchDefault_(op, std::forward<Args>(args)...);
+ }
+ virtual R Dispatch_(const SBlockRealizeNode* op, Args... args) {
+ return this->DispatchDefault_(op, std::forward<Args>(args)...);
+ }
+
+ protected:
+ using VTable = typename Parent::VTable;
+ explicit StmtFunctor(const VTable* vtable) : Parent(vtable) {}
+ static void InitVTable(VTable* vtable) {
+ Parent::InitVTable(vtable);
+ Parent::template SetDispatch<StmtFunctor, SBlockNode>(vtable);
+ Parent::template SetDispatch<StmtFunctor, SBlockRealizeNode>(vtable);
+ }
+};
+
+/*!
+ * \brief Extend native TIRX traversal with schedulable block semantics.
+ *
+ * Block iterator binders and annotations are not expression uses. Buffer
+ * definitions precede their regions; ordinary statements reuse TIRX hooks.
+ * Structural traversal remains available independently with its full field walk.
+ * Generic TIRX visitors traverse these nodes structurally, including whole
+ * iterators and annotations. Both paths visit allocation and match-buffer
+ * definitions before region uses. Only this S-TIR subclass supplies native
+ * block hooks; the core TIRX table does not register dialect nodes.
+ */
+class TVM_DLL StmtExprVisitor : public tirx::StmtExprVisitor {
+ public:
+ TVM_DEFINE_OBJECT_FUNCTOR_DEFAULT_CONSTRUCTOR(StmtExprVisitor, tirx::StmtExprVisitor)
+ using tirx::StmtExprVisitor::Visit;
+ using tirx::StmtExprVisitor::Visit_;
+
+ virtual ffi::Optional<VisitInterrupt> Visit_(const SBlockNode* op);
+ virtual ffi::Optional<VisitInterrupt> Visit_(const SBlockRealizeNode* op);
+
+ // Shared native traversal for specialized TIRX helpers extended by S-TIR.
+ static ffi::Optional<VisitInterrupt> VisitBlock(tirx::StmtExprVisitor* visitor,
+ const SBlockNode* op);
+ static ffi::Optional<VisitInterrupt> VisitBlockRealize(tirx::StmtExprVisitor* visitor,
+ const SBlockRealizeNode* op);
+
+ protected:
+ explicit StmtExprVisitor(const VTable* vtable) : tirx::StmtExprVisitor(vtable) {}
+ static void InitVTable(VTable* vtable);
+};
+
+/*!
+ * \brief Extend native TIRX mutation while preserving block iterator binders.
+ *
+ * Reuses inherited remapping and ownership checks. Block annotations are left
+ * intact; structural mutation separately provides the full field rewrite.
+ * Generic TIRX mutators instead use the full structural rewrite, including
+ * iterator definitions and annotations. Both paths establish allocation and
+ * match-buffer remaps before visiting region uses and preserve copy-on-write.
+ */
+class TVM_DLL StmtExprMutator : public tirx::StmtExprMutator {
+ public:
+ TVM_DEFINE_OBJECT_FUNCTOR_DEFAULT_CONSTRUCTOR(StmtExprMutator, tirx::StmtExprMutator)
+ using tirx::StmtExprMutator::Mutate;
+ using tirx::StmtExprMutator::Mutate_;
+
+ virtual UnchangedOr<tirx::Stmt> Mutate_(const SBlockNode* op, InplaceMode inplace_mode);
+ virtual UnchangedOr<tirx::Stmt> Mutate_(const SBlockRealizeNode* op, InplaceMode inplace_mode);
+
+ // Share block ownership and binder rules with specialized TIRX helpers.
+ static UnchangedOr<tirx::Stmt> MutateBlock(tirx::StmtExprMutator* mutator, const SBlockNode* op,
+ InplaceMode inplace_mode);
+ static UnchangedOr<tirx::Stmt> MutateBlockRealize(tirx::StmtExprMutator* mutator,
+ const SBlockRealizeNode* op,
+ InplaceMode inplace_mode);
+
+ protected:
+ explicit StmtExprMutator(const VTable* vtable) : tirx::StmtExprMutator(vtable) {}
+ static void InitVTable(VTable* vtable);
+};
+
+} // namespace s_tir
+} // namespace tvm
+#endif // TVM_S_TIR_STMT_FUNCTOR_H_
diff --git a/include/tvm/s_tir/tensor_intrin.h b/include/tvm/s_tir/tensor_intrin.h
new file mode 100644
index 0000000000..cff27d56a1
--- /dev/null
+++ b/include/tvm/s_tir/tensor_intrin.h
@@ -0,0 +1,90 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file tvm/s_tir/tensor_intrin.h
+ * \brief Tensor intrinsics for schedulable TIR.
+ */
+#ifndef TVM_S_TIR_TENSOR_INTRIN_H_
+#define TVM_S_TIR_TENSOR_INTRIN_H_
+
+#include <tvm/tirx/function.h>
+
+namespace tvm {
+namespace s_tir {
+
+/*!
+ * \brief Tensor intrinsics for tensorization
+ */
+class TensorIntrinNode : public ffi::Object {
+ public:
+ /*! \brief The function to describe the computation. */
+ tirx::PrimFunc desc;
+ /*! \brief The function of the implementation for the execution. */
+ tirx::PrimFunc impl;
+
+ static void RegisterReflection() {
+ namespace refl = tvm::ffi::reflection;
+ refl::ObjectDef<TensorIntrinNode>()
+ .def_ro("desc", &TensorIntrinNode::desc)
+ .def_ro("impl", &TensorIntrinNode::impl);
+ }
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("s_tir.TensorIntrin", TensorIntrinNode, ffi::Object);
+};
+
+/*!
+ * \brief Managed reference to TensorIntrinNode.
+ */
+class TensorIntrin : public ffi::ObjectRef {
+ public:
+ /*!
+ * \brief Constructor
+ * \param desc The function to describe the computation.
+ * \param impl The function of the implementation for the execution.
+ */
+ TVM_DLL explicit TensorIntrin(tirx::PrimFunc desc, tirx::PrimFunc impl);
+
+ /*!
+ * \brief Create and register a TensorIntrin. After registration, the TensorIntrin can be looked
+ * up with its name.
+ * \param name The name of the TensorIntrin to register
+ * \param intrin The TensorIntrin to register.
+ * \param override Whether override existing intrinsic.
+ * \throws This method throws an exception if the TensorIntrin with the specified name already
+ * exists.
+ */
+ TVM_DLL static void Register(ffi::String name, TensorIntrin intrin, bool override = false);
+
+ /*!
+ * \brief Look up TensorIntrin by name. Raises an exception if not found.
+ * \param name The name of the TensorIntrin.
+ * \param allow_missing Whether to allow missing tensor intrin. If false, an exception is raised
+ * if the tensor intrin is not found.
+ * \return The TensorIntrin with the specified name.
+ * \throws This method throws an exception if the TensorIntrin does not exist and allow_missing is
+ * false.
+ */
+ TVM_DLL static ffi::Optional<TensorIntrin> Get(ffi::String name, bool allow_missing = false);
+
+ TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(TensorIntrin, ffi::ObjectRef, TensorIntrinNode);
+};
+
+} // namespace s_tir
+} // namespace tvm
+#endif // TVM_S_TIR_TENSOR_INTRIN_H_
diff --git a/include/tvm/s_tir/transform.h b/include/tvm/s_tir/transform.h
index 6bbc91edfb..0a92264c4c 100644
--- a/include/tvm/s_tir/transform.h
+++ b/include/tvm/s_tir/transform.h
@@ -26,6 +26,7 @@
#include <tvm/ir/prim/expr.h>
#include <tvm/ir/transform.h>
+#include <tvm/s_tir/stmt.h>
#include <tvm/target/target.h>
#include <tvm/tirx/transform.h>
@@ -49,6 +50,13 @@ using tirx::transform::CreatePrimFuncPass;
using tvm::transform::Pass;
using tvm::transform::PassContext;
+/*! \brief De-duplicate definitions, including schedulable block iterators, across PrimFuncs. */
+TVM_DLL Pass ConvertSSA();
+
+/*! \brief Simplify schedulable TIR using block iteration constraints and shared simplifier options.
+ */
+TVM_DLL Pass StmtSimplify();
+
/*!
* \brief Canonicalize loop to start from zero .
* \return The pass.
diff --git a/include/tvm/s_tir/utils.h b/include/tvm/s_tir/utils.h
index cbbe930574..4d127818e9 100644
--- a/include/tvm/s_tir/utils.h
+++ b/include/tvm/s_tir/utils.h
@@ -20,6 +20,7 @@
#define TVM_S_TIR_UTILS_H_
#include <tvm/s_tir/sblock_scope.h>
+#include <tvm/s_tir/stmt.h>
#include <tvm/tirx/stmt.h>
#include <unordered_map>
@@ -32,7 +33,7 @@ namespace tirx {
* then check if the downcasting succeeded.
* \param Result The result variable, used for checking
* \param SRef The SRef to be cast
- * \param Type The type to be cast to, can be SBlock or For
+ * \param Type The type to be cast to, can be s_tir::SBlock or For
*/
#define TVM_SREF_AS_OR_ERR(Result, SRef, Type) \
SRef->StmtAs<Type>(); \
@@ -48,7 +49,7 @@ namespace tirx {
*/
#define TVM_SREF_TO_SBLOCK(SRef) \
[&]() { \
- auto result = TVM_SREF_AS_OR_ERR(result, (SRef), ::tvm::tirx::SBlockNode) \
+ auto result = TVM_SREF_AS_OR_ERR(result, (SRef), ::tvm::s_tir::SBlockNode) \
<< "Expects StmtSRef `" << #SRef << "` points to `Block`, but gets: " \
<< ((SRef)->stmt ? (SRef)->stmt->GetTypeKey() : "None"); \
return result; \
@@ -101,15 +102,15 @@ namespace tirx {
* \param stmt The statement, or the realize node of the statement whose sref to be set
* \param seq_index The seq_index to be set
* \param include_loops Ignore ForNodes if this value is false
- * \note The method is NOP for statements that are not schedulable, i.e. not For or SBlock
+ * \note The method is NOP for statements that are not schedulable, i.e. not For or s_tir::SBlock
*/
inline void SetSeqIndex(std::unordered_map<const StmtNode*, StmtSRef>& stmt2ref, // NOLINT(*)
const Stmt& stmt, int seq_index, bool include_loops = true) {
- if (const auto* realize = stmt.as<SBlockRealizeNode>()) {
- const SBlockNode* block = realize->block.get();
+ if (const auto* realize = stmt.as<s_tir::SBlockRealizeNode>()) {
+ const s_tir::SBlockNode* block = realize->block.get();
TVM_FFI_ICHECK(stmt2ref.count(block));
stmt2ref.at(block)->seq_index = seq_index;
- } else if (const auto* block = stmt.as<SBlockNode>()) {
+ } else if (const auto* block = stmt.as<s_tir::SBlockNode>()) {
TVM_FFI_ICHECK(stmt2ref.count(block));
stmt2ref.at(block)->seq_index = seq_index;
} else if (const auto* loop = stmt.as<ForNode>()) {
diff --git a/include/tvm/script/ir_builder/base.h b/include/tvm/script/ir_builder/base.h
index 85c8f991d9..417e9b0c38 100644
--- a/include/tvm/script/ir_builder/base.h
+++ b/include/tvm/script/ir_builder/base.h
@@ -49,7 +49,7 @@ namespace ir_builder {
*
* \endcode
*
- * The `T::MatchBuffer` below instead generates `MatchBufferRegion` in a TIR block:
+ * The `T::MatchBuffer` below instead generates `s_tir::MatchBufferRegion` in a TIR block:
*
* \code {.cpp}
*
diff --git a/include/tvm/tirx/analysis.h b/include/tvm/tirx/analysis.h
index 1379046a21..11fa65f0b0 100644
--- a/include/tvm/tirx/analysis.h
+++ b/include/tvm/tirx/analysis.h
@@ -27,7 +27,6 @@
#include <tvm/ir/module.h>
#include <tvm/ir/prim/expr.h>
#include <tvm/ir/transform.h>
-#include <tvm/s_tir/analysis.h>
#include <tvm/target/target.h>
#include <tvm/tirx/function.h>
#include <tvm/tirx/op_attr_types.h>
@@ -131,12 +130,7 @@ TVM_DLL size_t CalculateWorkspaceBytes(const PrimFunc& func, int64_t workspace_b
*
* - Each variable has a single point of definition.
*
- * - Expressions within a tirx::SBlock may not reference variables
- * defined outside the block. For example, for a block with iter
- * vars `vi, vj = T.axis.remap('SS', [i,j])`, the statement
- * `B[i,j] = A[i,j]` would be ill-formed, because it uses the loop
- * variables `i` and `j` instead of the block variables `vi` and
- * `vj`.
+ * Dialect statements require their dialect-specific verifier.
*
* \param func The PrimFunc to be verified.
* \param assert_mode The indicator if it raises an error when the function is not well-formed.
diff --git a/include/tvm/tirx/function.h b/include/tvm/tirx/function.h
index 8f73cc15cf..67e43a3c02 100644
--- a/include/tvm/tirx/function.h
+++ b/include/tvm/tirx/function.h
@@ -124,62 +124,6 @@ class PrimFunc : public BaseFunc {
TVM_DEFINE_OBJECT_REF_COW_METHOD(PrimFuncNode);
};
-/*!
- * \brief Tensor intrinsics for tensorization
- */
-class TensorIntrinNode : public ffi::Object {
- public:
- /*! \brief The function to describe the computation. */
- PrimFunc desc;
- /*! \brief The function of the implementation for the execution. */
- PrimFunc impl;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<TensorIntrinNode>()
- .def_ro("desc", &TensorIntrinNode::desc)
- .def_ro("impl", &TensorIntrinNode::impl);
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.TensorIntrin", TensorIntrinNode, ffi::Object);
-};
-
-/*!
- * \brief Managed reference to TensorIntrinNode.
- */
-class TensorIntrin : public ffi::ObjectRef {
- public:
- /*!
- * \brief Constructor
- * \param desc The function to describe the computation.
- * \param impl The function of the implementation for the execution.
- */
- TVM_DLL explicit TensorIntrin(PrimFunc desc, PrimFunc impl);
-
- /*!
- * \brief Create and register a TensorIntrin. After registration, the TensorIntrin can be looked
- * up with its name.
- * \param name The name of the TensorIntrin to register
- * \param intrin The TensorIntrin to register.
- * \param override Whether override existing intrinsic.
- * \throws This method throws an exception if the TensorIntrin with the specified name already
- * exists.
- */
- TVM_DLL static void Register(ffi::String name, TensorIntrin intrin, bool override = false);
-
- /*!
- * \brief Look up TensorIntrin by name. Raises an exception if not found.
- * \param name The name of the TensorIntrin.
- * \param allow_missing Whether to allow missing tensor intrin. If false, an exception is raised
- * if the tensor intrin is not found.
- * \return The TensorIntrin with the specified name.
- * \throws This method throws an exception if the TensorIntrin does not exist and allow_missing is
- * false.
- */
- TVM_DLL static ffi::Optional<TensorIntrin> Get(ffi::String name, bool allow_missing = false);
-
- TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(TensorIntrin, ffi::ObjectRef, TensorIntrinNode);
-};
-
/*!
* \brief Specialize parameters of PrimFunc.
* \param func The PrimFunc to be specialized.
diff --git a/include/tvm/tirx/script/builder/frame.h b/include/tvm/tirx/script/builder/frame.h
index 38259a1a57..b5c1774808 100644
--- a/include/tvm/tirx/script/builder/frame.h
+++ b/include/tvm/tirx/script/builder/frame.h
@@ -19,6 +19,7 @@
#ifndef TVM_SCRIPT_IR_BUILDER_TIR_FRAME_H_
#define TVM_SCRIPT_IR_BUILDER_TIR_FRAME_H_
+#include <tvm/s_tir/stmt.h>
#include <tvm/script/ir_builder/base.h>
#include <tvm/script/ir_builder/ir/frame.h>
#include <tvm/tirx/exec_scope.h>
@@ -87,7 +88,7 @@ class PrimFuncFrameNode : public TIRFrameNode {
ffi::Array<tvm::tirx::BufferVar> root_alloc_buffers;
// TIR utils
- /*! \brief Whether this PrimFunc uses s_tir semantics (root SBlock wrap,
+ /*! \brief Whether this PrimFunc uses s_tir semantics (root s_tir::SBlock wrap,
* parser layout default = None). Default (false) = tirx semantics. */
bool s_tir;
/*! \brief Whether it is a persistent kernel. */
@@ -152,7 +153,7 @@ class SBlockFrameNode : public TIRFrameNode {
/*! \brief The buffer allocated in the block. */
ffi::Array<tvm::tirx::BufferVar> alloc_buffers;
/*! \brief The match buffer regions. */
- ffi::Array<tvm::tirx::MatchBufferRegion> match_buffers;
+ ffi::Array<tvm::s_tir::MatchBufferRegion> match_buffers;
/*! \brief The annotation of the block. */
ffi::Optional<ffi::Map<ffi::String, Any>> annotations;
/*! \brief The corresponding values of the iter vars. */
diff --git a/include/tvm/tirx/script/builder/ir.h b/include/tvm/tirx/script/builder/ir.h
index 9e3e78a88c..e306c6b276 100644
--- a/include/tvm/tirx/script/builder/ir.h
+++ b/include/tvm/tirx/script/builder/ir.h
@@ -125,7 +125,7 @@ BufferVar MatchBuffer(ffi::ObjectRef param, ffi::Array<PrimExpr> shape,
/*!
* \brief The block declaration statement.
* \param name The name of the block.
- * \param no_realize The flag whether to construct SBlockRealize or SBlock.
+ * \param no_realize The flag whether to construct s_tir::SBlockRealize or s_tir::SBlock.
* \return The SBlockFrame.
*/
SBlockFrame Block(ffi::String name, bool no_realize = false, ffi::String exec_scope = "");
diff --git a/include/tvm/tirx/stmt.h b/include/tvm/tirx/stmt.h
index 65baf862e1..5c3f4f2cd9 100644
--- a/include/tvm/tirx/stmt.h
+++ b/include/tvm/tirx/stmt.h
@@ -769,172 +769,6 @@ class Continue : public Stmt {
TVM_DEFINE_OBJECT_REF_COW_METHOD(ContinueNode);
};
-/*!
- * \brief Match introduces a constraint that the source buffer region can be remapped to the data
- * layout specified by the buffer field. The constraint can be checked in later part of lowering (or
- * optionally during runtime).
- *
- * MatchBufferRegion provides a mechanism to represent data layout and compactness constraints in
- * low-level hardware primitives in the IR and defer the check after the sequence of
- * transformations.
- */
-class MatchBufferRegionNode : public ffi::Object {
- public:
- /*! \brief The target buffer. */
- BufferVar buffer;
- /*! \brief The source buffer region. */
- TensorRegion source;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<MatchBufferRegionNode>()
- .def_ro("buffer", &MatchBufferRegionNode::buffer, refl::AttachFieldFlag::SEqHashDefSimple())
- .def_ro("source", &MatchBufferRegionNode::source);
- }
-
- static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.MatchBufferRegion", MatchBufferRegionNode, ffi::Object);
-};
-
-/*!
- * \brief Managed reference to MatchBufferRegionNode.
- * \sa MatchBufferRegionNode
- */
-class MatchBufferRegion : public ffi::ObjectRef {
- public:
- TVM_DLL explicit MatchBufferRegion(BufferVar buffer, TensorRegion source);
-
- TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(MatchBufferRegion, ffi::ObjectRef,
- MatchBufferRegionNode);
- TVM_DEFINE_OBJECT_REF_COW_METHOD(MatchBufferRegionNode);
-};
-
-/*!
- * \brief A block is a basic schedule unit in TIR.
- * \note SBlock's body is parameterized by iter vars.
- * \code
- *
- * with T.sblock(name):
- * v0 = T.axis.S(domain, value0)
- * v1 = T.axis.R(domain, value1)
- * ...
- * T.reads([buffer0[start:end, ...], ...])
- * T.writes([buffer1[start:end, ...], ...])
- * T.where(predicate)
- * buffer2 = T.alloc_buffer(shape, dtype)
- * buffer3 = T.match_buffer(source_buffer[start:end, ...])
- * T.attr({attr_key: attr_value, ...})
- * with T.init():
- * // init body
- * // body
- *
- * \endcode
- */
-class SBlockNode : public StmtNode {
- public:
- /*! \brief The variables of the block. */
- ffi::Array<IterVar> iter_vars;
- /*! \brief The read buffer regions of the block. */
- ffi::Array<TensorRegion> reads;
- /*! \brief The write buffer regions of the block. */
- ffi::Array<TensorRegion> writes;
- /*! \brief The name_hint of the block. */
- ffi::String name_hint;
- /*! \brief The buffer allocated in the block. */
- ffi::Array<BufferVar> alloc_buffers;
- /*! \brief The match buffer regions. */
- ffi::Array<MatchBufferRegion> match_buffers;
- /*! \brief The annotation of the block. */
- ffi::Map<ffi::String, ffi::Any> annotations;
- /*!
- * \brief The init statement is executed during the first iteration of reduction loops in a
- * reduction block. The optional init field allows us to represent initialization and
- * reduction update in a single block and transform them collectively.
- * We also provide primitives to decompose the init into a separate block during scheduling.
- * Init field is `std::nullopt` if there is no reduction iter_vars
- */
- ffi::Optional<Stmt> init;
- /*! \brief The body of the block. */
- Stmt body;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<SBlockNode>()
- .def_ro("iter_vars", &SBlockNode::iter_vars)
- .def_ro("reads", &SBlockNode::reads)
- .def_ro("writes", &SBlockNode::writes)
- .def_ro("name_hint", &SBlockNode::name_hint, refl::AttachFieldFlag::SEqHashIgnore())
- .def_ro("alloc_buffers", &SBlockNode::alloc_buffers,
- refl::AttachFieldFlag::SEqHashDefSimple())
- .def_ro("match_buffers", &SBlockNode::match_buffers)
- .def_ro("annotations", &SBlockNode::annotations)
- .def_ro("init", &SBlockNode::init)
- .def_ro("body", &SBlockNode::body);
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.SBlock", SBlockNode, StmtNode);
-};
-
-/*!
- * \brief Managed reference to SBlockNode.
- * \sa SBlockNode
- */
-class SBlock : public Stmt {
- public:
- TVM_DLL explicit SBlock(
- ffi::Array<IterVar> iter_vars, ffi::Array<TensorRegion> reads,
- ffi::Array<TensorRegion> writes, ffi::String name_hint, Stmt body,
- ffi::Optional<Stmt> init = std::nullopt,
- ffi::Array<BufferVar> alloc_buffers = ffi::Array<BufferVar>(),
- ffi::Array<MatchBufferRegion> match_buffers = ffi::Array<MatchBufferRegion>(),
- ffi::Map<ffi::String, ffi::Any> annotations = ffi::Map<ffi::String, ffi::Any>(),
- Span span = Span());
-
- TVM_DLL explicit SBlock(ffi::String name_hint, Stmt body,
- ffi::Array<BufferVar> alloc_buffers = ffi::Array<BufferVar>(),
- Span span = Span());
-
- TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(SBlock, Stmt, SBlockNode);
- TVM_DEFINE_OBJECT_REF_COW_METHOD(SBlockNode);
-};
-
-/*!
- * \brief A block realization node represents execution of the block at the binding values.
- */
-class SBlockRealizeNode : public StmtNode {
- public:
- /*! \brief The corresponding values of the iter vars. */
- ffi::Array<PrimExpr> iter_values;
- /*!
- * \brief The predicate of the block realization, the block will only be executed when the
- * predicate is true.
- */
- PrimExpr predicate;
- /*! \brief The block to be realized. */
- SBlock block;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<SBlockRealizeNode>()
- .def_ro("iter_values", &SBlockRealizeNode::iter_values)
- .def_ro("predicate", &SBlockRealizeNode::predicate)
- .def_ro("block", &SBlockRealizeNode::block);
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.SBlockRealize", SBlockRealizeNode, StmtNode);
-};
-
-/*!
- * \brief Managed reference to BlockRealizeNode
- * \sa BlockRealizeNode
- */
-class SBlockRealize : public Stmt {
- public:
- TVM_DLL explicit SBlockRealize(ffi::Array<PrimExpr> iter_values, PrimExpr predicate, SBlock block,
- Span span = Span());
-
- TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(SBlockRealize, Stmt, SBlockRealizeNode);
- TVM_DEFINE_OBJECT_REF_COW_METHOD(SBlockRealizeNode);
-};
-
/*!
* \brief Standalone statement that declares a scope-id binding (e.g. cta_id,
* warp_id, lane_id). Carries a ``ScopeIdDef`` value.
@@ -995,6 +829,11 @@ constexpr const char* pragma_unroll_explicit = "pragma_unroll_explicit";
constexpr const char* storage_alignment = "storage_alignment";
/*! \brief Mark launching extent of thread, used by device API. */
constexpr const char* thread_extent = "thread_extent";
+
+/*! \brief Shared execution attributes consumed before and after block lowering. */
+constexpr const char* virtual_thread = "virtual_thread";
+constexpr const char* async_wait_queue_scope = "async_wait_queue_scope";
+constexpr const char* async_wait_inflight_count = "async_wait_inflight_count";
/*! \brief Annotation key on AllocBuffer marking the allocation as volatile. */
constexpr const char* kVolatile = "tirx.volatile";
/*! \brief Mark buffer initial addr alignment in bytes */
diff --git a/include/tvm/tirx/stmt_functor.h b/include/tvm/tirx/stmt_functor.h
index c1147f5f80..17c93ece72 100644
--- a/include/tvm/tirx/stmt_functor.h
+++ b/include/tvm/tirx/stmt_functor.h
@@ -116,12 +116,6 @@ class StmtFunctor<R(const Stmt&, Args...)> {
virtual R Dispatch_(const EvaluateNode* node, Args... args) {
return DispatchDefault_(node, std::forward<Args>(args)...);
}
- virtual R Dispatch_(const SBlockNode* node, Args... args) {
- return DispatchDefault_(node, std::forward<Args>(args)...);
- }
- virtual R Dispatch_(const SBlockRealizeNode* node, Args... args) {
- return DispatchDefault_(node, std::forward<Args>(args)...);
- }
virtual R Dispatch_(const ScopeIdDefStmtNode* node, Args... args) {
return DispatchDefault_(node, std::forward<Args>(args)...);
}
@@ -141,6 +135,10 @@ class StmtFunctor<R(const Stmt&, Args...)> {
explicit StmtFunctor(const VTable* vtable) : vtable_(vtable) {}
/*! \brief Register statement hooks in a fresh mutable table. */
static void InitVTable(VTable* vtable) {
+ vtable->template SetDispatch<StmtNode>(
+ [](const ffi::ObjectRef& node, TSelf* self, Args... args) -> R {
+ return self->DispatchDefault_(node.get(), std::forward<Args>(args)...);
+ });
SetDispatch<TSelf, BindNode>(vtable);
SetDispatch<TSelf, AttrStmtNode>(vtable);
SetDispatch<TSelf, IfThenElseNode>(vtable);
@@ -155,8 +153,6 @@ class StmtFunctor<R(const Stmt&, Args...)> {
SetDispatch<TSelf, AssertStmtNode>(vtable);
SetDispatch<TSelf, SeqStmtNode>(vtable);
SetDispatch<TSelf, EvaluateNode>(vtable);
- SetDispatch<TSelf, SBlockNode>(vtable);
- SetDispatch<TSelf, SBlockRealizeNode>(vtable);
SetDispatch<TSelf, ScopeIdDefStmtNode>(vtable);
SetDispatch<TSelf, tirx::TilePrimitiveCallNode>(vtable);
}
@@ -216,8 +212,6 @@ class TVM_DLL StmtExprVisitor : public tvm::ExprVisitor {
virtual ffi::Optional<VisitInterrupt> Visit_(const AssertStmtNode* op);
virtual ffi::Optional<VisitInterrupt> Visit_(const SeqStmtNode* op);
virtual ffi::Optional<VisitInterrupt> Visit_(const EvaluateNode* op);
- virtual ffi::Optional<VisitInterrupt> Visit_(const SBlockNode* op);
- virtual ffi::Optional<VisitInterrupt> Visit_(const SBlockRealizeNode* op);
virtual ffi::Optional<VisitInterrupt> Visit_(const ScopeIdDefStmtNode* op);
virtual ffi::Optional<VisitInterrupt> Visit_(const TilePrimitiveCallNode* op);
@@ -233,10 +227,10 @@ class TVM_DLL StmtExprVisitor : public tvm::ExprVisitor {
ffi::Optional<VisitInterrupt> Visit_(const prim::BroadcastNode* op) override;
ffi::Optional<VisitInterrupt> Visit_(const prim::ShuffleNode* op) override;
- protected:
- // Visit definition metadata as uses, separately from the buffer Var definition.
+ /*! \brief Visit definition metadata as uses, separately from the buffer Var definition. */
ffi::Optional<VisitInterrupt> VisitBufferMetadata(const BufferVar& buffer);
+ protected:
explicit StmtExprVisitor(const VTable* vtable) : tvm::ExprVisitor(vtable) {}
static void InitVTable(VTable* vtable);
};
@@ -282,8 +276,6 @@ class TVM_DLL StmtExprMutator : public tvm::ExprMutator {
virtual UnchangedOr<Stmt> Mutate_(const AssertStmtNode* op, InplaceMode inplace_mode);
virtual UnchangedOr<Stmt> Mutate_(const SeqStmtNode* op, InplaceMode inplace_mode);
virtual UnchangedOr<Stmt> Mutate_(const EvaluateNode* op, InplaceMode inplace_mode);
- virtual UnchangedOr<Stmt> Mutate_(const SBlockNode* op, InplaceMode inplace_mode);
- virtual UnchangedOr<Stmt> Mutate_(const SBlockRealizeNode* op, InplaceMode inplace_mode);
virtual UnchangedOr<Stmt> Mutate_(const ScopeIdDefStmtNode* op, InplaceMode inplace_mode);
virtual UnchangedOr<Stmt> Mutate_(const TilePrimitiveCallNode* op, InplaceMode inplace_mode);
@@ -297,7 +289,11 @@ class TVM_DLL StmtExprMutator : public tvm::ExprMutator {
* \param stmt The source statement to be substituted
* \param vmap returns a new value if re-mapping is needed, otherwise returns nullptr.
*
- * Substitution may change the data type of the expression.
+ * This statement overload legalizes only core TIRX nodes. Dtype-changing
+ * substitutions must be applied after lowering dialect blocks: structural
+ * traversal of extension statements does not legalize their iterator domains or
+ * bindings. Type-preserving mappings continue to traverse schedulable blocks
+ * structurally before lowering.
*
* \return The result.
*/
@@ -334,14 +330,13 @@ bool ContainsNode(const Stmt& stmt) {
if (contains_node || value.as<ExprNode>()) {
return std::nullopt;
}
+ if (value.as<Node>()) {
+ contains_node = true;
+ return std::nullopt;
+ }
return StmtExprVisitor::Visit(value);
}
- ffi::Optional<VisitInterrupt> Visit_(const Node* block) override {
- contains_node = true;
- return std::nullopt;
- }
-
bool contains_node{false};
};
diff --git a/python/tvm/backend/trn/tile_primitive/binary/default.py b/python/tvm/backend/trn/tile_primitive/binary/default.py
index 3a432c3fad..85f19d11c5 100644
--- a/python/tvm/backend/trn/tile_primitive/binary/default.py
+++ b/python/tvm/backend/trn/tile_primitive/binary/default.py
@@ -72,7 +72,8 @@ def binary_trn(
return _func(*args, reverse[0]) if inst_types[0] == InstType.TENSOR_SCALAR else _func(*args)
# Define the implementation function
- @T.prim_func
+ # This fragment captures buffers and indices from its insertion scope.
+ @T.prim_func(check_well_formed=False)
def impl():
for b_loop in T.serial(0, b_extent):
with T.attr(0, "tensorized_nki_instruction", 1):
diff --git a/python/tvm/backend/trn/tile_primitive/compose_op/binary_chain.py b/python/tvm/backend/trn/tile_primitive/compose_op/binary_chain.py
index 802227c17e..7c29db3794 100644
--- a/python/tvm/backend/trn/tile_primitive/compose_op/binary_chain.py
+++ b/python/tvm/backend/trn/tile_primitive/compose_op/binary_chain.py
@@ -91,7 +91,8 @@ def binary_chain_trn(op: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc |
# Create implementation
# fmt: off
- @T.prim_func
+ # This fragment captures buffers and indices from its insertion scope.
+ @T.prim_func(check_well_formed=False)
def impl():
for b_loop in T.serial(0, b_extent):
with T.attr(0, "tensorized_nki_instruction", 1):
diff --git a/python/tvm/backend/trn/tile_primitive/compose_op/binary_reduce.py b/python/tvm/backend/trn/tile_primitive/compose_op/binary_reduce.py
index 165ce1a2ce..a3a83bb3a4 100644
--- a/python/tvm/backend/trn/tile_primitive/compose_op/binary_reduce.py
+++ b/python/tvm/backend/trn/tile_primitive/compose_op/binary_reduce.py
@@ -101,7 +101,8 @@ def binary_reduce_trn(op: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc
if reduction_b_extent == 1:
# Direct implementation without intermediate buffer
# fmt: off
- @T.prim_func
+ # This fragment captures buffers and indices from its insertion scope.
+ @T.prim_func(check_well_formed=False)
def impl():
for b_loop in T.serial(0, spatial_b_extent):
with T.attr(0, "tensorized_nki_instruction", 1):
@@ -121,7 +122,8 @@ def binary_reduce_trn(op: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc
else:
# Implementation with intermediate buffer
# fmt: off
- @T.prim_func
+ # This fragment captures buffers and indices from its insertion scope.
+ @T.prim_func(check_well_formed=False)
def impl():
for b_loop in T.serial(0, spatial_b_extent):
for reduction_b_loop in T.serial(0, reduction_b_extent):
diff --git a/python/tvm/backend/trn/tile_primitive/compose_op/unary_reduce.py b/python/tvm/backend/trn/tile_primitive/compose_op/unary_reduce.py
index 9f622cf404..2cc80e57c2 100644
--- a/python/tvm/backend/trn/tile_primitive/compose_op/unary_reduce.py
+++ b/python/tvm/backend/trn/tile_primitive/compose_op/unary_reduce.py
@@ -100,7 +100,8 @@ def unary_reduce_trn(op: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc |
if reduction_b_extent == 1:
# Direct implementation without intermediate buffer
# fmt: off
- @T.prim_func
+ # This fragment captures buffers and indices from its insertion scope.
+ @T.prim_func(check_well_formed=False)
def impl():
for b_loop in T.serial(0, spatial_b_extent):
with T.attr(0, "tensorized_nki_instruction", 1):
@@ -125,7 +126,8 @@ def unary_reduce_trn(op: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc |
return mod["main"]
else:
# fmt: off
- @T.prim_func
+ # This fragment captures buffers and indices from its insertion scope.
+ @T.prim_func(check_well_formed=False)
def impl():
for b_loop in T.serial(0, spatial_b_extent):
for reduction_b_loop in T.serial(0, reduction_b_extent):
diff --git a/python/tvm/backend/trn/tile_primitive/copy/default.py b/python/tvm/backend/trn/tile_primitive/copy/default.py
index e22b3800b6..561fca20c9 100644
--- a/python/tvm/backend/trn/tile_primitive/copy/default.py
+++ b/python/tvm/backend/trn/tile_primitive/copy/default.py
@@ -97,7 +97,8 @@ def transpose_schedule(
)
sctx.add_alloc_buffer(identity_tensor)
- @T.prim_func
+ # This fragment captures buffers and indices from its insertion scope.
+ @T.prim_func(check_well_formed=False)
def identity_init():
with T.attr(0, "tensorized_nki_instruction", 1):
for p_loop in T.serial(0, p_size, annotations={nki_dim: "P"}):
@@ -113,8 +114,8 @@ def transpose_schedule(
dst_buffer = dst_region.source
src_buffer = src_region.source
if dst_buffer.scope() == "trn.psum":
-
- @T.prim_func
+ # This fragment captures buffers and indices from its insertion scope.
+ @T.prim_func(check_well_formed=False)
def transpose_psum_output():
for b_loop in T.serial(0, b_extent):
with T.attr(0, "tensorized_nki_instruction", 1):
@@ -164,7 +165,8 @@ def transpose_schedule(
max_psum_slots = acc_psum.ty.shape[0]
# fmt: off
- @T.prim_func
+ # This fragment captures buffers and indices from its insertion scope.
+ @T.prim_func(check_well_formed=False)
def transpose_sbuf_output():
for b_loop in T.serial(0, b_extent):
for extend_b_loop in T.serial(0, extend_len):
@@ -271,7 +273,8 @@ def copy_trn(op: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc | None:
b_extent = inst_gen.fill_in_block_dim(from_region, b_var)
# fmt: off
- @T.prim_func
+ # This fragment captures buffers and indices from its insertion scope.
+ @T.prim_func(check_well_formed=False)
def impl():
# the additional b loop is to satisfy hardware instuction size limit
for b_loop in T.serial(0, b_extent):
diff --git a/python/tvm/backend/trn/tile_primitive/gemm/default.py b/python/tvm/backend/trn/tile_primitive/gemm/default.py
index fa9a67a7c1..75467dffc8 100644
--- a/python/tvm/backend/trn/tile_primitive/gemm/default.py
+++ b/python/tvm/backend/trn/tile_primitive/gemm/default.py
@@ -241,7 +241,8 @@ def matmul_trn(op: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc | None:
T.evaluate(T.nki.matmul(acc[b_idx % max_psum_slots, lhs_f_loop, rhs_f_loop], A[lhs_indices], B[rhs_indices])) # noqa: E501
if C.scope() == "trn.psum":
- @T.prim_func
+ # This fragment captures buffers and indices from its insertion scope.
+ @T.prim_func(check_well_formed=False)
def impl_C_psum():
for lhs_b_loop, rhs_b_loop, reduction_b_loop in T.grid(lhs_b_extent, rhs_b_extent, reduction_b_extent): # noqa: E501
matmul_inst_macro(lhs_b_loop, rhs_b_loop, reduction_b_loop, C, True, None)
@@ -270,7 +271,8 @@ def matmul_trn(op: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc | None:
check_workspace_buffer(acc_psum, (p_size, largest_psum_per_bank), "trn.psum")
max_psum_slots = acc_psum.ty.shape[0]
- @T.prim_func
+ # This fragment captures buffers and indices from its insertion scope.
+ @T.prim_func(check_well_formed=False)
def impl_C_sbuf():
for lhs_b_loop, rhs_b_loop in T.grid(lhs_b_extent, rhs_b_extent):
for reduction_b_loop in T.serial(0, reduction_b_extent):
diff --git a/python/tvm/backend/trn/tile_primitive/private_alloc.py b/python/tvm/backend/trn/tile_primitive/private_alloc.py
index 81ab5827e5..582b492756 100644
--- a/python/tvm/backend/trn/tile_primitive/private_alloc.py
+++ b/python/tvm/backend/trn/tile_primitive/private_alloc.py
@@ -67,7 +67,8 @@ def alloc_const_bias_trn(
new_shape, dtype=_scalar_dtype(bias), scope="trn.sbuf", buffer_name="const_bias"
)
- @T.prim_func
+ # This fragment captures buffers and indices from its insertion scope.
+ @T.prim_func(check_well_formed=False)
def const_bias_init():
with T.attr(0, "tensorized_nki_instruction", 1):
for p_loop in T.serial(0, par_size, annotations={"nki_dim": "P"}):
@@ -117,7 +118,8 @@ def alloc_identity_trn(
new_shape, dtype=op.srcs[0].source.ty.dtype, scope="trn.sbuf", buffer_name="identity"
)
- @T.prim_func
+ # This fragment captures buffers and indices from its insertion scope.
+ @T.prim_func(check_well_formed=False)
def identity_init():
with T.attr(0, "tensorized_nki_instruction", 1):
for p_loop in T.serial(0, par_size, annotations={nki_dim: "P"}):
diff --git a/python/tvm/backend/trn/tile_primitive/reduction/utils.py b/python/tvm/backend/trn/tile_primitive/reduction/utils.py
index 2d791908b2..78cbe25de1 100644
--- a/python/tvm/backend/trn/tile_primitive/reduction/utils.py
+++ b/python/tvm/backend/trn/tile_primitive/reduction/utils.py
@@ -130,7 +130,8 @@ def reduction_trn(
# fmt: off
# Single-stage reduction implementation
if reduction_b_extent == 1:
- @T.prim_func
+ # This fragment captures buffers and indices from its insertion scope.
+ @T.prim_func(check_well_formed=False)
def impl():
for b_loop in T.serial(0, spatial_b_extent):
with T.attr(0, "tensorized_nki_instruction", 1):
@@ -144,7 +145,8 @@ def reduction_trn(
return impl
# Two-stage reduction implementation
else:
- @T.prim_func
+ # This fragment captures buffers and indices from its insertion scope.
+ @T.prim_func(check_well_formed=False)
def two_stage_reduction():
for b_loop in T.serial(0, spatial_b_extent):
for reduction_b_loop in T.serial(0, reduction_b_extent):
diff --git a/python/tvm/backend/trn/tile_primitive/select/default.py b/python/tvm/backend/trn/tile_primitive/select/default.py
index 9106f3c8a6..8e281c8e9f 100644
--- a/python/tvm/backend/trn/tile_primitive/select/default.py
+++ b/python/tvm/backend/trn/tile_primitive/select/default.py
@@ -109,7 +109,8 @@ def select_trn(op: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc | None:
true_value_buffer = true_value.source
# fmt: off
- @T.prim_func
+ # This fragment captures buffers and indices from its insertion scope.
+ @T.prim_func(check_well_formed=False)
def impl():
for b_loop in T.serial(0, b_extent):
with T.attr(0, "tensorized_nki_instruction", 1):
diff --git a/python/tvm/backend/trn/tile_primitive/unary/utils.py b/python/tvm/backend/trn/tile_primitive/unary/utils.py
index f2589a7a7d..40eabd5602 100644
--- a/python/tvm/backend/trn/tile_primitive/unary/utils.py
+++ b/python/tvm/backend/trn/tile_primitive/unary/utils.py
@@ -105,7 +105,8 @@ def get_const_bias_tensor(bias, shape, dtype, workspace, sctx):
bias_buffer = T.buffer(shape, dtype, scope="trn.sbuf", buffer_name="const_bias")
sctx.add_alloc_buffer(bias_buffer)
- @T.prim_func
+ # This fragment captures buffers and indices from its insertion scope.
+ @T.prim_func(check_well_formed=False)
def const_bias_init():
with T.attr(0, "tensorized_nki_instruction", 1):
for p_loop in T.serial(0, shape[0], annotations={nki_dim: "P"}):
@@ -166,7 +167,8 @@ def generate_unary_func(
bias_buffer = bias.source
# fmt: off
- @T.prim_func
+ # This fragment captures buffers and indices from its insertion scope.
+ @T.prim_func(check_well_formed=False)
def impl():
for b_loop in T.serial(0, b_extent):
with T.attr(0, "tensorized_nki_instruction", 1):
diff --git a/python/tvm/ir/json_compact.py b/python/tvm/ir/json_compact.py
index a7986a1491..e26aab947a 100644
--- a/python/tvm/ir/json_compact.py
+++ b/python/tvm/ir/json_compact.py
@@ -19,6 +19,11 @@
import json
_PRIM_TYPE_KEY_RENAMES = {
+ "tirx.BufferRegion": "ir.TensorRegion",
+ "tirx.SBlock": "s_tir.SBlock",
+ "tirx.SBlockRealize": "s_tir.SBlockRealize",
+ "tirx.MatchBufferRegion": "s_tir.MatchBufferRegion",
+ "tirx.TensorIntrin": "s_tir.TensorIntrin",
"tirx.StringImm": "ir.prim.StringImm",
"tirx.Cast": "ir.prim.Cast",
"tirx.Add": "ir.prim.Add",
@@ -119,7 +124,23 @@ def upgrade_json(json_str):
# compatible with the pre-unification Relax/TIRx schemas and with graphs
# written before the canonical Var field was renamed to `name`. Rewriting
# nodes in place preserves node indices and shared references.
- for node in data.get("nodes", []):
+ nodes = data.get("nodes", [])
+ buffer_region_type = None
+ for node in nodes:
+ if node.get("type") == "tirx.BufferRegion":
+ fields = node.get("data")
+ if not isinstance(fields, dict) or "buffer" not in fields:
+ raise ValueError("Legacy tirx.BufferRegion requires a buffer field")
+ fields["source"] = fields.pop("buffer")
+ # Typed BufferRegion already carries type/span. Before it became
+ # an Expr, it had only buffer/region; supply that form's defaults
+ # by appending a type node so existing graph indices stay intact.
+ if "ty" not in fields:
+ if buffer_region_type is None:
+ buffer_region_type = len(nodes)
+ nodes.append({"type": "tirx.BufferRegionType", "data": {"span": 0}})
+ fields["ty"] = buffer_region_type
+ fields.setdefault("span", 0)
node["type"] = _PRIM_TYPE_KEY_RENAMES.get(node.get("type"), node.get("type"))
if node.get("type") == "relax.expr.Var":
node["type"] = "ir.Var"
diff --git a/python/tvm/relax/analysis/analysis.py b/python/tvm/relax/analysis/analysis.py
index 7f0a6c2c39..c4ef71493a 100644
--- a/python/tvm/relax/analysis/analysis.py
+++ b/python/tvm/relax/analysis/analysis.py
@@ -29,7 +29,8 @@ from tvm import IRModule, tirx
from tvm.ir import Call, Type
from tvm.relax.expr import Binding, DataflowBlock, Expr, Function, GlobalVar, Var
from tvm.relax.type import FuncType
-from tvm.tirx import Buffer, IndexMap, PrimFunc, SBlock
+from tvm.s_tir import SBlock
+from tvm.tirx import Buffer, IndexMap, PrimFunc
from . import _ffi_api
diff --git a/python/tvm/s_tir/__init__.py b/python/tvm/s_tir/__init__.py
index 164dcc9901..7bb80b3419 100644
--- a/python/tvm/s_tir/__init__.py
+++ b/python/tvm/s_tir/__init__.py
@@ -18,7 +18,8 @@
# pylint: disable=invalid-name
"""S-TIR namespace for scheduable TensorIR"""
-from tvm.tirx.function import TensorIntrin
+from ._tensor_intrin import TensorIntrin
+from .stmt import MatchBufferRegion, SBlock, SBlockRealize
# dlight depends on compiler-only C++ functions (e.g. s_tir.schedule.GetSBlockRealize),
# so skip it in runtime-only builds.
diff --git a/python/tvm/s_tir/_tensor_intrin.py b/python/tvm/s_tir/_tensor_intrin.py
new file mode 100644
index 0000000000..00967c08ee
--- /dev/null
+++ b/python/tvm/s_tir/_tensor_intrin.py
@@ -0,0 +1,80 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Tensor intrinsics for schedulable TIR."""
+
+from typing import Optional
+
+import tvm_ffi
+
+from tvm.runtime import Object
+from tvm.tirx.function import PrimFunc
+
+from . import _ffi_api
+
+
+@tvm_ffi.register_object("s_tir.TensorIntrin")
+class TensorIntrin(Object):
+ """A tensor intrinsic.
+
+ Parameters
+ ----------
+ desc : PrimFunc
+ The function to describe the computation.
+
+ impl : PrimFunc
+ The function of the implementation for the execution.
+ """
+
+ def __init__(self, desc, impl):
+ self.__init_handle_by_constructor__(_ffi_api.TensorIntrin, desc, impl)
+
+ @staticmethod
+ def register(name: str, desc: PrimFunc, impl: PrimFunc, override: bool = False):
+ """Register a tensor intrinsic with its name.
+
+ Parameters
+ ----------
+ name : str
+ The name of the TensorIntrin to register.
+ desc : PrimFunc
+ The function to describe the computation.
+ impl : PrimFunc
+ The function of the implementation for the execution.
+ override: bool
+ Whether override existing intrinsic.
+ """
+ return _ffi_api.TensorIntrinRegister(name, TensorIntrin(desc, impl), override) # type: ignore
+
+ @staticmethod
+ def get(name: str, allow_missing: bool = False) -> Optional["TensorIntrin"]:
+ """Look up a tensor intrinsic by its name.
+
+ Parameters
+ ----------
+ name : str
+ The name of the TensorIntrin to look up.
+
+ allow_missing : bool
+ Whether to allow missing tensor intrin. If False, raise an error if the tensor intrin
+ doesn't exist.
+
+ Returns
+ -------
+ result : Optional[TensorIntrin]
+ The TensorIntrin with the specified name, or None if not found.
+ """
+ return _ffi_api.TensorIntrinGet(name, allow_missing) # pylint: type: ignore
diff --git a/python/tvm/s_tir/analysis/__init__.py b/python/tvm/s_tir/analysis/__init__.py
index 07021f092d..24732261a5 100644
--- a/python/tvm/s_tir/analysis/__init__.py
+++ b/python/tvm/s_tir/analysis/__init__.py
@@ -23,7 +23,7 @@ from typing import Optional, Union
import tvm
from tvm.ir import IRModule, TensorRegion
from tvm.tirx.expr import Var
-from tvm.tirx.stmt import SBlock
+from tvm.s_tir import SBlock
from tvm.tirx import Buffer, Stmt
from tvm.tirx.function import PrimFunc
@@ -38,7 +38,7 @@ def get_sblock_access_region(
Parameters
----------
- block: tvm.tirx.SBlock
+ block: tvm.s_tir.SBlock
The block in which we are detecting read/write regions.
buffer_var_map : Dict[Var, Buffer]
@@ -63,7 +63,7 @@ def get_sblock_read_write_region(
Parameters
----------
- block: tvm.tirx.SBlock
+ block: tvm.s_tir.SBlock
The block in which we are detecting read/write regions.
buffer_var_map : Dict[Var, Buffer]
@@ -214,3 +214,13 @@ def is_pure_function(func: PrimFunc) -> bool:
def assert_pure_function(func: PrimFunc) -> bool:
"""Asserts that the function is a pure function"""
return _ffi_api.is_pure_function(func, True) # type: ignore # pylint: disable=no-member
+
+
+def verify_well_formed(obj: PrimFunc | IRModule, assert_mode: bool = True) -> bool:
+ """Verify definitions, buffer loads and S-TIR block boundaries.
+
+ Modules may contain both S-TIR and ordinary PrimFuncs. Shared variable
+ identities are checked across function boundaries. Use the TIRX-specific
+ verifier separately for execution-scope restrictions on ordinary PrimFuncs.
+ """
+ return _ffi_api.VerifyWellFormed(obj, assert_mode)
diff --git a/python/tvm/s_tir/backend/adreno/pipeline.py b/python/tvm/s_tir/backend/adreno/pipeline.py
index 7dd65a4b30..1845f2f16e 100644
--- a/python/tvm/s_tir/backend/adreno/pipeline.py
+++ b/python/tvm/s_tir/backend/adreno/pipeline.py
@@ -44,7 +44,7 @@ def default_tir_pipeline():
s_tir.transform.LowerAutoCopy(),
s_tir.transform.UnifyThreadBinding(),
s_tir.transform.LowerMatchBuffer(),
- tirx.transform.StmtSimplify(),
+ s_tir.transform.StmtSimplify(),
s_tir.transform.InjectPermutedLayout(),
s_tir.transform.AnnotateIrregularLoop(),
s_tir.transform.InjectSoftwarePipeline(),
diff --git a/python/tvm/s_tir/dlight/analysis/common_analysis.py b/python/tvm/s_tir/dlight/analysis/common_analysis.py
index 254c30ee5f..70d14131a3 100644
--- a/python/tvm/s_tir/dlight/analysis/common_analysis.py
+++ b/python/tvm/s_tir/dlight/analysis/common_analysis.py
@@ -406,7 +406,7 @@ def get_root_block(sch: Schedule, func_name: str = "main") -> SBlockRV:
def collect_block_iter_vars_used_in_access_region(
- block: tirx.SBlock, region: list[ir.Range]
+ block: s_tir.SBlock, region: list[ir.Range]
) -> set[tirx.Var]:
"""Collect the block iter variables used in the access region of a buffer region."""
tir_vars = set()
@@ -428,7 +428,7 @@ def collect_vars_used_in_prim_expr(expr: tirx.Expr) -> set[tirx.Var]:
return tir_vars
-def detect_dominant_read(block: tirx.SBlock) -> tirx.Expr:
+def detect_dominant_read(block: s_tir.SBlock) -> tirx.Expr:
"""Detect the dominant read indices in the block."""
dominant_read = None
num_read_iters = -1
diff --git a/python/tvm/s_tir/dlight/analysis/gemv.py b/python/tvm/s_tir/dlight/analysis/gemv.py
index 6c31c2057f..33a83b38c7 100644
--- a/python/tvm/s_tir/dlight/analysis/gemv.py
+++ b/python/tvm/s_tir/dlight/analysis/gemv.py
@@ -28,7 +28,7 @@ from .common_analysis import (
)
-def get_reduction_expr(block: tirx.SBlock) -> tirx.Expr | None:
+def get_reduction_expr(block: s_tir.SBlock) -> tirx.Expr | None:
"""Extracts the reduction expression from a TIR block.
This function checks whether the given TIR block follows a reduction pattern
@@ -36,7 +36,7 @@ def get_reduction_expr(block: tirx.SBlock) -> tirx.Expr | None:
Parameters:
----------
- block : tirx.SBlock
+ block : s_tir.SBlock
The TIR block to analyze.
Returns:
@@ -106,7 +106,7 @@ def normalize(
block_info: SBlockInfo,
) -> bool | None:
"""Normalize the main block."""
- block_stmt: tirx.SBlock = sch.get(block_info.block_rv)
+ block_stmt: s_tir.SBlock = sch.get(block_info.block_rv)
access = arith.normalize_to_iter_sum(
detect_dominant_read(block_stmt),
input_iters={i.var: i.dom for i in block_stmt.iter_vars},
diff --git a/python/tvm/s_tir/dlight/gpu/low_batch_gemv.py b/python/tvm/s_tir/dlight/gpu/low_batch_gemv.py
index c997e32355..2732fef989 100644
--- a/python/tvm/s_tir/dlight/gpu/low_batch_gemv.py
+++ b/python/tvm/s_tir/dlight/gpu/low_batch_gemv.py
@@ -37,7 +37,7 @@ from ..base import auto_vectorize, get_bytes, get_extent, try_inline_contiguous_
from .base import GPUScheduleRule
-def _get_reduction_expr(block: tirx.SBlock) -> tirx.Expr | None:
+def _get_reduction_expr(block: s_tir.SBlock) -> tirx.Expr | None:
# Detect and return `Y` in `X[...] = X[...] + Y`
buffer_store = block.body
if not isinstance(buffer_store, tirx.BufferStore):
@@ -53,7 +53,7 @@ def _get_reduction_expr(block: tirx.SBlock) -> tirx.Expr | None:
return buffer_store.value.b
-def _has_pad_einsum_compatible_access(block: tirx.SBlock) -> bool:
+def _has_pad_einsum_compatible_access(block: s_tir.SBlock) -> bool:
"""Check the point-access restriction required by ``Schedule.pad_einsum``."""
return all(
isinstance(dim.extent, tirx.IntImm)
@@ -126,7 +126,7 @@ def is_gemv(sch: s_tir.Schedule, block_info: SBlockInfo) -> list[tirx.Buffer] |
return ret if 0 < len(ret) < len(block_stmt.reads) else None
-def detect_dominant_read(block: tirx.SBlock, const_iter_vars: set[tirx.Var]) -> tirx.Expr:
+def detect_dominant_read(block: s_tir.SBlock, const_iter_vars: set[tirx.Var]) -> tirx.Expr:
"""Detect the dominant read indices in the block."""
dominant_read = None
num_read_iters = -1
@@ -148,7 +148,7 @@ def normalize(
block_info: SBlockInfo,
) -> bool | None:
"""Normalize the main block."""
- block_stmt: tirx.SBlock = sch.get(block_info.block_rv)
+ block_stmt: s_tir.SBlock = sch.get(block_info.block_rv)
const_iter_vars = set(
iter_var.var
for iter_var in block_stmt.iter_vars
diff --git a/python/tvm/s_tir/dlight/gpu/matmul.py b/python/tvm/s_tir/dlight/gpu/matmul.py
index 48141e182f..b4bac8a2c7 100644
--- a/python/tvm/s_tir/dlight/gpu/matmul.py
+++ b/python/tvm/s_tir/dlight/gpu/matmul.py
@@ -166,12 +166,12 @@ def make_iter_fusion_index_map(
return tirx.IndexMap(input_iters, final_indices, None)
-def detect_iter_traits(block: tirx.SBlock) -> tuple[list[IterTrait]] | None:
+def detect_iter_traits(block: s_tir.SBlock) -> tuple[list[IterTrait]] | None:
"""Detect iter traits based on the pattern C[S, I, J] += A[S, I, K] * B[S, J, K]
Parameters
----------
- block : tirx.SBlock
+ block : s_tir.SBlock
The block to be analyzed
Returns
@@ -236,12 +236,12 @@ def detect_iter_traits(block: tirx.SBlock) -> tuple[list[IterTrait]] | None:
return A_traits, B_traits, C_traits, block_traits
-def get_index_map(block: tirx.SBlock) -> tuple[tirx.IndexMap, ...] | None:
+def get_index_map(block: s_tir.SBlock) -> tuple[tirx.IndexMap, ...] | None:
"""Get index maps for the block
Parameters
----------
- block : tirx.SBlock
+ block : s_tir.SBlock
The block to be analyzed
Returns
@@ -326,7 +326,7 @@ def get_reduction_blocks(sch, blocks) -> bool:
return reduction_blocks
-def get_in_out_dtypes(block: tirx.SBlock) -> tuple[str]:
+def get_in_out_dtypes(block: s_tir.SBlock) -> tuple[str]:
"""
Detect In/Out data types for the given block based on the analysis if read/write buffers.
"""
diff --git a/python/tvm/s_tir/dlight/gpu/reduction.py b/python/tvm/s_tir/dlight/gpu/reduction.py
index ff22cdfe98..6c010ade69 100644
--- a/python/tvm/s_tir/dlight/gpu/reduction.py
+++ b/python/tvm/s_tir/dlight/gpu/reduction.py
@@ -34,7 +34,7 @@ from ..base import suggest_threads_per_block, try_inline_contiguous_spatial
from .base import GPUScheduleRule
-def _get_reduction_expr(block: tirx.SBlock) -> tirx.Expr | None:
+def _get_reduction_expr(block: s_tir.SBlock) -> tirx.Expr | None:
# Detect and return `Y` in `X[...] = X[...] + Y`
buffer_store = block.body
if not isinstance(buffer_store, tirx.BufferStore):
diff --git a/python/tvm/s_tir/dlight/gpu/rmsnorm.py b/python/tvm/s_tir/dlight/gpu/rmsnorm.py
index 82888fdcf7..8245c6930a 100644
--- a/python/tvm/s_tir/dlight/gpu/rmsnorm.py
+++ b/python/tvm/s_tir/dlight/gpu/rmsnorm.py
@@ -20,8 +20,9 @@
import tvm
from tvm import tirx
from tvm.ir import Call, TensorLoad
+from tvm.s_tir import SBlock
from tvm.target import Target
-from tvm.tirx import BufferStore, SBlock
+from tvm.tirx import BufferStore
from tvm.tirx.expr import Cast
from ..base import ScheduleRule
diff --git a/python/tvm/s_tir/pipeline.py b/python/tvm/s_tir/pipeline.py
index 81a9fad21f..283ec8fc92 100644
--- a/python/tvm/s_tir/pipeline.py
+++ b/python/tvm/s_tir/pipeline.py
@@ -45,7 +45,7 @@ def default_s_tir_pipeline():
s_tir.transform.LowerAutoCopy(),
s_tir.transform.UnifyThreadBinding(),
s_tir.transform.LowerMatchBuffer(),
- tirx.transform.StmtSimplify(),
+ s_tir.transform.StmtSimplify(),
s_tir.transform.InjectPermutedLayout(),
s_tir.transform.AnnotateIrregularLoop(),
s_tir.transform.InjectSoftwarePipeline(),
diff --git a/python/tvm/s_tir/sblock_dependence_info.py b/python/tvm/s_tir/sblock_dependence_info.py
index 0376d2ca45..9d6bd542c6 100644
--- a/python/tvm/s_tir/sblock_dependence_info.py
+++ b/python/tvm/s_tir/sblock_dependence_info.py
@@ -21,7 +21,8 @@ from tvm_ffi import register_object
from tvm.ir.module import IRModule
from tvm.runtime import Object
-from tvm.tirx import PrimFunc, SBlock
+from tvm.s_tir import SBlock
+from tvm.tirx import PrimFunc
from . import _ffi_api
from .sblock_scope import SBlockScope, StmtSRef
diff --git a/python/tvm/s_tir/sblock_scope.py b/python/tvm/s_tir/sblock_scope.py
index bd5bb369fa..1b37b8c288 100644
--- a/python/tvm/s_tir/sblock_scope.py
+++ b/python/tvm/s_tir/sblock_scope.py
@@ -22,7 +22,8 @@ from typing import Optional
from tvm_ffi import register_object
from tvm.runtime import Object
-from tvm.tirx import For, SBlock
+from tvm.s_tir import SBlock
+from tvm.tirx import For
from . import _ffi_api
diff --git a/python/tvm/s_tir/schedule/schedule.py b/python/tvm/s_tir/schedule/schedule.py
index e97195ad54..42ac24df3d 100644
--- a/python/tvm/s_tir/schedule/schedule.py
+++ b/python/tvm/s_tir/schedule/schedule.py
@@ -25,7 +25,8 @@ from tvm_ffi import register_object as _register_object
from tvm.error import register_error
from tvm.ir import Expr, GlobalVar, IRModule, is_prim_expr
from tvm.runtime import DataTypeCode, Object
-from tvm.tirx import Buffer, FloatImm, For, IntImm, PrimFunc, SBlock, is_buffer_var
+from tvm.s_tir import SBlock
+from tvm.tirx import Buffer, FloatImm, For, IntImm, PrimFunc, is_buffer_var
from tvm.tirx.function import IndexMap
from . import _ffi_api
@@ -3043,7 +3044,7 @@ class Schedule(Object):
)
)
- tirx.TensorIntrin.register("test_mma_intrin", mma_desc, mma_intrin)
+ tvm.s_tir.TensorIntrin.register("test_mma_intrin", mma_desc, mma_intrin)
Create the schedule and do tensorize:
diff --git a/python/tvm/s_tir/schedule/state.py b/python/tvm/s_tir/schedule/state.py
index 6702930faa..98396fef6b 100644
--- a/python/tvm/s_tir/schedule/state.py
+++ b/python/tvm/s_tir/schedule/state.py
@@ -24,7 +24,8 @@ from tvm_ffi import register_object
from tvm.ir import IRModule
from tvm.runtime import Object
-from tvm.tirx import For, PrimFunc, SBlock, SBlockRealize
+from tvm.s_tir import SBlock, SBlockRealize
+from tvm.tirx import For, PrimFunc
from ..sblock_scope import SBlockScope, StmtSRef
from . import _ffi_api
diff --git a/python/tvm/s_tir/stmt.py b/python/tvm/s_tir/stmt.py
new file mode 100644
index 0000000000..3c84f5a59c
--- /dev/null
+++ b/python/tvm/s_tir/stmt.py
@@ -0,0 +1,179 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Schedulable TensorIR statement nodes."""
+
+from collections.abc import Mapping
+
+import tvm_ffi
+
+from tvm.ir import Expr, Span, TensorRegion
+from tvm.runtime import Object, Scriptable, const
+from tvm.tirx.buffer import Buffer
+from tvm.tirx.expr import IterVar
+from tvm.tirx.stmt import Stmt, _normalize_legacy_stmt
+
+from . import _ffi_api
+
+
+@tvm_ffi.register_object("s_tir.MatchBufferRegion")
+class MatchBufferRegion(Object, Scriptable):
+ """MatchBufferRegion node.
+
+ Parameters
+ ----------
+ buffer : Buffer
+ The target buffer
+
+ source : TensorRegion
+ The region of source buffer
+ """
+
+ buffer: Buffer
+ source: TensorRegion
+
+ def __init__(self, buffer: Buffer, source: TensorRegion) -> None:
+ self.__init_handle_by_constructor__(
+ _ffi_api.MatchBufferRegion,
+ buffer,
+ source, # type: ignore
+ )
+
+
+@tvm_ffi.register_object("s_tir.SBlock")
+class SBlock(Stmt):
+ """SBlock node.
+
+ Parameters
+ ----------
+ iter_vars : List[IterVar]
+ The block Variable.
+
+ reads : List[TensorRegion]
+ The read buffer regions of the block.
+
+ writes: List[TensorRegion]
+ The write buffer regions of the block.
+
+ name_hint: str
+ the name_hint of the block.
+
+ body: Stmt
+ The body of the block.
+
+ init: Optional[Stmt]
+ The init block of the reduction block
+
+ alloc_buffers: Optional[list[Buffer]]
+ The buffer allocations
+
+ match_buffers: Optional[List[MatchBufferRegion]]
+ The subregion buffer match
+
+ annotations: Optional[Mapping[str, Object]]
+ Additional annotation hints.
+
+ span : Optional[Span]
+ The location of this block in the source code.
+ """
+
+ iter_vars: list[IterVar]
+ reads: list[TensorRegion]
+ writes: list[TensorRegion]
+ name_hint: str
+ body: Stmt
+ init: Stmt | None
+ alloc_buffers: list[Buffer]
+ match_buffers: list[MatchBufferRegion]
+ annotations: Mapping[str, Object]
+ span: Span | None
+
+ def __init__(
+ self,
+ iter_vars: list[IterVar],
+ reads: list[TensorRegion],
+ writes: list[TensorRegion],
+ name_hint: str,
+ body: Stmt,
+ init: Stmt | None = None,
+ alloc_buffers: list[Buffer] | None = None,
+ match_buffers: list[MatchBufferRegion] | None = None,
+ annotations: Mapping[str, Object] | None = None,
+ span: Span | None = None,
+ ) -> None:
+ if alloc_buffers is None:
+ alloc_buffers = []
+ if match_buffers is None:
+ match_buffers = []
+ if annotations is None:
+ annotations = {}
+ body = _normalize_legacy_stmt(body)
+ init = _normalize_legacy_stmt(init)
+ self.__init_handle_by_constructor__(
+ _ffi_api.SBlock, # type: ignore
+ iter_vars,
+ reads,
+ writes,
+ name_hint,
+ body,
+ init,
+ alloc_buffers,
+ match_buffers,
+ annotations,
+ span,
+ ) # type: ignore
+
+
+@tvm_ffi.register_object("s_tir.SBlockRealize")
+class SBlockRealize(Stmt):
+ """SBlockRealize node.
+
+ Parameters
+ ----------
+ iter_values : List[Expr]
+ The binding values of the block var.
+
+ predicate : Union[Expr, bool]
+ The predicate of the block.
+
+ block : SBlock
+ The block to realize
+
+ span : Optional[Span]
+ The location of this block_realize in the source code.
+ """
+
+ iter_values: list[Expr]
+ predicate: Expr
+ block: SBlock
+ span: Span | None
+
+ def __init__(
+ self,
+ iter_values: list[Expr],
+ predicate: Expr | bool,
+ block: SBlock,
+ span: Span | None = None,
+ ) -> None:
+ if isinstance(predicate, bool):
+ predicate = const(predicate, "bool")
+ self.__init_handle_by_constructor__(
+ _ffi_api.SBlockRealize, # type: ignore
+ iter_values,
+ predicate,
+ block,
+ span,
+ ) # type: ignore
diff --git a/python/tvm/s_tir/tensor_intrin/arm_cpu.py b/python/tvm/s_tir/tensor_intrin/arm_cpu.py
index 984830b957..b5fbe3a0ac 100644
--- a/python/tvm/s_tir/tensor_intrin/arm_cpu.py
+++ b/python/tvm/s_tir/tensor_intrin/arm_cpu.py
@@ -487,7 +487,7 @@ def get_transpose_interleave_intrin_name(in_dtype, out_dtype, extent_cols, exten
sme_transpose_interleave_intrin_name = (
ARM_SME_2SVLx2SVL_FP32_TRANSPOSE_INTERLEAVE + f"_{extent_cols}_{extent_rows}"
)
- tirx.TensorIntrin.register(
+ TensorIntrin.register(
sme_transpose_interleave_intrin_name,
*get_sme_transpose_interleave_2svlx2svl_fp32_intrin(extent_cols, extent_rows),
override=True,
diff --git a/python/tvm/s_tir/tensor_intrin/cuda.py b/python/tvm/s_tir/tensor_intrin/cuda.py
index ae3ad7cd55..fb28a0c359 100644
--- a/python/tvm/s_tir/tensor_intrin/cuda.py
+++ b/python/tvm/s_tir/tensor_intrin/cuda.py
@@ -23,8 +23,9 @@ from typing import Literal
from tvm_ffi import register_global_func
from tvm.runtime import convert
+from tvm.s_tir import TensorIntrin
from tvm.script import tirx as T
-from tvm.tirx import Cast, IntImm, TensorIntrin
+from tvm.tirx import Cast, IntImm
from tvm.tirx.function import PrimFunc
diff --git a/python/tvm/s_tir/tensor_intrin/metal.py b/python/tvm/s_tir/tensor_intrin/metal.py
index 5b19793ddf..1750338044 100644
--- a/python/tvm/s_tir/tensor_intrin/metal.py
+++ b/python/tvm/s_tir/tensor_intrin/metal.py
@@ -19,8 +19,9 @@
from typing import Literal
+from tvm.s_tir import TensorIntrin
from tvm.script import tirx as T
-from tvm.tirx import Buffer, Expr, PrimFunc, TensorIntrin
+from tvm.tirx import Buffer, Expr, PrimFunc
######## simdgroup matrix intrinsics ########
diff --git a/python/tvm/s_tir/transform/transform.py b/python/tvm/s_tir/transform/transform.py
index 0a56fe5bea..6bbaeb0f33 100644
--- a/python/tvm/s_tir/transform/transform.py
+++ b/python/tvm/s_tir/transform/transform.py
@@ -22,6 +22,16 @@ import tvm_ffi as _ffi
from . import _ffi_api
+def ConvertSSA():
+ """De-duplicate definitions, including schedulable block iterators, across PrimFuncs."""
+ return _ffi_api.ConvertSSA()
+
+
+def StmtSimplify():
+ """Simplify schedulable TIR with block constraints and tirx.StmtSimplify options."""
+ return _ffi_api.StmtSimplify()
+
+
def CanonicalizeLoop():
"""Canonicalize the loop to start from zero and use trivial step
diff --git a/python/tvm/script/parser/core/entry.py b/python/tvm/script/parser/core/entry.py
index ec5b5d7ab1..afd155f3be 100644
--- a/python/tvm/script/parser/core/entry.py
+++ b/python/tvm/script/parser/core/entry.py
@@ -94,6 +94,11 @@ def parse(
check_well_formed : bool
Whether to check well-formedness after parsing.
+ s_tir : bool
+ Compatibility argument accepted by parse/from_source. It no longer
+ selects verification: each PrimFunc's s_tir attribute determines its
+ dialect checks, and common well-formedness checks cover the full module.
+
absent_params : Optional[Dict[str, None]]
Function parameters removed by a compile-time specialization. The
dialect-specific function parser decides how to bind these names.
@@ -142,10 +147,12 @@ def parse(
parser.report_error(source_ast, err=WELL_FORMED_ERROR_MESSAGE)
try:
- if s_tir:
- tvm.tirx.analysis.verify_well_formed(check_ret)
- else:
- tvm.tirx.analysis.verify_tirx_well_formed(check_ret)
+ # Keep shared-definition checks across both dialects, then apply
+ # ordinary TIRX execution restrictions only to its own functions.
+ tvm.s_tir.analysis.verify_well_formed(check_ret)
+ for func in check_ret.functions.values():
+ if isinstance(func, tvm.tirx.PrimFunc) and not func.attrs.get("s_tir", False):
+ tvm.tirx.analysis.verify_tirx_well_formed(func)
except Exception as err: # pylint: disable=broad-exception-caught
parser.report_error(
source_ast,
diff --git a/python/tvm/script/parser/ir/entry.py b/python/tvm/script/parser/ir/entry.py
index 4cfe60b77c..5191ce0f87 100644
--- a/python/tvm/script/parser/ir/entry.py
+++ b/python/tvm/script/parser/ir/entry.py
@@ -61,7 +61,7 @@ def ir_module(
extra_vars = utils.inspect_class_capture(mod)
# Resolve closure variables hidden by PEP 563 (annotation-only names)
utils.resolve_closure_vars(mod, extra_vars, outer_stack)
- m = parse(mod, extra_vars, check_well_formed=check_well_formed, s_tir=s_tir)
+ m = parse(mod, extra_vars, check_well_formed=check_well_formed)
if base_py_module_inherited:
# Lazy import: tvm.relax cannot be imported at module level in tvm.script.parser
diff --git a/python/tvm/tirx/__init__.py b/python/tvm/tirx/__init__.py
index af4b7161b0..e5a08aed0c 100644
--- a/python/tvm/tirx/__init__.py
+++ b/python/tvm/tirx/__init__.py
@@ -52,11 +52,11 @@ from .stmt import BufferStore, AllocBuffer, AttrStmt, DeclBuffer
from .stmt import SeqStmt
from .stmt import IfThenElse, Evaluate, stmt_seq, stmt_list
-from .stmt import BufferRegion, BufferRegionType, MatchBufferRegion, SBlock, SBlockRealize
+from .stmt import BufferRegion, BufferRegionType
from .stmt import ScopeIdDefStmt
from .tile_primitive import DispatchContext, LambdaExpr, TilePrimitiveCall
-from .function import PrimFunc, TensorIntrin, IndexMap
+from .function import PrimFunc, IndexMap
from .op import call_packed_lowered, call_cpacked_lowered, register_intrin_lowering
from .op import call_packed, call_cpacked, call_intrin, call_pure_extern, call_extern
diff --git a/python/tvm/tirx/analysis/analysis.py b/python/tvm/tirx/analysis/analysis.py
index f48710a683..faf2bf0b47 100644
--- a/python/tvm/tirx/analysis/analysis.py
+++ b/python/tvm/tirx/analysis/analysis.py
@@ -80,8 +80,9 @@ def undefined_vars(node: Stmt | Expr, defs: list[Var] | None = None) -> list[Var
def verify_well_formed(obj: PrimFunc | IRModule, assert_mode: bool = True) -> bool:
- """Verify if the given TIR is well-formed. The verification includes:
- - Check if expressions not contain vars that is defined outside the block.
+ """Verify definitions and buffer-load types in ordinary TIRX.
+
+ Use ``tvm.s_tir.analysis.verify_well_formed`` for schedulable blocks.
Parameters
----------
diff --git a/python/tvm/tirx/function.py b/python/tvm/tirx/function.py
index 38c701075f..d6d58c2d79 100644
--- a/python/tvm/tirx/function.py
+++ b/python/tvm/tirx/function.py
@@ -119,19 +119,17 @@ class PrimFunc(BaseFunc, Scriptable):
Examples
--------
- We can define a Meta TIR function with symbolic shape:
+ We can define a TIRX function with symbolic shape:
.. code-block:: python
- @T.prim_func(s_tir=True)
+ @T.prim_func
def mem_copy(a: T.handle, b: T.handle, m: T.int32, n: T.int32) -> None:
A = T.match_buffer(a, (m, n), "float32")
B = T.match_buffer(b, (m, n), "float32")
for i, j in T.grid(m, n):
- with T.sblock():
- vi, vj = T.axis.remap("SS", [i, j])
- B[vi, vj] = A[vi, vj]
+ B[i, j] = A[i, j]
Then we can make it specialized with given shapes or buffers.
@@ -146,15 +144,13 @@ class PrimFunc(BaseFunc, Scriptable):
.. code-block:: python
- @T.prim_func(s_tir=True)
+ @T.prim_func
def mem_copy_16_16(a: T.handle, b: T.handle) -> None:
A = T.match_buffer(a, (16, 16), "float32")
B = T.match_buffer(b, (16, 16), "float32")
for i, j in T.grid(16, 16):
- with T.sblock():
- vi, vj = T.axis.remap("SS", [i, j])
- B[vi, vj] = A[vi, vj]
+ B[i, j] = A[i, j]
Returns
-------
@@ -164,60 +160,6 @@ class PrimFunc(BaseFunc, Scriptable):
return _ffi_api.Specialize(self, param_map) # type: ignore
-@tvm_ffi.register_object("tirx.TensorIntrin")
-class TensorIntrin(Object):
- """A tensor intrinsic.
-
- Parameters
- ----------
- desc : PrimFunc
- The function to describe the computation.
-
- impl : PrimFunc
- The function of the implementation for the execution.
- """
-
- def __init__(self, desc, impl):
- self.__init_handle_by_constructor__(_ffi_api.TensorIntrin, desc, impl)
-
- @staticmethod
- def register(name: str, desc: PrimFunc, impl: PrimFunc, override: bool = False):
- """Register a tensor intrinsic with its name.
-
- Parameters
- ----------
- name : str
- The name of the TensorIntrin to register.
- desc : PrimFunc
- The function to describe the computation.
- impl : PrimFunc
- The function of the implementation for the execution.
- override: bool
- Whether override existing intrinsic.
- """
- return _ffi_api.TensorIntrinRegister(name, TensorIntrin(desc, impl), override) # type: ignore
-
- @staticmethod
- def get(name: str, allow_missing: bool = False) -> Optional["TensorIntrin"]:
- """Look up a tensor intrinsic by its name.
-
- Parameters
- ----------
- name : str
- The name of the TensorIntrin to look up.
-
- allow_missing : bool
- Whether to allow missing tensor intrin. If False, raise an error if the tensor intrin
- doesn't exist.
-
- Returns
- -------
- result : Optional[TensorIntrin]
- The TensorIntrin with the specified name, or None if not found.
- """
- return _ffi_api.TensorIntrinGet(name, allow_missing) # pylint: type: ignore
-
-
@tvm_ffi.register_object("tirx.IndexMap")
class IndexMap(Object):
"""A mapping from multi-dimensional indices to another set of multi-dimensional indices
diff --git a/python/tvm/tirx/script/parser/entry.py b/python/tvm/tirx/script/parser/entry.py
index d708aa621a..03bcc9d643 100644
--- a/python/tvm/tirx/script/parser/entry.py
+++ b/python/tvm/tirx/script/parser/entry.py
@@ -70,7 +70,7 @@ def prim_func(
return func
extra_vars = utils.inspect_function_capture(func)
utils.resolve_closure_vars(func, extra_vars, outer_stack)
- f = parse(func, extra_vars, check_well_formed=check_well_formed, s_tir=s_tir)
+ f = parse(func, extra_vars, check_well_formed=check_well_formed)
setattr(f, "__name__", func.__name__)
return f
@@ -335,7 +335,6 @@ class TIRJit:
self.func,
extra_vars,
check_well_formed=self.check_well_formed,
- s_tir=self.is_stir,
absent_params=absent_params,
)
setattr(prim_func, "__name__", self.func.__name__)
diff --git a/python/tvm/tirx/stmt.py b/python/tvm/tirx/stmt.py
index 5d15b227af..f609d4439f 100644
--- a/python/tvm/tirx/stmt.py
+++ b/python/tvm/tirx/stmt.py
@@ -34,7 +34,7 @@ from typing import Any
import tvm_ffi
from tvm.ir import Expr, Range, Span, TensorRegion, Type
-from tvm.runtime import Object, Scriptable, const
+from tvm.runtime import Object, Scriptable
from . import _ffi_api
from .buffer import Buffer
@@ -628,156 +628,6 @@ def BufferRegion(buffer: Buffer, region: list[Range]) -> TensorRegion:
return _ffi_api.BufferRegion(buffer, region)
-@tvm_ffi.register_object("tirx.MatchBufferRegion")
-class MatchBufferRegion(Object, Scriptable):
- """MatchBufferRegion node.
-
- Parameters
- ----------
- buffer : Buffer
- The target buffer
-
- source : TensorRegion
- The region of source buffer
- """
-
- buffer: Buffer
- source: TensorRegion
-
- def __init__(self, buffer: Buffer, source: TensorRegion) -> None:
- self.__init_handle_by_constructor__(
- _ffi_api.MatchBufferRegion,
- buffer,
- source, # type: ignore
- )
-
-
-@tvm_ffi.register_object("tirx.SBlock")
-class SBlock(Stmt):
- """SBlock node.
-
- Parameters
- ----------
- iter_vars : List[IterVar]
- The block Variable.
-
- reads : List[TensorRegion]
- The read buffer regions of the block.
-
- writes: List[TensorRegion]
- The write buffer regions of the block.
-
- name_hint: str
- the name_hint of the block.
-
- body: Stmt
- The body of the block.
-
- init: Optional[Stmt]
- The init block of the reduction block
-
- alloc_buffers: Optional[list[Buffer]]
- The buffer allocations
-
- match_buffers: Optional[List[MatchBufferRegion]]
- The subregion buffer match
-
- annotations: Optional[Mapping[str, Object]]
- Additional annotation hints.
-
- span : Optional[Span]
- The location of this block in the source code.
- """
-
- iter_vars: list[IterVar]
- reads: list[TensorRegion]
- writes: list[TensorRegion]
- name_hint: str
- body: Stmt
- init: Stmt | None
- alloc_buffers: list[Buffer]
- match_buffers: list[MatchBufferRegion]
- annotations: Mapping[str, Object]
- span: Span | None
-
- def __init__(
- self,
- iter_vars: list[IterVar],
- reads: list[TensorRegion],
- writes: list[TensorRegion],
- name_hint: str,
- body: Stmt,
- init: Stmt | None = None,
- alloc_buffers: list[Buffer] | None = None,
- match_buffers: list[MatchBufferRegion] | None = None,
- annotations: Mapping[str, Object] | None = None,
- span: Span | None = None,
- ) -> None:
- if alloc_buffers is None:
- alloc_buffers = []
- if match_buffers is None:
- match_buffers = []
- if annotations is None:
- annotations = {}
- body = _normalize_legacy_stmt(body)
- init = _normalize_legacy_stmt(init)
- self.__init_handle_by_constructor__(
- _ffi_api.SBlock, # type: ignore
- iter_vars,
- reads,
- writes,
- name_hint,
- body,
- init,
- alloc_buffers,
- match_buffers,
- annotations,
- span,
- ) # type: ignore
-
-
-@tvm_ffi.register_object("tirx.SBlockRealize")
-class SBlockRealize(Stmt):
- """SBlockRealize node.
-
- Parameters
- ----------
- iter_values : List[Expr]
- The binding values of the block var.
-
- predicate : Union[Expr, bool]
- The predicate of the block.
-
- block : SBlock
- The block to realize
-
- span : Optional[Span]
- The location of this block_realize in the source code.
- """
-
- iter_values: list[Expr]
- predicate: Expr
- block: SBlock
- span: Span | None
-
- def __init__(
- self,
- iter_values: list[Expr],
- predicate: Expr | bool,
- block: SBlock,
- span: Span | None = None,
- ) -> None:
- if isinstance(predicate, bool):
- predicate = const(predicate, "bool")
- self.__init_handle_by_constructor__(
- _ffi_api.SBlockRealize, # type: ignore
- iter_values,
- predicate,
- block,
- span,
- ) # type: ignore
-
-
@tvm_ffi.register_object("tirx.ScopeIdDefStmt")
class ScopeIdDefStmt(Stmt):
"""ScopeIdDefStmt node.
diff --git a/src/relax/analysis/layout_transformation.cc b/src/relax/analysis/layout_transformation.cc
index 92cf537bee..38af950dd2 100644
--- a/src/relax/analysis/layout_transformation.cc
+++ b/src/relax/analysis/layout_transformation.cc
@@ -28,9 +28,10 @@
#include <tvm/ffi/reflection/registry.h>
#include <tvm/relax/analysis.h>
#include <tvm/runtime/logging.h>
+#include <tvm/s_tir/stmt.h>
+#include <tvm/s_tir/stmt_functor.h>
#include <tvm/tirx/analysis.h>
#include <tvm/tirx/index_map.h>
-#include <tvm/tirx/stmt_functor.h>
namespace tvm {
namespace relax {
@@ -59,7 +60,7 @@ static bool IsBijectiveAffine(const IndexMap& m, const ffi::Array<Range>& ranges
* are used in it. This is important to get which spatial iterators are accessed in each index
* of buffer access.
*/
-class IndexAnalyzer : public tirx::StmtExprVisitor {
+class IndexAnalyzer : public s_tir::StmtExprVisitor {
public:
ffi::Array<tirx::Var> Analyze(const arith::IterSumExpr& expr) {
Visit(expr);
@@ -79,7 +80,7 @@ class IndexAnalyzer : public tirx::StmtExprVisitor {
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Visit(op->extent));
return Visit(op->scale);
}
- return tirx::StmtExprVisitor::Visit(value);
+ return s_tir::StmtExprVisitor::Visit(value);
}
ffi::Optional<VisitInterrupt> VisitIterMark(const arith::IterMark& op) {
@@ -323,9 +324,9 @@ static ffi::Optional<IndexMap> InferLayoutTransformation(const SpatialLayout& sr
* 2. Expects write buffer access to be affine and only use spatial iterators of the block.
* 3. Proposes transformations to a read buffer if all access to it are affine.
*/
-class BlockAnalyzer : public StmtExprVisitor {
+class BlockAnalyzer : public s_tir::StmtExprVisitor {
public:
- explicit BlockAnalyzer(const SBlock& block,
+ explicit BlockAnalyzer(const s_tir::SBlock& block,
const ffi::Map<BufferVar, IndexMap>& transformation_cache,
IndexMap write_transformation)
: can_transform_block_(true),
@@ -482,7 +483,7 @@ class BlockAnalyzer : public StmtExprVisitor {
}
}
- ffi::Optional<VisitInterrupt> Visit_(const SBlockNode* op) final {
+ ffi::Optional<VisitInterrupt> Visit_(const s_tir::SBlockNode* op) final {
// Blocks with nested blocks cannot be handled yet.
LOG(WARNING) << "[LayoutInference] Nested blocks are not supported for layout inference yet";
can_transform_block_ = false;
@@ -490,7 +491,7 @@ class BlockAnalyzer : public StmtExprVisitor {
return std::nullopt;
}
ffi::Optional<VisitInterrupt> Visit_(const BufferStoreNode* op) final {
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(op));
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(s_tir::StmtExprVisitor::Visit_(op));
BufferAccessInfo& access_info = buffer_access_info_[op->buffer];
@@ -551,7 +552,7 @@ class BlockAnalyzer : public StmtExprVisitor {
ffi::Map<PrimVar, Range> spatial_dom_;
arith::Analyzer arith_analyzer_;
- SBlock block_;
+ s_tir::SBlock block_;
IndexMap block_transformation_;
ffi::Map<BufferVar, IndexMap> read_buffer_transformations_;
@@ -568,7 +569,7 @@ class BlockAnalyzer : public StmtExprVisitor {
* possible that the PrimFunc is too complex for analysis. In such a case, no transformations are
* proposed.
*/
-class PrimFuncAnalyzer : public StmtExprVisitor {
+class PrimFuncAnalyzer : public s_tir::StmtExprVisitor {
public:
explicit PrimFuncAnalyzer(const PrimFunc& func, ffi::Array<IndexMap> write_transformations) {
TVM_FFI_ICHECK_LE(write_transformations.size(), func->params.size())
@@ -585,8 +586,8 @@ class PrimFuncAnalyzer : public StmtExprVisitor {
buffer_transformation_cache_.Set(param_buf.value(), write_transformations[i]);
}
}
- ffi::Map<SBlock, ffi::Map<ffi::ObjectRef, IndexMap>> GetSuggestedTransforms() {
- ffi::Map<SBlock, ffi::Map<ffi::ObjectRef, IndexMap>> result;
+ ffi::Map<s_tir::SBlock, ffi::Map<ffi::ObjectRef, IndexMap>> GetSuggestedTransforms() {
+ ffi::Map<s_tir::SBlock, ffi::Map<ffi::ObjectRef, IndexMap>> result;
for (const auto& [block, index_map] : block_transformations_) {
ffi::Map<ffi::ObjectRef, IndexMap> block_transformations;
block_transformations.Set(block, index_map);
@@ -599,13 +600,13 @@ class PrimFuncAnalyzer : public StmtExprVisitor {
}
private:
- ffi::Optional<VisitInterrupt> Visit_(const SBlockNode* op) final {
+ ffi::Optional<VisitInterrupt> Visit_(const s_tir::SBlockNode* op) final {
if (op->name_hint == "root") {
// Skip the root block
- return StmtExprVisitor::Visit_(op);
+ return s_tir::StmtExprVisitor::Visit_(op);
}
- SBlock block = ffi::GetRef<SBlock>(op);
+ s_tir::SBlock block = ffi::GetRef<s_tir::SBlock>(op);
// Get block write buffer transformation.
if (block->writes.size() != 1) return std::nullopt;
auto write_buffer = block->writes[0]->source.as_or_throw<tvm::tirx::BufferVar>();
@@ -632,12 +633,12 @@ class PrimFuncAnalyzer : public StmtExprVisitor {
private:
ffi::Map<BufferVar, IndexMap> buffer_transformation_cache_;
- ffi::Map<SBlock, IndexMap> block_transformations_;
- std::unordered_map<SBlock, ffi::Array<BufferVar>, ffi::ObjectPtrHash, ffi::ObjectPtrEqual>
+ ffi::Map<s_tir::SBlock, IndexMap> block_transformations_;
+ std::unordered_map<s_tir::SBlock, ffi::Array<BufferVar>, ffi::ObjectPtrHash, ffi::ObjectPtrEqual>
block_to_buffer_;
};
-ffi::Map<tirx::SBlock, ffi::Map<ffi::ObjectRef, tirx::IndexMap>> SuggestLayoutTransforms(
+ffi::Map<s_tir::SBlock, ffi::Map<ffi::ObjectRef, tirx::IndexMap>> SuggestLayoutTransforms(
const PrimFunc& prim_func, ffi::Array<IndexMap> write_buffer_transformations) {
// No changes to the PrimFunc are required if no transformations on output buffers.
if (write_buffer_transformations.empty()) return {};
diff --git a/src/relax/analysis/tir_op_pattern_kind.cc b/src/relax/analysis/tir_op_pattern_kind.cc
index 21a4e45bc6..f9f428c33a 100644
--- a/src/relax/analysis/tir_op_pattern_kind.cc
+++ b/src/relax/analysis/tir_op_pattern_kind.cc
@@ -24,11 +24,12 @@
#include <tvm/ffi/reflection/registry.h>
#include <tvm/relax/analysis.h>
#include <tvm/relax/op_attr_types.h>
+#include <tvm/s_tir/stmt.h>
+#include <tvm/s_tir/stmt_functor.h>
#include <tvm/tirx/analysis.h>
#include <tvm/tirx/expr_functor.h>
#include <tvm/tirx/function.h>
#include <tvm/tirx/op.h>
-#include <tvm/tirx/stmt_functor.h>
#include <algorithm>
@@ -38,7 +39,7 @@ using namespace tvm::prim;
using namespace tirx;
-class PatternKindAnalyzer : public StmtExprVisitor {
+class PatternKindAnalyzer : public s_tir::StmtExprVisitor {
public:
explicit PatternKindAnalyzer(const tirx::PrimFunc& func) {
for (const tirx::Var& param : func->params) {
@@ -50,7 +51,7 @@ class PatternKindAnalyzer : public StmtExprVisitor {
}
private:
- bool IsOutputBlock(const SBlockNode* block) {
+ bool IsOutputBlock(const s_tir::SBlockNode* block) {
for (const TensorRegion& write_region : block->writes) {
if (param_buffers_.count(write_region->source.as_or_throw<tvm::tirx::BufferVar>())) {
return true;
@@ -67,25 +68,25 @@ class PatternKindAnalyzer : public StmtExprVisitor {
return std::nullopt;
}
store_ = ffi::GetRef<BufferStore>(op);
- return StmtExprVisitor::Visit_(op);
+ return s_tir::StmtExprVisitor::Visit_(op);
}
ffi::Optional<VisitInterrupt> Visit_(const TensorLoadNode* op) final {
loads_.push_back(ffi::GetRef<TensorLoad>(op));
- return StmtExprVisitor::Visit_(op);
+ return s_tir::StmtExprVisitor::Visit_(op);
}
- ffi::Optional<VisitInterrupt> Visit_(const SBlockNode* op) final {
+ ffi::Optional<VisitInterrupt> Visit_(const s_tir::SBlockNode* op) final {
if (op->name_hint == "root") {
// Skip the root block
- return StmtExprVisitor::Visit(op->body);
+ return s_tir::StmtExprVisitor::Visit(op->body);
}
// Step 1. Clear loads and store
loads_.clear();
store_ = std::nullopt;
// Step 2. Visit block body.
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit(op->body));
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(s_tir::StmtExprVisitor::Visit(op->body));
// We support exactly one buffer store in a block (usually generated by TE compute)
// If we have not seen any store in the current block, classify as Opaque.
@@ -370,11 +371,11 @@ OpPatternKind AnalyzeOpPatternKind(const PrimFunc& func) {
}
bool HasReshapePattern(const PrimFunc& func) {
- class ReshapeDetector : public StmtExprVisitor {
+ class ReshapeDetector : public s_tir::StmtExprVisitor {
public:
ffi::Optional<VisitInterrupt> Visit(ffi::AnyView value) override {
if (value.as<tvm::ExprNode>()) return std::nullopt;
- return StmtExprVisitor::Visit(value);
+ return s_tir::StmtExprVisitor::Visit(value);
}
static bool Detect(const BufferVar& src_buffer, const BufferVar& dst_buffer, Stmt stmt) {
@@ -391,17 +392,18 @@ bool HasReshapePattern(const PrimFunc& func) {
ana_->Bind(loop->loop_var, Range::FromMinExtent(loop->min, loop->extent));
// To detect the reshape pattern, we require each For to have
// either another For or a BlockRealize as body.
- if (!(loop->body->IsInstance<ForNode>() || loop->body->IsInstance<SBlockRealizeNode>())) {
+ if (!(loop->body->IsInstance<ForNode>() ||
+ loop->body->IsInstance<s_tir::SBlockRealizeNode>())) {
return std::nullopt;
}
return this->Visit(loop->body);
}
- ffi::Optional<VisitInterrupt> Visit_(const SBlockRealizeNode* block_realize) final {
+ ffi::Optional<VisitInterrupt> Visit_(const s_tir::SBlockRealizeNode* block_realize) final {
// Constructing the mapping from block iterators to iterator
// binding values. The mapping will be used in the substitution of
// the flattened buffer access index.
- const SBlock& block = block_realize->block;
+ const s_tir::SBlock& block = block_realize->block;
const ffi::Array<IterVar>& block_iter = block->iter_vars;
const ffi::Array<PrimExpr>& iter_values = block_realize->iter_values;
TVM_FFI_ICHECK_EQ(block_iter.size(), iter_values.size());
@@ -417,7 +419,7 @@ bool HasReshapePattern(const PrimFunc& func) {
return this->Visit(block);
}
- ffi::Optional<VisitInterrupt> Visit_(const SBlockNode* block) final {
+ ffi::Optional<VisitInterrupt> Visit_(const s_tir::SBlockNode* block) final {
// Step 0. If the block body is a ForNode, recurse into it.
if (block->body->IsInstance<ForNode>()) {
return this->Visit(block->body);
@@ -572,7 +574,7 @@ bool HasReshapePattern(const PrimFunc& func) {
// To detect the reshape pattern, we require each For to have
// either another For or a BlockRealize as body.
- TVM_FFI_ICHECK(func->body->IsInstance<SBlockRealizeNode>());
+ TVM_FFI_ICHECK(func->body->IsInstance<s_tir::SBlockRealizeNode>());
return ReshapeDetector::Detect(src_buffer, dst_buffer, func->body);
}
diff --git a/src/relax/backend/task_extraction.cc b/src/relax/backend/task_extraction.cc
index 31d7f7f655..108aa6ffb1 100644
--- a/src/relax/backend/task_extraction.cc
+++ b/src/relax/backend/task_extraction.cc
@@ -22,9 +22,10 @@
#include <tvm/relax/expr.h>
#include <tvm/relax/expr_functor.h>
#include <tvm/s_tir/meta_schedule/extracted_task.h>
+#include <tvm/s_tir/stmt.h>
+#include <tvm/s_tir/stmt_functor.h>
#include <tvm/target/target.h>
#include <tvm/tirx/function.h>
-#include <tvm/tirx/stmt_functor.h>
#include "../../s_tir/meta_schedule/module_equality.h"
@@ -50,11 +51,11 @@ using s_tir::meta_schedule::ModuleHash;
* Then we will have a ExtractedTask for all three functions, whose weight
* is 5 + 3 + 2 = 10.
*/
-class BlockCounter : public tirx::StmtExprVisitor {
+class BlockCounter : public s_tir::StmtExprVisitor {
public:
ffi::Optional<VisitInterrupt> Visit(ffi::AnyView value) override {
if (value.as<tvm::ExprNode>()) return std::nullopt;
- return tirx::StmtExprVisitor::Visit(value);
+ return s_tir::StmtExprVisitor::Visit(value);
}
static size_t GetSBlockCount(const tirx::PrimFunc& func) {
@@ -64,9 +65,9 @@ class BlockCounter : public tirx::StmtExprVisitor {
}
private:
- ffi::Optional<VisitInterrupt> Visit_(const tirx::SBlockNode* op) final {
+ ffi::Optional<VisitInterrupt> Visit_(const s_tir::SBlockNode* op) final {
++count;
- return tirx::StmtExprVisitor::Visit_(op);
+ return s_tir::StmtExprVisitor::Visit_(op);
}
size_t count{0};
};
diff --git a/src/relax/distributed/transform/lower_global_view_to_local_view.cc b/src/relax/distributed/transform/lower_global_view_to_local_view.cc
index 42efd11af5..a7ab6cf621 100644
--- a/src/relax/distributed/transform/lower_global_view_to_local_view.cc
+++ b/src/relax/distributed/transform/lower_global_view_to_local_view.cc
@@ -27,8 +27,9 @@
#include <tvm/relax/distributed/axis_group_graph.h>
#include <tvm/relax/distributed/transform.h>
#include <tvm/relax/expr_functor.h>
+#include <tvm/s_tir/stmt.h>
+#include <tvm/s_tir/stmt_functor.h>
#include <tvm/s_tir/transform.h>
-#include <tvm/tirx/stmt_functor.h>
#include "../../../s_tir/schedule/transform.h"
#include "utils.h"
@@ -38,7 +39,7 @@ using namespace tvm::prim;
using namespace tvm::relax::distributed;
-class DistBufferReplacer : public StmtExprMutator {
+class DistBufferReplacer : public s_tir::StmtExprMutator {
public:
static Stmt BufferReplace(Stmt stmt, ffi::Map<BufferVar, BufferVar> buffer_map) {
auto replacer = ffi::make_object<DistBufferReplacer>(buffer_map);
@@ -52,26 +53,26 @@ class DistBufferReplacer : public StmtExprMutator {
}
};
-class DistSBlockInfoCollector : public StmtExprVisitor {
+class DistSBlockInfoCollector : public s_tir::StmtExprVisitor {
private:
ffi::Optional<VisitInterrupt> Visit_(const BufferStoreNode* op) final {
buffer_access_indices[op->buffer].push_back(op->indices);
- return StmtExprVisitor::Visit_(op);
+ return s_tir::StmtExprVisitor::Visit_(op);
}
ffi::Optional<VisitInterrupt> Visit_(const TensorLoadNode* op) final {
buffer_access_indices[op->source.as_or_throw<tvm::tirx::BufferVar>()].push_back(op->indices);
- return StmtExprVisitor::Visit_(op);
+ return s_tir::StmtExprVisitor::Visit_(op);
}
- ffi::Optional<VisitInterrupt> Visit_(const SBlockNode* op) final {
+ ffi::Optional<VisitInterrupt> Visit_(const s_tir::SBlockNode* op) final {
for (const auto& iter_var : op->iter_vars) {
if (iter_var->iter_type == kCommReduce) {
TVM_FFI_ICHECK(op->writes.size() == 1);
reduce_buffer_ = op->writes[0]->source.as_or_throw<tvm::tirx::BufferVar>();
}
}
- return StmtExprVisitor::Visit_(op);
+ return s_tir::StmtExprVisitor::Visit_(op);
}
bool IsReduceBufferAccess(const PrimExpr& expr) {
@@ -85,28 +86,28 @@ class DistSBlockInfoCollector : public StmtExprVisitor {
if (IsReduceBufferAccess(op->a) || IsReduceBufferAccess(op->b)) {
reduce_kind = "sum";
}
- return StmtExprVisitor::Visit_(op);
+ return s_tir::StmtExprVisitor::Visit_(op);
}
ffi::Optional<VisitInterrupt> Visit_(const prim::MulNode* op) final {
if (IsReduceBufferAccess(op->a) || IsReduceBufferAccess(op->b)) {
reduce_kind = "prod";
}
- return StmtExprVisitor::Visit_(op);
+ return s_tir::StmtExprVisitor::Visit_(op);
}
ffi::Optional<VisitInterrupt> Visit_(const prim::MinNode* op) final {
if (IsReduceBufferAccess(op->a) || IsReduceBufferAccess(op->b)) {
reduce_kind = "min";
}
- return StmtExprVisitor::Visit_(op);
+ return s_tir::StmtExprVisitor::Visit_(op);
}
ffi::Optional<VisitInterrupt> Visit_(const prim::MaxNode* op) final {
if (IsReduceBufferAccess(op->a) || IsReduceBufferAccess(op->b)) {
reduce_kind = "max";
}
- return StmtExprVisitor::Visit_(op);
+ return s_tir::StmtExprVisitor::Visit_(op);
}
BufferVar reduce_buffer_;
@@ -118,7 +119,7 @@ class DistSBlockInfoCollector : public StmtExprVisitor {
std::string reduce_kind;
};
-class DistributedBufferCompactor : public StmtExprMutator {
+class DistributedBufferCompactor : public s_tir::StmtExprMutator {
// FIXME: change to use unordered_map<int, AxisShardingSpec> (represent dim and sharding spec)
// Currently we assume device mesh is only 1d, but when we support 2d, we need to change this
using DimShard = std::unordered_map<int, int>;
@@ -183,7 +184,7 @@ class DistributedBufferCompactor : public StmtExprMutator {
}
ffi::Array<IterVar> ShardIterVar(
- SBlock block,
+ s_tir::SBlock block,
const std::unordered_map<BufferVar, ffi::Array<ffi::Array<PrimExpr>>, ffi::ObjectPtrHash,
ffi::ObjectPtrEqual>& buffer_access_indices) {
std::vector<BufferVar> buffers;
@@ -255,10 +256,10 @@ class DistributedBufferCompactor : public StmtExprMutator {
return BufferVar(buffer.name(), std::move(new_type), buffer.span());
}
- UnchangedOr<Stmt> Mutate_(const SBlockNode* op, InplaceMode inplace_mode) final {
- SBlock block = StmtExprMutator::Mutate_(op, inplace_mode)
- .ValueOrUnchanged(ffi::GetRef<Stmt>(op))
- .as_or_throw<SBlock>();
+ UnchangedOr<Stmt> Mutate_(const s_tir::SBlockNode* op, InplaceMode inplace_mode) final {
+ s_tir::SBlock block = s_tir::StmtExprMutator::Mutate_(op, inplace_mode)
+ .ValueOrUnchanged(ffi::GetRef<Stmt>(op))
+ .as_or_throw<s_tir::SBlock>();
auto collector = ffi::make_object<DistSBlockInfoCollector>();
collector->Visit(block);
ffi::Array<IterVar> new_iter_vars = ShardIterVar(block, collector->buffer_access_indices);
@@ -280,7 +281,8 @@ class DistributedBufferCompactor : public StmtExprMutator {
break;
}
}
- ffi::ObjectPtr<SBlockNode> new_block = ffi::make_object<SBlockNode>(*block.operator->());
+ ffi::ObjectPtr<s_tir::SBlockNode> new_block =
+ ffi::make_object<s_tir::SBlockNode>(*block.operator->());
new_block->iter_vars = new_iter_vars;
new_block->alloc_buffers = new_alloc_buffers;
if (new_block->name_hint == "root") {
@@ -289,15 +291,15 @@ class DistributedBufferCompactor : public StmtExprMutator {
allocated_buffer_under_root.end());
}
new_block->body = DistBufferReplacer::BufferReplace(block->body, buffer_map);
- return SBlock(new_block);
+ return s_tir::SBlock(new_block);
}
void AddAllReduceBlock(std::string reduce_kind) { add_allreduce_kind_ = reduce_kind; }
- UnchangedOr<Stmt> Mutate_(const SBlockRealizeNode* op, InplaceMode inplace_mode) final {
- SBlockRealize realize = StmtExprMutator::Mutate_(op, inplace_mode)
- .ValueOrUnchanged(ffi::GetRef<Stmt>(op))
- .as_or_throw<SBlockRealize>();
+ UnchangedOr<Stmt> Mutate_(const s_tir::SBlockRealizeNode* op, InplaceMode inplace_mode) final {
+ s_tir::SBlockRealize realize = s_tir::StmtExprMutator::Mutate_(op, inplace_mode)
+ .ValueOrUnchanged(ffi::GetRef<Stmt>(op))
+ .as_or_throw<s_tir::SBlockRealize>();
for (int i = 0; i < static_cast<int>(realize->iter_values.size()); i++) {
PrimExpr iter_value = realize->iter_values[i];
@@ -313,7 +315,7 @@ class DistributedBufferCompactor : public StmtExprMutator {
}
UnchangedOr<Stmt> Mutate_(const ForNode* op, InplaceMode inplace_mode) final {
- For new_loop = StmtExprMutator::Mutate_(op, inplace_mode)
+ For new_loop = s_tir::StmtExprMutator::Mutate_(op, inplace_mode)
.ValueOrUnchanged(ffi::GetRef<Stmt>(op))
.as_or_throw<For>();
if (loop_var_shards_.count(op->loop_var)) {
diff --git a/src/relax/transform/fuse_tir.cc b/src/relax/transform/fuse_tir.cc
index 6b9eabe198..925bbda693 100644
--- a/src/relax/transform/fuse_tir.cc
+++ b/src/relax/transform/fuse_tir.cc
@@ -24,9 +24,10 @@
#include <tvm/relax/expr_functor.h>
#include <tvm/relax/transform.h>
#include <tvm/relax/type.h>
+#include <tvm/s_tir/stmt.h>
+#include <tvm/s_tir/stmt_functor.h>
#include <tvm/s_tir/transform.h>
#include <tvm/tirx/op.h>
-#include <tvm/tirx/stmt_functor.h>
#include <unordered_map>
#include <unordered_set>
@@ -170,7 +171,7 @@ class SymbolicMatcher : ExprFunctor<void(const Expr& n, const PrimExpr& other)>
/*!
* \brief Substitute a given source buffer with a given target buffer in statements or expressions.
*/
-class FuseTIRBufferSubstitutor : public StmtExprMutator {
+class FuseTIRBufferSubstitutor : public s_tir::StmtExprMutator {
public:
explicit FuseTIRBufferSubstitutor(const ffi::Map<BufferVar, BufferVar>& buffer_map,
const ffi::Map<Var, PrimExpr>& var_map) {
@@ -190,10 +191,10 @@ class FuseTIRBufferSubstitutor : public StmtExprMutator {
}
private:
- UnchangedOr<Stmt> Mutate_(const SBlockNode* op, InplaceMode inplace_mode) final {
- SBlock block = StmtExprMutator::Mutate_(op, inplace_mode)
- .ValueOrUnchanged(ffi::GetRef<Stmt>(op))
- .as_or_throw<SBlock>();
+ UnchangedOr<Stmt> Mutate_(const s_tir::SBlockNode* op, InplaceMode inplace_mode) final {
+ s_tir::SBlock block = s_tir::StmtExprMutator::Mutate_(op, inplace_mode)
+ .ValueOrUnchanged(ffi::GetRef<Stmt>(op))
+ .as_or_throw<s_tir::SBlock>();
ffi::Array<TensorRegion> reads = UnionAccessRegion(block->reads);
ffi::Array<TensorRegion> writes = UnionAccessRegion(block->writes);
if (!reads.same_as(block->reads) || !writes.same_as(block->writes)) {
@@ -230,19 +231,19 @@ class FuseTIRBufferSubstitutor : public StmtExprMutator {
};
/*! \brief A mutator which detect block name duplication and deduplicate the names. */
-class SBlockNameDeduplicator : public tirx::StmtExprMutator {
+class SBlockNameDeduplicator : public s_tir::StmtExprMutator {
public:
- using StmtExprMutator::Mutate;
+ using s_tir::StmtExprMutator::Mutate;
UnchangedOr<ffi::Any> Mutate(ffi::AnyView value, InplaceMode inplace_mode) final {
if (value.as<ExprNode>()) return ffi::Unchanged();
- return StmtExprMutator::Mutate(value, inplace_mode);
+ return s_tir::StmtExprMutator::Mutate(value, inplace_mode);
}
private:
- UnchangedOr<Stmt> Mutate_(const SBlockNode* op, InplaceMode inplace_mode) final {
- SBlock block = tirx::StmtExprMutator::Mutate_(op, inplace_mode)
- .ValueOrUnchanged(ffi::GetRef<Stmt>(op))
- .as_or_throw<SBlock>();
+ UnchangedOr<Stmt> Mutate_(const s_tir::SBlockNode* op, InplaceMode inplace_mode) final {
+ s_tir::SBlock block = s_tir::StmtExprMutator::Mutate_(op, inplace_mode)
+ .ValueOrUnchanged(ffi::GetRef<Stmt>(op))
+ .as_or_throw<s_tir::SBlock>();
ffi::String name = GetUniqueName(block->name_hint);
@@ -573,10 +574,10 @@ class FusedTIRConstructor : public ExprVisitor {
// Step 3. Check functions are all schedulable funcs. i.e. the body of func is root block
// TODO(Siyuan): support un-schedulable functions.
- TVM_FFI_ICHECK(prim_func->body->IsInstance<tirx::SBlockRealizeNode>())
+ TVM_FFI_ICHECK(prim_func->body->IsInstance<s_tir::SBlockRealizeNode>())
<< "Only schedulable functions (whose body is the root block) can be fused";
- const tirx::SBlockRealize& root_realize = prim_func->body.as_or_throw<tirx::SBlockRealize>();
- const tirx::SBlock& root_block = root_realize->block;
+ const s_tir::SBlockRealize& root_realize = prim_func->body.as_or_throw<s_tir::SBlockRealize>();
+ const s_tir::SBlock& root_block = root_realize->block;
// Step 4. Add all the original alloc_buffers and body to the fused function.
func_info_.alloc_buffers.insert(func_info_.alloc_buffers.end(),
@@ -878,8 +879,8 @@ class FusedTIRConstructor : public ExprVisitor {
body = ffi::make_object<tirx::SBlockNameDeduplicator>()->Mutate(body).ValueOrUnchanged(body);
body = subst->Mutate(body).ValueOrUnchanged(body);
- body = tirx::SBlock({}, {}, {}, "root", std::move(body), std::nullopt, alloc_buffers);
- body = tirx::SBlockRealize({}, IntImm::Bool(true), body.as_or_throw<tirx::SBlock>());
+ body = s_tir::SBlock({}, {}, {}, "root", std::move(body), std::nullopt, alloc_buffers);
+ body = s_tir::SBlockRealize({}, IntImm::Bool(true), body.as_or_throw<s_tir::SBlock>());
ffi::Array<tirx::Var> params = func_info_.params.Map([&](const tirx::Var& param) {
if (auto buffer = func_info_.buffer_map.Get(param)) {
return buffer.value().var();
diff --git a/src/relax/transform/legalize_ops.cc b/src/relax/transform/legalize_ops.cc
index 12a5182792..ecc1f0e0db 100644
--- a/src/relax/transform/legalize_ops.cc
+++ b/src/relax/transform/legalize_ops.cc
@@ -31,6 +31,7 @@
#include <tvm/relax/transform.h>
#include <tvm/relax/type.h>
#include <tvm/runtime/logging.h>
+#include <tvm/s_tir/transform.h>
#include <tvm/tirx/transform.h>
#include <set>
@@ -102,7 +103,7 @@ class LegalizeMutator : public ExprMutator {
// Avoid accidental sharing of TIR variables in the legalized
// PrimFuncs, when kernels for multiple devices are generated
// from the same PrimFunc.
- output = tirx::transform::ConvertSSA()(output);
+ output = s_tir::transform::ConvertSSA()(output);
}
return output;
diff --git a/src/relax/transform/split_call_tir_by_pattern.cc b/src/relax/transform/split_call_tir_by_pattern.cc
index d179e1efe3..54bf928e71 100644
--- a/src/relax/transform/split_call_tir_by_pattern.cc
+++ b/src/relax/transform/split_call_tir_by_pattern.cc
@@ -29,10 +29,11 @@
#include <tvm/relax/tir_pattern.h>
#include <tvm/relax/transform.h>
#include <tvm/relax/type.h>
+#include <tvm/s_tir/stmt.h>
+#include <tvm/s_tir/stmt_functor.h>
#include <tvm/s_tir/transform.h>
#include <tvm/tirx/builtin.h>
#include <tvm/tirx/op.h>
-#include <tvm/tirx/stmt_functor.h>
#include "../../s_tir/schedule/ir_comparator.h"
@@ -65,7 +66,8 @@ class ForMatcher : public TensorizeComparator {
}
bool Match(const For& top) {
- const ForNode* pattern_top = pattern_->body.as<SBlockRealizeNode>()->block->body.as<ForNode>();
+ const ForNode* pattern_top =
+ pattern_->body.as<s_tir::SBlockRealizeNode>()->block->body.as<ForNode>();
TVM_FFI_ICHECK(pattern_top) << "Invalid pattern function";
if (!Dispatch(top, ffi::GetRef<Stmt>(pattern_top))) {
return false;
@@ -254,10 +256,10 @@ class ForMatcher : public TensorizeComparator {
loop_stack_lhs_.push_back(ffi::GetRef<For>(op));
loop_stack_rhs_.push_back(ffi::GetRef<For>(rhs));
// The body of loop must be loop or BlockRealize
- if (!op->body->IsInstance<SBlockRealizeNode>() && !op->body->IsInstance<ForNode>()) {
+ if (!op->body->IsInstance<s_tir::SBlockRealizeNode>() && !op->body->IsInstance<ForNode>()) {
return false;
}
- if (!rhs->body->IsInstance<SBlockRealizeNode>() && !rhs->body->IsInstance<ForNode>()) {
+ if (!rhs->body->IsInstance<s_tir::SBlockRealizeNode>() && !rhs->body->IsInstance<ForNode>()) {
return false;
}
// Build mapping between the loop vars
@@ -272,8 +274,8 @@ class ForMatcher : public TensorizeComparator {
return Dispatch(op->body, rhs->body);
}
- bool Dispatch_(const tirx::SBlockNode* op, const Stmt& other) final {
- const auto* rhs = other.as<SBlockNode>();
+ bool Dispatch_(const s_tir::SBlockNode* op, const Stmt& other) final {
+ const auto* rhs = other.as<s_tir::SBlockNode>();
// Check block equality.
// All iter vars and buffer regions including the order should match.
// When checking iter vars, DefEqual is used to remap variables.
@@ -301,8 +303,8 @@ class ForMatcher : public TensorizeComparator {
return Dispatch(op->body, rhs->body);
}
- bool Dispatch_(const SBlockRealizeNode* op, const Stmt& other) final {
- const auto* rhs = other.as<SBlockRealizeNode>();
+ bool Dispatch_(const s_tir::SBlockRealizeNode* op, const Stmt& other) final {
+ const auto* rhs = other.as<s_tir::SBlockRealizeNode>();
// Only allow trivial bindings
for (size_t i = 0; i < op->iter_values.size(); ++i) {
if (!op->iter_values[i].same_as(loop_stack_lhs_[i]->loop_var)) return false;
@@ -463,7 +465,7 @@ class TIRPatternMatcher {
/*! \brief helper class to partition a function into 2 parts. Return function information which we
* can use to construct the two partitioned parts.*/
-class FunctionPartitioner : public StmtExprVisitor {
+class FunctionPartitioner : public s_tir::StmtExprVisitor {
public:
explicit FunctionPartitioner(int num_matched_ops) : num_matched_ops_(num_matched_ops) {}
/*! \brief alloc_buffers for the first function */
@@ -471,7 +473,7 @@ class FunctionPartitioner : public StmtExprVisitor {
/*! \brief alloc_buffers for the second function */
std::unordered_set<BufferVar, ffi::ObjectPtrHash, ffi::ObjectPtrEqual> allocs2;
/*! \brief whether the current block is in the first function */
- ffi::Map<SBlock, bool> block_partition;
+ ffi::Map<s_tir::SBlock, bool> block_partition;
/*! \brief input buffers for the first function */
std::unordered_set<BufferVar, ffi::ObjectPtrHash, ffi::ObjectPtrEqual> input1;
/*! \brief input buffers for the second function */
@@ -484,7 +486,7 @@ class FunctionPartitioner : public StmtExprVisitor {
bool fail = false;
private:
- ffi::Optional<VisitInterrupt> Visit_(const SBlockNode* op) final {
+ ffi::Optional<VisitInterrupt> Visit_(const s_tir::SBlockNode* op) final {
block_counter_++;
bool is_matching_ = block_counter_ <= num_matched_ops_;
if (block_counter_ == num_matched_ops_) {
@@ -512,7 +514,7 @@ class FunctionPartitioner : public StmtExprVisitor {
input2.insert(write->source.as_or_throw<tvm::tirx::BufferVar>());
}
}
- block_partition.Set(ffi::GetRef<SBlock>(op), is_matching_);
+ block_partition.Set(ffi::GetRef<s_tir::SBlock>(op), is_matching_);
return std::nullopt;
}
@@ -522,30 +524,30 @@ class FunctionPartitioner : public StmtExprVisitor {
};
/*! \brief remove parts according to block partition, and update the alloc_buffers for blocks */
-class BlockRemover : public StmtExprMutator {
+class BlockRemover : public s_tir::StmtExprMutator {
public:
static Stmt RemoveBlockByPartition(
- Stmt stmt, const ffi::Map<SBlock, bool>& block_partition,
+ Stmt stmt, const ffi::Map<s_tir::SBlock, bool>& block_partition,
const std::unordered_set<BufferVar, ffi::ObjectPtrHash, ffi::ObjectPtrEqual>& allocs,
bool is_library_part) {
auto remover = ffi::make_object<BlockRemover>(block_partition, allocs, is_library_part);
return remover->Mutate(stmt, InplaceMode::kDisallow).ValueOrUnchanged(stmt);
}
- BlockRemover(const ffi::Map<SBlock, bool>& block_partition,
+ BlockRemover(const ffi::Map<s_tir::SBlock, bool>& block_partition,
const std::unordered_set<BufferVar, ffi::ObjectPtrHash, ffi::ObjectPtrEqual>& allocs,
bool is_library_part)
: block_partition(block_partition), allocs_(allocs), is_library_part_(is_library_part) {}
private:
- UnchangedOr<Stmt> Mutate_(const SBlockNode* op, InplaceMode inplace_mode) final {
- SBlock block = StmtExprMutator::Mutate_(op, inplace_mode)
- .ValueOrUnchanged(ffi::GetRef<Stmt>(op))
- .as_or_throw<SBlock>();
- ffi::ObjectPtr<SBlockNode> n = ffi::make_object<SBlockNode>(*block.operator->());
+ UnchangedOr<Stmt> Mutate_(const s_tir::SBlockNode* op, InplaceMode inplace_mode) final {
+ s_tir::SBlock block = s_tir::StmtExprMutator::Mutate_(op, inplace_mode)
+ .ValueOrUnchanged(ffi::GetRef<Stmt>(op))
+ .as_or_throw<s_tir::SBlock>();
+ ffi::ObjectPtr<s_tir::SBlockNode> n = ffi::make_object<s_tir::SBlockNode>(*block.operator->());
if (op->name_hint != "root") {
- TVM_FFI_ICHECK(block_partition.count(ffi::GetRef<SBlock>(op)));
- bool block_is_library = block_partition[ffi::GetRef<SBlock>(op)];
+ TVM_FFI_ICHECK(block_partition.count(ffi::GetRef<s_tir::SBlock>(op)));
+ bool block_is_library = block_partition[ffi::GetRef<s_tir::SBlock>(op)];
if (!(is_library_part_ ^ block_is_library)) {
n->body = block->body;
} else {
@@ -559,7 +561,7 @@ class BlockRemover : public StmtExprMutator {
}
}
n->alloc_buffers = alloc_buffers;
- return SBlock(n);
+ return s_tir::SBlock(n);
}
UnchangedOr<Stmt> Mutate_(const SeqStmtNode* op, InplaceMode inplace_mode) final {
@@ -576,7 +578,7 @@ class BlockRemover : public StmtExprMutator {
}
bool erased_ = false;
- ffi::Map<SBlock, bool> block_partition;
+ ffi::Map<s_tir::SBlock, bool> block_partition;
std::unordered_set<BufferVar, ffi::ObjectPtrHash, ffi::ObjectPtrEqual> allocs_;
bool is_library_part_ = false;
};
@@ -595,9 +597,9 @@ std::pair<PrimFunc, ffi::Optional<PrimFunc>> SplitFunctions(
PrimFunc func, std::vector<std::vector<int>>* arg_partition, ffi::Array<TIRPattern> patterns,
FCodegen f_codegen) {
// Step 1. Find the library kernel and the rest.
- Stmt body = func->body.as<SBlockRealizeNode>()->block->body;
+ Stmt body = func->body.as<s_tir::SBlockRealizeNode>()->block->body;
ffi::Array<MatchResult> match_results =
- TIRPatternMatcher::Match(patterns, func->body.as<SBlockRealizeNode>()->block->body);
+ TIRPatternMatcher::Match(patterns, func->body.as<s_tir::SBlockRealizeNode>()->block->body);
if (match_results.empty()) {
return {func, std::nullopt};
}
diff --git a/src/relax/transform/split_layout_rewrite_preproc.cc b/src/relax/transform/split_layout_rewrite_preproc.cc
index 845900ecc5..cd08d7d24e 100644
--- a/src/relax/transform/split_layout_rewrite_preproc.cc
+++ b/src/relax/transform/split_layout_rewrite_preproc.cc
@@ -26,8 +26,8 @@
#include <tvm/relax/expr_functor.h>
#include <tvm/relax/transform.h>
#include <tvm/s_tir/stmt.h>
+#include <tvm/s_tir/stmt_functor.h>
#include <tvm/s_tir/transform.h>
-#include <tvm/tirx/stmt_functor.h>
#include <algorithm>
#include <cstddef>
@@ -36,19 +36,19 @@ namespace tvm {
namespace tirx {
using namespace tvm::prim;
-class SplitPrimFuncLayoutRewrite : public StmtExprMutator {
+class SplitPrimFuncLayoutRewrite : public s_tir::StmtExprMutator {
public:
- using StmtExprMutator::Mutate;
+ using s_tir::StmtExprMutator::Mutate;
UnchangedOr<ffi::Any> Mutate(ffi::AnyView value, InplaceMode inplace_mode) final {
if (value.as<ExprNode>()) return ffi::Unchanged();
- return StmtExprMutator::Mutate(value, inplace_mode);
+ return s_tir::StmtExprMutator::Mutate(value, inplace_mode);
}
explicit SplitPrimFuncLayoutRewrite(const PrimFunc& func) : original_func_(func) {}
std::tuple<ffi::Optional<PrimFunc>, PrimFunc> Transform(const PrimFunc& func) {
- TVM_FFI_ICHECK(func->body.as<SBlockRealizeNode>())
+ TVM_FFI_ICHECK(func->body.as<s_tir::SBlockRealizeNode>())
<< "The body of the primfunc should be a root block.";
- const auto& block = func->body.as<SBlockRealizeNode>()->block;
+ const auto& block = func->body.as<s_tir::SBlockRealizeNode>()->block;
visit_root_block(block.get());
if (layout_rewrite_preproc_stmts_.size() > 0) {
return std::make_tuple(create_layout_rewrite_preproc_func(), create_compute_func());
@@ -83,12 +83,12 @@ class SplitPrimFuncLayoutRewrite : public StmtExprMutator {
<< "There should be at least one layout rewrite preproc stmt.";
Stmt body = layout_rewrite_preproc_stmts_.size() == 1 ? layout_rewrite_preproc_stmts_[0]
: SeqStmt(layout_rewrite_preproc_stmts_);
- body = SBlockRealize(
+ body = s_tir::SBlockRealize(
/*iter_values=*/ffi::Array<PrimExpr>(),
/*predicate=*/IntImm::Bool(true),
/*block=*/
- SBlock(/*iter_vars=*/{}, /*reads=*/{}, /*writes=*/{},
- /*name_hint=*/"root", body));
+ s_tir::SBlock(/*iter_vars=*/{}, /*reads=*/{}, /*writes=*/{},
+ /*name_hint=*/"root", body));
ffi::Map<ffi::String, ffi::Any> dict;
for (const auto& [key, original_value] : original_func_->attrs->dict) {
@@ -115,7 +115,7 @@ class SplitPrimFuncLayoutRewrite : public StmtExprMutator {
// Step 2: Create the body for the new PrimFunc
Stmt body = compute_stmts_.size() == 1 ? compute_stmts_[0] : SeqStmt(compute_stmts_);
- SBlock original_block = original_func_->body.as<SBlockRealizeNode>()->block;
+ s_tir::SBlock original_block = original_func_->body.as<s_tir::SBlockRealizeNode>()->block;
ffi::Array<BufferVar> alloc_buffers;
for (const auto& buffer : original_block->alloc_buffers) {
auto it =
@@ -126,14 +126,14 @@ class SplitPrimFuncLayoutRewrite : public StmtExprMutator {
}
}
- body = SBlockRealize(
+ body = s_tir::SBlockRealize(
/*iter_values=*/ffi::Array<PrimExpr>(),
/*predicate=*/IntImm::Bool(true),
/*block=*/
- SBlock(/*iter_vars=*/{}, /*reads=*/{}, /*writes=*/{},
- /*name_hint=*/"root", body,
- /*init=*/std::nullopt,
- /*alloc_buffers=*/alloc_buffers));
+ s_tir::SBlock(/*iter_vars=*/{}, /*reads=*/{}, /*writes=*/{},
+ /*name_hint=*/"root", body,
+ /*init=*/std::nullopt,
+ /*alloc_buffers=*/alloc_buffers));
ffi::Map<ffi::String, ffi::Any> dict;
for (const auto& [key, original_value] : original_func_->attrs->dict) {
@@ -149,7 +149,7 @@ class SplitPrimFuncLayoutRewrite : public StmtExprMutator {
return s_tir::RenewDefs(func);
}
- void visit_root_block(const SBlockNode* op) {
+ void visit_root_block(const s_tir::SBlockNode* op) {
Stmt body = op->body;
if (const auto* seq_stmt = body.as<SeqStmtNode>()) {
for (const auto& stmt : seq_stmt->seq) {
@@ -169,10 +169,10 @@ class SplitPrimFuncLayoutRewrite : public StmtExprMutator {
<< "There should be a compute block if there is only one subtree under the root.";
}
}
- UnchangedOr<Stmt> Mutate_(const SBlockNode* op, InplaceMode inplace_mode) final {
- SBlock block = StmtExprMutator::Mutate_(op, inplace_mode)
- .ValueOrUnchanged(ffi::GetRef<Stmt>(op))
- .as_or_throw<SBlock>();
+ UnchangedOr<Stmt> Mutate_(const s_tir::SBlockNode* op, InplaceMode inplace_mode) final {
+ s_tir::SBlock block = s_tir::StmtExprMutator::Mutate_(op, inplace_mode)
+ .ValueOrUnchanged(ffi::GetRef<Stmt>(op))
+ .as_or_throw<s_tir::SBlock>();
auto it = op->annotations.find(s_tir::attr::meta_schedule_layout_rewrite_preproc);
bool is_layout_rewrite_preproc =
it != op->annotations.end() && is_one((*it).second.cast<PrimExpr>());
@@ -213,9 +213,9 @@ class SplitPrimFuncLayoutRewrite : public StmtExprMutator {
auto new_annotations = op->annotations;
new_annotations.erase(s_tir::attr::meta_schedule_layout_rewrite_preproc);
- auto n = ffi::make_object<SBlockNode>(*block.get());
+ auto n = ffi::make_object<s_tir::SBlockNode>(*block.get());
n->annotations = new_annotations;
- return SBlock(n);
+ return s_tir::SBlock(n);
}
return block;
}
diff --git a/src/s_tir/analysis/calculate_allocated_memory.cc b/src/s_tir/analysis/calculate_allocated_memory.cc
index a5245595e8..c96a9c2666 100644
--- a/src/s_tir/analysis/calculate_allocated_memory.cc
+++ b/src/s_tir/analysis/calculate_allocated_memory.cc
@@ -25,10 +25,11 @@
#include <tvm/ffi/container/map.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/runtime/device_api.h>
+#include <tvm/s_tir/analysis.h>
+#include <tvm/s_tir/stmt_functor.h>
#include <tvm/s_tir/transform.h>
#include <tvm/tirx/analysis.h>
#include <tvm/tirx/function.h>
-#include <tvm/tirx/stmt_functor.h>
#include <tvm/tirx/transform.h>
#include <algorithm>
diff --git a/src/s_tir/analysis/conditional_bounds.cc b/src/s_tir/analysis/conditional_bounds.cc
index 20129a1411..448171b4da 100644
--- a/src/s_tir/analysis/conditional_bounds.cc
+++ b/src/s_tir/analysis/conditional_bounds.cc
@@ -28,6 +28,7 @@
#include <tvm/ffi/extra/structural_mutate.h>
#include <tvm/ir/expr_functor.h>
#include <tvm/ir/prim/builtin.h>
+#include <tvm/s_tir/analysis.h>
#include <tvm/tirx/analysis.h>
#include <tvm/tirx/op.h>
diff --git a/src/s_tir/analysis/domain_touched.cc b/src/s_tir/analysis/domain_touched.cc
index 2e67be88ea..bbd8c4430d 100644
--- a/src/s_tir/analysis/domain_touched.cc
+++ b/src/s_tir/analysis/domain_touched.cc
@@ -27,14 +27,14 @@
#include <tvm/ir/prim/expr.h>
#include <tvm/runtime/logging.h>
#include <tvm/s_tir/analysis.h>
+#include <tvm/s_tir/stmt_functor.h>
#include <tvm/te/tensor.h>
-#include <tvm/tirx/stmt_functor.h>
#include <tuple>
#include <unordered_map>
#include <unordered_set>
-#include "../../tirx/ir/ir_visitor_with_analyzer.h"
+#include "../../s_tir/ir/ir_visitor_with_analyzer.h"
namespace tvm {
namespace s_tir {
@@ -63,9 +63,9 @@ using BufferDomainAccess = std::tuple<LoadAccess, StoreAccess, CombinedAccess>;
} // namespace
// Find Read region of the tensor in the stmt.
-class BufferTouchedDomain final : public tirx::IRVisitorWithAnalyzer {
+class BufferTouchedDomain final : public s_tir::IRVisitorWithAnalyzer {
public:
- using tirx::IRVisitorWithAnalyzer::Visit_;
+ using s_tir::IRVisitorWithAnalyzer::Visit_;
std::unordered_map<const VarNode*, BufferDomainAccess>& GetAccessedBufferRegions() {
return buffer_access_map_;
@@ -99,7 +99,7 @@ class BufferTouchedDomain final : public tirx::IRVisitorWithAnalyzer {
}
private:
- using Parent = tirx::IRVisitorWithAnalyzer;
+ using Parent = s_tir::IRVisitorWithAnalyzer;
ffi::Optional<VisitInterrupt> Visit_(const TensorLoadNode* op) final {
BufferVar buffer = op->source.as_or_throw<tvm::tirx::BufferVar>();
diff --git a/src/s_tir/analysis/estimate_flops.cc b/src/s_tir/analysis/estimate_flops.cc
index 0da2814819..3665f5006d 100644
--- a/src/s_tir/analysis/estimate_flops.cc
+++ b/src/s_tir/analysis/estimate_flops.cc
@@ -17,8 +17,10 @@
* under the License.
*/
#include <tvm/ffi/reflection/registry.h>
+#include <tvm/s_tir/analysis.h>
+#include <tvm/s_tir/stmt.h>
+#include <tvm/s_tir/stmt_functor.h>
#include <tvm/tirx/analysis.h>
-#include <tvm/tirx/stmt_functor.h>
#include "tvm/arith/analyzer.h"
diff --git a/src/s_tir/analysis/find_anchor_sblock.cc b/src/s_tir/analysis/find_anchor_sblock.cc
index 666f4f4e2b..dc4f762102 100644
--- a/src/s_tir/analysis/find_anchor_sblock.cc
+++ b/src/s_tir/analysis/find_anchor_sblock.cc
@@ -25,19 +25,20 @@
#include <tvm/ffi/cast.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/s_tir/analysis.h>
+#include <tvm/s_tir/stmt.h>
+#include <tvm/s_tir/stmt_functor.h>
#include <tvm/tirx/analysis.h>
-#include <tvm/tirx/stmt_functor.h>
namespace tvm {
namespace tirx {
-Stmt GetEnclosingLoop(const SBlockNode* block, Stmt func_body) {
- struct GetRootSeqStmt : public StmtExprVisitor {
- using StmtExprVisitor::Visit_;
+Stmt GetEnclosingLoop(const s_tir::SBlockNode* block, Stmt func_body) {
+ struct GetRootSeqStmt : public s_tir::StmtExprVisitor {
+ using s_tir::StmtExprVisitor::Visit_;
ffi::Optional<VisitInterrupt> Visit(ffi::AnyView value) override {
if (value.as<ExprNode>()) return std::nullopt;
- return StmtExprVisitor::Visit(value);
+ return s_tir::StmtExprVisitor::Visit(value);
}
ffi::Optional<VisitInterrupt> Visit_(const SeqStmtNode* seq) override {
@@ -47,24 +48,24 @@ Stmt GetEnclosingLoop(const SBlockNode* block, Stmt func_body) {
const SeqStmtNode* result;
};
- struct BlockFinder : public StmtExprVisitor {
- using StmtExprVisitor::Visit_;
+ struct BlockFinder : public s_tir::StmtExprVisitor {
+ using s_tir::StmtExprVisitor::Visit_;
ffi::Optional<VisitInterrupt> Visit(ffi::AnyView value) override {
if (value.as<ExprNode>()) return std::nullopt;
- return StmtExprVisitor::Visit(value);
+ return s_tir::StmtExprVisitor::Visit(value);
}
- explicit BlockFinder(const SBlockNode* tgt) : target(tgt) {}
+ explicit BlockFinder(const s_tir::SBlockNode* tgt) : target(tgt) {}
- ffi::Optional<VisitInterrupt> Visit_(const SBlockNode* block) override {
+ ffi::Optional<VisitInterrupt> Visit_(const s_tir::SBlockNode* block) override {
if (block == target) {
found = true;
}
return std::nullopt;
}
- const SBlockNode* target;
+ const s_tir::SBlockNode* target;
bool found = false;
};
@@ -84,26 +85,26 @@ Stmt GetEnclosingLoop(const SBlockNode* block, Stmt func_body) {
}
TVM_FFI_THROW(InternalError) << "Enclosing loop not found for a block "
- << ffi::GetRef<SBlock>(block);
+ << ffi::GetRef<s_tir::SBlock>(block);
TVM_FFI_UNREACHABLE();
}
-const SBlockNode* FindAnchorBlock(const IRModule& mod) {
- struct ReductionSBlockCollector : public StmtExprVisitor {
- using StmtExprVisitor::Visit_;
+const s_tir::SBlockNode* FindAnchorBlock(const IRModule& mod) {
+ struct ReductionSBlockCollector : public s_tir::StmtExprVisitor {
+ using s_tir::StmtExprVisitor::Visit_;
ffi::Optional<VisitInterrupt> Visit(ffi::AnyView value) override {
if (value.as<ExprNode>()) return std::nullopt;
- return StmtExprVisitor::Visit(value);
+ return s_tir::StmtExprVisitor::Visit(value);
}
- ffi::Optional<VisitInterrupt> Visit_(const SBlockNode* block) override {
+ ffi::Optional<VisitInterrupt> Visit_(const s_tir::SBlockNode* block) override {
if (block->init) {
blocks.push_back(block);
}
- return StmtExprVisitor::Visit(block->body);
+ return s_tir::StmtExprVisitor::Visit(block->body);
}
- std::vector<const SBlockNode*> blocks;
+ std::vector<const s_tir::SBlockNode*> blocks;
};
if (auto prim_func = FindEntryFunc(mod, nullptr)) {
@@ -138,9 +139,9 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("s_tir.analysis.find_anchor_sblock", [](const IRModule& mod) {
auto ret = FindAnchorBlock(mod);
if (ret) {
- return ffi::Optional<SBlock>(ffi::GetRef<SBlock>(ret));
+ return ffi::Optional<s_tir::SBlock>(ffi::GetRef<s_tir::SBlock>(ret));
}
- return ffi::Optional<SBlock>(std::nullopt);
+ return ffi::Optional<s_tir::SBlock>(std::nullopt);
});
}
diff --git a/src/s_tir/analysis/identify_memcpy.cc b/src/s_tir/analysis/identify_memcpy.cc
index d80c957ee8..02cfdbe865 100644
--- a/src/s_tir/analysis/identify_memcpy.cc
+++ b/src/s_tir/analysis/identify_memcpy.cc
@@ -38,7 +38,7 @@
#include <string>
#include <variant>
-#include "../../tirx/ir/ir_visitor_with_analyzer.h"
+#include "../../s_tir/ir/ir_visitor_with_analyzer.h"
namespace tvm {
namespace s_tir {
@@ -297,9 +297,9 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("s_tir.analysis._identify_memcpy", [](const Stmt& stmt) {
ffi::Array<ffi::ObjectRef> output;
- struct Visitor : tirx::IRVisitorWithAnalyzer {
+ struct Visitor : s_tir::IRVisitorWithAnalyzer {
public:
- using tirx::IRVisitorWithAnalyzer::Visit_;
+ using s_tir::IRVisitorWithAnalyzer::Visit_;
explicit Visitor(ffi::Array<ffi::ObjectRef>* output) : output(output) {}
ffi::Array<ffi::ObjectRef>* output;
diff --git a/src/s_tir/analysis/is_pure_function.cc b/src/s_tir/analysis/is_pure_function.cc
index 3c7acb10e6..a5fb6a692c 100644
--- a/src/s_tir/analysis/is_pure_function.cc
+++ b/src/s_tir/analysis/is_pure_function.cc
@@ -24,10 +24,11 @@
#include <tvm/ffi/cast.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/op.h>
+#include <tvm/s_tir/analysis.h>
+#include <tvm/s_tir/stmt_functor.h>
#include <tvm/tirx/analysis.h>
-#include <tvm/tirx/stmt_functor.h>
-#include "../../tirx/ir/tir_visitor_with_path.h"
+#include "../ir/tir_visitor_with_path.h"
namespace tvm {
namespace s_tir {
diff --git a/src/s_tir/analysis/oob_checker.cc b/src/s_tir/analysis/oob_checker.cc
index 55bf73e344..42e3161b51 100644
--- a/src/s_tir/analysis/oob_checker.cc
+++ b/src/s_tir/analysis/oob_checker.cc
@@ -24,7 +24,7 @@
#include <tvm/ffi/reflection/registry.h>
#include <tvm/tirx/transform.h>
-#include "../../tirx/ir/ir_visitor_with_analyzer.h"
+#include "../../s_tir/ir/ir_visitor_with_analyzer.h"
#include "../schedule/error.h"
namespace tvm {
@@ -69,9 +69,9 @@ class OOBError : public s_tir::ScheduleErrorContextObj {
IRModule mod_;
std::vector<OOBLocation> locations_;
};
-class OOBCheckerVisitor final : public tirx::IRVisitorWithAnalyzer {
+class OOBCheckerVisitor final : public s_tir::IRVisitorWithAnalyzer {
public:
- using tirx::IRVisitorWithAnalyzer::Visit_;
+ using s_tir::IRVisitorWithAnalyzer::Visit_;
ffi::Optional<VisitInterrupt> Visit_(const BufferStoreNode* node) final {
for (size_t i = 0; i < node->buffer->shape.size(); i++) {
diff --git a/src/s_tir/analysis/sblock_access_region_detector.cc b/src/s_tir/analysis/sblock_access_region_detector.cc
index 574401da95..b293b0d41d 100644
--- a/src/s_tir/analysis/sblock_access_region_detector.cc
+++ b/src/s_tir/analysis/sblock_access_region_detector.cc
@@ -26,13 +26,15 @@
#include <tvm/ffi/cast.h>
#include <tvm/ffi/extra/structural_mutate.h>
#include <tvm/ffi/reflection/registry.h>
+#include <tvm/s_tir/stmt.h>
+#include <tvm/s_tir/stmt_functor.h>
#include <tvm/tirx/op.h>
-#include <tvm/tirx/stmt_functor.h>
#include <unordered_map>
#include <unordered_set>
#include "../../tirx/transform/ir_utils.h"
+#include "../transform/ir_utils.h"
#include "conditional_bounds.h"
namespace tvm {
@@ -43,9 +45,9 @@ namespace tirx {
* by order of appearance in the AST. \note This detector can only visit blocks and will not visit
* child blocks recursively
*/
-class BlockReadWriteDetector : public StmtExprVisitor {
+class BlockReadWriteDetector : public s_tir::StmtExprVisitor {
public:
- using StmtExprVisitor::Visit_;
+ using s_tir::StmtExprVisitor::Visit_;
explicit BlockReadWriteDetector(const ffi::Map<Var, BufferVar>& buffer_var_map)
: buffer_var_map_(buffer_var_map) {
@@ -92,7 +94,7 @@ class BlockReadWriteDetector : public StmtExprVisitor {
/*! \brief The outside buffer data mapping to its buffer */
ffi::Map<Var, BufferVar> buffer_var_map_;
/*! \brief The target buffer var mapping to its matching */
- std::unordered_map<const VarNode*, MatchBufferRegion> match_buffers_;
+ std::unordered_map<const VarNode*, s_tir::MatchBufferRegion> match_buffers_;
/*! \brief let bindings inside the block */
std::unordered_map<const VarNode*, PrimExpr> let_bindings_;
/*!\ brief Internal analyzer. */
@@ -115,7 +117,7 @@ class BlockReadWriteDetector : public StmtExprVisitor {
const std::unordered_set<const VarNode*>* excluded_buffers = nullptr);
/*! \brief Helper function to convert matched access region to source region. */
- std::vector<arith::IntSet> ConvertMatchedRegion(const MatchBufferRegion& match_buffer,
+ std::vector<arith::IntSet> ConvertMatchedRegion(const s_tir::MatchBufferRegion& match_buffer,
const std::vector<arith::IntSet>& int_sets) const;
/*! \brief Helper function to update a opaque access. */
@@ -136,7 +138,7 @@ class BlockReadWriteDetector : public StmtExprVisitor {
ffi::Optional<VisitInterrupt> Visit_(const ForNode* op) override;
ffi::Optional<VisitInterrupt> Visit_(const IfThenElseNode* op) override;
- ffi::Optional<VisitInterrupt> Visit_(const SBlockRealizeNode* op) override;
+ ffi::Optional<VisitInterrupt> Visit_(const s_tir::SBlockRealizeNode* op) override;
ffi::Optional<VisitInterrupt> Visit_(const DeclBufferNode* op) override;
ffi::Optional<VisitInterrupt> Visit_(const BufferStoreNode* op) override;
ffi::Optional<VisitInterrupt> Visit_(const BindNode* op) override;
@@ -146,10 +148,10 @@ class BlockReadWriteDetector : public StmtExprVisitor {
};
void BlockReadWriteDetector::operator()(const Stmt& stmt) {
- const auto* block = stmt.as<SBlockNode>();
+ const auto* block = stmt.as<s_tir::SBlockNode>();
TVM_FFI_ICHECK(block != nullptr)
<< "Only visiting Blocks is allowed, but got " << stmt->GetTypeKey();
- for (const MatchBufferRegion& match_buffer : block->match_buffers) {
+ for (const s_tir::MatchBufferRegion& match_buffer : block->match_buffers) {
const Var target_var = match_buffer->buffer.var();
const Var source_var = match_buffer->source->source.as_or_throw<tvm::tirx::BufferVar>().var();
if (buffer_var_map_.find(source_var) != buffer_var_map_.end()) {
@@ -157,7 +159,7 @@ void BlockReadWriteDetector::operator()(const Stmt& stmt) {
buffer_var_map_.Set(target_var, match_buffer->buffer);
}
}
- StmtExprVisitor::Visit(stmt);
+ s_tir::StmtExprVisitor::Visit(stmt);
}
ffi::Array<TensorRegion> BlockReadWriteDetector::CollectReads(
@@ -209,7 +211,7 @@ ffi::Optional<VisitInterrupt> BlockReadWriteDetector::Visit_(const TensorLoadNod
ffi::Optional<VisitInterrupt> BlockReadWriteDetector::Visit_(const ForNode* op) {
Range range = Range::FromMinExtent(op->min, op->extent);
dom_map_[op->loop_var.get()] = arith::IntSet::FromRange(range);
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(op));
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(s_tir::StmtExprVisitor::Visit_(op));
dom_map_.erase(op->loop_var.get());
return std::nullopt;
}
@@ -220,13 +222,13 @@ ffi::Optional<VisitInterrupt> BlockReadWriteDetector::Visit_(const IfThenElseNod
// Visit then branch
With<s_tir::ConditionalBoundsContext> ctx(op->condition, &dom_map_, &hint_map_,
&pending_conditions_);
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit(op->then_case));
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(s_tir::StmtExprVisitor::Visit(op->then_case));
}
if (op->else_case) {
// Visit else branch
With<s_tir::ConditionalBoundsContext> ctx(!op->condition, &dom_map_, &hint_map_,
&pending_conditions_);
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit(op->else_case.value()));
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(s_tir::StmtExprVisitor::Visit(op->else_case.value()));
}
return std::nullopt;
}
@@ -243,7 +245,7 @@ ffi::Optional<VisitInterrupt> BlockReadWriteDetector::Visit_(const BindNode* op)
if (auto value = op->value.as<PrimExpr>()) {
let_bindings_[op->var.get()] = value.value();
}
- return StmtExprVisitor::Visit_(op);
+ return s_tir::StmtExprVisitor::Visit_(op);
}
ffi::Optional<VisitInterrupt> BlockReadWriteDetector::Visit_(const CallNode* op) {
@@ -312,7 +314,7 @@ ffi::Optional<VisitInterrupt> BlockReadWriteDetector::Visit_(const CallNode* op)
}
}
} else {
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(op));
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(s_tir::StmtExprVisitor::Visit_(op));
}
return std::nullopt;
}
@@ -324,18 +326,18 @@ ffi::Optional<VisitInterrupt> BlockReadWriteDetector::Visit_(const CallNode* op)
With<s_tir::ConditionalBoundsContext> ctx(condition, &dom_map_, &hint_map_,
&pending_conditions_);
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(
- StmtExprVisitor::Visit(op->args[1].as_or_throw<PrimExpr>()));
+ s_tir::StmtExprVisitor::Visit(op->args[1].as_or_throw<PrimExpr>()));
}
{
// Visit else branch
With<s_tir::ConditionalBoundsContext> ctx(!condition, &dom_map_, &hint_map_,
&pending_conditions_);
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(
- StmtExprVisitor::Visit(op->args[2].as_or_throw<PrimExpr>()));
+ s_tir::StmtExprVisitor::Visit(op->args[2].as_or_throw<PrimExpr>()));
}
return std::nullopt;
}
- return StmtExprVisitor::Visit_(op);
+ return s_tir::StmtExprVisitor::Visit_(op);
}
ffi::Optional<VisitInterrupt> BlockReadWriteDetector::Visit_(const BufferStoreNode* op) {
@@ -364,7 +366,7 @@ ffi::Optional<VisitInterrupt> BlockReadWriteDetector::Visit_(const BufferStoreNo
return std::nullopt;
}
-ffi::Optional<VisitInterrupt> BlockReadWriteDetector::Visit_(const SBlockRealizeNode* op) {
+ffi::Optional<VisitInterrupt> BlockReadWriteDetector::Visit_(const s_tir::SBlockRealizeNode* op) {
/*! \note detector will not visit child block recursively, so it will stop here */
std::unordered_map<const VarNode*, PrimExpr> vmap;
for (size_t i = 0; i < op->block->iter_vars.size(); ++i) {
@@ -404,7 +406,8 @@ ffi::Optional<VisitInterrupt> BlockReadWriteDetector::Visit_(const SBlockRealize
}
std::vector<arith::IntSet> BlockReadWriteDetector::ConvertMatchedRegion(
- const MatchBufferRegion& match_buffer, const std::vector<arith::IntSet>& int_sets) const {
+ const s_tir::MatchBufferRegion& match_buffer,
+ const std::vector<arith::IntSet>& int_sets) const {
const BufferVar& buffer = match_buffer->buffer;
ffi::Array<Range> region;
@@ -432,7 +435,7 @@ void BlockReadWriteDetector::Update(std::vector<BufferVar>* buffers,
// Handle match_buffer remap
auto it = match_buffers_.find(buffer.get());
if (it != match_buffers_.end()) {
- const MatchBufferRegion& match_buffer = it->second;
+ const s_tir::MatchBufferRegion& match_buffer = it->second;
buffer = match_buffer->source->source.as_or_throw<tvm::tirx::BufferVar>();
region = ConvertMatchedRegion(match_buffer, std::move(region));
}
@@ -495,7 +498,7 @@ void BlockReadWriteDetector::UpdateOpaque(const Var& buffer_var) {
}
ffi::Array<ffi::Array<TensorRegion>> GetSBlockAccessRegion(
- const SBlock& block, const ffi::Map<Var, BufferVar>& buffer_var_map) {
+ const s_tir::SBlock& block, const ffi::Map<Var, BufferVar>& buffer_var_map) {
auto detector = ffi::make_object<BlockReadWriteDetector>(buffer_var_map);
detector->operator()(block);
ffi::Array<TensorRegion> writes = detector->CollectWrites();
@@ -512,7 +515,7 @@ ffi::Array<ffi::Array<TensorRegion>> GetSBlockAccessRegion(
}
ffi::Array<ffi::Array<TensorRegion>> GetSBlockReadWriteRegion(
- const SBlock& block, const ffi::Map<Var, BufferVar>& buffer_var_map) {
+ const s_tir::SBlock& block, const ffi::Map<Var, BufferVar>& buffer_var_map) {
auto detector = ffi::make_object<BlockReadWriteDetector>(buffer_var_map);
detector->operator()(block);
ffi::Array<TensorRegion> opaques = detector->CollectOpaques();
diff --git a/src/s_tir/analysis/sblock_buffer_access_lca_detector.cc b/src/s_tir/analysis/sblock_buffer_access_lca_detector.cc
index 568a330dd5..0d73fb0765 100644
--- a/src/s_tir/analysis/sblock_buffer_access_lca_detector.cc
+++ b/src/s_tir/analysis/sblock_buffer_access_lca_detector.cc
@@ -25,8 +25,10 @@
#include <tvm/ffi/cast.h>
#include <tvm/ffi/extra/structural_visit.h>
#include <tvm/ffi/reflection/registry.h>
+#include <tvm/s_tir/analysis.h>
+#include <tvm/s_tir/stmt.h>
+#include <tvm/s_tir/stmt_functor.h>
#include <tvm/tirx/analysis.h>
-#include <tvm/tirx/stmt_functor.h>
#include "../../runtime/thread_storage_scope.h"
#include "../../support/arena.h"
@@ -42,9 +44,9 @@ namespace tirx {
* global memory will have its buffer access LCA outside all launch sites of `blockIdx`, in order to
* prevent conflicts between buffer memory scopes and CUDA hierarchy.
*/
-class LCADetector : public StmtExprVisitor {
+class LCADetector : public s_tir::StmtExprVisitor {
public:
- using StmtExprVisitor::Visit_;
+ using s_tir::StmtExprVisitor::Visit_;
static ffi::Map<BufferVar, ffi::Optional<Stmt>> Detect(const PrimFunc& func) {
auto detector = ffi::make_object<LCADetector>();
@@ -107,14 +109,14 @@ class LCADetector : public StmtExprVisitor {
ancestor_scopes_.push_back(current_scope);
loop_scope_map_.insert({op->loop_var.get(), current_scope});
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(op));
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(s_tir::StmtExprVisitor::Visit_(op));
ancestor_scopes_.pop_back();
loop_scope_map_.erase(op->loop_var.get());
return std::nullopt;
}
- ffi::Optional<VisitInterrupt> Visit_(const SBlockRealizeNode* op) final {
- const SBlockNode* block = op->block.get();
+ ffi::Optional<VisitInterrupt> Visit_(const s_tir::SBlockRealizeNode* op) final {
+ const s_tir::SBlockNode* block = op->block.get();
int n = ancestor_scopes_.size();
for (const BufferVar& buf : block->alloc_buffers) {
buffer_var_map_.emplace(buf.get(), buf.get());
@@ -134,18 +136,18 @@ class LCADetector : public StmtExprVisitor {
UpdateDominateScopeOfNonDataParIter(op);
// Update match_buffers
- for (const MatchBufferRegion& match_buffer : block->match_buffers) {
+ for (const s_tir::MatchBufferRegion& match_buffer : block->match_buffers) {
UpdateBufferLCA(match_buffer->source->source.as_or_throw<tvm::tirx::BufferVar>().get(),
ancestor_scopes_.back());
match_buffers_.insert(match_buffer->buffer.get());
}
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(op));
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(s_tir::StmtExprVisitor::Visit_(op));
ancestor_scopes_.pop_back();
return std::nullopt;
}
- void UpdateDominateScopeOfNonDataParIter(const SBlockRealizeNode* block_realize) {
+ void UpdateDominateScopeOfNonDataParIter(const s_tir::SBlockRealizeNode* block_realize) {
// map iter var to the scope which dominate all loop carried dependencies.
std::unordered_map<const VarNode*, const ScopeInfo*> opaque_var_scope;
// maintain highest scope which dominate all reduce loop iters. null denotes non-reduce block.
@@ -179,7 +181,7 @@ class LCADetector : public StmtExprVisitor {
// collect non-data-parallel block iteration's dominate scope.
// for reduction iter type, we maintain the highest dominate scope for all reduce iters.
// for other iter type, we maintain the dict for each individual iter.
- const SBlock& block = block_realize->block;
+ const s_tir::SBlock& block = block_realize->block;
bool is_reduce_block = false;
for (size_t i = 0; i < block_realize->iter_values.size(); ++i) {
const IterVar& iter_var = block->iter_vars[i];
@@ -263,12 +265,12 @@ class LCADetector : public StmtExprVisitor {
blockidx_scopes_.push_back(ancestor_scopes_.back());
}
}
- return StmtExprVisitor::Visit_(op);
+ return s_tir::StmtExprVisitor::Visit_(op);
}
// Declared regions carry bounds, not opaque runtime accesses.
ffi::Optional<VisitInterrupt> Visit_(const TensorRegionNode* op) final {
- if (!op->source.as<BufferVar>()) return StmtExprVisitor::Visit_(op);
+ if (!op->source.as<BufferVar>()) return s_tir::StmtExprVisitor::Visit_(op);
for (const Range& range : op->region) {
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Visit(range->min));
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Visit(range->extent));
@@ -354,7 +356,7 @@ class LCADetector : public StmtExprVisitor {
return lhs;
}
- /*! \brief The ancestor scope stacks info (SBlock and For). The
+ /*! \brief The ancestor scope stacks info (s_tir::SBlock and For). The
* first element is initialized in LCADetector::Detect to represent
* the root scope.
*/
diff --git a/src/s_tir/analysis/verify_gpu_code.cc b/src/s_tir/analysis/verify_gpu_code.cc
index 17f068c751..b56756ad23 100644
--- a/src/s_tir/analysis/verify_gpu_code.cc
+++ b/src/s_tir/analysis/verify_gpu_code.cc
@@ -27,10 +27,11 @@
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/prim/expr.h>
+#include <tvm/s_tir/analysis.h>
#include <tvm/s_tir/stmt.h>
+#include <tvm/s_tir/stmt_functor.h>
#include <tvm/tirx/analysis.h>
#include <tvm/tirx/stmt.h>
-#include <tvm/tirx/stmt_functor.h>
#include "../../runtime/thread_storage_scope.h"
#include "../../tirx/transform/ir_utils.h"
diff --git a/src/s_tir/analysis/verify_well_formed.cc b/src/s_tir/analysis/verify_well_formed.cc
new file mode 100644
index 0000000000..b1a6439fae
--- /dev/null
+++ b/src/s_tir/analysis/verify_well_formed.cc
@@ -0,0 +1,168 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+#include "../../tirx/analysis/verify_well_formed.h"
+
+#include <tvm/ffi/reflection/registry.h>
+#include <tvm/s_tir/analysis.h>
+#include <tvm/s_tir/stmt_functor.h>
+
+#include "../ir/tir_visitor_with_path.h"
+
+namespace tvm {
+namespace s_tir {
+using tirx::ForNode;
+using tirx::PrimFunc;
+
+/*! \brief Verify all Expr inside the block does not contain:
+ * 1. loop vars outside the current block.
+ * 2. block vars of parent blocks.
+ */
+class BlockVarAccessVerifier : public StmtExprVisitor {
+ public:
+ static bool Verify(const PrimFunc& func, bool assert_mode) {
+ auto verifier = ffi::make_object<BlockVarAccessVerifier>(assert_mode);
+ verifier->Visit(func->body);
+ return !verifier->has_error_;
+ }
+
+ explicit BlockVarAccessVerifier(bool assert_mode) : assert_mode_(assert_mode) {}
+
+ private:
+ ffi::Optional<VisitInterrupt> Visit(ffi::AnyView stmt) final {
+ if (!has_error_) {
+ return StmtExprVisitor::Visit(stmt);
+ }
+ return std::nullopt;
+ }
+
+ ffi::Optional<VisitInterrupt> Visit_(const VarNode* op) final {
+ auto it = loop_vars_.find(op);
+ if (it != loop_vars_.end() && it->second < block_stack_.size()) {
+ has_error_ = true;
+ if (assert_mode_) {
+ if (it->second == 0) {
+ TVM_FFI_THROW(InternalError)
+ << "Well-formedness check failed: "
+ << "Loop iterator var " << op->name << " is defined outside of any block, "
+ << "but is used inside the non-opaque current block \""
+ << block_stack_.back()->name_hint << "\".";
+ } else {
+ TVM_FFI_THROW(InternalError)
+ << "Well-formedness check failed: "
+ << "Loop iterator var " << op->name << " is defined in block \""
+ << block_stack_[it->second - 1]->name_hint << "\", "
+ << "but is used inside the non-opaque current block \""
+ << block_stack_.back()->name_hint << "\".";
+ }
+ }
+ }
+ return std::nullopt;
+ }
+
+ ffi::Optional<VisitInterrupt> Visit_(const ForNode* op) final {
+ TVM_FFI_ICHECK(loop_vars_.find(op->loop_var.get()) == loop_vars_.end());
+ loop_vars_[op->loop_var.get()] = block_stack_.size();
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(op));
+ loop_vars_.erase(op->loop_var.get());
+ return std::nullopt;
+ }
+
+ ffi::Optional<VisitInterrupt> Visit_(const SBlockNode* op) final {
+ // Do not check boundary if it's a opaque block.
+ bool is_non_opaque = op->iter_vars.size();
+ if (is_non_opaque) {
+ block_stack_.push_back(op);
+ }
+
+ // Step 0. Skip block iter var's domain
+
+ // Step 1. Visit read/write regions
+ auto fvisit_buffer_region = [this](const TensorRegion& s) -> ffi::Optional<VisitInterrupt> {
+ for (const auto& range : s->region) {
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->Visit(range->min));
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->Visit(range->extent));
+ }
+ return std::nullopt;
+ };
+ for (const auto& region : op->reads) {
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(fvisit_buffer_region(region));
+ }
+ for (const auto& region : op->writes) {
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(fvisit_buffer_region(region));
+ }
+
+ // Step 2. Visit match buffers
+ for (const auto& match_buffer_region : op->match_buffers) {
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(fvisit_buffer_region(match_buffer_region->source));
+ }
+
+ // Step 3. Visit init and body
+ if (op->init.has_value()) {
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->Visit(op->init.value()));
+ }
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->Visit(op->body));
+
+ if (is_non_opaque) {
+ block_stack_.pop_back();
+ }
+ return std::nullopt;
+ }
+
+ private:
+ /*! \brief The map from outside loop vars to its corresponding block level. */
+ std::unordered_map<const VarNode*, size_t> loop_vars_;
+ /*! \brief Whether it's in assert mode. */
+ bool assert_mode_;
+ /*! \brief Current nested block stack level. */
+ std::vector<const SBlockNode*> block_stack_;
+ /*! \brief Whether there is error. */
+ bool has_error_{false};
+};
+
+bool VerifyWellFormed(const tirx::PrimFunc& func, bool assert_mode) {
+ return BlockVarAccessVerifier::Verify(func, assert_mode) &&
+ tirx::VerifyWellFormedCommon<TIRVisitorWithPath>(func, assert_mode);
+}
+
+bool VerifyWellFormed(const IRModule& mod, bool assert_mode) {
+ for (const auto& [gvar, base_func] : mod->functions) {
+ if (auto func = base_func.as<tirx::PrimFunc>()) {
+ if (!BlockVarAccessVerifier::Verify(func.value(), assert_mode)) return false;
+ }
+ }
+ return tirx::VerifyWellFormedCommon<TIRVisitorWithPath>(mod, assert_mode);
+}
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+ ffi::reflection::GlobalDef().def(
+ "s_tir.analysis.VerifyWellFormed", [](const ffi::ObjectRef& obj, bool assert_mode) {
+ if (auto func = obj.as<tirx::PrimFunc>()) {
+ return s_tir::VerifyWellFormed(func.value(), assert_mode);
+ }
+ if (auto mod = obj.as<IRModule>()) {
+ return s_tir::VerifyWellFormed(mod.value(), assert_mode);
+ }
+ TVM_FFI_THROW(TypeError) << "Expected a PrimFunc or IRModule, but received "
+ << obj->GetTypeKey();
+ TVM_FFI_UNREACHABLE();
+ });
+}
+} // namespace s_tir
+} // namespace tvm
diff --git a/src/s_tir/backend/adreno/inject_texture_alloc.cc b/src/s_tir/backend/adreno/inject_texture_alloc.cc
index fa3c2925f6..e9403acaeb 100644
--- a/src/s_tir/backend/adreno/inject_texture_alloc.cc
+++ b/src/s_tir/backend/adreno/inject_texture_alloc.cc
@@ -22,12 +22,13 @@
*/
#include <tvm/arith/iter_affine_map.h>
+#include <tvm/s_tir/analysis.h>
#include <tvm/s_tir/backend/adreno/transform.h>
+#include <tvm/s_tir/stmt_functor.h>
#include <tvm/tirx/analysis.h>
-#include <tvm/tirx/stmt_functor.h>
#include "../../../backend/opencl/runtime/texture.h"
-#include "../../../tirx/ir/ir_mutator_with_analyzer.h"
+#include "../../../s_tir/ir/ir_mutator_with_analyzer.h"
#include "../../../tirx/transform/ir_utils.h"
namespace tvm {
@@ -42,10 +43,10 @@ using runtime::IsTextureStorage;
/*!
* \brief Inject Texture Alloc Intrinsic right after AllocBufferNode are realized.
*/
-class TextureAllocInjector : public tirx::IRMutatorWithAnalyzer {
+class TextureAllocInjector : public s_tir::IRMutatorWithAnalyzer {
public:
- using tirx::IRMutatorWithAnalyzer::Mutate;
- using tirx::IRMutatorWithAnalyzer::Mutate_;
+ using s_tir::IRMutatorWithAnalyzer::Mutate;
+ using s_tir::IRMutatorWithAnalyzer::Mutate_;
static PrimFunc Inject(PrimFunc func) {
arith::Analyzer ana;
diff --git a/src/s_tir/backend/adreno/texture_flatten.cc b/src/s_tir/backend/adreno/texture_flatten.cc
index d6abeef17f..a03f6ced7a 100644
--- a/src/s_tir/backend/adreno/texture_flatten.cc
+++ b/src/s_tir/backend/adreno/texture_flatten.cc
@@ -35,7 +35,7 @@
#include "../../../backend/opencl/runtime/texture.h"
#include "../../../runtime/thread_storage_scope.h"
-#include "../../../tirx/ir/ir_visitor_with_analyzer.h"
+#include "../../../s_tir/ir/ir_visitor_with_analyzer.h"
namespace tvm {
namespace s_tir {
diff --git a/src/s_tir/data_layout.cc b/src/s_tir/data_layout.cc
index 19819d1c87..9000d298e5 100644
--- a/src/s_tir/data_layout.cc
+++ b/src/s_tir/data_layout.cc
@@ -29,9 +29,10 @@
#include <tvm/ir/expr.h>
#include <tvm/ir/prim/expr.h>
#include <tvm/runtime/logging.h>
+#include <tvm/s_tir/analysis.h>
#include <tvm/s_tir/data_layout.h>
+#include <tvm/s_tir/stmt_functor.h>
#include <tvm/tirx/analysis.h>
-#include <tvm/tirx/stmt_functor.h>
#include <tvm/tirx/var.h>
#include <algorithm>
diff --git a/src/s_tir/ir/data_type_rewriter.cc b/src/s_tir/ir/data_type_rewriter.cc
new file mode 100644
index 0000000000..edf356e1e5
--- /dev/null
+++ b/src/s_tir/ir/data_type_rewriter.cc
@@ -0,0 +1,241 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+#include "data_type_rewriter.h"
+
+#include <tvm/tirx/op.h>
+
+#include <functional>
+
+namespace tvm {
+namespace s_tir {
+using namespace tvm::tirx;
+using namespace tvm::prim;
+
+PrimFunc IndexDataTypeNormalizer::Rewrite(PrimFunc func) {
+ // Keep this short setup local so its collector uses S-TIR block semantics
+ // without adding a dialect-specific collector hook to the TIRX normalizer.
+ // Collect scalar dtype requirements without changing types. Buffer definitions
+ // are rewritten only after every scalar replacement has been seeded.
+ class IndexVarCollector : public IndexDataTypeNormalizer {
+ public:
+ explicit IndexVarCollector(std::function<void(const VarNode*)> collect)
+ : IndexDataTypeNormalizer(PrimType::Int(64)), collect_(std::move(collect)) {}
+ using IndexDataTypeNormalizer::Mutate;
+ using IndexDataTypeNormalizer::Mutate_;
+ UnchangedOr<Expr> Mutate_(const VarNode* op, InplaceMode mode) final {
+ if (def_region_kind() == kTVMFFIDefRegionKindNone && is_enabled_) collect_(op);
+ return IndexDataTypeNormalizer::Mutate_(op, mode);
+ }
+
+ protected:
+ bool CanRewriteDType(PrimType dtype) const final { return false; }
+
+ private:
+ std::function<void(const VarNode*)> collect_;
+ };
+ auto seed = [this](const VarNode* var) {
+ auto dtype = var->ty.as<PrimType>();
+ if (dtype && CanRewriteDType(dtype.value()) && dtype.value() != target_data_type_ &&
+ VarRemapGet(ffi::AnyView(var)) == nullptr) {
+ VarRemapSet(ffi::AnyView(var), ffi::GetRef<Var>(var).CopyWithDType(target_data_type_));
+ }
+ };
+ auto collector = ffi::make_object<IndexVarCollector>(seed);
+ collector->Mutate(func->body);
+ for (const Var& param : func->params) {
+ if (param.as<BufferVar>()) {
+ collector->WithDefRegionKind(kTVMFFIDefRegionKindSimple,
+ [&] { return collector->Mutate(param); });
+ } else {
+ seed(param.get());
+ }
+ }
+ ffi::Array<Var> params = func->params.Map([this](const Var& param) {
+ return WithDefRegionKind(kTVMFFIDefRegionKindSimple, [&] {
+ return Mutate(param).ValueOrUnchanged(param).as_or_throw<Var>();
+ });
+ });
+ PrimFuncNode* new_func = func.CopyOnWrite();
+ new_func->params = std::move(params);
+ new_func->body = Mutate(new_func->body).ValueOrUnchanged(new_func->body);
+ return func;
+}
+
+UnchangedOr<Stmt> IndexDataTypeNormalizer::Mutate_(const SBlockRealizeNode* op,
+ InplaceMode inplace_mode) {
+ bool is_condition = this->is_condition_;
+ this->is_condition_ = true;
+ auto new_predicate_result = this->Mutate(op->predicate, inplace_mode);
+ bool new_predicate_unchanged = new_predicate_result.UnchangedOrSameAs(op->predicate);
+ auto new_predicate = std::move(new_predicate_result).ValueOrUnchanged(op->predicate);
+ this->is_condition_ = is_condition;
+
+ bool is_enabled = this->is_enabled_;
+ this->is_enabled_ = true;
+ auto new_iter_values = this->Mutate(op->iter_values, inplace_mode)
+ .as_or_throw<UnchangedOr<ffi::Array<PrimExpr>>>()
+ .ValueOrUnchanged(op->iter_values);
+ this->is_enabled_ = is_enabled;
+ SBlock new_body =
+ this->Mutate(op->block, inplace_mode).ValueOrUnchanged(op->block).as_or_throw<SBlock>();
+ if (!new_predicate_unchanged || !new_iter_values.same_as(op->iter_values) ||
+ !new_body.same_as(op->block)) {
+ SBlockRealize new_block_realize = ffi::GetRef<SBlockRealize>(op);
+ auto* n = new_block_realize.CopyOnWrite();
+ n->predicate = std::move(new_predicate);
+ n->iter_values = std::move(new_iter_values);
+ n->block = std::move(new_body);
+ return new_block_realize;
+
+ } else {
+ return ffi::Unchanged();
+ }
+}
+
+UnchangedOr<Stmt> IndexDataTypeNormalizer::Mutate_(const SBlockNode* op, InplaceMode inplace_mode) {
+ auto new_alloc_buffers = this->WithDefRegionKind(kTVMFFIDefRegionKindSimple, [&] {
+ return this->Mutate(op->alloc_buffers, inplace_mode)
+ .as_or_throw<UnchangedOr<ffi::Array<BufferVar>>>()
+ .ValueOrUnchanged(op->alloc_buffers);
+ });
+ auto new_match_buffers = op->match_buffers.Map([this](const MatchBufferRegion& match) {
+ BufferVar buffer = this->WithDefRegionKind(kTVMFFIDefRegionKindSimple, [&] {
+ return this->Mutate(match->buffer, InplaceMode::kDisallow)
+ .as_or_throw<UnchangedOr<BufferVar>>()
+ .ValueOrUnchanged(match->buffer);
+ });
+ TensorRegion source = VisitBufferRegion(match->source);
+ if (buffer.same_as(match->buffer) && source.same_as(match->source)) return match;
+ return MatchBufferRegion(buffer, source);
+ });
+ ffi::Array<TensorRegion> new_reads = op->reads.Map(
+ [this](const TensorRegion& buffer_region) { return VisitBufferRegion(buffer_region); });
+ ffi::Array<TensorRegion> new_writes = op->writes.Map(
+ [this](const TensorRegion& buffer_region) { return VisitBufferRegion(buffer_region); });
+ ffi::Array<IterVar> new_iter_vars =
+ op->iter_vars.Map([this](const IterVar& iter_var) { return VisitIterVar(iter_var); });
+ ffi::Optional<Stmt> new_init = std::nullopt;
+ if (op->init.has_value()) {
+ new_init = this->Mutate(op->init.value(), inplace_mode).ValueOrUnchanged(op->init.value());
+ }
+ ffi::Map<ffi::String, ffi::Any> new_annotations = VisitBlockAnnotations(op->annotations);
+ auto new_body_result = this->Mutate(op->body, inplace_mode);
+ bool new_body_unchanged = new_body_result.UnchangedOrSameAs(op->body);
+ Stmt new_body = std::move(new_body_result).ValueOrUnchanged(op->body);
+
+ if (!new_init.same_as(op->init) || !new_body_unchanged ||
+ !new_alloc_buffers.same_as(op->alloc_buffers) ||
+ !new_match_buffers.same_as(op->match_buffers) || !new_reads.same_as(op->reads) ||
+ !new_writes.same_as(op->writes) || !new_iter_vars.same_as(op->iter_vars) ||
+ !new_annotations.same_as(op->annotations)) {
+ SBlock new_block = ffi::GetRef<SBlock>(op);
+ SBlockNode* n = new_block.CopyOnWrite();
+ n->alloc_buffers = std::move(new_alloc_buffers);
+ n->match_buffers = std::move(new_match_buffers);
+ n->reads = std::move(new_reads);
+ n->writes = std::move(new_writes);
+ n->iter_vars = std::move(new_iter_vars);
+ n->init = std::move(new_init);
+ n->annotations = std::move(new_annotations);
+ n->body = std::move(new_body);
+ return new_block;
+ }
+ return ffi::Unchanged();
+}
+
+ffi::Map<ffi::String, ffi::Any> IndexDataTypeNormalizer::VisitBlockAnnotations(
+ const ffi::Map<ffi::String, ffi::Any>& annotations) {
+ auto new_annotations = annotations;
+
+ std::function<Any(const Any&)> f_mutate_obj = [this, &f_mutate_obj](const Any& obj) -> Any {
+ if (obj == nullptr) {
+ return obj;
+ }
+ if (auto var = obj.as<Var>(); var && var.value()->ty.as<BufferTypeNode>()) {
+ BufferVar buffer(var.value());
+ if (BufferVar new_buffer = this->Mutate(buffer, InplaceMode::kDisallow)
+ .as_or_throw<UnchangedOr<BufferVar>>()
+ .ValueOrUnchanged(buffer);
+ !new_buffer.same_as(buffer)) {
+ return new_buffer;
+ }
+ } else if (obj.as<ffi::ArrayObj>()) {
+ return obj.as_or_throw<ffi::Array<Any>>().Map(f_mutate_obj);
+ }
+ return obj;
+ };
+ for (const auto& [key, value] : annotations) {
+ if (auto opt_object_ref = value.as<ffi::ObjectRef>()) {
+ auto new_value = f_mutate_obj(*opt_object_ref);
+ if (!new_value.same_as(*opt_object_ref)) {
+ new_annotations.Set(key, new_value);
+ }
+ }
+ }
+ return new_annotations;
+}
+
+IterVar IndexDataTypeNormalizer::VisitIterVar(const IterVar& iter_var) {
+ bool is_enabled = this->is_enabled_;
+ this->is_enabled_ = true;
+ PrimVar new_var = this->Mutate(iter_var->var, InplaceMode::kDisallow)
+ .ValueOrUnchanged(iter_var->var)
+ .as_or_throw<PrimVar>();
+ PrimExpr min =
+ this->Mutate(iter_var->dom->min, InplaceMode::kDisallow).ValueOrUnchanged(iter_var->dom->min);
+ PrimExpr extent = this->Mutate(iter_var->dom->extent, InplaceMode::kDisallow)
+ .ValueOrUnchanged(iter_var->dom->extent);
+ this->is_enabled_ = is_enabled;
+ if (!new_var.same_as(iter_var->var) || !min.same_as(iter_var->dom->min) ||
+ !extent.same_as(iter_var->dom->extent)) {
+ IterVar new_iter_var = iter_var;
+ IterVarNode* n = new_iter_var.CopyOnWrite();
+ n->var = std::move(new_var);
+ n->dom = Range(min, extent);
+ return new_iter_var;
+ }
+ return iter_var;
+}
+
+TensorRegion IndexDataTypeNormalizer::VisitBufferRegion(const TensorRegion& buffer_region) {
+ BufferVar remapped_buffer =
+ this->Mutate(buffer_region->source.as_or_throw<BufferVar>(), InplaceMode::kDisallow)
+ .as_or_throw<UnchangedOr<BufferVar>>()
+ .ValueOrUnchanged(buffer_region->source.as_or_throw<BufferVar>());
+
+ bool is_enabled = this->is_enabled_;
+ this->is_enabled_ = true;
+ auto new_region = buffer_region->region.Map([&](const Range& range) {
+ return Range::FromMinExtent(
+ this->Mutate(range->min, InplaceMode::kDisallow).ValueOrUnchanged(range->min),
+ this->Mutate(range->extent, InplaceMode::kDisallow).ValueOrUnchanged(range->extent));
+ });
+ this->is_enabled_ = is_enabled;
+
+ if (!remapped_buffer.same_as(buffer_region->source.as_or_throw<BufferVar>()) ||
+ !new_region.same_as(buffer_region->region)) {
+ return BufferRegion(remapped_buffer, new_region);
+ } else {
+ return buffer_region;
+ }
+}
+
+} // namespace s_tir
+} // namespace tvm
diff --git a/src/s_tir/ir/data_type_rewriter.h b/src/s_tir/ir/data_type_rewriter.h
new file mode 100644
index 0000000000..cf0c21e796
--- /dev/null
+++ b/src/s_tir/ir/data_type_rewriter.h
@@ -0,0 +1,69 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+#ifndef TVM_S_TIR_IR_DATA_TYPE_REWRITER_H_
+#define TVM_S_TIR_IR_DATA_TYPE_REWRITER_H_
+
+#include <tvm/s_tir/stmt.h>
+
+#include "../../tirx/ir/data_type_rewriter.h"
+
+namespace tvm {
+namespace s_tir {
+
+// Explicit normalization for schedulable TE functions and tensor intrinsic bodies.
+// Ordinary TIRX dtype passes run after block lowering.
+class IndexDataTypeNormalizer : public tirx::IndexDataTypeNormalizer {
+ public:
+ using Parent = tirx::IndexDataTypeNormalizer;
+ using Parent::Mutate;
+ using Parent::Mutate_;
+ explicit IndexDataTypeNormalizer(PrimType target_data_type)
+ : Parent(std::move(target_data_type), GlobalVTable()) {}
+ tirx::PrimFunc Rewrite(tirx::PrimFunc func);
+
+ virtual UnchangedOr<tirx::Stmt> Mutate_(const SBlockNode* op, InplaceMode inplace_mode);
+ virtual UnchangedOr<tirx::Stmt> Mutate_(const SBlockRealizeNode* op, InplaceMode inplace_mode);
+
+ protected:
+ static void InitVTable(VTable* table) {
+ Parent::InitVTable(table);
+ SetDispatch<IndexDataTypeNormalizer, SBlockNode>(table);
+ SetDispatch<IndexDataTypeNormalizer, SBlockRealizeNode>(table);
+ }
+ static const VTable* GlobalVTable() {
+ static const VTable table = [] {
+ VTable table;
+ InitVTable(&table);
+ table.Finalize();
+ return table;
+ }();
+ return &table;
+ }
+
+ private:
+ ffi::Map<ffi::String, ffi::Any> VisitBlockAnnotations(
+ const ffi::Map<ffi::String, ffi::Any>& annotations);
+ tirx::IterVar VisitIterVar(const tirx::IterVar& iter_var);
+ TensorRegion VisitBufferRegion(const TensorRegion& buffer_region);
+};
+
+} // namespace s_tir
+} // namespace tvm
+#endif // TVM_S_TIR_IR_DATA_TYPE_REWRITER_H_
diff --git a/src/tirx/transform/stmt_simplify.h b/src/s_tir/ir/ir_mutator_with_analyzer.cc
similarity index 61%
copy from src/tirx/transform/stmt_simplify.h
copy to src/s_tir/ir/ir_mutator_with_analyzer.cc
index 224df0ed8b..bd68223887 100644
--- a/src/tirx/transform/stmt_simplify.h
+++ b/src/s_tir/ir/ir_mutator_with_analyzer.cc
@@ -17,25 +17,20 @@
* under the License.
*/
-/*!
- * \file stmt_simplify.h
- * \brief Statement-level simplification of TIR PrimFuncs.
- */
-#ifndef TVM_TIR_TRANSFORM_STMT_SIMPLIFY_H_
-#define TVM_TIR_TRANSFORM_STMT_SIMPLIFY_H_
-
-#include <tvm/arith/analyzer.h>
-#include <tvm/tirx/function.h>
+#include "ir_mutator_with_analyzer.h"
namespace tvm {
-namespace tirx {
-
-/* \brief Simplify statements in the prim func
- *
- * Applies the same behavior as the tirx.transform.StmtSimplify pass.
- */
-PrimFunc StmtSimplify(PrimFunc func, const arith::Analyzer& analyzer);
+namespace s_tir {
+using namespace tirx;
+UnchangedOr<Stmt> IRMutatorWithAnalyzer::Mutate_(const SBlockNode* op, InplaceMode inplace_mode) {
+ return constraint_scope_.WithNewScope([&]() -> UnchangedOr<Stmt> {
+ for (const auto& iter_var : op->iter_vars) {
+ analyzer_->Bind(iter_var->var, iter_var->dom);
+ iter_vars_.Set(iter_var->var, iter_var->dom);
+ }
+ return s_tir::StmtExprMutator::MutateBlock(this, op, inplace_mode);
+ });
+}
-} // namespace tirx
+} // namespace s_tir
} // namespace tvm
-#endif // TVM_TIR_TRANSFORM_STMT_SIMPLIFY_H_
diff --git a/src/s_tir/ir/ir_mutator_with_analyzer.h b/src/s_tir/ir/ir_mutator_with_analyzer.h
new file mode 100644
index 0000000000..49179ab43c
--- /dev/null
+++ b/src/s_tir/ir/ir_mutator_with_analyzer.h
@@ -0,0 +1,60 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+#ifndef TVM_S_TIR_IR_MUTATOR_WITH_ANALYZER_H_
+#define TVM_S_TIR_IR_MUTATOR_WITH_ANALYZER_H_
+
+#include <tvm/s_tir/stmt_functor.h>
+
+#include "../../tirx/ir/ir_mutator_with_analyzer.h"
+
+namespace tvm {
+namespace s_tir {
+class IRMutatorWithAnalyzer : public tirx::IRMutatorWithAnalyzer {
+ public:
+ using Parent = tirx::IRMutatorWithAnalyzer;
+ using Parent::Mutate;
+ using Parent::Mutate_;
+ explicit IRMutatorWithAnalyzer(const arith::Analyzer& analyzer)
+ : IRMutatorWithAnalyzer(analyzer.get()) {}
+ explicit IRMutatorWithAnalyzer(arith::AnalyzerObj* analyzer) : Parent(analyzer, GlobalVTable()) {}
+ virtual UnchangedOr<tirx::Stmt> Mutate_(const SBlockNode* op, InplaceMode inplace_mode);
+ virtual UnchangedOr<tirx::Stmt> Mutate_(const SBlockRealizeNode* op, InplaceMode inplace_mode) {
+ return s_tir::StmtExprMutator::MutateBlockRealize(this, op, inplace_mode);
+ }
+
+ protected:
+ static void InitVTable(VTable* vtable) {
+ Parent::InitVTable(vtable);
+ SetDispatch<IRMutatorWithAnalyzer, SBlockNode>(vtable);
+ SetDispatch<IRMutatorWithAnalyzer, SBlockRealizeNode>(vtable);
+ }
+ static const VTable* GlobalVTable() {
+ static const VTable table = [] {
+ VTable table;
+ InitVTable(&table);
+ table.Finalize();
+ return table;
+ }();
+ return &table;
+ }
+};
+} // namespace s_tir
+} // namespace tvm
+#endif // TVM_S_TIR_IR_MUTATOR_WITH_ANALYZER_H_
diff --git a/src/tirx/transform/stmt_simplify.h b/src/s_tir/ir/ir_visitor_with_analyzer.cc
similarity index 61%
copy from src/tirx/transform/stmt_simplify.h
copy to src/s_tir/ir/ir_visitor_with_analyzer.cc
index 224df0ed8b..56834a8b98 100644
--- a/src/tirx/transform/stmt_simplify.h
+++ b/src/s_tir/ir/ir_visitor_with_analyzer.cc
@@ -17,25 +17,19 @@
* under the License.
*/
-/*!
- * \file stmt_simplify.h
- * \brief Statement-level simplification of TIR PrimFuncs.
- */
-#ifndef TVM_TIR_TRANSFORM_STMT_SIMPLIFY_H_
-#define TVM_TIR_TRANSFORM_STMT_SIMPLIFY_H_
-
-#include <tvm/arith/analyzer.h>
-#include <tvm/tirx/function.h>
+#include "ir_visitor_with_analyzer.h"
namespace tvm {
-namespace tirx {
-
-/* \brief Simplify statements in the prim func
- *
- * Applies the same behavior as the tirx.transform.StmtSimplify pass.
- */
-PrimFunc StmtSimplify(PrimFunc func, const arith::Analyzer& analyzer);
+namespace s_tir {
+using namespace tirx;
+ffi::Optional<VisitInterrupt> IRVisitorWithAnalyzer::Visit_(const SBlockNode* op) {
+ return constraint_scope_.WithNewScope([&]() -> ffi::Optional<VisitInterrupt> {
+ for (const auto& iter_var : op->iter_vars) {
+ analyzer_->Bind(iter_var->var, iter_var->dom);
+ }
+ return s_tir::StmtExprVisitor::VisitBlock(this, op);
+ });
+}
-} // namespace tirx
+} // namespace s_tir
} // namespace tvm
-#endif // TVM_TIR_TRANSFORM_STMT_SIMPLIFY_H_
diff --git a/src/s_tir/ir/ir_visitor_with_analyzer.h b/src/s_tir/ir/ir_visitor_with_analyzer.h
new file mode 100644
index 0000000000..ea52dba33f
--- /dev/null
+++ b/src/s_tir/ir/ir_visitor_with_analyzer.h
@@ -0,0 +1,49 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+#ifndef TVM_S_TIR_IR_VISITOR_WITH_ANALYZER_H_
+#define TVM_S_TIR_IR_VISITOR_WITH_ANALYZER_H_
+
+#include <tvm/s_tir/stmt_functor.h>
+
+#include "../../tirx/ir/ir_visitor_with_analyzer.h"
+
+namespace tvm {
+namespace s_tir {
+class IRVisitorWithAnalyzer : public tirx::IRVisitorWithAnalyzer {
+ public:
+ using Parent = tirx::IRVisitorWithAnalyzer;
+ using Parent::Visit;
+ using Parent::Visit_;
+ TVM_DEFINE_OBJECT_FUNCTOR_DEFAULT_CONSTRUCTOR(IRVisitorWithAnalyzer, Parent)
+ virtual ffi::Optional<VisitInterrupt> Visit_(const SBlockNode* op);
+ virtual ffi::Optional<VisitInterrupt> Visit_(const SBlockRealizeNode* op) {
+ return s_tir::StmtExprVisitor::VisitBlockRealize(this, op);
+ }
+
+ protected:
+ static void InitVTable(VTable* vtable) {
+ Parent::InitVTable(vtable);
+ SetDispatch<IRVisitorWithAnalyzer, SBlockNode>(vtable);
+ SetDispatch<IRVisitorWithAnalyzer, SBlockRealizeNode>(vtable);
+ }
+};
+} // namespace s_tir
+} // namespace tvm
+#endif // TVM_S_TIR_IR_VISITOR_WITH_ANALYZER_H_
diff --git a/src/s_tir/ir/tir_visitor_with_path.cc b/src/s_tir/ir/tir_visitor_with_path.cc
new file mode 100644
index 0000000000..af69a4b4fe
--- /dev/null
+++ b/src/s_tir/ir/tir_visitor_with_path.cc
@@ -0,0 +1,81 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+#include "tir_visitor_with_path.h"
+
+#include <tvm/s_tir/stmt.h>
+
+namespace tvm {
+namespace s_tir {
+using namespace tirx;
+using AccessPath = ffi::reflection::AccessPath;
+
+void TIRVisitorWithPath::Dispatch_(const SBlockNode* op, AccessPath path) {
+ std::vector<std::variant<DefContext<Var>, DefContext<IterVar>, DefContext<BufferVar>>> context;
+
+ {
+ auto iter_path = path->Attr("iter_vars");
+ for (size_t i = 0; i < op->iter_vars.size(); i++) {
+ context.push_back(WithDef(op->iter_vars[i], iter_path->ArrayItem(i)));
+ }
+ }
+
+ // Define alloc_buffers before visiting reads/writes, since reads/writes
+ // may reference buffers from alloc_buffers (e.g. after transform_layout).
+ {
+ auto alloc_path = path->Attr("alloc_buffers");
+ for (size_t i = 0; i < op->alloc_buffers.size(); i++) {
+ auto buffer_path = alloc_path->ArrayItem(i);
+ auto buf = op->alloc_buffers[i];
+ context.push_back(WithDef(buf, buffer_path));
+ }
+ }
+
+ {
+ auto match_path = path->Attr("match_buffers");
+ for (size_t i = 0; i < op->match_buffers.size(); i++) {
+ Visit(op->match_buffers[i]->source, match_path->ArrayItem(i)->Attr("source"));
+ auto buf = op->match_buffers[i]->buffer;
+ auto buffer_path = match_path->ArrayItem(i)->Attr("buffer");
+
+ for (auto& def : WithMatchBufferDefs(buf, buffer_path)) {
+ context.push_back(std::move(def));
+ }
+ context.push_back(WithDef(buf, buffer_path));
+ }
+ }
+
+ // Regions may use allocation and match-buffer definitions in this block.
+ Visit(op->reads, path->Attr("reads"));
+ Visit(op->writes, path->Attr("writes"));
+
+ bind_scope_.WithNewScope([&]() { Visit(op->init, path->Attr("init")); });
+ bind_scope_.WithNewScope([&]() { Visit(op->body, path->Attr("body")); });
+
+ while (context.size()) context.pop_back();
+}
+
+void TIRVisitorWithPath::Dispatch_(const SBlockRealizeNode* op, AccessPath path) {
+ Visit(op->iter_values, path->Attr("iter_values"));
+ Visit(op->predicate, path->Attr("predicate"));
+ Visit(op->block, path->Attr("block"));
+}
+
+} // namespace s_tir
+} // namespace tvm
diff --git a/src/s_tir/ir/tir_visitor_with_path.h b/src/s_tir/ir/tir_visitor_with_path.h
new file mode 100644
index 0000000000..6556d2eec7
--- /dev/null
+++ b/src/s_tir/ir/tir_visitor_with_path.h
@@ -0,0 +1,54 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+#ifndef TVM_S_TIR_IR_TIR_VISITOR_WITH_PATH_H_
+#define TVM_S_TIR_IR_TIR_VISITOR_WITH_PATH_H_
+#include <tvm/s_tir/stmt.h>
+
+#include "../../tirx/ir/tir_visitor_with_path.h"
+namespace tvm {
+namespace s_tir {
+class TIRVisitorWithPath : public tirx::TIRVisitorWithPath {
+ public:
+ using Parent = tirx::TIRVisitorWithPath;
+ TVM_DEFINE_OBJECT_FUNCTOR_DEFAULT_CONSTRUCTOR(TIRVisitorWithPath, Parent)
+ protected:
+ using AccessPath = ffi::reflection::AccessPath;
+ using Parent::Dispatch_;
+ virtual void Dispatch_(const SBlockNode* op, AccessPath path);
+ virtual void Dispatch_(const SBlockRealizeNode* op, AccessPath path);
+ static void InitVTable(VTable* vtable) {
+ Parent::InitVTable(vtable);
+ // StmtVisitor is a protected base of the path visitor, so create the
+ // downcast thunks here instead of in StmtVisitor::SetDispatch.
+ vtable->SetDispatch<SBlockNode>(
+ [](const ffi::ObjectRef& node, StmtVisitor* self, AccessPath path) {
+ static_cast<TIRVisitorWithPath*>(self)->Dispatch_(
+ static_cast<const SBlockNode*>(node.get()), path);
+ });
+ vtable->SetDispatch<SBlockRealizeNode>(
+ [](const ffi::ObjectRef& node, StmtVisitor* self, AccessPath path) {
+ static_cast<TIRVisitorWithPath*>(self)->Dispatch_(
+ static_cast<const SBlockRealizeNode*>(node.get()), path);
+ });
+ }
+};
+} // namespace s_tir
+} // namespace tvm
+#endif // TVM_S_TIR_IR_TIR_VISITOR_WITH_PATH_H_
diff --git a/src/s_tir/meta_schedule/feature_extractor/per_store_feature.cc b/src/s_tir/meta_schedule/feature_extractor/per_store_feature.cc
index 9ad9fc5a71..123abd0d84 100644
--- a/src/s_tir/meta_schedule/feature_extractor/per_store_feature.cc
+++ b/src/s_tir/meta_schedule/feature_extractor/per_store_feature.cc
@@ -19,6 +19,7 @@
#include <tvm/ffi/cast.h>
#include <tvm/ffi/extra/structural_visit.h>
#include <tvm/ffi/reflection/registry.h>
+#include <tvm/s_tir/stmt.h>
#include <tvm/s_tir/transform.h>
#include <tvm/tirx/transform.h>
@@ -326,11 +327,11 @@ tvm::transform::Sequential PassListForPerStoreFeature() {
s_tir::transform::PlanAndUpdateBufferAllocationLocation(),
s_tir::transform::ConvertBlocksToOpaque(),
s_tir::transform::CompactBufferAllocation(),
- tirx::transform::StmtSimplify(),
+ s_tir::transform::StmtSimplify(),
s_tir::transform::LowerAutoCopy(),
s_tir::transform::UnifyThreadBinding(),
s_tir::transform::LowerMatchBuffer(),
- tirx::transform::StmtSimplify(),
+ s_tir::transform::StmtSimplify(),
});
}
diff --git a/src/s_tir/meta_schedule/module_equality.cc b/src/s_tir/meta_schedule/module_equality.cc
index 428dc551eb..7a590353ed 100644
--- a/src/s_tir/meta_schedule/module_equality.cc
+++ b/src/s_tir/meta_schedule/module_equality.cc
@@ -22,6 +22,8 @@
#include <tvm/ffi/extra/structural_equal.h>
#include <tvm/ffi/extra/structural_hash.h>
#include <tvm/ir/module.h>
+#include <tvm/s_tir/analysis.h>
+#include <tvm/s_tir/stmt.h>
#include <tvm/tirx/analysis.h>
#include <memory>
@@ -56,7 +58,7 @@ class ModuleEqualityAnchorBlock : public ModuleEquality {
size_t Hash(IRModule mod) const {
auto anchor_block = tirx::FindAnchorBlock(mod);
if (anchor_block) {
- return ffi::StructuralHash::Hash(ffi::GetRef<tirx::SBlock>(anchor_block),
+ return ffi::StructuralHash::Hash(ffi::GetRef<s_tir::SBlock>(anchor_block),
/*map_free_vars=*/false,
/*skip_tensor_content=*/true);
}
@@ -66,8 +68,8 @@ class ModuleEqualityAnchorBlock : public ModuleEquality {
auto anchor_block_lhs = tirx::FindAnchorBlock(lhs);
auto anchor_block_rhs = tirx::FindAnchorBlock(rhs);
if (anchor_block_lhs && anchor_block_rhs) {
- return tvm::ffi::StructuralEqual::Equal(ffi::GetRef<tirx::SBlock>(anchor_block_lhs),
- ffi::GetRef<tirx::SBlock>(anchor_block_rhs),
+ return tvm::ffi::StructuralEqual::Equal(ffi::GetRef<s_tir::SBlock>(anchor_block_lhs),
+ ffi::GetRef<s_tir::SBlock>(anchor_block_rhs),
/*map_free_vars=*/false,
/*skip_tensor_content=*/true);
}
diff --git a/src/s_tir/meta_schedule/postproc/rewrite_tensorize.cc b/src/s_tir/meta_schedule/postproc/rewrite_tensorize.cc
index 9ac8b2e699..a617ed6901 100644
--- a/src/s_tir/meta_schedule/postproc/rewrite_tensorize.cc
+++ b/src/s_tir/meta_schedule/postproc/rewrite_tensorize.cc
@@ -37,9 +37,9 @@ void CollectTensorizationJobs(
const s_tir::Schedule& sch, const ffi::String& func_name, const tirx::PrimFuncNode* func,
bool vectorize_init_loop,
std::vector<std::tuple<ffi::String, ffi::String, std::function<void(s_tir::SBlockRV)>>>* jobs) {
- auto walk_fn = [=, &jobs](const tirx::SBlock& block) -> ffi::Expected<ffi::WalkResult> {
+ auto walk_fn = [=, &jobs](const s_tir::SBlock& block) -> ffi::Expected<ffi::WalkResult> {
tirx::StmtSRef block_sref = sch->GetSRef(block.get());
- std::string block_name = block_sref->StmtAs<tirx::SBlockNode>()->name_hint;
+ std::string block_name = block_sref->StmtAs<s_tir::SBlockNode>()->name_hint;
if (ffi::Optional<ffi::String> intrin_name =
s_tir::GetAnn<ffi::String>(block_sref, s_tir::attr::meta_schedule_auto_tensorize)) {
if (intrin_name.value() != "") {
diff --git a/src/s_tir/meta_schedule/postproc/rewrite_unbound_block.cc b/src/s_tir/meta_schedule/postproc/rewrite_unbound_block.cc
index e57b5624d6..abbf13cc48 100644
--- a/src/s_tir/meta_schedule/postproc/rewrite_unbound_block.cc
+++ b/src/s_tir/meta_schedule/postproc/rewrite_unbound_block.cc
@@ -18,6 +18,7 @@
*/
#include <tvm/ffi/reflection/registry.h>
#include <tvm/s_tir/meta_schedule/schedule/cuda/thread_bind.h>
+#include <tvm/s_tir/stmt.h>
#include "../utils.h"
diff --git a/src/s_tir/meta_schedule/postproc/verify_gpu_code.cc b/src/s_tir/meta_schedule/postproc/verify_gpu_code.cc
index a40d18bde7..cd707c1d72 100644
--- a/src/s_tir/meta_schedule/postproc/verify_gpu_code.cc
+++ b/src/s_tir/meta_schedule/postproc/verify_gpu_code.cc
@@ -176,7 +176,7 @@ class VerifyGPUCodeNode : public PostprocNode {
pass_list.push_back(s_tir::transform::LiftThreadBinding());
pass_list.push_back(s_tir::transform::ManifestSharedMemoryLocalStage());
pass_list.push_back(s_tir::transform::CompactBufferAllocation());
- pass_list.push_back(tirx::transform::StmtSimplify());
+ pass_list.push_back(s_tir::transform::StmtSimplify());
pass_list.push_back(s_tir::transform::LowerAutoCopy());
pass_list.push_back(s_tir::transform::UnifyThreadBinding());
pass_list.push_back(s_tir::transform::LowerMatchBuffer());
diff --git a/src/s_tir/meta_schedule/schedule_rule/cross_thread_reduction.cc b/src/s_tir/meta_schedule/schedule_rule/cross_thread_reduction.cc
index d1d8e633e0..65451de337 100644
--- a/src/s_tir/meta_schedule/schedule_rule/cross_thread_reduction.cc
+++ b/src/s_tir/meta_schedule/schedule_rule/cross_thread_reduction.cc
@@ -17,6 +17,7 @@
* under the License.
*/
#include <tvm/ffi/reflection/registry.h>
+#include <tvm/s_tir/stmt.h>
#include "../utils.h"
@@ -225,7 +226,7 @@ class CrossThreadReductionNode : public ScheduleRuleNode {
// - If the lowest common ancestor is a loop, the target block is also the first consumer.
const tirx::StmtSRef& lca_sref =
s_tir::GetSRefLowestCommonAncestor(s_tir::SBlockRVs2StmtSRefs(sch, consumers));
- if (consumers.size() > 1 && lca_sref->StmtAs<tirx::SBlockNode>() != nullptr) {
+ if (consumers.size() > 1 && lca_sref->StmtAs<s_tir::SBlockNode>() != nullptr) {
return std::make_tuple(false, s_tir::LoopRV{ffi::UnsafeInit()},
s_tir::SBlockRV{ffi::UnsafeInit()}, s_tir::LoopRV{ffi::UnsafeInit()});
}
diff --git a/src/s_tir/meta_schedule/schedule_rule/multi_level_tiling.cc b/src/s_tir/meta_schedule/schedule_rule/multi_level_tiling.cc
index d0e8db98d9..305e614ce5 100644
--- a/src/s_tir/meta_schedule/schedule_rule/multi_level_tiling.cc
+++ b/src/s_tir/meta_schedule/schedule_rule/multi_level_tiling.cc
@@ -369,7 +369,7 @@ std::vector<State> MultiLevelTilingNode::AddAsyncPipeline(State state) const {
void MultiLevelTilingNode::AnnotateCooperativeFetching(Schedule* sch,
const s_tir::SBlockRV& block) const {
// Filter out invalid vector lanes according to the data type.
- const tirx::SBlockNode* block_node = (*sch)->GetSRef(block)->StmtAs<tirx::SBlockNode>();
+ const s_tir::SBlockNode* block_node = (*sch)->GetSRef(block)->StmtAs<s_tir::SBlockNode>();
TVM_FFI_ICHECK_EQ(block_node->writes.size(), 1);
const DLDataType dtype =
block_node->writes[0]->source.as_or_throw<tvm::tirx::BufferVar>()->dtype->dtype;
diff --git a/src/s_tir/meta_schedule/schedule_rule/multi_level_tiling_tensor_core.cc b/src/s_tir/meta_schedule/schedule_rule/multi_level_tiling_tensor_core.cc
index 7c4e9238a9..bf5e1d4f1a 100644
--- a/src/s_tir/meta_schedule/schedule_rule/multi_level_tiling_tensor_core.cc
+++ b/src/s_tir/meta_schedule/schedule_rule/multi_level_tiling_tensor_core.cc
@@ -20,6 +20,7 @@
#include <tvm/ffi/reflection/registry.h>
#include <tvm/s_tir/meta_schedule/schedule_rule.h>
#include <tvm/s_tir/stmt.h>
+#include <tvm/s_tir/tensor_intrin.h>
#include <tvm/tirx/op.h>
#include <algorithm>
@@ -65,7 +66,7 @@ TensorCoreIntrinGroup TensorCoreIntrinGroup::FromConfig(
TVM_FFI_CHECK(config.count(key_name), ValueError) << key_name << " is not set.";
*intrin_name = config.at(key_name);
// Check the existence of the intrin
- tirx::TensorIntrin::Get(*intrin_name);
+ TensorIntrin::Get(*intrin_name);
};
TensorCoreIntrinGroup intrin_group;
f_initialize_intrin("init", &intrin_group.init_intrin);
@@ -226,7 +227,7 @@ ffi::Array<Schedule> MultiLevelTilingTensorCoreNode::Apply(const Schedule& sch,
ffi::Optional<s_tir::AutoTensorizeMappingInfo> mapping_info =
s_tir::GetAutoTensorizeMappingInfo(
sch->state(), sch->GetSRef(block_rv),
- tirx::TensorIntrin::Get(intrin_groups[i].compute_intrin).value()->desc);
+ TensorIntrin::Get(intrin_groups[i].compute_intrin).value()->desc);
if (mapping_info.has_value()) {
intrin_group_to_mapping_info.emplace(i, mapping_info.value());
}
@@ -447,10 +448,10 @@ std::vector<State> MultiLevelTilingTensorCoreNode::TransformIntermediateOutputLa
// Get the shape of the wmma accumulator
auto [frag_shape_m, frag_shape_n] = [&]() {
- tirx::SBlock intrin_block = tirx::TensorIntrin::Get(state->intrin_group.init_intrin)
- .value()
- ->desc->body.as_or_throw<tirx::SBlockRealize>()
- ->block;
+ s_tir::SBlock intrin_block = TensorIntrin::Get(state->intrin_group.init_intrin)
+ .value()
+ ->desc->body.as_or_throw<s_tir::SBlockRealize>()
+ ->block;
tirx::For loop_m = intrin_block->body.as_or_throw<tirx::For>();
tirx::For loop_n = loop_m->body.as_or_throw<tirx::For>();
return std::make_tuple(loop_m->extent, loop_n->extent);
@@ -627,9 +628,10 @@ std::vector<State> MultiLevelTilingTensorCoreNode::AddReadReuseTensorCore(
const s_tir::SBlockRV cache_read = state->read_reuse.at(i);
// Inline the reindex / padding block
sch->ComputeInline(sch->GetProducers(cache_read)[0]);
- const tirx::SBlockNode* cache_read_block = sch->GetSRef(cache_read)->StmtAs<tirx::SBlockNode>();
+ const s_tir::SBlockNode* cache_read_block =
+ sch->GetSRef(cache_read)->StmtAs<s_tir::SBlockNode>();
tirx::BufferVar cache_read_buffer =
- s_tir::GetNthAccessBuffer(sch->state(), ffi::GetRef<tirx::SBlock>(cache_read_block), 0,
+ s_tir::GetNthAccessBuffer(sch->state(), ffi::GetRef<s_tir::SBlock>(cache_read_block), 0,
s_tir::BufferIndexType::kWrite);
const DLDataType dtype = cache_read_buffer->dtype->dtype;
// Storage alignment is chosen from element storage width; this schedule rule uses scalar
@@ -782,9 +784,9 @@ ffi::Optional<LoopRV> MultiLevelTilingTensorCoreNode::TransformWithTensorIntrin(
tirx::StmtSRef block_sref = state->sch->GetSRef(state->block_rv);
// Add reindex stages
- const tirx::SBlockNode* block = TVM_SREF_TO_SBLOCK(block_sref);
+ const s_tir::SBlockNode* block = TVM_SREF_TO_SBLOCK(block_sref);
// Hold the reference of the block before reindex
- const tirx::SBlock block_before_reindex = ffi::GetRef<tirx::SBlock>(block);
+ const s_tir::SBlock block_before_reindex = ffi::GetRef<s_tir::SBlock>(block);
if (block->reads.size() != 2 || block->writes.size() != 1) {
// only matmul-like computation is allowed
return std::nullopt;
@@ -867,7 +869,7 @@ ffi::Optional<LoopRV> MultiLevelTilingTensorCoreNode::TransformWithTensorIntrin(
// Refresh block pointer (block sref is not invalidated)
block = TVM_SREF_TO_SBLOCK(block_sref);
const tvm::TensorRegion& reindexed_buffer_region = s_tir::GetNthAccessBufferRegion(
- state->sch->state(), ffi::GetRef<tirx::SBlock>(block), buffer_index, index_type);
+ state->sch->state(), ffi::GetRef<s_tir::SBlock>(block), buffer_index, index_type);
auto sub_index_map = f_get_sub_index_map(lhs_buffer, reindexed_buffer_region->region);
buffer_sub_index_map.Set(lhs_buffer, sub_index_map);
state->sch->TransformLayout(state->block_rv, buffer_index, index_type, sub_index_map,
diff --git a/src/s_tir/meta_schedule/schedule_rule/multi_level_tiling_wide_vector.cc b/src/s_tir/meta_schedule/schedule_rule/multi_level_tiling_wide_vector.cc
index 51351fbd89..5b8d16e28e 100644
--- a/src/s_tir/meta_schedule/schedule_rule/multi_level_tiling_wide_vector.cc
+++ b/src/s_tir/meta_schedule/schedule_rule/multi_level_tiling_wide_vector.cc
@@ -18,6 +18,7 @@
*/
#include <tvm/ffi/reflection/registry.h>
+#include <tvm/s_tir/stmt.h>
#include "../../schedule/analysis.h"
#include "../../schedule/transform.h"
@@ -67,8 +68,8 @@ MultiLevelTilingWideVectorNode::SplitLoop(const Schedule& sch, SBlockRV block_rv
int n_tiles) const {
const tirx::ForNode* loop = TVM_SREF_TO_FOR(sch->GetSRef(loop_rv));
const tirx::StmtSRef block_sref = sch->GetSRef(block_rv);
- const tirx::SBlockNode* block_node = block_sref->StmtAs<tirx::SBlockNode>();
- const tirx::SBlockRealize block_realize = s_tir::GetSBlockRealize(sch->state(), block_sref);
+ const s_tir::SBlockNode* block_node = block_sref->StmtAs<s_tir::SBlockNode>();
+ const s_tir::SBlockRealize block_realize = s_tir::GetSBlockRealize(sch->state(), block_sref);
TVM_FFI_ICHECK(block_node && block_node->writes.size() == 1);
const auto out_dtype = block_node->writes[0]->source.as_or_throw<tvm::tirx::BufferVar>()->dtype;
diff --git a/src/s_tir/meta_schedule/schedule_rule/multi_level_tiling_with_intrin.cc b/src/s_tir/meta_schedule/schedule_rule/multi_level_tiling_with_intrin.cc
index fab8937ae5..c499820395 100644
--- a/src/s_tir/meta_schedule/schedule_rule/multi_level_tiling_with_intrin.cc
+++ b/src/s_tir/meta_schedule/schedule_rule/multi_level_tiling_with_intrin.cc
@@ -19,6 +19,7 @@
#include <tvm/ffi/reflection/registry.h>
#include <tvm/s_tir/stmt.h>
+#include <tvm/s_tir/tensor_intrin.h>
#include "../../schedule/analysis.h"
#include "../../schedule/transform.h"
@@ -52,7 +53,7 @@ class MultiLevelTilingWithIntrinNode : public MultiLevelTilingNode {
protected:
ffi::Array<s_tir::Schedule> Apply(const s_tir::Schedule& sch,
const s_tir::SBlockRV& block_rv) final {
- auto desc_func = tirx::TensorIntrin::Get(intrin_name).value()->desc;
+ auto desc_func = TensorIntrin::Get(intrin_name).value()->desc;
if (!CheckAutoTensorizeApplicable(sch, block_rv, desc_func)) {
TVM_PY_LOG(INFO, logger) << "The workload cannot be tensorized.";
return {sch};
@@ -106,7 +107,7 @@ ScheduleRule ScheduleRule::MultiLevelTilingWithIntrin(
ffi::Optional<ffi::Array<int64_t>> vector_load_lens,
ffi::Optional<ffi::Map<ffi::String, ffi::Any>> reuse_read,
ffi::Optional<ffi::Map<ffi::String, ffi::Any>> reuse_write) {
- TVM_FFI_ICHECK(tirx::TensorIntrin::Get(intrin_name).has_value())
+ TVM_FFI_ICHECK(TensorIntrin::Get(intrin_name).has_value())
<< "Provided tensor intrinsic " << intrin_name << " is not registered.";
auto node = MultiLevelTilingInitCommon<MultiLevelTilingWithIntrinNode>(
structure, tile_binds, max_innermost_factor, vector_load_lens, reuse_read, reuse_write);
diff --git a/src/s_tir/meta_schedule/schedule_rule/schedule_rule.cc b/src/s_tir/meta_schedule/schedule_rule/schedule_rule.cc
index eee9ef2685..4a5467247a 100644
--- a/src/s_tir/meta_schedule/schedule_rule/schedule_rule.cc
+++ b/src/s_tir/meta_schedule/schedule_rule/schedule_rule.cc
@@ -18,6 +18,7 @@
*/
#include <tvm/ffi/dtype.h>
#include <tvm/ffi/reflection/registry.h>
+#include <tvm/s_tir/tensor_intrin.h>
#include "../utils.h"
@@ -325,7 +326,7 @@ ffi::Array<ScheduleRule> ScheduleRule::DefaultRISCV(const int vlen) {
const auto rvv_kernels_inventory = reg_rvv_intrinsics(current_target, /* inventory_only */ true)
.cast<ffi::Map<ffi::String, int>>();
for (const auto& intrin : rvv_kernels_inventory) {
- if (!tirx::TensorIntrin::Get(intrin.first, /*allow_missing*/ true)) {
+ if (!TensorIntrin::Get(intrin.first, /*allow_missing*/ true)) {
// on demand intrinsic register
reg_rvv_intrinsics(current_target, /* inventory_only */ false);
}
diff --git a/src/s_tir/meta_schedule/trace_apply.cc b/src/s_tir/meta_schedule/trace_apply.cc
index 34678aac84..2bfc430562 100644
--- a/src/s_tir/meta_schedule/trace_apply.cc
+++ b/src/s_tir/meta_schedule/trace_apply.cc
@@ -19,8 +19,9 @@
#include "trace_apply.h"
#include <tvm/ffi/reflection/registry.h>
+#include <tvm/s_tir/analysis.h>
+#include <tvm/s_tir/stmt_functor.h>
#include <tvm/tirx/analysis.h>
-#include <tvm/tirx/stmt_functor.h>
#include <optional>
#include <string>
diff --git a/src/s_tir/meta_schedule/utils.h b/src/s_tir/meta_schedule/utils.h
index 02ac858b1b..d8d9e39a79 100644
--- a/src/s_tir/meta_schedule/utils.h
+++ b/src/s_tir/meta_schedule/utils.h
@@ -41,6 +41,7 @@
#include <tvm/s_tir/meta_schedule/task_scheduler.h>
#include <tvm/s_tir/meta_schedule/tune_context.h>
#include <tvm/s_tir/schedule/schedule.h>
+#include <tvm/s_tir/stmt.h>
#include <tvm/support/io.h>
#include <tvm/support/serializer.h>
#include <tvm/tirx/transform.h>
@@ -301,7 +302,7 @@ inline std::string Concat(const ffi::Array<ffi::String>& strs, const std::string
*/
inline s_tir::SBlockRV GetRVFromSRef(const s_tir::Schedule& sch, const tirx::StmtSRef& block_sref,
const ffi::String& global_var_name) {
- const tirx::SBlockNode* block = TVM_SREF_TO_SBLOCK(block_sref);
+ const s_tir::SBlockNode* block = TVM_SREF_TO_SBLOCK(block_sref);
return sch->GetSBlock(block->name_hint, global_var_name);
}
@@ -612,13 +613,13 @@ inline double Sum(const ffi::Array<FloatImm>& arr) {
}
/*! \brief Collecting all the blocks */
-class SBlockCollector : public tirx::StmtExprVisitor {
+class SBlockCollector : public s_tir::StmtExprVisitor {
public:
- using tirx::StmtExprVisitor::Visit_;
+ using s_tir::StmtExprVisitor::Visit_;
ffi::Optional<VisitInterrupt> Visit(ffi::AnyView value) override {
if (value.as<ExprNode>()) return std::nullopt;
- return tirx::StmtExprVisitor::Visit(value);
+ return s_tir::StmtExprVisitor::Visit(value);
}
static ffi::Array<s_tir::SBlockRV> Collect(const s_tir::Schedule& sch,
@@ -662,8 +663,8 @@ class SBlockCollector : public tirx::StmtExprVisitor {
private:
/*! \brief Override the Stmt visiting behaviour */
- ffi::Optional<VisitInterrupt> Visit_(const tirx::SBlockNode* block) override {
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(tirx::StmtExprVisitor::Visit_(block));
+ ffi::Optional<VisitInterrupt> Visit_(const s_tir::SBlockNode* block) override {
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(s_tir::StmtExprVisitor::Visit_(block));
TVM_FFI_ICHECK(block_names_.count(block->name_hint) == 0)
<< "Duplicated block name " << block->name_hint << " in function " << func_name_
<< " not supported!";
@@ -673,7 +674,7 @@ class SBlockCollector : public tirx::StmtExprVisitor {
// Otherwise collect all blocks.
bool collect_block = true;
if (f_block_filter_ != nullptr) {
- collect_block = f_block_filter_(ffi::GetRef<tirx::SBlock>(block)).cast<IntImm>()->value != 0;
+ collect_block = f_block_filter_(ffi::GetRef<s_tir::SBlock>(block)).cast<IntImm>()->value != 0;
}
if (collect_block) {
blocks_to_collect_.push_back(block->name_hint);
diff --git a/src/s_tir/sblock_dependence_info.cc b/src/s_tir/sblock_dependence_info.cc
index 9dbba04aaa..55ba3ef8f4 100644
--- a/src/s_tir/sblock_dependence_info.cc
+++ b/src/s_tir/sblock_dependence_info.cc
@@ -19,6 +19,7 @@
#include <tvm/ffi/reflection/registry.h>
#include <tvm/s_tir/sblock_dependence_info.h>
+#include <tvm/s_tir/stmt.h>
#include <tvm/s_tir/utils.h>
namespace tvm {
@@ -27,15 +28,15 @@ namespace tirx {
TVM_FFI_STATIC_INIT_BLOCK() { SBlockDependenceInfoNode::RegisterReflection(); }
/**
- * @brief A helper class to collect and build SBlock Dependences using SBlockScope class
+ * @brief A helper class to collect and build s_tir::SBlock Dependences using SBlockScope class
*/
-class SBlockDependenceInfoCollector : public StmtExprVisitor {
+class SBlockDependenceInfoCollector : public s_tir::StmtExprVisitor {
public:
- using StmtExprVisitor::Visit_;
+ using s_tir::StmtExprVisitor::Visit_;
ffi::Optional<VisitInterrupt> Visit(ffi::AnyView value) override {
if (value.as<ExprNode>()) return std::nullopt;
- return StmtExprVisitor::Visit(value);
+ return s_tir::StmtExprVisitor::Visit(value);
}
static void Collect(SBlockDependenceInfoNode* self, const Stmt& stmt) {
@@ -53,9 +54,9 @@ class SBlockDependenceInfoCollector : public StmtExprVisitor {
self_->sref2scope[scope] = SBlockScope(child_block_srefs);
}
- ffi::Optional<VisitInterrupt> Visit_(const SBlockRealizeNode* realize) final {
+ ffi::Optional<VisitInterrupt> Visit_(const s_tir::SBlockRealizeNode* realize) final {
block_frames_.emplace_back();
- const SBlockNode* block = realize->block.get();
+ const s_tir::SBlockNode* block = realize->block.get();
// Recursive visit
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Visit(block->body)); // `block->init` is not visited
// Create SBlockInfo for the block
@@ -69,7 +70,7 @@ class SBlockDependenceInfoCollector : public StmtExprVisitor {
ffi::Optional<VisitInterrupt> Visit_(const SeqStmtNode* seq_stmt) final {
// Set `seq_index` information for SeqStmtNode
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(seq_stmt));
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(s_tir::StmtExprVisitor::Visit_(seq_stmt));
SetSeqIndexInChildren(self_->stmt2ref, seq_stmt, false);
return std::nullopt;
}
diff --git a/src/s_tir/sblock_scope.cc b/src/s_tir/sblock_scope.cc
index f4f2f0c815..12ef4368a8 100644
--- a/src/s_tir/sblock_scope.cc
+++ b/src/s_tir/sblock_scope.cc
@@ -19,6 +19,7 @@
#include <tvm/ffi/cast.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/s_tir/sblock_scope.h>
+#include <tvm/s_tir/stmt.h>
#include <tvm/s_tir/utils.h>
namespace tvm {
@@ -85,7 +86,7 @@ SBlockScope::SBlockScope(const ffi::Array<StmtSRef>& child_block_srefs) {
SMap<BufferVar, ffi::Array<StmtSRef>> buffer_readers;
SMap<BufferVar, ffi::Array<StmtSRef>>& buffer_writers = n->buffer_writers;
for (const StmtSRef& child_block_sref : child_block_srefs) {
- const SBlockNode* child_block = TVM_SREF_TO_SBLOCK(child_block_sref);
+ const s_tir::SBlockNode* child_block = TVM_SREF_TO_SBLOCK(child_block_sref);
// Step 1. Update `buffer_readers` and `buffer_writers` for each buffer
for (const TensorRegion& region : child_block->reads) {
buffer_readers[region->source.as_or_throw<tvm::tirx::BufferVar>()].push_back(
@@ -182,8 +183,8 @@ ffi::Optional<VisitInterrupt> SRefTreeCreator::Visit_(const ForNode* loop) {
return std::nullopt;
}
-ffi::Optional<VisitInterrupt> SRefTreeCreator::Visit_(const SBlockRealizeNode* realize) {
- const SBlockNode* block = realize->block.get();
+ffi::Optional<VisitInterrupt> SRefTreeCreator::Visit_(const s_tir::SBlockRealizeNode* realize) {
+ const s_tir::SBlockNode* block = realize->block.get();
PushSRef(block);
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Visit(block->body)); // `block->init` is not visited
PopAndRecordSRef();
@@ -192,7 +193,7 @@ ffi::Optional<VisitInterrupt> SRefTreeCreator::Visit_(const SBlockRealizeNode* r
ffi::Optional<VisitInterrupt> SRefTreeCreator::Visit_(const SeqStmtNode* seq_stmt) {
// Set `seq_index` information for SeqStmtNode
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(seq_stmt));
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(s_tir::StmtExprVisitor::Visit_(seq_stmt));
SetSeqIndexInChildren(stmt2ref_, seq_stmt, include_loops_);
return std::nullopt;
}
diff --git a/src/s_tir/schedule/analysis.h b/src/s_tir/schedule/analysis.h
index 54a77e93d1..2ee6240b72 100644
--- a/src/s_tir/schedule/analysis.h
+++ b/src/s_tir/schedule/analysis.h
@@ -24,6 +24,7 @@
#include <tvm/ir/prim/expr.h>
#include <tvm/s_tir/schedule/schedule.h>
#include <tvm/s_tir/schedule/state.h>
+#include <tvm/s_tir/stmt.h>
#include <tvm/te/operation.h>
#include <tvm/tirx/index_map.h>
diff --git a/src/s_tir/schedule/analysis/analysis.cc b/src/s_tir/schedule/analysis/analysis.cc
index 57a098a31c..789dcbff73 100644
--- a/src/s_tir/schedule/analysis/analysis.cc
+++ b/src/s_tir/schedule/analysis/analysis.cc
@@ -1707,7 +1707,7 @@ bool NeedsRFactorOrCrossThreadReduction(const s_tir::ScheduleState& self, //
return false;
}
} else {
- const auto* block_realize = loop_i->body.as<tirx::SBlockRealizeNode>();
+ const auto* block_realize = loop_i->body.as<s_tir::SBlockRealizeNode>();
if (!block_realize || block_realize->block.get() != block) {
return false;
}
@@ -1791,7 +1791,7 @@ ffi::Optional<TensorizeInfo> GetTensorizeLoopMapping(const s_tir::ScheduleState&
const tirx::PrimFunc& desc_func,
bool allow_padding) {
arith::Analyzer analyzer;
- const tirx::SBlockRealize& block = GetSBlockRealize(self, block_sref);
+ const s_tir::SBlockRealize& block = GetSBlockRealize(self, block_sref);
// Step 1. Analyze desc_func, extract its block, loops and loop vars
TensorIntrinDescInfo desc_info = ExtractTensorIntrinDescInfo(analyzer.get(), desc_func);
// Step 2. Collect loops from block_sref
diff --git a/src/s_tir/schedule/analysis/reducer.cc b/src/s_tir/schedule/analysis/reducer.cc
index 9de832d805..680d556df6 100644
--- a/src/s_tir/schedule/analysis/reducer.cc
+++ b/src/s_tir/schedule/analysis/reducer.cc
@@ -19,6 +19,7 @@
#include <tvm/ffi/cast.h>
#include <tvm/ffi/extra/structural_visit.h>
#include <tvm/ir/prim/expr.h>
+#include <tvm/s_tir/stmt.h>
#include <tvm/te/operation.h>
#include "../utils.h"
diff --git a/src/s_tir/schedule/analysis/verify.cc b/src/s_tir/schedule/analysis/verify.cc
index d83d5df4f1..11d32ae235 100644
--- a/src/s_tir/schedule/analysis/verify.cc
+++ b/src/s_tir/schedule/analysis/verify.cc
@@ -17,6 +17,7 @@
* under the License.
*/
#include <tvm/ffi/cast.h>
+#include <tvm/s_tir/stmt.h>
#include "../utils.h"
diff --git a/src/s_tir/schedule/concrete_schedule.cc b/src/s_tir/schedule/concrete_schedule.cc
index bb3e1e7701..903203e5a8 100644
--- a/src/s_tir/schedule/concrete_schedule.cc
+++ b/src/s_tir/schedule/concrete_schedule.cc
@@ -20,6 +20,8 @@
#include <tvm/ffi/cast.h>
#include <tvm/runtime/logging.h>
+#include <tvm/s_tir/stmt.h>
+#include <tvm/s_tir/tensor_intrin.h>
#include <random>
@@ -927,7 +929,7 @@ SBlockRV ConcreteScheduleNode::Blockize(const ffi::Array<SBlockRV>& blocks,
void ConcreteScheduleNode::Tensorize(const LoopRV& loop_rv, const ffi::String& intrin,
bool preserve_unit_iters) {
TVM_TIR_SCHEDULE_BEGIN();
- s_tir::Tensorize(state_, this->GetSRef(loop_rv), tirx::TensorIntrin::Get(intrin).value(),
+ s_tir::Tensorize(state_, this->GetSRef(loop_rv), TensorIntrin::Get(intrin).value(),
preserve_unit_iters);
this->state_->DebugVerify();
TVM_TIR_SCHEDULE_END("tensorize", this->error_render_level_);
@@ -936,7 +938,7 @@ void ConcreteScheduleNode::Tensorize(const LoopRV& loop_rv, const ffi::String& i
void ConcreteScheduleNode::Tensorize(const SBlockRV& block_rv, const ffi::String& intrin,
bool preserve_unit_iters) {
TVM_TIR_SCHEDULE_BEGIN();
- s_tir::Tensorize(state_, this->GetSRef(block_rv), tirx::TensorIntrin::Get(intrin).value(),
+ s_tir::Tensorize(state_, this->GetSRef(block_rv), TensorIntrin::Get(intrin).value(),
preserve_unit_iters);
this->state_->DebugVerify();
TVM_TIR_SCHEDULE_END("tensorize", this->error_render_level_);
diff --git a/src/s_tir/schedule/concrete_schedule.h b/src/s_tir/schedule/concrete_schedule.h
index 3a1501ed7c..205eda800a 100644
--- a/src/s_tir/schedule/concrete_schedule.h
+++ b/src/s_tir/schedule/concrete_schedule.h
@@ -22,6 +22,7 @@
#include <tvm/ffi/cast.h>
#include <tvm/ffi/extra/structural_mutate.h>
#include <tvm/ir/prim/expr.h>
+#include <tvm/s_tir/stmt.h>
#include <memory>
#include <utility>
diff --git a/src/s_tir/schedule/error.h b/src/s_tir/schedule/error.h
index 006304ea5a..0bdc14837f 100644
--- a/src/s_tir/schedule/error.h
+++ b/src/s_tir/schedule/error.h
@@ -23,6 +23,7 @@
#include <tvm/ffi/memory.h>
#include <tvm/ir/prim/expr.h>
#include <tvm/s_tir/schedule/state.h>
+#include <tvm/s_tir/stmt.h>
#include <string>
#include <type_traits>
diff --git a/src/s_tir/schedule/ir_comparator.cc b/src/s_tir/schedule/ir_comparator.cc
index 1b8b1851cf..015ba132f6 100644
--- a/src/s_tir/schedule/ir_comparator.cc
+++ b/src/s_tir/schedule/ir_comparator.cc
@@ -20,6 +20,7 @@
#include <tvm/ffi/cast.h>
#include <tvm/ir/prim/builtin.h>
+#include <tvm/s_tir/stmt.h>
#include <tvm/tirx/builtin.h>
#include "../../tirx/analysis/check_contains.h"
diff --git a/src/s_tir/schedule/ir_comparator.h b/src/s_tir/schedule/ir_comparator.h
index 75303df5b8..58db4f8ebd 100644
--- a/src/s_tir/schedule/ir_comparator.h
+++ b/src/s_tir/schedule/ir_comparator.h
@@ -20,6 +20,7 @@
#define TVM_S_TIR_SCHEDULE_IR_COMPARATOR_H_
#include <tvm/ir/prim/expr.h>
+#include <tvm/s_tir/stmt.h>
#include <string>
#include <unordered_map>
diff --git a/src/s_tir/schedule/primitive.h b/src/s_tir/schedule/primitive.h
index 83ab03ec57..89c058f9ba 100644
--- a/src/s_tir/schedule/primitive.h
+++ b/src/s_tir/schedule/primitive.h
@@ -22,6 +22,8 @@
#include <tvm/ir/prim/expr.h>
#include <tvm/s_tir/random_engine.h>
#include <tvm/s_tir/schedule/state.h>
+#include <tvm/s_tir/stmt.h>
+#include <tvm/s_tir/tensor_intrin.h>
#include <vector>
diff --git a/src/s_tir/schedule/primitive/annotate.cc b/src/s_tir/schedule/primitive/annotate.cc
index 8b22e99e9f..9fd1bc2752 100644
--- a/src/s_tir/schedule/primitive/annotate.cc
+++ b/src/s_tir/schedule/primitive/annotate.cc
@@ -17,6 +17,7 @@
* under the License.
*/
#include <tvm/ffi/cast.h>
+#include <tvm/s_tir/stmt.h>
#include "../utils.h"
diff --git a/src/s_tir/schedule/primitive/block_annotate.cc b/src/s_tir/schedule/primitive/block_annotate.cc
index 97288e803e..6e6c0f4804 100644
--- a/src/s_tir/schedule/primitive/block_annotate.cc
+++ b/src/s_tir/schedule/primitive/block_annotate.cc
@@ -21,7 +21,7 @@
#include <tvm/ir/prim/expr.h>
#include <tvm/s_tir/stmt.h>
-#include "../../../tirx/transform/ir_utils.h"
+#include "../../transform/ir_utils.h"
#include "../utils.h"
namespace tvm {
diff --git a/src/s_tir/schedule/primitive/blockize_tensorize.cc b/src/s_tir/schedule/primitive/blockize_tensorize.cc
index 572f8c4804..fc88118224 100644
--- a/src/s_tir/schedule/primitive/blockize_tensorize.cc
+++ b/src/s_tir/schedule/primitive/blockize_tensorize.cc
@@ -21,11 +21,13 @@
#include <tvm/ffi/extra/structural_mutate.h>
#include <tvm/ffi/extra/structural_visit.h>
#include <tvm/runtime/logging.h>
+#include <tvm/s_tir/stmt.h>
+#include <tvm/s_tir/tensor_intrin.h>
#include <functional>
-#include "../../../tirx/ir/data_type_rewriter.h"
-#include "../../../tirx/transform/stmt_simplify.h"
+#include "../../ir/data_type_rewriter.h"
+#include "../../transform/stmt_simplify.h"
#include "../ir_comparator.h"
#include "../utils.h"
@@ -817,7 +819,7 @@ void Tensorize(ScheduleState self, const StmtSRef& sref, const TensorIntrin& int
}
arith::Analyzer analyzer;
- PrimFunc intrin_desc = StmtSimplify(intrin->desc, analyzer);
+ PrimFunc intrin_desc = s_tir::StmtSimplify(intrin->desc, analyzer);
PrimFunc intrin_impl = DeepCopy(intrin->impl);
int index_dtype_bits = -1;
diff --git a/src/s_tir/schedule/primitive/cache_index.cc b/src/s_tir/schedule/primitive/cache_index.cc
index c656d4ed74..630c4a6814 100644
--- a/src/s_tir/schedule/primitive/cache_index.cc
+++ b/src/s_tir/schedule/primitive/cache_index.cc
@@ -20,6 +20,7 @@
#include <tvm/ffi/cast.h>
#include <tvm/ffi/extra/structural_mutate.h>
#include <tvm/ffi/extra/structural_visit.h>
+#include <tvm/s_tir/stmt.h>
#include "../../../tirx/transform/replace_selected_expr.h"
#include "../utils.h"
diff --git a/src/s_tir/schedule/primitive/cache_index_helpers.cc b/src/s_tir/schedule/primitive/cache_index_helpers.cc
index d34b103be6..c88a2b9b32 100644
--- a/src/s_tir/schedule/primitive/cache_index_helpers.cc
+++ b/src/s_tir/schedule/primitive/cache_index_helpers.cc
@@ -28,10 +28,11 @@
#include <tvm/arith/analyzer.h> // For the arith::Analyzer::Simplify() method simplifying terms
#include <tvm/ffi/cast.h>
#include <tvm/ir/prim/expr.h>
+#include <tvm/s_tir/analysis.h>
+#include <tvm/s_tir/stmt_functor.h>
#include <tvm/tirx/analysis.h>
#include <tvm/tirx/expr_functor.h>
#include <tvm/tirx/stmt.h>
-#include <tvm/tirx/stmt_functor.h>
#include <algorithm> // For std::find_if
#include <unordered_map> // For the hashtable datatype
@@ -214,7 +215,7 @@ ffi::Optional<VisitInterrupt> ComputationsDoneBy::Visit(ffi::AnyView expr_value)
auto opt_expr = expr_value.as<PrimExpr>();
if (!opt_expr) {
- return StmtExprVisitor::Visit(expr_value);
+ return s_tir::StmtExprVisitor::Visit(expr_value);
}
PrimExpr expr = opt_expr.value();
if (expr.as<IntImmNode>() != nullptr || expr.as<FloatImmNode>() != nullptr ||
@@ -322,7 +323,7 @@ ComputationTable ComputationsDoneBy::ComputationsDoneByChildrenOf(
std::function<bool(const PrimExpr&)> can_contain_computations) {
auto computations_done_by =
ffi::make_object<ComputationsDoneBy>(is_eligible_computation, can_contain_computations);
- computations_done_by->StmtExprVisitor::Visit(expr);
+ computations_done_by->s_tir::StmtExprVisitor::Visit(expr);
cache_.cache_expr_table_computations_[expr] = computations_done_by->table_of_computations_;
return computations_done_by->table_of_computations_;
@@ -336,7 +337,7 @@ ComputationTable ComputationsDoneBy::ComputationsDoneByChildrenOf(
std::function<bool(const PrimExpr&)> can_contain_computations) {
auto computations_done_by =
ffi::make_object<ComputationsDoneBy>(is_eligible_computation, can_contain_computations);
- computations_done_by->StmtExprVisitor::Visit(stmt);
+ computations_done_by->s_tir::StmtExprVisitor::Visit(stmt);
cache_.cache_stmt_table_computations_[stmt] = computations_done_by->table_of_computations_;
return computations_done_by->table_of_computations_;
@@ -371,7 +372,7 @@ DirectSubexpr::DirectSubexpr(std::function<bool(const PrimExpr&)> is_eligible_co
*/
ffi::Optional<VisitInterrupt> DirectSubexpr::Visit(ffi::AnyView expr_value) {
auto opt_expr = expr_value.as<PrimExpr>();
- if (!opt_expr) return StmtExprVisitor::Visit(expr_value);
+ if (!opt_expr) return s_tir::StmtExprVisitor::Visit(expr_value);
PrimExpr expr = opt_expr.value();
if (entered_) {
if (is_eligible_computation_(expr)) {
@@ -379,7 +380,7 @@ ffi::Optional<VisitInterrupt> DirectSubexpr::Visit(ffi::AnyView expr_value) {
return std::nullopt;
} else {
if (can_contain_computations_(expr)) {
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit(expr));
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(s_tir::StmtExprVisitor::Visit(expr));
}
return std::nullopt;
}
@@ -387,7 +388,7 @@ ffi::Optional<VisitInterrupt> DirectSubexpr::Visit(ffi::AnyView expr_value) {
if (can_contain_computations_(expr)) {
entered_ = true;
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit(expr));
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(s_tir::StmtExprVisitor::Visit(expr));
}
return std::nullopt;
}
diff --git a/src/s_tir/schedule/primitive/cache_index_helpers.h b/src/s_tir/schedule/primitive/cache_index_helpers.h
index 2159381816..af54b6415e 100644
--- a/src/s_tir/schedule/primitive/cache_index_helpers.h
+++ b/src/s_tir/schedule/primitive/cache_index_helpers.h
@@ -29,9 +29,9 @@
#include <tvm/ffi/extra/structural_hash.h>
#include <tvm/ffi/string.h>
#include <tvm/ir/prim/expr.h>
+#include <tvm/s_tir/stmt_functor.h> // For the class s_tir::StmtExprVisitor
#include <tvm/tirx/expr_functor.h>
#include <tvm/tirx/stmt.h>
-#include <tvm/tirx/stmt_functor.h> // For the class StmtExprVisitor
#include <unordered_map>
#include <utility> // For pairs datatype
@@ -77,9 +77,9 @@ struct ComputationCache {
* \note Computations here are considered syntactically, meaning that semantically equivalent
computations that are not syntactically the same are not merged together.
*/
-class ComputationsDoneBy : public StmtExprVisitor {
+class ComputationsDoneBy : public s_tir::StmtExprVisitor {
public:
- using StmtExprVisitor::Visit_;
+ using s_tir::StmtExprVisitor::Visit_;
// Toplevel (static) methods
static ComputationTable GetComputationsDoneBy(
@@ -124,9 +124,9 @@ class ComputationsDoneBy : public StmtExprVisitor {
So for instance, for (A+(B+C)) it will return A and (B+C) if they are eligible,
but not B and C.
*/
-class DirectSubexpr : public StmtExprVisitor {
+class DirectSubexpr : public s_tir::StmtExprVisitor {
public:
- using StmtExprVisitor::Visit_;
+ using s_tir::StmtExprVisitor::Visit_;
// Toplevel (static) function
static std::vector<PrimExpr> GetDirectSubexpressions(
diff --git a/src/s_tir/schedule/primitive/cache_read_write.cc b/src/s_tir/schedule/primitive/cache_read_write.cc
index a113f09427..d0d9ec36ef 100644
--- a/src/s_tir/schedule/primitive/cache_read_write.cc
+++ b/src/s_tir/schedule/primitive/cache_read_write.cc
@@ -22,6 +22,7 @@
#include <tvm/ffi/extra/structural_mutate.h>
#include <tvm/ffi/extra/structural_visit.h>
#include <tvm/ir/prim/expr.h>
+#include <tvm/s_tir/stmt.h>
#include <tvm/tirx/builtin.h>
#include <unordered_set>
diff --git a/src/s_tir/schedule/primitive/compute_at.cc b/src/s_tir/schedule/primitive/compute_at.cc
index efccaa7b3e..9c73b9a41e 100644
--- a/src/s_tir/schedule/primitive/compute_at.cc
+++ b/src/s_tir/schedule/primitive/compute_at.cc
@@ -18,6 +18,7 @@
*/
#include <tvm/ffi/cast.h>
#include <tvm/ffi/extra/structural_mutate.h>
+#include <tvm/s_tir/stmt.h>
#include "../utils.h"
diff --git a/src/s_tir/schedule/primitive/decompose_padding.cc b/src/s_tir/schedule/primitive/decompose_padding.cc
index 6a4d03ab24..283a5609a3 100644
--- a/src/s_tir/schedule/primitive/decompose_padding.cc
+++ b/src/s_tir/schedule/primitive/decompose_padding.cc
@@ -19,6 +19,7 @@
#include <tvm/ffi/cast.h>
#include <tvm/ffi/extra/structural_mutate.h>
#include <tvm/ffi/reflection/registry.h>
+#include <tvm/s_tir/stmt.h>
#include "../../../tirx/transform/ir_utils.h"
#include "../utils.h"
diff --git a/src/s_tir/schedule/primitive/for_kind.cc b/src/s_tir/schedule/primitive/for_kind.cc
index b9d06ed277..558aff410d 100644
--- a/src/s_tir/schedule/primitive/for_kind.cc
+++ b/src/s_tir/schedule/primitive/for_kind.cc
@@ -18,6 +18,7 @@
*/
#include <tvm/ffi/cast.h>
#include <tvm/ffi/extra/structural_visit.h>
+#include <tvm/s_tir/stmt.h>
#include "../utils.h"
diff --git a/src/s_tir/schedule/primitive/get_block_loop.cc b/src/s_tir/schedule/primitive/get_block_loop.cc
index ee3a489141..b2e0823679 100644
--- a/src/s_tir/schedule/primitive/get_block_loop.cc
+++ b/src/s_tir/schedule/primitive/get_block_loop.cc
@@ -17,6 +17,7 @@
* under the License.
*/
#include <tvm/ffi/cast.h>
+#include <tvm/s_tir/stmt.h>
#include "../analysis.h"
#include "../utils.h"
diff --git a/src/s_tir/schedule/primitive/hide_buffer_access.cc b/src/s_tir/schedule/primitive/hide_buffer_access.cc
index 2f9f9950f5..43649044ea 100644
--- a/src/s_tir/schedule/primitive/hide_buffer_access.cc
+++ b/src/s_tir/schedule/primitive/hide_buffer_access.cc
@@ -17,6 +17,7 @@
* under the License.
*/
#include <tvm/ffi/cast.h>
+#include <tvm/s_tir/stmt.h>
#include <set>
diff --git a/src/s_tir/schedule/primitive/layout_transformation.cc b/src/s_tir/schedule/primitive/layout_transformation.cc
index dff34341da..c6122710b2 100644
--- a/src/s_tir/schedule/primitive/layout_transformation.cc
+++ b/src/s_tir/schedule/primitive/layout_transformation.cc
@@ -23,11 +23,12 @@
#include <tvm/ffi/extra/structural_visit.h>
#include <tvm/ir/prim/expr.h>
#include <tvm/runtime/logging.h>
+#include <tvm/s_tir/stmt.h>
#include <optional>
#include <variant>
-#include "../../../tirx/ir/ir_mutator_with_analyzer.h"
+#include "../../../s_tir/ir/ir_mutator_with_analyzer.h"
#include "../utils.h"
namespace tvm {
@@ -761,13 +762,13 @@ class TransformLayoutPlanner : public StmtExprVisitor {
* \brief Collect blocks that are part of root block to be passed to ScheduleState::Replace for SRef
* reuse
*/
-class ReuseBlocksCollector : public tirx::StmtExprVisitor {
+class ReuseBlocksCollector : public s_tir::StmtExprVisitor {
public:
- using tirx::StmtExprVisitor::Visit_;
+ using s_tir::StmtExprVisitor::Visit_;
ffi::Optional<VisitInterrupt> Visit(ffi::AnyView value) override {
if (value.as<ExprNode>()) return std::nullopt;
- return tirx::StmtExprVisitor::Visit(value);
+ return s_tir::StmtExprVisitor::Visit(value);
}
static ffi::Map<SBlock, SBlock> Collect(SBlock result,
@@ -788,7 +789,7 @@ class ReuseBlocksCollector : public tirx::StmtExprVisitor {
private:
/*! \brief Override the Stmt visiting behaviour */
- ffi::Optional<VisitInterrupt> Visit_(const tirx::SBlockNode* block) override {
+ ffi::Optional<VisitInterrupt> Visit_(const s_tir::SBlockNode* block) override {
SBlock block_ref = ffi::GetRef<SBlock>(block);
auto it = new_block_to_old_.find(block_ref);
if (it != new_block_to_old_.end()) {
@@ -803,10 +804,10 @@ class ReuseBlocksCollector : public tirx::StmtExprVisitor {
ffi::Map<SBlock, SBlock> new_block_to_old_;
};
-class TransformLayoutRewriter : public tirx::IRMutatorWithAnalyzer {
+class TransformLayoutRewriter : public s_tir::IRMutatorWithAnalyzer {
public:
- using tirx::IRMutatorWithAnalyzer::Mutate;
- using tirx::IRMutatorWithAnalyzer::Mutate_;
+ using s_tir::IRMutatorWithAnalyzer::Mutate;
+ using s_tir::IRMutatorWithAnalyzer::Mutate_;
/*!
* \brief Rewrite the access to the buffer after the transformation
@@ -864,7 +865,7 @@ class TransformLayoutRewriter : public tirx::IRMutatorWithAnalyzer {
*indices = this->IterMapSimplifyWithContext(*indices, true);
}
- using Parent = tirx::IRMutatorWithAnalyzer;
+ using Parent = s_tir::IRMutatorWithAnalyzer;
UnchangedOr<ffi::Any> Mutate(ffi::AnyView value, InplaceMode inplace_mode) final {
const auto* stmt = value.as<StmtNode>();
diff --git a/src/s_tir/schedule/primitive/loop_transformation.cc b/src/s_tir/schedule/primitive/loop_transformation.cc
index a7f9cbbec4..f2987eae47 100644
--- a/src/s_tir/schedule/primitive/loop_transformation.cc
+++ b/src/s_tir/schedule/primitive/loop_transformation.cc
@@ -19,6 +19,7 @@
#include <tvm/ffi/cast.h>
#include <tvm/ffi/extra/structural_mutate.h>
#include <tvm/ffi/extra/structural_visit.h>
+#include <tvm/s_tir/stmt.h>
#include "../utils.h"
diff --git a/src/s_tir/schedule/primitive/pad_einsum.cc b/src/s_tir/schedule/primitive/pad_einsum.cc
index 05de3ac334..7bfd671779 100644
--- a/src/s_tir/schedule/primitive/pad_einsum.cc
+++ b/src/s_tir/schedule/primitive/pad_einsum.cc
@@ -19,6 +19,7 @@
#include <tvm/ffi/cast.h>
#include <tvm/ffi/extra/structural_visit.h>
+#include <tvm/s_tir/stmt.h>
#include <tvm/tirx/op.h>
#include "../utils.h"
diff --git a/src/s_tir/schedule/primitive/reduction.cc b/src/s_tir/schedule/primitive/reduction.cc
index bb596dc6ab..5caa524fb9 100644
--- a/src/s_tir/schedule/primitive/reduction.cc
+++ b/src/s_tir/schedule/primitive/reduction.cc
@@ -20,6 +20,7 @@
#include <tvm/ffi/extra/structural_mutate.h>
#include <tvm/ffi/extra/structural_visit.h>
#include <tvm/ffi/reflection/registry.h>
+#include <tvm/s_tir/stmt.h>
#include <tvm/te/operation.h>
#include "../utils.h"
diff --git a/src/s_tir/schedule/primitive/reorder_block_iter_var.cc b/src/s_tir/schedule/primitive/reorder_block_iter_var.cc
index 9eab2390c7..58bc95e036 100644
--- a/src/s_tir/schedule/primitive/reorder_block_iter_var.cc
+++ b/src/s_tir/schedule/primitive/reorder_block_iter_var.cc
@@ -17,6 +17,7 @@
* under the License.
*/
#include <tvm/ffi/cast.h>
+#include <tvm/s_tir/stmt.h>
#include <set>
diff --git a/src/s_tir/schedule/primitive/rolling_buffer.cc b/src/s_tir/schedule/primitive/rolling_buffer.cc
index 0014836b8a..875a34b5f7 100644
--- a/src/s_tir/schedule/primitive/rolling_buffer.cc
+++ b/src/s_tir/schedule/primitive/rolling_buffer.cc
@@ -18,6 +18,7 @@
*/
#include <tvm/ffi/cast.h>
#include <tvm/ffi/extra/structural_mutate.h>
+#include <tvm/s_tir/stmt.h>
#include <functional>
diff --git a/src/s_tir/schedule/schedule.cc b/src/s_tir/schedule/schedule.cc
index e8e666672f..ef6af17c3a 100644
--- a/src/s_tir/schedule/schedule.cc
+++ b/src/s_tir/schedule/schedule.cc
@@ -17,6 +17,7 @@
* under the License.
*/
#include <tvm/ffi/reflection/registry.h>
+#include <tvm/s_tir/stmt.h>
#include "./utils.h"
namespace tvm {
diff --git a/src/s_tir/schedule/state.cc b/src/s_tir/schedule/state.cc
index ab5e6b370b..f07a88c7ba 100644
--- a/src/s_tir/schedule/state.cc
+++ b/src/s_tir/schedule/state.cc
@@ -20,6 +20,7 @@
#include <tvm/ffi/cast.h>
#include <tvm/ffi/extra/structural_mutate.h>
#include <tvm/ffi/reflection/registry.h>
+#include <tvm/s_tir/stmt.h>
#include "./utils.h"
namespace tvm {
@@ -431,7 +432,7 @@ ScheduleState::ScheduleState(IRModule mod, int debug_mask, bool enable_check) {
const BaseFunc& base_func = kv.second;
if (auto opt = base_func.as<PrimFunc>()) {
auto func = opt.value();
- VerifyWellFormed(func);
+ s_tir::VerifyWellFormed(func);
SBlockInfoCollector::Collect(self, func->body);
}
}
diff --git a/src/s_tir/schedule/traced_schedule.cc b/src/s_tir/schedule/traced_schedule.cc
index 646269abfb..8cea5fc2bb 100644
--- a/src/s_tir/schedule/traced_schedule.cc
+++ b/src/s_tir/schedule/traced_schedule.cc
@@ -18,6 +18,8 @@
*/
#include "./traced_schedule.h"
+#include <tvm/s_tir/stmt.h>
+
namespace tvm {
namespace s_tir {
using namespace tvm::tirx;
diff --git a/src/s_tir/schedule/traced_schedule.h b/src/s_tir/schedule/traced_schedule.h
index 52c1842bc0..0ddd2378ba 100644
--- a/src/s_tir/schedule/traced_schedule.h
+++ b/src/s_tir/schedule/traced_schedule.h
@@ -20,6 +20,7 @@
#define TVM_S_TIR_SCHEDULE_TRACED_SCHEDULE_H_
#include <tvm/ir/prim/expr.h>
+#include <tvm/s_tir/stmt.h>
#include "./concrete_schedule.h"
diff --git a/src/s_tir/schedule/transform.cc b/src/s_tir/schedule/transform.cc
index c962e1f0f8..bb0e1565c3 100644
--- a/src/s_tir/schedule/transform.cc
+++ b/src/s_tir/schedule/transform.cc
@@ -20,6 +20,8 @@
#include <tvm/ffi/cast.h>
#include <tvm/ffi/extra/structural_equal.h>
#include <tvm/ffi/reflection/registry.h>
+#include <tvm/s_tir/stmt.h>
+#include <tvm/s_tir/tensor_intrin.h>
#include <tvm/tirx/builtin.h>
#include "../../tirx/transform/ir_utils.h"
@@ -316,7 +318,7 @@ ffi::Optional<LoopRV> TileWithTensorIntrin(const s_tir::Schedule& sch,
const ffi::String& intrin_name, bool allow_padding) {
ffi::Optional<TensorizeInfo> opt_tensorize_info =
GetTensorizeLoopMapping(sch->state(), sch->GetSRef(block_rv),
- tirx::TensorIntrin::Get(intrin_name).value()->desc, allow_padding);
+ TensorIntrin::Get(intrin_name).value()->desc, allow_padding);
if (!opt_tensorize_info) return std::nullopt;
const TensorizeInfoNode* info = opt_tensorize_info.value().get();
if (info->block_iter_paddings.has_value()) {
@@ -468,7 +470,7 @@ void BlockBufferAccessSimplifier::SimplifyBufferIndices(ffi::Array<PrimExpr>* in
UnchangedOr<Stmt> BlockBufferAccessSimplifier::Mutate_(const SBlockNode* op,
InplaceMode inplace_mode) {
- SBlock block = tirx::IRMutatorWithAnalyzer::Mutate_(op, inplace_mode)
+ SBlock block = s_tir::IRMutatorWithAnalyzer::Mutate_(op, inplace_mode)
.ValueOrUnchanged(ffi::GetRef<Stmt>(op))
.as_or_throw<SBlock>();
auto* n = block.CopyOnWrite();
@@ -479,7 +481,7 @@ UnchangedOr<Stmt> BlockBufferAccessSimplifier::Mutate_(const SBlockNode* op,
UnchangedOr<Stmt> BlockBufferAccessSimplifier::Mutate_(const BufferStoreNode* op,
InplaceMode inplace_mode) {
- BufferStore node = tirx::IRMutatorWithAnalyzer::Mutate_(op, inplace_mode)
+ BufferStore node = s_tir::IRMutatorWithAnalyzer::Mutate_(op, inplace_mode)
.ValueOrUnchanged(ffi::GetRef<Stmt>(op))
.as_or_throw<BufferStore>();
SimplifyBufferIndices(&node.CopyOnWrite()->indices);
@@ -488,7 +490,7 @@ UnchangedOr<Stmt> BlockBufferAccessSimplifier::Mutate_(const BufferStoreNode* op
UnchangedOr<PrimExpr> BlockBufferAccessSimplifier::Mutate_(const TensorLoadNode* op,
InplaceMode inplace_mode) {
- TensorLoad node = tirx::IRMutatorWithAnalyzer::Mutate_(op, inplace_mode)
+ TensorLoad node = s_tir::IRMutatorWithAnalyzer::Mutate_(op, inplace_mode)
.ValueOrUnchanged(ffi::GetRef<PrimExpr>(op))
.as_or_throw<TensorLoad>();
SimplifyBufferIndices(&node.CopyOnWrite()->indices);
diff --git a/src/s_tir/schedule/transform.h b/src/s_tir/schedule/transform.h
index 2392937dd1..5cf61192ce 100644
--- a/src/s_tir/schedule/transform.h
+++ b/src/s_tir/schedule/transform.h
@@ -22,12 +22,13 @@
#include <tvm/ir/prim/expr.h>
#include <tvm/s_tir/schedule/schedule.h>
#include <tvm/s_tir/schedule/state.h>
-#include <tvm/tirx/stmt_functor.h>
+#include <tvm/s_tir/stmt.h>
+#include <tvm/s_tir/stmt_functor.h>
#include <unordered_map>
#include <utility>
-#include "../../tirx/ir/ir_mutator_with_analyzer.h"
+#include "../../s_tir/ir/ir_mutator_with_analyzer.h"
namespace tvm {
namespace s_tir {
@@ -210,10 +211,10 @@ ffi::Optional<s_tir::LoopRV> TileWithTensorIntrin(const s_tir::Schedule& sch,
/*!
* \brief Simplifier for indices of buffer access and block buffer access regions.
*/
-class BlockBufferAccessSimplifier : public tirx::IRMutatorWithAnalyzer {
+class BlockBufferAccessSimplifier : public s_tir::IRMutatorWithAnalyzer {
public:
- using tirx::IRMutatorWithAnalyzer::Mutate;
- using tirx::IRMutatorWithAnalyzer::Mutate_;
+ using s_tir::IRMutatorWithAnalyzer::Mutate;
+ using s_tir::IRMutatorWithAnalyzer::Mutate_;
/*!
* \brief Simplify indices of buffer access and block buffer access regions in the statement
diff --git a/src/s_tir/schedule/utils.h b/src/s_tir/schedule/utils.h
index 38330f31c1..9ac1aaec06 100644
--- a/src/s_tir/schedule/utils.h
+++ b/src/s_tir/schedule/utils.h
@@ -26,15 +26,17 @@
#include <tvm/ffi/extra/json.h>
#include <tvm/ffi/extra/serialization.h>
#include <tvm/ir/prim/expr.h>
+#include <tvm/s_tir/analysis.h>
#include <tvm/s_tir/schedule/instruction.h>
#include <tvm/s_tir/schedule/schedule.h>
#include <tvm/s_tir/schedule/state.h>
#include <tvm/s_tir/schedule/trace.h>
+#include <tvm/s_tir/stmt.h>
+#include <tvm/s_tir/stmt_functor.h>
#include <tvm/s_tir/utils.h>
#include <tvm/tirx/analysis.h>
#include <tvm/tirx/function.h>
#include <tvm/tirx/op.h>
-#include <tvm/tirx/stmt_functor.h>
#include <string>
#include <unordered_map>
@@ -369,15 +371,15 @@ inline ffi::String BufferIndexType2Str(BufferIndexType buffer_index_type) {
/*! \brief Returns the names of the blocks in the provided module. */
inline std::unordered_set<std::string> GetSBlockNames(const IRModule& mod) {
- struct BlockNameCollector : public tirx::StmtExprVisitor {
- using tirx::StmtExprVisitor::Visit_;
+ struct BlockNameCollector : public s_tir::StmtExprVisitor {
+ using s_tir::StmtExprVisitor::Visit_;
ffi::Optional<VisitInterrupt> Visit(ffi::AnyView value) override {
if (value.as<ExprNode>()) return std::nullopt;
- return tirx::StmtExprVisitor::Visit(value);
+ return s_tir::StmtExprVisitor::Visit(value);
}
- ffi::Optional<VisitInterrupt> Visit_(const tirx::SBlockNode* block) override {
+ ffi::Optional<VisitInterrupt> Visit_(const s_tir::SBlockNode* block) override {
block_names.insert(block->name_hint);
return StmtExprVisitor::Visit(block->body);
}
diff --git a/src/s_tir/stmt.cc b/src/s_tir/stmt.cc
new file mode 100644
index 0000000000..98753c031b
--- /dev/null
+++ b/src/s_tir/stmt.cc
@@ -0,0 +1,426 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file tvm/s_tir/stmt.cc
+ * \brief Schedulable block definitions and structural traversal.
+ */
+#include <tvm/arith/analyzer.h>
+#include <tvm/ffi/extra/structural_mutate.h>
+#include <tvm/ffi/extra/structural_visit.h>
+#include <tvm/ffi/reflection/registry.h>
+#include <tvm/ir/prim/op.h>
+#include <tvm/s_tir/stmt.h>
+
+namespace tvm {
+namespace s_tir {
+using namespace tvm::tirx;
+using namespace tvm::prim;
+
+namespace {
+
+TVMFFIAny MatchBufferRegionVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
+ const MatchBufferRegionNode* self =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const MatchBufferRegionNode>(
+ value);
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->WithDefRegionKind(
+ kTVMFFIDefRegionKindSimple, [&]() { return visitor->VisitExpected(self->buffer); }));
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->source));
+ return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+}
+
+TVMFFIAny MatchBufferRegionMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
+ const MatchBufferRegionNode* self =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const MatchBufferRegionNode>(
+ value);
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<BufferVar>, mapped_buffer,
+ mutator->WithDefRegionKind(kTVMFFIDefRegionKindSimple, [&]() {
+ return mutator->MutateExpected(self->buffer);
+ }));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<TensorRegion>, mapped_source,
+ mutator->MutateExpected(self->source));
+ if (mapped_buffer.UnchangedOrSameAs(self->buffer) &&
+ mapped_source.UnchangedOrSameAs(self->source)) {
+ return ffi::Unchanged().CopyToTVMFFIAny();
+ }
+ ffi::ObjectPtr<MatchBufferRegionNode> copy = ffi::make_object<MatchBufferRegionNode>(*self);
+ copy->buffer = std::move(mapped_buffer).ValueOrUnchanged(std::move(copy->buffer));
+ copy->source = std::move(mapped_source).ValueOrUnchanged(std::move(copy->source));
+ return ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(copy)));
+}
+
+TVMFFIAny MatchBufferRegionMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
+ ffi::AnyView value) noexcept {
+ MatchBufferRegionNode* self = const_cast<MatchBufferRegionNode*>(
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const MatchBufferRegionNode>(
+ value));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<BufferVar>, mapped_buffer,
+ mutator->WithDefRegionKind(kTVMFFIDefRegionKindSimple, [&]() {
+ return mutator->MutateExpected(self->buffer,
+ ffi::InplaceMode::kAllow);
+ }));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(
+ ffi::UnchangedOr<TensorRegion>, mapped_source,
+ mutator->MutateExpected(self->source, ffi::InplaceMode::kAllow));
+ if (mapped_buffer.UnchangedOrSameAs(self->buffer) &&
+ mapped_source.UnchangedOrSameAs(self->source)) {
+ return ffi::Unchanged().CopyToTVMFFIAny();
+ }
+ if (!mapped_buffer.IsUnchanged()) self->buffer = std::move(mapped_buffer).ValueUnchecked();
+ if (!mapped_source.IsUnchanged()) self->source = std::move(mapped_source).ValueUnchecked();
+ return ffi::Unchanged().CopyToTVMFFIAny();
+}
+
+TVMFFIAny SBlockVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
+ // Establish allocation and match-buffer definitions before their region uses.
+ // Whole iterators and annotations remain part of structural traversal.
+ // skips: name_hint
+ const SBlockNode* self =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const SBlockNode>(value);
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->iter_vars));
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->WithDefRegionKind(
+ kTVMFFIDefRegionKindSimple, [&]() { return visitor->VisitExpected(self->alloc_buffers); }));
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->match_buffers));
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->reads));
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->writes));
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->annotations));
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->init));
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->body));
+ return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+}
+
+TVMFFIAny SBlockMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
+ // Establish allocation and match-buffer definitions before their region uses.
+ // Whole iterators and annotations remain part of structural traversal.
+ // skips: name_hint
+ const SBlockNode* self =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const SBlockNode>(value);
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<ffi::Array<IterVar>>, mapped_iter_vars,
+ mutator->MutateExpected(self->iter_vars));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<ffi::Array<BufferVar>>, mapped_alloc_buffers,
+ mutator->WithDefRegionKind(kTVMFFIDefRegionKindSimple, [&]() {
+ return mutator->MutateExpected(self->alloc_buffers);
+ }));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<ffi::Array<MatchBufferRegion>>,
+ mapped_match_buffers,
+ mutator->MutateExpected(self->match_buffers));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<ffi::Array<TensorRegion>>, mapped_reads,
+ mutator->MutateExpected(self->reads));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<ffi::Array<TensorRegion>>, mapped_writes,
+ mutator->MutateExpected(self->writes));
+ using AnnotationMap = ffi::Map<ffi::String, ffi::Any>;
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<AnnotationMap>, mapped_annotations,
+ mutator->MutateExpected(self->annotations));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<ffi::Optional<Stmt>>, mapped_init,
+ mutator->MutateExpected(self->init));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Stmt>, mapped_body,
+ mutator->MutateExpected(self->body));
+ if (mapped_iter_vars.UnchangedOrSameAs(self->iter_vars) &&
+ mapped_reads.UnchangedOrSameAs(self->reads) &&
+ mapped_writes.UnchangedOrSameAs(self->writes) &&
+ mapped_alloc_buffers.UnchangedOrSameAs(self->alloc_buffers) &&
+ mapped_match_buffers.UnchangedOrSameAs(self->match_buffers) &&
+ mapped_annotations.UnchangedOrSameAs(self->annotations) &&
+ mapped_init.UnchangedOrSameAs(self->init) && mapped_body.UnchangedOrSameAs(self->body)) {
+ return ffi::Unchanged().CopyToTVMFFIAny();
+ }
+ ffi::ObjectPtr<SBlockNode> copy = ffi::make_object<SBlockNode>(*self);
+ copy->iter_vars = std::move(mapped_iter_vars).ValueOrUnchanged(std::move(copy->iter_vars));
+ copy->reads = std::move(mapped_reads).ValueOrUnchanged(std::move(copy->reads));
+ copy->writes = std::move(mapped_writes).ValueOrUnchanged(std::move(copy->writes));
+ copy->alloc_buffers =
+ std::move(mapped_alloc_buffers).ValueOrUnchanged(std::move(copy->alloc_buffers));
+ copy->match_buffers =
+ std::move(mapped_match_buffers).ValueOrUnchanged(std::move(copy->match_buffers));
+ copy->annotations = std::move(mapped_annotations).ValueOrUnchanged(std::move(copy->annotations));
+ copy->init = std::move(mapped_init).ValueOrUnchanged(std::move(copy->init));
+ copy->body = std::move(mapped_body).ValueOrUnchanged(std::move(copy->body));
+ return ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(copy)));
+}
+
+TVMFFIAny SBlockMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
+ ffi::AnyView value) noexcept {
+ // Establish allocation and match-buffer definitions before their region uses.
+ // Whole iterators and annotations remain part of structural traversal.
+ // skips: name_hint
+ SBlockNode* self = const_cast<SBlockNode*>(
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const SBlockNode>(value));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(
+ ffi::UnchangedOr<ffi::Array<IterVar>>, mapped_iter_vars,
+ mutator->MutateExpected(self->iter_vars, ffi::InplaceMode::kAllow));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<ffi::Array<BufferVar>>, mapped_alloc_buffers,
+ mutator->WithDefRegionKind(kTVMFFIDefRegionKindSimple, [&]() {
+ return mutator->MutateExpected(self->alloc_buffers,
+ ffi::InplaceMode::kAllow);
+ }));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(
+ ffi::UnchangedOr<ffi::Array<MatchBufferRegion>>, mapped_match_buffers,
+ mutator->MutateExpected(self->match_buffers, ffi::InplaceMode::kAllow));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<ffi::Array<TensorRegion>>, mapped_reads,
+ mutator->MutateExpected(self->reads, ffi::InplaceMode::kAllow));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(
+ ffi::UnchangedOr<ffi::Array<TensorRegion>>, mapped_writes,
+ mutator->MutateExpected(self->writes, ffi::InplaceMode::kAllow));
+ using AnnotationMap = ffi::Map<ffi::String, ffi::Any>;
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(
+ ffi::UnchangedOr<AnnotationMap>, mapped_annotations,
+ mutator->MutateExpected(self->annotations, ffi::InplaceMode::kAllow));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<ffi::Optional<Stmt>>, mapped_init,
+ mutator->MutateExpected(self->init, ffi::InplaceMode::kAllow));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Stmt>, mapped_body,
+ mutator->MutateExpected(self->body, ffi::InplaceMode::kAllow));
+ if (mapped_iter_vars.UnchangedOrSameAs(self->iter_vars) &&
+ mapped_reads.UnchangedOrSameAs(self->reads) &&
+ mapped_writes.UnchangedOrSameAs(self->writes) &&
+ mapped_alloc_buffers.UnchangedOrSameAs(self->alloc_buffers) &&
+ mapped_match_buffers.UnchangedOrSameAs(self->match_buffers) &&
+ mapped_annotations.UnchangedOrSameAs(self->annotations) &&
+ mapped_init.UnchangedOrSameAs(self->init) && mapped_body.UnchangedOrSameAs(self->body)) {
+ return ffi::Unchanged().CopyToTVMFFIAny();
+ }
+ if (!mapped_iter_vars.IsUnchanged())
+ self->iter_vars = std::move(mapped_iter_vars).ValueUnchecked();
+ if (!mapped_reads.IsUnchanged()) self->reads = std::move(mapped_reads).ValueUnchecked();
+ if (!mapped_writes.IsUnchanged()) self->writes = std::move(mapped_writes).ValueUnchecked();
+ if (!mapped_alloc_buffers.IsUnchanged()) {
+ self->alloc_buffers = std::move(mapped_alloc_buffers).ValueUnchecked();
+ }
+ if (!mapped_match_buffers.IsUnchanged()) {
+ self->match_buffers = std::move(mapped_match_buffers).ValueUnchecked();
+ }
+ if (!mapped_annotations.IsUnchanged())
+ self->annotations = std::move(mapped_annotations).ValueUnchecked();
+ if (!mapped_init.IsUnchanged()) self->init = std::move(mapped_init).ValueUnchecked();
+ if (!mapped_body.IsUnchanged()) self->body = std::move(mapped_body).ValueUnchecked();
+ return ffi::Unchanged().CopyToTVMFFIAny();
+}
+
+TVMFFIAny SBlockRealizeVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
+ const SBlockRealizeNode* self =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const SBlockRealizeNode>(value);
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->iter_values));
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->predicate));
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->block));
+ return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+}
+
+TVMFFIAny SBlockRealizeMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
+ const SBlockRealizeNode* self =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const SBlockRealizeNode>(value);
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<ffi::Array<PrimExpr>>, mapped_iter_values,
+ mutator->MutateExpected(self->iter_values));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<PrimExpr>, mapped_predicate,
+ mutator->MutateExpected(self->predicate));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<SBlock>, mapped_block,
+ mutator->MutateExpected(self->block));
+ if (mapped_iter_values.UnchangedOrSameAs(self->iter_values) &&
+ mapped_predicate.UnchangedOrSameAs(self->predicate) &&
+ mapped_block.UnchangedOrSameAs(self->block)) {
+ return ffi::Unchanged().CopyToTVMFFIAny();
+ }
+ ffi::ObjectPtr<SBlockRealizeNode> copy = ffi::make_object<SBlockRealizeNode>(*self);
+ copy->iter_values = std::move(mapped_iter_values).ValueOrUnchanged(std::move(copy->iter_values));
+ copy->predicate = std::move(mapped_predicate).ValueOrUnchanged(std::move(copy->predicate));
+ copy->block = std::move(mapped_block).ValueOrUnchanged(std::move(copy->block));
+ return ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(copy)));
+}
+
+TVMFFIAny SBlockRealizeMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
+ ffi::AnyView value) noexcept {
+ SBlockRealizeNode* self = const_cast<SBlockRealizeNode*>(
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const SBlockRealizeNode>(value));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(
+ ffi::UnchangedOr<ffi::Array<PrimExpr>>, mapped_iter_values,
+ mutator->MutateExpected(self->iter_values, ffi::InplaceMode::kAllow));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(
+ ffi::UnchangedOr<PrimExpr>, mapped_predicate,
+ mutator->MutateExpected(self->predicate, ffi::InplaceMode::kAllow));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<SBlock>, mapped_block,
+ mutator->MutateExpected(self->block, ffi::InplaceMode::kAllow));
+ if (mapped_iter_values.UnchangedOrSameAs(self->iter_values) &&
+ mapped_predicate.UnchangedOrSameAs(self->predicate) &&
+ mapped_block.UnchangedOrSameAs(self->block)) {
+ return ffi::Unchanged().CopyToTVMFFIAny();
+ }
+ if (!mapped_iter_values.IsUnchanged())
+ self->iter_values = std::move(mapped_iter_values).ValueUnchecked();
+ if (!mapped_predicate.IsUnchanged())
+ self->predicate = std::move(mapped_predicate).ValueUnchecked();
+ if (!mapped_block.IsUnchanged()) self->block = std::move(mapped_block).ValueUnchecked();
+ return ffi::Unchanged().CopyToTVMFFIAny();
+}
+
+} // namespace
+
+// MatchBufferRegion
+MatchBufferRegion::MatchBufferRegion(BufferVar buffer, TensorRegion source) {
+ const BufferVar& source_buffer = source->source.as_or_throw<BufferVar>();
+ TVM_FFI_ICHECK_EQ(source_buffer->shape.size(), source->region.size())
+ << "MatchBufferRegion source must match its buffer rank";
+ arith::Analyzer analyzer;
+ // Check scope and dtype
+ TVM_FFI_ICHECK_EQ(buffer.scope(), source_buffer.scope())
+ << "MatchBuffer " << buffer << " scope mismatch:" << buffer.scope() << " vs. "
+ << source_buffer.scope();
+ TVM_FFI_ICHECK_EQ(buffer->dtype, source_buffer->dtype)
+ << "MatchBuffer " << buffer << " data type mismatch:" << buffer->dtype << " vs. "
+ << source_buffer->dtype;
+
+ // Check data_alignment
+ TVM_FFI_ICHECK(source_buffer->data_alignment % buffer->data_alignment == 0)
+ << "Trying to match buffer to another one with lower alignment requirement "
+ << " required alignment=" << buffer->data_alignment
+ << ", provided alignment=" << source_buffer->data_alignment;
+
+ // Validate shape
+ TVM_FFI_ICHECK(source->region.size() >= buffer->shape.size())
+ << "Dimension of source ffi::Array<Range> expected to be larger or equal than target buffer "
+ "shape, but "
+ "got "
+ << source->region.size() << " vs. " << buffer->shape.size();
+ size_t offset = source->region.size() - buffer->shape.size();
+ for (size_t i = 0; i < offset; ++i) {
+ TVM_FFI_ICHECK(analyzer->CanProve(source->region[i]->extent == 1))
+ << "The higher dimension should be 1, but got " << source->region[i]->extent << ".";
+ }
+ for (size_t i = 0; i < buffer->shape.size(); ++i) {
+ const Range& source_range = source->region[i + offset];
+ const PrimExpr& buffer_shape = buffer->shape[i];
+ if (!buffer_shape.as<PrimVar>()) {
+ TVM_FFI_ICHECK(analyzer->CanProve(source_range->extent == buffer_shape))
+ << "The dimension mismatched between source region and target buffer shape, got "
+ << source_range->extent << " vs. " << buffer_shape << ".";
+ }
+ }
+ // Note that we do not check elem_offset and strides in this function
+ ffi::ObjectPtr<MatchBufferRegionNode> node = ffi::make_object<MatchBufferRegionNode>();
+ node->buffer = std::move(buffer);
+ node->source = std::move(source);
+ data_ = std::move(node);
+}
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+ namespace refl = tvm::ffi::reflection;
+ MatchBufferRegionNode::RegisterReflection();
+ refl::TypeAttrDef<MatchBufferRegionNode>()
+ .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&MatchBufferRegionVisit))
+ .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&MatchBufferRegionMutate))
+ .attr(refl::type_attr::kStructuralMaybeInplaceMutate,
+ reinterpret_cast<void*>(&MatchBufferRegionMaybeInplaceMutate));
+
+ refl::GlobalDef().def("s_tir.MatchBufferRegion", [](BufferVar buffer, TensorRegion source) {
+ return MatchBufferRegion(buffer, source);
+ });
+}
+
+// Block
+SBlock::SBlock(ffi::Array<IterVar> iter_vars, ffi::Array<TensorRegion> reads,
+ ffi::Array<TensorRegion> writes, ffi::String name_hint, Stmt body,
+ ffi::Optional<Stmt> init, ffi::Array<BufferVar> alloc_buffers,
+ ffi::Array<MatchBufferRegion> match_buffers, ffi::Map<ffi::String, Any> annotations,
+ Span span) {
+ for (const auto& regions : {reads, writes}) {
+ for (const TensorRegion& region : regions) {
+ const auto buffer = region->source.as_or_throw<BufferVar>();
+ TVM_FFI_ICHECK_EQ(buffer->shape.size(), region->region.size())
+ << "SBlock region must match its buffer rank";
+ }
+ }
+ ffi::ObjectPtr<SBlockNode> node = ffi::make_object<SBlockNode>();
+ node->iter_vars = std::move(iter_vars);
+ node->reads = std::move(reads);
+ node->writes = std::move(writes);
+ node->name_hint = std::move(name_hint);
+ node->body = std::move(body);
+ node->init = std::move(init);
+ node->alloc_buffers = std::move(alloc_buffers);
+ node->match_buffers = std::move(match_buffers);
+ node->annotations = std::move(annotations);
+ node->span = std::move(span);
+ data_ = std::move(node);
+}
+
+SBlock::SBlock(ffi::String name_hint, Stmt body, ffi::Array<BufferVar> alloc_buffers, Span span) {
+ ffi::ObjectPtr<SBlockNode> node = ffi::make_object<SBlockNode>();
+ node->iter_vars = {};
+ node->reads = {};
+ node->writes = {};
+ node->name_hint = std::move(name_hint);
+ node->body = std::move(body);
+ node->init = std::nullopt;
+ node->alloc_buffers = std::move(alloc_buffers);
+ node->match_buffers = {};
+ node->annotations = {};
+ node->span = std::move(span);
+ data_ = std::move(node);
+}
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+ namespace refl = tvm::ffi::reflection;
+ SBlockNode::RegisterReflection();
+ refl::TypeAttrDef<SBlockNode>()
+ .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&SBlockVisit))
+ .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&SBlockMutate))
+ .attr(refl::type_attr::kStructuralMaybeInplaceMutate,
+ reinterpret_cast<void*>(&SBlockMaybeInplaceMutate));
+
+ refl::GlobalDef().def("s_tir.SBlock",
+ [](ffi::Array<IterVar> iter_vars, ffi::Array<TensorRegion> reads,
+ ffi::Array<TensorRegion> writes, ffi::String name_hint, Stmt body,
+ ffi::Optional<Stmt> init, ffi::Array<BufferVar> alloc_buffers,
+ ffi::Array<MatchBufferRegion> match_buffers,
+ ffi::Map<ffi::String, Any> annotations, Span span) {
+ return SBlock(iter_vars, reads, writes, name_hint, body, init,
+ alloc_buffers, match_buffers, annotations, span);
+ });
+}
+
+// BlockRealize
+SBlockRealize::SBlockRealize(ffi::Array<PrimExpr> values, PrimExpr predicate, SBlock block,
+ Span span) {
+ TVM_FFI_CHECK_EQ(block->iter_vars.size(), values.size(), ValueError)
+ << "BlockRealize needs to have the same number of iter_vars and binding values";
+ PrimType predicate_ty = predicate.ty();
+ TVM_FFI_CHECK(predicate_ty.MatchesCode(DLDataTypeCode::kDLBool), TypeError)
+ << "Expect Block.predicate to be a bool expression";
+ ffi::ObjectPtr<SBlockRealizeNode> node = ffi::make_object<SBlockRealizeNode>();
+ node->iter_values = std::move(values);
+ node->predicate = std::move(predicate);
+ node->block = std::move(block);
+ node->span = std::move(span);
+ data_ = std::move(node);
+}
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+ namespace refl = tvm::ffi::reflection;
+ SBlockRealizeNode::RegisterReflection();
+ refl::TypeAttrDef<SBlockRealizeNode>()
+ .attr(refl::type_attr::kStructuralVisit, reinterpret_cast<void*>(&SBlockRealizeVisit))
+ .attr(refl::type_attr::kStructuralMutate, reinterpret_cast<void*>(&SBlockRealizeMutate))
+ .attr(refl::type_attr::kStructuralMaybeInplaceMutate,
+ reinterpret_cast<void*>(&SBlockRealizeMaybeInplaceMutate));
+
+ refl::GlobalDef().def("s_tir.SBlockRealize", [](ffi::Array<PrimExpr> iter_values,
+ PrimExpr predicate, SBlock block, Span span) {
+ return SBlockRealize(iter_values, predicate, block, span);
+ });
+}
+
+} // namespace s_tir
+} // namespace tvm
diff --git a/src/s_tir/stmt_functor.cc b/src/s_tir/stmt_functor.cc
new file mode 100644
index 0000000000..924a8efb87
--- /dev/null
+++ b/src/s_tir/stmt_functor.cc
@@ -0,0 +1,201 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+/*!
+ * \file stmt_functor.cc
+ * \brief Native traversal of schedulable TIR nodes.
+ */
+#include <tvm/s_tir/stmt_functor.h>
+
+#include <utility>
+#include <vector>
+
+namespace tvm {
+namespace s_tir {
+
+using namespace tirx;
+
+void StmtExprVisitor::InitVTable(VTable* vtable) {
+ tirx::StmtExprVisitor::InitVTable(vtable);
+ SetDispatch<StmtExprVisitor, SBlockNode>(vtable);
+ SetDispatch<StmtExprVisitor, SBlockRealizeNode>(vtable);
+}
+
+void StmtExprMutator::InitVTable(VTable* vtable) {
+ tirx::StmtExprMutator::InitVTable(vtable);
+ SetDispatch<StmtExprMutator, SBlockNode>(vtable);
+ SetDispatch<StmtExprMutator, SBlockRealizeNode>(vtable);
+}
+
+ffi::Optional<VisitInterrupt> StmtExprVisitor::Visit_(const SBlockNode* op) {
+ return VisitBlock(this, op);
+}
+
+ffi::Optional<VisitInterrupt> StmtExprVisitor::VisitBlock(tirx::StmtExprVisitor* visitor,
+ const SBlockNode* op) {
+ for (const IterVar& iter_var : op->iter_vars) {
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->Visit(iter_var->dom->min));
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->Visit(iter_var->dom->extent));
+ }
+ for (const BufferVar& buf : op->alloc_buffers) {
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->WithDefRegionKind(
+ kTVMFFIDefRegionKindSimple, [&]() { return visitor->Visit(buf); }));
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitBufferMetadata(buf));
+ }
+ // Define match-buffer targets before visiting reads/writes that may use them.
+ // This differs from the old TIRX native order (reads/writes before matches)
+ // and agrees with structural traversal's definition-before-use contract.
+ for (const MatchBufferRegion& match_buffer_region : op->match_buffers) {
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->WithDefRegionKind(
+ kTVMFFIDefRegionKindSimple, [&]() { return visitor->Visit(match_buffer_region->buffer); }));
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitBufferMetadata(match_buffer_region->buffer));
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->Visit(match_buffer_region->source));
+ }
+ for (const TensorRegion& region : op->reads) {
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->Visit(region));
+ }
+ for (const TensorRegion& region : op->writes) {
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->Visit(region));
+ }
+ if (op->init.has_value()) {
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->Visit(op->init.value()));
+ }
+ return visitor->Visit(op->body);
+}
+
+ffi::Optional<VisitInterrupt> StmtExprVisitor::Visit_(const SBlockRealizeNode* op) {
+ return VisitBlockRealize(this, op);
+}
+
+ffi::Optional<VisitInterrupt> StmtExprVisitor::VisitBlockRealize(tirx::StmtExprVisitor* visitor,
+ const SBlockRealizeNode* op) {
+ for (const auto& child : op->iter_values) {
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->Visit(child));
+ }
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->Visit(op->predicate));
+ return visitor->Visit(op->block);
+}
+
+UnchangedOr<Stmt> StmtExprMutator::Mutate_(const SBlockNode* op, InplaceMode inplace_mode) {
+ return MutateBlock(this, op, inplace_mode);
+}
+
+UnchangedOr<Stmt> StmtExprMutator::MutateBlock(tirx::StmtExprMutator* mutator, const SBlockNode* op,
+ InplaceMode inplace_mode) {
+ // SBlock iteration variables keep their binders; only their domains are expressions here.
+ const auto* iters = op->iter_vars.GetArrayObj();
+ InplaceMode iter_mode = iters->unique() ? inplace_mode : InplaceMode::kDisallow;
+ std::vector<std::pair<size_t, IterVar>> replacements;
+ for (size_t i = 0; i < iters->size(); ++i) {
+ const auto* iter = (*iters)[i].as<IterVarNode>();
+ InplaceMode domain_mode = iter->unique() ? iter_mode : InplaceMode::kDisallow;
+ auto domain = mutator->Mutate(iter->dom, domain_mode).as_or_throw<UnchangedOr<Range>>();
+ if (domain.UnchangedOrSameAs(iter->dom)) continue;
+ if (domain_mode == InplaceMode::kAllow) {
+ const_cast<IterVarNode*>(iter)->dom = std::move(domain).ValueUnchecked();
+ } else {
+ auto updated = ffi::make_object<IterVarNode>(*iter);
+ updated->dom = std::move(domain).ValueUnchecked();
+ replacements.emplace_back(i, IterVar(std::move(updated)));
+ }
+ }
+ UnchangedOr<ffi::Array<IterVar>> iter_vars = ffi::Unchanged();
+ if (!replacements.empty()) {
+ if (iter_mode == InplaceMode::kAllow) {
+ for (auto& [i, iter] : replacements) {
+ const_cast<ffi::ArrayObj*>(iters)->SetItem(i, std::move(iter));
+ }
+ } else {
+ ffi::Array<IterVar> updated = op->iter_vars;
+ for (auto& [i, iter] : replacements) updated.Set(i, std::move(iter));
+ iter_vars = std::move(updated);
+ }
+ }
+ auto alloc_buffers =
+ mutator
+ ->WithDefRegionKind(kTVMFFIDefRegionKindSimple,
+ [&] { return mutator->Mutate(op->alloc_buffers, inplace_mode); })
+ .as_or_throw<UnchangedOr<ffi::Array<BufferVar>>>();
+ auto match_buffers = mutator->Mutate(op->match_buffers, inplace_mode)
+ .as_or_throw<UnchangedOr<ffi::Array<MatchBufferRegion>>>();
+ auto reads =
+ mutator->Mutate(op->reads, inplace_mode).as_or_throw<UnchangedOr<ffi::Array<TensorRegion>>>();
+ auto writes = mutator->Mutate(op->writes, inplace_mode)
+ .as_or_throw<UnchangedOr<ffi::Array<TensorRegion>>>();
+ auto init =
+ mutator->Mutate(op->init, inplace_mode).as_or_throw<UnchangedOr<ffi::Optional<Stmt>>>();
+ auto body = mutator->Mutate(op->body, inplace_mode);
+ if (iter_vars.UnchangedOrSameAs(op->iter_vars) &&
+ alloc_buffers.UnchangedOrSameAs(op->alloc_buffers) && reads.UnchangedOrSameAs(op->reads) &&
+ writes.UnchangedOrSameAs(op->writes) && match_buffers.UnchangedOrSameAs(op->match_buffers) &&
+ init.UnchangedOrSameAs(op->init) && body.UnchangedOrSameAs(op->body))
+ return ffi::Unchanged();
+ if (inplace_mode == InplaceMode::kAllow) {
+ auto* writable = const_cast<SBlockNode*>(op);
+ if (!iter_vars.IsUnchanged()) writable->iter_vars = std::move(iter_vars).ValueUnchecked();
+ if (!alloc_buffers.IsUnchanged())
+ writable->alloc_buffers = std::move(alloc_buffers).ValueUnchecked();
+ if (!reads.IsUnchanged()) writable->reads = std::move(reads).ValueUnchecked();
+ if (!writes.IsUnchanged()) writable->writes = std::move(writes).ValueUnchecked();
+ if (!match_buffers.IsUnchanged())
+ writable->match_buffers = std::move(match_buffers).ValueUnchecked();
+ if (!init.IsUnchanged()) writable->init = std::move(init).ValueUnchecked();
+ if (!body.IsUnchanged()) writable->body = std::move(body).ValueUnchecked();
+ return ffi::Unchanged();
+ }
+ auto copy = ffi::make_object<SBlockNode>(*op);
+ if (!iter_vars.IsUnchanged()) copy->iter_vars = std::move(iter_vars).ValueUnchecked();
+ if (!alloc_buffers.IsUnchanged()) copy->alloc_buffers = std::move(alloc_buffers).ValueUnchecked();
+ if (!reads.IsUnchanged()) copy->reads = std::move(reads).ValueUnchecked();
+ if (!writes.IsUnchanged()) copy->writes = std::move(writes).ValueUnchecked();
+ if (!match_buffers.IsUnchanged()) copy->match_buffers = std::move(match_buffers).ValueUnchecked();
+ if (!init.IsUnchanged()) copy->init = std::move(init).ValueUnchecked();
+ if (!body.IsUnchanged()) copy->body = std::move(body).ValueUnchecked();
+ return Stmt(std::move(copy));
+}
+
+UnchangedOr<Stmt> StmtExprMutator::Mutate_(const SBlockRealizeNode* op, InplaceMode inplace_mode) {
+ return MutateBlockRealize(this, op, inplace_mode);
+}
+
+UnchangedOr<Stmt> StmtExprMutator::MutateBlockRealize(tirx::StmtExprMutator* mutator,
+ const SBlockRealizeNode* op,
+ InplaceMode inplace_mode) {
+ auto iter_values = mutator->Mutate(op->iter_values, inplace_mode)
+ .as_or_throw<UnchangedOr<ffi::Array<PrimExpr>>>();
+ auto predicate = mutator->Mutate(op->predicate, inplace_mode);
+ auto block = mutator->Mutate(op->block, inplace_mode).as_or_throw<UnchangedOr<SBlock>>();
+ if (iter_values.UnchangedOrSameAs(op->iter_values) &&
+ predicate.UnchangedOrSameAs(op->predicate) && block.UnchangedOrSameAs(op->block))
+ return ffi::Unchanged();
+ if (inplace_mode == InplaceMode::kAllow) {
+ auto* writable = const_cast<SBlockRealizeNode*>(op);
+ if (!iter_values.IsUnchanged()) writable->iter_values = std::move(iter_values).ValueUnchecked();
+ if (!predicate.IsUnchanged()) writable->predicate = std::move(predicate).ValueUnchecked();
+ if (!block.IsUnchanged()) writable->block = std::move(block).ValueUnchecked();
+ return ffi::Unchanged();
+ }
+ auto copy = ffi::make_object<SBlockRealizeNode>(*op);
+ if (!iter_values.IsUnchanged()) copy->iter_values = std::move(iter_values).ValueUnchecked();
+ if (!predicate.IsUnchanged()) copy->predicate = std::move(predicate).ValueUnchecked();
+ if (!block.IsUnchanged()) copy->block = std::move(block).ValueUnchecked();
+ return Stmt(std::move(copy));
+}
+
+} // namespace s_tir
+} // namespace tvm
diff --git a/src/s_tir/tensor_intrin.cc b/src/s_tir/tensor_intrin.cc
new file mode 100644
index 0000000000..e381448198
--- /dev/null
+++ b/src/s_tir/tensor_intrin.cc
@@ -0,0 +1,100 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file src/s_tir/tensor_intrin.cc
+ * \brief Tensor intrinsic registry and construction.
+ */
+#include <tvm/ffi/reflection/registry.h>
+#include <tvm/s_tir/tensor_intrin.h>
+
+namespace tvm {
+namespace s_tir {
+using tirx::BufferTypeNode;
+using tirx::PrimFunc;
+
+TVM_FFI_STATIC_INIT_BLOCK() { TensorIntrinNode::RegisterReflection(); }
+
+class TensorIntrinManager {
+ public:
+ ffi::Map<ffi::String, TensorIntrin> reg;
+
+ static TensorIntrinManager* Global() {
+ static TensorIntrinManager* inst = new TensorIntrinManager();
+ return inst;
+ }
+};
+
+TensorIntrin::TensorIntrin(PrimFunc desc, PrimFunc impl) {
+ // Check the number of func var is equal
+ TVM_FFI_CHECK_EQ(desc->params.size(), impl->params.size(), ValueError)
+ << "The number of parameters of the description and the implementation of the "
+ "tensor intrinsic doesn't match.";
+ auto is_handle = [](const Var& param) {
+ return param->ty.as<PointerTypeNode>() != nullptr || param->ty.as<BufferTypeNode>() != nullptr;
+ };
+ for (size_t i = 0; i < desc->params.size(); i++) {
+ TVM_FFI_CHECK(is_handle(desc->params[i]), ValueError)
+ << "Parameters of the description of the "
+ "tensor intrinsic should be handle only.";
+ TVM_FFI_CHECK(is_handle(impl->params[i]), ValueError)
+ << "Parameters of the implementation of "
+ "the tensor intrinsic should be handle only.";
+ }
+ ffi::ObjectPtr<TensorIntrinNode> n = ffi::make_object<TensorIntrinNode>();
+ n->desc = std::move(desc);
+ n->impl = std::move(impl);
+ data_ = std::move(n);
+}
+
+void TensorIntrin::Register(ffi::String name, TensorIntrin intrin, bool override) {
+ TensorIntrinManager* manager = TensorIntrinManager::Global();
+ if (!override) {
+ TVM_FFI_CHECK_EQ(manager->reg.count(name), 0, ValueError)
+ << "TensorIntrin '" << name << "' has already been registered";
+ }
+ manager->reg.Set(name, intrin);
+}
+
+ffi::Optional<TensorIntrin> TensorIntrin::Get(ffi::String name, bool allow_missing) {
+ const TensorIntrinManager* manager = TensorIntrinManager::Global();
+ auto it = manager->reg.find(name);
+ if (it == manager->reg.end()) {
+ if (allow_missing) {
+ return std::nullopt;
+ } else {
+ TVM_FFI_THROW(ValueError) << "TensorIntrin '" << name << "' is not registered";
+ }
+ }
+ return (*it).second;
+}
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+ namespace refl = tvm::ffi::reflection;
+ refl::GlobalDef()
+ .def("s_tir.TensorIntrin",
+ [](PrimFunc desc_func, PrimFunc intrin_func) {
+ return TensorIntrin(desc_func, intrin_func);
+ })
+ .def("s_tir.TensorIntrinRegister", TensorIntrin::Register)
+ .def("s_tir.TensorIntrinGet", TensorIntrin::Get);
+}
+
+} // namespace s_tir
+} // namespace tvm
diff --git a/src/s_tir/transform/annotate_irregular_loop.cc b/src/s_tir/transform/annotate_irregular_loop.cc
index d3067da5b8..b7ac58630d 100644
--- a/src/s_tir/transform/annotate_irregular_loop.cc
+++ b/src/s_tir/transform/annotate_irregular_loop.cc
@@ -24,9 +24,9 @@
#include <tvm/ir/prim/expr.h>
#include <tvm/ir/transform.h>
#include <tvm/s_tir/stmt.h>
+#include <tvm/s_tir/stmt_functor.h>
#include <tvm/s_tir/transform.h>
#include <tvm/tirx/builtin.h>
-#include <tvm/tirx/stmt_functor.h>
namespace tvm {
namespace s_tir {
diff --git a/src/s_tir/transform/bound_checker.cc b/src/s_tir/transform/bound_checker.cc
index 5b6a9fbcf5..00d8b1f8b1 100644
--- a/src/s_tir/transform/bound_checker.cc
+++ b/src/s_tir/transform/bound_checker.cc
@@ -29,10 +29,10 @@
#include <tvm/ir/prim/builtin.h>
#include <tvm/ir/prim/expr.h>
#include <tvm/s_tir/stmt.h>
+#include <tvm/s_tir/stmt_functor.h>
#include <tvm/s_tir/transform.h>
#include <tvm/tirx/builtin.h>
#include <tvm/tirx/op.h>
-#include <tvm/tirx/stmt_functor.h>
#include <unordered_map>
#include <utility>
diff --git a/src/s_tir/transform/canonicalize_loop.cc b/src/s_tir/transform/canonicalize_loop.cc
index 048595c1dd..101acfd07f 100644
--- a/src/s_tir/transform/canonicalize_loop.cc
+++ b/src/s_tir/transform/canonicalize_loop.cc
@@ -25,10 +25,10 @@
#include <tvm/ffi/cast.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/runtime/device_api.h>
+#include <tvm/s_tir/stmt_functor.h>
#include <tvm/s_tir/transform.h>
#include <tvm/tirx/function.h>
#include <tvm/tirx/op.h>
-#include <tvm/tirx/stmt_functor.h>
#include <utility>
diff --git a/src/s_tir/transform/compact_buffer_region.cc b/src/s_tir/transform/compact_buffer_region.cc
index bd5e5042e4..4d3535fd4f 100644
--- a/src/s_tir/transform/compact_buffer_region.cc
+++ b/src/s_tir/transform/compact_buffer_region.cc
@@ -27,19 +27,19 @@
#include <tvm/ffi/extra/structural_visit.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/s_tir/stmt.h>
+#include <tvm/s_tir/stmt_functor.h>
#include <tvm/s_tir/transform.h>
#include <tvm/tirx/op.h>
-#include <tvm/tirx/stmt_functor.h>
#include <numeric>
#include <stack>
#include "../../support/arena.h"
#include "../../support/utils.h"
-#include "../../tirx/transform/ir_utils.h"
#include "../analysis/conditional_bounds.h"
#include "../schedule/utils.h"
#include "../support/nd_int_set.h"
+#include "ir_utils.h"
namespace tvm {
namespace s_tir {
diff --git a/src/s_tir/transform/convert_blocks_to_opaque.cc b/src/s_tir/transform/convert_blocks_to_opaque.cc
index 2c8db5af32..c129170101 100644
--- a/src/s_tir/transform/convert_blocks_to_opaque.cc
+++ b/src/s_tir/transform/convert_blocks_to_opaque.cc
@@ -24,8 +24,9 @@
#include <tvm/ffi/cast.h>
#include <tvm/ffi/reflection/registry.h>
+#include <tvm/s_tir/stmt.h>
+#include <tvm/s_tir/stmt_functor.h>
#include <tvm/s_tir/transform.h>
-#include <tvm/tirx/stmt_functor.h>
#include "../../tirx/transform/ir_utils.h"
diff --git a/src/s_tir/transform/default_gpu_schedule.cc b/src/s_tir/transform/default_gpu_schedule.cc
index 9839fd47d2..35686953f3 100644
--- a/src/s_tir/transform/default_gpu_schedule.cc
+++ b/src/s_tir/transform/default_gpu_schedule.cc
@@ -19,6 +19,7 @@
#include <tvm/ffi/cast.h>
#include <tvm/ffi/reflection/registry.h>
+#include <tvm/s_tir/stmt.h>
#include "../meta_schedule/utils.h"
@@ -119,7 +120,7 @@ IRModule MarkScheduled(const IRModule& mod) {
* iter_values/iter_vars counts consistent for downstream checks.
*/
tirx::PrimFunc WrapBareSBlockBody(const tirx::PrimFunc& func) {
- const auto* realize = func->body.as<tirx::SBlockRealizeNode>();
+ const auto* realize = func->body.as<s_tir::SBlockRealizeNode>();
if (realize == nullptr || !realize->block->iter_vars.empty()) {
return func;
}
@@ -128,7 +129,7 @@ tirx::PrimFunc WrapBareSBlockBody(const tirx::PrimFunc& func) {
// whose block body is a For loop (or a nested SBlockRealize) — that case
// already has somewhere to put thread bindings, so leave it alone.
const tirx::Stmt& inner = realize->block->body;
- if (inner->IsInstance<tirx::ForNode>() || inner->IsInstance<tirx::SBlockRealizeNode>()) {
+ if (inner->IsInstance<tirx::ForNode>() || inner->IsInstance<s_tir::SBlockRealizeNode>()) {
return func;
}
tvm::IntImm zero(tvm::PrimType::Int(32), 0);
@@ -137,19 +138,19 @@ tirx::PrimFunc WrapBareSBlockBody(const tirx::PrimFunc& func) {
tirx::Var iter_var_var("vu", tvm::PrimType::Int(32));
tirx::IterVar new_iter(tvm::Range::FromMinExtent(zero, one), iter_var_var.as_or_throw<PrimVar>(),
tirx::IterVarType::kDataPar);
- tirx::SBlock inner_block = realize->block;
+ s_tir::SBlock inner_block = realize->block;
inner_block.CopyOnWrite()->iter_vars = ffi::Array<tirx::IterVar>{new_iter};
- tirx::SBlockRealize inner_realize(
+ s_tir::SBlockRealize inner_realize(
/*iter_values=*/ffi::Array<tvm::PrimExpr>{loop_var.as_or_throw<tvm::PrimExpr>()},
/*predicate=*/realize->predicate, inner_block);
tirx::Stmt for_stmt =
tirx::For(loop_var.as_or_throw<PrimVar>(), zero, one, tirx::ForKind::kSerial, inner_realize);
- tirx::SBlock root_block(/*iter_vars=*/ffi::Array<tirx::IterVar>{},
- /*reads=*/ffi::Array<tvm::TensorRegion>{},
- /*writes=*/ffi::Array<tvm::TensorRegion>{},
- /*name_hint=*/"root", /*body=*/for_stmt);
- tirx::SBlockRealize root_realize(/*iter_values=*/ffi::Array<tvm::PrimExpr>{},
- /*predicate=*/IntImm::Bool(true), root_block);
+ s_tir::SBlock root_block(/*iter_vars=*/ffi::Array<tirx::IterVar>{},
+ /*reads=*/ffi::Array<tvm::TensorRegion>{},
+ /*writes=*/ffi::Array<tvm::TensorRegion>{},
+ /*name_hint=*/"root", /*body=*/for_stmt);
+ s_tir::SBlockRealize root_realize(/*iter_values=*/ffi::Array<tvm::PrimExpr>{},
+ /*predicate=*/IntImm::Bool(true), root_block);
tirx::PrimFunc result = func;
result.CopyOnWrite()->body = std::move(root_realize);
return result;
diff --git a/src/s_tir/transform/hoist_expression.cc b/src/s_tir/transform/hoist_expression.cc
index 9af41d25c5..3d4c52df1f 100644
--- a/src/s_tir/transform/hoist_expression.cc
+++ b/src/s_tir/transform/hoist_expression.cc
@@ -26,9 +26,10 @@
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/prim/expr.h>
+#include <tvm/s_tir/analysis.h>
+#include <tvm/s_tir/stmt_functor.h>
#include <tvm/s_tir/transform.h>
#include <tvm/tirx/analysis.h>
-#include <tvm/tirx/stmt_functor.h>
#include <queue>
#include <unordered_map>
@@ -37,8 +38,8 @@
#include "../../arith/interval_set.h"
#include "../../runtime/thread_storage_scope.h"
-#include "../../tirx/ir/ir_mutator_with_analyzer.h"
-#include "../../tirx/transform/ir_utils.h"
+#include "../../s_tir/ir/ir_mutator_with_analyzer.h"
+#include "ir_utils.h"
namespace tvm {
namespace s_tir {
@@ -455,10 +456,10 @@ class HoistInfoCollector : public StmtExprVisitor {
std::unordered_set<const VarNode*> active_loop_vars;
};
-class ExpressionHoister : public tirx::IRMutatorWithAnalyzer {
+class ExpressionHoister : public s_tir::IRMutatorWithAnalyzer {
public:
- using tirx::IRMutatorWithAnalyzer::Mutate;
- using tirx::IRMutatorWithAnalyzer::Mutate_;
+ using s_tir::IRMutatorWithAnalyzer::Mutate;
+ using s_tir::IRMutatorWithAnalyzer::Mutate_;
static Stmt Hoist(Stmt stmt, HoistExpressionConfig config) {
auto loop_info = HoistInfoCollector::Collect(stmt, config);
@@ -466,12 +467,12 @@ class ExpressionHoister : public tirx::IRMutatorWithAnalyzer {
arith::Analyzer analyzer;
auto hoister = ffi::make_object<ExpressionHoister>(std::move(loop_info), config, analyzer);
stmt = hoister->Mutate(stmt, InplaceMode::kAllow).ValueOrUnchanged(std::move(stmt));
- stmt = ConvertSSA(std::move(stmt));
+ stmt = s_tir::ConvertSSA(std::move(stmt));
return stmt;
}
private:
- using Parent = tirx::IRMutatorWithAnalyzer;
+ using Parent = s_tir::IRMutatorWithAnalyzer;
public:
explicit ExpressionHoister(std::vector<HoistInfoCollector::HoistInfo> loop_info,
@@ -592,7 +593,7 @@ Pass HoistExpression() {
return tvm::transform::Sequential(
{
insertion_pass,
- tirx::transform::StmtSimplify(),
+ s_tir::transform::StmtSimplify(),
tirx::transform::RemoveNoOp(),
},
"s_tir.HoistExpression");
@@ -630,7 +631,7 @@ static Pass HoistIfThenElseImpl() {
return tvm::transform::Sequential(
{
insertion_pass,
- tirx::transform::StmtSimplify(),
+ s_tir::transform::StmtSimplify(),
tirx::transform::RemoveNoOp(),
},
"s_tir.HoistIfThenElse");
@@ -648,7 +649,7 @@ static Pass HoistIfThenElseBasicImpl() {
return tvm::transform::Sequential(
{
insertion_pass,
- tirx::transform::StmtSimplify(),
+ s_tir::transform::StmtSimplify(),
tirx::transform::RemoveNoOp(),
},
"s_tir.HoistIfThenElseBasic");
diff --git a/src/s_tir/transform/inject_double_buffer.cc b/src/s_tir/transform/inject_double_buffer.cc
index 2100a82780..7474c552ca 100644
--- a/src/s_tir/transform/inject_double_buffer.cc
+++ b/src/s_tir/transform/inject_double_buffer.cc
@@ -27,11 +27,11 @@
#include <tvm/ffi/reflection/registry.h>
#include <tvm/runtime/logging.h>
#include <tvm/s_tir/stmt.h>
+#include <tvm/s_tir/stmt_functor.h>
#include <tvm/s_tir/transform.h>
#include <tvm/tirx/op.h>
-#include <tvm/tirx/stmt_functor.h>
-#include "../../tirx/transform/ir_utils.h"
+#include "ir_utils.h"
namespace tvm {
namespace s_tir {
@@ -161,7 +161,7 @@ class DoubleBufferInjector : public StmtExprMutator {
for (const VarNode* v : detector->touched_) {
dbuffer_info_[v] = StorageEntry();
}
- return ConvertSSA(Mutate(stmt, InplaceMode::kAllow).ValueOrUnchanged(std::move(stmt)));
+ return s_tir::ConvertSSA(Mutate(stmt, InplaceMode::kAllow).ValueOrUnchanged(std::move(stmt)));
}
UnchangedOr<Stmt> Mutate_(const AttrStmtNode* op, InplaceMode inplace_mode) final {
diff --git a/src/s_tir/transform/inject_permuted_layout.cc b/src/s_tir/transform/inject_permuted_layout.cc
index c106515890..4f5fe0c506 100644
--- a/src/s_tir/transform/inject_permuted_layout.cc
+++ b/src/s_tir/transform/inject_permuted_layout.cc
@@ -24,14 +24,15 @@
#include <tvm/arith/analyzer.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/op.h>
+#include <tvm/s_tir/stmt.h>
+#include <tvm/s_tir/stmt_functor.h>
#include <tvm/s_tir/transform.h>
#include <tvm/tirx/function.h>
#include <tvm/tirx/op.h>
-#include <tvm/tirx/stmt_functor.h>
#include "../../runtime/thread_storage_scope.h"
+#include "../../s_tir/ir/ir_mutator_with_analyzer.h"
#include "../../support/utils.h"
-#include "../../tirx/ir/ir_mutator_with_analyzer.h"
#include "../../tirx/transform/ir_utils.h"
namespace tvm {
diff --git a/src/s_tir/transform/inject_ptx_async_copy.cc b/src/s_tir/transform/inject_ptx_async_copy.cc
index 1d206840e3..58bf05be31 100644
--- a/src/s_tir/transform/inject_ptx_async_copy.cc
+++ b/src/s_tir/transform/inject_ptx_async_copy.cc
@@ -25,11 +25,12 @@
#include <tvm/ir/op.h>
#include <tvm/ir/prim/builtin.h>
#include <tvm/ir/prim/expr.h>
+#include <tvm/s_tir/analysis.h>
+#include <tvm/s_tir/stmt_functor.h>
#include <tvm/s_tir/transform.h>
#include <tvm/tirx/analysis.h>
#include <tvm/tirx/builtin.h>
#include <tvm/tirx/op.h>
-#include <tvm/tirx/stmt_functor.h>
#include "../../tirx/ir/buffer_common.h"
#include "storage_access.h"
diff --git a/src/s_tir/transform/inject_ptx_ldg32.cc b/src/s_tir/transform/inject_ptx_ldg32.cc
index 8eb4fe739a..97f34a29cd 100644
--- a/src/s_tir/transform/inject_ptx_ldg32.cc
+++ b/src/s_tir/transform/inject_ptx_ldg32.cc
@@ -22,11 +22,12 @@
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/op.h>
+#include <tvm/s_tir/analysis.h>
+#include <tvm/s_tir/stmt_functor.h>
#include <tvm/s_tir/transform.h>
#include <tvm/tirx/analysis.h>
#include <tvm/tirx/op.h>
#include <tvm/tirx/stmt.h>
-#include <tvm/tirx/stmt_functor.h>
#include "../../arith/const_fold.h"
#include "../../arith/pattern_match.h"
diff --git a/src/s_tir/transform/inject_software_pipeline.cc b/src/s_tir/transform/inject_software_pipeline.cc
index 80fe163f3d..0dd4c8a0a6 100644
--- a/src/s_tir/transform/inject_software_pipeline.cc
+++ b/src/s_tir/transform/inject_software_pipeline.cc
@@ -36,8 +36,8 @@
#include <unordered_set>
#include "../../support/utils.h"
-#include "../../tirx/transform/ir_utils.h"
#include "../schedule/utils.h"
+#include "ir_utils.h"
namespace tvm {
namespace s_tir {
@@ -1354,7 +1354,7 @@ Pass InjectSoftwarePipeline() {
auto pass_func = [=](PrimFunc f, IRModule m, PassContext ctx) {
auto* fptr = f.CopyOnWrite();
fptr->body = software_pipeline::PipelineInjector::Inject(f);
- fptr->body = ConvertSSA(std::move(fptr->body));
+ fptr->body = s_tir::ConvertSSA(std::move(fptr->body));
return f;
};
return CreatePrimFuncPass(pass_func, 0, "s_tir.InjectSoftwarePipeline", {});
diff --git a/src/s_tir/transform/inject_virtual_thread.cc b/src/s_tir/transform/inject_virtual_thread.cc
index 99ccdeea8b..b80cee1f49 100644
--- a/src/s_tir/transform/inject_virtual_thread.cc
+++ b/src/s_tir/transform/inject_virtual_thread.cc
@@ -27,14 +27,14 @@
#include <tvm/ir/prim/builtin.h>
#include <tvm/ir/prim/expr.h>
#include <tvm/s_tir/stmt.h>
+#include <tvm/s_tir/stmt_functor.h>
#include <tvm/s_tir/transform.h>
#include <tvm/tirx/builtin.h>
-#include <tvm/tirx/stmt_functor.h>
#include <unordered_set>
-#include "../../tirx/ir/ir_mutator_with_analyzer.h"
-#include "../../tirx/transform/ir_utils.h"
+#include "../../s_tir/ir/ir_mutator_with_analyzer.h"
+#include "ir_utils.h"
namespace tvm {
namespace s_tir {
@@ -230,10 +230,10 @@ class VarTouchedAnalysis : public StmtExprVisitor {
// Inject virtual thread loop
// rewrite the buffer access pattern when necessary.
-class VTInjector : public tirx::IRMutatorWithAnalyzer {
+class VTInjector : public s_tir::IRMutatorWithAnalyzer {
public:
- using tirx::IRMutatorWithAnalyzer::Mutate;
- using tirx::IRMutatorWithAnalyzer::Mutate_;
+ using s_tir::IRMutatorWithAnalyzer::Mutate;
+ using s_tir::IRMutatorWithAnalyzer::Mutate_;
// constructor
VTInjector(arith::AnalyzerObj* analyzer, Var var, int num_threads,
@@ -686,10 +686,10 @@ class VTInjector : public tirx::IRMutatorWithAnalyzer {
*/
};
-class VirtualThreadInjector : public tirx::IRMutatorWithAnalyzer {
+class VirtualThreadInjector : public s_tir::IRMutatorWithAnalyzer {
public:
- using tirx::IRMutatorWithAnalyzer::Mutate;
- using tirx::IRMutatorWithAnalyzer::Mutate_;
+ using s_tir::IRMutatorWithAnalyzer::Mutate;
+ using s_tir::IRMutatorWithAnalyzer::Mutate_;
using IRMutatorWithAnalyzer::IRMutatorWithAnalyzer;
@@ -722,7 +722,7 @@ Pass InjectVirtualThread() {
n->body = ffi::make_object<VirtualThreadInjector>(analyzer)
->Mutate(n->body, InplaceMode::kAllow)
.ValueOrUnchanged(std::move(n->body));
- n->body = ConvertSSA(std::move(n->body));
+ n->body = s_tir::ConvertSSA(std::move(n->body));
return f;
};
return CreatePrimFuncPass(pass_func, 0, "s_tir.InjectVirtualThread", {});
diff --git a/src/s_tir/transform/ir_utils.cc b/src/s_tir/transform/ir_utils.cc
new file mode 100644
index 0000000000..adb3e54c5f
--- /dev/null
+++ b/src/s_tir/transform/ir_utils.cc
@@ -0,0 +1,201 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+#include "ir_utils.h"
+
+#include <tvm/arith/analyzer.h>
+#include <tvm/s_tir/stmt_functor.h>
+#include <tvm/s_tir/transform.h>
+#include <tvm/tirx/op.h>
+
+namespace tvm {
+namespace s_tir {
+using namespace tirx;
+using namespace tvm::prim;
+
+namespace {
+class SIRConvertSSA final : public tirx::IRConvertSSA {
+ public:
+ TVM_DEFINE_OBJECT_FUNCTOR_DEFAULT_CONSTRUCTOR(SIRConvertSSA, tirx::IRConvertSSA)
+ using tirx::IRConvertSSA::Mutate_;
+
+ protected:
+ static void InitVTable(VTable* table) {
+ tirx::IRConvertSSA::InitVTable(table);
+ SetDispatch<SIRConvertSSA, SBlockNode>(table);
+ SetDispatch<SIRConvertSSA, SBlockRealizeNode>(table);
+ }
+
+ public:
+ UnchangedOr<Stmt> Mutate_(const SBlockNode* op, InplaceMode mode) {
+ SBlock block = ffi::GetRef<SBlock>(op);
+ return WithScope([&]() -> Stmt {
+ auto iter_vars = op->iter_vars.Map([&](IterVar iter) {
+ Var var = DefineVar(iter->var);
+ if (!var.same_as(iter->var)) iter.CopyOnWrite()->var = var.as_or_throw<PrimVar>();
+ return iter;
+ });
+ auto remap_region = [&](TensorRegion region) {
+ BufferVar buffer = GetRemappedBuffer(region->source.as_or_throw<BufferVar>());
+ if (!buffer.same_as(region->source.as_or_throw<BufferVar>())) {
+ region.CopyOnWrite()->source = buffer.var();
+ }
+ return region;
+ };
+ auto reads = block->reads.Map(remap_region);
+ auto writes = block->writes.Map(remap_region);
+ if (!reads.same_as(block->reads) || !writes.same_as(block->writes) ||
+ !iter_vars.same_as(op->iter_vars)) {
+ auto* writer = block.CopyOnWrite();
+ writer->reads = reads;
+ writer->writes = writes;
+ writer->iter_vars = iter_vars;
+ }
+ return s_tir::StmtExprMutator::MutateBlock(this, block.get(),
+ block.unique() ? mode : InplaceMode::kDisallow)
+ .ValueOrUnchanged(block);
+ });
+ }
+
+ UnchangedOr<Stmt> Mutate_(const SBlockRealizeNode* op, InplaceMode mode) {
+ return s_tir::StmtExprMutator::MutateBlockRealize(this, op, mode);
+ }
+};
+} // namespace
+
+Stmt ConvertSSA(Stmt stmt) {
+ return ffi::make_object<SIRConvertSSA>()
+ ->Mutate(stmt, InplaceMode::kAllow)
+ .ValueOrUnchanged(stmt);
+}
+
+IRModule ConvertSSA(IRModule mod) {
+ return ffi::make_object<SIRConvertSSA>()->VisitIRModule(std::move(mod));
+}
+
+namespace transform {
+Pass ConvertSSA() {
+ auto pass_func = [](IRModule mod, PassContext ctx) { return s_tir::ConvertSSA(std::move(mod)); };
+ return tvm::transform::CreateModulePass(pass_func, 0, "s_tir.ConvertSSA", {});
+}
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+ ffi::reflection::GlobalDef().def("s_tir.transform.ConvertSSA", ConvertSSA);
+}
+} // namespace transform
+
+ffi::Array<PrimExpr> ConvertIndices(const MatchBufferRegion& match_buffer,
+ const ffi::Array<PrimExpr>& indices) {
+ const BufferVar& target = match_buffer->buffer;
+ const TensorRegion& source = match_buffer->source;
+ TVM_FFI_ICHECK_EQ(indices.size(), target->shape.size());
+
+ arith::Analyzer analyzer;
+ ffi::Array<PrimExpr> result;
+ result.reserve(source->region.size());
+ size_t offset = source->region.size() - indices.size();
+ for (size_t i = 0; i < offset; ++i) {
+ const Range& range = source->region[i];
+ TVM_FFI_ICHECK(analyzer->CanProve(range->extent == 1));
+ result.push_back(range->min);
+ }
+ for (size_t i = 0; i < indices.size(); ++i) {
+ const Range& range = source->region[i + offset];
+ const PrimExpr& index = indices[i];
+ result.push_back(range->min + index);
+ }
+ return result;
+}
+
+Region ConvertRegion(const MatchBufferRegion& match_buffer, const Region& region) {
+ const BufferVar& target = match_buffer->buffer;
+ const TensorRegion& source = match_buffer->source;
+ TVM_FFI_ICHECK_EQ(region.size(), target->shape.size());
+
+ arith::Analyzer analyzer;
+ Region result;
+ result.reserve(source->region.size());
+ size_t offset = source->region.size() - region.size();
+ for (size_t i = 0; i < offset; ++i) {
+ const Range& source_range = source->region[i];
+ TVM_FFI_ICHECK(analyzer->CanProve(source_range->extent == 1));
+ result.push_back(Range::FromMinExtent(source_range->min, 1));
+ }
+ for (size_t i = 0; i < region.size(); ++i) {
+ const Range& source_range = source->region[i + offset];
+ const Range& target_range = region[i];
+ result.push_back(
+ Range::FromMinExtent(source_range->min + target_range->min, target_range->extent));
+ }
+ return result;
+}
+
+/*! \brief Collect storage alignment information from annotations. */
+class StorageAlignCollector : public StmtExprVisitor {
+ public:
+ ffi::Optional<VisitInterrupt> Visit(ffi::AnyView value) override {
+ if (value.as<ExprNode>()) return std::nullopt;
+ return StmtExprVisitor::Visit(value);
+ }
+
+ private:
+ friend std::unordered_map<Var, StorageAlignAnnotation> CollectStorageAlignAnnotation(
+ const Stmt& body);
+
+ /*! \brief SBlock: resolve each annotation's buffer index through the write regions. */
+ ffi::Optional<VisitInterrupt> Visit_(const SBlockNode* op) final {
+ auto it = op->annotations.find(attr::buffer_dim_align);
+ if (it != op->annotations.end()) {
+ auto annotation = (*it).second.as_or_throw<StorageAlignAnnotation>();
+ for (const auto& item : annotation) {
+ storage_align_[op->writes[item.get<0>()]->source.as_or_throw<BufferVar>().var()].push_back(
+ item);
+ }
+ }
+ return StmtExprVisitor::Visit_(op);
+ }
+
+ /*! \brief AllocBuffer: check for buffer_dim_align annotations. */
+ ffi::Optional<VisitInterrupt> Visit_(const AllocBufferNode* op) final {
+ auto it = op->annotations.find(attr::buffer_dim_align);
+ if (it != op->annotations.end()) {
+ auto storage_align_annotation = (*it).second.as_or_throw<StorageAlignAnnotation>();
+ for (const auto& storage_align_tuple : storage_align_annotation) {
+ int buffer_index = storage_align_tuple.get<0>();
+ // the first buffer idx info is meaningless for alloc
+ // stmt and should set as negative intentionally.
+ TVM_FFI_ICHECK_EQ(buffer_index, -1);
+ storage_align_[op->buffer.var()].push_back(storage_align_tuple);
+ }
+ }
+ return StmtExprVisitor::Visit_(op);
+ }
+
+ /*! \brief The map from buffer var to its storage alignment information. */
+ std::unordered_map<Var, StorageAlignAnnotation> storage_align_;
+};
+
+std::unordered_map<Var, StorageAlignAnnotation> CollectStorageAlignAnnotation(const Stmt& body) {
+ auto collector = ffi::make_object<StorageAlignCollector>();
+ collector->Visit(body);
+ return std::move(collector->storage_align_);
+}
+
+} // namespace s_tir
+} // namespace tvm
diff --git a/src/s_tir/transform/ir_utils.h b/src/s_tir/transform/ir_utils.h
new file mode 100644
index 0000000000..e9d6c6cc2d
--- /dev/null
+++ b/src/s_tir/transform/ir_utils.h
@@ -0,0 +1,66 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+#ifndef TVM_S_TIR_TRANSFORM_IR_UTILS_H_
+#define TVM_S_TIR_TRANSFORM_IR_UTILS_H_
+
+#include <tvm/ffi/container/tuple.h>
+#include <tvm/s_tir/stmt.h>
+
+#include <unordered_map>
+
+#include "../../tirx/transform/ir_utils.h"
+
+namespace tvm {
+namespace s_tir {
+
+/*! \brief Convert a schedulable statement or module to SSA form. */
+tirx::Stmt ConvertSSA(tirx::Stmt stmt);
+IRModule ConvertSSA(IRModule mod);
+
+/*!
+ * \brief Convert match buffer target buffer access indices to original one.
+ * \param indices The indices of the target buffer
+ * \return The indices of source buffer.
+ */
+ffi::Array<PrimExpr> ConvertIndices(const MatchBufferRegion& match_buffer,
+ const ffi::Array<PrimExpr>& indices);
+
+/*!
+ * \brief Convert match buffer target buffer region to original one.
+ * \param region The sub-region of the target buffer
+ * \return The region of source buffer.
+ */
+tirx::Region ConvertRegion(const MatchBufferRegion& match_buffer, const tirx::Region& region);
+
+/*! \brief The quad used by StorageAlign for (buffer_idx, axis, factor, offset) */
+using StorageAlignTuple = ffi::Tuple<int32_t, int32_t, int32_t, int32_t>;
+/*! \brief A list of StorageAlignTuple, used by StorageAlign */
+using StorageAlignAnnotation = ffi::Array<StorageAlignTuple>;
+/*!
+ * \brief Collect storage alignment annotations for all buffer vars within body.
+ * \param body The stmt to collect.
+ * \return The result dict from buffer var to storage align annotations.
+ */
+std::unordered_map<tirx::Var, StorageAlignAnnotation> CollectStorageAlignAnnotation(
+ const tirx::Stmt& body);
+
+} // namespace s_tir
+} // namespace tvm
+#endif // TVM_S_TIR_TRANSFORM_IR_UTILS_H_
diff --git a/src/s_tir/transform/lift_thread_binding.cc b/src/s_tir/transform/lift_thread_binding.cc
index f46047e010..c252563b7f 100644
--- a/src/s_tir/transform/lift_thread_binding.cc
+++ b/src/s_tir/transform/lift_thread_binding.cc
@@ -24,8 +24,8 @@
#include <tvm/ffi/cast.h>
#include <tvm/ffi/reflection/registry.h>
+#include <tvm/s_tir/stmt_functor.h>
#include <tvm/s_tir/transform.h>
-#include <tvm/tirx/stmt_functor.h>
#include "../../runtime/thread_storage_scope.h"
#include "../../tirx/transform/ir_utils.h"
diff --git a/src/s_tir/transform/loop_partition.cc b/src/s_tir/transform/loop_partition.cc
index 16668e8373..392d821b71 100644
--- a/src/s_tir/transform/loop_partition.cc
+++ b/src/s_tir/transform/loop_partition.cc
@@ -30,11 +30,12 @@
#include <tvm/ir/prim/builtin.h>
#include <tvm/ir/prim/expr.h>
#include <tvm/runtime/logging.h>
+#include <tvm/s_tir/analysis.h>
#include <tvm/s_tir/stmt.h>
+#include <tvm/s_tir/stmt_functor.h>
#include <tvm/s_tir/transform.h>
#include <tvm/tirx/analysis.h>
#include <tvm/tirx/builtin.h>
-#include <tvm/tirx/stmt_functor.h>
#include <optional>
#include <unordered_map>
@@ -42,7 +43,7 @@
#include "../../arith/interval_set.h"
#include "../../runtime/thread_storage_scope.h"
-#include "../../tirx/transform/ir_utils.h"
+#include "ir_utils.h"
namespace tvm {
namespace s_tir {
@@ -843,7 +844,7 @@ Stmt LoopPartitioner::TryPartition(const Stmt& stmt, Var var, PrimExpr min, Prim
->Mutate(stmt)
.ValueOrUnchanged(stmt);
}
- s = ConvertSSA(s);
+ s = s_tir::ConvertSSA(s);
return s;
}
diff --git a/src/s_tir/transform/lower_async_dma.cc b/src/s_tir/transform/lower_async_dma.cc
index 389c888773..bc1dfc6d43 100644
--- a/src/s_tir/transform/lower_async_dma.cc
+++ b/src/s_tir/transform/lower_async_dma.cc
@@ -28,25 +28,25 @@
#include <tvm/runtime/logging.h>
#include <tvm/s_tir/analysis.h>
#include <tvm/s_tir/stmt.h>
+#include <tvm/s_tir/stmt_functor.h>
#include <tvm/s_tir/transform.h>
#include <tvm/tirx/buffer.h>
#include <tvm/tirx/stmt.h>
-#include <tvm/tirx/stmt_functor.h>
#include <optional>
#include <set>
-#include "../../tirx/ir/ir_mutator_with_analyzer.h"
+#include "../../s_tir/ir/ir_mutator_with_analyzer.h"
#include "../../tirx/transform/ir_utils.h"
namespace tvm {
namespace s_tir {
using namespace tvm::tirx;
-class AsyncDMALowerer : public tirx::IRMutatorWithAnalyzer {
+class AsyncDMALowerer : public s_tir::IRMutatorWithAnalyzer {
public:
- using tirx::IRMutatorWithAnalyzer::Mutate;
- using tirx::IRMutatorWithAnalyzer::Mutate_;
+ using s_tir::IRMutatorWithAnalyzer::Mutate;
+ using s_tir::IRMutatorWithAnalyzer::Mutate_;
explicit AsyncDMALowerer(bool dma_bypass_cache, const arith::Analyzer& analyzer)
: IRMutatorWithAnalyzer(analyzer), dma_bypass_cache_(dma_bypass_cache) {}
@@ -55,7 +55,7 @@ class AsyncDMALowerer : public tirx::IRMutatorWithAnalyzer {
UnchangedOr<Stmt> Mutate_(const ForNode* loop, InplaceMode inplace_mode) final {
// if for loop is not within async_commit_queue_scope
if (!async_queue_id_.has_value()) {
- return tirx::IRMutatorWithAnalyzer::Mutate_(loop, inplace_mode);
+ return s_tir::IRMutatorWithAnalyzer::Mutate_(loop, inplace_mode);
}
// if for loop is not a memcpy of a contiguous region, it might be a cuda cp.async behavior
@@ -63,7 +63,7 @@ class AsyncDMALowerer : public tirx::IRMutatorWithAnalyzer {
s_tir::IdentifyMemCpy(ffi::GetRef<For>(loop), ffi::GetRef<arith::Analyzer>(analyzer_));
if (!mem_copy.has_value() || mem_copy->dest->region.size() != 1 ||
mem_copy->source->region.size() != 1) {
- return tirx::IRMutatorWithAnalyzer::Mutate_(loop, inplace_mode);
+ return s_tir::IRMutatorWithAnalyzer::Mutate_(loop, inplace_mode);
}
// now that we are about to perform the `copy` transform
@@ -94,7 +94,7 @@ class AsyncDMALowerer : public tirx::IRMutatorWithAnalyzer {
UnchangedOr<Stmt> Mutate_(const AttrStmtNode* op, InplaceMode inplace_mode) final {
// populate analyzer knowledge of loop iterators
- auto previsit = tirx::IRMutatorWithAnalyzer::Mutate_(op, inplace_mode)
+ auto previsit = s_tir::IRMutatorWithAnalyzer::Mutate_(op, inplace_mode)
.ValueOrUnchanged(ffi::GetRef<Stmt>(op));
if (!op->unique()) inplace_mode = InplaceMode::kDisallow;
@@ -136,10 +136,10 @@ class AsyncDMALowerer : public tirx::IRMutatorWithAnalyzer {
// The nested attribute is skipped by this descent.
InplaceMode body_mode = async_wait->unique() ? inplace_mode : InplaceMode::kDisallow;
// concatenate the call with the body and return
- return SeqStmt({call_dma_wait,
- tirx::IRMutatorWithAnalyzer::Mutate(ffi::AnyView(async_wait->body), body_mode)
- .ValueOrUnchanged(async_wait->body)
- .as_or_throw<Stmt>()});
+ return SeqStmt({call_dma_wait, s_tir::IRMutatorWithAnalyzer::Mutate(
+ ffi::AnyView(async_wait->body), body_mode)
+ .ValueOrUnchanged(async_wait->body)
+ .as_or_throw<Stmt>()});
// Convert this, for example:
// attr [0] "async_commit_queue_scope" = 0;
@@ -161,7 +161,7 @@ class AsyncDMALowerer : public tirx::IRMutatorWithAnalyzer {
auto queue_id_node = op->value.as<IntImmNode>();
TVM_FFI_ICHECK(queue_id_node);
async_queue_id_ = queue_id_node->value.as<int>().value();
- auto result = tirx::IRMutatorWithAnalyzer::Mutate_(op, inplace_mode)
+ auto result = s_tir::IRMutatorWithAnalyzer::Mutate_(op, inplace_mode)
.ValueOrUnchanged(ffi::GetRef<Stmt>(op));
if (dmas_in_group_ > 1) {
auto call_dma_start_group =
@@ -178,7 +178,7 @@ class AsyncDMALowerer : public tirx::IRMutatorWithAnalyzer {
dmas_in_group_ = 0;
return result;
}
- return tirx::IRMutatorWithAnalyzer::Mutate_(op, inplace_mode);
+ return s_tir::IRMutatorWithAnalyzer::Mutate_(op, inplace_mode);
}
private:
diff --git a/src/s_tir/transform/lower_cross_thread_reduction.cc b/src/s_tir/transform/lower_cross_thread_reduction.cc
index 4926d9161f..51ff831942 100644
--- a/src/s_tir/transform/lower_cross_thread_reduction.cc
+++ b/src/s_tir/transform/lower_cross_thread_reduction.cc
@@ -25,11 +25,12 @@
#include <tvm/ffi/extra/structural_mutate.h>
#include <tvm/ffi/extra/structural_visit.h>
#include <tvm/ffi/reflection/registry.h>
+#include <tvm/s_tir/analysis.h>
#include <tvm/s_tir/stmt.h>
+#include <tvm/s_tir/stmt_functor.h>
#include <tvm/s_tir/transform.h>
#include <tvm/te/operation.h>
#include <tvm/tirx/analysis.h>
-#include <tvm/tirx/stmt_functor.h>
#include "../../runtime/thread_storage_scope.h"
#include "../../support/utils.h"
diff --git a/src/s_tir/transform/lower_init_block.cc b/src/s_tir/transform/lower_init_block.cc
index f230ba9ea1..a6441af39d 100644
--- a/src/s_tir/transform/lower_init_block.cc
+++ b/src/s_tir/transform/lower_init_block.cc
@@ -22,9 +22,10 @@
* \file lower_reduction.cc
*/
#include <tvm/ffi/reflection/registry.h>
+#include <tvm/s_tir/stmt.h>
+#include <tvm/s_tir/stmt_functor.h>
#include <tvm/s_tir/transform.h>
#include <tvm/tirx/op.h>
-#include <tvm/tirx/stmt_functor.h>
#include "../../tirx/transform/ir_utils.h"
diff --git a/src/s_tir/transform/lower_match_buffer.cc b/src/s_tir/transform/lower_match_buffer.cc
index eecfae6c36..c17dda686e 100644
--- a/src/s_tir/transform/lower_match_buffer.cc
+++ b/src/s_tir/transform/lower_match_buffer.cc
@@ -27,12 +27,14 @@
#include <tvm/ffi/extra/structural_mutate.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/runtime/logging.h>
+#include <tvm/s_tir/stmt.h>
+#include <tvm/s_tir/stmt_functor.h>
#include <tvm/s_tir/transform.h>
#include <tvm/tirx/function.h>
#include <tvm/tirx/op.h>
-#include <tvm/tirx/stmt_functor.h>
#include "../../tirx/transform/ir_utils.h"
+#include "../transform/ir_utils.h"
namespace tvm {
namespace s_tir {
diff --git a/src/s_tir/transform/lower_opaque_block.cc b/src/s_tir/transform/lower_opaque_block.cc
index e069856a32..1ac40d1832 100644
--- a/src/s_tir/transform/lower_opaque_block.cc
+++ b/src/s_tir/transform/lower_opaque_block.cc
@@ -24,10 +24,10 @@
#include <tvm/ffi/cast.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/s_tir/stmt.h>
+#include <tvm/s_tir/stmt_functor.h>
#include <tvm/s_tir/transform.h>
-#include <tvm/tirx/stmt_functor.h>
-#include "../../tirx/transform/ir_utils.h"
+#include "ir_utils.h"
namespace tvm {
namespace s_tir {
diff --git a/src/s_tir/transform/lower_thread_allreduce.cc b/src/s_tir/transform/lower_thread_allreduce.cc
index 4c0b28ff5f..6f4691e26e 100644
--- a/src/s_tir/transform/lower_thread_allreduce.cc
+++ b/src/s_tir/transform/lower_thread_allreduce.cc
@@ -27,11 +27,11 @@
#include <tvm/ir/prim/builtin.h>
#include <tvm/ir/prim/expr.h>
#include <tvm/s_tir/stmt.h>
+#include <tvm/s_tir/stmt_functor.h>
#include <tvm/s_tir/transform.h>
#include <tvm/target/target.h>
#include <tvm/te/operation.h>
#include <tvm/tirx/builtin.h>
-#include <tvm/tirx/stmt_functor.h>
#include <unordered_set>
diff --git a/src/s_tir/transform/lower_vtcm_alloc.cc b/src/s_tir/transform/lower_vtcm_alloc.cc
index 526819ec7d..54f4058d84 100644
--- a/src/s_tir/transform/lower_vtcm_alloc.cc
+++ b/src/s_tir/transform/lower_vtcm_alloc.cc
@@ -19,10 +19,10 @@
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/prim/builtin.h>
+#include <tvm/s_tir/stmt_functor.h>
#include <tvm/s_tir/transform.h>
#include <tvm/tirx/builtin.h>
#include <tvm/tirx/stmt.h>
-#include <tvm/tirx/stmt_functor.h>
namespace tvm {
namespace s_tir {
diff --git a/src/s_tir/transform/manifest_shared_memory_local_stage.cc b/src/s_tir/transform/manifest_shared_memory_local_stage.cc
index f7484ce295..67ad132466 100644
--- a/src/s_tir/transform/manifest_shared_memory_local_stage.cc
+++ b/src/s_tir/transform/manifest_shared_memory_local_stage.cc
@@ -32,9 +32,9 @@
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/prim/expr.h>
#include <tvm/s_tir/stmt.h>
+#include <tvm/s_tir/stmt_functor.h>
#include <tvm/s_tir/transform.h>
#include <tvm/tirx/op.h>
-#include <tvm/tirx/stmt_functor.h>
#include <unordered_set>
diff --git a/src/s_tir/transform/memhammer_lower_auto_copy.cc b/src/s_tir/transform/memhammer_lower_auto_copy.cc
index fea18924dc..f1a8a052c2 100644
--- a/src/s_tir/transform/memhammer_lower_auto_copy.cc
+++ b/src/s_tir/transform/memhammer_lower_auto_copy.cc
@@ -25,10 +25,10 @@
#include <tvm/ir/op.h>
#include <tvm/ir/prim/expr.h>
#include <tvm/s_tir/stmt.h>
+#include <tvm/s_tir/stmt_functor.h>
#include <tvm/s_tir/transform.h>
#include <tvm/target/target.h>
#include <tvm/tirx/op.h>
-#include <tvm/tirx/stmt_functor.h>
#include <array>
#include <stack>
diff --git a/src/s_tir/transform/memhammer_rewrite_rule.h b/src/s_tir/transform/memhammer_rewrite_rule.h
index 83b96e49e1..4c0a136f37 100644
--- a/src/s_tir/transform/memhammer_rewrite_rule.h
+++ b/src/s_tir/transform/memhammer_rewrite_rule.h
@@ -22,10 +22,10 @@
#include <tvm/arith/iter_affine_map.h>
#include <tvm/ffi/function.h>
#include <tvm/ir/prim/expr.h>
+#include <tvm/s_tir/stmt_functor.h>
#include <tvm/s_tir/transform.h>
#include <tvm/target/target.h>
#include <tvm/tirx/op.h>
-#include <tvm/tirx/stmt_functor.h>
#include <utility>
diff --git a/src/s_tir/transform/memhammer_tensorcore_rewrite.cc b/src/s_tir/transform/memhammer_tensorcore_rewrite.cc
index 7363572751..e21a086e34 100644
--- a/src/s_tir/transform/memhammer_tensorcore_rewrite.cc
+++ b/src/s_tir/transform/memhammer_tensorcore_rewrite.cc
@@ -21,6 +21,7 @@
#include <tvm/ffi/extra/structural_mutate.h>
#include <tvm/ffi/extra/structural_visit.h>
#include <tvm/ir/op.h>
+#include <tvm/s_tir/stmt.h>
#include "./memhammer_rewrite_rule.h"
diff --git a/src/s_tir/transform/merge_shared_memory_allocations.cc b/src/s_tir/transform/merge_shared_memory_allocations.cc
index d1d72eb368..321db8f39d 100644
--- a/src/s_tir/transform/merge_shared_memory_allocations.cc
+++ b/src/s_tir/transform/merge_shared_memory_allocations.cc
@@ -30,9 +30,9 @@
#include <tvm/ir/prim/expr.h>
#include <tvm/runtime/logging.h>
#include <tvm/s_tir/stmt.h>
+#include <tvm/s_tir/stmt_functor.h>
#include <tvm/s_tir/transform.h>
#include <tvm/tirx/op.h>
-#include <tvm/tirx/stmt_functor.h>
#include <list>
#include <map>
diff --git a/src/s_tir/transform/plan_update_buffer_allocation_location.cc b/src/s_tir/transform/plan_update_buffer_allocation_location.cc
index f49ba03a24..b7412857b4 100644
--- a/src/s_tir/transform/plan_update_buffer_allocation_location.cc
+++ b/src/s_tir/transform/plan_update_buffer_allocation_location.cc
@@ -23,9 +23,11 @@
*/
#include <tvm/ffi/reflection/registry.h>
+#include <tvm/s_tir/analysis.h>
+#include <tvm/s_tir/stmt.h>
+#include <tvm/s_tir/stmt_functor.h>
#include <tvm/s_tir/transform.h>
#include <tvm/tirx/analysis.h>
-#include <tvm/tirx/stmt_functor.h>
#include <tvm/tirx/var.h>
#include "../../tirx/transform/ir_utils.h"
diff --git a/src/s_tir/transform/profile_instrumentation.cc b/src/s_tir/transform/profile_instrumentation.cc
index 448b883f6f..458f0e06de 100644
--- a/src/s_tir/transform/profile_instrumentation.cc
+++ b/src/s_tir/transform/profile_instrumentation.cc
@@ -27,10 +27,10 @@
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/prim/builtin.h>
#include <tvm/ir/prim/expr.h>
+#include <tvm/s_tir/stmt_functor.h>
#include <tvm/s_tir/transform.h>
#include <tvm/tirx/builtin.h>
#include <tvm/tirx/stmt.h>
-#include <tvm/tirx/stmt_functor.h>
namespace tvm {
namespace s_tir {
diff --git a/src/s_tir/transform/remove_store_undef.cc b/src/s_tir/transform/remove_store_undef.cc
index 9fa243cd37..aa70ab33fe 100644
--- a/src/s_tir/transform/remove_store_undef.cc
+++ b/src/s_tir/transform/remove_store_undef.cc
@@ -24,12 +24,13 @@
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/prim/builtin.h>
+#include <tvm/s_tir/analysis.h>
+#include <tvm/s_tir/stmt_functor.h>
#include <tvm/s_tir/transform.h>
#include <tvm/tirx/analysis.h>
#include <tvm/tirx/builtin.h>
#include <tvm/tirx/op.h>
#include <tvm/tirx/stmt.h>
-#include <tvm/tirx/stmt_functor.h>
namespace tvm {
namespace s_tir {
diff --git a/src/s_tir/transform/remove_weight_layout_rewrite_block.cc b/src/s_tir/transform/remove_weight_layout_rewrite_block.cc
index 1632099a0d..d377bfa6c4 100644
--- a/src/s_tir/transform/remove_weight_layout_rewrite_block.cc
+++ b/src/s_tir/transform/remove_weight_layout_rewrite_block.cc
@@ -24,10 +24,10 @@
#include <tvm/ffi/reflection/registry.h>
#include <tvm/s_tir/stmt.h>
+#include <tvm/s_tir/stmt_functor.h>
#include <tvm/s_tir/transform.h>
#include <tvm/tirx/index_map.h>
#include <tvm/tirx/op.h>
-#include <tvm/tirx/stmt_functor.h>
#include <unordered_set>
diff --git a/src/s_tir/transform/renew_defs.cc b/src/s_tir/transform/renew_defs.cc
index 9740975e11..4776777ca3 100644
--- a/src/s_tir/transform/renew_defs.cc
+++ b/src/s_tir/transform/renew_defs.cc
@@ -25,8 +25,9 @@
#include <tvm/ffi/cast.h>
#include <tvm/ffi/extra/structural_visit.h>
#include <tvm/ffi/reflection/registry.h>
+#include <tvm/s_tir/stmt.h>
+#include <tvm/s_tir/stmt_functor.h>
#include <tvm/s_tir/transform.h>
-#include <tvm/tirx/stmt_functor.h>
namespace tvm {
namespace s_tir {
diff --git a/src/s_tir/transform/renormalize_split_pattern.cc b/src/s_tir/transform/renormalize_split_pattern.cc
index e290d3e2bd..8b5f6cff4b 100644
--- a/src/s_tir/transform/renormalize_split_pattern.cc
+++ b/src/s_tir/transform/renormalize_split_pattern.cc
@@ -23,14 +23,15 @@
*/
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
+#include <tvm/s_tir/analysis.h>
+#include <tvm/s_tir/stmt_functor.h>
#include <tvm/s_tir/transform.h>
#include <tvm/tirx/analysis.h>
#include <tvm/tirx/op.h>
#include <tvm/tirx/stmt.h>
-#include <tvm/tirx/stmt_functor.h>
#include "../../arith/pattern_match.h"
-#include "../../tirx/ir/ir_mutator_with_analyzer.h"
+#include "../../s_tir/ir/ir_mutator_with_analyzer.h"
namespace tvm {
namespace s_tir {
diff --git a/src/s_tir/transform/rewrite_unsafe_select.cc b/src/s_tir/transform/rewrite_unsafe_select.cc
index 7922185df0..66ff4bab8b 100644
--- a/src/s_tir/transform/rewrite_unsafe_select.cc
+++ b/src/s_tir/transform/rewrite_unsafe_select.cc
@@ -25,10 +25,10 @@
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/prim/builtin.h>
#include <tvm/ir/prim/expr.h>
+#include <tvm/s_tir/stmt_functor.h>
#include <tvm/s_tir/transform.h>
#include <tvm/tirx/builtin.h>
#include <tvm/tirx/op_attr_types.h>
-#include <tvm/tirx/stmt_functor.h>
namespace tvm {
namespace s_tir {
diff --git a/src/s_tir/transform/stmt_simplify.cc b/src/s_tir/transform/stmt_simplify.cc
new file mode 100644
index 0000000000..36d8cf4328
--- /dev/null
+++ b/src/s_tir/transform/stmt_simplify.cc
@@ -0,0 +1,94 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+#include "stmt_simplify.h"
+
+#include <tvm/ffi/reflection/registry.h>
+#include <tvm/s_tir/stmt_functor.h>
+#include <tvm/s_tir/transform.h>
+
+#include "../../tirx/transform/stmt_simplify.h"
+
+namespace tvm {
+namespace s_tir {
+using namespace tirx;
+
+// Reuse ordinary TIRX simplification, adding scoped constraints for S-TIR blocks.
+class StmtSimplifier final : public tirx::StmtSimplifier {
+ public:
+ using Parent = tirx::StmtSimplifier;
+ StmtSimplifier(const arith::Analyzer& analyzer, tirx::StmtSimplifyConfig config)
+ : Parent(GlobalVTable(), analyzer, config) {}
+ using Parent::Mutate_;
+ using Parent::Run;
+
+ public:
+ UnchangedOr<Stmt> Mutate_(const SBlockNode* op, InplaceMode inplace_mode) {
+ // This small binding step stays local: the shared simplifier is TIRX-only,
+ // while the analyzer state is protected by its owning base class.
+ return constraint_scope_.WithNewScope([&]() -> UnchangedOr<Stmt> {
+ for (const auto& iter_var : op->iter_vars) {
+ analyzer_->Bind(iter_var->var, iter_var->dom);
+ iter_vars_.Set(iter_var->var, iter_var->dom);
+ }
+ return s_tir::StmtExprMutator::MutateBlock(this, op, inplace_mode);
+ });
+ }
+ UnchangedOr<Stmt> Mutate_(const SBlockRealizeNode* op, InplaceMode inplace_mode) {
+ return s_tir::StmtExprMutator::MutateBlockRealize(this, op, inplace_mode);
+ }
+
+ protected:
+ static void InitVTable(VTable* vtable) {
+ Parent::InitVTable(vtable);
+ SetDispatch<StmtSimplifier, SBlockNode>(vtable);
+ SetDispatch<StmtSimplifier, SBlockRealizeNode>(vtable);
+ }
+ static const VTable* GlobalVTable() {
+ static const VTable table = [] {
+ VTable table;
+ InitVTable(&table);
+ table.Finalize();
+ return table;
+ }();
+ return &table;
+ }
+};
+
+PrimFunc StmtSimplify(PrimFunc func, const arith::Analyzer& analyzer) {
+ auto config = tvm::transform::PassConfigWithDefaults<tirx::StmtSimplifyConfig>();
+ return ffi::make_object<StmtSimplifier>(analyzer, config)->Run(std::move(func));
+}
+
+namespace transform {
+Pass StmtSimplify() {
+ auto pass_func = [](PrimFunc func, IRModule, tvm::transform::PassContext ctx) {
+ arith::Analyzer analyzer;
+ auto config = ctx->GetConfig<tirx::StmtSimplifyConfig>("tirx.StmtSimplify")
+ .value_or(tvm::transform::PassConfigWithDefaults<tirx::StmtSimplifyConfig>());
+ return ffi::make_object<s_tir::StmtSimplifier>(analyzer, config)->Run(std::move(func));
+ };
+ return tirx::transform::CreatePrimFuncPass(pass_func, 0, "s_tir.StmtSimplify", {});
+}
+TVM_FFI_STATIC_INIT_BLOCK() {
+ ffi::reflection::GlobalDef().def("s_tir.transform.StmtSimplify", StmtSimplify);
+}
+} // namespace transform
+} // namespace s_tir
+} // namespace tvm
diff --git a/src/tirx/transform/stmt_simplify.h b/src/s_tir/transform/stmt_simplify.h
similarity index 66%
copy from src/tirx/transform/stmt_simplify.h
copy to src/s_tir/transform/stmt_simplify.h
index 224df0ed8b..8dcff5bcc6 100644
--- a/src/tirx/transform/stmt_simplify.h
+++ b/src/s_tir/transform/stmt_simplify.h
@@ -17,25 +17,13 @@
* under the License.
*/
-/*!
- * \file stmt_simplify.h
- * \brief Statement-level simplification of TIR PrimFuncs.
- */
-#ifndef TVM_TIR_TRANSFORM_STMT_SIMPLIFY_H_
-#define TVM_TIR_TRANSFORM_STMT_SIMPLIFY_H_
-
+#ifndef TVM_S_TIR_TRANSFORM_STMT_SIMPLIFY_H_
+#define TVM_S_TIR_TRANSFORM_STMT_SIMPLIFY_H_
#include <tvm/arith/analyzer.h>
#include <tvm/tirx/function.h>
-
namespace tvm {
-namespace tirx {
-
-/* \brief Simplify statements in the prim func
- *
- * Applies the same behavior as the tirx.transform.StmtSimplify pass.
- */
-PrimFunc StmtSimplify(PrimFunc func, const arith::Analyzer& analyzer);
-
-} // namespace tirx
+namespace s_tir {
+tirx::PrimFunc StmtSimplify(tirx::PrimFunc func, const arith::Analyzer& analyzer);
+} // namespace s_tir
} // namespace tvm
-#endif // TVM_TIR_TRANSFORM_STMT_SIMPLIFY_H_
+#endif // TVM_S_TIR_TRANSFORM_STMT_SIMPLIFY_H_
diff --git a/src/s_tir/transform/storage_access.h b/src/s_tir/transform/storage_access.h
index 291e0da804..83b56d3520 100644
--- a/src/s_tir/transform/storage_access.h
+++ b/src/s_tir/transform/storage_access.h
@@ -27,7 +27,7 @@
#include <tvm/arith/int_set.h>
#include <tvm/ir/attrs.h>
#include <tvm/ir/prim/expr.h>
-#include <tvm/tirx/stmt_functor.h>
+#include <tvm/s_tir/stmt_functor.h>
#include <unordered_map>
#include <vector>
diff --git a/src/s_tir/transform/tensorcore_infer_fragment.cc b/src/s_tir/transform/tensorcore_infer_fragment.cc
index 8944dc2cc1..bf7d5414ce 100644
--- a/src/s_tir/transform/tensorcore_infer_fragment.cc
+++ b/src/s_tir/transform/tensorcore_infer_fragment.cc
@@ -27,8 +27,8 @@
#include <tvm/ir/op.h>
#include <tvm/ir/prim/expr.h>
#include <tvm/s_tir/stmt.h>
+#include <tvm/s_tir/stmt_functor.h>
#include <tvm/s_tir/transform.h>
-#include <tvm/tirx/stmt_functor.h>
#include <unordered_map>
#include <unordered_set>
@@ -53,11 +53,11 @@ const VarNode* GetBufferVarFromData(const Expr& data) {
}
// Get fragment information from tensor intrinsics
-class FragmentGetter : public StmtExprVisitor {
+class FragmentGetter : public s_tir::StmtExprVisitor {
public:
- using StmtExprVisitor::Visit_;
+ using s_tir::StmtExprVisitor::Visit_;
ffi::Optional<VisitInterrupt> Visit_(const CallNode* op) final {
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(op));
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(s_tir::StmtExprVisitor::Visit_(op));
static const Op& tvm_load_matrix_sync_op = Op::Get("tirx.tvm_load_matrix_sync");
static const Op& tvm_store_matrix_sync_op = Op::Get("tirx.tvm_store_matrix_sync");
@@ -145,13 +145,13 @@ std::unordered_map<const VarNode*, FragmentInfo> GetTensorCoreFragmentInfo(const
namespace s_tir {
// Check shape of fragment making sure it is a valid shape for tvm_mma_sync
-class FragmentChecker : public StmtExprVisitor {
+class FragmentChecker : public s_tir::StmtExprVisitor {
public:
- using StmtExprVisitor::Visit_;
+ using s_tir::StmtExprVisitor::Visit_;
explicit FragmentChecker(const FragmentGetter& getter) : fragment_getter(getter) {}
ffi::Optional<VisitInterrupt> Visit_(const CallNode* op) final {
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(op));
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(s_tir::StmtExprVisitor::Visit_(op));
// Check shape when calling tvm_mma_sync
static const Op& tvm_mma_sync_op = Op::Get("tirx.tvm_mma_sync");
static const Op& tvm_bmma_sync_op = Op::Get("tirx.tvm_bmma_sync");
@@ -194,19 +194,20 @@ class FragmentChecker : public StmtExprVisitor {
};
// Store the metadata into attributes
-class InferFragmenter : public StmtExprMutator {
+class InferFragmenter : public s_tir::StmtExprMutator {
public:
- using StmtExprMutator::Mutate;
- using StmtExprMutator::Mutate_;
+ using s_tir::StmtExprMutator::Mutate;
+ using s_tir::StmtExprMutator::Mutate_;
UnchangedOr<ffi::Any> Mutate(ffi::AnyView value, InplaceMode inplace_mode) override {
if (value.as<ExprNode>()) return ffi::Unchanged();
- return StmtExprMutator::Mutate(value, inplace_mode);
+ return s_tir::StmtExprMutator::Mutate(value, inplace_mode);
}
explicit InferFragmenter(const FragmentGetter& getter) : fragment_getter(getter) {}
UnchangedOr<Stmt> Mutate_(const AllocBufferNode* op, InplaceMode inplace_mode) final {
- Stmt stmt = StmtExprMutator::Mutate_(op, inplace_mode).ValueOrUnchanged(ffi::GetRef<Stmt>(op));
+ Stmt stmt =
+ s_tir::StmtExprMutator::Mutate_(op, inplace_mode).ValueOrUnchanged(ffi::GetRef<Stmt>(op));
const VarNode* buffer = op->buffer.get();
if (fragment_getter.fragments.count(buffer)) {
FragmentInfo info = fragment_getter.fragments.at(buffer);
diff --git a/src/s_tir/transform/thread_storage_sync.cc b/src/s_tir/transform/thread_storage_sync.cc
index 07c6576bf7..08e644f5c1 100644
--- a/src/s_tir/transform/thread_storage_sync.cc
+++ b/src/s_tir/transform/thread_storage_sync.cc
@@ -25,12 +25,13 @@
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/prim/builtin.h>
#include <tvm/ir/prim/expr.h>
+#include <tvm/s_tir/analysis.h>
#include <tvm/s_tir/stmt.h>
+#include <tvm/s_tir/stmt_functor.h>
#include <tvm/s_tir/transform.h>
#include <tvm/tirx/analysis.h>
#include <tvm/tirx/builtin.h>
#include <tvm/tirx/op.h>
-#include <tvm/tirx/stmt_functor.h>
#include <unordered_set>
diff --git a/src/s_tir/transform/transform_mma_buffer_layout.cc b/src/s_tir/transform/transform_mma_buffer_layout.cc
index 4eb31678f3..3784db34e9 100644
--- a/src/s_tir/transform/transform_mma_buffer_layout.cc
+++ b/src/s_tir/transform/transform_mma_buffer_layout.cc
@@ -20,10 +20,12 @@
#include <tvm/arith/analyzer.h>
#include <tvm/ffi/cast.h>
#include <tvm/ffi/reflection/registry.h>
+#include <tvm/s_tir/analysis.h>
+#include <tvm/s_tir/stmt.h>
+#include <tvm/s_tir/stmt_functor.h>
#include <tvm/s_tir/transform.h>
#include <tvm/tirx/analysis.h>
#include <tvm/tirx/index_map.h>
-#include <tvm/tirx/stmt_functor.h>
#include "../../tirx/transform/ir_utils.h"
diff --git a/src/s_tir/transform/unify_thread_binding.cc b/src/s_tir/transform/unify_thread_binding.cc
index ebdbefc399..658b9fe77a 100644
--- a/src/s_tir/transform/unify_thread_binding.cc
+++ b/src/s_tir/transform/unify_thread_binding.cc
@@ -24,10 +24,11 @@
#include <tvm/arith/analyzer.h>
#include <tvm/ffi/cast.h>
#include <tvm/ffi/reflection/registry.h>
+#include <tvm/s_tir/analysis.h>
#include <tvm/s_tir/stmt.h>
+#include <tvm/s_tir/stmt_functor.h>
#include <tvm/s_tir/transform.h>
#include <tvm/tirx/analysis.h>
-#include <tvm/tirx/stmt_functor.h>
#include "../../support/utils.h"
#include "../../tirx/transform/ir_utils.h"
diff --git a/src/s_tir/transform/using_assume_to_reduce_branches.cc b/src/s_tir/transform/using_assume_to_reduce_branches.cc
index 8d8e9e5d54..95cc2fbfff 100644
--- a/src/s_tir/transform/using_assume_to_reduce_branches.cc
+++ b/src/s_tir/transform/using_assume_to_reduce_branches.cc
@@ -40,16 +40,16 @@
#include <tvm/ir/prim/builtin.h>
#include <tvm/relax/expr.h>
#include <tvm/relax/op_attr_types.h>
+#include <tvm/s_tir/stmt_functor.h>
#include <tvm/s_tir/transform.h>
#include <tvm/tirx/builtin.h>
#include <tvm/tirx/function.h>
#include <tvm/tirx/op.h>
-#include <tvm/tirx/stmt_functor.h>
#include <unordered_map>
#include "../../arith/constraint_extract.h"
-#include "../../tirx/ir/ir_mutator_with_analyzer.h"
+#include "../../s_tir/ir/ir_mutator_with_analyzer.h"
#include "tvm/ir/expr.h"
namespace tvm {
namespace s_tir {
diff --git a/src/te/operation/create_primfunc.cc b/src/te/operation/create_primfunc.cc
index 67532c035e..ca0b2045d6 100644
--- a/src/te/operation/create_primfunc.cc
+++ b/src/te/operation/create_primfunc.cc
@@ -27,11 +27,11 @@
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/unique_name_supply.h>
#include <tvm/s_tir/stmt.h>
+#include <tvm/s_tir/stmt_functor.h>
#include <tvm/te/operation.h>
#include <tvm/tirx/analysis.h>
#include <tvm/tirx/function.h>
#include <tvm/tirx/op.h>
-#include <tvm/tirx/stmt_functor.h>
#include <algorithm>
#include <set>
@@ -40,7 +40,7 @@
#include <utility>
#include <vector>
-#include "../../tirx/ir/data_type_rewriter.h"
+#include "../../s_tir/ir/data_type_rewriter.h"
#include "graph.h"
namespace tvm {
@@ -70,7 +70,7 @@ void VerifyNoOpaqueArtifacts(const PrimFunc& func) {
} // namespace
/*! \brief The helper mutator that transforms Tensor-callee Calls to BufferLoad. */
-class TensorLoadToBufferTransformer : public StmtExprMutator {
+class TensorLoadToBufferTransformer : public s_tir::StmtExprMutator {
public:
explicit TensorLoadToBufferTransformer(
const std::unordered_map<te::Tensor, BufferVar>& tensor2buffers)
@@ -80,7 +80,7 @@ class TensorLoadToBufferTransformer : public StmtExprMutator {
const auto* reduce =
op->IsInstance<te::ReduceNode>() ? static_cast<const te::ReduceNode*>(op) : nullptr;
if (reduce == nullptr) {
- return StmtExprMutator::Mutate_(op, inplace_mode);
+ return s_tir::StmtExprMutator::Mutate_(op, inplace_mode);
}
auto axis = reduce->axis.Map([this](const IterVar& iter_var) {
@@ -117,7 +117,7 @@ class TensorLoadToBufferTransformer : public StmtExprMutator {
}
UnchangedOr<Expr> Mutate_(const CallNode* op, InplaceMode inplace_mode) final {
- Call call = StmtExprMutator::Mutate_(op, inplace_mode)
+ Call call = s_tir::StmtExprMutator::Mutate_(op, inplace_mode)
.ValueOrUnchanged(ffi::GetRef<Expr>(op))
.as_or_throw<Call>();
if (!te::IsTensorLoad(call)) {
@@ -136,7 +136,7 @@ class TensorLoadToBufferTransformer : public StmtExprMutator {
};
/*! \brief The helper mutator to rewrite buffer and buffer var accessed by block body */
-class BufferSubstituter : public StmtExprMutator {
+class BufferSubstituter : public s_tir::StmtExprMutator {
public:
explicit BufferSubstituter(const std::unordered_map<const VarNode*, Expr>& var_map,
const std::unordered_map<const VarNode*, BufferVar>& buffer_map) {
@@ -174,12 +174,12 @@ struct CreateFuncInfo {
}
};
-class LayoutFreePlaceholdersNormalizer : public StmtExprMutator {
+class LayoutFreePlaceholdersNormalizer : public s_tir::StmtExprMutator {
public:
- using StmtExprMutator::Mutate;
+ using s_tir::StmtExprMutator::Mutate;
UnchangedOr<ffi::Any> Mutate(ffi::AnyView value, InplaceMode inplace_mode) final {
if (value.as<ExprNode>()) return ffi::Unchanged();
- return StmtExprMutator::Mutate(value, inplace_mode);
+ return s_tir::StmtExprMutator::Mutate(value, inplace_mode);
}
PrimFunc Process(PrimFunc func) {
@@ -201,11 +201,11 @@ class LayoutFreePlaceholdersNormalizer : public StmtExprMutator {
return WithAttr(std::move(func), s_tir::attr::layout_free_buffers, indices);
}
- UnchangedOr<Stmt> Mutate_(const SBlockNode* _block, InplaceMode inplace_mode) final {
- SBlock block = StmtExprMutator::Mutate_(_block, inplace_mode)
- .ValueOrUnchanged(ffi::GetRef<Stmt>(_block))
- .as_or_throw<SBlock>();
- SBlockNode* n = block.CopyOnWrite();
+ UnchangedOr<Stmt> Mutate_(const s_tir::SBlockNode* _block, InplaceMode inplace_mode) final {
+ s_tir::SBlock block = s_tir::StmtExprMutator::Mutate_(_block, inplace_mode)
+ .ValueOrUnchanged(ffi::GetRef<Stmt>(_block))
+ .as_or_throw<s_tir::SBlock>();
+ s_tir::SBlockNode* n = block.CopyOnWrite();
if (auto opt_ann = n->annotations.Get(topi_attr)) {
ffi::Array<BufferVar> new_buffers;
for (BufferVar buffer : opt_ann.value().as_or_throw<ffi::Array<BufferVar>>()) {
@@ -628,18 +628,19 @@ Stmt GenerateStmtFromCompute(const te::ComputeOp& compute_op, CreateFuncInfo* in
}
Stmt body =
GenerateBodyStmt(leaf.store_indices, buffers, leaf.axes_remap, expr_body, info, analyzer);
- seq_stmt.push_back(SBlockRealize(/*iter_values=*/leaf.bindings,
- /*predicate=*/IntImm::Bool(true),
- /*block=*/
- SBlock(/*iter_vars=*/leaf.block_iters,
- /*reads=*/{},
- /*writes=*/{},
- /*name_hint=*/info->FreshName(compute_op->name),
- /*body=*/body,
- /*init=*/init,
- /*alloc_buffers=*/{},
- /*match_buffers=*/{},
- /*annotations=*/annotations)));
+ seq_stmt.push_back(
+ s_tir::SBlockRealize(/*iter_values=*/leaf.bindings,
+ /*predicate=*/IntImm::Bool(true),
+ /*block=*/
+ s_tir::SBlock(/*iter_vars=*/leaf.block_iters,
+ /*reads=*/{},
+ /*writes=*/{},
+ /*name_hint=*/info->FreshName(compute_op->name),
+ /*body=*/body,
+ /*init=*/init,
+ /*alloc_buffers=*/{},
+ /*match_buffers=*/{},
+ /*annotations=*/annotations)));
} else {
for (int i = 0; i < compute_op->num_outputs(); ++i) {
@@ -650,18 +651,19 @@ Stmt GenerateStmtFromCompute(const te::ComputeOp& compute_op, CreateFuncInfo* in
PrimExpr expr_body = compute_op->body[i];
Stmt body = GenerateBodyStmt(leaf.store_indices, {buffers[i]}, leaf.axes_remap, expr_body,
info, analyzer);
- seq_stmt.push_back(SBlockRealize(/*iter_values=*/leaf.bindings,
- /*predicate=*/IntImm::Bool(true),
- /*block=*/
- SBlock(/*iter_vars=*/leaf.block_iters,
- /*reads=*/{},
- /*writes=*/{},
- /*name_hint=*/info->FreshName(buffers[i].name()),
- /*body=*/body,
- /*init=*/std::nullopt,
- /*alloc_buffers=*/{},
- /*match_buffers=*/{},
- /*annotations=*/annotations)));
+ seq_stmt.push_back(
+ s_tir::SBlockRealize(/*iter_values=*/leaf.bindings,
+ /*predicate=*/IntImm::Bool(true),
+ /*block=*/
+ s_tir::SBlock(/*iter_vars=*/leaf.block_iters,
+ /*reads=*/{},
+ /*writes=*/{},
+ /*name_hint=*/info->FreshName(buffers[i].name()),
+ /*body=*/body,
+ /*init=*/std::nullopt,
+ /*alloc_buffers=*/{},
+ /*match_buffers=*/{},
+ /*annotations=*/annotations)));
}
}
Stmt body = SeqStmt::Flatten(seq_stmt);
@@ -679,18 +681,18 @@ Stmt GenerateStmtFromCompute(const te::ComputeOp& compute_op, CreateFuncInfo* in
}
// wrap nested block
- body = SBlockRealize(/*iter_values=*/cur.bindings,
- /*predicate=*/IntImm::Bool(true),
- /*block=*/
- SBlock(/*iter_vars=*/block_iters,
- /*reads=*/{},
- /*writes=*/{},
- /*name_hint=*/block_name,
- /*body=*/body,
- /*init=*/init,
- /*alloc_buffers=*/{},
- /*match_buffers=*/{},
- /*annotations=*/annotations));
+ body = s_tir::SBlockRealize(/*iter_values=*/cur.bindings,
+ /*predicate=*/IntImm::Bool(true),
+ /*block=*/
+ s_tir::SBlock(/*iter_vars=*/block_iters,
+ /*reads=*/{},
+ /*writes=*/{},
+ /*name_hint=*/block_name,
+ /*body=*/body,
+ /*init=*/init,
+ /*alloc_buffers=*/{},
+ /*match_buffers=*/{},
+ /*annotations=*/annotations));
}
for (size_t j = cur.loop_vars.size(); j > 0; --j) {
const auto& [loop_var, dom] = cur.loop_vars[j - 1];
@@ -739,7 +741,7 @@ Stmt GenerateStmtFromExternOp(const te::ExternOp& extern_op, CreateFuncInfo* inf
// be generated with the later application of "script.Complete" in
// GenerateAndCompletePrimFunc. Waiting until later also handles
// the case where there is only a single BlockNode, which then
- // becomes the root SBlock of the function, and should not have
+ // becomes the root s_tir::SBlock of the function, and should not have
// reads/writes filled in.
auto substituter = ffi::make_object<BufferSubstituter>(var_map, input_buffer_map);
@@ -751,18 +753,18 @@ Stmt GenerateStmtFromExternOp(const te::ExternOp& extern_op, CreateFuncInfo* inf
.ValueOrUnchanged(substituted_body);
// Step 4. Generate opaque block as body.
- return SBlockRealize(/*iter_values=*/{},
- /*predicate=*/IntImm::Bool(true),
- /*block=*/
- SBlock(/*iter_vars=*/{},
- /*reads=*/{},
- /*writes=*/{},
- /*name_hint=*/info->FreshName(extern_op->name),
- /*body=*/std::move(body),
- /*init=*/std::nullopt,
- /*alloc_buffers=*/{},
- /*match_buffers=*/{},
- /*annotations=*/extern_op->attrs));
+ return s_tir::SBlockRealize(/*iter_values=*/{},
+ /*predicate=*/IntImm::Bool(true),
+ /*block=*/
+ s_tir::SBlock(/*iter_vars=*/{},
+ /*reads=*/{},
+ /*writes=*/{},
+ /*name_hint=*/info->FreshName(extern_op->name),
+ /*body=*/std::move(body),
+ /*init=*/std::nullopt,
+ /*alloc_buffers=*/{},
+ /*match_buffers=*/{},
+ /*annotations=*/extern_op->attrs));
}
ffi::Array<te::Operation> CollectOrderedOps(const ffi::Array<te::Tensor>& arg_list) {
@@ -872,7 +874,7 @@ PrimFunc CreatePrimFunc(const ffi::Array<te::Tensor>& arg_list,
// Step 4. Create func and complete prim func.
auto func = GenerateAndCompletePrimFunc(arg_list, root_stmts, &info);
if (index_dtype_override.has_value()) {
- func = ffi::make_object<IndexDataTypeNormalizer>(index_dtype_override.value())
+ func = ffi::make_object<s_tir::IndexDataTypeNormalizer>(index_dtype_override.value())
->Rewrite(std::move(func));
}
auto result = ffi::make_object<LayoutFreePlaceholdersNormalizer>()->Process(std::move(func));
@@ -948,7 +950,7 @@ PrimFunc CreatePrimFunc(const ffi::Array<ffi::ObjectRef>& arg_list,
}
auto func = GenerateAndCompletePrimFunc(arg_list, root_stmts, &info);
if (index_dtype_override.has_value()) {
- func = ffi::make_object<IndexDataTypeNormalizer>(index_dtype_override.value())
+ func = ffi::make_object<s_tir::IndexDataTypeNormalizer>(index_dtype_override.value())
->Rewrite(std::move(func));
}
auto result = ffi::make_object<LayoutFreePlaceholdersNormalizer>()->Process(std::move(func));
diff --git a/src/tirx/analysis/verify_tirx_well_formed.cc b/src/tirx/analysis/verify_tirx_well_formed.cc
index 66adcebb0b..8407bdb671 100644
--- a/src/tirx/analysis/verify_tirx_well_formed.cc
+++ b/src/tirx/analysis/verify_tirx_well_formed.cc
@@ -50,14 +50,6 @@ class ExecScopeVerifier : public Verifier<ExecScopeVerifier> {
private:
using Verifier::Visit;
- void Dispatch_(const SBlockNode* op, ffi::reflection::AccessPath path) override {
- Verify(false) << "TIRxError: SBlock is not allowed in tirx=True mode at " << path;
- }
-
- void Dispatch_(const SBlockRealizeNode* op, ffi::reflection::AccessPath path) override {
- Verify(false) << "TIRxError: SBlockRealize is not allowed in tirx=True mode at " << path;
- }
-
void Dispatch_(const tirx::TilePrimitiveCallNode* op, ffi::reflection::AccessPath path) override {
static const auto& category_map = Op::GetAttrMap<tirx::TIRxOpCategory>("TIRxOpCategory");
Verify(category_map.get(op->op, ffi::String("")) == "tile_primitive")
@@ -131,14 +123,6 @@ class LayoutVerifier : public Verifier<LayoutVerifier> {
private:
using Verifier::Visit;
-
- void Dispatch_(const SBlockNode* op, ffi::reflection::AccessPath path) override {
- Verify(false) << "TIRxError: SBlock is not allowed in tirx=True mode at " << path;
- }
-
- void Dispatch_(const SBlockRealizeNode* op, ffi::reflection::AccessPath path) override {
- Verify(false) << "TIRxError: SBlockRealize is not allowed in tirx=True mode at " << path;
- }
};
class AsyncStructsVerifier : public Verifier<AsyncStructsVerifier> {
@@ -147,14 +131,6 @@ class AsyncStructsVerifier : public Verifier<AsyncStructsVerifier> {
private:
using Verifier::Visit;
-
- void Dispatch_(const SBlockNode* op, ffi::reflection::AccessPath path) override {
- Verify(false) << "TIRxError: SBlock is not allowed in tirx=True mode at " << path;
- }
-
- void Dispatch_(const SBlockRealizeNode* op, ffi::reflection::AccessPath path) override {
- Verify(false) << "TIRxError: SBlockRealize is not allowed in tirx=True mode at " << path;
- }
};
class DeviceFuncVerifier : public Verifier<DeviceFuncVerifier> {
@@ -163,14 +139,6 @@ class DeviceFuncVerifier : public Verifier<DeviceFuncVerifier> {
private:
using Verifier::Visit;
-
- void Dispatch_(const SBlockNode* op, ffi::reflection::AccessPath path) override {
- Verify(false) << "TIRxError: SBlock is not allowed in tirx=True mode at " << path;
- }
-
- void Dispatch_(const SBlockRealizeNode* op, ffi::reflection::AccessPath path) override {
- Verify(false) << "TIRxError: SBlockRealize is not allowed in tirx=True mode at " << path;
- }
};
bool VerifyTIRxWellFormed(const PrimFunc& func, bool assert_mode, bool device_func) {
@@ -197,11 +165,6 @@ bool VerifyTIRxWellFormed(const PrimFunc& func, bool assert_mode, bool device_fu
bool VerifyTIRxWellFormed(const IRModule& mod, bool assert_mode, bool device_func) {
for (const auto& [gvar, base_func] : mod->functions) {
if (auto prim_func = base_func.as<PrimFunc>()) {
- // s_tir=True PrimFuncs use s_tir semantics — defer to VerifyWellFormed.
- if (prim_func.value()->attrs->dict.count(tvm::attr::kSTir)) {
- if (!VerifyWellFormed(prim_func.value(), assert_mode)) return false;
- continue;
- }
bool res = VerifyTIRxWellFormed(prim_func.value(), assert_mode, device_func);
if (!res) {
return false;
diff --git a/src/tirx/analysis/verify_well_formed.cc b/src/tirx/analysis/verify_well_formed.cc
index 0123c463b6..adef48e1c1 100644
--- a/src/tirx/analysis/verify_well_formed.cc
+++ b/src/tirx/analysis/verify_well_formed.cc
@@ -22,6 +22,8 @@
* \brief Check if schedulable tirx is well-formed.
*/
+#include "verify_well_formed.h"
+
#include <tvm/ffi/cast.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
@@ -41,344 +43,6 @@ namespace tirx {
using AccessPath = ffi::reflection::AccessPath;
-/*! \brief Verify all Expr inside the block does not contain:
- * 1. loop vars outside the current block.
- * 2. block vars of parent blocks.
- */
-class BlockVarAccessVerifier : public StmtExprVisitor {
- public:
- static bool Verify(const PrimFunc& func, bool assert_mode) {
- auto verifier = ffi::make_object<BlockVarAccessVerifier>(assert_mode);
- verifier->Visit(func->body);
- return !verifier->has_error_;
- }
-
- explicit BlockVarAccessVerifier(bool assert_mode) : assert_mode_(assert_mode) {}
-
- private:
- ffi::Optional<VisitInterrupt> Visit(ffi::AnyView stmt) final {
- if (!has_error_) {
- return StmtExprVisitor::Visit(stmt);
- }
- return std::nullopt;
- }
-
- ffi::Optional<VisitInterrupt> Visit_(const VarNode* op) final {
- auto it = loop_vars_.find(op);
- if (it != loop_vars_.end() && it->second < block_stack_.size()) {
- has_error_ = true;
- if (assert_mode_) {
- if (it->second == 0) {
- TVM_FFI_THROW(InternalError)
- << "Well-formedness check failed: "
- << "Loop iterator var " << op->name << " is defined outside of any block, "
- << "but is used inside the non-opaque current block \""
- << block_stack_.back()->name_hint << "\".";
- } else {
- TVM_FFI_THROW(InternalError)
- << "Well-formedness check failed: "
- << "Loop iterator var " << op->name << " is defined in block \""
- << block_stack_[it->second - 1]->name_hint << "\", "
- << "but is used inside the non-opaque current block \""
- << block_stack_.back()->name_hint << "\".";
- }
- }
- }
- return std::nullopt;
- }
-
- ffi::Optional<VisitInterrupt> Visit_(const ForNode* op) final {
- TVM_FFI_ICHECK(loop_vars_.find(op->loop_var.get()) == loop_vars_.end());
- loop_vars_[op->loop_var.get()] = block_stack_.size();
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(op));
- loop_vars_.erase(op->loop_var.get());
- return std::nullopt;
- }
-
- ffi::Optional<VisitInterrupt> Visit_(const SBlockNode* op) final {
- // Do not check boundary if it's a opaque block.
- bool is_non_opaque = op->iter_vars.size();
- if (is_non_opaque) {
- block_stack_.push_back(op);
- }
-
- // Step 0. Skip block iter var's domain
-
- // Step 1. Visit read/write regions
- auto fvisit_buffer_region = [this](const TensorRegion& s) -> ffi::Optional<VisitInterrupt> {
- for (const auto& range : s->region) {
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->Visit(range->min));
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->Visit(range->extent));
- }
- return std::nullopt;
- };
- for (const auto& region : op->reads) {
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(fvisit_buffer_region(region));
- }
- for (const auto& region : op->writes) {
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(fvisit_buffer_region(region));
- }
-
- // Step 2. Visit match buffers
- for (const auto& match_buffer_region : op->match_buffers) {
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(fvisit_buffer_region(match_buffer_region->source));
- }
-
- // Step 3. Visit init and body
- if (op->init.has_value()) {
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->Visit(op->init.value()));
- }
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->Visit(op->body));
-
- if (is_non_opaque) {
- block_stack_.pop_back();
- }
- return std::nullopt;
- }
-
- private:
- /*! \brief The map from outside loop vars to its corresponding block level. */
- std::unordered_map<const VarNode*, size_t> loop_vars_;
- /*! \brief Whether it's in assert mode. */
- bool assert_mode_;
- /*! \brief Current nested block stack level. */
- std::vector<const SBlockNode*> block_stack_;
- /*! \brief Whether there is error. */
- bool has_error_{false};
-};
-
-class UndefinedVarVerifier : public Verifier<UndefinedVarVerifier> {
- public:
- // Until templated-this arrives in C++23, the CRTP can't inject a
- // constructor into the child class. Therefore, must explicitly add
- // the constructor.
- using Verifier::Verifier;
-
- private:
- using Verifier::Visit;
- void Visit(const PrimFunc& prim_func, AccessPath path) override {
- Verifier::Visit(prim_func, path);
- redefine_allowed_within_function_.clear();
- }
-
- void EnterDef(const IterVar& iter_var, AccessPath path) override {
- Verifier::EnterDef(iter_var, path);
- if (iter_var->iter_type == IterVarType::kThreadIndex) {
- redefine_allowed_within_function_.insert(iter_var->var);
- }
- }
-
- void EnterDef(const BufferVar& buffer, AccessPath path) override {
- Verifier::EnterDef(buffer, path);
- }
-
- void EnterDef(const Var& var, AccessPath path) override {
- bool redefine_is_allowed = redefine_allowed_within_function_.count(var);
- {
- auto it = currently_defined_.find(var);
- auto verify = Verify(it == currently_defined_.end() || redefine_is_allowed);
- verify << "ValueError: "
- << "TIR is ill-formed, "
- << "due to multiple nested definitions of variable " << var->name << ".";
- if (it != currently_defined_.end()) {
- verify << " It was first defined at " << it->second << ", and was re-defined at " << path;
- }
- }
-
- {
- auto it = previously_defined_.find(var);
- auto verify = Verify(it == previously_defined_.end() || redefine_is_allowed);
- verify << "ValueError: "
- << "TIR is ill-formed, "
- << "due to multiple definitions of variable " << var->name << ".";
- if (it != previously_defined_.end()) {
- verify << " It was first defined at " << it->second << ", and was later re-defined at "
- << path;
- }
- }
-
- currently_defined_.insert({var, path});
- }
-
- void ExitDef(const Var& var, AccessPath path) override {
- auto active_def = currently_defined_.find(var);
-
- currently_defined_.erase(active_def);
- previously_defined_.insert({var, path});
- }
-
- void Dispatch_(const VarNode* op, AccessPath path) override {
- auto var = ffi::GetRef<Var>(op);
-
- auto active_def = currently_defined_.find(var);
- auto verify = Verify(active_def != currently_defined_.end());
- verify << "ValueError: "
- << "Invalid use of undefined variable " << var->name << " at " << path << ".";
-
- // Check if there was a previous definition, and append the
- // location to the error message if there was. This is to aid in
- // debugging, by distinguishing between a variable that is
- // currently out-of-scope, and a variable that never had a
- // definition in the first place.
- if (auto prev_def = previously_defined_.find(var); prev_def != previously_defined_.end()) {
- verify << ". While this variable was previously defined at " << prev_def->second
- << ", this definition is no longer in-scope.";
- }
- }
-
- // Variables that are defined in the currently-visited scope.
- std::unordered_map<Var, AccessPath> currently_defined_;
-
- // Variables that were previously defined, and are now out of scope.
- std::unordered_map<Var, AccessPath> previously_defined_;
-
- // Special variables that are allowed to be re-defined, so long as
- // that re-definition occurs within the same PrimFunc. For example
- std::unordered_set<Var> redefine_allowed_within_function_;
-};
-
-/*! \brief Verify that buffers with a declaration are not used outside their declared scope.
- *
- * When a buffer is declared via one of the following sites:
- * - BufferType-annotated PrimFunc parameters
- * - DeclBuffer statement
- * - SBlock::alloc_buffers
- * - SBlock::match_buffers
- *
- * it must not appear in a BufferLoad, BufferStore, or TensorRegion outside that declaration's
- * scope.
- *
- * All buffers that appear in TensorLoad or BufferStore must have a prior declaration.
- */
-class UndefinedBufferVerifier : public Verifier<UndefinedBufferVerifier> {
- public:
- using Verifier::Verifier;
-
- private:
- using Verifier::Visit;
-
- void Visit(const PrimFunc& prim_func, AccessPath path) override {
- Verifier::Visit(prim_func, path);
- // Clear per-function state (buffers should not cross function boundaries).
- currently_defined_.clear();
- previously_defined_.clear();
- }
-
- void EnterDef(const BufferVar& buffer, AccessPath path) override {
- // Call the base class to visit buffer's internal vars (shape, strides, etc.)
- Verifier::EnterDef(buffer, path);
- currently_defined_.insert({buffer, path});
- }
-
- void ExitDef(const BufferVar& buffer, AccessPath path) override {
- auto active_def = currently_defined_.find(buffer);
- if (active_def != currently_defined_.end()) {
- currently_defined_.erase(active_def);
- }
- previously_defined_.insert({buffer, path});
- }
-
- void VisitBufferUse(const BufferVar& buffer, AccessPath path) override {
- bool is_declared = currently_defined_.count(buffer);
- bool was_declared = previously_defined_.count(buffer);
-
- if (was_declared && !is_declared) {
- // BufferVar was previously declared but is now out of scope — always an error.
- auto prev_def = previously_defined_.find(buffer);
- Verify(false) << "TIR is ill-formed: buffer " << buffer.name() << " is used at " << path
- << " but its declaration is no longer in-scope. "
- << "It was declared at " << prev_def->second << ".";
- } else if (!is_declared && !was_declared) {
- // BufferVar was never declared — error.
- Verify(false) << "TIR is ill-formed: buffer " << buffer.name() << " is used at " << path
- << " without a prior DeclBuffer or other declaration.";
- }
- // BufferVar fields are visited at definition site (EnterDef), not here.
- Verifier::VisitBufferUse(buffer, path);
- }
-
- // Buffers defined in the currently-visited scope.
- std::unordered_map<BufferVar, AccessPath, ffi::ObjectPtrHash, ffi::ObjectPtrEqual>
- currently_defined_;
- // Buffers that were previously defined and are now out of scope.
- std::unordered_map<BufferVar, AccessPath, ffi::ObjectPtrHash, ffi::ObjectPtrEqual>
- previously_defined_;
-};
-
-/*! \brief Verify the asserted type of each tirx buffer load. */
-class TensorLoadTypeVerifier : public Verifier<TensorLoadTypeVerifier> {
- public:
- using Verifier::Verifier;
-
- private:
- void Dispatch_(const TensorLoadNode* op, AccessPath path) override {
- auto buffer = op->source.as<BufferVar>();
- auto valid_source = Verify(buffer.has_value());
- valid_source << "TypeError: TIR TensorLoad source at " << path->Attr("source")
- << " must be a BufferVar.";
- if (!buffer.has_value()) {
- Visit(op->indices, path->Attr("indices"));
- return;
- }
-
- bool valid_indices = buffer.value()->shape.size() == op->indices.size();
- auto valid_rank = Verify(valid_indices);
- valid_rank << "ValueError: TIR TensorLoad at " << path << " indexes "
- << buffer.value()->shape.size() << "-dimensional buffer " << buffer.value().name()
- << " with " << op->indices.size() << " indices.";
- if (!valid_indices) {
- Visit(op->indices, path->Attr("indices"));
- return;
- }
-
- for (size_t i = 0; i + 1 < op->indices.size(); ++i) {
- bool is_scalar = op->indices[i].ty().IsScalar();
- auto valid_index = Verify(is_scalar);
- valid_index << "TypeError: TIR TensorLoad index " << i << " at "
- << path->Attr("indices")->ArrayItem(i)
- << " must be scalar because only the final index may be vector-valued.";
- valid_indices = valid_indices && is_scalar;
- }
- if (!valid_indices) {
- Visit(op->indices, path->Attr("indices"));
- return;
- }
-
- ffi::Optional<PrimType> index_ty = op->indices.empty()
- ? ffi::Optional<PrimType>(buffer.value()->dtype)
- : op->indices.back().ty().as<PrimType>();
- AccessPath final_index_path = op->indices.empty()
- ? path->Attr("indices")
- : path->Attr("indices")->ArrayItem(op->indices.size() - 1);
- auto valid_index_type = Verify(index_ty.has_value());
- valid_index_type << "TypeError: TIR TensorLoad final index at " << final_index_path
- << " must have a primitive type.";
- if (!index_ty.has_value()) {
- Visit(op->indices, path->Attr("indices"));
- return;
- }
-
- bool scalable_compatible = op->indices.empty() || !(buffer.value()->dtype.IsScalableVector() &&
- index_ty.value().IsScalableVector());
- auto valid_scalability = Verify(scalable_compatible);
- valid_scalability << "TypeError: TIR TensorLoad at " << path
- << " cannot combine a scalable buffer dtype with a scalable index.";
- if (!scalable_compatible) {
- Visit(op->indices, path->Attr("indices"));
- return;
- }
-
- TensorLoad expected = BufferLoad(buffer.value(), op->indices, op->span);
- ffi::Optional<PrimType> asserted_ty = op->ty.as<PrimType>();
- ffi::Optional<PrimType> expected_ty = expected->ty.as<PrimType>();
- auto valid_type = Verify(asserted_ty.has_value() && expected_ty.has_value() &&
- asserted_ty.value() == expected_ty.value());
- valid_type << "TypeError: TIR TensorLoad at " << path << " asserts result type " << op->ty
- << ", but its source and indices imply " << expected->ty << ".";
- TIRVisitorWithPath::Dispatch_(op, path);
- }
-};
-
/* \brief Verify unique tirx::Var for each environment thread
*
* Environment threads, such as CUDA's `threadIdx.x`, are defined in
@@ -421,37 +85,11 @@ class SingleEnvThreadVerifier : public Verifier<SingleEnvThreadVerifier> {
};
bool VerifyWellFormed(const PrimFunc& func, bool assert_mode) {
- if (!BlockVarAccessVerifier::Verify(func, assert_mode)) {
- return false;
- }
-
- if (!UndefinedVarVerifier::Verify(func, assert_mode)) return false;
-
- if (!UndefinedBufferVerifier::Verify(func, assert_mode)) return false;
-
- if (!TensorLoadTypeVerifier::Verify(func, assert_mode)) return false;
-
- // TODO(Siyuan): add more checks here.
- return true;
+ return VerifyWellFormedCommon<TIRVisitorWithPath>(func, assert_mode);
}
bool VerifyWellFormed(const IRModule& mod, bool assert_mode) {
- for (const auto& [gvar, base_func] : mod->functions) {
- if (auto prim_func = base_func.as<PrimFunc>()) {
- bool res = VerifyWellFormed(prim_func.value(), assert_mode);
- if (!res) {
- return false;
- }
- }
- }
-
- if (!UndefinedVarVerifier::Verify(mod, assert_mode)) return false;
-
- if (!UndefinedBufferVerifier::Verify(mod, assert_mode)) return false;
-
- if (!TensorLoadTypeVerifier::Verify(mod, assert_mode)) return false;
-
- return true;
+ return VerifyWellFormedCommon<TIRVisitorWithPath>(mod, assert_mode);
}
TVM_FFI_STATIC_INIT_BLOCK() {
diff --git a/src/tirx/analysis/verify_well_formed.cc b/src/tirx/analysis/verify_well_formed.h
similarity index 57%
copy from src/tirx/analysis/verify_well_formed.cc
copy to src/tirx/analysis/verify_well_formed.h
index 0123c463b6..16b57274ff 100644
--- a/src/tirx/analysis/verify_well_formed.cc
+++ b/src/tirx/analysis/verify_well_formed.h
@@ -17,142 +17,26 @@
* under the License.
*/
-/*!
- * \file tirx/analysis/verify_well_formed.cc
- * \brief Check if schedulable tirx is well-formed.
- */
-
-#include <tvm/ffi/cast.h>
-#include <tvm/ffi/function.h>
-#include <tvm/ffi/reflection/registry.h>
-#include <tvm/tirx/stmt.h>
-#include <tvm/tirx/stmt_functor.h>
+#ifndef TVM_TIRX_ANALYSIS_VERIFY_WELL_FORMED_H_
+#define TVM_TIRX_ANALYSIS_VERIFY_WELL_FORMED_H_
-#include <exception>
-#include <optional>
-#include <tuple>
-#include <variant>
+#include <tvm/tirx/analysis.h>
#include "../ir/tir_visitor_with_path.h"
-#include "tvm/ir/module.h"
-
namespace tvm {
namespace tirx {
-
using AccessPath = ffi::reflection::AccessPath;
-/*! \brief Verify all Expr inside the block does not contain:
- * 1. loop vars outside the current block.
- * 2. block vars of parent blocks.
- */
-class BlockVarAccessVerifier : public StmtExprVisitor {
- public:
- static bool Verify(const PrimFunc& func, bool assert_mode) {
- auto verifier = ffi::make_object<BlockVarAccessVerifier>(assert_mode);
- verifier->Visit(func->body);
- return !verifier->has_error_;
- }
-
- explicit BlockVarAccessVerifier(bool assert_mode) : assert_mode_(assert_mode) {}
+template <typename PathVisitor>
+class UndefinedVarVerifier : public Verifier<UndefinedVarVerifier<PathVisitor>, PathVisitor> {
+ using Verifier = tirx::Verifier<UndefinedVarVerifier<PathVisitor>, PathVisitor>;
- private:
- ffi::Optional<VisitInterrupt> Visit(ffi::AnyView stmt) final {
- if (!has_error_) {
- return StmtExprVisitor::Visit(stmt);
- }
- return std::nullopt;
- }
-
- ffi::Optional<VisitInterrupt> Visit_(const VarNode* op) final {
- auto it = loop_vars_.find(op);
- if (it != loop_vars_.end() && it->second < block_stack_.size()) {
- has_error_ = true;
- if (assert_mode_) {
- if (it->second == 0) {
- TVM_FFI_THROW(InternalError)
- << "Well-formedness check failed: "
- << "Loop iterator var " << op->name << " is defined outside of any block, "
- << "but is used inside the non-opaque current block \""
- << block_stack_.back()->name_hint << "\".";
- } else {
- TVM_FFI_THROW(InternalError)
- << "Well-formedness check failed: "
- << "Loop iterator var " << op->name << " is defined in block \""
- << block_stack_[it->second - 1]->name_hint << "\", "
- << "but is used inside the non-opaque current block \""
- << block_stack_.back()->name_hint << "\".";
- }
- }
- }
- return std::nullopt;
- }
-
- ffi::Optional<VisitInterrupt> Visit_(const ForNode* op) final {
- TVM_FFI_ICHECK(loop_vars_.find(op->loop_var.get()) == loop_vars_.end());
- loop_vars_[op->loop_var.get()] = block_stack_.size();
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(op));
- loop_vars_.erase(op->loop_var.get());
- return std::nullopt;
- }
-
- ffi::Optional<VisitInterrupt> Visit_(const SBlockNode* op) final {
- // Do not check boundary if it's a opaque block.
- bool is_non_opaque = op->iter_vars.size();
- if (is_non_opaque) {
- block_stack_.push_back(op);
- }
-
- // Step 0. Skip block iter var's domain
-
- // Step 1. Visit read/write regions
- auto fvisit_buffer_region = [this](const TensorRegion& s) -> ffi::Optional<VisitInterrupt> {
- for (const auto& range : s->region) {
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->Visit(range->min));
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->Visit(range->extent));
- }
- return std::nullopt;
- };
- for (const auto& region : op->reads) {
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(fvisit_buffer_region(region));
- }
- for (const auto& region : op->writes) {
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(fvisit_buffer_region(region));
- }
-
- // Step 2. Visit match buffers
- for (const auto& match_buffer_region : op->match_buffers) {
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(fvisit_buffer_region(match_buffer_region->source));
- }
-
- // Step 3. Visit init and body
- if (op->init.has_value()) {
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->Visit(op->init.value()));
- }
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->Visit(op->body));
-
- if (is_non_opaque) {
- block_stack_.pop_back();
- }
- return std::nullopt;
- }
-
- private:
- /*! \brief The map from outside loop vars to its corresponding block level. */
- std::unordered_map<const VarNode*, size_t> loop_vars_;
- /*! \brief Whether it's in assert mode. */
- bool assert_mode_;
- /*! \brief Current nested block stack level. */
- std::vector<const SBlockNode*> block_stack_;
- /*! \brief Whether there is error. */
- bool has_error_{false};
-};
-
-class UndefinedVarVerifier : public Verifier<UndefinedVarVerifier> {
public:
// Until templated-this arrives in C++23, the CRTP can't inject a
// constructor into the child class. Therefore, must explicitly add
// the constructor.
using Verifier::Verifier;
+ using Verifier::Verify;
private:
using Verifier::Visit;
@@ -242,17 +126,20 @@ class UndefinedVarVerifier : public Verifier<UndefinedVarVerifier> {
* When a buffer is declared via one of the following sites:
* - BufferType-annotated PrimFunc parameters
* - DeclBuffer statement
- * - SBlock::alloc_buffers
- * - SBlock::match_buffers
+ * - Dialect-specific definitions exposed by PathVisitor
*
- * it must not appear in a BufferLoad, BufferStore, or TensorRegion outside that declaration's
+ * it must not appear in a BufferLoad, BufferStore, or BufferRegion outside that declaration's
* scope.
*
* All buffers that appear in TensorLoad or BufferStore must have a prior declaration.
*/
-class UndefinedBufferVerifier : public Verifier<UndefinedBufferVerifier> {
+template <typename PathVisitor>
+class UndefinedBufferVerifier : public Verifier<UndefinedBufferVerifier<PathVisitor>, PathVisitor> {
+ using Verifier = tirx::Verifier<UndefinedBufferVerifier<PathVisitor>, PathVisitor>;
+
public:
using Verifier::Verifier;
+ using Verifier::Verify;
private:
using Verifier::Visit;
@@ -306,11 +193,16 @@ class UndefinedBufferVerifier : public Verifier<UndefinedBufferVerifier> {
};
/*! \brief Verify the asserted type of each tirx buffer load. */
-class TensorLoadTypeVerifier : public Verifier<TensorLoadTypeVerifier> {
+template <typename PathVisitor>
+class TensorLoadTypeVerifier : public Verifier<TensorLoadTypeVerifier<PathVisitor>, PathVisitor> {
+ using Verifier = tirx::Verifier<TensorLoadTypeVerifier<PathVisitor>, PathVisitor>;
+
public:
using Verifier::Verifier;
+ using Verifier::Verify;
private:
+ using Verifier::Visit;
void Dispatch_(const TensorLoadNode* op, AccessPath path) override {
auto buffer = op->source.as<BufferVar>();
auto valid_source = Verify(buffer.has_value());
@@ -375,100 +267,16 @@ class TensorLoadTypeVerifier : public Verifier<TensorLoadTypeVerifier> {
asserted_ty.value() == expected_ty.value());
valid_type << "TypeError: TIR TensorLoad at " << path << " asserts result type " << op->ty
<< ", but its source and indices imply " << expected->ty << ".";
- TIRVisitorWithPath::Dispatch_(op, path);
- }
-};
-
-/* \brief Verify unique tirx::Var for each environment thread
- *
- * Environment threads, such as CUDA's `threadIdx.x`, are defined in
- * TIR using an `AttrStmt` with the key `attr::thread_extent`. A
- * `PrimFunc` may contain multiple such attributes for the same
- * environment thread. However, all such attributes must use the same
- * `tirx::Var` for a given thread.
- */
-class SingleEnvThreadVerifier : public Verifier<SingleEnvThreadVerifier> {
- public:
- using Verifier::Verifier;
-
- private:
- void Visit(const PrimFunc& prim_func, AccessPath path) override {
- Verifier::Visit(prim_func, path);
- env_thread_vars_.clear();
- }
-
- void EnterDef(const IterVar& iter_var, AccessPath path) override {
- if (iter_var->iter_type == IterVarType::kThreadIndex) {
- if (auto it = env_thread_vars_.find(iter_var->thread_tag); it != env_thread_vars_.end()) {
- const auto& [prev_var, prev_path] = it->second;
- Verify(prev_var.same_as(iter_var->var))
- << "ValueError: "
- << "PrimFunc uses multiple distinct TIR variables "
- << " for the environment thread \"" << iter_var->thread_tag << "\". "
- << "While multiple tirx::AttrStmt may define the same environment thread, "
- << "all definitions within a single PrimFunc must share the same tirx::Var. "
- << "Binding of environment thread \"" << iter_var->thread_tag
- << "\" to the TIR variable " << iter_var->var->name << " at " << path
- << " conflicts with the previous binding to the TIR variable " << prev_var->name
- << " at " << path;
- } else {
- env_thread_vars_.insert({iter_var->thread_tag, {iter_var->var, path}});
- }
- }
+ PathVisitor::Dispatch_(op, path);
}
-
- std::unordered_map<ffi::String, std::tuple<Var, AccessPath>> env_thread_vars_;
};
-bool VerifyWellFormed(const PrimFunc& func, bool assert_mode) {
- if (!BlockVarAccessVerifier::Verify(func, assert_mode)) {
- return false;
- }
-
- if (!UndefinedVarVerifier::Verify(func, assert_mode)) return false;
-
- if (!UndefinedBufferVerifier::Verify(func, assert_mode)) return false;
-
- if (!TensorLoadTypeVerifier::Verify(func, assert_mode)) return false;
-
- // TODO(Siyuan): add more checks here.
- return true;
-}
-
-bool VerifyWellFormed(const IRModule& mod, bool assert_mode) {
- for (const auto& [gvar, base_func] : mod->functions) {
- if (auto prim_func = base_func.as<PrimFunc>()) {
- bool res = VerifyWellFormed(prim_func.value(), assert_mode);
- if (!res) {
- return false;
- }
- }
- }
-
- if (!UndefinedVarVerifier::Verify(mod, assert_mode)) return false;
-
- if (!UndefinedBufferVerifier::Verify(mod, assert_mode)) return false;
-
- if (!TensorLoadTypeVerifier::Verify(mod, assert_mode)) return false;
-
- return true;
-}
-
-TVM_FFI_STATIC_INIT_BLOCK() {
- namespace refl = tvm::ffi::reflection;
- refl::GlobalDef().def(
- "tirx.analysis.VerifyWellFormed", [](const ffi::ObjectRef& obj, bool assert_mode) {
- if (auto opt = obj.as<PrimFunc>()) {
- return VerifyWellFormed(opt.value(), assert_mode);
- } else if (auto opt = obj.as<IRModule>()) {
- return VerifyWellFormed(opt.value(), assert_mode);
- } else {
- TVM_FFI_THROW(InternalError)
- << "Expected VerifyWellFormed argument to be a PrimFunc or IRModule, but found "
- << obj->GetTypeKey();
- }
- });
+template <typename PathVisitor, typename NodeRef>
+bool VerifyWellFormedCommon(const NodeRef& node, bool assert_mode) {
+ return UndefinedVarVerifier<PathVisitor>::Verify(node, assert_mode) &&
+ UndefinedBufferVerifier<PathVisitor>::Verify(node, assert_mode) &&
+ TensorLoadTypeVerifier<PathVisitor>::Verify(node, assert_mode);
}
-
} // namespace tirx
} // namespace tvm
+#endif // TVM_TIRX_ANALYSIS_VERIFY_WELL_FORMED_H_
diff --git a/src/tirx/ir/data_type_rewriter.cc b/src/tirx/ir/data_type_rewriter.cc
index 2b7504513d..8f79792d18 100644
--- a/src/tirx/ir/data_type_rewriter.cc
+++ b/src/tirx/ir/data_type_rewriter.cc
@@ -27,14 +27,12 @@
#include <tvm/ffi/cast.h>
#include <tvm/ir/op.h>
#include <tvm/ir/prim/builtin.h>
-#include <tvm/s_tir/stmt.h>
#include <tvm/tirx/builtin.h>
#include <tvm/tirx/op.h>
#include <algorithm>
#include <functional>
#include <utility>
-#include <vector>
#include "tvm/ir/expr.h"
#include "tvm/ir/prim/expr.h"
@@ -44,7 +42,6 @@
namespace tvm {
namespace tirx {
using namespace tvm::prim;
-
UnchangedOr<Stmt> DataTypeLegalizer::Mutate_(const ForNode* op, InplaceMode inplace_mode) {
auto result = StmtExprMutator::Mutate_(op, inplace_mode);
if (!result.IsUnchanged()) {
@@ -74,51 +71,8 @@ UnchangedOr<Stmt> DataTypeLegalizer::Mutate_(const ForNode* op, InplaceMode inpl
return For(n);
}
-UnchangedOr<Stmt> DataTypeLegalizer::Mutate_(const SBlockRealizeNode* op,
- InplaceMode inplace_mode) {
- SBlockRealize realize = StmtExprMutator::Mutate_(op, inplace_mode)
- .ValueOrUnchanged(ffi::GetRef<Stmt>(op))
- .as_or_throw<SBlockRealize>();
- ffi::Array<PrimExpr> new_iter_values;
- bool changed = false;
- for (int i = 0; i < static_cast<int>(op->iter_values.size()); ++i) {
- PrimType dtype = realize->block->iter_vars[i]->var.ty();
- if (op->iter_values[i].ty() != dtype) {
- new_iter_values.push_back(prim::cast(dtype, realize->iter_values[i]));
- changed = true;
- } else {
- new_iter_values.push_back(realize->iter_values[i]);
- }
- }
- if (changed) {
- realize.CopyOnWrite()->iter_values = std::move(new_iter_values);
- }
- return realize;
-}
-
-UnchangedOr<Stmt> DataTypeLegalizer::Mutate_(const SBlockNode* op, InplaceMode inplace_mode) {
- SBlock new_block = StmtExprMutator::Mutate_(op, inplace_mode)
- .ValueOrUnchanged(ffi::GetRef<Stmt>(op))
- .as_or_throw<SBlock>();
- ffi::Array<IterVar> new_iter_vars = new_block->iter_vars.Map([](const IterVar& iter) {
- PrimType dtype = iter->var.ty();
- if (iter->dom->min.ty() != dtype || iter->dom->extent.ty() != dtype) {
- IterVar new_iter = iter;
- new_iter.CopyOnWrite()->dom =
- Range(prim::cast(dtype, iter->dom->min), prim::cast(dtype, iter->dom->extent));
- return new_iter;
- } else {
- return iter;
- }
- });
- if (!op->iter_vars.same_as(new_iter_vars)) {
- new_block.CopyOnWrite()->iter_vars = std::move(new_iter_vars);
- }
- return new_block;
-}
-
UnchangedOr<Stmt> DataTypeLegalizer::Mutate_(const AttrStmtNode* op, InplaceMode inplace_mode) {
- if (op->attr_key == attr::thread_extent || op->attr_key == s_tir::attr::virtual_thread) {
+ if (op->attr_key == attr::thread_extent || op->attr_key == tvm::tirx::attr::virtual_thread) {
Stmt s = StmtExprMutator::Mutate_(op, inplace_mode).ValueOrUnchanged(ffi::GetRef<Stmt>(op));
op = s.as<AttrStmtNode>();
TVM_FFI_ICHECK(op != nullptr) << "Expected type to be AttrStmtNode"
@@ -374,7 +328,7 @@ UnchangedOr<Expr> DataTypeLegalizer::Mutate_(const CallNode* op, InplaceMode inp
}
UnchangedOr<Stmt> IndexDataTypeRewriter::Mutate_(const AttrStmtNode* op, InplaceMode inplace_mode) {
- if (op->attr_key == attr::thread_extent || op->attr_key == s_tir::attr::virtual_thread) {
+ if (op->attr_key == attr::thread_extent || op->attr_key == tvm::tirx::attr::virtual_thread) {
bool is_enabled = is_enabled_;
is_enabled_ = true;
auto stmt = DataTypeLegalizer::Mutate_(op, inplace_mode);
@@ -392,165 +346,6 @@ UnchangedOr<ffi::Any> IndexDataTypeRewriter::Mutate(ffi::AnyView value, InplaceM
return result;
}
-UnchangedOr<Stmt> IndexDataTypeRewriter::Mutate_(const SBlockRealizeNode* op,
- InplaceMode inplace_mode) {
- bool is_condition = is_condition_;
- is_condition_ = true;
- auto new_predicate_result = Mutate(op->predicate, inplace_mode);
- bool new_predicate_unchanged = new_predicate_result.UnchangedOrSameAs(op->predicate);
- auto new_predicate = std::move(new_predicate_result).ValueOrUnchanged(op->predicate);
- is_condition_ = is_condition;
-
- bool is_enabled = is_enabled_;
- is_enabled_ = true;
- auto new_iter_values = Mutate(op->iter_values, inplace_mode)
- .as_or_throw<UnchangedOr<ffi::Array<PrimExpr>>>()
- .ValueOrUnchanged(op->iter_values);
- is_enabled_ = is_enabled;
- SBlock new_body =
- this->Mutate(op->block, inplace_mode).ValueOrUnchanged(op->block).as_or_throw<SBlock>();
- if (!new_predicate_unchanged || !new_iter_values.same_as(op->iter_values) ||
- !new_body.same_as(op->block)) {
- SBlockRealize new_block_realize = ffi::GetRef<SBlockRealize>(op);
- auto* n = new_block_realize.CopyOnWrite();
- n->predicate = std::move(new_predicate);
- n->iter_values = std::move(new_iter_values);
- n->block = std::move(new_body);
- return new_block_realize;
-
- } else {
- return ffi::Unchanged();
- }
-}
-
-UnchangedOr<Stmt> IndexDataTypeRewriter::Mutate_(const SBlockNode* op, InplaceMode inplace_mode) {
- auto new_alloc_buffers = WithDefRegionKind(kTVMFFIDefRegionKindSimple, [&] {
- return Mutate(op->alloc_buffers, inplace_mode)
- .as_or_throw<UnchangedOr<ffi::Array<BufferVar>>>()
- .ValueOrUnchanged(op->alloc_buffers);
- });
- auto new_match_buffers = op->match_buffers.Map([this](const MatchBufferRegion& match) {
- BufferVar buffer = WithDefRegionKind(kTVMFFIDefRegionKindSimple, [&] {
- return Mutate(match->buffer, InplaceMode::kDisallow)
- .as_or_throw<UnchangedOr<BufferVar>>()
- .ValueOrUnchanged(match->buffer);
- });
- TensorRegion source = VisitBufferRegion(match->source);
- if (buffer.same_as(match->buffer) && source.same_as(match->source)) return match;
- return MatchBufferRegion(buffer, source);
- });
- ffi::Array<TensorRegion> new_reads = op->reads.Map(
- [this](const TensorRegion& buffer_region) { return this->VisitBufferRegion(buffer_region); });
- ffi::Array<TensorRegion> new_writes = op->writes.Map(
- [this](const TensorRegion& buffer_region) { return this->VisitBufferRegion(buffer_region); });
- ffi::Array<IterVar> new_iter_vars =
- op->iter_vars.Map([this](const IterVar& iter_var) { return this->VisitIterVar(iter_var); });
- ffi::Optional<Stmt> new_init = std::nullopt;
- if (op->init.has_value()) {
- new_init = this->Mutate(op->init.value(), inplace_mode).ValueOrUnchanged(op->init.value());
- }
- ffi::Map<ffi::String, ffi::Any> new_annotations = VisitBlockAnnotations(op->annotations);
- auto new_body_result = this->Mutate(op->body, inplace_mode);
- bool new_body_unchanged = new_body_result.UnchangedOrSameAs(op->body);
- Stmt new_body = std::move(new_body_result).ValueOrUnchanged(op->body);
-
- if (!new_init.same_as(op->init) || !new_body_unchanged ||
- !new_alloc_buffers.same_as(op->alloc_buffers) ||
- !new_match_buffers.same_as(op->match_buffers) || !new_reads.same_as(op->reads) ||
- !new_writes.same_as(op->writes) || !new_iter_vars.same_as(op->iter_vars) ||
- !new_annotations.same_as(op->annotations)) {
- SBlock new_block = ffi::GetRef<SBlock>(op);
- SBlockNode* n = new_block.CopyOnWrite();
- n->alloc_buffers = std::move(new_alloc_buffers);
- n->match_buffers = std::move(new_match_buffers);
- n->reads = std::move(new_reads);
- n->writes = std::move(new_writes);
- n->iter_vars = std::move(new_iter_vars);
- n->init = std::move(new_init);
- n->annotations = std::move(new_annotations);
- n->body = std::move(new_body);
- return new_block;
- }
- return ffi::Unchanged();
-}
-
-ffi::Map<ffi::String, ffi::Any> IndexDataTypeRewriter::VisitBlockAnnotations(
- const ffi::Map<ffi::String, ffi::Any>& annotations) {
- auto new_annotations = annotations;
-
- std::function<Any(const Any&)> f_mutate_obj = [this, &f_mutate_obj](const Any& obj) -> Any {
- if (obj == nullptr) {
- return obj;
- }
- if (auto var = obj.as<Var>(); var && var.value()->ty.as<BufferTypeNode>()) {
- BufferVar buffer(var.value());
- if (BufferVar new_buffer = Mutate(buffer, InplaceMode::kDisallow)
- .as_or_throw<UnchangedOr<BufferVar>>()
- .ValueOrUnchanged(buffer);
- !new_buffer.same_as(buffer)) {
- return new_buffer;
- }
- } else if (obj.as<ffi::ArrayObj>()) {
- return obj.as_or_throw<ffi::Array<Any>>().Map(f_mutate_obj);
- }
- return obj;
- };
- for (const auto& [key, value] : annotations) {
- if (auto opt_object_ref = value.as<ffi::ObjectRef>()) {
- auto new_value = f_mutate_obj(*opt_object_ref);
- if (!new_value.same_as(*opt_object_ref)) {
- new_annotations.Set(key, new_value);
- }
- }
- }
- return new_annotations;
-}
-
-IterVar IndexDataTypeRewriter::VisitIterVar(const IterVar& iter_var) {
- bool is_enabled = is_enabled_;
- is_enabled_ = true;
- PrimVar new_var = Mutate(iter_var->var, InplaceMode::kDisallow)
- .ValueOrUnchanged(iter_var->var)
- .as_or_throw<PrimVar>();
- PrimExpr min =
- Mutate(iter_var->dom->min, InplaceMode::kDisallow).ValueOrUnchanged(iter_var->dom->min);
- PrimExpr extent =
- Mutate(iter_var->dom->extent, InplaceMode::kDisallow).ValueOrUnchanged(iter_var->dom->extent);
- is_enabled_ = is_enabled;
- if (!new_var.same_as(iter_var->var) || !min.same_as(iter_var->dom->min) ||
- !extent.same_as(iter_var->dom->extent)) {
- IterVar new_iter_var = iter_var;
- IterVarNode* n = new_iter_var.CopyOnWrite();
- n->var = std::move(new_var);
- n->dom = Range(min, extent);
- return new_iter_var;
- }
- return iter_var;
-}
-
-TensorRegion IndexDataTypeRewriter::VisitBufferRegion(const TensorRegion& buffer_region) {
- BufferVar remapped_buffer =
- Mutate(buffer_region->source.as_or_throw<BufferVar>(), InplaceMode::kDisallow)
- .as_or_throw<UnchangedOr<BufferVar>>()
- .ValueOrUnchanged(buffer_region->source.as_or_throw<BufferVar>());
-
- bool is_enabled = is_enabled_;
- is_enabled_ = true;
- auto new_region = buffer_region->region.Map([&](const Range& range) {
- return Range::FromMinExtent(
- this->Mutate(range->min, InplaceMode::kDisallow).ValueOrUnchanged(range->min),
- this->Mutate(range->extent, InplaceMode::kDisallow).ValueOrUnchanged(range->extent));
- });
- is_enabled_ = is_enabled;
-
- if (!remapped_buffer.same_as(buffer_region->source.as_or_throw<BufferVar>()) ||
- !new_region.same_as(buffer_region->region)) {
- return BufferRegion(remapped_buffer, new_region);
- } else {
- return buffer_region;
- }
-}
-
UnchangedOr<Stmt> IndexDataTypeRewriter::Mutate_(const BufferStoreNode* op,
InplaceMode inplace_mode) {
BufferStore store = ffi::GetRef<BufferStore>(op);
@@ -755,6 +550,9 @@ UnchangedOr<PrimExpr> IndexDataTypeRewriter::Mutate_(const prim::SelectNode* op,
IndexDataTypeNormalizer::IndexDataTypeNormalizer(PrimType target_data_type)
: target_data_type_(std::move(target_data_type)) {}
+IndexDataTypeNormalizer::IndexDataTypeNormalizer(PrimType target_data_type, const VTable* vtable)
+ : IndexDataTypeRewriter(vtable), target_data_type_(std::move(target_data_type)) {}
+
PrimFunc IndexDataTypeNormalizer::Rewrite(PrimFunc func) {
// Collect scalar dtype requirements without changing types. Buffer definitions
// are rewritten only after every scalar replacement has been seeded.
diff --git a/src/tirx/ir/data_type_rewriter.h b/src/tirx/ir/data_type_rewriter.h
index c2cca01ed8..7d1e09d85b 100644
--- a/src/tirx/ir/data_type_rewriter.h
+++ b/src/tirx/ir/data_type_rewriter.h
@@ -40,8 +40,6 @@ namespace tirx {
* bounds.
* - The data type of the binary and ternary expressions must be consistent with the data types of
* each of their operands.
- * - The data type of the bounds and binding values of block iter vars must be consistent with the
- * data type of the block iter vars.
*
* Usually we enforce the consistency of data types when constructing the IR nodes. However, such
* inconsistency may happen as a result of IR mutation in some passes. This class can be used as
@@ -49,14 +47,14 @@ namespace tirx {
*/
class DataTypeLegalizer : public StmtExprMutator {
public:
+ TVM_DEFINE_OBJECT_FUNCTOR_DEFAULT_CONSTRUCTOR(DataTypeLegalizer, StmtExprMutator)
using StmtExprMutator::Mutate;
using StmtExprMutator::Mutate_;
protected:
+ explicit DataTypeLegalizer(const VTable* vtable) : StmtExprMutator(vtable) {}
UnchangedOr<Stmt> Mutate_(const ForNode* op, InplaceMode inplace_mode) override;
UnchangedOr<Stmt> Mutate_(const AttrStmtNode* op, InplaceMode inplace_mode) override;
- UnchangedOr<Stmt> Mutate_(const SBlockRealizeNode* op, InplaceMode inplace_mode) override;
- UnchangedOr<Stmt> Mutate_(const SBlockNode* op, InplaceMode inplace_mode) override;
UnchangedOr<Stmt> Mutate_(const BindNode* op, InplaceMode inplace_mode) override;
UnchangedOr<PrimExpr> Mutate_(const prim::SelectNode* op, InplaceMode inplace_mode) override;
UnchangedOr<PrimExpr> Mutate_(const prim::RampNode* op, InplaceMode inplace_mode) override;
@@ -99,14 +97,14 @@ class DataTypeLegalizer : public StmtExprMutator {
*/
class IndexDataTypeRewriter : public DataTypeLegalizer {
public:
+ TVM_DEFINE_OBJECT_FUNCTOR_DEFAULT_CONSTRUCTOR(IndexDataTypeRewriter, DataTypeLegalizer)
using DataTypeLegalizer::Mutate;
using DataTypeLegalizer::Mutate_;
protected:
+ explicit IndexDataTypeRewriter(const VTable* vtable) : DataTypeLegalizer(vtable) {}
using Parent = DataTypeLegalizer;
UnchangedOr<ffi::Any> Mutate(ffi::AnyView value, InplaceMode inplace_mode) override;
- UnchangedOr<Stmt> Mutate_(const SBlockRealizeNode* op, InplaceMode inplace_mode) override;
- UnchangedOr<Stmt> Mutate_(const SBlockNode* op, InplaceMode inplace_mode) override;
UnchangedOr<Stmt> Mutate_(const BufferStoreNode* op, InplaceMode inplace_mode) override;
UnchangedOr<Stmt> Mutate_(const AttrStmtNode* op, InplaceMode inplace_mode) override;
UnchangedOr<PrimExpr> Mutate_(const TensorLoadNode* op, InplaceMode inplace_mode) override;
@@ -124,10 +122,6 @@ class IndexDataTypeRewriter : public DataTypeLegalizer {
UnchangedOr<Stmt> Mutate_(const ForNode* op, InplaceMode inplace_mode) override;
- ffi::Map<ffi::String, ffi::Any> VisitBlockAnnotations(
- const ffi::Map<ffi::String, ffi::Any>& annotations);
- TensorRegion VisitBufferRegion(const TensorRegion& region);
- IterVar VisitIterVar(const IterVar& iter_var);
// indicator of index expr to rewrite
bool is_enabled_{false};
// indicator of condition
@@ -149,6 +143,7 @@ class IndexDataTypeNormalizer : public IndexDataTypeRewriter {
PrimFunc Rewrite(PrimFunc func);
protected:
+ IndexDataTypeNormalizer(PrimType target_data_type, const VTable* vtable);
using Parent = IndexDataTypeRewriter;
UnchangedOr<PrimExpr> Mutate_(const IntImmNode* op, InplaceMode inplace_mode) override;
diff --git a/src/tirx/ir/function.cc b/src/tirx/ir/function.cc
index 0670113e71..3b87c1a62c 100644
--- a/src/tirx/ir/function.cc
+++ b/src/tirx/ir/function.cc
@@ -140,8 +140,6 @@ TVMFFIAny PrimFuncMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
} // namespace
-TVM_FFI_STATIC_INIT_BLOCK() { TensorIntrinNode::RegisterReflection(); }
-
// Get the function type of a PrimFunc
PrimFunc::PrimFunc(ffi::Array<tirx::Var> params, Stmt body, Type ret_type, DictAttrs attrs,
Span span) {
@@ -183,70 +181,5 @@ FuncType PrimFuncNode::func_type_annotation() const {
return FuncType(param_types, ret_type);
}
-class TensorIntrinManager {
- public:
- ffi::Map<ffi::String, tirx::TensorIntrin> reg;
-
- static TensorIntrinManager* Global() {
- static TensorIntrinManager* inst = new TensorIntrinManager();
- return inst;
- }
-};
-
-TensorIntrin::TensorIntrin(PrimFunc desc, PrimFunc impl) {
- // Check the number of func var is equal
- TVM_FFI_CHECK_EQ(desc->params.size(), impl->params.size(), ValueError)
- << "The number of parameters of the description and the implementation of the "
- "tensor intrinsic doesn't match.";
- auto is_handle = [](const Var& param) {
- return param->ty.as<PointerTypeNode>() != nullptr || param->ty.as<BufferTypeNode>() != nullptr;
- };
- for (size_t i = 0; i < desc->params.size(); i++) {
- TVM_FFI_CHECK(is_handle(desc->params[i]), ValueError)
- << "Parameters of the description of the "
- "tensor intrinsic should be handle only.";
- TVM_FFI_CHECK(is_handle(impl->params[i]), ValueError)
- << "Parameters of the implementation of "
- "the tensor intrinsic should be handle only.";
- }
- ffi::ObjectPtr<TensorIntrinNode> n = ffi::make_object<TensorIntrinNode>();
- n->desc = std::move(desc);
- n->impl = std::move(impl);
- data_ = std::move(n);
-}
-
-void TensorIntrin::Register(ffi::String name, TensorIntrin intrin, bool override) {
- TensorIntrinManager* manager = TensorIntrinManager::Global();
- if (!override) {
- TVM_FFI_CHECK_EQ(manager->reg.count(name), 0, ValueError)
- << "TensorIntrin '" << name << "' has already been registered";
- }
- manager->reg.Set(name, intrin);
-}
-
-ffi::Optional<TensorIntrin> TensorIntrin::Get(ffi::String name, bool allow_missing) {
- const TensorIntrinManager* manager = TensorIntrinManager::Global();
- auto it = manager->reg.find(name);
- if (it == manager->reg.end()) {
- if (allow_missing) {
- return std::nullopt;
- } else {
- TVM_FFI_THROW(ValueError) << "TensorIntrin '" << name << "' is not registered";
- }
- }
- return (*it).second;
-}
-
-TVM_FFI_STATIC_INIT_BLOCK() {
- namespace refl = tvm::ffi::reflection;
- refl::GlobalDef()
- .def("tirx.TensorIntrin",
- [](PrimFunc desc_func, PrimFunc intrin_func) {
- return TensorIntrin(desc_func, intrin_func);
- })
- .def("tirx.TensorIntrinRegister", TensorIntrin::Register)
- .def("tirx.TensorIntrinGet", TensorIntrin::Get);
-}
-
} // namespace tirx
} // namespace tvm
diff --git a/src/tirx/ir/ir_mutator_with_analyzer.cc b/src/tirx/ir/ir_mutator_with_analyzer.cc
index e26fdf5282..1c65f94874 100644
--- a/src/tirx/ir/ir_mutator_with_analyzer.cc
+++ b/src/tirx/ir/ir_mutator_with_analyzer.cc
@@ -25,7 +25,6 @@
#include <tvm/arith/iter_affine_map.h>
#include <tvm/ffi/cast.h>
#include <tvm/ir/op.h>
-#include <tvm/s_tir/stmt.h>
#include <tvm/tirx/analysis.h>
#include <tvm/tirx/builtin.h>
#include <tvm/tirx/op.h>
@@ -34,6 +33,18 @@
namespace tvm {
namespace tirx {
+void IRMutatorWithAnalyzer::InitVTable(VTable* vtable) { StmtExprMutator::InitVTable(vtable); }
+
+const IRMutatorWithAnalyzer::VTable* IRMutatorWithAnalyzer::GlobalVTable() {
+ static const VTable table = [] {
+ VTable table;
+ InitVTable(&table);
+ table.Finalize();
+ return table;
+ }();
+ return &table;
+}
+
using namespace tvm::prim;
using arith::detail::EnterConstraintFacts;
@@ -113,16 +124,6 @@ UnchangedOr<Stmt> IRMutatorWithAnalyzer::Mutate_(const ForNode* op, InplaceMode
});
}
-UnchangedOr<Stmt> IRMutatorWithAnalyzer::Mutate_(const SBlockNode* op, InplaceMode inplace_mode) {
- return constraint_scope_.WithNewScope([&]() -> UnchangedOr<Stmt> {
- for (const auto& iter_var : op->iter_vars) {
- analyzer_->Bind(iter_var->var, iter_var->dom);
- iter_vars_.Set(iter_var->var, iter_var->dom);
- }
- return StmtExprMutator::Mutate_(op, inplace_mode);
- });
-}
-
UnchangedOr<Stmt> IRMutatorWithAnalyzer::Mutate_(const BindNode* op, InplaceMode inplace_mode) {
auto value_result = this->Mutate(op->value, inplace_mode);
bool value_unchanged = value_result.UnchangedOrSameAs(op->value);
@@ -201,7 +202,8 @@ UnchangedOr<Stmt> IRMutatorWithAnalyzer::Mutate_(const IfThenElseNode* op,
UnchangedOr<Stmt> IRMutatorWithAnalyzer::Mutate_(const AttrStmtNode* op, InplaceMode inplace_mode) {
return constraint_scope_.WithNewScope([&]() -> UnchangedOr<Stmt> {
- if (op->attr_key == tirx::attr::thread_extent || op->attr_key == s_tir::attr::virtual_thread) {
+ if (op->attr_key == tirx::attr::thread_extent ||
+ op->attr_key == tvm::tirx::attr::virtual_thread) {
IterVar iv = op->node.as_or_throw<IterVar>();
TVM_FFI_ICHECK_NE(iv->thread_tag.length(), 0U);
Range dom = Range::FromMinExtent(IntImm(op->value.ty(), 0), op->value);
diff --git a/src/tirx/ir/ir_mutator_with_analyzer.h b/src/tirx/ir/ir_mutator_with_analyzer.h
index 20a6dd2587..bcc976c57a 100644
--- a/src/tirx/ir/ir_mutator_with_analyzer.h
+++ b/src/tirx/ir/ir_mutator_with_analyzer.h
@@ -51,12 +51,13 @@ class IRMutatorWithAnalyzer : public StmtExprMutator {
public:
using StmtExprMutator::Mutate;
using StmtExprMutator::Mutate_;
- explicit IRMutatorWithAnalyzer(const arith::Analyzer& analyzer) : analyzer_(analyzer.get()) {}
- explicit IRMutatorWithAnalyzer(arith::AnalyzerObj* analyzer) : analyzer_(analyzer) {}
+ explicit IRMutatorWithAnalyzer(const arith::Analyzer& analyzer)
+ : IRMutatorWithAnalyzer(analyzer.get()) {}
+ explicit IRMutatorWithAnalyzer(arith::AnalyzerObj* analyzer)
+ : IRMutatorWithAnalyzer(analyzer, GlobalVTable()) {}
// override functions that need to populate the context information.
UnchangedOr<Stmt> Mutate_(const ForNode* op, InplaceMode inplace_mode) override;
- UnchangedOr<Stmt> Mutate_(const SBlockNode* op, InplaceMode inplace_mode) override;
UnchangedOr<Stmt> Mutate_(const BindNode* op, InplaceMode inplace_mode) override;
UnchangedOr<Stmt> Mutate_(const IfThenElseNode* op, InplaceMode inplace_mode) override;
UnchangedOr<Stmt> Mutate_(const AttrStmtNode* op, InplaceMode inplace_mode) override;
@@ -66,6 +67,10 @@ class IRMutatorWithAnalyzer : public StmtExprMutator {
UnchangedOr<Expr> Mutate_(const CallNode* op, InplaceMode inplace_mode) override;
protected:
+ static void InitVTable(VTable* vtable);
+ IRMutatorWithAnalyzer(arith::AnalyzerObj* analyzer, const VTable* vtable)
+ : StmtExprMutator(vtable), analyzer_(analyzer) {}
+ static const VTable* GlobalVTable();
/*!
* \brief Mark all buffer-parameter shape values as positive values.
*
diff --git a/src/tirx/ir/ir_visitor_with_analyzer.cc b/src/tirx/ir/ir_visitor_with_analyzer.cc
index 9bdc895722..ae823b3dde 100644
--- a/src/tirx/ir/ir_visitor_with_analyzer.cc
+++ b/src/tirx/ir/ir_visitor_with_analyzer.cc
@@ -24,13 +24,13 @@
#include <tvm/ir/op.h>
#include <tvm/ir/prim/builtin.h>
-#include <tvm/s_tir/stmt.h>
#include <tvm/tirx/analysis.h>
... 4954 lines suppressed ...