From ba2606debe23cf4763c3a2f8403c75e63e22332e Mon Sep 17 00:00:00 2001 From: Tianqi Chen Date: Thu, 17 Sep 2026 00:44:36 +0000 Subject: [PATCH 01/18] [REFACTOR][S-TIR] Own TensorIntrin definitions and registry --- include/tvm/s_tir/function.h | 90 ++++++++++++++++ include/tvm/tirx/function.h | 56 ---------- python/tvm/ir/json_compact.py | 1 + python/tvm/s_tir/__init__.py | 2 +- python/tvm/s_tir/function.py | 80 ++++++++++++++ python/tvm/s_tir/schedule/schedule.py | 2 +- 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/tirx/__init__.py | 2 +- python/tvm/tirx/function.py | 54 ---------- src/s_tir/function.cc | 100 ++++++++++++++++++ .../multi_level_tiling_tensor_core.cc | 7 +- .../multi_level_tiling_with_intrin.cc | 5 +- .../schedule_rule/schedule_rule.cc | 3 +- src/s_tir/schedule/concrete_schedule.cc | 5 +- src/s_tir/schedule/primitive.h | 1 + .../schedule/primitive/blockize_tensorize.cc | 1 + src/s_tir/schedule/transform.cc | 3 +- src/tirx/ir/function.cc | 67 ------------ .../schedule/test_tir_schedule_analysis.py | 3 +- .../schedule/test_tir_schedule_tensorize.py | 18 ++-- tests/python/s_tir/test_tensor_intrin.py | 71 +++++++++++++ 23 files changed, 376 insertions(+), 203 deletions(-) create mode 100644 include/tvm/s_tir/function.h create mode 100644 python/tvm/s_tir/function.py create mode 100644 src/s_tir/function.cc create mode 100644 tests/python/s_tir/test_tensor_intrin.py diff --git a/include/tvm/s_tir/function.h b/include/tvm/s_tir/function.h new file mode 100644 index 000000000000..73ae70084fcb --- /dev/null +++ b/include/tvm/s_tir/function.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/function.h + * \brief Tensor intrinsics for schedulable TIR. + */ +#ifndef TVM_S_TIR_FUNCTION_H_ +#define TVM_S_TIR_FUNCTION_H_ + +#include + +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() + .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 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_FUNCTION_H_ diff --git a/include/tvm/tirx/function.h b/include/tvm/tirx/function.h index 8f73cc15cff8..67e43a3c02b7 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() - .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 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/python/tvm/ir/json_compact.py b/python/tvm/ir/json_compact.py index a7986a1491e7..b1e547887df9 100644 --- a/python/tvm/ir/json_compact.py +++ b/python/tvm/ir/json_compact.py @@ -19,6 +19,7 @@ import json _PRIM_TYPE_KEY_RENAMES = { + "tirx.TensorIntrin": "s_tir.TensorIntrin", "tirx.StringImm": "ir.prim.StringImm", "tirx.Cast": "ir.prim.Cast", "tirx.Add": "ir.prim.Add", diff --git a/python/tvm/s_tir/__init__.py b/python/tvm/s_tir/__init__.py index 164dcc99019b..1f8beb1faa9d 100644 --- a/python/tvm/s_tir/__init__.py +++ b/python/tvm/s_tir/__init__.py @@ -18,7 +18,7 @@ # pylint: disable=invalid-name """S-TIR namespace for scheduable TensorIR""" -from tvm.tirx.function import TensorIntrin +from .function import TensorIntrin # 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/function.py b/python/tvm/s_tir/function.py new file mode 100644 index 000000000000..00967c08eea6 --- /dev/null +++ b/python/tvm/s_tir/function.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/schedule/schedule.py b/python/tvm/s_tir/schedule/schedule.py index e97195ad54f4..cf3b6c613e4d 100644 --- a/python/tvm/s_tir/schedule/schedule.py +++ b/python/tvm/s_tir/schedule/schedule.py @@ -3043,7 +3043,7 @@ def mma_intrin(a: T.handle, b: T.handle, c: T.handle) -> None: ) ) - 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/tensor_intrin/arm_cpu.py b/python/tvm/s_tir/tensor_intrin/arm_cpu.py index 984830b957af..b5fbe3a0ac19 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 ae3ad7cd55a4..fb28a0c359d6 100644 --- a/python/tvm/s_tir/tensor_intrin/cuda.py +++ b/python/tvm/s_tir/tensor_intrin/cuda.py @@ -23,8 +23,9 @@ 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 5b19793ddfc0..175033804422 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/tirx/__init__.py b/python/tvm/tirx/__init__.py index af4b7161b0bf..051a26d9362b 100644 --- a/python/tvm/tirx/__init__.py +++ b/python/tvm/tirx/__init__.py @@ -56,7 +56,7 @@ 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/function.py b/python/tvm/tirx/function.py index 38c701075fde..708f55864976 100644 --- a/python/tvm/tirx/function.py +++ b/python/tvm/tirx/function.py @@ -164,60 +164,6 @@ def mem_copy_16_16(a: T.handle, b: T.handle) -> None: 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/src/s_tir/function.cc b/src/s_tir/function.cc new file mode 100644 index 000000000000..c33a538790bf --- /dev/null +++ b/src/s_tir/function.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/function.cc + * \brief Tensor intrinsic registry and construction. + */ +#include +#include + +namespace tvm { +namespace s_tir { +using tirx::BufferTypeNode; +using tirx::PrimFunc; + +TVM_FFI_STATIC_INIT_BLOCK() { TensorIntrinNode::RegisterReflection(); } + +class TensorIntrinManager { + public: + ffi::Map 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() != nullptr || param->ty.as() != 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 n = ffi::make_object(); + 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::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/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 7c4e9238a902..bf1ea750993d 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 @@ -18,6 +18,7 @@ */ #include #include +#include #include #include #include @@ -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 MultiLevelTilingTensorCoreNode::Apply(const Schedule& sch, ffi::Optional 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,7 +448,7 @@ std::vector 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) + tirx::SBlock intrin_block = TensorIntrin::Get(state->intrin_group.init_intrin) .value() ->desc->body.as_or_throw() ->block; 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 fab8937ae56c..571820508a75 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 @@ -18,6 +18,7 @@ */ #include +#include #include #include "../../schedule/analysis.h" @@ -52,7 +53,7 @@ class MultiLevelTilingWithIntrinNode : public MultiLevelTilingNode { protected: ffi::Array 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> vector_load_lens, ffi::Optional> reuse_read, ffi::Optional> 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( 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 eee9ef2685b8..0a34cd417795 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 #include +#include #include "../utils.h" @@ -325,7 +326,7 @@ ffi::Array ScheduleRule::DefaultRISCV(const int vlen) { const auto rvv_kernels_inventory = reg_rvv_intrinsics(current_target, /* inventory_only */ true) .cast>(); 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/schedule/concrete_schedule.cc b/src/s_tir/schedule/concrete_schedule.cc index bb3e1e77019f..145e054e731f 100644 --- a/src/s_tir/schedule/concrete_schedule.cc +++ b/src/s_tir/schedule/concrete_schedule.cc @@ -20,6 +20,7 @@ #include #include +#include #include @@ -927,7 +928,7 @@ SBlockRV ConcreteScheduleNode::Blockize(const ffi::Array& 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 +937,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/primitive.h b/src/s_tir/schedule/primitive.h index 83ab03ec5780..af589e22c174 100644 --- a/src/s_tir/schedule/primitive.h +++ b/src/s_tir/schedule/primitive.h @@ -20,6 +20,7 @@ #define TVM_S_TIR_SCHEDULE_PRIMITIVE_H_ #include +#include #include #include diff --git a/src/s_tir/schedule/primitive/blockize_tensorize.cc b/src/s_tir/schedule/primitive/blockize_tensorize.cc index 572f8c48048e..9b68f6e19962 100644 --- a/src/s_tir/schedule/primitive/blockize_tensorize.cc +++ b/src/s_tir/schedule/primitive/blockize_tensorize.cc @@ -21,6 +21,7 @@ #include #include #include +#include #include diff --git a/src/s_tir/schedule/transform.cc b/src/s_tir/schedule/transform.cc index c962e1f0f8fb..db58cf3865f7 100644 --- a/src/s_tir/schedule/transform.cc +++ b/src/s_tir/schedule/transform.cc @@ -20,6 +20,7 @@ #include #include #include +#include #include #include "../../tirx/transform/ir_utils.h" @@ -316,7 +317,7 @@ ffi::Optional TileWithTensorIntrin(const s_tir::Schedule& sch, const ffi::String& intrin_name, bool allow_padding) { ffi::Optional 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()) { diff --git a/src/tirx/ir/function.cc b/src/tirx/ir/function.cc index 0670113e7143..3b87c1a62c5d 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 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 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() != nullptr || param->ty.as() != 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 n = ffi::make_object(); - 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::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/tests/python/s_tir/schedule/test_tir_schedule_analysis.py b/tests/python/s_tir/schedule/test_tir_schedule_analysis.py index c8beaaa90c7d..f49c2cd14091 100644 --- a/tests/python/s_tir/schedule/test_tir_schedule_analysis.py +++ b/tests/python/s_tir/schedule/test_tir_schedule_analysis.py @@ -23,7 +23,7 @@ import tvm import tvm.testing from tvm.ir.prim import expr_deep_equal -from tvm.s_tir import Schedule +from tvm.s_tir import Schedule, TensorIntrin from tvm.s_tir.meta_schedule.testing import te_workload from tvm.s_tir.schedule.analysis import ( TensorizeInfo, @@ -49,7 +49,6 @@ floordiv, floormod, ) -from tvm.tirx.function import TensorIntrin def _make_vars(*args: str) -> list[Var]: diff --git a/tests/python/s_tir/schedule/test_tir_schedule_tensorize.py b/tests/python/s_tir/schedule/test_tir_schedule_tensorize.py index 3f5d98723bcd..4e07160d151b 100644 --- a/tests/python/s_tir/schedule/test_tir_schedule_tensorize.py +++ b/tests/python/s_tir/schedule/test_tir_schedule_tensorize.py @@ -21,7 +21,7 @@ import tvm import tvm.testing -from tvm import te, tirx +from tvm import s_tir, te, tirx from tvm.s_tir.schedule.testing import ( assert_structural_equal_ignore_global_symbol, verify_trace_roundtrip, @@ -496,11 +496,11 @@ def annotated_tensorized_matmul(a: T.handle, b: T.handle, c: T.handle) -> None: # fmt: off # pylint: disable=no-member,invalid-name,unused-variable,line-too-long,redefined-outer-name,unexpected-keyword-arg,too-many-nested-blocks -tirx.TensorIntrin.register("test_mma_intrin", mma_desc, mma_intrin) -tirx.TensorIntrin.register("test_annotated_mma_intrin", annotated_mma_desc, mma_intrin) -tirx.TensorIntrin.register("test_dot_product_intrin", dot_product_desc, dot_product_intrin) -tirx.TensorIntrin.register("test_outer_product_intrin", outer_product_desc, outer_product_intrin) -tirx.TensorIntrin.register("test_dot_product_intrin_annotated", dot_product_desc, dot_product_intrin_annotated) +s_tir.TensorIntrin.register("test_mma_intrin", mma_desc, mma_intrin) +s_tir.TensorIntrin.register("test_annotated_mma_intrin", annotated_mma_desc, mma_intrin) +s_tir.TensorIntrin.register("test_dot_product_intrin", dot_product_desc, dot_product_intrin) +s_tir.TensorIntrin.register("test_outer_product_intrin", outer_product_desc, outer_product_intrin) +s_tir.TensorIntrin.register("test_dot_product_intrin_annotated", dot_product_desc, dot_product_intrin_annotated) def test_tensorize_matmul(): @@ -749,9 +749,9 @@ def fetch_to_shared(block, idx): def test_tensor_intrin_look_up(): intrin_name = 'non_existent_intrin' - assert tirx.TensorIntrin.get(intrin_name, allow_missing=True) is None + assert s_tir.TensorIntrin.get(intrin_name, allow_missing=True) is None with pytest.raises(ValueError): - tirx.TensorIntrin.get(intrin_name) + s_tir.TensorIntrin.get(intrin_name) def test_tensorize_matmul_mixed_dtype(): @@ -911,7 +911,7 @@ def decode_i4s_to_f16_impl(compressed: T.handle, decompressed: T.handle) -> None 8, ) -tirx.TensorIntrin.register("test_decode_i4s_to_f16_intrin", decode_i4s_to_f16_desc, decode_i4s_to_f16_impl) +s_tir.TensorIntrin.register("test_decode_i4s_to_f16_intrin", decode_i4s_to_f16_desc, decode_i4s_to_f16_impl) def test_tensorize_arith_simplification(): # fmt: off diff --git a/tests/python/s_tir/test_tensor_intrin.py b/tests/python/s_tir/test_tensor_intrin.py new file mode 100644 index 000000000000..29b9824db849 --- /dev/null +++ b/tests/python/s_tir/test_tensor_intrin.py @@ -0,0 +1,71 @@ +# 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 intrinsic ownership and runtime compatibility.""" + +import json + +import pytest + +import tvm +import tvm.testing +from tvm import s_tir, tirx + + +def test_tensor_intrin_serialization(): + func = tirx.PrimFunc([], tirx.Evaluate(0)) + intrin = s_tir.TensorIntrin(func, func) + graph = json.loads(tvm.ir.save_json(intrin)) + assert any(node.get("type") == "s_tir.TensorIntrin" for node in graph["nodes"]) + for legacy in (False, True): + if legacy: + for node in graph["nodes"]: + if node.get("type") == "s_tir.TensorIntrin": + node["type"] = "tirx.TensorIntrin" + restored = tvm.ir.load_json(json.dumps(graph)) + assert isinstance(restored, s_tir.TensorIntrin) + assert restored.desc.same_as(restored.impl) + tvm.ir.assert_structural_equal(restored.desc, func) + + +def test_tensor_intrin_registration(): + func = tirx.PrimFunc([], tirx.Evaluate(0)) + name = "test_s_tir_tensor_intrin_registration" + s_tir.TensorIntrin.register(name, func, func, override=True) + assert s_tir.TensorIntrin.get(name).desc.same_as(func) + with pytest.raises(ValueError, match="already been registered"): + s_tir.TensorIntrin.register(name, func, func) + replacement = tirx.PrimFunc([], tirx.Evaluate(1)) + s_tir.TensorIntrin.register(name, func, replacement, override=True) + assert s_tir.TensorIntrin.get(name).impl.same_as(replacement) + + +def test_tensor_intrin_constructor_constraints(): + empty = tirx.PrimFunc([], tirx.Evaluate(0)) + scalar = tirx.PrimFunc([tirx.Var("x", "int32")], tirx.Evaluate(0)) + with pytest.raises(ValueError, match="number of parameters"): + s_tir.TensorIntrin(empty, scalar) + with pytest.raises(ValueError, match="description.*handle only"): + s_tir.TensorIntrin(scalar, scalar) + pointer = tirx.PrimFunc( + [tirx.Var("p", tvm.ir.PointerType(tvm.ir.PrimType("float32")))], tirx.Evaluate(0) + ) + with pytest.raises(ValueError, match="implementation.*handle only"): + s_tir.TensorIntrin(pointer, scalar) + + +if __name__ == "__main__": + tvm.testing.main() From ef86a73e77e367590bfb8c68f463d5b9980fa02a Mon Sep 17 00:00:00 2001 From: Tianqi Chen Date: Thu, 17 Sep 2026 01:24:49 +0000 Subject: [PATCH 02/18] [S-TIR] Move schedulable block ownership out of TIRX Own block, realization, and matching-region nodes in S-TIR, including reflection, structural callbacks, Python APIs, and serialized identities. Extend inherited traversal and dialect-aware helpers while retaining generic TIRX pass behavior and existing script syntax. --- include/tvm/relax/analysis.h | 3 +- .../tvm/relax/distributed/axis_group_graph.h | 15 +- include/tvm/s_tir/analysis.h | 11 +- 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 | 171 +++++++- include/tvm/s_tir/stmt_functor.h | 127 ++++++ include/tvm/s_tir/transform.h | 1 + include/tvm/s_tir/utils.h | 13 +- include/tvm/script/ir_builder/base.h | 2 +- include/tvm/tirx/analysis.h | 3 +- include/tvm/tirx/script/builder/frame.h | 5 +- include/tvm/tirx/script/builder/ir.h | 2 +- include/tvm/tirx/stmt.h | 166 ------- include/tvm/tirx/stmt_functor.h | 205 +++++---- python/tvm/ir/json_compact.py | 3 + python/tvm/relax/analysis/analysis.py | 3 +- python/tvm/s_tir/__init__.py | 1 + python/tvm/s_tir/analysis/__init__.py | 7 +- .../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/sblock_dependence_info.py | 3 +- python/tvm/s_tir/sblock_scope.py | 3 +- python/tvm/s_tir/schedule/schedule.py | 3 +- python/tvm/s_tir/schedule/state.py | 3 +- python/tvm/s_tir/stmt.py | 179 ++++++++ python/tvm/tirx/__init__.py | 2 +- python/tvm/tirx/stmt.py | 154 +------ src/relax/analysis/layout_transformation.cc | 35 +- src/relax/analysis/tir_op_pattern_kind.cc | 36 +- src/relax/backend/task_extraction.cc | 11 +- .../lower_global_view_to_local_view.cc | 50 ++- src/relax/transform/fuse_tir.cc | 41 +- .../transform/split_call_tir_by_pattern.cc | 54 +-- .../transform/split_layout_rewrite_preproc.cc | 44 +- .../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 | 3 +- src/s_tir/analysis/oob_checker.cc | 6 +- .../analysis/sblock_access_region_detector.cc | 53 +-- .../sblock_buffer_access_lca_detector.cc | 29 +- src/s_tir/analysis/verify_gpu_code.cc | 3 +- src/s_tir/analysis/verify_well_formed.cc | 138 ++++++ .../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 | 307 +++++++++++++ src/s_tir/ir/ir_mutator_with_analyzer.cc | 48 ++ src/s_tir/ir/ir_mutator_with_analyzer.h | 67 +++ src/s_tir/ir/ir_visitor_with_analyzer.cc | 45 ++ src/s_tir/ir/ir_visitor_with_analyzer.h | 56 +++ src/s_tir/ir/specialize.cc | 70 +++ src/s_tir/ir/tir_visitor_with_path.cc | 114 +++++ .../feature_extractor/per_store_feature.cc | 1 + src/s_tir/meta_schedule/module_equality.cc | 8 +- .../postproc/rewrite_tensorize.cc | 4 +- .../postproc/rewrite_unbound_block.cc | 1 + .../schedule_rule/cross_thread_reduction.cc | 3 +- .../schedule_rule/multi_level_tiling.cc | 2 +- .../multi_level_tiling_tensor_core.cc | 21 +- .../multi_level_tiling_wide_vector.cc | 5 +- 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 | 1 + 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 | 1 + src/s_tir/schedule/primitive/annotate.cc | 1 + .../schedule/primitive/blockize_tensorize.cc | 1 + src/s_tir/schedule/primitive/cache_index.cc | 1 + .../schedule/primitive/cache_index_helpers.cc | 15 +- .../schedule/primitive/cache_index_helpers.h | 10 +- .../schedule/primitive/cache_read_write.cc | 1 + src/s_tir/schedule/primitive/compute_at.cc | 1 + .../schedule/primitive/decompose_padding.cc | 1 + src/s_tir/schedule/primitive/for_kind.cc | 1 + .../schedule/primitive/get_block_loop.cc | 1 + .../schedule/primitive/hide_buffer_access.cc | 1 + .../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 + .../primitive/reorder_block_iter_var.cc | 1 + .../schedule/primitive/rolling_buffer.cc | 1 + src/s_tir/schedule/schedule.cc | 1 + src/s_tir/schedule/state.cc | 1 + src/s_tir/schedule/traced_schedule.cc | 2 + src/s_tir/schedule/traced_schedule.h | 1 + src/s_tir/schedule/transform.cc | 7 +- src/s_tir/schedule/transform.h | 11 +- src/s_tir/schedule/utils.h | 12 +- src/s_tir/stmt.cc | 412 ++++++++++++++++++ src/s_tir/stmt_functor.cc | 234 ++++++++++ .../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 | 2 +- .../transform/convert_blocks_to_opaque.cc | 3 +- src/s_tir/transform/default_gpu_schedule.cc | 21 +- src/s_tir/transform/hoist_expression.cc | 13 +- src/s_tir/transform/inject_double_buffer.cc | 2 +- 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_virtual_thread.cc | 16 +- src/s_tir/transform/ir_utils.cc | 77 ++++ src/s_tir/transform/ir_utils.h | 47 ++ src/s_tir/transform/lift_thread_binding.cc | 2 +- src/s_tir/transform/loop_partition.cc | 3 +- 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 | 2 +- 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 +- .../transform/memhammer_lower_auto_copy.cc | 2 +- src/s_tir/transform/memhammer_rewrite_rule.h | 2 +- .../transform/memhammer_tensorcore_rewrite.cc | 1 + .../merge_shared_memory_allocations.cc | 2 +- .../plan_update_buffer_allocation_location.cc | 4 +- .../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 +- .../transform/renormalize_split_pattern.cc | 5 +- src/s_tir/transform/rewrite_unsafe_select.cc | 2 +- src/s_tir/transform/stmt_extension.cc | 132 ++++++ src/s_tir/transform/storage_access.h | 2 +- .../transform/tensorcore_infer_fragment.cc | 25 +- src/s_tir/transform/thread_storage_sync.cc | 3 +- .../transform/transform_mma_buffer_layout.cc | 4 +- src/s_tir/transform/unify_thread_binding.cc | 3 +- .../using_assume_to_reduce_branches.cc | 4 +- src/te/operation/create_primfunc.cc | 126 +++--- src/tirx/analysis/verify_tirx_well_formed.cc | 40 +- src/tirx/analysis/verify_well_formed.cc | 120 +---- src/tirx/analysis/verify_well_formed.h | 29 ++ src/tirx/ir/data_type_rewriter.cc | 234 ++-------- src/tirx/ir/data_type_rewriter.h | 28 +- src/tirx/ir/ir_mutator_with_analyzer.cc | 39 +- src/tirx/ir/ir_mutator_with_analyzer.h | 17 +- src/tirx/ir/ir_visitor_with_analyzer.cc | 27 +- src/tirx/ir/ir_visitor_with_analyzer.h | 12 +- src/tirx/ir/script/script_complete.cc | 32 +- src/tirx/ir/specialize.cc | 60 ++- src/tirx/ir/specialize.h | 39 ++ src/tirx/ir/stmt.cc | 385 ++-------------- src/tirx/ir/stmt_functor.cc | 164 ++----- src/tirx/ir/tir_visitor_with_path.cc | 85 +--- src/tirx/ir/tir_visitor_with_path.h | 50 ++- src/tirx/script/builder/frame.cc | 16 +- src/tirx/script/builder/ir.cc | 7 +- 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/utils.h | 2 +- src/tirx/transform/flatten_buffer.cc | 42 +- .../transform/force_narrow_index_to_i32.cc | 17 +- .../transform/inline_private_functions.cc | 8 +- src/tirx/transform/ir_utils.cc | 120 +---- src/tirx/transform/ir_utils.h | 16 - src/tirx/transform/narrow_datatype.cc | 15 +- src/tirx/transform/stmt_extension.h | 159 +++++++ tests/cpp/ir_functor_test.cc | 37 +- tests/cpp/s_tir_functor_test.cc | 316 ++++++++++++++ ...test_analysis_suggest_layout_transforms.py | 4 +- .../analysis/test_sblock_access_region.py | 2 +- .../s_tir/base/test_sblock_dependence_info.py | 4 +- ...chedule_schedule_rule_apply_custom_rule.py | 2 +- .../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 +- tests/python/s_tir/test_s_tir_renew_defs.py | 2 +- tests/python/s_tir/test_stmt.py | 66 +++ ...tir_transform_memhammer_lower_auto_copy.py | 2 +- .../test_s_tir_transform_mma_buffer_layout.py | 4 +- .../test_tir_analysis_verify_well_formed.py | 6 +- .../tvmscript/test_tvmscript_complete.py | 10 +- .../tvmscript/test_tvmscript_error_report.py | 4 +- .../test_tvmscript_ir_builder_tir.py | 16 +- .../tvmscript/test_tvmscript_printer_tir.py | 4 +- .../tvmscript/test_tvmscript_roundtrip.py | 40 +- 201 files changed, 3997 insertions(+), 2192 deletions(-) create mode 100644 include/tvm/s_tir/stmt_functor.h create mode 100644 python/tvm/s_tir/stmt.py create mode 100644 src/s_tir/analysis/verify_well_formed.cc create mode 100644 src/s_tir/ir/data_type_rewriter.cc create mode 100644 src/s_tir/ir/ir_mutator_with_analyzer.cc create mode 100644 src/s_tir/ir/ir_mutator_with_analyzer.h create mode 100644 src/s_tir/ir/ir_visitor_with_analyzer.cc create mode 100644 src/s_tir/ir/ir_visitor_with_analyzer.h create mode 100644 src/s_tir/ir/specialize.cc create mode 100644 src/s_tir/ir/tir_visitor_with_path.cc create mode 100644 src/s_tir/stmt.cc create mode 100644 src/s_tir/stmt_functor.cc create mode 100644 src/s_tir/transform/ir_utils.cc create mode 100644 src/s_tir/transform/ir_utils.h create mode 100644 src/s_tir/transform/stmt_extension.cc create mode 100644 src/tirx/analysis/verify_well_formed.h create mode 100644 src/tirx/ir/specialize.h create mode 100644 src/tirx/transform/stmt_extension.h create mode 100644 tests/cpp/s_tir_functor_test.cc create mode 100644 tests/python/s_tir/test_stmt.py diff --git a/include/tvm/relax/analysis.h b/include/tvm/relax/analysis.h index 454e218b945d..092745f1718e 100644 --- a/include/tvm/relax/analysis.h +++ b/include/tvm/relax/analysis.h @@ -29,6 +29,7 @@ #include #include #include +#include #include #include @@ -646,7 +647,7 @@ TVM_DLL bool CheckWellFormed(ffi::Variant obj, bool check_ty * from the object (block or buffer) to it's index map transformation. */ -TVM_DLL ffi::Map> SuggestLayoutTransforms( +TVM_DLL ffi::Map> SuggestLayoutTransforms( const Function& fn, ffi::Array 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 1b94df4abc4e..9590394297f7 100644 --- a/include/tvm/relax/distributed/axis_group_graph.h +++ b/include/tvm/relax/distributed/axis_group_graph.h @@ -23,8 +23,9 @@ #include #include #include +#include +#include #include -#include #include #include @@ -65,7 +66,7 @@ Var GetShardingVarFromIndex(PrimExpr index, ffi::Map 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> GetTIRVarAxisGraph(const PrimFunc& prim_func) { auto extractor = ffi::make_object(); @@ -119,14 +120,14 @@ class BufferAxisGraphExtractor : public StmtExprVisitor { private: ffi::Optional 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 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(), op->indices}); return std::nullopt; @@ -158,12 +159,12 @@ class BufferAxisGraphExtractor : public StmtExprVisitor { return true; } - ffi::Optional Visit_(const SBlockNode* op) final { + ffi::Optional 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 4a7c72d386c6..48fdcded6705 100644 --- a/include/tvm/s_tir/analysis.h +++ b/include/tvm/s_tir/analysis.h @@ -27,6 +27,7 @@ #include #include #include +#include #include #include #include @@ -47,8 +48,8 @@ namespace tirx { * - second: write regions * - third: opaque regions */ -TVM_DLL ffi::Array> GetSBlockAccessRegion( - const SBlock& block, const ffi::Map& buffer_var_map); +TVM_DLL ffi::Array> GetSBlockAccessRegion( + const s_tir::SBlock& block, const ffi::Map& buffer_var_map); /*! * \brief Auto detect the block read/write region according to its body stmt. An opaque access will @@ -58,8 +59,8 @@ TVM_DLL ffi::Array> GetSBlockAccessRegion( * It is a map from buffer var to the buffer * \return An array only consisting of the read regions and write regions of the input block */ -TVM_DLL ffi::Array> GetSBlockReadWriteRegion( - const SBlock& block, const ffi::Map& buffer_var_map); +TVM_DLL ffi::Array> GetSBlockReadWriteRegion( + const s_tir::SBlock& block, const ffi::Map& 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> 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 diff --git a/include/tvm/s_tir/sblock_scope.h b/include/tvm/s_tir/sblock_scope.h index 34ff87025f9f..3509f2a6db2b 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 +#include +#include #include #include -#include #include #include @@ -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 @@ -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 Visit(ffi::AnyView value) override { if (value.as()) return std::nullopt; - return StmtExprVisitor::Visit(value); + return s_tir::StmtExprVisitor::Visit(value); } /*! @@ -185,7 +186,7 @@ class SRefTreeCreator : public StmtExprVisitor { ffi::Optional Visit_(const ForNode* loop) final; - ffi::Optional Visit_(const SBlockRealizeNode* realize) final; + ffi::Optional Visit_(const s_tir::SBlockRealizeNode* realize) final; ffi::Optional 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 b149eb3b2c75..7af08a1630f3 100644 --- a/include/tvm/s_tir/schedule/schedule.h +++ b/include/tvm/s_tir/schedule/schedule.h @@ -23,6 +23,7 @@ #include #include #include +#include #include namespace tvm { diff --git a/include/tvm/s_tir/schedule/state.h b/include/tvm/s_tir/schedule/state.h index bbaeb3039652..feb1650ee299 100644 --- a/include/tvm/s_tir/schedule/state.h +++ b/include/tvm/s_tir/schedule/state.h @@ -27,6 +27,7 @@ #include #include #include +#include #include #include diff --git a/include/tvm/s_tir/stmt.h b/include/tvm/s_tir/stmt.h index 57dcd70aabaa..818776242030 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,10 +27,177 @@ #ifndef TVM_S_TIR_STMT_H_ #define TVM_S_TIR_STMT_H_ -#include +#include 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. */ + tirx::BufferRegion source; + + static void RegisterReflection() { + namespace refl = tvm::ffi::reflection; + refl::ObjectDef() + .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, tirx::BufferRegion 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 iter_vars; + /*! \brief The read buffer regions of the block. */ + ffi::Array reads; + /*! \brief The write buffer regions of the block. */ + ffi::Array writes; + /*! \brief The name_hint of the block. */ + ffi::String name_hint; + /*! \brief The buffer allocated in the block. */ + ffi::Array alloc_buffers; + /*! \brief The match buffer regions. */ + ffi::Array match_buffers; + /*! \brief The annotation of the block. */ + ffi::Map 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 init; + /*! \brief The body of the block. */ + tirx::Stmt body; + + static void RegisterReflection() { + namespace refl = tvm::ffi::reflection; + refl::ObjectDef() + .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 iter_vars, ffi::Array reads, + ffi::Array writes, ffi::String name_hint, tirx::Stmt body, + ffi::Optional init = std::nullopt, + ffi::Array alloc_buffers = ffi::Array(), + ffi::Array match_buffers = ffi::Array(), + ffi::Map annotations = ffi::Map(), + Span span = Span()); + + TVM_DLL explicit SBlock(ffi::String name_hint, tirx::Stmt body, + ffi::Array alloc_buffers = ffi::Array(), + 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 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() + .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 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 { /*! diff --git a/include/tvm/s_tir/stmt_functor.h b/include/tvm/s_tir/stmt_functor.h new file mode 100644 index 000000000000..80d6da180bd9 --- /dev/null +++ b/include/tvm/s_tir/stmt_functor.h @@ -0,0 +1,127 @@ +/* + * 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 +#include + +namespace tvm { +namespace s_tir { + +/*! + * \brief Extend TIRX statement dispatch with schedulable blocks. + * \tparam FType The statement signature, retaining the VisitStmt API. + */ +template +class StmtFunctor; + +template +class StmtFunctor + : public tirx::StmtFunctor { + using Parent = tirx::StmtFunctor; + + public: + TVM_DEFINE_OBJECT_FUNCTOR_DEFAULT_CONSTRUCTOR(StmtFunctor, Parent) + using Parent::VisitStmt_; + + virtual R VisitStmt_(const SBlockNode* op, Args... args) { + return this->VisitStmtDefault_(op, std::forward(args)...); + } + virtual R VisitStmt_(const SBlockRealizeNode* op, Args... args) { + return this->VisitStmtDefault_(op, std::forward(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(vtable); + Parent::template SetDispatch(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. + * S-TIR registers this same native policy into generic TIRX visitors. This + * subclass additionally exposes virtual block hooks for block-aware passes. + * Other foreign nodes without a native policy still use structural fallback. + */ +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 Visit_(const SBlockNode* op); + virtual ffi::Optional Visit_(const SBlockRealizeNode* op); + + // Shared native traversal for specialized TIRX helpers extended by S-TIR. + static ffi::Optional VisitBlock(tirx::StmtExprVisitor* visitor, + const SBlockNode* op); + static ffi::Optional 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. + * S-TIR registers the same policy into generic TIRX mutators so allocation + * remaps precede region uses and block binders retain their existing identity. + * Foreign nodes without a registered policy still use structural fallback. + */ +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 Mutate_(const SBlockNode* op, InplaceMode inplace_mode); + virtual UnchangedOr Mutate_(const SBlockRealizeNode* op, InplaceMode inplace_mode); + + // Share block ownership and binder rules with specialized TIRX helpers. + static UnchangedOr MutateBlock(tirx::StmtExprMutator* mutator, const SBlockNode* op, + InplaceMode inplace_mode); + static UnchangedOr 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/transform.h b/include/tvm/s_tir/transform.h index 6bbc91edfbf3..f898be2e467e 100644 --- a/include/tvm/s_tir/transform.h +++ b/include/tvm/s_tir/transform.h @@ -26,6 +26,7 @@ #include #include +#include #include #include diff --git a/include/tvm/s_tir/utils.h b/include/tvm/s_tir/utils.h index cbbe930574bc..4d127818e958 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 +#include #include #include @@ -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(); \ @@ -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& stmt2ref, // NOLINT(*) const Stmt& stmt, int seq_index, bool include_loops = true) { - if (const auto* realize = stmt.as()) { - const SBlockNode* block = realize->block.get(); + if (const auto* realize = stmt.as()) { + 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()) { + } else if (const auto* block = stmt.as()) { TVM_FFI_ICHECK(stmt2ref.count(block)); stmt2ref.at(block)->seq_index = seq_index; } else if (const auto* loop = stmt.as()) { diff --git a/include/tvm/script/ir_builder/base.h b/include/tvm/script/ir_builder/base.h index 85c8f991d9f4..417e9b0c386a 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 1379046a219d..ee163d5a26e6 100644 --- a/include/tvm/tirx/analysis.h +++ b/include/tvm/tirx/analysis.h @@ -27,7 +27,6 @@ #include #include #include -#include #include #include #include @@ -131,7 +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 + * - Expressions within a s_tir::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 diff --git a/include/tvm/tirx/script/builder/frame.h b/include/tvm/tirx/script/builder/frame.h index 38259a1a57cc..b5c1774808f3 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 #include #include #include @@ -87,7 +88,7 @@ class PrimFuncFrameNode : public TIRFrameNode { ffi::Array 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 alloc_buffers; /*! \brief The match buffer regions. */ - ffi::Array match_buffers; + ffi::Array match_buffers; /*! \brief The annotation of the block. */ ffi::Optional> 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 9e3e78a88c11..e306c6b27648 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 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 65baf862e1ca..912088a5d022 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() - .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 iter_vars; - /*! \brief The read buffer regions of the block. */ - ffi::Array reads; - /*! \brief The write buffer regions of the block. */ - ffi::Array writes; - /*! \brief The name_hint of the block. */ - ffi::String name_hint; - /*! \brief The buffer allocated in the block. */ - ffi::Array alloc_buffers; - /*! \brief The match buffer regions. */ - ffi::Array match_buffers; - /*! \brief The annotation of the block. */ - ffi::Map 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 init; - /*! \brief The body of the block. */ - Stmt body; - - static void RegisterReflection() { - namespace refl = tvm::ffi::reflection; - refl::ObjectDef() - .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 iter_vars, ffi::Array reads, - ffi::Array writes, ffi::String name_hint, Stmt body, - ffi::Optional init = std::nullopt, - ffi::Array alloc_buffers = ffi::Array(), - ffi::Array match_buffers = ffi::Array(), - ffi::Map annotations = ffi::Map(), - Span span = Span()); - - TVM_DLL explicit SBlock(ffi::String name_hint, Stmt body, - ffi::Array alloc_buffers = ffi::Array(), - 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 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() - .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 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. diff --git a/include/tvm/tirx/stmt_functor.h b/include/tvm/tirx/stmt_functor.h index c1147f5f8010..6d3a8ef3ffbc 100644 --- a/include/tvm/tirx/stmt_functor.h +++ b/include/tvm/tirx/stmt_functor.h @@ -52,124 +52,112 @@ namespace tirx { template class StmtFunctor; +#define STMT_FUNCTOR_DEFAULT \ + { \ + return VisitStmtDefault_(op, std::forward(args)...); \ + } + +#define IR_STMT_FUNCTOR_DISPATCH(OP) \ + vtable->template SetDispatch([](const ffi::ObjectRef& n, TSelf* self, Args... args) { \ + return self->VisitStmt_(static_cast(n.get()), std::forward(args)...); \ + }); + template class StmtFunctor { private: - using TSelf = StmtFunctor; + using TSelf = StmtFunctor; public: /*! \brief The result type of this functor. */ using result_type = R; - /*! \brief Construct a functor with the TIRx statement hooks. */ StmtFunctor() : StmtFunctor(GlobalVTable()) {} - /*! \brief Destroy through the statement functor base. */ - virtual ~StmtFunctor() = default; - /*! \brief Dispatch a statement, forwarding additional arguments to its hook. */ - TVM_FFI_INLINE R operator()(const Stmt& node, Args... args) { - return Dispatch(node, std::forward(args)...); - } - /*! \brief Dispatch to a node hook, including registered ancestor hooks. */ - TVM_FFI_INLINE virtual R Dispatch(const Stmt& node, Args... args) { - TVM_FFI_ICHECK(node.defined()) << "Cannot dispatch a null statement"; - return (*vtable_)(node, this, std::forward(args)...); - } - - virtual R Dispatch_(const BindNode* node, Args... args) { - return DispatchDefault_(node, std::forward(args)...); - } - virtual R Dispatch_(const AttrStmtNode* node, Args... args) { - return DispatchDefault_(node, std::forward(args)...); - } - virtual R Dispatch_(const IfThenElseNode* node, Args... args) { - return DispatchDefault_(node, std::forward(args)...); - } - virtual R Dispatch_(const ForNode* node, Args... args) { - return DispatchDefault_(node, std::forward(args)...); - } - virtual R Dispatch_(const WhileNode* node, Args... args) { - return DispatchDefault_(node, std::forward(args)...); - } - virtual R Dispatch_(const ReturnNode* node, Args... args) { - return DispatchDefault_(node, std::forward(args)...); - } - virtual R Dispatch_(const BreakNode* node, Args... args) { - return DispatchDefault_(node, std::forward(args)...); - } - virtual R Dispatch_(const ContinueNode* node, Args... args) { - return DispatchDefault_(node, std::forward(args)...); - } - virtual R Dispatch_(const AllocBufferNode* node, Args... args) { - return DispatchDefault_(node, std::forward(args)...); - } - virtual R Dispatch_(const DeclBufferNode* node, Args... args) { - return DispatchDefault_(node, std::forward(args)...); - } - virtual R Dispatch_(const BufferStoreNode* node, Args... args) { - return DispatchDefault_(node, std::forward(args)...); - } - virtual R Dispatch_(const AssertStmtNode* node, Args... args) { - return DispatchDefault_(node, std::forward(args)...); - } - virtual R Dispatch_(const SeqStmtNode* node, Args... args) { - return DispatchDefault_(node, std::forward(args)...); - } - virtual R Dispatch_(const EvaluateNode* node, Args... args) { - return DispatchDefault_(node, std::forward(args)...); - } - virtual R Dispatch_(const SBlockNode* node, Args... args) { - return DispatchDefault_(node, std::forward(args)...); - } - virtual R Dispatch_(const SBlockRealizeNode* node, Args... args) { - return DispatchDefault_(node, std::forward(args)...); - } - virtual R Dispatch_(const ScopeIdDefStmtNode* node, Args... args) { - return DispatchDefault_(node, std::forward(args)...); - } - virtual R Dispatch_(const tirx::TilePrimitiveCallNode* node, Args... args) { - return DispatchDefault_(node, std::forward(args)...); + /*! \brief virtual destructor */ + virtual ~StmtFunctor() {} + /*! + * \brief Same as call. + * \param n The stmt node. + * \param args Additional arguments. + * \return The result of the call + */ + R operator()(const Stmt& n, Args... args) { return VisitStmt(n, std::forward(args)...); } + /*! + * \brief The functor call. + * \param n The stmt node. + * \param args Additional arguments. + * \return The result of the call + */ + virtual R VisitStmt(const Stmt& n, Args... args) { + return (*vtable_)(n, this, std::forward(args)...); } - /*! \brief Default behavior for statement hooks not overridden by a subclass. */ - virtual R DispatchDefault_(const ffi::Object* node, Args...) { - TVM_FFI_THROW(InternalError) << "Do not have a default for " << node->GetTypeKey(); + // Functions that can be overriden by subclass + virtual R VisitStmt_(const BindNode* op, Args... args) STMT_FUNCTOR_DEFAULT; + virtual R VisitStmt_(const AttrStmtNode* op, Args... args) STMT_FUNCTOR_DEFAULT; + virtual R VisitStmt_(const IfThenElseNode* op, Args... args) STMT_FUNCTOR_DEFAULT; + virtual R VisitStmt_(const ForNode* op, Args... args) STMT_FUNCTOR_DEFAULT; + virtual R VisitStmt_(const WhileNode* op, Args... args) STMT_FUNCTOR_DEFAULT; + virtual R VisitStmt_(const ReturnNode* op, Args... args) STMT_FUNCTOR_DEFAULT; + virtual R VisitStmt_(const BreakNode* op, Args... args) STMT_FUNCTOR_DEFAULT; + virtual R VisitStmt_(const ContinueNode* op, Args... args) STMT_FUNCTOR_DEFAULT; + virtual R VisitStmt_(const AllocBufferNode* op, Args... args) STMT_FUNCTOR_DEFAULT; + virtual R VisitStmt_(const DeclBufferNode* op, Args... args) STMT_FUNCTOR_DEFAULT; + virtual R VisitStmt_(const BufferStoreNode* op, Args... args) STMT_FUNCTOR_DEFAULT; + virtual R VisitStmt_(const AssertStmtNode* op, Args... args) STMT_FUNCTOR_DEFAULT; + virtual R VisitStmt_(const SeqStmtNode* op, Args... args) STMT_FUNCTOR_DEFAULT; + virtual R VisitStmt_(const EvaluateNode* op, Args... args) STMT_FUNCTOR_DEFAULT; + virtual R VisitStmt_(const ScopeIdDefStmtNode* op, Args... args) STMT_FUNCTOR_DEFAULT; + virtual R VisitStmt_(const tirx::TilePrimitiveCallNode* op, Args... args) STMT_FUNCTOR_DEFAULT; + virtual R VisitStmtDefault_(const ffi::Object* op, Args...) { + TVM_FFI_THROW(InternalError) << "Do not have a default for " << op->GetTypeKey(); TVM_FFI_UNREACHABLE(); } protected: - /*! \brief Dispatch table shared by this signature and its subclasses. */ using VTable = ObjectFunctor; - /*! \brief Construct with a finalized table that outlives the functor. */ + explicit StmtFunctor(const VTable* vtable) : vtable_(vtable) {} - /*! \brief Register statement hooks in a fresh mutable table. */ + + // Register inherited hooks in a fresh table before adding dialect nodes. static void InitVTable(VTable* vtable) { - SetDispatch(vtable); - SetDispatch(vtable); - SetDispatch(vtable); - SetDispatch(vtable); - SetDispatch(vtable); - SetDispatch(vtable); - SetDispatch(vtable); - SetDispatch(vtable); - SetDispatch(vtable); - SetDispatch(vtable); - SetDispatch(vtable); - SetDispatch(vtable); - SetDispatch(vtable); - SetDispatch(vtable); - SetDispatch(vtable); - SetDispatch(vtable); - SetDispatch(vtable); - SetDispatch(vtable); + IR_STMT_FUNCTOR_DISPATCH(BindNode); + IR_STMT_FUNCTOR_DISPATCH(AttrStmtNode); + IR_STMT_FUNCTOR_DISPATCH(IfThenElseNode); + IR_STMT_FUNCTOR_DISPATCH(ForNode); + IR_STMT_FUNCTOR_DISPATCH(WhileNode); + IR_STMT_FUNCTOR_DISPATCH(ReturnNode); + IR_STMT_FUNCTOR_DISPATCH(BreakNode); + IR_STMT_FUNCTOR_DISPATCH(ContinueNode); + IR_STMT_FUNCTOR_DISPATCH(AllocBufferNode); + IR_STMT_FUNCTOR_DISPATCH(DeclBufferNode); + IR_STMT_FUNCTOR_DISPATCH(AssertStmtNode); + IR_STMT_FUNCTOR_DISPATCH(SeqStmtNode); + IR_STMT_FUNCTOR_DISPATCH(EvaluateNode); + IR_STMT_FUNCTOR_DISPATCH(BufferStoreNode); + IR_STMT_FUNCTOR_DISPATCH(ScopeIdDefStmtNode); + IR_STMT_FUNCTOR_DISPATCH(tirx::TilePrimitiveCallNode); } - /*! \brief Register an additional node hook implemented by Self. */ + template static void SetDispatch(VTable* vtable) { - vtable->template SetDispatch( - [](const ffi::ObjectRef& node, TSelf* self, Args... args) -> R { - return static_cast(self)->Dispatch_(static_cast(node.get()), - std::forward(args)...); - }); + vtable->template SetDispatch([](const ffi::ObjectRef& node, TSelf* self, Args... args) { + return static_cast(self)->VisitStmt_(static_cast(node.get()), + std::forward(args)...); + }); } + private: + static const VTable* GlobalVTable() { + static const VTable table = [] { + VTable table; + InitVTable(&table); + table.Finalize(); + return table; + }(); + return &table; + } + + const VTable* const vtable_; +}; + private: static const VTable* GlobalVTable() { static const VTable table = [] { @@ -197,6 +185,9 @@ class StmtFunctor { */ class TVM_DLL StmtExprVisitor : public tvm::ExprVisitor { public: + using tvm::ExprVisitor::VTable; + // Register dialect hooks during library initialization, before first table use. + static void RegisterExtension(void (*init)(VTable*)); TVM_DEFINE_OBJECT_FUNCTOR_DEFAULT_CONSTRUCTOR(StmtExprVisitor, tvm::ExprVisitor) using tvm::ExprVisitor::Visit; @@ -216,8 +207,6 @@ class TVM_DLL StmtExprVisitor : public tvm::ExprVisitor { virtual ffi::Optional Visit_(const AssertStmtNode* op); virtual ffi::Optional Visit_(const SeqStmtNode* op); virtual ffi::Optional Visit_(const EvaluateNode* op); - virtual ffi::Optional Visit_(const SBlockNode* op); - virtual ffi::Optional Visit_(const SBlockRealizeNode* op); virtual ffi::Optional Visit_(const ScopeIdDefStmtNode* op); virtual ffi::Optional Visit_(const TilePrimitiveCallNode* op); @@ -233,10 +222,10 @@ class TVM_DLL StmtExprVisitor : public tvm::ExprVisitor { ffi::Optional Visit_(const prim::BroadcastNode* op) override; ffi::Optional 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 VisitBufferMetadata(const BufferVar& buffer); + protected: explicit StmtExprVisitor(const VTable* vtable) : tvm::ExprVisitor(vtable) {} static void InitVTable(VTable* vtable); }; @@ -251,6 +240,9 @@ class TVM_DLL StmtExprVisitor : public tvm::ExprVisitor { */ class TVM_DLL StmtExprMutator : public tvm::ExprMutator { public: + using tvm::ExprMutator::VTable; + // Register dialect hooks during library initialization, before first table use. + static void RegisterExtension(void (*init)(VTable*)); TVM_DEFINE_OBJECT_FUNCTOR_DEFAULT_CONSTRUCTOR(StmtExprMutator, tvm::ExprMutator) using tvm::ExprMutator::Mutate; using tvm::ExprMutator::Mutate_; @@ -282,8 +274,6 @@ class TVM_DLL StmtExprMutator : public tvm::ExprMutator { virtual UnchangedOr Mutate_(const AssertStmtNode* op, InplaceMode inplace_mode); virtual UnchangedOr Mutate_(const SeqStmtNode* op, InplaceMode inplace_mode); virtual UnchangedOr Mutate_(const EvaluateNode* op, InplaceMode inplace_mode); - virtual UnchangedOr Mutate_(const SBlockNode* op, InplaceMode inplace_mode); - virtual UnchangedOr Mutate_(const SBlockRealizeNode* op, InplaceMode inplace_mode); virtual UnchangedOr Mutate_(const ScopeIdDefStmtNode* op, InplaceMode inplace_mode); virtual UnchangedOr Mutate_(const TilePrimitiveCallNode* op, InplaceMode inplace_mode); @@ -334,14 +324,13 @@ bool ContainsNode(const Stmt& stmt) { if (contains_node || value.as()) { return std::nullopt; } + if (value.as()) { + contains_node = true; + return std::nullopt; + } return StmtExprVisitor::Visit(value); } - ffi::Optional Visit_(const Node* block) override { - contains_node = true; - return std::nullopt; - } - bool contains_node{false}; }; diff --git a/python/tvm/ir/json_compact.py b/python/tvm/ir/json_compact.py index b1e547887df9..63d882d8b39f 100644 --- a/python/tvm/ir/json_compact.py +++ b/python/tvm/ir/json_compact.py @@ -19,6 +19,9 @@ import json _PRIM_TYPE_KEY_RENAMES = { + "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", diff --git a/python/tvm/relax/analysis/analysis.py b/python/tvm/relax/analysis/analysis.py index 7f0a6c2c3950..c4ef71493ae2 100644 --- a/python/tvm/relax/analysis/analysis.py +++ b/python/tvm/relax/analysis/analysis.py @@ -29,7 +29,8 @@ 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 1f8beb1faa9d..6586b90f4d35 100644 --- a/python/tvm/s_tir/__init__.py +++ b/python/tvm/s_tir/__init__.py @@ -19,6 +19,7 @@ """S-TIR namespace for scheduable TensorIR""" from .function 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/analysis/__init__.py b/python/tvm/s_tir/analysis/__init__.py index 07021f092d95..0961ebcafad4 100644 --- a/python/tvm/s_tir/analysis/__init__.py +++ b/python/tvm/s_tir/analysis/__init__.py @@ -23,7 +23,8 @@ import tvm from tvm.ir import IRModule, TensorRegion from tvm.tirx.expr import Var -from tvm.tirx.stmt import SBlock +from tvm.tirx.stmt import BufferRegion +from tvm.s_tir import SBlock from tvm.tirx import Buffer, Stmt from tvm.tirx.function import PrimFunc @@ -38,7 +39,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 +64,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] diff --git a/python/tvm/s_tir/dlight/analysis/common_analysis.py b/python/tvm/s_tir/dlight/analysis/common_analysis.py index 254c30ee5fcd..70d14131a306 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_tir_var(expr: 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 6c31c2057f0a..33a83b38c7ca 100644 --- a/python/tvm/s_tir/dlight/analysis/gemv.py +++ b/python/tvm/s_tir/dlight/analysis/gemv.py @@ -28,7 +28,7 @@ ) -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 c997e3235510..2732fef989f7 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 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 48141e182fed..b4bac8a2c790 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 get_access_axes(region: list[Range]) -> set[Var]: 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 is_spatial(block: SBlockRV) -> 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 ff22cdfe98fe..6c010ade697c 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 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 82888fdcf78e..8245c6930a54 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/sblock_dependence_info.py b/python/tvm/s_tir/sblock_dependence_info.py index 0376d2ca4560..9d6bd542c622 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.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 bd5bb369faa7..1b37b8c28808 100644 --- a/python/tvm/s_tir/sblock_scope.py +++ b/python/tvm/s_tir/sblock_scope.py @@ -22,7 +22,8 @@ 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 cf3b6c613e4d..42ac24df3dce 100644 --- a/python/tvm/s_tir/schedule/schedule.py +++ b/python/tvm/s_tir/schedule/schedule.py @@ -25,7 +25,8 @@ 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 diff --git a/python/tvm/s_tir/schedule/state.py b/python/tvm/s_tir/schedule/state.py index 6702930faa1b..98396fef6b75 100644 --- a/python/tvm/s_tir/schedule/state.py +++ b/python/tvm/s_tir/schedule/state.py @@ -24,7 +24,8 @@ 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 000000000000..329dc3c6e219 --- /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 +from tvm.runtime import Object, Scriptable, const +from tvm.tirx.buffer import Buffer +from tvm.tirx.expr import IterVar +from tvm.tirx.stmt import BufferRegion, 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 : BufferRegion + The region of source buffer + """ + + buffer: Buffer + source: BufferRegion + + def __init__(self, buffer: Buffer, source: BufferRegion) -> 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[BufferRegion] + The read buffer regions of the block. + + writes: List[BufferRegion] + 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[BufferRegion] + writes: list[BufferRegion] + 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[BufferRegion], + writes: list[BufferRegion], + 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/tirx/__init__.py b/python/tvm/tirx/__init__.py index 051a26d9362b..e5a08aed0cf8 100644 --- a/python/tvm/tirx/__init__.py +++ b/python/tvm/tirx/__init__.py @@ -52,7 +52,7 @@ 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 diff --git a/python/tvm/tirx/stmt.py b/python/tvm/tirx/stmt.py index 5d15b227af16..2787a7b2b897 100644 --- a/python/tvm/tirx/stmt.py +++ b/python/tvm/tirx/stmt.py @@ -33,8 +33,8 @@ import tvm_ffi -from tvm.ir import Expr, Range, Span, TensorRegion, Type -from tvm.runtime import Object, Scriptable, const +from tvm.ir import Expr, Range, Span, Type +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 92cf537bee6c..38af950dd20b 100644 --- a/src/relax/analysis/layout_transformation.cc +++ b/src/relax/analysis/layout_transformation.cc @@ -28,9 +28,10 @@ #include #include #include +#include +#include #include #include -#include namespace tvm { namespace relax { @@ -59,7 +60,7 @@ static bool IsBijectiveAffine(const IndexMap& m, const ffi::Array& 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 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 VisitIterMark(const arith::IterMark& op) { @@ -323,9 +324,9 @@ static ffi::Optional 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& transformation_cache, IndexMap write_transformation) : can_transform_block_(true), @@ -482,7 +483,7 @@ class BlockAnalyzer : public StmtExprVisitor { } } - ffi::Optional Visit_(const SBlockNode* op) final { + ffi::Optional 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 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 spatial_dom_; arith::Analyzer arith_analyzer_; - SBlock block_; + s_tir::SBlock block_; IndexMap block_transformation_; ffi::Map 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 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> GetSuggestedTransforms() { - ffi::Map> result; + ffi::Map> GetSuggestedTransforms() { + ffi::Map> result; for (const auto& [block, index_map] : block_transformations_) { ffi::Map block_transformations; block_transformations.Set(block, index_map); @@ -599,13 +600,13 @@ class PrimFuncAnalyzer : public StmtExprVisitor { } private: - ffi::Optional Visit_(const SBlockNode* op) final { + ffi::Optional 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(op); + s_tir::SBlock block = ffi::GetRef(op); // Get block write buffer transformation. if (block->writes.size() != 1) return std::nullopt; auto write_buffer = block->writes[0]->source.as_or_throw(); @@ -632,12 +633,12 @@ class PrimFuncAnalyzer : public StmtExprVisitor { private: ffi::Map buffer_transformation_cache_; - ffi::Map block_transformations_; - std::unordered_map, ffi::ObjectPtrHash, ffi::ObjectPtrEqual> + ffi::Map block_transformations_; + std::unordered_map, ffi::ObjectPtrHash, ffi::ObjectPtrEqual> block_to_buffer_; }; -ffi::Map> SuggestLayoutTransforms( +ffi::Map> SuggestLayoutTransforms( const PrimFunc& prim_func, ffi::Array 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 21a4e45bc655..f49a800eacd6 100644 --- a/src/relax/analysis/tir_op_pattern_kind.cc +++ b/src/relax/analysis/tir_op_pattern_kind.cc @@ -24,11 +24,12 @@ #include #include #include +#include +#include #include #include #include #include -#include #include @@ -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,9 +51,9 @@ class PatternKindAnalyzer : public StmtExprVisitor { } private: - bool IsOutputBlock(const SBlockNode* block) { - for (const TensorRegion& write_region : block->writes) { - if (param_buffers_.count(write_region->source.as_or_throw())) { + bool IsOutputBlock(const s_tir::SBlockNode* block) { + for (const BufferRegion& write_region : block->writes) { + if (param_buffers_.count(write_region->buffer)) { return true; } } @@ -67,25 +68,25 @@ class PatternKindAnalyzer : public StmtExprVisitor { return std::nullopt; } store_ = ffi::GetRef(op); - return StmtExprVisitor::Visit_(op); + return s_tir::StmtExprVisitor::Visit_(op); } ffi::Optional Visit_(const TensorLoadNode* op) final { loads_.push_back(ffi::GetRef(op)); - return StmtExprVisitor::Visit_(op); + return s_tir::StmtExprVisitor::Visit_(op); } - ffi::Optional Visit_(const SBlockNode* op) final { + ffi::Optional 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 Visit(ffi::AnyView value) override { if (value.as()) 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() || loop->body->IsInstance())) { + if (!(loop->body->IsInstance() || + loop->body->IsInstance())) { return std::nullopt; } return this->Visit(loop->body); } - ffi::Optional Visit_(const SBlockRealizeNode* block_realize) final { + ffi::Optional 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& block_iter = block->iter_vars; const ffi::Array& 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 Visit_(const SBlockNode* block) final { + ffi::Optional Visit_(const s_tir::SBlockNode* block) final { // Step 0. If the block body is a ForNode, recurse into it. if (block->body->IsInstance()) { 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()); + TVM_FFI_ICHECK(func->body->IsInstance()); 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 31d7f7f65505..108aa6ffb1e4 100644 --- a/src/relax/backend/task_extraction.cc +++ b/src/relax/backend/task_extraction.cc @@ -22,9 +22,10 @@ #include #include #include +#include +#include #include #include -#include #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 Visit(ffi::AnyView value) override { if (value.as()) 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 Visit_(const tirx::SBlockNode* op) final { + ffi::Optional 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 42efd11af5ff..a7ab6cf621a6 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 #include #include +#include +#include #include -#include #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 buffer_map) { auto replacer = ffi::make_object(buffer_map); @@ -52,26 +53,26 @@ class DistBufferReplacer : public StmtExprMutator { } }; -class DistSBlockInfoCollector : public StmtExprVisitor { +class DistSBlockInfoCollector : public s_tir::StmtExprVisitor { private: ffi::Optional 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 Visit_(const TensorLoadNode* op) final { buffer_access_indices[op->source.as_or_throw()].push_back(op->indices); - return StmtExprVisitor::Visit_(op); + return s_tir::StmtExprVisitor::Visit_(op); } - ffi::Optional Visit_(const SBlockNode* op) final { + ffi::Optional 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(); } } - 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 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 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 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 (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; @@ -183,7 +184,7 @@ class DistributedBufferCompactor : public StmtExprMutator { } ffi::Array ShardIterVar( - SBlock block, + s_tir::SBlock block, const std::unordered_map>, ffi::ObjectPtrHash, ffi::ObjectPtrEqual>& buffer_access_indices) { std::vector buffers; @@ -255,10 +256,10 @@ class DistributedBufferCompactor : public StmtExprMutator { return BufferVar(buffer.name(), std::move(new_type), buffer.span()); } - UnchangedOr Mutate_(const SBlockNode* op, InplaceMode inplace_mode) final { - SBlock block = StmtExprMutator::Mutate_(op, inplace_mode) - .ValueOrUnchanged(ffi::GetRef(op)) - .as_or_throw(); + UnchangedOr Mutate_(const s_tir::SBlockNode* op, InplaceMode inplace_mode) final { + s_tir::SBlock block = s_tir::StmtExprMutator::Mutate_(op, inplace_mode) + .ValueOrUnchanged(ffi::GetRef(op)) + .as_or_throw(); auto collector = ffi::make_object(); collector->Visit(block); ffi::Array new_iter_vars = ShardIterVar(block, collector->buffer_access_indices); @@ -280,7 +281,8 @@ class DistributedBufferCompactor : public StmtExprMutator { break; } } - ffi::ObjectPtr new_block = ffi::make_object(*block.operator->()); + ffi::ObjectPtr new_block = + ffi::make_object(*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 Mutate_(const SBlockRealizeNode* op, InplaceMode inplace_mode) final { - SBlockRealize realize = StmtExprMutator::Mutate_(op, inplace_mode) - .ValueOrUnchanged(ffi::GetRef(op)) - .as_or_throw(); + UnchangedOr Mutate_(const s_tir::SBlockRealizeNode* op, InplaceMode inplace_mode) final { + s_tir::SBlockRealize realize = s_tir::StmtExprMutator::Mutate_(op, inplace_mode) + .ValueOrUnchanged(ffi::GetRef(op)) + .as_or_throw(); for (int i = 0; i < static_cast(realize->iter_values.size()); i++) { PrimExpr iter_value = realize->iter_values[i]; @@ -313,7 +315,7 @@ class DistributedBufferCompactor : public StmtExprMutator { } UnchangedOr 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(op)) .as_or_throw(); 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 6b9eabe19838..e2e9a4287d1f 100644 --- a/src/relax/transform/fuse_tir.cc +++ b/src/relax/transform/fuse_tir.cc @@ -24,9 +24,10 @@ #include #include #include +#include +#include #include #include -#include #include #include @@ -170,7 +171,7 @@ class SymbolicMatcher : ExprFunctor /*! * \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& buffer_map, const ffi::Map& var_map) { @@ -190,12 +191,12 @@ class FuseTIRBufferSubstitutor : public StmtExprMutator { } private: - UnchangedOr Mutate_(const SBlockNode* op, InplaceMode inplace_mode) final { - SBlock block = StmtExprMutator::Mutate_(op, inplace_mode) - .ValueOrUnchanged(ffi::GetRef(op)) - .as_or_throw(); - ffi::Array reads = UnionAccessRegion(block->reads); - ffi::Array writes = UnionAccessRegion(block->writes); + UnchangedOr Mutate_(const s_tir::SBlockNode* op, InplaceMode inplace_mode) final { + s_tir::SBlock block = s_tir::StmtExprMutator::Mutate_(op, inplace_mode) + .ValueOrUnchanged(ffi::GetRef(op)) + .as_or_throw(); + ffi::Array reads = UnionAccessRegion(block->reads); + ffi::Array writes = UnionAccessRegion(block->writes); if (!reads.same_as(block->reads) || !writes.same_as(block->writes)) { auto* n = block.CopyOnWrite(); n->reads = std::move(reads); @@ -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 Mutate(ffi::AnyView value, InplaceMode inplace_mode) final { if (value.as()) return ffi::Unchanged(); - return StmtExprMutator::Mutate(value, inplace_mode); + return s_tir::StmtExprMutator::Mutate(value, inplace_mode); } private: - UnchangedOr Mutate_(const SBlockNode* op, InplaceMode inplace_mode) final { - SBlock block = tirx::StmtExprMutator::Mutate_(op, inplace_mode) - .ValueOrUnchanged(ffi::GetRef(op)) - .as_or_throw(); + UnchangedOr Mutate_(const s_tir::SBlockNode* op, InplaceMode inplace_mode) final { + s_tir::SBlock block = s_tir::StmtExprMutator::Mutate_(op, inplace_mode) + .ValueOrUnchanged(ffi::GetRef(op)) + .as_or_throw(); 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()) + TVM_FFI_ICHECK(prim_func->body->IsInstance()) << "Only schedulable functions (whose body is the root block) can be fused"; - const tirx::SBlockRealize& root_realize = prim_func->body.as_or_throw(); - const tirx::SBlock& root_block = root_realize->block; + const s_tir::SBlockRealize& root_realize = prim_func->body.as_or_throw(); + 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()->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()); + body = s_tir::SBlock({}, {}, {}, "root", std::move(body), std::nullopt, alloc_buffers); + body = s_tir::SBlockRealize({}, IntImm::Bool(true), body.as_or_throw()); ffi::Array 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/split_call_tir_by_pattern.cc b/src/relax/transform/split_call_tir_by_pattern.cc index d179e1efe30c..4d29bd47267f 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 #include #include +#include +#include #include #include #include -#include #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()->block->body.as(); + const ForNode* pattern_top = + pattern_->body.as()->block->body.as(); TVM_FFI_ICHECK(pattern_top) << "Invalid pattern function"; if (!Dispatch(top, ffi::GetRef(pattern_top))) { return false; @@ -254,10 +256,10 @@ class ForMatcher : public TensorizeComparator { loop_stack_lhs_.push_back(ffi::GetRef(op)); loop_stack_rhs_.push_back(ffi::GetRef(rhs)); // The body of loop must be loop or BlockRealize - if (!op->body->IsInstance() && !op->body->IsInstance()) { + if (!op->body->IsInstance() && !op->body->IsInstance()) { return false; } - if (!rhs->body->IsInstance() && !rhs->body->IsInstance()) { + if (!rhs->body->IsInstance() && !rhs->body->IsInstance()) { 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(); + bool VisitStmt_(const s_tir::SBlockNode* op, const Stmt& other) final { + const auto* rhs = other.as(); // 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(); + bool VisitStmt_(const s_tir::SBlockRealizeNode* op, const Stmt& other) final { + const auto* rhs = other.as(); // 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 allocs2; /*! \brief whether the current block is in the first function */ - ffi::Map block_partition; + ffi::Map block_partition; /*! \brief input buffers for the first function */ std::unordered_set input1; /*! \brief input buffers for the second function */ @@ -484,7 +486,7 @@ class FunctionPartitioner : public StmtExprVisitor { bool fail = false; private: - ffi::Optional Visit_(const SBlockNode* op) final { + ffi::Optional 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()); } } - block_partition.Set(ffi::GetRef(op), is_matching_); + block_partition.Set(ffi::GetRef(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& block_partition, + Stmt stmt, const ffi::Map& block_partition, const std::unordered_set& allocs, bool is_library_part) { auto remover = ffi::make_object(block_partition, allocs, is_library_part); return remover->Mutate(stmt, InplaceMode::kDisallow).ValueOrUnchanged(stmt); } - BlockRemover(const ffi::Map& block_partition, + BlockRemover(const ffi::Map& block_partition, const std::unordered_set& allocs, bool is_library_part) : block_partition(block_partition), allocs_(allocs), is_library_part_(is_library_part) {} private: - UnchangedOr Mutate_(const SBlockNode* op, InplaceMode inplace_mode) final { - SBlock block = StmtExprMutator::Mutate_(op, inplace_mode) - .ValueOrUnchanged(ffi::GetRef(op)) - .as_or_throw(); - ffi::ObjectPtr n = ffi::make_object(*block.operator->()); + UnchangedOr Mutate_(const s_tir::SBlockNode* op, InplaceMode inplace_mode) final { + s_tir::SBlock block = s_tir::StmtExprMutator::Mutate_(op, inplace_mode) + .ValueOrUnchanged(ffi::GetRef(op)) + .as_or_throw(); + ffi::ObjectPtr n = ffi::make_object(*block.operator->()); if (op->name_hint != "root") { - TVM_FFI_ICHECK(block_partition.count(ffi::GetRef(op))); - bool block_is_library = block_partition[ffi::GetRef(op)]; + TVM_FFI_ICHECK(block_partition.count(ffi::GetRef(op))); + bool block_is_library = block_partition[ffi::GetRef(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 Mutate_(const SeqStmtNode* op, InplaceMode inplace_mode) final { @@ -576,7 +578,7 @@ class BlockRemover : public StmtExprMutator { } bool erased_ = false; - ffi::Map block_partition; + ffi::Map block_partition; std::unordered_set allocs_; bool is_library_part_ = false; }; @@ -595,9 +597,9 @@ std::pair> SplitFunctions( PrimFunc func, std::vector>* arg_partition, ffi::Array patterns, FCodegen f_codegen) { // Step 1. Find the library kernel and the rest. - Stmt body = func->body.as()->block->body; + Stmt body = func->body.as()->block->body; ffi::Array match_results = - TIRPatternMatcher::Match(patterns, func->body.as()->block->body); + TIRPatternMatcher::Match(patterns, func->body.as()->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 845900ecc5fb..cd08d7d24e17 100644 --- a/src/relax/transform/split_layout_rewrite_preproc.cc +++ b/src/relax/transform/split_layout_rewrite_preproc.cc @@ -26,8 +26,8 @@ #include #include #include +#include #include -#include #include #include @@ -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 Mutate(ffi::AnyView value, InplaceMode inplace_mode) final { if (value.as()) 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, PrimFunc> Transform(const PrimFunc& func) { - TVM_FFI_ICHECK(func->body.as()) + TVM_FFI_ICHECK(func->body.as()) << "The body of the primfunc should be a root block."; - const auto& block = func->body.as()->block; + const auto& block = func->body.as()->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(), /*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 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()->block; + s_tir::SBlock original_block = original_func_->body.as()->block; ffi::Array 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(), /*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 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()) { 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 Mutate_(const SBlockNode* op, InplaceMode inplace_mode) final { - SBlock block = StmtExprMutator::Mutate_(op, inplace_mode) - .ValueOrUnchanged(ffi::GetRef(op)) - .as_or_throw(); + UnchangedOr Mutate_(const s_tir::SBlockNode* op, InplaceMode inplace_mode) final { + s_tir::SBlock block = s_tir::StmtExprMutator::Mutate_(op, inplace_mode) + .ValueOrUnchanged(ffi::GetRef(op)) + .as_or_throw(); 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()); @@ -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(*block.get()); + auto n = ffi::make_object(*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 a5245595e8a7..c96a9c26667b 100644 --- a/src/s_tir/analysis/calculate_allocated_memory.cc +++ b/src/s_tir/analysis/calculate_allocated_memory.cc @@ -25,10 +25,11 @@ #include #include #include +#include +#include #include #include #include -#include #include #include diff --git a/src/s_tir/analysis/conditional_bounds.cc b/src/s_tir/analysis/conditional_bounds.cc index 20129a1411bd..448171b4da16 100644 --- a/src/s_tir/analysis/conditional_bounds.cc +++ b/src/s_tir/analysis/conditional_bounds.cc @@ -28,6 +28,7 @@ #include #include #include +#include #include #include diff --git a/src/s_tir/analysis/domain_touched.cc b/src/s_tir/analysis/domain_touched.cc index 2e67be88ea6b..bbd8c4430da0 100644 --- a/src/s_tir/analysis/domain_touched.cc +++ b/src/s_tir/analysis/domain_touched.cc @@ -27,14 +27,14 @@ #include #include #include +#include #include -#include #include #include #include -#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; } // 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& 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 Visit_(const TensorLoadNode* op) final { BufferVar buffer = op->source.as_or_throw(); diff --git a/src/s_tir/analysis/estimate_flops.cc b/src/s_tir/analysis/estimate_flops.cc index 0da281481918..3665f5006d48 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 +#include +#include +#include #include -#include #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 666f4f4e2b42..dc4f76210262 100644 --- a/src/s_tir/analysis/find_anchor_sblock.cc +++ b/src/s_tir/analysis/find_anchor_sblock.cc @@ -25,19 +25,20 @@ #include #include #include +#include +#include #include -#include 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 Visit(ffi::AnyView value) override { if (value.as()) return std::nullopt; - return StmtExprVisitor::Visit(value); + return s_tir::StmtExprVisitor::Visit(value); } ffi::Optional 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 Visit(ffi::AnyView value) override { if (value.as()) 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 Visit_(const SBlockNode* block) override { + ffi::Optional 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(block); + << ffi::GetRef(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 Visit(ffi::AnyView value) override { if (value.as()) return std::nullopt; - return StmtExprVisitor::Visit(value); + return s_tir::StmtExprVisitor::Visit(value); } - ffi::Optional Visit_(const SBlockNode* block) override { + ffi::Optional 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 blocks; + std::vector 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(ffi::GetRef(ret)); + return ffi::Optional(ffi::GetRef(ret)); } - return ffi::Optional(std::nullopt); + return ffi::Optional(std::nullopt); }); } diff --git a/src/s_tir/analysis/identify_memcpy.cc b/src/s_tir/analysis/identify_memcpy.cc index d80c957ee8be..02cfdbe865e3 100644 --- a/src/s_tir/analysis/identify_memcpy.cc +++ b/src/s_tir/analysis/identify_memcpy.cc @@ -38,7 +38,7 @@ #include #include -#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 output; - struct Visitor : tirx::IRVisitorWithAnalyzer { + struct Visitor : s_tir::IRVisitorWithAnalyzer { public: - using tirx::IRVisitorWithAnalyzer::Visit_; + using s_tir::IRVisitorWithAnalyzer::Visit_; explicit Visitor(ffi::Array* output) : output(output) {} ffi::Array* output; diff --git a/src/s_tir/analysis/is_pure_function.cc b/src/s_tir/analysis/is_pure_function.cc index 3c7acb10e60b..47a7f3deec2a 100644 --- a/src/s_tir/analysis/is_pure_function.cc +++ b/src/s_tir/analysis/is_pure_function.cc @@ -24,8 +24,9 @@ #include #include #include +#include +#include #include -#include #include "../../tirx/ir/tir_visitor_with_path.h" diff --git a/src/s_tir/analysis/oob_checker.cc b/src/s_tir/analysis/oob_checker.cc index 55bf73e3443b..42e3161b5113 100644 --- a/src/s_tir/analysis/oob_checker.cc +++ b/src/s_tir/analysis/oob_checker.cc @@ -24,7 +24,7 @@ #include #include -#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 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 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 574401da95d5..2331078005cd 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 #include #include +#include +#include #include -#include #include #include #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& 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 buffer_var_map_; /*! \brief The target buffer var mapping to its matching */ - std::unordered_map match_buffers_; + std::unordered_map match_buffers_; /*! \brief let bindings inside the block */ std::unordered_map let_bindings_; /*!\ brief Internal analyzer. */ @@ -115,7 +117,7 @@ class BlockReadWriteDetector : public StmtExprVisitor { const std::unordered_set* excluded_buffers = nullptr); /*! \brief Helper function to convert matched access region to source region. */ - std::vector ConvertMatchedRegion(const MatchBufferRegion& match_buffer, + std::vector ConvertMatchedRegion(const s_tir::MatchBufferRegion& match_buffer, const std::vector& int_sets) const; /*! \brief Helper function to update a opaque access. */ @@ -136,7 +138,7 @@ class BlockReadWriteDetector : public StmtExprVisitor { ffi::Optional Visit_(const ForNode* op) override; ffi::Optional Visit_(const IfThenElseNode* op) override; - ffi::Optional Visit_(const SBlockRealizeNode* op) override; + ffi::Optional Visit_(const s_tir::SBlockRealizeNode* op) override; ffi::Optional Visit_(const DeclBufferNode* op) override; ffi::Optional Visit_(const BufferStoreNode* op) override; ffi::Optional Visit_(const BindNode* op) override; @@ -146,10 +148,10 @@ class BlockReadWriteDetector : public StmtExprVisitor { }; void BlockReadWriteDetector::operator()(const Stmt& stmt) { - const auto* block = stmt.as(); + const auto* block = stmt.as(); 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().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 BlockReadWriteDetector::CollectReads( @@ -209,7 +211,7 @@ ffi::Optional BlockReadWriteDetector::Visit_(const TensorLoadNod ffi::Optional 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 BlockReadWriteDetector::Visit_(const IfThenElseNod // Visit then branch With 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 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 BlockReadWriteDetector::Visit_(const BindNode* op) if (auto value = op->value.as()) { let_bindings_[op->var.get()] = value.value(); } - return StmtExprVisitor::Visit_(op); + return s_tir::StmtExprVisitor::Visit_(op); } ffi::Optional BlockReadWriteDetector::Visit_(const CallNode* op) { @@ -312,7 +314,7 @@ ffi::Optional 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 BlockReadWriteDetector::Visit_(const CallNode* op) With ctx(condition, &dom_map_, &hint_map_, &pending_conditions_); TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN( - StmtExprVisitor::Visit(op->args[1].as_or_throw())); + s_tir::StmtExprVisitor::Visit(op->args[1].as_or_throw())); } { // Visit else branch With ctx(!condition, &dom_map_, &hint_map_, &pending_conditions_); TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN( - StmtExprVisitor::Visit(op->args[2].as_or_throw())); + s_tir::StmtExprVisitor::Visit(op->args[2].as_or_throw())); } return std::nullopt; } - return StmtExprVisitor::Visit_(op); + return s_tir::StmtExprVisitor::Visit_(op); } ffi::Optional BlockReadWriteDetector::Visit_(const BufferStoreNode* op) { @@ -364,7 +366,7 @@ ffi::Optional BlockReadWriteDetector::Visit_(const BufferStoreNo return std::nullopt; } -ffi::Optional BlockReadWriteDetector::Visit_(const SBlockRealizeNode* op) { +ffi::Optional BlockReadWriteDetector::Visit_(const s_tir::SBlockRealizeNode* op) { /*! \note detector will not visit child block recursively, so it will stop here */ std::unordered_map vmap; for (size_t i = 0; i < op->block->iter_vars.size(); ++i) { @@ -404,7 +406,8 @@ ffi::Optional BlockReadWriteDetector::Visit_(const SBlockRealize } std::vector BlockReadWriteDetector::ConvertMatchedRegion( - const MatchBufferRegion& match_buffer, const std::vector& int_sets) const { + const s_tir::MatchBufferRegion& match_buffer, + const std::vector& int_sets) const { const BufferVar& buffer = match_buffer->buffer; ffi::Array region; @@ -432,8 +435,8 @@ void BlockReadWriteDetector::Update(std::vector* buffers, // Handle match_buffer remap auto it = match_buffers_.find(buffer.get()); if (it != match_buffers_.end()) { - const MatchBufferRegion& match_buffer = it->second; - buffer = match_buffer->source->source.as_or_throw(); + const s_tir::MatchBufferRegion& match_buffer = it->second; + buffer = match_buffer->source->buffer; region = ConvertMatchedRegion(match_buffer, std::move(region)); } TVM_FFI_ICHECK_EQ(buffers->size(), regions->size()) @@ -494,8 +497,8 @@ void BlockReadWriteDetector::UpdateOpaque(const Var& buffer_var) { } } -ffi::Array> GetSBlockAccessRegion( - const SBlock& block, const ffi::Map& buffer_var_map) { +ffi::Array> GetSBlockAccessRegion( + const s_tir::SBlock& block, const ffi::Map& buffer_var_map) { auto detector = ffi::make_object(buffer_var_map); detector->operator()(block); ffi::Array writes = detector->CollectWrites(); @@ -511,8 +514,8 @@ ffi::Array> GetSBlockAccessRegion( return {reads, writes, opaques}; } -ffi::Array> GetSBlockReadWriteRegion( - const SBlock& block, const ffi::Map& buffer_var_map) { +ffi::Array> GetSBlockReadWriteRegion( + const s_tir::SBlock& block, const ffi::Map& buffer_var_map) { auto detector = ffi::make_object(buffer_var_map); detector->operator()(block); ffi::Array 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 568a330dd5a9..9ec400dbc7f0 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 #include #include +#include +#include +#include #include -#include #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> Detect(const PrimFunc& func) { auto detector = ffi::make_object(); @@ -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 Visit_(const SBlockRealizeNode* op) final { - const SBlockNode* block = op->block.get(); + ffi::Optional 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,17 @@ class LCADetector : public StmtExprVisitor { UpdateDominateScopeOfNonDataParIter(op); // Update match_buffers - for (const MatchBufferRegion& match_buffer : block->match_buffers) { - UpdateBufferLCA(match_buffer->source->source.as_or_throw().get(), - ancestor_scopes_.back()); + for (const s_tir::MatchBufferRegion& match_buffer : block->match_buffers) { + UpdateBufferLCA(match_buffer->source->buffer.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 opaque_var_scope; // maintain highest scope which dominate all reduce loop iters. null denotes non-reduce block. @@ -179,7 +180,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,7 +264,7 @@ 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. @@ -354,7 +355,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 17f068c75197..b56756ad2380 100644 --- a/src/s_tir/analysis/verify_gpu_code.cc +++ b/src/s_tir/analysis/verify_gpu_code.cc @@ -27,10 +27,11 @@ #include #include #include +#include #include +#include #include #include -#include #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 000000000000..f98b45eb0387 --- /dev/null +++ b/src/s_tir/analysis/verify_well_formed.cc @@ -0,0 +1,138 @@ +/* + * 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 + +namespace tvm { +namespace s_tir { +using tirx::BufferRegion; +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(assert_mode); + verifier->Visit(func->body); + return !verifier->has_error_; + } + + explicit BlockVarAccessVerifier(bool assert_mode) : assert_mode_(assert_mode) {} + + private: + ffi::Optional Visit(ffi::AnyView stmt) final { + if (!has_error_) { + return StmtExprVisitor::Visit(stmt); + } + return std::nullopt; + } + + ffi::Optional 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 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 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 BufferRegion& s) -> ffi::Optional { + 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 loop_vars_; + /*! \brief Whether it's in assert mode. */ + bool assert_mode_; + /*! \brief Current nested block stack level. */ + std::vector block_stack_; + /*! \brief Whether there is error. */ + bool has_error_{false}; +}; + +TVM_FFI_STATIC_INIT_BLOCK() { tirx::RegisterWellFormedExtension(BlockVarAccessVerifier::Verify); } +} // 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 fa3c2925f60e..e9403acaebc7 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 +#include #include +#include #include -#include #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 d6abeef17fcb..a03f6ced7ae8 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 19819d1c879c..9000d298e586 100644 --- a/src/s_tir/data_layout.cc +++ b/src/s_tir/data_layout.cc @@ -29,9 +29,10 @@ #include #include #include +#include #include +#include #include -#include #include #include 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 000000000000..4c1cefdcc939 --- /dev/null +++ b/src/s_tir/ir/data_type_rewriter.cc @@ -0,0 +1,307 @@ +/* + * 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/ir/data_type_rewriter.h" + +#include +#include + +#include + +namespace tvm { +namespace tirx { +using namespace tvm::prim; +using s_tir::MatchBufferRegion; +using s_tir::SBlock; +using s_tir::SBlockNode; +using s_tir::SBlockRealize; +using s_tir::SBlockRealizeNode; + +class DataTypeLegalizer::Extension { + public: + static void InitVTable(VTable* vtable); + static UnchangedOr MutateBlockRealize(DataTypeLegalizer* self, + const s_tir::SBlockRealizeNode* op, + InplaceMode inplace_mode); + static UnchangedOr MutateBlock(DataTypeLegalizer* self, const s_tir::SBlockNode* op, + InplaceMode inplace_mode); +}; +class IndexDataTypeRewriter::Extension { + public: + static void InitVTable(VTable* vtable); + static UnchangedOr MutateBlockRealize(IndexDataTypeRewriter* self, + const s_tir::SBlockRealizeNode* op, + InplaceMode inplace_mode); + static UnchangedOr MutateBlock(IndexDataTypeRewriter* self, const s_tir::SBlockNode* op, + InplaceMode inplace_mode); + static ffi::Map VisitBlockAnnotations( + IndexDataTypeRewriter* self, const ffi::Map& annotations); + static IterVar VisitIterVar(IndexDataTypeRewriter* self, const IterVar& iter_var); + static BufferRegion VisitBufferRegion(IndexDataTypeRewriter* self, + const BufferRegion& buffer_region); +}; + +UnchangedOr DataTypeLegalizer::Extension::MutateBlockRealize(DataTypeLegalizer* self, + const SBlockRealizeNode* op, + InplaceMode inplace_mode) { + SBlockRealize realize = s_tir::StmtExprMutator::MutateBlockRealize(self, op, inplace_mode) + .ValueOrUnchanged(ffi::GetRef(op)) + .as_or_throw(); + ffi::Array new_iter_values; + bool changed = false; + for (int i = 0; i < static_cast(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 DataTypeLegalizer::Extension::MutateBlock(DataTypeLegalizer* self, + const SBlockNode* op, + InplaceMode inplace_mode) { + SBlock new_block = s_tir::StmtExprMutator::MutateBlock(self, op, inplace_mode) + .ValueOrUnchanged(ffi::GetRef(op)) + .as_or_throw(); + ffi::Array 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; +} + +void DataTypeLegalizer::Extension::InitVTable(VTable* vtable) { + vtable->ClearDispatch(); + vtable->ClearDispatch(); + vtable->SetDispatch( + [](const ffi::Object* node, ObjectMutator* base, InplaceMode mode) -> UnchangedOr { + return MutateBlock(static_cast(base), + static_cast(node), mode); + }); + vtable->SetDispatch( + [](const ffi::Object* node, ObjectMutator* base, InplaceMode mode) -> UnchangedOr { + return MutateBlockRealize(static_cast(base), + static_cast(node), mode); + }); +} +TVM_FFI_STATIC_INIT_BLOCK() { + DataTypeLegalizer::RegisterExtension(DataTypeLegalizer::Extension::InitVTable); +} + +UnchangedOr IndexDataTypeRewriter::Extension::MutateBlockRealize(IndexDataTypeRewriter* self, + const SBlockRealizeNode* op, + InplaceMode inplace_mode) { + bool is_condition = self->is_condition_; + self->is_condition_ = true; + auto new_predicate_result = self->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); + self->is_condition_ = is_condition; + + bool is_enabled = self->is_enabled_; + self->is_enabled_ = true; + auto new_iter_values = self->Mutate(op->iter_values, inplace_mode) + .as_or_throw>>() + .ValueOrUnchanged(op->iter_values); + self->is_enabled_ = is_enabled; + SBlock new_body = + self->Mutate(op->block, inplace_mode).ValueOrUnchanged(op->block).as_or_throw(); + if (!new_predicate_unchanged || !new_iter_values.same_as(op->iter_values) || + !new_body.same_as(op->block)) { + SBlockRealize new_block_realize = ffi::GetRef(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 IndexDataTypeRewriter::Extension::MutateBlock(IndexDataTypeRewriter* self, + const SBlockNode* op, + InplaceMode inplace_mode) { + auto new_alloc_buffers = self->WithDefRegionKind(kTVMFFIDefRegionKindSimple, [&] { + return self->Mutate(op->alloc_buffers, inplace_mode) + .as_or_throw>>() + .ValueOrUnchanged(op->alloc_buffers); + }); + auto new_match_buffers = op->match_buffers.Map([self](const MatchBufferRegion& match) { + BufferVar buffer = self->WithDefRegionKind(kTVMFFIDefRegionKindSimple, [&] { + return self->Mutate(match->buffer, InplaceMode::kDisallow) + .as_or_throw>() + .ValueOrUnchanged(match->buffer); + }); + BufferRegion source = VisitBufferRegion(self, match->source); + if (buffer.same_as(match->buffer) && source.same_as(match->source)) return match; + return MatchBufferRegion(buffer, source); + }); + ffi::Array new_reads = op->reads.Map( + [self](const BufferRegion& buffer_region) { return VisitBufferRegion(self, buffer_region); }); + ffi::Array new_writes = op->writes.Map( + [self](const BufferRegion& buffer_region) { return VisitBufferRegion(self, buffer_region); }); + ffi::Array new_iter_vars = + op->iter_vars.Map([self](const IterVar& iter_var) { return VisitIterVar(self, iter_var); }); + ffi::Optional new_init = std::nullopt; + if (op->init.has_value()) { + new_init = self->Mutate(op->init.value(), inplace_mode).ValueOrUnchanged(op->init.value()); + } + ffi::Map new_annotations = VisitBlockAnnotations(self, op->annotations); + auto new_body_result = self->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(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); + for (const auto& buffer : new_block->alloc_buffers) self->ValidateAllocation(buffer); + return new_block; + } + for (const auto& buffer : op->alloc_buffers) self->ValidateAllocation(buffer); + return ffi::Unchanged(); +} + +ffi::Map IndexDataTypeRewriter::Extension::VisitBlockAnnotations( + IndexDataTypeRewriter* self, const ffi::Map& annotations) { + auto new_annotations = annotations; + + std::function f_mutate_obj = [self, &f_mutate_obj](const Any& obj) -> Any { + if (obj == nullptr) { + return obj; + } + if (auto var = obj.as(); var && var.value()->ty.as()) { + BufferVar buffer(var.value()); + if (BufferVar new_buffer = self->Mutate(buffer, InplaceMode::kDisallow) + .as_or_throw>() + .ValueOrUnchanged(buffer); + !new_buffer.same_as(buffer)) { + return new_buffer; + } + } else if (obj.as()) { + return obj.as_or_throw>().Map(f_mutate_obj); + } + return obj; + }; + for (const auto& [key, value] : annotations) { + if (auto opt_object_ref = value.as()) { + 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::Extension::VisitIterVar(IndexDataTypeRewriter* self, + const IterVar& iter_var) { + bool is_enabled = self->is_enabled_; + self->is_enabled_ = true; + PrimVar new_var = self->Mutate(iter_var->var, InplaceMode::kDisallow) + .ValueOrUnchanged(iter_var->var) + .as_or_throw(); + PrimExpr min = + self->Mutate(iter_var->dom->min, InplaceMode::kDisallow).ValueOrUnchanged(iter_var->dom->min); + PrimExpr extent = self->Mutate(iter_var->dom->extent, InplaceMode::kDisallow) + .ValueOrUnchanged(iter_var->dom->extent); + self->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; +} + +BufferRegion IndexDataTypeRewriter::Extension::VisitBufferRegion( + IndexDataTypeRewriter* self, const BufferRegion& buffer_region) { + BufferVar remapped_buffer = self->Mutate(buffer_region->buffer, InplaceMode::kDisallow) + .as_or_throw>() + .ValueOrUnchanged(buffer_region->buffer); + + bool is_enabled = self->is_enabled_; + self->is_enabled_ = true; + auto new_region = buffer_region->region.Map([&](const Range& range) { + return Range::FromMinExtent( + self->Mutate(range->min, InplaceMode::kDisallow).ValueOrUnchanged(range->min), + self->Mutate(range->extent, InplaceMode::kDisallow).ValueOrUnchanged(range->extent)); + }); + self->is_enabled_ = is_enabled; + + if (!remapped_buffer.same_as(buffer_region->buffer) || + !new_region.same_as(buffer_region->region)) { + return BufferRegion(remapped_buffer, new_region); + } else { + return buffer_region; + } +} + +void IndexDataTypeRewriter::Extension::InitVTable(VTable* vtable) { + vtable->ClearDispatch(); + vtable->SetDispatch( + [](const ffi::Object* node, ObjectMutator* base, InplaceMode mode) -> UnchangedOr { + return MutateBlock(static_cast(base), + static_cast(node), mode); + }); + vtable->ClearDispatch(); + vtable->SetDispatch( + [](const ffi::Object* node, ObjectMutator* base, InplaceMode mode) -> UnchangedOr { + return MutateBlockRealize(static_cast(base), + static_cast(node), mode); + }); +} +TVM_FFI_STATIC_INIT_BLOCK() { + IndexDataTypeRewriter::RegisterExtension(IndexDataTypeRewriter::Extension::InitVTable); +} +} // namespace tirx +} // namespace tvm diff --git a/src/s_tir/ir/ir_mutator_with_analyzer.cc b/src/s_tir/ir/ir_mutator_with_analyzer.cc new file mode 100644 index 000000000000..f97f340ece73 --- /dev/null +++ b/src/s_tir/ir/ir_mutator_with_analyzer.cc @@ -0,0 +1,48 @@ +/* + * 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_mutator_with_analyzer.h" + +namespace tvm { +namespace tirx { +UnchangedOr IRMutatorWithAnalyzer::Extension::MutateBlock(IRMutatorWithAnalyzer* self, + const s_tir::SBlockNode* op, + InplaceMode inplace_mode) { + return self->constraint_scope_.WithNewScope([&]() -> UnchangedOr { + for (const auto& iter_var : op->iter_vars) { + self->analyzer_->Bind(iter_var->var, iter_var->dom); + self->iter_vars_.Set(iter_var->var, iter_var->dom); + } + return s_tir::StmtExprMutator::MutateBlock(self, op, inplace_mode); + }); +} + +void IRMutatorWithAnalyzer::Extension::InitVTable(VTable* vtable) { + vtable->ClearDispatch(); + vtable->SetDispatch( + [](const ffi::Object* node, ObjectMutator* base, InplaceMode mode) -> UnchangedOr { + return MutateBlock(static_cast(base), + static_cast(node), mode); + }); +} +TVM_FFI_STATIC_INIT_BLOCK() { + IRMutatorWithAnalyzer::RegisterExtension(IRMutatorWithAnalyzer::Extension::InitVTable); +} +} // namespace tirx +} // namespace tvm 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 000000000000..12a0a4672c31 --- /dev/null +++ b/src/s_tir/ir/ir_mutator_with_analyzer.h @@ -0,0 +1,67 @@ +/* + * 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 + +#include "../../tirx/ir_mutator_with_analyzer.h" + +namespace tvm { +namespace tirx { +class IRMutatorWithAnalyzer::Extension { + public: + static UnchangedOr MutateBlock(IRMutatorWithAnalyzer* self, const s_tir::SBlockNode* op, + InplaceMode inplace_mode); + static void InitVTable(VTable* vtable); +}; +} // namespace tirx +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 Mutate_(const SBlockNode* op, InplaceMode inplace_mode) { + return Parent::Extension::MutateBlock(this, op, inplace_mode); + } + + protected: + static void InitVTable(VTable* vtable) { + Parent::InitVTable(vtable); + vtable->ClearDispatch(); + SetDispatch(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 diff --git a/src/s_tir/ir/ir_visitor_with_analyzer.cc b/src/s_tir/ir/ir_visitor_with_analyzer.cc new file mode 100644 index 000000000000..de5e01fe7095 --- /dev/null +++ b/src/s_tir/ir/ir_visitor_with_analyzer.cc @@ -0,0 +1,45 @@ +/* + * 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_visitor_with_analyzer.h" + +namespace tvm { +namespace tirx { +ffi::Optional IRVisitorWithAnalyzer::Extension::VisitBlock( + IRVisitorWithAnalyzer* self, const s_tir::SBlockNode* op) { + return self->constraint_scope_.WithNewScope([&]() -> ffi::Optional { + for (const auto& iter_var : op->iter_vars) { + self->analyzer_->Bind(iter_var->var, iter_var->dom); + } + return s_tir::StmtExprVisitor::VisitBlock(self, op); + }); +} + +void IRVisitorWithAnalyzer::Extension::InitVTable(VTable* vtable) { + vtable->ClearDispatch(); + vtable->SetDispatch([](const ffi::Object* node, ObjectVisitor* base) { + return VisitBlock(static_cast(base), + static_cast(node)); + }); +} +TVM_FFI_STATIC_INIT_BLOCK() { + IRVisitorWithAnalyzer::RegisterExtension(IRVisitorWithAnalyzer::Extension::InitVTable); +} +} // namespace tirx +} // namespace tvm 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 000000000000..87d20b82cb3f --- /dev/null +++ b/src/s_tir/ir/ir_visitor_with_analyzer.h @@ -0,0 +1,56 @@ +/* + * 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 + +#include "../../tirx/ir_visitor_with_analyzer.h" + +namespace tvm { +namespace tirx { +class IRVisitorWithAnalyzer::Extension { + public: + static ffi::Optional VisitBlock(IRVisitorWithAnalyzer* self, + const s_tir::SBlockNode* op); + static void InitVTable(VTable* vtable); +}; +} // namespace tirx +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 Visit_(const SBlockNode* op) { + return Parent::Extension::VisitBlock(this, op); + } + + protected: + static void InitVTable(VTable* vtable) { + Parent::InitVTable(vtable); + vtable->ClearDispatch(); + SetDispatch(vtable); + } +}; +} // namespace s_tir +} // namespace tvm +#endif diff --git a/src/s_tir/ir/specialize.cc b/src/s_tir/ir/specialize.cc new file mode 100644 index 000000000000..a46d00af8d5a --- /dev/null +++ b/src/s_tir/ir/specialize.cc @@ -0,0 +1,70 @@ +/* + * 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/ir/specialize.h" + +#include +#include + +namespace tvm { +namespace s_tir { +namespace { +ffi::Optional PlanBlockBuffers(tirx::StmtExprVisitor* planner, + const SBlockNode* op) { + // Block allocations were planned before all other block children by the specializer. + for (const tirx::BufferVar& buffer : op->alloc_buffers) { + TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(planner->WithDefRegionKind( + kTVMFFIDefRegionKindSimple, [&]() { return planner->Visit(buffer); })); + TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(planner->VisitBufferMetadata(buffer)); + } + for (const tirx::IterVar& iter : op->iter_vars) { + TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(planner->Visit(iter->dom->min)); + TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(planner->Visit(iter->dom->extent)); + } + for (const tirx::BufferRegion& region : op->reads) { + TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(planner->Visit(region)); + } + for (const tirx::BufferRegion& region : op->writes) { + TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(planner->Visit(region)); + } + for (const MatchBufferRegion& match : op->match_buffers) { + TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(planner->WithDefRegionKind( + kTVMFFIDefRegionKindSimple, [&]() { return planner->Visit(match->buffer); })); + TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(planner->VisitBufferMetadata(match->buffer)); + TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(planner->Visit(match->source)); + } + if (op->init.has_value()) { + TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(planner->Visit(op->init.value())); + } + return planner->Visit(op->body); +} + +void InitBufferPlanner(tirx::SpecializeVisitorVTable* vtable) { + vtable->ClearDispatch(); + vtable->SetDispatch([](const ffi::Object* node, ObjectVisitor* visitor) { + return PlanBlockBuffers(static_cast(visitor), + static_cast(node)); + }); +} + +} // namespace + +TVM_FFI_STATIC_INIT_BLOCK() { tirx::RegisterSpecializeBufferPlannerExtension(InitBufferPlanner); } +} // namespace s_tir +} // namespace tvm 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 000000000000..6aca37d59ab4 --- /dev/null +++ b/src/s_tir/ir/tir_visitor_with_path.cc @@ -0,0 +1,114 @@ +/* + * 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/ir/tir_visitor_with_path.h" + +#include + +namespace tvm { +namespace tirx { +using AccessPath = ffi::reflection::AccessPath; +using s_tir::SBlockNode; +using s_tir::SBlockRealizeNode; + +class TIRVisitorWithPath::Extension { + public: + static void InitVTable(VTable* vtable); + static void VisitBlock(TIRVisitorWithPath* self, const SBlockNode* op, AccessPath path); + static void VisitBlockRealize(TIRVisitorWithPath* self, const SBlockRealizeNode* op, + AccessPath path); +}; + +void TIRVisitorWithPath::Extension::VisitBlock(TIRVisitorWithPath* self, const SBlockNode* op, + AccessPath path) { + std::vector, DefContext, DefContext>> context; + + { + auto iter_path = path->Attr("iter_vars"); + for (size_t i = 0; i < op->iter_vars.size(); i++) { + context.push_back(self->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(self->WithDef(buf, buffer_path)); + } + } + + self->Visit(op->reads, path->Attr("reads")); + self->Visit(op->writes, path->Attr("writes")); + + { + auto match_path = path->Attr("match_buffers"); + for (size_t i = 0; i < op->match_buffers.size(); ++i) { + self->Visit(op->match_buffers[i]->source, match_path->ArrayItem(i)->Attr("source")); + } + + for (size_t i = 0; i < op->match_buffers.size(); i++) { + auto buf = op->match_buffers[i]->buffer; + auto buffer_path = match_path->ArrayItem(i)->Attr("buffer"); + + for (auto& def : self->WithMatchBufferDefs(buf, buffer_path)) { + context.push_back(std::move(def)); + } + context.push_back(self->WithDef(buf, buffer_path)); + } + } + + self->bind_scope_.WithNewScope([&]() { self->Visit(op->init, path->Attr("init")); }); + self->bind_scope_.WithNewScope([&]() { self->Visit(op->body, path->Attr("body")); }); + + while (context.size()) context.pop_back(); +} + +void TIRVisitorWithPath::Extension::VisitBlockRealize(TIRVisitorWithPath* self, + const SBlockRealizeNode* op, + AccessPath path) { + self->Visit(op->iter_values, path->Attr("iter_values")); + self->Visit(op->predicate, path->Attr("predicate")); + self->Visit(op->block, path->Attr("block")); +} + +void TIRVisitorWithPath::Extension::InitVTable(VTable* vtable) { + vtable->SetDispatch( + [](const ffi::ObjectRef& node, StmtVisitor* base, AccessPath path) { + auto* self = static_cast(base); + if (self->EnterExtensionStmt(node.get(), path)) { + VisitBlock(self, static_cast(node.get()), path); + } + }); + vtable->SetDispatch( + [](const ffi::ObjectRef& node, StmtVisitor* base, AccessPath path) { + auto* self = static_cast(base); + if (self->EnterExtensionStmt(node.get(), path)) { + VisitBlockRealize(self, static_cast(node.get()), path); + } + }); +} +TVM_FFI_STATIC_INIT_BLOCK() { + TIRVisitorWithPath::RegisterExtension(TIRVisitorWithPath::Extension::InitVTable); +} +} // namespace tirx +} // namespace tvm 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 9ad9fc5a7144..90fad9b3d3b2 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 #include #include +#include #include #include diff --git a/src/s_tir/meta_schedule/module_equality.cc b/src/s_tir/meta_schedule/module_equality.cc index 428dc551eb2f..7a590353edbc 100644 --- a/src/s_tir/meta_schedule/module_equality.cc +++ b/src/s_tir/meta_schedule/module_equality.cc @@ -22,6 +22,8 @@ #include #include #include +#include +#include #include #include @@ -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(anchor_block), + return ffi::StructuralHash::Hash(ffi::GetRef(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(anchor_block_lhs), - ffi::GetRef(anchor_block_rhs), + return tvm::ffi::StructuralEqual::Equal(ffi::GetRef(anchor_block_lhs), + ffi::GetRef(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 9ac8b2e699e4..a617ed690159 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>>* jobs) { - auto walk_fn = [=, &jobs](const tirx::SBlock& block) -> ffi::Expected { + auto walk_fn = [=, &jobs](const s_tir::SBlock& block) -> ffi::Expected { tirx::StmtSRef block_sref = sch->GetSRef(block.get()); - std::string block_name = block_sref->StmtAs()->name_hint; + std::string block_name = block_sref->StmtAs()->name_hint; if (ffi::Optional intrin_name = s_tir::GetAnn(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 e57b5624d648..abbf13cc48d3 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 #include +#include #include "../utils.h" 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 d1d8e633e035..65451de3374c 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 +#include #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() != nullptr) { + if (consumers.size() > 1 && lca_sref->StmtAs() != 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 d0e8db98d9f9..305e614ce56f 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 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(); + const s_tir::SBlockNode* block_node = (*sch)->GetSRef(block)->StmtAs(); TVM_FFI_ICHECK_EQ(block_node->writes.size(), 1); const DLDataType dtype = block_node->writes[0]->source.as_or_throw()->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 bf1ea750993d..ccf35589d727 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 @@ -448,10 +448,10 @@ std::vector MultiLevelTilingTensorCoreNode::TransformIntermediateOutputLa // Get the shape of the wmma accumulator auto [frag_shape_m, frag_shape_n] = [&]() { - tirx::SBlock intrin_block = TensorIntrin::Get(state->intrin_group.init_intrin) - .value() - ->desc->body.as_or_throw() - ->block; + s_tir::SBlock intrin_block = TensorIntrin::Get(state->intrin_group.init_intrin) + .value() + ->desc->body.as_or_throw() + ->block; tirx::For loop_m = intrin_block->body.as_or_throw(); tirx::For loop_n = loop_m->body.as_or_throw(); return std::make_tuple(loop_m->extent, loop_n->extent); @@ -628,9 +628,10 @@ std::vector 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(); + const s_tir::SBlockNode* cache_read_block = + sch->GetSRef(cache_read)->StmtAs(); tirx::BufferVar cache_read_buffer = - s_tir::GetNthAccessBuffer(sch->state(), ffi::GetRef(cache_read_block), 0, + s_tir::GetNthAccessBuffer(sch->state(), ffi::GetRef(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 @@ -783,9 +784,9 @@ ffi::Optional 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(block); + const s_tir::SBlock block_before_reindex = ffi::GetRef(block); if (block->reads.size() != 2 || block->writes.size() != 1) { // only matmul-like computation is allowed return std::nullopt; @@ -867,8 +868,8 @@ ffi::Optional MultiLevelTilingTensorCoreNode::TransformWithTensorIntrin( visited_buffers.insert(lhs_buffer); // 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(block), buffer_index, index_type); + const tirx::BufferRegion& reindexed_buffer_region = s_tir::GetNthAccessBufferRegion( + state->sch->state(), ffi::GetRef(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 51351fbd8937..5b8d16e28eea 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 +#include #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(); - const tirx::SBlockRealize block_realize = s_tir::GetSBlockRealize(sch->state(), block_sref); + const s_tir::SBlockNode* block_node = block_sref->StmtAs(); + 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()->dtype; diff --git a/src/s_tir/meta_schedule/trace_apply.cc b/src/s_tir/meta_schedule/trace_apply.cc index 34678aac8459..2bfc43056299 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 +#include +#include #include -#include #include #include diff --git a/src/s_tir/meta_schedule/utils.h b/src/s_tir/meta_schedule/utils.h index 02ac858b1b8a..d8d9e39a79ef 100644 --- a/src/s_tir/meta_schedule/utils.h +++ b/src/s_tir/meta_schedule/utils.h @@ -41,6 +41,7 @@ #include #include #include +#include #include #include #include @@ -301,7 +302,7 @@ inline std::string Concat(const ffi::Array& 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& 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 Visit(ffi::AnyView value) override { if (value.as()) return std::nullopt; - return tirx::StmtExprVisitor::Visit(value); + return s_tir::StmtExprVisitor::Visit(value); } static ffi::Array Collect(const s_tir::Schedule& sch, @@ -662,8 +663,8 @@ class SBlockCollector : public tirx::StmtExprVisitor { private: /*! \brief Override the Stmt visiting behaviour */ - ffi::Optional Visit_(const tirx::SBlockNode* block) override { - TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(tirx::StmtExprVisitor::Visit_(block)); + ffi::Optional 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(block)).cast()->value != 0; + collect_block = f_block_filter_(ffi::GetRef(block)).cast()->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 9dbba04aaafa..55ba3ef8f417 100644 --- a/src/s_tir/sblock_dependence_info.cc +++ b/src/s_tir/sblock_dependence_info.cc @@ -19,6 +19,7 @@ #include #include +#include #include 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 Visit(ffi::AnyView value) override { if (value.as()) 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 Visit_(const SBlockRealizeNode* realize) final { + ffi::Optional 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 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 f4f2f0c815a7..12ef4368a82a 100644 --- a/src/s_tir/sblock_scope.cc +++ b/src/s_tir/sblock_scope.cc @@ -19,6 +19,7 @@ #include #include #include +#include #include namespace tvm { @@ -85,7 +86,7 @@ SBlockScope::SBlockScope(const ffi::Array& child_block_srefs) { SMap> buffer_readers; SMap>& 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()].push_back( @@ -182,8 +183,8 @@ ffi::Optional SRefTreeCreator::Visit_(const ForNode* loop) { return std::nullopt; } -ffi::Optional SRefTreeCreator::Visit_(const SBlockRealizeNode* realize) { - const SBlockNode* block = realize->block.get(); +ffi::Optional 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 SRefTreeCreator::Visit_(const SBlockRealizeNode* r ffi::Optional 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 54a77e93d19e..2ee6240b7268 100644 --- a/src/s_tir/schedule/analysis.h +++ b/src/s_tir/schedule/analysis.h @@ -24,6 +24,7 @@ #include #include #include +#include #include #include diff --git a/src/s_tir/schedule/analysis/analysis.cc b/src/s_tir/schedule/analysis/analysis.cc index 57a098a31ca5..789dcbff7314 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(); + const auto* block_realize = loop_i->body.as(); if (!block_realize || block_realize->block.get() != block) { return false; } @@ -1791,7 +1791,7 @@ ffi::Optional 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 9de832d8051d..680d556df6f4 100644 --- a/src/s_tir/schedule/analysis/reducer.cc +++ b/src/s_tir/schedule/analysis/reducer.cc @@ -19,6 +19,7 @@ #include #include #include +#include #include #include "../utils.h" diff --git a/src/s_tir/schedule/analysis/verify.cc b/src/s_tir/schedule/analysis/verify.cc index d83d5df4f1a9..11d32ae23525 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 +#include #include "../utils.h" diff --git a/src/s_tir/schedule/concrete_schedule.cc b/src/s_tir/schedule/concrete_schedule.cc index 145e054e731f..ac94bd5c1ee3 100644 --- a/src/s_tir/schedule/concrete_schedule.cc +++ b/src/s_tir/schedule/concrete_schedule.cc @@ -21,6 +21,7 @@ #include #include #include +#include #include diff --git a/src/s_tir/schedule/concrete_schedule.h b/src/s_tir/schedule/concrete_schedule.h index 3a1501ed7c2a..205eda800a3e 100644 --- a/src/s_tir/schedule/concrete_schedule.h +++ b/src/s_tir/schedule/concrete_schedule.h @@ -22,6 +22,7 @@ #include #include #include +#include #include #include diff --git a/src/s_tir/schedule/error.h b/src/s_tir/schedule/error.h index 006304ea5a0a..0bdc14837f77 100644 --- a/src/s_tir/schedule/error.h +++ b/src/s_tir/schedule/error.h @@ -23,6 +23,7 @@ #include #include #include +#include #include #include diff --git a/src/s_tir/schedule/ir_comparator.cc b/src/s_tir/schedule/ir_comparator.cc index 1b8b1851cf4f..015ba132f661 100644 --- a/src/s_tir/schedule/ir_comparator.cc +++ b/src/s_tir/schedule/ir_comparator.cc @@ -20,6 +20,7 @@ #include #include +#include #include #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 75303df5b81e..58db4f8ebdd2 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 +#include #include #include diff --git a/src/s_tir/schedule/primitive.h b/src/s_tir/schedule/primitive.h index af589e22c174..d7f0a77c0f1d 100644 --- a/src/s_tir/schedule/primitive.h +++ b/src/s_tir/schedule/primitive.h @@ -23,6 +23,7 @@ #include #include #include +#include #include diff --git a/src/s_tir/schedule/primitive/annotate.cc b/src/s_tir/schedule/primitive/annotate.cc index 8b22e99e9fff..9fd1bc275246 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 +#include #include "../utils.h" diff --git a/src/s_tir/schedule/primitive/blockize_tensorize.cc b/src/s_tir/schedule/primitive/blockize_tensorize.cc index 9b68f6e19962..f8343288d26a 100644 --- a/src/s_tir/schedule/primitive/blockize_tensorize.cc +++ b/src/s_tir/schedule/primitive/blockize_tensorize.cc @@ -22,6 +22,7 @@ #include #include #include +#include #include diff --git a/src/s_tir/schedule/primitive/cache_index.cc b/src/s_tir/schedule/primitive/cache_index.cc index c656d4ed74d9..630c4a681482 100644 --- a/src/s_tir/schedule/primitive/cache_index.cc +++ b/src/s_tir/schedule/primitive/cache_index.cc @@ -20,6 +20,7 @@ #include #include #include +#include #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 d34b103be628..c88a2b9b3284 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 // For the arith::Analyzer::Simplify() method simplifying terms #include #include +#include +#include #include #include #include -#include #include // For std::find_if #include // For the hashtable datatype @@ -214,7 +215,7 @@ ffi::Optional ComputationsDoneBy::Visit(ffi::AnyView expr_value) auto opt_expr = expr_value.as(); if (!opt_expr) { - return StmtExprVisitor::Visit(expr_value); + return s_tir::StmtExprVisitor::Visit(expr_value); } PrimExpr expr = opt_expr.value(); if (expr.as() != nullptr || expr.as() != nullptr || @@ -322,7 +323,7 @@ ComputationTable ComputationsDoneBy::ComputationsDoneByChildrenOf( std::function can_contain_computations) { auto computations_done_by = ffi::make_object(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 can_contain_computations) { auto computations_done_by = ffi::make_object(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 is_eligible_co */ ffi::Optional DirectSubexpr::Visit(ffi::AnyView expr_value) { auto opt_expr = expr_value.as(); - 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 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 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 2159381816c7..af54b6415e5a 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 #include #include +#include // For the class s_tir::StmtExprVisitor #include #include -#include // For the class StmtExprVisitor #include #include // 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 GetDirectSubexpressions( diff --git a/src/s_tir/schedule/primitive/cache_read_write.cc b/src/s_tir/schedule/primitive/cache_read_write.cc index a113f094271c..d0d9ec36ef98 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 #include #include +#include #include #include diff --git a/src/s_tir/schedule/primitive/compute_at.cc b/src/s_tir/schedule/primitive/compute_at.cc index efccaa7b3e00..9c73b9a41ed4 100644 --- a/src/s_tir/schedule/primitive/compute_at.cc +++ b/src/s_tir/schedule/primitive/compute_at.cc @@ -18,6 +18,7 @@ */ #include #include +#include #include "../utils.h" diff --git a/src/s_tir/schedule/primitive/decompose_padding.cc b/src/s_tir/schedule/primitive/decompose_padding.cc index 6a4d03ab24e5..283a5609a350 100644 --- a/src/s_tir/schedule/primitive/decompose_padding.cc +++ b/src/s_tir/schedule/primitive/decompose_padding.cc @@ -19,6 +19,7 @@ #include #include #include +#include #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 b9d06ed27774..558aff410d43 100644 --- a/src/s_tir/schedule/primitive/for_kind.cc +++ b/src/s_tir/schedule/primitive/for_kind.cc @@ -18,6 +18,7 @@ */ #include #include +#include #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 ee3a489141ca..b2e08236795b 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 +#include #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 2f9f9950f5fa..43649044ea00 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 +#include #include diff --git a/src/s_tir/schedule/primitive/layout_transformation.cc b/src/s_tir/schedule/primitive/layout_transformation.cc index dff34341dabf..c6122710b257 100644 --- a/src/s_tir/schedule/primitive/layout_transformation.cc +++ b/src/s_tir/schedule/primitive/layout_transformation.cc @@ -23,11 +23,12 @@ #include #include #include +#include #include #include -#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 Visit(ffi::AnyView value) override { if (value.as()) return std::nullopt; - return tirx::StmtExprVisitor::Visit(value); + return s_tir::StmtExprVisitor::Visit(value); } static ffi::Map Collect(SBlock result, @@ -788,7 +789,7 @@ class ReuseBlocksCollector : public tirx::StmtExprVisitor { private: /*! \brief Override the Stmt visiting behaviour */ - ffi::Optional Visit_(const tirx::SBlockNode* block) override { + ffi::Optional Visit_(const s_tir::SBlockNode* block) override { SBlock block_ref = ffi::GetRef(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 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 Mutate(ffi::AnyView value, InplaceMode inplace_mode) final { const auto* stmt = value.as(); diff --git a/src/s_tir/schedule/primitive/loop_transformation.cc b/src/s_tir/schedule/primitive/loop_transformation.cc index a7f9cbbec485..f2987eae4706 100644 --- a/src/s_tir/schedule/primitive/loop_transformation.cc +++ b/src/s_tir/schedule/primitive/loop_transformation.cc @@ -19,6 +19,7 @@ #include #include #include +#include #include "../utils.h" diff --git a/src/s_tir/schedule/primitive/pad_einsum.cc b/src/s_tir/schedule/primitive/pad_einsum.cc index 05de3ac334b8..7bfd67177913 100644 --- a/src/s_tir/schedule/primitive/pad_einsum.cc +++ b/src/s_tir/schedule/primitive/pad_einsum.cc @@ -19,6 +19,7 @@ #include #include +#include #include #include "../utils.h" diff --git a/src/s_tir/schedule/primitive/reduction.cc b/src/s_tir/schedule/primitive/reduction.cc index bb596dc6ab72..5caa524fb978 100644 --- a/src/s_tir/schedule/primitive/reduction.cc +++ b/src/s_tir/schedule/primitive/reduction.cc @@ -20,6 +20,7 @@ #include #include #include +#include #include #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 9eab2390c7e7..58bc95e03624 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 +#include #include diff --git a/src/s_tir/schedule/primitive/rolling_buffer.cc b/src/s_tir/schedule/primitive/rolling_buffer.cc index 0014836b8a36..875a34b5f7a2 100644 --- a/src/s_tir/schedule/primitive/rolling_buffer.cc +++ b/src/s_tir/schedule/primitive/rolling_buffer.cc @@ -18,6 +18,7 @@ */ #include #include +#include #include diff --git a/src/s_tir/schedule/schedule.cc b/src/s_tir/schedule/schedule.cc index e8e666672f08..ef6af17c3a3f 100644 --- a/src/s_tir/schedule/schedule.cc +++ b/src/s_tir/schedule/schedule.cc @@ -17,6 +17,7 @@ * under the License. */ #include +#include #include "./utils.h" namespace tvm { diff --git a/src/s_tir/schedule/state.cc b/src/s_tir/schedule/state.cc index ab5e6b370b8c..86d51bbca4ce 100644 --- a/src/s_tir/schedule/state.cc +++ b/src/s_tir/schedule/state.cc @@ -20,6 +20,7 @@ #include #include #include +#include #include "./utils.h" namespace tvm { diff --git a/src/s_tir/schedule/traced_schedule.cc b/src/s_tir/schedule/traced_schedule.cc index 646269abfb7e..8cea5fc2bb50 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 + 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 52c1842bc0d0..0ddd2378ba08 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 +#include #include "./concrete_schedule.h" diff --git a/src/s_tir/schedule/transform.cc b/src/s_tir/schedule/transform.cc index db58cf3865f7..9de1345711cb 100644 --- a/src/s_tir/schedule/transform.cc +++ b/src/s_tir/schedule/transform.cc @@ -21,6 +21,7 @@ #include #include #include +#include #include #include "../../tirx/transform/ir_utils.h" @@ -469,7 +470,7 @@ void BlockBufferAccessSimplifier::SimplifyBufferIndices(ffi::Array* in UnchangedOr 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(op)) .as_or_throw(); auto* n = block.CopyOnWrite(); @@ -480,7 +481,7 @@ UnchangedOr BlockBufferAccessSimplifier::Mutate_(const SBlockNode* op, UnchangedOr 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(op)) .as_or_throw(); SimplifyBufferIndices(&node.CopyOnWrite()->indices); @@ -489,7 +490,7 @@ UnchangedOr BlockBufferAccessSimplifier::Mutate_(const BufferStoreNode* op UnchangedOr 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(op)) .as_or_throw(); SimplifyBufferIndices(&node.CopyOnWrite()->indices); diff --git a/src/s_tir/schedule/transform.h b/src/s_tir/schedule/transform.h index 2392937dd1bd..5cf61192ce08 100644 --- a/src/s_tir/schedule/transform.h +++ b/src/s_tir/schedule/transform.h @@ -22,12 +22,13 @@ #include #include #include -#include +#include +#include #include #include -#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 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 38330f31c13e..9ac1aaec06dc 100644 --- a/src/s_tir/schedule/utils.h +++ b/src/s_tir/schedule/utils.h @@ -26,15 +26,17 @@ #include #include #include +#include #include #include #include #include +#include +#include #include #include #include #include -#include #include #include @@ -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 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 Visit(ffi::AnyView value) override { if (value.as()) return std::nullopt; - return tirx::StmtExprVisitor::Visit(value); + return s_tir::StmtExprVisitor::Visit(value); } - ffi::Optional Visit_(const tirx::SBlockNode* block) override { + ffi::Optional 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 000000000000..2f4a78886a0c --- /dev/null +++ b/src/s_tir/stmt.cc @@ -0,0 +1,412 @@ +/* + * 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 +#include +#include +#include +#include +#include + +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( + 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( + value); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, mapped_buffer, + mutator->WithDefRegionKind(kTVMFFIDefRegionKindSimple, [&]() { + return mutator->MutateExpected(self->buffer); + })); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, mapped_source, + mutator->MutateExpected(self->source)); + if (mapped_buffer.UnchangedOrSameAs(self->buffer) && + mapped_source.UnchangedOrSameAs(self->source)) { + return ffi::Unchanged().CopyToTVMFFIAny(); + } + ffi::ObjectPtr copy = ffi::make_object(*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( + ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck( + value)); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, mapped_buffer, + mutator->WithDefRegionKind(kTVMFFIDefRegionKindSimple, [&]() { + return mutator->MutateExpected(self->buffer, + ffi::InplaceMode::kAllow); + })); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN( + ffi::UnchangedOr, 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 { + // skips: name_hint + const SBlockNode* self = + ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value); + TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->iter_vars)); + 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->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->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 { + // skips: name_hint + const SBlockNode* self = + ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr>, mapped_iter_vars, + mutator->MutateExpected(self->iter_vars)); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr>, mapped_reads, + mutator->MutateExpected(self->reads)); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr>, mapped_writes, + mutator->MutateExpected(self->writes)); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr>, mapped_alloc_buffers, + mutator->WithDefRegionKind(kTVMFFIDefRegionKindSimple, [&]() { + return mutator->MutateExpected(self->alloc_buffers); + })); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr>, + mapped_match_buffers, + mutator->MutateExpected(self->match_buffers)); + using AnnotationMap = ffi::Map; + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, mapped_annotations, + mutator->MutateExpected(self->annotations)); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr>, mapped_init, + mutator->MutateExpected(self->init)); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, 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 copy = ffi::make_object(*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 { + // skips: name_hint + SBlockNode* self = const_cast( + ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value)); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN( + ffi::UnchangedOr>, mapped_iter_vars, + mutator->MutateExpected(self->iter_vars, ffi::InplaceMode::kAllow)); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr>, mapped_reads, + mutator->MutateExpected(self->reads, ffi::InplaceMode::kAllow)); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN( + ffi::UnchangedOr>, mapped_writes, + mutator->MutateExpected(self->writes, ffi::InplaceMode::kAllow)); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr>, mapped_alloc_buffers, + mutator->WithDefRegionKind(kTVMFFIDefRegionKindSimple, [&]() { + return mutator->MutateExpected(self->alloc_buffers, + ffi::InplaceMode::kAllow); + })); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN( + ffi::UnchangedOr>, mapped_match_buffers, + mutator->MutateExpected(self->match_buffers, ffi::InplaceMode::kAllow)); + using AnnotationMap = ffi::Map; + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN( + ffi::UnchangedOr, mapped_annotations, + mutator->MutateExpected(self->annotations, ffi::InplaceMode::kAllow)); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr>, mapped_init, + mutator->MutateExpected(self->init, ffi::InplaceMode::kAllow)); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, 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(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(value); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr>, mapped_iter_values, + mutator->MutateExpected(self->iter_values)); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, mapped_predicate, + mutator->MutateExpected(self->predicate)); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, 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 copy = ffi::make_object(*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( + ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value)); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN( + ffi::UnchangedOr>, mapped_iter_values, + mutator->MutateExpected(self->iter_values, ffi::InplaceMode::kAllow)); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN( + ffi::UnchangedOr, mapped_predicate, + mutator->MutateExpected(self->predicate, ffi::InplaceMode::kAllow)); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, 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, BufferRegion source) { + const BufferVar& source_buffer = source->buffer; + 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 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()) { + 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 node = ffi::make_object(); + 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() + .attr(refl::type_attr::kStructuralVisit, reinterpret_cast(&MatchBufferRegionVisit)) + .attr(refl::type_attr::kStructuralMutate, reinterpret_cast(&MatchBufferRegionMutate)) + .attr(refl::type_attr::kStructuralMaybeInplaceMutate, + reinterpret_cast(&MatchBufferRegionMaybeInplaceMutate)); + + refl::GlobalDef().def("s_tir.MatchBufferRegion", [](BufferVar buffer, BufferRegion source) { + return MatchBufferRegion(buffer, source); + }); +} + +// Block +SBlock::SBlock(ffi::Array iter_vars, ffi::Array reads, + ffi::Array writes, ffi::String name_hint, Stmt body, + ffi::Optional init, ffi::Array alloc_buffers, + ffi::Array match_buffers, ffi::Map annotations, + Span span) { + ffi::ObjectPtr node = ffi::make_object(); + 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 alloc_buffers, Span span) { + ffi::ObjectPtr node = ffi::make_object(); + 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() + .attr(refl::type_attr::kStructuralVisit, reinterpret_cast(&SBlockVisit)) + .attr(refl::type_attr::kStructuralMutate, reinterpret_cast(&SBlockMutate)) + .attr(refl::type_attr::kStructuralMaybeInplaceMutate, + reinterpret_cast(&SBlockMaybeInplaceMutate)); + + refl::GlobalDef().def("s_tir.SBlock", + [](ffi::Array iter_vars, ffi::Array reads, + ffi::Array writes, ffi::String name_hint, Stmt body, + ffi::Optional init, ffi::Array alloc_buffers, + ffi::Array match_buffers, + ffi::Map annotations, Span span) { + return SBlock(iter_vars, reads, writes, name_hint, body, init, + alloc_buffers, match_buffers, annotations, span); + }); +} + +// BlockRealize +SBlockRealize::SBlockRealize(ffi::Array 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 node = ffi::make_object(); + 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() + .def("tirx.prevent_inline", true) + .attr(refl::type_attr::kStructuralVisit, reinterpret_cast(&SBlockRealizeVisit)) + .attr(refl::type_attr::kStructuralMutate, reinterpret_cast(&SBlockRealizeMutate)) + .attr(refl::type_attr::kStructuralMaybeInplaceMutate, + reinterpret_cast(&SBlockRealizeMaybeInplaceMutate)); + + refl::GlobalDef().def("s_tir.SBlockRealize", [](ffi::Array 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 000000000000..6c8d66a65f8d --- /dev/null +++ b/src/s_tir/stmt_functor.cc @@ -0,0 +1,234 @@ +/* + * 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 +#include + +#include +#include + +namespace tvm { +namespace s_tir { + +using namespace tirx; + +// Generic TIRX passes keep native dialect traversal without owning dialect nodes. +// StructuralVisitor/ObjectVisitor and structural mutation retain the full field walk. +TVM_FFI_STATIC_INIT_BLOCK() { + tirx::StmtExprVisitor::RegisterExtension([](tirx::StmtExprVisitor::VTable* vtable) { + vtable->SetDispatch([](const ffi::Object* node, ObjectVisitor* visitor) { + return StmtExprVisitor::VisitBlock(static_cast(visitor), + static_cast(node)); + }); + vtable->SetDispatch([](const ffi::Object* node, ObjectVisitor* visitor) { + return StmtExprVisitor::VisitBlockRealize(static_cast(visitor), + static_cast(node)); + }); + }); + tirx::StmtExprMutator::RegisterExtension([](tirx::StmtExprMutator::VTable* vtable) { + vtable->SetDispatch( + [](const ffi::Object* node, ObjectMutator* mutator, InplaceMode mode) { + return ffi::details::UnchangedOrUnsafe::MoveFromTVMFFIAny( + ffi::details::UnchangedOrUnsafe::MoveToTVMFFIAny( + StmtExprMutator::MutateBlock(static_cast(mutator), + static_cast(node), mode))); + }); + vtable->SetDispatch( + [](const ffi::Object* node, ObjectMutator* mutator, InplaceMode mode) { + return ffi::details::UnchangedOrUnsafe::MoveFromTVMFFIAny( + ffi::details::UnchangedOrUnsafe::MoveToTVMFFIAny(StmtExprMutator::MutateBlockRealize( + static_cast(mutator), + static_cast(node), mode))); + }); + }); +} + +void StmtExprVisitor::InitVTable(VTable* vtable) { + tirx::StmtExprVisitor::InitVTable(vtable); + vtable->ClearDispatch(); + vtable->ClearDispatch(); + SetDispatch(vtable); + SetDispatch(vtable); +} + +void StmtExprMutator::InitVTable(VTable* vtable) { + tirx::StmtExprMutator::InitVTable(vtable); + vtable->ClearDispatch(); + vtable->ClearDispatch(); + SetDispatch(vtable); + SetDispatch(vtable); +} + +ffi::Optional StmtExprVisitor::Visit_(const SBlockNode* op) { + return VisitBlock(this, op); +} + +ffi::Optional 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)); + } + for (const BufferRegion& region : op->reads) { + TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->Visit(region)); + } + for (const BufferRegion& region : op->writes) { + TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->Visit(region)); + } + 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)); + } + if (op->init.has_value()) { + TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->Visit(op->init.value())); + } + return visitor->Visit(op->body); +} + +ffi::Optional StmtExprVisitor::Visit_(const SBlockRealizeNode* op) { + return VisitBlockRealize(this, op); +} + +ffi::Optional 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 StmtExprMutator::Mutate_(const SBlockNode* op, InplaceMode inplace_mode) { + return MutateBlock(this, op, inplace_mode); +} + +UnchangedOr 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> replacements; + for (size_t i = 0; i < iters->size(); ++i) { + const auto* iter = (*iters)[i].as(); + InplaceMode domain_mode = iter->unique() ? iter_mode : InplaceMode::kDisallow; + auto domain = mutator->Mutate(iter->dom, domain_mode).as_or_throw>(); + if (domain.UnchangedOrSameAs(iter->dom)) continue; + if (domain_mode == InplaceMode::kAllow) { + const_cast(iter)->dom = std::move(domain).ValueUnchecked(); + } else { + auto updated = ffi::make_object(*iter); + updated->dom = std::move(domain).ValueUnchecked(); + replacements.emplace_back(i, IterVar(std::move(updated))); + } + } + UnchangedOr> iter_vars = ffi::Unchanged(); + if (!replacements.empty()) { + if (iter_mode == InplaceMode::kAllow) { + for (auto& [i, iter] : replacements) { + const_cast(iters)->SetItem(i, std::move(iter)); + } + } else { + ffi::Array 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>>(); + auto reads = + mutator->Mutate(op->reads, inplace_mode).as_or_throw>>(); + auto writes = mutator->Mutate(op->writes, inplace_mode) + .as_or_throw>>(); + auto match_buffers = mutator->Mutate(op->match_buffers, inplace_mode) + .as_or_throw>>(); + auto init = + mutator->Mutate(op->init, inplace_mode).as_or_throw>>(); + 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(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(*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 StmtExprMutator::Mutate_(const SBlockRealizeNode* op, InplaceMode inplace_mode) { + return MutateBlockRealize(this, op, inplace_mode); +} + +UnchangedOr StmtExprMutator::MutateBlockRealize(tirx::StmtExprMutator* mutator, + const SBlockRealizeNode* op, + InplaceMode inplace_mode) { + auto iter_values = mutator->Mutate(op->iter_values, inplace_mode) + .as_or_throw>>(); + auto predicate = mutator->Mutate(op->predicate, inplace_mode); + auto block = mutator->Mutate(op->block, inplace_mode).as_or_throw>(); + 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(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(*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/transform/annotate_irregular_loop.cc b/src/s_tir/transform/annotate_irregular_loop.cc index d3067da5b83d..b7ac58630dc9 100644 --- a/src/s_tir/transform/annotate_irregular_loop.cc +++ b/src/s_tir/transform/annotate_irregular_loop.cc @@ -24,9 +24,9 @@ #include #include #include +#include #include #include -#include namespace tvm { namespace s_tir { diff --git a/src/s_tir/transform/bound_checker.cc b/src/s_tir/transform/bound_checker.cc index 5b6a9fbcf58d..00d8b1f8b1d0 100644 --- a/src/s_tir/transform/bound_checker.cc +++ b/src/s_tir/transform/bound_checker.cc @@ -29,10 +29,10 @@ #include #include #include +#include #include #include #include -#include #include #include diff --git a/src/s_tir/transform/canonicalize_loop.cc b/src/s_tir/transform/canonicalize_loop.cc index 048595c1dd35..101acfd07f70 100644 --- a/src/s_tir/transform/canonicalize_loop.cc +++ b/src/s_tir/transform/canonicalize_loop.cc @@ -25,10 +25,10 @@ #include #include #include +#include #include #include #include -#include #include diff --git a/src/s_tir/transform/compact_buffer_region.cc b/src/s_tir/transform/compact_buffer_region.cc index bd5e5042e430..f4062b3d236e 100644 --- a/src/s_tir/transform/compact_buffer_region.cc +++ b/src/s_tir/transform/compact_buffer_region.cc @@ -27,9 +27,9 @@ #include #include #include +#include #include #include -#include #include #include diff --git a/src/s_tir/transform/convert_blocks_to_opaque.cc b/src/s_tir/transform/convert_blocks_to_opaque.cc index 2c8db5af32fb..c12917010196 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 #include +#include +#include #include -#include #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 9839fd47d2c5..353326c21d67 100644 --- a/src/s_tir/transform/default_gpu_schedule.cc +++ b/src/s_tir/transform/default_gpu_schedule.cc @@ -19,6 +19,7 @@ #include #include +#include #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(); + const auto* realize = func->body.as(); 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() || inner->IsInstance()) { + if (inner->IsInstance() || inner->IsInstance()) { 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(), tirx::IterVarType::kDataPar); - tirx::SBlock inner_block = realize->block; + s_tir::SBlock inner_block = realize->block; inner_block.CopyOnWrite()->iter_vars = ffi::Array{new_iter}; - tirx::SBlockRealize inner_realize( + s_tir::SBlockRealize inner_realize( /*iter_values=*/ffi::Array{loop_var.as_or_throw()}, /*predicate=*/realize->predicate, inner_block); tirx::Stmt for_stmt = tirx::For(loop_var.as_or_throw(), zero, one, tirx::ForKind::kSerial, inner_realize); - tirx::SBlock root_block(/*iter_vars=*/ffi::Array{}, - /*reads=*/ffi::Array{}, - /*writes=*/ffi::Array{}, - /*name_hint=*/"root", /*body=*/for_stmt); - tirx::SBlockRealize root_realize(/*iter_values=*/ffi::Array{}, - /*predicate=*/IntImm::Bool(true), root_block); + s_tir::SBlock root_block(/*iter_vars=*/ffi::Array{}, + /*reads=*/ffi::Array{}, + /*writes=*/ffi::Array{}, + /*name_hint=*/"root", /*body=*/for_stmt); + s_tir::SBlockRealize root_realize(/*iter_values=*/ffi::Array{}, + /*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 9af41d25c551..4420def0ae86 100644 --- a/src/s_tir/transform/hoist_expression.cc +++ b/src/s_tir/transform/hoist_expression.cc @@ -26,9 +26,10 @@ #include #include #include +#include +#include #include #include -#include #include #include @@ -37,7 +38,7 @@ #include "../../arith/interval_set.h" #include "../../runtime/thread_storage_scope.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 { @@ -455,10 +456,10 @@ class HoistInfoCollector : public StmtExprVisitor { std::unordered_set 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); @@ -471,7 +472,7 @@ class ExpressionHoister : public tirx::IRMutatorWithAnalyzer { } private: - using Parent = tirx::IRMutatorWithAnalyzer; + using Parent = s_tir::IRMutatorWithAnalyzer; public: explicit ExpressionHoister(std::vector loop_info, diff --git a/src/s_tir/transform/inject_double_buffer.cc b/src/s_tir/transform/inject_double_buffer.cc index 2100a82780bb..f95dbd7b0c5a 100644 --- a/src/s_tir/transform/inject_double_buffer.cc +++ b/src/s_tir/transform/inject_double_buffer.cc @@ -27,9 +27,9 @@ #include #include #include +#include #include #include -#include #include "../../tirx/transform/ir_utils.h" diff --git a/src/s_tir/transform/inject_permuted_layout.cc b/src/s_tir/transform/inject_permuted_layout.cc index c10651589029..4f5fe0c50699 100644 --- a/src/s_tir/transform/inject_permuted_layout.cc +++ b/src/s_tir/transform/inject_permuted_layout.cc @@ -24,14 +24,15 @@ #include #include #include +#include +#include #include #include #include -#include #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 1d206840e365..58bf05be31cc 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 #include #include +#include +#include #include #include #include #include -#include #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 8eb4fe739ac3..97f34a29cd7c 100644 --- a/src/s_tir/transform/inject_ptx_ldg32.cc +++ b/src/s_tir/transform/inject_ptx_ldg32.cc @@ -22,11 +22,12 @@ #include #include #include +#include +#include #include #include #include #include -#include #include "../../arith/const_fold.h" #include "../../arith/pattern_match.h" diff --git a/src/s_tir/transform/inject_virtual_thread.cc b/src/s_tir/transform/inject_virtual_thread.cc index 99ccdeea8bfb..65ec0942ed77 100644 --- a/src/s_tir/transform/inject_virtual_thread.cc +++ b/src/s_tir/transform/inject_virtual_thread.cc @@ -27,13 +27,13 @@ #include #include #include +#include #include #include -#include #include -#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 { @@ -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; diff --git a/src/s_tir/transform/ir_utils.cc b/src/s_tir/transform/ir_utils.cc new file mode 100644 index 000000000000..e7bb0d9adb77 --- /dev/null +++ b/src/s_tir/transform/ir_utils.cc @@ -0,0 +1,77 @@ +/* + * 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 +#include + +namespace tvm { +namespace s_tir { +using namespace tirx; +using namespace tvm::prim; + +ffi::Array ConvertIndices(const MatchBufferRegion& match_buffer, + const ffi::Array& indices) { + const BufferVar& target = match_buffer->buffer; + const BufferRegion& source = match_buffer->source; + TVM_FFI_ICHECK_EQ(indices.size(), target->shape.size()); + + arith::Analyzer analyzer; + ffi::Array 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 BufferRegion& 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; +} + +} // 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 000000000000..dcc0b0bcf73e --- /dev/null +++ b/src/s_tir/transform/ir_utils.h @@ -0,0 +1,47 @@ +/* + * 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 + +#include "../../tirx/transform/ir_utils.h" + +namespace tvm { +namespace s_tir { + +/*! + * \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 ConvertIndices(const MatchBufferRegion& match_buffer, + const ffi::Array& 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); + +} // 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 f46047e01088..c252563b7f8a 100644 --- a/src/s_tir/transform/lift_thread_binding.cc +++ b/src/s_tir/transform/lift_thread_binding.cc @@ -24,8 +24,8 @@ #include #include +#include #include -#include #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 16668e837399..ad32131734dd 100644 --- a/src/s_tir/transform/loop_partition.cc +++ b/src/s_tir/transform/loop_partition.cc @@ -30,11 +30,12 @@ #include #include #include +#include #include +#include #include #include #include -#include #include #include diff --git a/src/s_tir/transform/lower_async_dma.cc b/src/s_tir/transform/lower_async_dma.cc index 389c88877347..bc1dfc6d437b 100644 --- a/src/s_tir/transform/lower_async_dma.cc +++ b/src/s_tir/transform/lower_async_dma.cc @@ -28,25 +28,25 @@ #include #include #include +#include #include #include #include -#include #include #include -#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 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(loop), ffi::GetRef(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 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(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()}); + return SeqStmt({call_dma_wait, s_tir::IRMutatorWithAnalyzer::Mutate( + ffi::AnyView(async_wait->body), body_mode) + .ValueOrUnchanged(async_wait->body) + .as_or_throw()}); // 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(); TVM_FFI_ICHECK(queue_id_node); async_queue_id_ = queue_id_node->value.as().value(); - auto result = tirx::IRMutatorWithAnalyzer::Mutate_(op, inplace_mode) + auto result = s_tir::IRMutatorWithAnalyzer::Mutate_(op, inplace_mode) .ValueOrUnchanged(ffi::GetRef(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 4926d9161fad..51ff831942df 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 #include #include +#include #include +#include #include #include #include -#include #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 f230ba9ea1b5..a6441af39d82 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 +#include +#include #include #include -#include #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 eecfae6c3629..c17dda686e5c 100644 --- a/src/s_tir/transform/lower_match_buffer.cc +++ b/src/s_tir/transform/lower_match_buffer.cc @@ -27,12 +27,14 @@ #include #include #include +#include +#include #include #include #include -#include #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 e069856a32ed..fabb258aed00 100644 --- a/src/s_tir/transform/lower_opaque_block.cc +++ b/src/s_tir/transform/lower_opaque_block.cc @@ -24,8 +24,8 @@ #include #include #include +#include #include -#include #include "../../tirx/transform/ir_utils.h" diff --git a/src/s_tir/transform/lower_thread_allreduce.cc b/src/s_tir/transform/lower_thread_allreduce.cc index 4c0b28ff5fa9..6f4691e26ec7 100644 --- a/src/s_tir/transform/lower_thread_allreduce.cc +++ b/src/s_tir/transform/lower_thread_allreduce.cc @@ -27,11 +27,11 @@ #include #include #include +#include #include #include #include #include -#include #include diff --git a/src/s_tir/transform/lower_vtcm_alloc.cc b/src/s_tir/transform/lower_vtcm_alloc.cc index 526819ec7d84..54f4058d841a 100644 --- a/src/s_tir/transform/lower_vtcm_alloc.cc +++ b/src/s_tir/transform/lower_vtcm_alloc.cc @@ -19,10 +19,10 @@ #include #include +#include #include #include #include -#include 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 f7484ce295c7..67ad132466cb 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 #include #include +#include #include #include -#include #include diff --git a/src/s_tir/transform/memhammer_lower_auto_copy.cc b/src/s_tir/transform/memhammer_lower_auto_copy.cc index fea18924dcd8..f1a8a052c2c5 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 #include #include +#include #include #include #include -#include #include #include diff --git a/src/s_tir/transform/memhammer_rewrite_rule.h b/src/s_tir/transform/memhammer_rewrite_rule.h index 83b96e49e1dd..4c0a136f37c4 100644 --- a/src/s_tir/transform/memhammer_rewrite_rule.h +++ b/src/s_tir/transform/memhammer_rewrite_rule.h @@ -22,10 +22,10 @@ #include #include #include +#include #include #include #include -#include #include diff --git a/src/s_tir/transform/memhammer_tensorcore_rewrite.cc b/src/s_tir/transform/memhammer_tensorcore_rewrite.cc index 736357275195..e21a086e34a1 100644 --- a/src/s_tir/transform/memhammer_tensorcore_rewrite.cc +++ b/src/s_tir/transform/memhammer_tensorcore_rewrite.cc @@ -21,6 +21,7 @@ #include #include #include +#include #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 d1d72eb36871..321db8f39d74 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 #include #include +#include #include #include -#include #include #include 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 f49ba03a24e0..b7412857b4f6 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 +#include +#include +#include #include #include -#include #include #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 448b883f6f75..458f0e06de92 100644 --- a/src/s_tir/transform/profile_instrumentation.cc +++ b/src/s_tir/transform/profile_instrumentation.cc @@ -27,10 +27,10 @@ #include #include #include +#include #include #include #include -#include 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 9fa243cd37db..aa70ab33fe80 100644 --- a/src/s_tir/transform/remove_store_undef.cc +++ b/src/s_tir/transform/remove_store_undef.cc @@ -24,12 +24,13 @@ #include #include #include +#include +#include #include #include #include #include #include -#include 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 1632099a0db7..d377bfa6c4fc 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 #include +#include #include #include #include -#include #include diff --git a/src/s_tir/transform/renew_defs.cc b/src/s_tir/transform/renew_defs.cc index 9740975e11d8..4776777ca3f8 100644 --- a/src/s_tir/transform/renew_defs.cc +++ b/src/s_tir/transform/renew_defs.cc @@ -25,8 +25,9 @@ #include #include #include +#include +#include #include -#include 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 e290d3e2bd7c..8b5f6cff4bdb 100644 --- a/src/s_tir/transform/renormalize_split_pattern.cc +++ b/src/s_tir/transform/renormalize_split_pattern.cc @@ -23,14 +23,15 @@ */ #include #include +#include +#include #include #include #include #include -#include #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 7922185df0bb..66ff4bab8b2b 100644 --- a/src/s_tir/transform/rewrite_unsafe_select.cc +++ b/src/s_tir/transform/rewrite_unsafe_select.cc @@ -25,10 +25,10 @@ #include #include #include +#include #include #include #include -#include namespace tvm { namespace s_tir { diff --git a/src/s_tir/transform/stmt_extension.cc b/src/s_tir/transform/stmt_extension.cc new file mode 100644 index 000000000000..bcd902270523 --- /dev/null +++ b/src/s_tir/transform/stmt_extension.cc @@ -0,0 +1,132 @@ +/* + * 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/transform/stmt_extension.h" + +#include +#include + +namespace tvm { +namespace s_tir { +using namespace tirx; +namespace { + +UnchangedOr ConvertSSABlock(tirx::SSAStmtMutator* self, const SBlockNode* op, + InplaceMode mode) { + SBlock block = ffi::GetRef(op); + return self->WithScope([&]() -> Stmt { + auto iter_vars = op->iter_vars.Map([&](IterVar iter) { + Var var = self->DefineVar(iter->var); + if (!var.same_as(iter->var)) iter.CopyOnWrite()->var = var.as_or_throw(); + return iter; + }); + auto remap_region = [&](BufferRegion region) { + BufferVar buffer = self->RemapBuffer(region->buffer); + if (!buffer.same_as(region->buffer)) region.CopyOnWrite()->buffer = buffer; + 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 StmtExprMutator::MutateBlock(self, block.get(), + block.unique() ? mode : InplaceMode::kDisallow) + .ValueOrUnchanged(block); + }); +} + +UnchangedOr FlattenBlock(tirx::FlattenStmtMutator* self, const SBlockNode* op, + InplaceMode mode) { + TVM_FFI_ICHECK_EQ(op->match_buffers.size(), 0) + << "Unexpected MatchBufferRegion found during tirx.transform.FlattenBuffer. " + << "All MatchBufferRegion should be removed in tirx.transform.LowerMatchBuffer."; + SBlock block = ffi::GetRef(op); + auto alloc_buffers = op->alloc_buffers; + alloc_buffers.MutateByApply([&](BufferVar buffer) { return self->DefineBuffer(buffer); }); + if (!alloc_buffers.same_as(op->alloc_buffers)) block.CopyOnWrite()->alloc_buffers = alloc_buffers; + auto reads = op->reads; + reads.MutateByApply([&](BufferRegion region) { return self->RewriteRegion(region); }); + if (!reads.same_as(op->reads)) block.CopyOnWrite()->reads = reads; + auto writes = op->writes; + writes.MutateByApply([&](BufferRegion region) { return self->RewriteRegion(region); }); + if (!writes.same_as(op->writes)) block.CopyOnWrite()->writes = writes; + return StmtExprMutator::MutateBlock(self, block.get(), + block.unique() ? mode : InplaceMode::kDisallow) + .ValueOrUnchanged(block); +} + +// These hooks extend each pass's existing native table before it is finalized. +// All ordinary statements and expression remapping remain in the TIRX pass. +TVM_FFI_STATIC_INIT_BLOCK() { + tirx::SSAStmtMutator::RegisterExtension([](tirx::SSAStmtMutator::VTable* table) { + table->ClearDispatch(); + table->SetDispatch( + [](const ffi::Object* node, ObjectMutator* self, InplaceMode mode) { + return ffi::details::UnchangedOrUnsafe::MoveFromTVMFFIAny( + ffi::details::UnchangedOrUnsafe::MoveToTVMFFIAny( + ConvertSSABlock(static_cast(self), + static_cast(node), mode))); + }); + }); + tirx::FlattenStmtMutator::RegisterExtension([](tirx::FlattenStmtMutator::VTable* table) { + table->ClearDispatch(); + table->SetDispatch( + [](const ffi::Object* node, ObjectMutator* self, InplaceMode mode) { + return ffi::details::UnchangedOrUnsafe::MoveFromTVMFFIAny( + ffi::details::UnchangedOrUnsafe::MoveToTVMFFIAny( + FlattenBlock(static_cast(self), + static_cast(node), mode))); + }); + }); + tirx::IndexDomainVisitor::RegisterExtension([](tirx::IndexDomainVisitor::VTable* table) { + table->ClearDispatch(); + table->SetDispatch([](const ffi::Object* node, ObjectVisitor* base) { + auto* self = static_cast(base); + auto* block = static_cast(node); + for (const auto& iter : block->iter_vars) { + self->BindDomain(iter->var, Range::FromMinExtent(iter->dom->min, iter->dom->extent)); + } + return StmtExprVisitor::VisitBlock(self, block); + }); + }); + tirx::StorageAlignVisitor::RegisterExtension([](tirx::StorageAlignVisitor::VTable* table) { + table->ClearDispatch(); + table->SetDispatch([](const ffi::Object* node, ObjectVisitor* base) { + auto* self = static_cast(base); + auto* block = static_cast(node); + auto it = block->annotations.find(attr::buffer_dim_align); + if (it != block->annotations.end()) { + auto annotation = (*it).second.as_or_throw(); + for (const auto& item : annotation) { + self->RecordAlignment(block->writes[item.get<0>()]->buffer.var(), item); + } + } + return StmtExprVisitor::VisitBlock(self, block); + }); + }); +} + +} // namespace +} // namespace s_tir +} // namespace tvm diff --git a/src/s_tir/transform/storage_access.h b/src/s_tir/transform/storage_access.h index 291e0da804b9..83b56d3520b6 100644 --- a/src/s_tir/transform/storage_access.h +++ b/src/s_tir/transform/storage_access.h @@ -27,7 +27,7 @@ #include #include #include -#include +#include #include #include diff --git a/src/s_tir/transform/tensorcore_infer_fragment.cc b/src/s_tir/transform/tensorcore_infer_fragment.cc index 8944dc2cc1b0..bf7d5414cee8 100644 --- a/src/s_tir/transform/tensorcore_infer_fragment.cc +++ b/src/s_tir/transform/tensorcore_infer_fragment.cc @@ -27,8 +27,8 @@ #include #include #include +#include #include -#include #include #include @@ -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 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 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 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 Mutate(ffi::AnyView value, InplaceMode inplace_mode) override { if (value.as()) 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 Mutate_(const AllocBufferNode* op, InplaceMode inplace_mode) final { - Stmt stmt = StmtExprMutator::Mutate_(op, inplace_mode).ValueOrUnchanged(ffi::GetRef(op)); + Stmt stmt = + s_tir::StmtExprMutator::Mutate_(op, inplace_mode).ValueOrUnchanged(ffi::GetRef(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 07c6576bf7bf..08e644f5c121 100644 --- a/src/s_tir/transform/thread_storage_sync.cc +++ b/src/s_tir/transform/thread_storage_sync.cc @@ -25,12 +25,13 @@ #include #include #include +#include #include +#include #include #include #include #include -#include #include diff --git a/src/s_tir/transform/transform_mma_buffer_layout.cc b/src/s_tir/transform/transform_mma_buffer_layout.cc index 4eb31678f3de..3784db34e929 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 #include #include +#include +#include +#include #include #include #include -#include #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 ebdbefc399ed..658b9fe77a47 100644 --- a/src/s_tir/transform/unify_thread_binding.cc +++ b/src/s_tir/transform/unify_thread_binding.cc @@ -24,10 +24,11 @@ #include #include #include +#include #include +#include #include #include -#include #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 8d8e9e5d5433..95cc2fbfff36 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 #include #include +#include #include #include #include #include -#include #include #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 67532c035e77..bd28eddf5659 100644 --- a/src/te/operation/create_primfunc.cc +++ b/src/te/operation/create_primfunc.cc @@ -27,11 +27,11 @@ #include #include #include +#include #include #include #include #include -#include #include #include @@ -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& tensor2buffers) @@ -80,7 +80,7 @@ class TensorLoadToBufferTransformer : public StmtExprMutator { const auto* reduce = op->IsInstance() ? static_cast(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 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(op)) .as_or_throw(); 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& var_map, const std::unordered_map& 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 Mutate(ffi::AnyView value, InplaceMode inplace_mode) final { if (value.as()) 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 Mutate_(const SBlockNode* _block, InplaceMode inplace_mode) final { - SBlock block = StmtExprMutator::Mutate_(_block, inplace_mode) - .ValueOrUnchanged(ffi::GetRef(_block)) - .as_or_throw(); - SBlockNode* n = block.CopyOnWrite(); + UnchangedOr Mutate_(const s_tir::SBlockNode* _block, InplaceMode inplace_mode) final { + s_tir::SBlock block = s_tir::StmtExprMutator::Mutate_(_block, inplace_mode) + .ValueOrUnchanged(ffi::GetRef(_block)) + .as_or_throw(); + s_tir::SBlockNode* n = block.CopyOnWrite(); if (auto opt_ann = n->annotations.Get(topi_attr)) { ffi::Array new_buffers; for (BufferVar buffer : opt_ann.value().as_or_throw>()) { @@ -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(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 CollectOrderedOps(const ffi::Array& arg_list) { diff --git a/src/tirx/analysis/verify_tirx_well_formed.cc b/src/tirx/analysis/verify_tirx_well_formed.cc index 66adcebb0b3e..d1c9c0d58573 100644 --- a/src/tirx/analysis/verify_tirx_well_formed.cc +++ b/src/tirx/analysis/verify_tirx_well_formed.cc @@ -50,12 +50,10 @@ class ExecScopeVerifier : public Verifier { 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 EnterExtensionStmt(const ffi::Object* op, ffi::reflection::AccessPath path) override { + Verify(false) << "TIRxError: " << op->GetTypeKey() << " is not allowed in tirx=True mode at " + << path; + return false; } void Dispatch_(const tirx::TilePrimitiveCallNode* op, ffi::reflection::AccessPath path) override { @@ -132,12 +130,10 @@ class LayoutVerifier : public Verifier { 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 EnterExtensionStmt(const ffi::Object* op, ffi::reflection::AccessPath path) override { + Verify(false) << "TIRxError: " << op->GetTypeKey() << " is not allowed in tirx=True mode at " + << path; + return false; } }; @@ -148,12 +144,10 @@ class AsyncStructsVerifier : public Verifier { 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 EnterExtensionStmt(const ffi::Object* op, ffi::reflection::AccessPath path) override { + Verify(false) << "TIRxError: " << op->GetTypeKey() << " is not allowed in tirx=True mode at " + << path; + return false; } }; @@ -164,12 +158,10 @@ class DeviceFuncVerifier : public Verifier { 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 EnterExtensionStmt(const ffi::Object* op, ffi::reflection::AccessPath path) override { + Verify(false) << "TIRxError: " << op->GetTypeKey() << " is not allowed in tirx=True mode at " + << path; + return false; } }; diff --git a/src/tirx/analysis/verify_well_formed.cc b/src/tirx/analysis/verify_well_formed.cc index 0123c463b64c..ea49da4c3984 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 #include #include @@ -41,111 +43,15 @@ 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(assert_mode); - verifier->Visit(func->body); - return !verifier->has_error_; - } - - explicit BlockVarAccessVerifier(bool assert_mode) : assert_mode_(assert_mode) {} - - private: - ffi::Optional Visit(ffi::AnyView stmt) final { - if (!has_error_) { - return StmtExprVisitor::Visit(stmt); - } - return std::nullopt; - } - - ffi::Optional 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 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 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 { - 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 loop_vars_; - /*! \brief Whether it's in assert mode. */ - bool assert_mode_; - /*! \brief Current nested block stack level. */ - std::vector block_stack_; - /*! \brief Whether there is error. */ - bool has_error_{false}; -}; +namespace { +std::vector& WellFormedExtensions() { + static std::vector extensions; + return extensions; +} +} // namespace +void RegisterWellFormedExtension(bool (*verify)(const PrimFunc&, bool)) { + WellFormedExtensions().push_back(verify); +} class UndefinedVarVerifier : public Verifier { public: @@ -421,8 +327,8 @@ class SingleEnvThreadVerifier : public Verifier { }; bool VerifyWellFormed(const PrimFunc& func, bool assert_mode) { - if (!BlockVarAccessVerifier::Verify(func, assert_mode)) { - return false; + for (auto verify : WellFormedExtensions()) { + if (!verify(func, assert_mode)) return false; } if (!UndefinedVarVerifier::Verify(func, assert_mode)) return false; diff --git a/src/tirx/analysis/verify_well_formed.h b/src/tirx/analysis/verify_well_formed.h new file mode 100644 index 000000000000..0c19522d2621 --- /dev/null +++ b/src/tirx/analysis/verify_well_formed.h @@ -0,0 +1,29 @@ +/* + * 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_TIRX_ANALYSIS_VERIFY_WELL_FORMED_H_ +#define TVM_TIRX_ANALYSIS_VERIFY_WELL_FORMED_H_ +#include +namespace tvm { +namespace tirx { +// Register supplementary dialect checks before verification is first invoked. +void RegisterWellFormedExtension(bool (*verify)(const PrimFunc&, bool)); +} // namespace tirx +} // namespace tvm +#endif diff --git a/src/tirx/ir/data_type_rewriter.cc b/src/tirx/ir/data_type_rewriter.cc index 2b7504513de7..3f6ed8956315 100644 --- a/src/tirx/ir/data_type_rewriter.cc +++ b/src/tirx/ir/data_type_rewriter.cc @@ -27,7 +27,6 @@ #include #include #include -#include #include #include @@ -44,6 +43,33 @@ namespace tvm { namespace tirx { using namespace tvm::prim; +namespace { +std::vector& IndexDataTypeRewriterExtensions() { + static std::vector extensions; + return extensions; +} +} // namespace +void IndexDataTypeRewriter::RegisterExtension(void (*init)(VTable*)) { + IndexDataTypeRewriterExtensions().push_back(init); +} +void IndexDataTypeRewriter::InitVTable(VTable* vtable) { + DataTypeLegalizer::InitVTable(vtable); + for (auto init : IndexDataTypeRewriterExtensions()) init(vtable); +} + +namespace { +std::vector& DataTypeLegalizerExtensions() { + static std::vector extensions; + return extensions; +} +} // namespace +void DataTypeLegalizer::RegisterExtension(void (*init)(VTable*)) { + DataTypeLegalizerExtensions().push_back(init); +} +void DataTypeLegalizer::InitVTable(VTable* vtable) { + StmtExprMutator::InitVTable(vtable); + for (auto init : DataTypeLegalizerExtensions()) init(vtable); +} UnchangedOr DataTypeLegalizer::Mutate_(const ForNode* op, InplaceMode inplace_mode) { auto result = StmtExprMutator::Mutate_(op, inplace_mode); @@ -74,51 +100,8 @@ UnchangedOr DataTypeLegalizer::Mutate_(const ForNode* op, InplaceMode inpl return For(n); } -UnchangedOr DataTypeLegalizer::Mutate_(const SBlockRealizeNode* op, - InplaceMode inplace_mode) { - SBlockRealize realize = StmtExprMutator::Mutate_(op, inplace_mode) - .ValueOrUnchanged(ffi::GetRef(op)) - .as_or_throw(); - ffi::Array new_iter_values; - bool changed = false; - for (int i = 0; i < static_cast(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 DataTypeLegalizer::Mutate_(const SBlockNode* op, InplaceMode inplace_mode) { - SBlock new_block = StmtExprMutator::Mutate_(op, inplace_mode) - .ValueOrUnchanged(ffi::GetRef(op)) - .as_or_throw(); - ffi::Array 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 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 == "virtual_thread") { Stmt s = StmtExprMutator::Mutate_(op, inplace_mode).ValueOrUnchanged(ffi::GetRef(op)); op = s.as(); TVM_FFI_ICHECK(op != nullptr) << "Expected type to be AttrStmtNode" @@ -374,7 +357,7 @@ UnchangedOr DataTypeLegalizer::Mutate_(const CallNode* op, InplaceMode inp } UnchangedOr 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 == "virtual_thread") { bool is_enabled = is_enabled_; is_enabled_ = true; auto stmt = DataTypeLegalizer::Mutate_(op, inplace_mode); @@ -392,165 +375,6 @@ UnchangedOr IndexDataTypeRewriter::Mutate(ffi::AnyView value, InplaceM return result; } -UnchangedOr 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>>() - .ValueOrUnchanged(op->iter_values); - is_enabled_ = is_enabled; - SBlock new_body = - this->Mutate(op->block, inplace_mode).ValueOrUnchanged(op->block).as_or_throw(); - if (!new_predicate_unchanged || !new_iter_values.same_as(op->iter_values) || - !new_body.same_as(op->block)) { - SBlockRealize new_block_realize = ffi::GetRef(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 IndexDataTypeRewriter::Mutate_(const SBlockNode* op, InplaceMode inplace_mode) { - auto new_alloc_buffers = WithDefRegionKind(kTVMFFIDefRegionKindSimple, [&] { - return Mutate(op->alloc_buffers, inplace_mode) - .as_or_throw>>() - .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>() - .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 new_reads = op->reads.Map( - [this](const TensorRegion& buffer_region) { return this->VisitBufferRegion(buffer_region); }); - ffi::Array new_writes = op->writes.Map( - [this](const TensorRegion& buffer_region) { return this->VisitBufferRegion(buffer_region); }); - ffi::Array new_iter_vars = - op->iter_vars.Map([this](const IterVar& iter_var) { return this->VisitIterVar(iter_var); }); - ffi::Optional new_init = std::nullopt; - if (op->init.has_value()) { - new_init = this->Mutate(op->init.value(), inplace_mode).ValueOrUnchanged(op->init.value()); - } - ffi::Map 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(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 IndexDataTypeRewriter::VisitBlockAnnotations( - const ffi::Map& annotations) { - auto new_annotations = annotations; - - std::function f_mutate_obj = [this, &f_mutate_obj](const Any& obj) -> Any { - if (obj == nullptr) { - return obj; - } - if (auto var = obj.as(); var && var.value()->ty.as()) { - BufferVar buffer(var.value()); - if (BufferVar new_buffer = Mutate(buffer, InplaceMode::kDisallow) - .as_or_throw>() - .ValueOrUnchanged(buffer); - !new_buffer.same_as(buffer)) { - return new_buffer; - } - } else if (obj.as()) { - return obj.as_or_throw>().Map(f_mutate_obj); - } - return obj; - }; - for (const auto& [key, value] : annotations) { - if (auto opt_object_ref = value.as()) { - 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(); - 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(), InplaceMode::kDisallow) - .as_or_throw>() - .ValueOrUnchanged(buffer_region->source.as_or_throw()); - - 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()) || - !new_region.same_as(buffer_region->region)) { - return BufferRegion(remapped_buffer, new_region); - } else { - return buffer_region; - } -} - UnchangedOr IndexDataTypeRewriter::Mutate_(const BufferStoreNode* op, InplaceMode inplace_mode) { BufferStore store = ffi::GetRef(op); diff --git a/src/tirx/ir/data_type_rewriter.h b/src/tirx/ir/data_type_rewriter.h index c2cca01ed8d6..584539e02910 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,20 @@ namespace tirx { */ class DataTypeLegalizer : public StmtExprMutator { public: + using StmtExprMutator::VTable; + // Dialect-owned handlers use nested access to the active traversal context. + class Extension; + // Register during library initialization, before the finalized table is first used. + static void RegisterExtension(void (*init)(VTable*)); + TVM_DEFINE_OBJECT_FUNCTOR_DEFAULT_CONSTRUCTOR(DataTypeLegalizer, StmtExprMutator) using StmtExprMutator::Mutate; using StmtExprMutator::Mutate_; protected: + explicit DataTypeLegalizer(const VTable* vtable) : StmtExprMutator(vtable) {} + static void InitVTable(VTable* vtable); UnchangedOr Mutate_(const ForNode* op, InplaceMode inplace_mode) override; UnchangedOr Mutate_(const AttrStmtNode* op, InplaceMode inplace_mode) override; - UnchangedOr Mutate_(const SBlockRealizeNode* op, InplaceMode inplace_mode) override; - UnchangedOr Mutate_(const SBlockNode* op, InplaceMode inplace_mode) override; UnchangedOr Mutate_(const BindNode* op, InplaceMode inplace_mode) override; UnchangedOr Mutate_(const prim::SelectNode* op, InplaceMode inplace_mode) override; UnchangedOr Mutate_(const prim::RampNode* op, InplaceMode inplace_mode) override; @@ -99,14 +103,19 @@ class DataTypeLegalizer : public StmtExprMutator { */ class IndexDataTypeRewriter : public DataTypeLegalizer { public: + using DataTypeLegalizer::VTable; + // Dialect-owned handlers use nested access to the active traversal context. + class Extension; + // Register during library initialization, before the finalized table is first used. + static void RegisterExtension(void (*init)(VTable*)); + TVM_DEFINE_OBJECT_FUNCTOR_DEFAULT_CONSTRUCTOR(IndexDataTypeRewriter, DataTypeLegalizer) using DataTypeLegalizer::Mutate; using DataTypeLegalizer::Mutate_; protected: + static void InitVTable(VTable* vtable); using Parent = DataTypeLegalizer; UnchangedOr Mutate(ffi::AnyView value, InplaceMode inplace_mode) override; - UnchangedOr Mutate_(const SBlockRealizeNode* op, InplaceMode inplace_mode) override; - UnchangedOr Mutate_(const SBlockNode* op, InplaceMode inplace_mode) override; UnchangedOr Mutate_(const BufferStoreNode* op, InplaceMode inplace_mode) override; UnchangedOr Mutate_(const AttrStmtNode* op, InplaceMode inplace_mode) override; UnchangedOr Mutate_(const TensorLoadNode* op, InplaceMode inplace_mode) override; @@ -124,10 +133,9 @@ class IndexDataTypeRewriter : public DataTypeLegalizer { UnchangedOr Mutate_(const ForNode* op, InplaceMode inplace_mode) override; - ffi::Map VisitBlockAnnotations( - const ffi::Map& annotations); - TensorRegion VisitBufferRegion(const TensorRegion& region); - IterVar VisitIterVar(const IterVar& iter_var); + // Dialect allocation hooks retain common buffer validation in the active rewriter. + virtual void ValidateAllocation(const BufferVar& buffer) {} + // indicator of index expr to rewrite bool is_enabled_{false}; // indicator of condition diff --git a/src/tirx/ir/ir_mutator_with_analyzer.cc b/src/tirx/ir/ir_mutator_with_analyzer.cc index e26fdf5282e1..d11c6ee224df 100644 --- a/src/tirx/ir/ir_mutator_with_analyzer.cc +++ b/src/tirx/ir/ir_mutator_with_analyzer.cc @@ -25,7 +25,6 @@ #include #include #include -#include #include #include #include @@ -34,6 +33,32 @@ namespace tvm { namespace tirx { +namespace { +std::vector& IRMutatorWithAnalyzerExtensions() { + static std::vector extensions; + return extensions; +} +} // namespace + +void IRMutatorWithAnalyzer::RegisterExtension(void (*init)(VTable*)) { + IRMutatorWithAnalyzerExtensions().push_back(init); +} + +void IRMutatorWithAnalyzer::InitVTable(VTable* vtable) { + StmtExprMutator::InitVTable(vtable); + for (auto init : IRMutatorWithAnalyzerExtensions()) init(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 +138,6 @@ UnchangedOr IRMutatorWithAnalyzer::Mutate_(const ForNode* op, InplaceMode }); } -UnchangedOr IRMutatorWithAnalyzer::Mutate_(const SBlockNode* op, InplaceMode inplace_mode) { - return constraint_scope_.WithNewScope([&]() -> UnchangedOr { - 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 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 +216,7 @@ UnchangedOr IRMutatorWithAnalyzer::Mutate_(const IfThenElseNode* op, UnchangedOr IRMutatorWithAnalyzer::Mutate_(const AttrStmtNode* op, InplaceMode inplace_mode) { return constraint_scope_.WithNewScope([&]() -> UnchangedOr { - 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 == "virtual_thread") { IterVar iv = op->node.as_or_throw(); 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 20a6dd258788..9cb079f4a3b9 100644 --- a/src/tirx/ir/ir_mutator_with_analyzer.h +++ b/src/tirx/ir/ir_mutator_with_analyzer.h @@ -34,6 +34,7 @@ #include #include +#include namespace tvm { namespace tirx { @@ -49,14 +50,20 @@ namespace tirx { */ class IRMutatorWithAnalyzer : public StmtExprMutator { public: + using StmtExprMutator::VTable; + // Dialect-owned handlers use nested access to the active traversal context. + class Extension; + // Extensions register during library initialization, before constructing a visitor. + static void RegisterExtension(void (*init)(VTable*)); 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 Mutate_(const ForNode* op, InplaceMode inplace_mode) override; - UnchangedOr Mutate_(const SBlockNode* op, InplaceMode inplace_mode) override; UnchangedOr Mutate_(const BindNode* op, InplaceMode inplace_mode) override; UnchangedOr Mutate_(const IfThenElseNode* op, InplaceMode inplace_mode) override; UnchangedOr Mutate_(const AttrStmtNode* op, InplaceMode inplace_mode) override; @@ -66,6 +73,10 @@ class IRMutatorWithAnalyzer : public StmtExprMutator { UnchangedOr 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 9bdc895722cc..c944a02ba5f0 100644 --- a/src/tirx/ir/ir_visitor_with_analyzer.cc +++ b/src/tirx/ir/ir_visitor_with_analyzer.cc @@ -24,13 +24,27 @@ #include #include -#include #include #include #include namespace tvm { namespace tirx { +namespace { +std::vector& IRVisitorWithAnalyzerExtensions() { + static std::vector extensions; + return extensions; +} +} // namespace + +void IRVisitorWithAnalyzer::RegisterExtension(void (*init)(VTable*)) { + IRVisitorWithAnalyzerExtensions().push_back(init); +} + +void IRVisitorWithAnalyzer::InitVTable(VTable* vtable) { + StmtExprVisitor::InitVTable(vtable); + for (auto init : IRVisitorWithAnalyzerExtensions()) init(vtable); +} ffi::Optional IRVisitorWithAnalyzer::Visit_(const ForNode* op) { return constraint_scope_.WithNewScope([&]() -> ffi::Optional { @@ -47,15 +61,6 @@ ffi::Optional IRVisitorWithAnalyzer::Visit_(const ForNode* op) { }); } -ffi::Optional IRVisitorWithAnalyzer::Visit_(const SBlockNode* op) { - return constraint_scope_.WithNewScope([&]() -> ffi::Optional { - for (const auto& iter_var : op->iter_vars) { - analyzer_->Bind(iter_var->var, iter_var->dom); - } - return StmtExprVisitor::Visit_(op); - }); -} - ffi::Optional IRVisitorWithAnalyzer::Visit_(const BindNode* op) { TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->Visit(op->value)); if (ffi::Optional value = op->value.as()) { @@ -89,7 +94,7 @@ ffi::Optional IRVisitorWithAnalyzer::Visit_(const IfThenElseNode ffi::Optional IRVisitorWithAnalyzer::Visit_(const AttrStmtNode* op) { return constraint_scope_.WithNewScope([&]() -> ffi::Optional { - 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 == "virtual_thread") { IterVar iv = op->node.as_or_throw(); TVM_FFI_ICHECK_NE(iv->thread_tag.length(), 0U); analyzer_->Bind(iv->var, Range::FromMinExtent(IntImm(op->value.ty(), 0), op->value)); diff --git a/src/tirx/ir/ir_visitor_with_analyzer.h b/src/tirx/ir/ir_visitor_with_analyzer.h index 33a895acf6bd..619d388d4d7f 100644 --- a/src/tirx/ir/ir_visitor_with_analyzer.h +++ b/src/tirx/ir/ir_visitor_with_analyzer.h @@ -31,17 +31,25 @@ #include #include +#include + namespace tvm { namespace tirx { class IRVisitorWithAnalyzer : public StmtExprVisitor { public: + using StmtExprVisitor::VTable; + // Dialect-owned handlers use nested access to the active traversal context. + class Extension; + // Extensions register during library initialization, before constructing a visitor. + static void RegisterExtension(void (*init)(VTable*)); + TVM_DEFINE_OBJECT_FUNCTOR_DEFAULT_CONSTRUCTOR(IRVisitorWithAnalyzer, StmtExprVisitor) + PrimExpr Simplify(const PrimExpr& expr) { return analyzer_->Simplify(expr); } using StmtExprVisitor::Visit_; ffi::Optional Visit_(const ForNode* op); - ffi::Optional Visit_(const SBlockNode* op); ffi::Optional Visit_(const BindNode* op); ffi::Optional Visit_(const IfThenElseNode* op); ffi::Optional Visit_(const AttrStmtNode* op); @@ -54,6 +62,8 @@ class IRVisitorWithAnalyzer : public StmtExprVisitor { // condition. protected: + static void InitVTable(VTable* vtable); + explicit IRVisitorWithAnalyzer(const VTable* vtable) : StmtExprVisitor(vtable) {} /*! \brief internal analyzer field. */ arith::Analyzer analyzer_; diff --git a/src/tirx/ir/script/script_complete.cc b/src/tirx/ir/script/script_complete.cc index 7c39a8e09889..e37c2e10a497 100644 --- a/src/tirx/ir/script/script_complete.cc +++ b/src/tirx/ir/script/script_complete.cc @@ -26,7 +26,9 @@ #include #include +#include #include +#include #include #include @@ -36,12 +38,12 @@ namespace tvm { namespace tirx { /*! \brief Generate surrounding loops automatically */ -class ScriptCompleter : public StmtExprMutator { +class ScriptCompleter : public s_tir::StmtExprMutator { public: - using StmtExprMutator::Mutate; + using s_tir::StmtExprMutator::Mutate; UnchangedOr Mutate(ffi::AnyView value, InplaceMode inplace_mode) final { if (value.as()) return ffi::Unchanged(); - return StmtExprMutator::Mutate(value, inplace_mode); + return s_tir::StmtExprMutator::Mutate(value, inplace_mode); } explicit ScriptCompleter(ffi::Map* buffer_var_map, bool s_tir = false) @@ -49,16 +51,16 @@ class ScriptCompleter : public StmtExprMutator { private: ffi::Map* buffer_var_map_; - UnchangedOr Mutate_(const SBlockRealizeNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const s_tir::SBlockRealizeNode* op, InplaceMode inplace_mode) final { for (const PrimExpr& value : op->iter_values) { PrimType value_ty = value.ty(); TVM_FFI_ICHECK(value_ty.code() == DLDataTypeCode::kDLInt) << "BlockRealize iter_value expected a IntImm, but got " << value_ty->dtype; } - return StmtExprMutator::Mutate_(op, inplace_mode); + return s_tir::StmtExprMutator::Mutate_(op, inplace_mode); } - UnchangedOr Mutate_(const SBlockNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const s_tir::SBlockNode* op, InplaceMode inplace_mode) final { // Buffers allocated in the block can be accessed by its body. for (const auto& alloc_buffer : op->alloc_buffers) { buffer_var_map_->Set(alloc_buffer.var(), alloc_buffer); @@ -70,9 +72,9 @@ class ScriptCompleter : public StmtExprMutator { bool is_root_block = this->is_root_block_; this->is_root_block_ = false; - SBlock block = StmtExprMutator::Mutate_(op, inplace_mode) - .ValueOrUnchanged(ffi::GetRef(op)) - .as_or_throw(); + s_tir::SBlock block = s_tir::StmtExprMutator::Mutate_(op, inplace_mode) + .ValueOrUnchanged(ffi::GetRef(op)) + .as_or_throw(); this->is_root_block_ = is_root_block; // Remove buffers allocated inside block to detect its access region @@ -117,7 +119,7 @@ class ScriptCompleter : public StmtExprMutator { if (!buffer_var_map_->count(op->buffer.var())) { buffer_var_map_->Set(op->buffer.var(), op->buffer); } - return StmtExprMutator::Mutate_(op, inplace_mode); + return s_tir::StmtExprMutator::Mutate_(op, inplace_mode); } UnchangedOr Mutate_(const DeclBufferNode* op, InplaceMode inplace_mode) final { @@ -125,7 +127,7 @@ class ScriptCompleter : public StmtExprMutator { if (!buffer_var_map_->count(op->buffer.var())) { buffer_var_map_->Set(op->buffer.var(), op->buffer); } - return StmtExprMutator::Mutate_(op, inplace_mode); + return s_tir::StmtExprMutator::Mutate_(op, inplace_mode); } bool is_root_block_ = true; @@ -152,19 +154,19 @@ PrimFunc ScriptComplete(PrimFunc func, const ffi::Array& root_allocat if (root_allocates.size()) { return true; } - auto* block_realize = func->body.as(); + auto* block_realize = func->body.as(); if (block_realize && block_realize->block->iter_vars.size()) { return true; } - if (!block_realize && ContainsNode(func->body)) { + if (!block_realize && ContainsNode(func->body)) { return true; } return false; }(); if (s_tir && should_insert_root) { - SBlock root_block({}, {}, {}, "root", std::move(res), std::nullopt, root_allocates); - res = SBlockRealize({}, IntImm::Bool(true), std::move(root_block)); + s_tir::SBlock root_block({}, {}, {}, "root", std::move(res), std::nullopt, root_allocates); + res = s_tir::SBlockRealize({}, IntImm::Bool(true), std::move(root_block)); } // generate surrounding loops automatically diff --git a/src/tirx/ir/specialize.cc b/src/tirx/ir/specialize.cc index 899b32a3a04d..9f5d6ee61718 100644 --- a/src/tirx/ir/specialize.cc +++ b/src/tirx/ir/specialize.cc @@ -21,12 +21,13 @@ * \file src/tirx/ir/specialize.cc * \brief Specialize parameters of PrimFunc. */ +#include "specialize.h" + #include #include #include #include #include -#include #include #include #include @@ -35,6 +36,7 @@ #include #include +#include #include "../transform/ir_utils.h" @@ -43,6 +45,17 @@ namespace tirx { using VarMap = std::unordered_map; +namespace { +std::vector& BufferPlannerExtensions() { + static std::vector extensions; + return extensions; +} +} // namespace + +void RegisterSpecializeBufferPlannerExtension(void (*init)(SpecializeVisitorVTable*)) { + BufferPlannerExtensions().push_back(init); +} + /**************** Helper functions ****************/ /*! \brief Helper function to check whether the given var is in function parameter list. */ @@ -155,9 +168,21 @@ class PrimFuncSpecializer : public StmtExprMutator { public: using StmtExprVisitor::Visit_; - explicit BufferPlanner(PrimFuncSpecializer* specializer) : specializer_(specializer) {} + explicit BufferPlanner(PrimFuncSpecializer* specializer) + : StmtExprVisitor(GlobalVTable()), specializer_(specializer) {} private: + static const VTable* GlobalVTable() { + static const VTable table = [] { + VTable table; + StmtExprVisitor::InitVTable(&table); + for (auto init : BufferPlannerExtensions()) init(&table); + table.Finalize(); + return table; + }(); + return &table; + } + ffi::Optional Visit_(const VarNode* op) final { if (op->ty.as()) { if (def_region_kind() == kTVMFFIDefRegionKindSimple) { @@ -177,35 +202,6 @@ class PrimFuncSpecializer : public StmtExprMutator { return Visit(op->data); } - ffi::Optional Visit_(const SBlockNode* op) final { - // Block allocations were planned before all other block children by the specializer. - for (const BufferVar& buffer : op->alloc_buffers) { - TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->WithDefRegionKind( - kTVMFFIDefRegionKindSimple, [&]() { return this->Visit(buffer); })); - TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(VisitBufferMetadata(buffer)); - } - for (const IterVar& iter : op->iter_vars) { - TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Visit(iter->dom->min)); - TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Visit(iter->dom->extent)); - } - for (const TensorRegion& region : op->reads) { - TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Visit(region)); - } - for (const TensorRegion& region : op->writes) { - TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Visit(region)); - } - for (const MatchBufferRegion& match : op->match_buffers) { - TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->WithDefRegionKind( - kTVMFFIDefRegionKindSimple, [&]() { return this->Visit(match->buffer); })); - TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(VisitBufferMetadata(match->buffer)); - TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Visit(match->source)); - } - if (op->init.has_value()) { - TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Visit(op->init.value())); - } - return Visit(op->body); - } - PrimFuncSpecializer* specializer_; }; @@ -347,7 +343,7 @@ class PrimFuncSpecializer : public StmtExprMutator { << "(see discussion on https://github.com/apache/tvm/pull/14565 for more details). " << "Please add a definition for this buffer, " << "either as a BufferType-annotated PrimFunc parameter, " - << "in a tirx::SBlock's alloc_buffer, " + << "in a block's buffer allocations, " << "or in a DeclBuffer statement."; } diff --git a/src/tirx/ir/specialize.h b/src/tirx/ir/specialize.h new file mode 100644 index 000000000000..4e104f61c3ee --- /dev/null +++ b/src/tirx/ir/specialize.h @@ -0,0 +1,39 @@ +/* + * 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_TIRX_IR_SPECIALIZE_H_ +#define TVM_TIRX_IR_SPECIALIZE_H_ + +#include + +namespace tvm { +namespace tirx { + +using SpecializeVisitorVTable = + ObjectFunctor(const ffi::Object*, ObjectVisitor*)>; + +// Register dialect-specific buffer planning before any specialization is run. +// Initializers extend the inherited native traversal table without exposing the +// specializer's private buffer remapping and declaration state. +void RegisterSpecializeBufferPlannerExtension(void (*init)(SpecializeVisitorVTable*)); + +} // namespace tirx +} // namespace tvm + +#endif // TVM_TIRX_IR_SPECIALIZE_H_ diff --git a/src/tirx/ir/stmt.cc b/src/tirx/ir/stmt.cc index 1699a8b686b6..8d7823594a85 100644 --- a/src/tirx/ir/stmt.cc +++ b/src/tirx/ir/stmt.cc @@ -721,173 +721,47 @@ TVMFFIAny BufferRegionTypeMaybeInplaceMutate(ffi::StructuralMutatorObj*, ffi::An return ffi::Unchanged().CopyToTVMFFIAny(); } -TVMFFIAny MatchBufferRegionVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept { - const MatchBufferRegionNode* self = - ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck( - 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)); +TVMFFIAny BufferRegionVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept { + const BufferRegionNode* self = + ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value); + TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->buffer)); + TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->region)); return ffi::AnyView(nullptr).CopyToTVMFFIAny(); } -TVMFFIAny MatchBufferRegionMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept { - const MatchBufferRegionNode* self = - ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck( - value); +TVMFFIAny BufferRegionMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept { + const BufferRegionNode* self = + ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value); TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, mapped_buffer, - mutator->WithDefRegionKind(kTVMFFIDefRegionKindSimple, [&]() { - return mutator->MutateExpected(self->buffer); - })); - TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, mapped_source, - mutator->MutateExpected(self->source)); + mutator->MutateExpected(self->buffer)); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr>, mapped_region, + mutator->MutateExpected(self->region)); if (mapped_buffer.UnchangedOrSameAs(self->buffer) && - mapped_source.UnchangedOrSameAs(self->source)) { + mapped_region.UnchangedOrSameAs(self->region)) { return ffi::Unchanged().CopyToTVMFFIAny(); } - ffi::ObjectPtr copy = ffi::make_object(*self); + ffi::ObjectPtr copy = ffi::make_object(*self); copy->buffer = std::move(mapped_buffer).ValueOrUnchanged(std::move(copy->buffer)); - copy->source = std::move(mapped_source).ValueOrUnchanged(std::move(copy->source)); + copy->region = std::move(mapped_region).ValueOrUnchanged(std::move(copy->region)); return ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(copy))); } -TVMFFIAny MatchBufferRegionMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator, - ffi::AnyView value) noexcept { - MatchBufferRegionNode* self = const_cast( - ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck( - value)); - TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, mapped_buffer, - mutator->WithDefRegionKind(kTVMFFIDefRegionKindSimple, [&]() { - return mutator->MutateExpected(self->buffer, - ffi::InplaceMode::kAllow); - })); +TVMFFIAny BufferRegionMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator, + ffi::AnyView value) noexcept { + BufferRegionNode* self = const_cast( + ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value)); TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN( - ffi::UnchangedOr, mapped_source, - mutator->MutateExpected(self->source, ffi::InplaceMode::kAllow)); + ffi::UnchangedOr, mapped_buffer, + mutator->MutateExpected(self->buffer, ffi::InplaceMode::kAllow)); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN( + ffi::UnchangedOr>, mapped_region, + mutator->MutateExpected(self->region, ffi::InplaceMode::kAllow)); if (mapped_buffer.UnchangedOrSameAs(self->buffer) && - mapped_source.UnchangedOrSameAs(self->source)) { + mapped_region.UnchangedOrSameAs(self->region)) { 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 { - // skips: name_hint - const SBlockNode* self = - ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value); - TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->iter_vars)); - 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->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->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 { - // skips: name_hint - const SBlockNode* self = - ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value); - TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr>, mapped_iter_vars, - mutator->MutateExpected(self->iter_vars)); - TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr>, mapped_reads, - mutator->MutateExpected(self->reads)); - TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr>, mapped_writes, - mutator->MutateExpected(self->writes)); - TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr>, mapped_alloc_buffers, - mutator->WithDefRegionKind(kTVMFFIDefRegionKindSimple, [&]() { - return mutator->MutateExpected(self->alloc_buffers); - })); - TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr>, - mapped_match_buffers, - mutator->MutateExpected(self->match_buffers)); - using AnnotationMap = ffi::Map; - TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, mapped_annotations, - mutator->MutateExpected(self->annotations)); - TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr>, mapped_init, - mutator->MutateExpected(self->init)); - TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, 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 copy = ffi::make_object(*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 { - // skips: name_hint - SBlockNode* self = const_cast( - ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value)); - TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN( - ffi::UnchangedOr>, mapped_iter_vars, - mutator->MutateExpected(self->iter_vars, ffi::InplaceMode::kAllow)); - TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr>, mapped_reads, - mutator->MutateExpected(self->reads, ffi::InplaceMode::kAllow)); - TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN( - ffi::UnchangedOr>, mapped_writes, - mutator->MutateExpected(self->writes, ffi::InplaceMode::kAllow)); - TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr>, mapped_alloc_buffers, - mutator->WithDefRegionKind(kTVMFFIDefRegionKindSimple, [&]() { - return mutator->MutateExpected(self->alloc_buffers, - ffi::InplaceMode::kAllow); - })); - TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN( - ffi::UnchangedOr>, mapped_match_buffers, - mutator->MutateExpected(self->match_buffers, ffi::InplaceMode::kAllow)); - using AnnotationMap = ffi::Map; - TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN( - ffi::UnchangedOr, mapped_annotations, - mutator->MutateExpected(self->annotations, ffi::InplaceMode::kAllow)); - TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr>, mapped_init, - mutator->MutateExpected(self->init, ffi::InplaceMode::kAllow)); - TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, 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(); + if (!mapped_region.IsUnchanged()) self->region = std::move(mapped_region).ValueUnchecked(); return ffi::Unchanged().CopyToTVMFFIAny(); } @@ -925,61 +799,6 @@ TVMFFIAny ScopeIdDefStmtMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator, return ffi::Unchanged().CopyToTVMFFIAny(); } -TVMFFIAny SBlockRealizeVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept { - const SBlockRealizeNode* self = - ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(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(value); - TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr>, mapped_iter_values, - mutator->MutateExpected(self->iter_values)); - TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, mapped_predicate, - mutator->MutateExpected(self->predicate)); - TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, 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 copy = ffi::make_object(*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( - ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value)); - TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN( - ffi::UnchangedOr>, mapped_iter_values, - mutator->MutateExpected(self->iter_values, ffi::InplaceMode::kAllow)); - TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN( - ffi::UnchangedOr, mapped_predicate, - mutator->MutateExpected(self->predicate, ffi::InplaceMode::kAllow)); - TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, 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 TVM_FFI_STATIC_INIT_BLOCK() { StmtNode::RegisterReflection(); } @@ -1553,129 +1372,6 @@ TensorRegion BufferRegionFromPoint(BufferVar buffer, ffi::Array indice return BufferRegion(buffer, region); } -// MatchBufferRegion -MatchBufferRegion::MatchBufferRegion(BufferVar buffer, TensorRegion source) { - const BufferVar& source_buffer = source->source.as_or_throw(); - 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 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()) { - 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 node = ffi::make_object(); - 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() - .attr(refl::type_attr::kStructuralVisit, reinterpret_cast(&MatchBufferRegionVisit)) - .attr(refl::type_attr::kStructuralMutate, reinterpret_cast(&MatchBufferRegionMutate)) - .attr(refl::type_attr::kStructuralMaybeInplaceMutate, - reinterpret_cast(&MatchBufferRegionMaybeInplaceMutate)); - - refl::GlobalDef().def("tirx.MatchBufferRegion", [](BufferVar buffer, TensorRegion source) { - return MatchBufferRegion(buffer, source); - }); -} - -// Block -SBlock::SBlock(ffi::Array iter_vars, ffi::Array reads, - ffi::Array writes, ffi::String name_hint, Stmt body, - ffi::Optional init, ffi::Array alloc_buffers, - ffi::Array match_buffers, ffi::Map annotations, - Span span) { - for (const auto& regions : {reads, writes}) { - for (const TensorRegion& region : regions) { - const auto buffer = region->source.as_or_throw(); - TVM_FFI_ICHECK_EQ(buffer->shape.size(), region->region.size()) - << "SBlock region must match its buffer rank"; - } - } - ffi::ObjectPtr node = ffi::make_object(); - 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 alloc_buffers, Span span) { - ffi::ObjectPtr node = ffi::make_object(); - 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() - .attr(refl::type_attr::kStructuralVisit, reinterpret_cast(&SBlockVisit)) - .attr(refl::type_attr::kStructuralMutate, reinterpret_cast(&SBlockMutate)) - .attr(refl::type_attr::kStructuralMaybeInplaceMutate, - reinterpret_cast(&SBlockMaybeInplaceMutate)); - - refl::GlobalDef().def("tirx.SBlock", - [](ffi::Array iter_vars, ffi::Array reads, - ffi::Array writes, ffi::String name_hint, Stmt body, - ffi::Optional init, ffi::Array alloc_buffers, - ffi::Array match_buffers, - ffi::Map annotations, Span span) { - return SBlock(iter_vars, reads, writes, name_hint, body, init, - alloc_buffers, match_buffers, annotations, span); - }); -} - // ScopeIdDefStmt ScopeIdDefStmt::ScopeIdDefStmt(ScopeIdDef def, Span span) { TVM_FFI_ICHECK(def.defined()); @@ -1698,37 +1394,6 @@ TVM_FFI_STATIC_INIT_BLOCK() { [](ScopeIdDef def, Span span) { return ScopeIdDefStmt(def, span); }); } -// BlockRealize -SBlockRealize::SBlockRealize(ffi::Array 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 node = ffi::make_object(); - 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() - .attr(refl::type_attr::kStructuralVisit, reinterpret_cast(&SBlockRealizeVisit)) - .attr(refl::type_attr::kStructuralMutate, reinterpret_cast(&SBlockRealizeMutate)) - .attr(refl::type_attr::kStructuralMaybeInplaceMutate, - reinterpret_cast(&SBlockRealizeMaybeInplaceMutate)); - - refl::GlobalDef().def("tirx.SBlockRealize", [](ffi::Array iter_values, - PrimExpr predicate, SBlock block, Span span) { - return SBlockRealize(iter_values, predicate, block, span); - }); -} - PrimExpr TypeAnnotation(PrimType dtype, Span span) { static const Op& type_annotation_op = Op::Get("tirx.type_annotation"); return Call(dtype, type_annotation_op, {}, {}, {}, span).as_or_throw(); diff --git a/src/tirx/ir/stmt_functor.cc b/src/tirx/ir/stmt_functor.cc index a0a1fc823c5a..5e19a6348c42 100644 --- a/src/tirx/ir/stmt_functor.cc +++ b/src/tirx/ir/stmt_functor.cc @@ -32,6 +32,7 @@ #include #include #include +#include #include "data_type_rewriter.h" #include "seq_stmt_mutate.h" @@ -39,6 +40,17 @@ namespace tvm { namespace tirx { +namespace { +std::vector& StmtExprVisitorExtensions() { + static std::vector extensions; + return extensions; +} +} // namespace + +void StmtExprVisitor::RegisterExtension(void (*init)(VTable*)) { + StmtExprVisitorExtensions().push_back(init); +} + void StmtExprVisitor::InitVTable(VTable* vtable) { tvm::ExprVisitor::InitVTable(vtable); SetDispatch(vtable); @@ -55,10 +67,10 @@ void StmtExprVisitor::InitVTable(VTable* vtable) { SetDispatch(vtable); SetDispatch(vtable); SetDispatch(vtable); - SetDispatch(vtable); - SetDispatch(vtable); SetDispatch(vtable); SetDispatch(vtable); + SetDispatch(vtable); + for (auto init : StmtExprVisitorExtensions()) init(vtable); } ffi::Optional StmtExprVisitor::Visit_(const VarNode* op) { return std::nullopt; } @@ -228,42 +240,6 @@ ffi::Optional StmtExprVisitor::Visit_(const EvaluateNode* op) { return this->Visit(op->value); } -ffi::Optional StmtExprVisitor::Visit_(const SBlockNode* op) { - for (const IterVar& iter_var : op->iter_vars) { - TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->Visit(iter_var->dom->min)); - TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->Visit(iter_var->dom->extent)); - } - for (const BufferVar& buf : op->alloc_buffers) { - TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN( - this->WithDefRegionKind(kTVMFFIDefRegionKindSimple, [&]() { return this->Visit(buf); })); - TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(VisitBufferMetadata(buf)); - } - for (const TensorRegion& region : op->reads) { - TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->Visit(region)); - } - for (const TensorRegion& region : op->writes) { - TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->Visit(region)); - } - for (const MatchBufferRegion& match_buffer_region : op->match_buffers) { - TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->WithDefRegionKind( - kTVMFFIDefRegionKindSimple, [&]() { return this->Visit(match_buffer_region->buffer); })); - TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(VisitBufferMetadata(match_buffer_region->buffer)); - TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->Visit(match_buffer_region->source)); - } - if (op->init.has_value()) { - TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->Visit(op->init.value())); - } - return this->Visit(op->body); -} - -ffi::Optional StmtExprVisitor::Visit_(const SBlockRealizeNode* op) { - for (const auto& child : op->iter_values) { - TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->Visit(child)); - } - TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->Visit(op->predicate)); - return this->Visit(op->block); -} - ffi::Optional StmtExprVisitor::Visit_(const ScopeIdDefStmtNode* op) { // Flat stmt -- no body. Visit extents (skip deferred defs whose extents // are NullOpt) and any preferred_extents. @@ -308,6 +284,17 @@ ffi::Optional StmtExprVisitor::Visit_(const TilePrimitiveCallNod return std::nullopt; } +namespace { +std::vector& StmtExprMutatorExtensions() { + static std::vector extensions; + return extensions; +} +} // namespace + +void StmtExprMutator::RegisterExtension(void (*init)(VTable*)) { + StmtExprMutatorExtensions().push_back(init); +} + void StmtExprMutator::InitVTable(VTable* vtable) { tvm::ExprMutator::InitVTable(vtable); SetDispatch(vtable); @@ -324,10 +311,10 @@ void StmtExprMutator::InitVTable(VTable* vtable) { SetDispatch(vtable); SetDispatch(vtable); SetDispatch(vtable); - SetDispatch(vtable); - SetDispatch(vtable); SetDispatch(vtable); SetDispatch(vtable); + SetDispatch(vtable); + for (auto init : StmtExprMutatorExtensions()) init(vtable); } UnchangedOr StmtExprMutator::Mutate_(const BindNode* op, InplaceMode inplace_mode) { @@ -467,28 +454,6 @@ UnchangedOr StmtExprMutator::Mutate_(const EvaluateNode* op, InplaceMode i return Stmt(std::move(copy)); } -UnchangedOr StmtExprMutator::Mutate_(const SBlockRealizeNode* op, InplaceMode inplace_mode) { - auto iter_values = - Mutate(op->iter_values, inplace_mode).as_or_throw>>(); - auto predicate = Mutate(op->predicate, inplace_mode); - auto block = Mutate(op->block, inplace_mode).as_or_throw>(); - 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(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(*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)); -} - UnchangedOr StmtExprMutator::Mutate_(const BreakNode* op, InplaceMode inplace_mode) { return ffi::Unchanged(); } @@ -552,73 +517,22 @@ UnchangedOr StmtExprMutator::Mutate_(const BufferStoreNode* op, InplaceMod return Stmt(std::move(copy)); } -UnchangedOr StmtExprMutator::Mutate_(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> replacements; - for (size_t i = 0; i < iters->size(); ++i) { - const auto* iter = (*iters)[i].as(); - InplaceMode domain_mode = iter->unique() ? iter_mode : InplaceMode::kDisallow; - auto domain = Mutate(iter->dom, domain_mode).as_or_throw>(); - if (domain.UnchangedOrSameAs(iter->dom)) continue; - if (domain_mode == InplaceMode::kAllow) { - const_cast(iter)->dom = std::move(domain).ValueUnchecked(); - } else { - auto updated = ffi::make_object(*iter); - updated->dom = std::move(domain).ValueUnchecked(); - replacements.emplace_back(i, IterVar(std::move(updated))); - } - } - UnchangedOr> iter_vars = ffi::Unchanged(); - if (!replacements.empty()) { - if (iter_mode == InplaceMode::kAllow) { - for (auto& [i, iter] : replacements) { - const_cast(iters)->SetItem(i, std::move(iter)); - } - } else { - ffi::Array updated = op->iter_vars; - for (auto& [i, iter] : replacements) updated.Set(i, std::move(iter)); - iter_vars = std::move(updated); - } - } - auto alloc_buffers = WithDefRegionKind(kTVMFFIDefRegionKindSimple, [&] { - return Mutate(op->alloc_buffers, inplace_mode); - }).as_or_throw>>(); - auto reads = Mutate(op->reads, inplace_mode).as_or_throw>>(); - auto writes = - Mutate(op->writes, inplace_mode).as_or_throw>>(); - auto match_buffers = Mutate(op->match_buffers, inplace_mode) - .as_or_throw>>(); - auto init = Mutate(op->init, inplace_mode).as_or_throw>>(); - auto body = 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)) +UnchangedOr StmtExprMutator::Mutate_(const BufferRegionNode* op, InplaceMode inplace_mode) { + auto buffer = Mutate(op->buffer, inplace_mode).as_or_throw>(); + auto region = Mutate(op->region, inplace_mode).as_or_throw>>(); + if (buffer.UnchangedOrSameAs(op->buffer) && region.UnchangedOrSameAs(op->region)) { return ffi::Unchanged(); + } if (inplace_mode == InplaceMode::kAllow) { - auto* writable = const_cast(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(); + auto* writable = const_cast(op); + if (!buffer.IsUnchanged()) writable->buffer = std::move(buffer).ValueUnchecked(); + if (!region.IsUnchanged()) writable->region = std::move(region).ValueUnchecked(); return ffi::Unchanged(); } - auto copy = ffi::make_object(*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)); + auto copy = ffi::make_object(*op); + if (!buffer.IsUnchanged()) copy->buffer = std::move(buffer).ValueUnchecked(); + if (!region.IsUnchanged()) copy->region = std::move(region).ValueUnchecked(); + return Expr(std::move(copy)); } UnchangedOr StmtExprMutator::Mutate_(const SeqStmtNode* op, InplaceMode inplace_mode) { diff --git a/src/tirx/ir/tir_visitor_with_path.cc b/src/tirx/ir/tir_visitor_with_path.cc index 0f0f684a50fa..d7061f93cc52 100644 --- a/src/tirx/ir/tir_visitor_with_path.cc +++ b/src/tirx/ir/tir_visitor_with_path.cc @@ -24,7 +24,6 @@ #include "tir_visitor_with_path.h" #include -#include #include #include @@ -34,6 +33,26 @@ namespace tvm { namespace tirx { +namespace { +std::vector& PathVisitorExtensions() { + static std::vector extensions; + return extensions; +} +} // namespace +void TIRVisitorWithPath::RegisterExtension(void (*init)(VTable*)) { + PathVisitorExtensions().push_back(init); +} +TIRVisitorWithPath::TIRVisitorWithPath() + : StmtVisitor([] { + static const VTable table = [] { + VTable table; + StmtVisitor::InitVTable(&table); + for (auto init : PathVisitorExtensions()) init(&table); + table.Finalize(); + return table; + }(); + return &table; + }()) {} using AccessPath = ffi::reflection::AccessPath; @@ -159,14 +178,6 @@ void TIRVisitorWithPath::Visit(const TensorRegion& region, AccessPath path) { Visit(region->region, path->Attr("region")); } -void TIRVisitorWithPath::Visit(const MatchBufferRegion& match, AccessPath path) { - Visit(match->source, path->Attr("source")); - - // MatchBufferRegion define the match->buffer, but do not own the - // body in which the match->buffer is defined. Therefore, the - // definitions are handled in the BlockNode visitor. -} - void TIRVisitorWithPath::Visit(const IterVar& iter_var, AccessPath path) { if (iter_var->dom.defined()) { Visit(iter_var->dom, path->Attr("dom")); @@ -191,8 +202,7 @@ void TIRVisitorWithPath::Dispatch_(const AttrStmtNode* op, AccessPath path) { std::vector, DefContext, DefContext>> context; if (auto iter_var = op->node.as(); - iter_var && - (op->attr_key == attr::thread_extent || op->attr_key == s_tir::attr::virtual_thread)) { + iter_var && (op->attr_key == attr::thread_extent || op->attr_key == "virtual_thread")) { // Some attributes serve as a source of definition for the // tirx::Var they annotate. context.push_back(WithDef(iter_var.value(), path->Attr("node"))); @@ -268,58 +278,7 @@ void TIRVisitorWithPath::Dispatch_(const EvaluateNode* op, AccessPath path) { Visit(op->value, path->Attr("value")); } -void TIRVisitorWithPath::Dispatch_(const SBlockNode* op, AccessPath path) { - std::vector, DefContext, DefContext>> 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)); - } - } - - Visit(op->reads, path->Attr("reads")); - Visit(op->writes, path->Attr("writes")); - - { - auto match_path = path->Attr("match_buffers"); - Visit(op->match_buffers, match_path); - - for (size_t i = 0; i < op->match_buffers.size(); i++) { - 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)); - } - } - - 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")); -} - -void TIRVisitorWithPath::Dispatch_(const tirx::TilePrimitiveCallNode* op, AccessPath path) { +void TIRVisitorWithPath::VisitStmt_(const tirx::TilePrimitiveCallNode* op, AccessPath path) { for (size_t i = 0; i < op->args.size(); i++) { if (op->args[i] == nullptr) { continue; diff --git a/src/tirx/ir/tir_visitor_with_path.h b/src/tirx/ir/tir_visitor_with_path.h index f4547becd6cc..aea29aac0c9f 100644 --- a/src/tirx/ir/tir_visitor_with_path.h +++ b/src/tirx/ir/tir_visitor_with_path.h @@ -45,12 +45,23 @@ namespace tirx { class TIRVisitorWithPath : protected ExprFunctor, protected StmtFunctor { public: + using StmtVisitor = StmtFunctor; + using VTable = StmtVisitor::VTable; + // Dialect-owned handlers use nested access to the active traversal context. + class Extension; + TIRVisitorWithPath(); + // Extensions register at library initialization, before first visitor construction. + static void RegisterExtension(void (*init)(VTable*)); template void operator()(TObjectRef&& obj) { Visit(std::forward(obj), ffi::reflection::AccessPath::Root()); } protected: + // A dialect restriction can reject extension statements without knowing concrete node types. + virtual bool EnterExtensionStmt(const ffi::Object* op, ffi::reflection::AccessPath path) { + return true; + } // Delegate to ExprFunctor::Dispatch for PrimExpr, and any subclasses virtual inline void Visit(const PrimExpr& obj, ffi::reflection::AccessPath path) { Dispatch(obj, path); @@ -92,8 +103,7 @@ class TIRVisitorWithPath : protected ExprFunctor #include #include +#include +#include #include #include #include #include #include #include -#include #include "../../../tirx/ir/script/script_complete.h" #include "./utils.h" @@ -50,7 +51,7 @@ namespace { // // This normalizer runs at PrimFunc construction time: it strips any defined // layout from buffers in `buffer_map` / `root_alloc_buffers` and rewrites -// matching body references through the StmtExprMutator's built-in +// matching body references through the s_tir::StmtExprMutator's built-in // variable remapping, so the body remains well-formed. class STirBufferLayoutNormalizer : public tvm::tirx::StmtExprMutator { public: @@ -202,7 +203,7 @@ void PrimFuncFrameNode::ExitWithScope() { void SBlockFrameNode::ExitWithScope() { TIRFrameNode::ExitWithScope(); - // Allow SBlock construction in raw IRBuilder context (no enclosing PrimFuncFrame) + // Allow s_tir::SBlock construction in raw IRBuilder context (no enclosing PrimFuncFrame) // so test fixtures can construct blocks/block-realizes directly. ffi::Array tir_alloc_buffers; @@ -213,9 +214,10 @@ void SBlockFrameNode::ExitWithScope() { if (int detect_access = (!reads.has_value()) | (!writes.has_value() << 1)) { attrs.Set("tirx.script_parsing_detect_access", tvm::IntImm::Int64(detect_access)); } - tvm::tirx::SBlock block(iter_vars, reads.value_or(ffi::Array()), - writes.value_or(ffi::Array()), name, AsStmt(stmts), - init, tir_alloc_buffers, match_buffers, attrs, tvm::Span()); + tvm::s_tir::SBlock block(iter_vars, reads.value_or(ffi::Array()), + writes.value_or(ffi::Array()), name, + AsStmt(stmts), init, tir_alloc_buffers, match_buffers, attrs, + tvm::Span()); if (no_realize) { TVM_FFI_CHECK(iter_values.empty(), ValueError) << "Block bindings are not allowed when `no_realize=True`"; @@ -224,7 +226,7 @@ void SBlockFrameNode::ExitWithScope() { AddToParent(block); } else { AddToParent( - tvm::tirx::SBlockRealize(iter_values, predicate.value_or(IntImm::Bool(true)), block)); + tvm::s_tir::SBlockRealize(iter_values, predicate.value_or(IntImm::Bool(true)), block)); } } diff --git a/src/tirx/script/builder/ir.cc b/src/tirx/script/builder/ir.cc index 53bab5646b72..b776e48c6b20 100644 --- a/src/tirx/script/builder/ir.cc +++ b/src/tirx/script/builder/ir.cc @@ -27,6 +27,7 @@ #include #include #include +#include #include #include #include @@ -157,12 +158,12 @@ BufferVar MatchBuffer(ffi::ObjectRef param, ffi::Array shape, PrimType TVM_FFI_THROW(InternalError) << "ValueError: Can not bind non-input param to buffer."; } else if (const auto* buffer_load = param.as()) { SBlockFrame frame = FindSBlockFrame("T.match_buffer"); - frame->match_buffers.push_back(tvm::tirx::MatchBufferRegion( + frame->match_buffers.push_back(tvm::s_tir::MatchBufferRegion( buffer, BufferRegionFromLoad(ffi::GetRef(buffer_load)))); } else if (const auto* buffer_region = param.as()) { SBlockFrame frame = FindSBlockFrame("T.match_buffer"); frame->match_buffers.push_back( - tvm::tirx::MatchBufferRegion(buffer, ffi::GetRef(buffer_region))); + tvm::s_tir::MatchBufferRegion(buffer, ffi::GetRef(buffer_region))); } else { TVM_FFI_THROW(InternalError) << "ValueError: Unexpected type for TIR MatchBuffer."; } @@ -399,7 +400,7 @@ ffi::Variant SBlockAllocBuffer( "Use `T.alloc_buffer()` inside default (tirx) PrimFuncs."; } - // Walk up the frame stack: attach to the innermost enclosing SBlock (lifting + // Walk up the frame stack: attach to the innermost enclosing s_tir::SBlock (lifting // the allocation past any intermediate For/If/While frames). Fall back to the // PrimFunc root when no sblock is in scope. When neither is present (raw // IRBuilder construction used by tests), just return the buffer. diff --git a/src/tirx/script/printer/block.cc b/src/tirx/script/printer/block.cc index 6d1fb5698d3c..d56b9a7921ca 100644 --- a/src/tirx/script/printer/block.cc +++ b/src/tirx/script/printer/block.cc @@ -16,6 +16,8 @@ * specific language governing permissions and limitations * under the License. */ +#include + #include "./utils.h" namespace tvm { @@ -23,12 +25,12 @@ namespace script { namespace printer { -Doc PrintBlock(IRDocsifier d, tirx::SBlock block, AccessPath block_p, // - ffi::Optional opt_realize, +Doc PrintBlock(IRDocsifier d, s_tir::SBlock block, AccessPath block_p, // + ffi::Optional opt_realize, ffi::Optional opt_realize_p) { With frame(d, block); TVM_FFI_ICHECK_EQ(opt_realize.has_value(), opt_realize_p.has_value()); - const tirx::SBlockRealizeNode* realize = + const s_tir::SBlockRealizeNode* realize = opt_realize.has_value() ? opt_realize.value().get() : nullptr; AccessPath realize_p = *opt_realize_p; @@ -189,7 +191,7 @@ Doc PrintBlock(IRDocsifier d, tirx::SBlock block, AccessPath block_p, // } // Step 6. Handle `match_buffer` for (int i = 0, n = block->match_buffers.size(); i < n; ++i) { - tirx::MatchBufferRegion buffer_region = block->match_buffers[i]; + s_tir::MatchBufferRegion buffer_region = block->match_buffers[i]; AccessPath buffer_region_p = block_p->Attr("match_buffers")->ArrayItem(i); StmtDoc doc = d->AsDoc(buffer_region, buffer_region_p); (*frame)->stmts.push_back(doc); @@ -218,8 +220,8 @@ Doc PrintBlock(IRDocsifier d, tirx::SBlock block, AccessPath block_p, // } TVM_FFI_STATIC_INIT_BLOCK() { - IRDocsifier::vtable().set_dispatch( - "", [](tirx::SBlockRealize realize, AccessPath p, IRDocsifier d) -> Doc { + IRDocsifier::vtable().set_dispatch( + "", [](s_tir::SBlockRealize realize, AccessPath p, IRDocsifier d) -> Doc { Doc doc = PrintBlock(d, realize->block, p->Attr("block"), realize, p); // since we do not have d->AsDoc for realize->block, // we should add possible doc decoration manually. @@ -229,14 +231,14 @@ TVM_FFI_STATIC_INIT_BLOCK() { } TVM_FFI_STATIC_INIT_BLOCK() { - IRDocsifier::vtable().set_dispatch( - "", [](tirx::SBlock block, AccessPath p, IRDocsifier d) -> Doc { + IRDocsifier::vtable().set_dispatch( + "", [](s_tir::SBlock block, AccessPath p, IRDocsifier d) -> Doc { return PrintBlock(d, block, p, std::nullopt, std::nullopt); }); } -TVM_REGISTER_SCRIPT_AS_REPR(tirx::SBlockNode, ReprPrintTIR); -TVM_REGISTER_SCRIPT_AS_REPR(tirx::SBlockRealizeNode, ReprPrintTIR); +TVM_REGISTER_SCRIPT_AS_REPR(s_tir::SBlockNode, ReprPrintTIR); +TVM_REGISTER_SCRIPT_AS_REPR(s_tir::SBlockRealizeNode, ReprPrintTIR); TVM_FFI_STATIC_INIT_BLOCK() { IRDocsifier::vtable().set_dispatch( diff --git a/src/tirx/script/printer/buffer.cc b/src/tirx/script/printer/buffer.cc index a69b018f5775..e1b43645bcc2 100644 --- a/src/tirx/script/printer/buffer.cc +++ b/src/tirx/script/printer/buffer.cc @@ -17,6 +17,7 @@ * under the License. */ #include // For `kAllocAlignment` +#include #include #include @@ -575,8 +576,8 @@ TVM_FFI_STATIC_INIT_BLOCK() { } TVM_FFI_STATIC_INIT_BLOCK() { - IRDocsifier::vtable().set_dispatch( - "", [](tirx::MatchBufferRegion stmt, AccessPath p, IRDocsifier d) -> Doc { + IRDocsifier::vtable().set_dispatch( + "", [](s_tir::MatchBufferRegion stmt, AccessPath p, IRDocsifier d) -> Doc { Frame frame = d->frames.back(); ExprDoc lhs = DefineBuffer(stmt->buffer, frame, d); ExprDoc src_buffer = d->AsDoc(stmt->source, p->Attr("source")); @@ -593,7 +594,7 @@ TVM_SCRIPT_REPR(tirx::BufferTypeNode, ReprPrintTIR); TVM_SCRIPT_REPR(tirx::IterNode, ReprPrintTIR); TVM_SCRIPT_REPR(tirx::TileLayoutNode, ReprPrintTIR); TVM_SCRIPT_REPR(tirx::ComposeLayoutNode, ReprPrintTIR); -TVM_SCRIPT_REPR(tirx::MatchBufferRegionNode, ReprPrintTIR); +TVM_SCRIPT_REPR(s_tir::MatchBufferRegionNode, ReprPrintTIR); } // namespace printer } // namespace script diff --git a/src/tirx/script/printer/function.cc b/src/tirx/script/printer/function.cc index dc03b26eb66a..7e37317a7479 100644 --- a/src/tirx/script/printer/function.cc +++ b/src/tirx/script/printer/function.cc @@ -17,6 +17,7 @@ * under the License. */ #include +#include #include @@ -161,21 +162,21 @@ TVM_FFI_STATIC_INIT_BLOCK() { } } // Step 3. Handle `func->body` - ffi::Optional implicit_root_block = [&]() -> ffi::Optional { - const tirx::SBlockRealizeNode* root_block_realize = - func->body.as(); + ffi::Optional implicit_root_block = [&]() -> ffi::Optional { + const s_tir::SBlockRealizeNode* root_block_realize = + func->body.as(); if (root_block_realize && !root_block_realize->iter_values.size() && tvm::prim::is_one(root_block_realize->predicate)) { - tirx::SBlock root_block = root_block_realize->block; + s_tir::SBlock root_block = root_block_realize->block; if (!root_block->annotations.size() && !root_block->match_buffers.size() && !root_block->reads.size() && !root_block->writes.size() && !root_block->init.has_value()) { - const tirx::SBlockRealizeNode* block_realize = - root_block->body.as(); + const s_tir::SBlockRealizeNode* block_realize = + root_block->body.as(); if (root_block->alloc_buffers.size() || (block_realize && block_realize->block->iter_vars.size()) || (!block_realize && - tirx::ContainsNode(root_block->body))) { + tirx::ContainsNode(root_block->body))) { return root_block; } } @@ -183,7 +184,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { return std::nullopt; }(); if (d->cfg->syntax_sugar && implicit_root_block) { - tirx::SBlock root_block = implicit_root_block.value(); + s_tir::SBlock root_block = implicit_root_block.value(); AccessPath root_block_p = p->Attr("body")->Attr("block"); (*f)->stmts.push_back(CommentDoc("with T.sblock(\"root\"):")); // Handle root block `alloc_buffer` diff --git a/src/tirx/script/printer/utils.h b/src/tirx/script/printer/utils.h index 1417febd7b9c..f5b4b4cdd84c 100644 --- a/src/tirx/script/printer/utils.h +++ b/src/tirx/script/printer/utils.h @@ -23,6 +23,7 @@ #include #include #include +#include #include #include #include @@ -31,7 +32,6 @@ #include #include #include -#include #include #include diff --git a/src/tirx/transform/flatten_buffer.cc b/src/tirx/transform/flatten_buffer.cc index 0bf7e3cb436c..f41cc14fe0ba 100644 --- a/src/tirx/transform/flatten_buffer.cc +++ b/src/tirx/transform/flatten_buffer.cc @@ -34,6 +34,7 @@ #include "../ir/ir_mutator_with_analyzer.h" #include "ir_utils.h" +#include "stmt_extension.h" namespace tvm { namespace tirx { @@ -58,7 +59,7 @@ using namespace tvm::prim; * Every use site then only looks the pair up; a use before its definition is * a hard error instead of a silently stale reference. */ -class BufferFlattener : public IRMutatorWithAnalyzer { +class BufferFlattener : public FlattenStmtMutator { public: using IRMutatorWithAnalyzer::Mutate; using IRMutatorWithAnalyzer::Mutate_; @@ -98,7 +99,7 @@ class BufferFlattener : public IRMutatorWithAnalyzer { } public: - explicit BufferFlattener(const arith::Analyzer& ana) : IRMutatorWithAnalyzer(ana) {} + explicit BufferFlattener(const arith::Analyzer& ana) : FlattenStmtMutator(ana) {} private: struct FlatInfo { @@ -178,36 +179,7 @@ class BufferFlattener : public IRMutatorWithAnalyzer { return it->second; } - UnchangedOr Mutate_(const SBlockNode* op, InplaceMode inplace_mode) final { - TVM_FFI_ICHECK_EQ(op->match_buffers.size(), 0) - << "Unexpected MatchBufferRegion found during tirx.transform.FlattenBuffer. " - << "All MatchBufferRegion should be removed in tirx.transform.LowerMatchBuffer."; - - SBlock block = ffi::GetRef(op); - - ffi::Array alloc_buffers = op->alloc_buffers; - alloc_buffers.MutateByApply([this](BufferVar buf) { return Define(buf).flattened; }); - if (!alloc_buffers.same_as(op->alloc_buffers)) { - block.CopyOnWrite()->alloc_buffers = alloc_buffers; - } - - ffi::Array reads = op->reads; - reads.MutateByApply([this](TensorRegion region) { return MutateBufferRegion(region); }); - if (!reads.same_as(op->reads)) { - block.CopyOnWrite()->reads = reads; - } - - ffi::Array writes = op->writes; - writes.MutateByApply([this](TensorRegion region) { return MutateBufferRegion(region); }); - if (!writes.same_as(op->writes)) { - block.CopyOnWrite()->writes = writes; - } - - // The retained or rebuilt block bypasses the generic entry's current-node check. - return StmtExprMutator::Mutate_(block.get(), - block.unique() ? inplace_mode : InplaceMode::kDisallow) - .ValueOrUnchanged(block); - } + BufferVar DefineBuffer(BufferVar buffer) final { return Define(buffer).flattened; } UnchangedOr Mutate_(const AllocBufferNode* op, InplaceMode inplace_mode) final { const FlatInfo& info = Define(op->buffer); @@ -323,9 +295,9 @@ class BufferFlattener : public IRMutatorWithAnalyzer { return BufferLoad(info.flattened, FoldIndices(info, node->indices), node->span); } - TensorRegion MutateBufferRegion(TensorRegion region) { - const FlatInfo& info = Lookup(region->source.as_or_throw()); - if (info.flattened.same_as(region->source.as_or_throw())) { + BufferRegion RewriteRegion(BufferRegion region) final { + const FlatInfo& info = Lookup(region->buffer); + if (info.flattened.same_as(region->buffer)) { return region; } diff --git a/src/tirx/transform/force_narrow_index_to_i32.cc b/src/tirx/transform/force_narrow_index_to_i32.cc index 91e0d35e5a6b..956ba262baf8 100644 --- a/src/tirx/transform/force_narrow_index_to_i32.cc +++ b/src/tirx/transform/force_narrow_index_to_i32.cc @@ -69,19 +69,12 @@ class Int32DTypeNarrower : public IndexDataTypeNormalizer { return ffi::Unchanged(); } - UnchangedOr Mutate_(const SBlockNode* block, InplaceMode inplace_mode) final { - SBlock block_ = IndexDataTypeNormalizer::Mutate_(block, inplace_mode) - .ValueOrUnchanged(ffi::GetRef(block)) - .as_or_throw(); - // Check if the allocated integer buffers have dtype other than int32. - for (const BufferVar& buf : block_->alloc_buffers) { - if (buf->dtype.MatchesCode(DLDataTypeCode::kDLInt) && buf->dtype.bits() > 32) { - TVM_FFI_THROW(InternalError) - << "The buffer " << buf << " allocated in the function has dtype " << buf->dtype - << ". The function is " << func_; - } + void ValidateAllocation(const BufferVar& buf) final { + if (buf->dtype.MatchesCode(DLDataTypeCode::kDLInt) && buf->dtype.bits() > 32) { + TVM_FFI_THROW(InternalError) + << "The buffer " << buf << " allocated in the function has dtype " << buf->dtype + << ". The function is " << func_; } - return block_; } PrimFunc func_; diff --git a/src/tirx/transform/inline_private_functions.cc b/src/tirx/transform/inline_private_functions.cc index e385963b2200..af52564a71f5 100644 --- a/src/tirx/transform/inline_private_functions.cc +++ b/src/tirx/transform/inline_private_functions.cc @@ -123,10 +123,11 @@ bool IsInlinablePrimFunc(const GlobalVar& gvar, const PrimFunc& prim_func, // We do not currently support inlining of schedulable TIR // functions. To support this use case, repeated names in - // `tirx::SBlock` nodes resulting from multiple calls to the same + // schedulable block nodes resulting from multiple calls to the same // inlined function will need to be de-duplicated. - bool has_block_node = prim_func->body.as(); - if (has_block_node) return false; + static ffi::reflection::TypeAttrColumn prevent_inline("tirx.prevent_inline"); + ffi::AnyView value = prevent_inline[prim_func->body->type_index()]; + if (value != nullptr && value.cast()) return false; return true; } @@ -303,6 +304,7 @@ Pass InlinePrivateFunctions() { } TVM_FFI_STATIC_INIT_BLOCK() { + ffi::reflection::EnsureTypeAttrColumn("tirx.prevent_inline"); namespace refl = tvm::ffi::reflection; refl::GlobalDef().def("tirx.transform.InlinePrivateFunctions", InlinePrivateFunctions); } diff --git a/src/tirx/transform/ir_utils.cc b/src/tirx/transform/ir_utils.cc index 121b12ba8af1..81f7d000c06d 100644 --- a/src/tirx/transform/ir_utils.cc +++ b/src/tirx/transform/ir_utils.cc @@ -28,7 +28,6 @@ #include #include #include -#include #include #include #include @@ -38,6 +37,8 @@ #include #include +#include "stmt_extension.h" + namespace tvm { namespace tirx { using namespace tvm::prim; @@ -88,7 +89,7 @@ Stmt MergeNest(const std::vector>& nest, Stmt body) { return body; } -class IRConvertSSA final : public StmtExprMutator { +class IRConvertSSA final : public SSAStmtMutator { public: using StmtExprMutator::Mutate; using StmtExprMutator::Mutate_; @@ -253,43 +254,20 @@ class IRConvertSSA final : public StmtExprMutator { return decl; } - UnchangedOr Mutate_(const SBlockNode* op, InplaceMode inplace_mode) final { - SBlock block = ffi::GetRef(op); - - // The SBlockNode is the point of definition for the IterVar - // instances. These re-defines must be present before visiting - // the body of the SBlockNode. - return scope_.WithNewScope([&]() -> Stmt { - ffi::Array iter_vars = op->iter_vars.Map([&](IterVar iter_var) { - if (defined_.count(iter_var->var.get())) { - Var new_var = MakeNewVar(iter_var->var); - PushVarRemap(iter_var->var, new_var); - iter_var.CopyOnWrite()->var = new_var.as_or_throw(); - } else { - defined_.insert(iter_var->var.get()); - } - return iter_var; - }); - ffi::Array reads = - block->reads.Map([&](const auto& region) { return VisitBufferAccess(region); }); - ffi::Array writes = - block->writes.Map([&](const auto& region) { return VisitBufferAccess(region); }); - - if (!reads.same_as(block->reads) || !writes.same_as(block->writes) || - !iter_vars.same_as(op->iter_vars)) { - auto write_ptr = block.CopyOnWrite(); - write_ptr->reads = reads; - write_ptr->writes = writes; - write_ptr->iter_vars = iter_vars; - } + Stmt WithScope(const std::function& body) final { return scope_.WithNewScope(body); } - return StmtExprMutator::Mutate_(block.get(), - block.unique() ? inplace_mode : InplaceMode::kDisallow) - .ValueOrUnchanged(block) - .as_or_throw(); - }); + Var DefineVar(Var var) final { + if (defined_.count(var.get())) { + Var new_var = MakeNewVar(var); + PushVarRemap(var, new_var); + return new_var; + } + defined_.insert(var.get()); + return var; } + BufferVar RemapBuffer(BufferVar buffer) final { return GetRemappedBuffer(buffer); } + template Node VisitBufferAccess(Node node) { BufferVar new_buf = GetRemappedBuffer(node->buffer); @@ -784,61 +762,16 @@ ffi::Array GetBufferAllocationShape(const BufferVar& buffer) { return alloc_shape; } -ffi::Array ConvertIndices(const MatchBufferRegion& match_buffer, - const ffi::Array& 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 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; -} - +// Attribute strings are the metadata protocol shared by lowered and schedulable statements. std::pair GetAsyncWaitAttributes(const AttrStmtNode* op) { - TVM_FFI_ICHECK(op && op->attr_key == s_tir::attr::async_wait_queue_scope); + TVM_FFI_ICHECK(op && op->attr_key == "async_wait_queue_scope"); auto inner = op->body.as(); - TVM_FFI_ICHECK(inner && inner->attr_key == s_tir::attr::async_wait_inflight_count); + TVM_FFI_ICHECK(inner && inner->attr_key == "async_wait_inflight_count"); return std::make_pair(op->value, inner->value); } /*! \brief Collect storage alignment information from annotations. */ -class StorageAlignCollector : public StmtExprVisitor { +class StorageAlignCollector : public StorageAlignVisitor { public: ffi::Optional Visit(ffi::AnyView value) override { if (value.as()) return std::nullopt; @@ -849,24 +782,13 @@ class StorageAlignCollector : public StmtExprVisitor { friend std::unordered_map CollectStorageAlignAnnotation( const Stmt& body); - /*! \brief For s-stir, the alignment annotations reside in block annotations. */ - ffi::Optional Visit_(const SBlockNode* op) final { - auto it = op->annotations.find(s_tir::attr::buffer_dim_align); - if (it != op->annotations.end()) { - auto storage_align_annotation = (*it).second.as_or_throw(); - for (const auto& storage_align_tuple : storage_align_annotation) { - int buffer_index = storage_align_tuple.get<0>(); - const BufferVar& buffer = - op->writes[buffer_index]->source.as_or_throw(); - storage_align_[buffer.var()].push_back(storage_align_tuple); - } - } - return StmtExprVisitor::Visit_(op); + void RecordAlignment(const Var& buffer, const StorageAlignTuple& annotation) final { + storage_align_[buffer].push_back(annotation); } /*! \brief AllocBuffer: check for buffer_dim_align annotations. */ ffi::Optional Visit_(const AllocBufferNode* op) final { - auto it = op->annotations.find(s_tir::attr::buffer_dim_align); + auto it = op->annotations.find("buffer_dim_align"); if (it != op->annotations.end()) { auto storage_align_annotation = (*it).second.as_or_throw(); for (const auto& storage_align_tuple : storage_align_annotation) { diff --git a/src/tirx/transform/ir_utils.h b/src/tirx/transform/ir_utils.h index 1b24e462ee9a..89895f63c9d4 100644 --- a/src/tirx/transform/ir_utils.h +++ b/src/tirx/transform/ir_utils.h @@ -30,7 +30,6 @@ #include #include #include -#include #include #include #include @@ -232,21 +231,6 @@ Stmt ConvertSSA(Stmt stmt); */ ffi::String GetPtrStorageScope(Var buffer_var); -/*! - * \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 ConvertIndices(const MatchBufferRegion& match_buffer, - const ffi::Array& 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. - */ -Region ConvertRegion(const MatchBufferRegion& match_buffer, const Region& region); - /*! * \brief Get stride aware buffer allocation shape from buffer. * \param buffer The buffer object. diff --git a/src/tirx/transform/narrow_datatype.cc b/src/tirx/transform/narrow_datatype.cc index 6e0d65d7f440..99bc06b62ce9 100644 --- a/src/tirx/transform/narrow_datatype.cc +++ b/src/tirx/transform/narrow_datatype.cc @@ -27,13 +27,13 @@ #include #include #include -#include #include #include #include #include #include "../ir/data_type_rewriter.h" +#include "stmt_extension.h" namespace tvm { namespace tirx { @@ -75,7 +75,7 @@ using arith::ConstIntBound; // then we narrow `var` into `target_bits_`. That is, // `vmap[var] = min(target_bits_, var.dtype.bits())` // Otherwise, `var` is not narrowed, that is, `vmap[var] = var.dtype.bits()` -class DataTypeVisitor final : public StmtExprVisitor { +class DataTypeVisitor final : public IndexDomainVisitor { public: explicit DataTypeVisitor(int target_bits) : bits_(target_bits), target_bits_(target_bits) {} @@ -124,16 +124,13 @@ class DataTypeVisitor final : public StmtExprVisitor { return StmtExprVisitor::Visit_(op); } - ffi::Optional Visit_(const SBlockNode* op) { - for (const IterVar& iter : op->iter_vars) { - analyzer_->Bind(iter->var, Range::FromMinExtent(iter->dom->min, iter->dom->extent)); - vextent_.insert_or_assign(iter->var.as(), iter->dom->extent.ty()); - } - return StmtExprVisitor::Visit_(op); + void BindDomain(const Var& var, const Range& domain) final { + analyzer_->Bind(var, domain); + vextent_.insert_or_assign(var.as(), domain->extent.ty()); } ffi::Optional Visit_(const AttrStmtNode* op) { - 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 == "virtual_thread") { IterVar iv = op->node.as_or_throw(); TVM_FFI_ICHECK_NE(iv->thread_tag.length(), 0U); analyzer_->Bind(iv->var, Range::FromMinExtent(0, op->value)); diff --git a/src/tirx/transform/stmt_extension.h b/src/tirx/transform/stmt_extension.h new file mode 100644 index 000000000000..75beacf81b61 --- /dev/null +++ b/src/tirx/transform/stmt_extension.h @@ -0,0 +1,159 @@ +/* + * 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_extension.h + * \brief Narrow operation contexts for dialect-owned statement transform hooks. + */ +#ifndef TVM_TIRX_TRANSFORM_STMT_EXTENSION_H_ +#define TVM_TIRX_TRANSFORM_STMT_EXTENSION_H_ + +#include +#include + +#include "../ir_mutator_with_analyzer.h" +#include "ir_utils.h" + +namespace tvm { +namespace tirx { + +// Expose only the scope and remap operations needed at extension definition sites. +class SSAStmtMutator : public StmtExprMutator { + public: + using StmtExprMutator::VTable; + static void RegisterExtension(void (*init)(VTable*)) { Extensions().push_back(init); } + SSAStmtMutator() : StmtExprMutator(GlobalVTable()) {} + virtual Stmt WithScope(const std::function& body) = 0; + virtual Var DefineVar(Var var) = 0; + virtual BufferVar RemapBuffer(BufferVar buffer) = 0; + + protected: + static void InitVTable(VTable* table) { + StmtExprMutator::InitVTable(table); + for (auto init : Extensions()) init(table); + } + + private: + static std::vector& Extensions() { + static std::vector extensions; + return extensions; + } + static const VTable* GlobalVTable() { + static const VTable table = [] { + VTable table; + InitVTable(&table); + table.Finalize(); + return table; + }(); + return &table; + } +}; + +// Buffer geometry stays in the flattener; extensions identify definitions and regions. +class FlattenStmtMutator : public IRMutatorWithAnalyzer { + public: + using IRMutatorWithAnalyzer::VTable; + static void RegisterExtension(void (*init)(VTable*)) { Extensions().push_back(init); } + explicit FlattenStmtMutator(const arith::Analyzer& analyzer) + : IRMutatorWithAnalyzer(analyzer.get(), GlobalVTable()) {} + virtual BufferVar DefineBuffer(BufferVar buffer) = 0; + virtual BufferRegion RewriteRegion(BufferRegion region) = 0; + + protected: + static void InitVTable(VTable* table) { + IRMutatorWithAnalyzer::InitVTable(table); + for (auto init : Extensions()) init(table); + } + + private: + static std::vector& Extensions() { + static std::vector extensions; + return extensions; + } + static const VTable* GlobalVTable() { + static const VTable table = [] { + VTable table; + InitVTable(&table); + table.Finalize(); + return table; + }(); + return &table; + } +}; + +class IndexDomainVisitor : public StmtExprVisitor { + public: + using StmtExprVisitor::VTable; + static void RegisterExtension(void (*init)(VTable*)) { Extensions().push_back(init); } + IndexDomainVisitor() : StmtExprVisitor(GlobalVTable()) {} + virtual void BindDomain(const Var& var, const Range& domain) = 0; + + protected: + static void InitVTable(VTable* table) { + StmtExprVisitor::InitVTable(table); + for (auto init : Extensions()) init(table); + } + + private: + static std::vector& Extensions() { + static std::vector extensions; + return extensions; + } + static const VTable* GlobalVTable() { + static const VTable table = [] { + VTable table; + InitVTable(&table); + table.Finalize(); + return table; + }(); + return &table; + } +}; + +class StorageAlignVisitor : public StmtExprVisitor { + public: + using StmtExprVisitor::VTable; + static void RegisterExtension(void (*init)(VTable*)) { Extensions().push_back(init); } + StorageAlignVisitor() : StmtExprVisitor(GlobalVTable()) {} + virtual void RecordAlignment(const Var& buffer, const StorageAlignTuple& annotation) = 0; + + protected: + static void InitVTable(VTable* table) { + StmtExprVisitor::InitVTable(table); + for (auto init : Extensions()) init(table); + } + + private: + static std::vector& Extensions() { + static std::vector extensions; + return extensions; + } + static const VTable* GlobalVTable() { + static const VTable table = [] { + VTable table; + InitVTable(&table); + table.Finalize(); + return table; + }(); + return &table; + } +}; + +} // namespace tirx +} // namespace tvm +#endif // TVM_TIRX_TRANSFORM_STMT_EXTENSION_H_ diff --git a/tests/cpp/ir_functor_test.cc b/tests/cpp/ir_functor_test.cc index b62b5aa134d6..d21225ee3014 100644 --- a/tests/cpp/ir_functor_test.cc +++ b/tests/cpp/ir_functor_test.cc @@ -26,12 +26,13 @@ #include #include #include +#include +#include #include #include #include #include #include -#include #include #include @@ -120,9 +121,9 @@ TEST(IRF, PreOrderStructuralWalk) { Stmt init = IfThenElse(IntImm::Bool(true), Evaluate(IntImm::Int32(0)), Evaluate(IntImm::Int32(0))); Stmt body = Evaluate(IntImm::Int32(1)); - SBlock block(/*iter_vars=*/{}, /*reads=*/{}, - /*writes=*/{}, /*name_hint=*/"block", /*body=*/body, - /*init=*/init); + s_tir::SBlock block(/*iter_vars=*/{}, /*reads=*/{}, + /*writes=*/{}, /*name_hint=*/"block", /*body=*/body, + /*init=*/init); bool init_visited = false; bool stopped_at_if = true; bool body_visited = false; @@ -202,7 +203,7 @@ TEST(IRF, StmtVisitor) { using namespace tvm; using namespace tvm::tirx; PrimVar x("x"); - class MyVisitor : public StmtExprVisitor { + class MyVisitor : public s_tir::StmtExprVisitor { public: int count = 0; // implementation @@ -233,13 +234,13 @@ TEST(IRF, StmtVisitor) { tirx::Var buf_var("b", PointerType(dtype)); BufferVar buffer = decl_buffer({16}); body = SeqStmt({DeclBuffer(buffer, buf_var), std::move(body)}); - TensorRegion buffer_region = BufferRegion(buffer, {Range::FromMinExtent(x + 1, 1)}); - MatchBufferRegion match_buffer_region(decl_buffer({1}), buffer_region); + BufferRegion buffer_region(buffer, {Range::FromMinExtent(x + 1, 1)}); + s_tir::MatchBufferRegion match_buffer_region(decl_buffer({1}), buffer_region); // construct block and block_realize - SBlock block = SBlock({}, {buffer_region}, {buffer_region}, "block", body, body, {}, - {match_buffer_region}); - Stmt block_realize = SBlockRealize({}, IntImm::Bool(true), block); + s_tir::SBlock block = s_tir::SBlock({}, {buffer_region}, {buffer_region}, "block", body, body, + {}, {match_buffer_region}); + Stmt block_realize = s_tir::SBlockRealize({}, IntImm::Bool(true), block); v->count = 0; v->Visit(block_realize); @@ -256,9 +257,9 @@ TEST(IRF, StmtExprMutator) { using namespace tvm::tirx; PrimVar x("x"); - class MyMutator : public tirx::StmtExprMutator { + class MyMutator : public s_tir::StmtExprMutator { public: - using StmtExprMutator::Mutate_; + using s_tir::StmtExprMutator::Mutate_; UnchangedOr Mutate_(const prim::AddNode* op, InplaceMode) final { return op->a; } }; auto fmakealloc = [&]() { @@ -361,15 +362,15 @@ TEST(IRF, StmtExprMutator) { Stmt alloc = fmakealloc(); // body is: DeclBuffer, AllocBuffer, Evaluate Stmt body = SeqStmt({decl, alloc, eval_body}); - TensorRegion buffer_region = BufferRegion(buffer, {Range::FromMinExtent(x + 1, 1)}); - MatchBufferRegion match_buffer_region(decl_buffer({1}), buffer_region); + BufferRegion buffer_region(buffer, {Range::FromMinExtent(x + 1, 1)}); + s_tir::MatchBufferRegion match_buffer_region(decl_buffer({1}), buffer_region); // construct block and block_realize - SBlock block = SBlock({}, {buffer_region}, {buffer_region}, "block", body, body, {}, - {match_buffer_region}); - Stmt block_realize = SBlockRealize({}, IntImm::Bool(true), block); + s_tir::SBlock block = s_tir::SBlock({}, {buffer_region}, {buffer_region}, "block", body, body, + {}, {match_buffer_region}); + Stmt block_realize = s_tir::SBlockRealize({}, IntImm::Bool(true), block); body = v->Mutate(block_realize).ValueOrUnchanged(std::move(block_realize)); // the body should be changed - SBlock new_block = body.as()->block; + s_tir::SBlock new_block = body.as()->block; // body is a SeqStmt; the Evaluate(x+1) -> Evaluate(x) auto* seq = new_block->body.as(); TVM_FFI_ICHECK(seq != nullptr); diff --git a/tests/cpp/s_tir_functor_test.cc b/tests/cpp/s_tir_functor_test.cc new file mode 100644 index 000000000000..9f9e5734e3db --- /dev/null +++ b/tests/cpp/s_tir_functor_test.cc @@ -0,0 +1,316 @@ +/* + * 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 +#include +#include +#include + +#include +#include +#include + +namespace tvm { +namespace s_tir { +namespace { + +using namespace tirx; + +TEST(STIRFunctor, LegacyInheritedDispatchAndContainsNode) { + class Dispatch : public StmtFunctor { + public: + using StmtFunctor::VisitStmt_; + int VisitStmt_(const SBlockNode*, int value) final { return value + 1; } + int VisitStmt_(const SBlockRealizeNode*, int value) final { return value + 2; } + int VisitStmt_(const EvaluateNode*, int value) final { return value + 3; } + } dispatch; + Stmt body = Evaluate(0); + SBlock block({}, {}, {}, "block", body); + Stmt realize = SBlockRealize({}, IntImm::Bool(true), block); + EXPECT_EQ(dispatch(block, 10), 11); + EXPECT_EQ(dispatch(realize, 10), 12); + EXPECT_EQ(dispatch(body, 10), 13); + EXPECT_TRUE(ContainsNode(realize)); + EXPECT_TRUE(ContainsNode(realize)); + EXPECT_FALSE(ContainsNode(realize)); +} + +TEST(STIRFunctor, NativeBlockOverrideReusesInheritedCoreHooks) { + class Visitor : public StmtExprVisitor { + public: + using StmtExprVisitor::Visit_; + ffi::Optional Visit_(const SBlockNode* op) final { + ++blocks; + return StmtExprVisitor::Visit_(op); + } + ffi::Optional Visit_(const EvaluateNode* op) final { + ++evaluates; + return StmtExprVisitor::Visit_(op); + } + int blocks = 0; + int evaluates = 0; + }; + SBlock block({}, {}, {}, "block", Evaluate(0)); + auto visitor = ffi::make_object(); + visitor->Visit(SBlockRealize({}, IntImm::Bool(true), block)); + EXPECT_EQ(visitor->blocks, 1); + EXPECT_EQ(visitor->evaluates, 1); +} + +TEST(STIRFunctor, NativeVisitPreservesBlockOrderAndBinders) { + PrimVar index("index"), extent("extent"), annotation("annotation"); + BufferVar buffer = decl_buffer({16}); + BufferRegion region(buffer, {Range::FromMinExtent(0, 16)}); + IterVar iter(Range::FromMinExtent(0, extent), index, IterVarType::kDataPar); + SBlock block({iter}, {region}, {}, "block", Evaluate(index), std::nullopt, {buffer}, {}, + {{"annotation", annotation}}); + class Visitor : public StmtExprVisitor { + public: + using StmtExprVisitor::Visit_; + ffi::Optional Visit_(const VarNode* var) final { + vars.push_back(var); + if (var->ty.as()) { + buffer_regions.push_back(def_region_kind()); + } + return std::nullopt; + } + std::vector vars; + std::vector buffer_regions; + }; + auto visitor = ffi::make_object(); + visitor->Visit(block); + EXPECT_EQ(std::count(visitor->vars.begin(), visitor->vars.end(), index.get()), 1); + EXPECT_EQ(std::count(visitor->vars.begin(), visitor->vars.end(), annotation.get()), 0); + ASSERT_EQ(visitor->buffer_regions.size(), 2); + EXPECT_EQ(visitor->buffer_regions[0], kTVMFFIDefRegionKindSimple); + EXPECT_EQ(visitor->buffer_regions[1], kTVMFFIDefRegionKindNone); + + // A full structural walk retains its distinct binder/annotation traversal. + int structural_index = 0; + int structural_annotation = 0; + ffi::StructuralWalk( + block, [&](const Var& var) -> ffi::Expected { + structural_index += var.same_as(index); + structural_annotation += var.same_as(annotation); + return ffi::WalkResult::Advance(); + }); + EXPECT_EQ(structural_index, 2); + EXPECT_EQ(structural_annotation, 1); +} + +TEST(STIRFunctor, NativeMutationKeepsAnnotationsAndSharedIteratorBinders) { + PrimVar index("index"), extent("extent"); + PrimExpr expression = extent + 1; + IterVar iter(Range::FromMinExtent(0, expression), index, IterVarType::kDataPar); + SBlock block({iter}, {}, {}, "block", Evaluate(expression), std::nullopt, {}, {}, + {{"annotation", expression}}); + SBlock retained = block; + class Mutator : public StmtExprMutator { + public: + using StmtExprMutator::Mutate_; + UnchangedOr Mutate_(const prim::AddNode* op, InplaceMode) final { return op->a; } + }; + auto mutator = ffi::make_object(); + Stmt result = mutator->Mutate(block, InplaceMode::kAllow).ValueOrUnchanged(block); + const auto* changed = result.as(); + ASSERT_NE(changed, nullptr); + EXPECT_NE(changed, retained.get()); + EXPECT_TRUE(changed->iter_vars[0]->var.same_as(index)); + EXPECT_TRUE(changed->iter_vars[0]->dom->extent.same_as(extent)); + EXPECT_TRUE(changed->body.as()->value.same_as(extent)); + EXPECT_TRUE(changed->annotations.at("annotation").cast().same_as(expression)); + EXPECT_TRUE(retained->iter_vars[0]->dom->extent.same_as(expression)); + EXPECT_TRUE(iter->dom->extent.same_as(expression)); + + // A sole owning block can update in place while its shared iterator is copied. + SBlock unique({iter}, {}, {}, "unique", Evaluate(expression)); + const auto* original = unique.get(); + auto update = mutator->Mutate(unique, InplaceMode::kAllow); + EXPECT_TRUE(update.IsUnchanged()); + EXPECT_EQ(unique.get(), original); + EXPECT_TRUE(unique->iter_vars[0]->dom->extent.same_as(extent)); + EXPECT_TRUE(iter->dom->extent.same_as(expression)); +} + +template +void CheckMutationRemapsBufferDefinitionsAndUses() { + PrimVar extent("extent"); + BufferVar allocated = decl_buffer({extent + 1}, PrimType::Int(32)); + BufferVar matched = decl_buffer({extent + 1}, PrimType::Int(32)); + BufferRegion region(allocated, {Range::FromMinExtent(0, extent + 1)}); + MatchBufferRegion match(matched, region); + Stmt body = SeqStmt({BufferStore(allocated, 0, {0}), BufferStore(matched, 0, {0})}); + SBlock block({}, {region}, {region}, "block", body, std::nullopt, {allocated}, {match}); + class Mutator : public Base { + public: + using Base::Mutate_; + UnchangedOr Mutate_(const prim::AddNode* op, InplaceMode) final { return op->a; } + }; + auto mutator = ffi::make_object(); + Stmt result = mutator->Mutate(block).ValueOrUnchanged(block); + const auto* changed = result.as(); + ASSERT_NE(changed, nullptr); + BufferVar new_allocated = changed->alloc_buffers[0]; + BufferVar new_matched = changed->match_buffers[0]->buffer; + EXPECT_FALSE(new_allocated.same_as(allocated)); + EXPECT_FALSE(new_matched.same_as(matched)); + EXPECT_TRUE(new_allocated->shape[0].same_as(extent)); + EXPECT_TRUE(new_matched->shape[0].same_as(extent)); + EXPECT_TRUE(changed->reads[0]->buffer.same_as(new_allocated)); + EXPECT_TRUE(changed->writes[0]->buffer.same_as(new_allocated)); + EXPECT_TRUE(changed->match_buffers[0]->source->buffer.same_as(new_allocated)); + const auto* statements = changed->body.as(); + ASSERT_NE(statements, nullptr); + EXPECT_TRUE(statements->seq[0].as()->buffer.same_as(new_allocated)); + EXPECT_TRUE(statements->seq[1].as()->buffer.same_as(new_matched)); + EXPECT_TRUE(block->alloc_buffers[0].same_as(allocated)); + EXPECT_TRUE(block->match_buffers[0]->buffer.same_as(matched)); +} + +TEST(STIRFunctor, NativeMutationRemapsBufferDefinitionsAndUses) { + CheckMutationRemapsBufferDefinitionsAndUses(); +} + +TEST(STIRFunctor, GenericTIRXMutationRemapsBufferDefinitionsAndUses) { + CheckMutationRemapsBufferDefinitionsAndUses(); +} + +TEST(STIRFunctor, GenericTIRXVisitorUsesRegisteredNativePolicy) { + PrimVar index("index"), annotation("annotation"); + BufferVar buffer = decl_buffer({16}); + BufferRegion region(buffer, {Range::FromMinExtent(0, 16)}); + IterVar iter(Range::FromMinExtent(0, 16), index, IterVarType::kDataPar); + SBlock block({iter}, {region}, {}, "block", Evaluate(index), std::nullopt, {buffer}, {}, + {{"annotation", annotation}}); + class Visitor : public tirx::StmtExprVisitor { + public: + using tirx::StmtExprVisitor::Visit_; + ffi::Optional Visit_(const VarNode* var) final { + vars.push_back(var); + if (var->ty.as()) { + buffer_regions.push_back(def_region_kind()); + } + return std::nullopt; + } + std::vector vars; + std::vector buffer_regions; + }; + auto visitor = ffi::make_object(); + visitor->Visit(block); + EXPECT_EQ(std::count(visitor->vars.begin(), visitor->vars.end(), index.get()), 1); + EXPECT_EQ(std::count(visitor->vars.begin(), visitor->vars.end(), annotation.get()), 0); + ASSERT_EQ(visitor->buffer_regions.size(), 2); + EXPECT_EQ(visitor->buffer_regions[0], kTVMFFIDefRegionKindSimple); + EXPECT_EQ(visitor->buffer_regions[1], kTVMFFIDefRegionKindNone); +} + +TEST(STIRFunctor, GenericTIRXMutationPreservesBindersAndAnnotations) { + PrimVar index("index"), replacement("replacement"), extent("extent"); + PrimExpr expression = extent + 1; + IterVar iter(Range::FromMinExtent(0, expression), index, IterVarType::kDataPar); + BufferVar buffer = decl_buffer({expression}, PrimType::Int(32)); + SBlock block({iter}, {}, {}, "block", BufferStore(buffer, index, {0}), std::nullopt, {buffer}, {}, + {{"annotation", expression}}); + class Mutator : public tirx::StmtExprMutator { + public: + using tirx::StmtExprMutator::Mutate_; + UnchangedOr Mutate_(const prim::AddNode* op, InplaceMode) final { return op->a; } + }; + auto mutator = ffi::make_object(); + mutator->VarRemapSet(index, replacement); + SBlock retained = block; + Stmt result = mutator->Mutate(block, InplaceMode::kAllow).ValueOrUnchanged(block); + const auto* changed = result.as(); + ASSERT_NE(changed, nullptr); + EXPECT_TRUE(changed->iter_vars[0]->var.same_as(index)); + EXPECT_TRUE(changed->iter_vars[0]->dom->extent.same_as(extent)); + EXPECT_TRUE(changed->annotations.at("annotation").cast().same_as(expression)); + EXPECT_TRUE(changed->body.as()->value.same_as(replacement)); + EXPECT_TRUE(changed->body.as()->buffer.same_as(changed->alloc_buffers[0])); + EXPECT_TRUE(changed->alloc_buffers[0]->shape[0].same_as(extent)); + EXPECT_TRUE(block->iter_vars[0]->var.same_as(index)); + EXPECT_TRUE(block->annotations.at("annotation").cast().same_as(expression)); + + SBlock unique({}, {}, {}, "unique", Evaluate(expression), std::nullopt, {}, {}, + {{"annotation", expression}}); + const auto* original = unique.get(); + auto update = mutator->Mutate(unique, InplaceMode::kAllow); + EXPECT_TRUE(update.IsUnchanged()); + EXPECT_EQ(unique.get(), original); + EXPECT_TRUE(unique->body.as()->value.same_as(extent)); + EXPECT_TRUE(unique->annotations.at("annotation").cast().same_as(expression)); +} + +TEST(STIRFunctor, GenericTIRXPolicyPreservesInterruptAndErrorIdentity) { + PrimVar annotation("annotation"), body("body"); + SBlock block({}, {}, {}, "block", Evaluate(body), std::nullopt, {}, {}, + {{"annotation", annotation}}); + class Visitor : public tirx::StmtExprVisitor { + public: + using tirx::StmtExprVisitor::Visit_; + ffi::Optional Visit_(const VarNode* op) final { + ++count; + return VisitInterrupt(ffi::GetRef(op)); + } + int count = 0; + }; + auto visitor = ffi::make_object(); + auto interrupt = visitor->Visit(block); + ASSERT_TRUE(interrupt.has_value()); + EXPECT_TRUE(interrupt.value()->value.cast().same_as(body)); + EXPECT_EQ(visitor->count, 1); + + class Mutator : public tirx::StmtExprMutator { + public: + using tirx::StmtExprMutator::Mutate_; + ffi::Error error{"ValueError", "block child mutation error", ""}; + UnchangedOr Mutate_(const VarNode*, InplaceMode) final { throw error; } + }; + auto mutator = ffi::make_object(); + auto result = mutator->MutateExpected(block); + ASSERT_TRUE(result.is_err()); + EXPECT_TRUE(result.error().same_as(mutator->error)); + auto context = ffi::VisitErrorContext::TryGetFromError(result.error()); + ASSERT_TRUE(context.has_value()); + EXPECT_TRUE(context.value()->reverse_visit_pattern.back().same_as(block)); +} + +TEST(STIRFunctor, NativeInterruptStopsBeforeBlockBody) { + PrimVar stop("stop"), body("body"); + IterVar iter(Range::FromMinExtent(0, 16), PrimVar("index"), IterVarType::kDataPar); + SBlock block({iter}, {}, {}, "block", Evaluate(body)); + Stmt realize = SBlockRealize({stop}, IntImm::Bool(true), block); + class Visitor : public StmtExprVisitor { + public: + using StmtExprVisitor::Visit_; + ffi::Optional Visit_(const VarNode* op) final { + ++count; + return VisitInterrupt(ffi::GetRef(op)); + } + int count = 0; + }; + auto visitor = ffi::make_object(); + auto interrupt = visitor->Visit(realize); + ASSERT_TRUE(interrupt.has_value()); + EXPECT_TRUE(interrupt.value()->value.cast().same_as(stop)); + EXPECT_EQ(visitor->count, 1); +} + +} // namespace +} // namespace s_tir +} // namespace tvm diff --git a/tests/python/relax/test_analysis_suggest_layout_transforms.py b/tests/python/relax/test_analysis_suggest_layout_transforms.py index cd4e4e3d8c43..21e2f0b2d050 100644 --- a/tests/python/relax/test_analysis_suggest_layout_transforms.py +++ b/tests/python/relax/test_analysis_suggest_layout_transforms.py @@ -19,7 +19,7 @@ import pytest import tvm.testing -from tvm import relax, tirx +from tvm import relax, s_tir, tirx from tvm.script import tirx as T @@ -28,7 +28,7 @@ def apply_transformations(func, suggested_transfoms, print_transformation=False) for block, per_block_transformations in suggested_transfoms.items(): blockrv = sch.get_sblock(block.name_hint) for obj, index_map in per_block_transformations.items(): - if isinstance(obj, tirx.SBlock): + if isinstance(obj, s_tir.SBlock): block_name = obj.name_hint if print_transformation: print("Block transformation: ", block_name, " :: ", index_map) diff --git a/tests/python/s_tir/analysis/test_sblock_access_region.py b/tests/python/s_tir/analysis/test_sblock_access_region.py index 268e621fede6..378ee7f3142a 100644 --- a/tests/python/s_tir/analysis/test_sblock_access_region.py +++ b/tests/python/s_tir/analysis/test_sblock_access_region.py @@ -504,7 +504,7 @@ def test_conditional_inequality_access_regions(case): ) for var, (minimum, extent) in reversed(list(zip(variables, domains))): body = tirx.For(var, minimum, extent, tirx.ForKind.SERIAL, body) - block = tirx.SBlock([], [], [], "conditional", body) + block = s_tir.SBlock([], [], [], "conditional", body) # Unbounded access sets conservatively cover the whole buffer. outside_expected = [(0, 256)] if case == "unbounded" else domains reads, writes, opaque = s_tir.analysis.get_sblock_access_region( diff --git a/tests/python/s_tir/base/test_sblock_dependence_info.py b/tests/python/s_tir/base/test_sblock_dependence_info.py index 3c750f9da754..eb5b29f20448 100644 --- a/tests/python/s_tir/base/test_sblock_dependence_info.py +++ b/tests/python/s_tir/base/test_sblock_dependence_info.py @@ -90,10 +90,10 @@ def get_sblocks(func: PrimFunc): blocks = {} def update_blocks(node): - if isinstance(node, tvm.tirx.SBlock): + if isinstance(node, tvm.s_tir.SBlock): blocks[node.name_hint] = node - # post_order_visit(func.body, lambda node: blocks[node.name_hint] = node if isinstance(node, tvm.tirx.SBlock) else None) + # post_order_visit(func.body, lambda node: blocks[node.name_hint] = node if isinstance(node, tvm.s_tir.SBlock) else None) structural_walk(func.body, update_blocks, order="post") return blocks diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_schedule_rule_apply_custom_rule.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_schedule_rule_apply_custom_rule.py index 155254491c8b..35146f391909 100644 --- a/tests/python/s_tir/meta_schedule/test_meta_schedule_schedule_rule_apply_custom_rule.py +++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_schedule_rule_apply_custom_rule.py @@ -43,7 +43,7 @@ def main(a: T.handle, b: T.handle, c: T.handle) -> None: @tvm.register_global_func("s_tir.meta_schedule.cpu.test_apply_custom_rule") -def sch_fn(sch: tvm.s_tir.Schedule, block: tvm.tirx.SBlock) -> list[tvm.s_tir.Schedule]: +def sch_fn(sch: tvm.s_tir.Schedule, block: tvm.s_tir.SBlock) -> list[tvm.s_tir.Schedule]: raise ValueError("Intended for s_tir.meta_schedule.cpu.test_apply_custom_rule") diff --git a/tests/python/s_tir/schedule/test_tir_schedule_block_scope.py b/tests/python/s_tir/schedule/test_tir_schedule_block_scope.py index 280b3987b0f3..6c27a7632938 100644 --- a/tests/python/s_tir/schedule/test_tir_schedule_block_scope.py +++ b/tests/python/s_tir/schedule/test_tir_schedule_block_scope.py @@ -85,12 +85,12 @@ def _get_sblock(s: s_tir.ScheduleState, name_hint: str) -> s_tir.StmtSRef: def f_visit(node): nonlocal result - if isinstance(node, tvm.tirx.SBlock) and node.name_hint == name_hint: + if isinstance(node, tvm.s_tir.SBlock) and node.name_hint == name_hint: result = node func = s.mod["main"] structural_walk(func.body, f_visit, order="post") - assert result is not None and isinstance(result, tvm.tirx.SBlock) + assert result is not None and isinstance(result, tvm.s_tir.SBlock) return s.get_sref(result) diff --git a/tests/python/s_tir/schedule/test_tir_schedule_state.py b/tests/python/s_tir/schedule/test_tir_schedule_state.py index 43a4a84f5ec0..23a55cfb81bd 100644 --- a/tests/python/s_tir/schedule/test_tir_schedule_state.py +++ b/tests/python/s_tir/schedule/test_tir_schedule_state.py @@ -95,7 +95,7 @@ def block_in_opaque_block(a: T.handle, b: T.handle) -> None: def replace_ir_builder(deep_copy=False, realize=False): new_func = tvm.script.from_source(elementwise.script()) s = tvm.s_tir.ScheduleState(new_func, debug_mask="all") - target = tvm.tirx.SBlock( + target = tvm.s_tir.SBlock( iter_vars=[], reads=[], writes=[], @@ -107,7 +107,7 @@ def replace_ir_builder(deep_copy=False, realize=False): annotations=None, ) if realize: - target = tvm.tirx.SBlockRealize( + target = tvm.s_tir.SBlockRealize( iter_values=[], predicate=True, block=target, @@ -123,7 +123,7 @@ def replace_ir_builder_module(deep_copy=False, realize=False): other_func = tvm.script.from_source(elementwise.script()) mod = IRModule(functions={"main": new_func, "other": other_func}) s = tvm.s_tir.ScheduleState(mod, debug_mask="all") - target = tvm.tirx.SBlock( + target = tvm.s_tir.SBlock( iter_vars=[], reads=[], writes=[], @@ -135,7 +135,7 @@ def replace_ir_builder_module(deep_copy=False, realize=False): annotations=None, ) if realize: - target = tvm.tirx.SBlockRealize( + target = tvm.s_tir.SBlockRealize( iter_values=[], predicate=True, block=target, diff --git a/tests/python/s_tir/schedule/test_tir_schedule_state_cached_flags.py b/tests/python/s_tir/schedule/test_tir_schedule_state_cached_flags.py index f4369b3ecbdc..2fd157b0c03b 100644 --- a/tests/python/s_tir/schedule/test_tir_schedule_state_cached_flags.py +++ b/tests/python/s_tir/schedule/test_tir_schedule_state_cached_flags.py @@ -468,12 +468,12 @@ def _get_sblock(s: s_tir.ScheduleState, name_hint: str) -> s_tir.StmtSRef: def f_visit(node): nonlocal result - if isinstance(node, tvm.tirx.SBlock) and node.name_hint == name_hint: + if isinstance(node, tvm.s_tir.SBlock) and node.name_hint == name_hint: result = node func = s.mod["main"] structural_walk(func.body, f_visit, order="post") - assert result is not None and isinstance(result, tvm.tirx.SBlock) + assert result is not None and isinstance(result, tvm.s_tir.SBlock) return s.get_sref(result) diff --git a/tests/python/s_tir/test_s_tir_renew_defs.py b/tests/python/s_tir/test_s_tir_renew_defs.py index 23495d7f42d3..e4ee8e2ba0da 100644 --- a/tests/python/s_tir/test_s_tir_renew_defs.py +++ b/tests/python/s_tir/test_s_tir_renew_defs.py @@ -17,10 +17,10 @@ import tvm import tvm.testing +from tvm.s_tir import SBlock from tvm.script import tirx as T from tvm.tirx.buffer import Buffer from tvm.tirx.function import PrimFunc -from tvm.tirx.stmt import SBlock def _check_func_signature_remap(lhs: PrimFunc, rhs: PrimFunc): diff --git a/tests/python/s_tir/test_stmt.py b/tests/python/s_tir/test_stmt.py new file mode 100644 index 000000000000..76b13b4501b2 --- /dev/null +++ b/tests/python/s_tir/test_stmt.py @@ -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. +"""S-TIR node ownership and serialization compatibility.""" + +import json + +import pytest + +import tvm +import tvm.testing +from tvm import s_tir, tirx + + +@pytest.mark.parametrize("legacy", [False, True]) +def test_sblock_serialization(legacy): + source = tirx.decl_buffer((4,), "float32", name="source") + target = tirx.decl_buffer((4,), "float32", name="target") + region = tirx.BufferRegion(source, [tvm.ir.Range(0, 4)]) + match = s_tir.MatchBufferRegion(target, region) + block = s_tir.SBlock([], [region], [region], "copy", tirx.Evaluate(0), match_buffers=[match]) + realize = s_tir.SBlockRealize([], True, block) + graph = json.loads(tvm.ir.save_json([block, realize, match])) + type_keys = {node.get("type") for node in graph["nodes"]} + for name in ("SBlock", "SBlockRealize", "MatchBufferRegion"): + assert f"s_tir.{name}" in type_keys + assert f"tirx.{name}" not in type_keys + assert getattr(s_tir, name).__module__ == "tvm.s_tir.stmt" + assert not hasattr(tirx, name) + assert not hasattr(tirx.stmt, name) + if legacy: + for node in graph["nodes"]: + if node.get("type") in { + "s_tir.SBlock", + "s_tir.SBlockRealize", + "s_tir.MatchBufferRegion", + }: + node["type"] = node["type"].replace("s_tir.", "tirx.") + restored_block, restored_realize, restored_match = tvm.ir.load_json(json.dumps(graph)) + assert isinstance(restored_block, s_tir.SBlock) + assert isinstance(restored_realize, s_tir.SBlockRealize) + assert isinstance(restored_match, s_tir.MatchBufferRegion) + assert isinstance(restored_block, tirx.Stmt) + assert isinstance(restored_realize, tirx.Stmt) + assert restored_realize.block.same_as(restored_block) + assert restored_block.match_buffers[0].same_as(restored_match) + assert restored_block.reads[0].same_as(restored_block.writes[0]) + assert restored_match.source.same_as(restored_block.reads[0]) + tvm.ir.assert_structural_equal(restored_realize, realize, map_free_vars=True) + + +if __name__ == "__main__": + tvm.testing.main() diff --git a/tests/python/s_tir/transform/test_s_tir_transform_memhammer_lower_auto_copy.py b/tests/python/s_tir/transform/test_s_tir_transform_memhammer_lower_auto_copy.py index 5d050149f8ef..0a5935c86a8c 100644 --- a/tests/python/s_tir/transform/test_s_tir_transform_memhammer_lower_auto_copy.py +++ b/tests/python/s_tir/transform/test_s_tir_transform_memhammer_lower_auto_copy.py @@ -1139,7 +1139,7 @@ def verify_single_allocation(stmt, alloc_size=None): def verify(n): if ( - isinstance(n, tvm.tirx.SBlock) + isinstance(n, tvm.s_tir.SBlock) and n.alloc_buffers is not None and (True in ((buf.scope() == "shared.dyn") for buf in n.alloc_buffers)) ): diff --git a/tests/python/s_tir/transform/test_s_tir_transform_mma_buffer_layout.py b/tests/python/s_tir/transform/test_s_tir_transform_mma_buffer_layout.py index 7d2db417f536..d37eafa7d196 100644 --- a/tests/python/s_tir/transform/test_s_tir_transform_mma_buffer_layout.py +++ b/tests/python/s_tir/transform/test_s_tir_transform_mma_buffer_layout.py @@ -32,8 +32,8 @@ def test_explicit_matrix_ab_access_is_rejected(scope, shape, access_kind): body = tirx.Evaluate(tirx.BufferLoad(buffer, [0, 0])) else: body = tirx.BufferStore(buffer, 0.0, [0, 0]) - block = tirx.SBlock([], [], [], "root", body, alloc_buffers=[buffer]) - func = tirx.PrimFunc([], tirx.SBlockRealize([], True, block)) + block = s_tir.SBlock([], [], [], "root", body, alloc_buffers=[buffer]) + func = tirx.PrimFunc([], s_tir.SBlockRealize([], True, block)) with pytest.raises(tvm.error.InternalError, match=f"{scope}.*explicit"): s_tir.transform.TransformMmaBufferLayout()(tvm.IRModule.from_expr(func)) diff --git a/tests/python/tirx-analysis/test_tir_analysis_verify_well_formed.py b/tests/python/tirx-analysis/test_tir_analysis_verify_well_formed.py index 5f1102ea6437..39d0e732cc12 100644 --- a/tests/python/tirx-analysis/test_tir_analysis_verify_well_formed.py +++ b/tests/python/tirx-analysis/test_tir_analysis_verify_well_formed.py @@ -54,7 +54,7 @@ def test_buffer_region_bounds_are_visited(): buffer = tvm.tirx.decl_buffer([4], "int32", data=data) undefined = tvm.tirx.Var("undefined", "int32") region = tvm.tirx.BufferRegion(buffer, [tvm.ir.Range.from_min_extent(undefined, 4)]) - block = tvm.tirx.SBlock([], [region], [], "region", tvm.tirx.Evaluate(0)) + block = tvm.s_tir.SBlock([], [region], [], "region", tvm.tirx.Evaluate(0)) func = tvm.tirx.PrimFunc([buffer], block) assert not tvm.tirx.analysis.verify_well_formed(func, assert_mode=False) @@ -465,14 +465,14 @@ def test_error_undeclared_buffer_in_schedulable_tir(): # Build a block that writes to B without any declaration of B. bi = tvm.tirx.Var("bi", "int32") - block = tvm.tirx.SBlock( + block = tvm.s_tir.SBlock( iter_vars=[tvm.tirx.IterVar(tvm.ir.Range(0, n), bi, 0)], # 0 = kDataPar reads=[tvm.tirx.BufferRegion(A, [tvm.ir.Range(bi, bi + 1)])], writes=[tvm.tirx.BufferRegion(B, [tvm.ir.Range(bi, bi + 1)])], body=tvm.tirx.BufferStore(B, tvm.tirx.BufferLoad(A, [bi]), [bi]), name_hint="write_B", ) - block_realize = tvm.tirx.SBlockRealize( + block_realize = tvm.s_tir.SBlockRealize( iter_values=[i], predicate=tvm.tirx.const(True), block=block, diff --git a/tests/python/tvmscript/test_tvmscript_complete.py b/tests/python/tvmscript/test_tvmscript_complete.py index b4b9736c9bb5..d6098ca52b7e 100644 --- a/tests/python/tvmscript/test_tvmscript_complete.py +++ b/tests/python/tvmscript/test_tvmscript_complete.py @@ -113,7 +113,7 @@ def test_complete_matmul(): A, B, C = [x for x in func.params if tvm.tirx.is_buffer_var(x)] block = func.body.block.body.body.body.body.block - assert isinstance(block, tvm.tirx.SBlock) + assert isinstance(block, tvm.s_tir.SBlock) vi, vj, vk = [x.var for x in block.iter_vars] access_A = tvm.tirx.BufferRegion( A, [Range.from_min_extent(vi, 1), Range.from_min_extent(vk, 1)] @@ -133,7 +133,7 @@ def test_complete_matmul_original(): A, B, C = [x for x in func.params if tvm.tirx.is_buffer_var(x)] block1 = func.body.block.body.body.body[0].block - assert isinstance(block1, tvm.tirx.SBlock) + assert isinstance(block1, tvm.s_tir.SBlock) vi, vj = [x.var for x in block1.iter_vars] access_C = tvm.tirx.BufferRegion( C, [Range.from_min_extent(vi * 4, 4), Range.from_min_extent(vj * 4, 4)] @@ -142,7 +142,7 @@ def test_complete_matmul_original(): tvm.ir.assert_structural_equal(block1.writes, [access_C]) block2 = func.body.block.body.body.body[1].body.block - assert isinstance(block2, tvm.tirx.SBlock) + assert isinstance(block2, tvm.s_tir.SBlock) vi, vj, vk = [x.var for x in block2.iter_vars] access_A = tvm.tirx.BufferRegion( A, [Range.from_min_extent(vi * 4, 4), Range.from_min_extent(vk * 4, 4)] @@ -165,7 +165,7 @@ def _check_elementwise(func): assert len(root_block.writes) == 0 block1 = func.body.block.body[0].body.body.block - assert isinstance(block1, tvm.tirx.SBlock) + assert isinstance(block1, tvm.s_tir.SBlock) vi, vj = [x.var for x in block1.iter_vars] tvm.ir.assert_structural_equal( @@ -178,7 +178,7 @@ def _check_elementwise(func): ) block2 = func.body.block.body[1].body.body.block - assert isinstance(block2, tvm.tirx.SBlock) + assert isinstance(block2, tvm.s_tir.SBlock) vi, vj = [x.var for x in block2.iter_vars] tvm.ir.assert_structural_equal( block2.reads, diff --git a/tests/python/tvmscript/test_tvmscript_error_report.py b/tests/python/tvmscript/test_tvmscript_error_report.py index 87133d2174f0..a61350c1aff6 100644 --- a/tests/python/tvmscript/test_tvmscript_error_report.py +++ b/tests/python/tvmscript/test_tvmscript_error_report.py @@ -439,7 +439,7 @@ def test_reorder_fail_block(): with pytest.raises(tvm.s_tir.ScheduleError) as execinfo: sch.reorder(l, i) expected_sub_error_message = ( - " # tirx.SBlock#0\n" + " # s_tir.SBlock#0\n" ' with T.sblock("B"):\n' " ^^^^^^^^^^^^^^^^^^^\n" ) @@ -482,7 +482,7 @@ def test_report_error_root_block(): with pytest.raises(tvm.s_tir.ScheduleError) as execinfo: sch.compute_inline(root) expected_sub_error_message = ( - ' # tirx.SBlock#0\n with T.sblock("root"):\n ^^^^^^^^^^^^^^^^^^^^^^\n' + ' # s_tir.SBlock#0\n with T.sblock("root"):\n ^^^^^^^^^^^^^^^^^^^^^^\n' ) assert expected_sub_error_message in str(execinfo.value) diff --git a/tests/python/tvmscript/test_tvmscript_ir_builder_tir.py b/tests/python/tvmscript/test_tvmscript_ir_builder_tir.py index 1927e07e3477..1a44ae2bc382 100644 --- a/tests/python/tvmscript/test_tvmscript_ir_builder_tir.py +++ b/tests/python/tvmscript/test_tvmscript_ir_builder_tir.py @@ -24,7 +24,7 @@ import tvm import tvm.runtime import tvm.testing -from tvm import tirx +from tvm import s_tir, tirx from tvm.ir.base import SourceName, Span, assert_structural_equal from tvm.script.ir_builder import IRBuilder from tvm.script.ir_builder import tirx as T @@ -108,7 +108,7 @@ def test_ir_builder_tir_block_base(): block_realize_actual = ib.get() # the expected block - block_expected = tirx.SBlock( + block_expected = s_tir.SBlock( iter_vars=[], reads=[], writes=[], @@ -118,7 +118,7 @@ def test_ir_builder_tir_block_base(): match_buffers=None, annotations={"tirx.script_parsing_detect_access": tirx.IntImm("int64", 3)}, ) - block_realize_expected = tirx.SBlockRealize( + block_realize_expected = s_tir.SBlockRealize( iter_values=[], predicate=True, block=block_expected, @@ -156,7 +156,7 @@ def test_ir_builder_tir_block_complete(): var_d = tirx.Var("d", "int32") buffer_e = tirx.decl_buffer((128, 128), "float32", name="c") var_f = tirx.Var("f", "int32") - block_expected = tirx.SBlock( + block_expected = s_tir.SBlock( iter_vars=[tirx.IterVar((0, 128), tirx.Var("", "int32"), iter_type=tirx.IterVar.DataPar)], reads=[buffer_b[0:16, 0:16]], writes=[buffer_c[var_d:128, var_d:128]], @@ -164,11 +164,11 @@ def test_ir_builder_tir_block_complete(): body=tirx.Evaluate(0), alloc_buffers=[tirx.decl_buffer((128, 128), "float32")], match_buffers=[ - tirx.MatchBufferRegion(tirx.decl_buffer((32, 32), "float32"), buffer_e[0:32, 0:32]) + s_tir.MatchBufferRegion(tirx.decl_buffer((32, 32), "float32"), buffer_e[0:32, 0:32]) ], annotations={"key": "value"}, ) - block_realize_expected = tirx.SBlockRealize( + block_realize_expected = s_tir.SBlockRealize( iter_values=[var_f], predicate=var_a > 1, block=block_expected, @@ -199,7 +199,7 @@ def test_ir_builder_tir_axis(): var_b = tirx.Var("b", "int32") var_c = tirx.Var("c", "int32") var_d = tirx.Var("d", "int32") - block_expected = tirx.SBlock( + block_expected = s_tir.SBlock( iter_vars=[ tirx.IterVar((0, 8), tirx.Var("", "int32"), iter_type=tirx.IterVar.DataPar), tirx.IterVar((0, 16), tirx.Var("", "int32"), iter_type=tirx.IterVar.CommReduce), @@ -212,7 +212,7 @@ def test_ir_builder_tir_axis(): body=tirx.Evaluate(0), annotations={"tirx.script_parsing_detect_access": tirx.IntImm("int64", 3)}, ) - block_realize_expected = tirx.SBlockRealize( + block_realize_expected = s_tir.SBlockRealize( iter_values=[var_a, var_b, var_c, var_d], predicate=True, block=block_expected, diff --git a/tests/python/tvmscript/test_tvmscript_printer_tir.py b/tests/python/tvmscript/test_tvmscript_printer_tir.py index 8be6466b2be8..a876bc86a60d 100644 --- a/tests/python/tvmscript/test_tvmscript_printer_tir.py +++ b/tests/python/tvmscript/test_tvmscript_printer_tir.py @@ -22,7 +22,7 @@ import pytest import tvm.testing -from tvm import ir, tirx +from tvm import ir, s_tir, tirx from tvm.ir import Range from tvm.script.ir_builder import IRBuilder from tvm.script.ir_builder import tirx as T @@ -208,7 +208,7 @@ def test_block(): def test_match_buffer_region(): src = tirx.decl_buffer((128, 128), "float32", name="src") tgt = tirx.decl_buffer((64, 64), "float32", name="tgt") - obj = tirx.MatchBufferRegion( + obj = s_tir.MatchBufferRegion( tgt, tirx.BufferRegion( src, diff --git a/tests/python/tvmscript/test_tvmscript_roundtrip.py b/tests/python/tvmscript/test_tvmscript_roundtrip.py index e41755a9a24b..3466246b69ba 100644 --- a/tests/python/tvmscript/test_tvmscript_roundtrip.py +++ b/tests/python/tvmscript/test_tvmscript_roundtrip.py @@ -23,7 +23,7 @@ import tvm import tvm.testing -from tvm import tirx +from tvm import s_tir, tirx from tvm.script import ir as I from tvm.script import relax as R from tvm.script import tirx as T @@ -1776,13 +1776,13 @@ def test_matmul_original(): rt_func = tvm.script.from_source(func.script()) tvm.ir.assert_structural_equal(func, rt_func) - assert isinstance(rt_func.body.block, tirx.stmt.SBlock) + assert isinstance(rt_func.body.block, s_tir.SBlock) assert isinstance(rt_func.body.block.body, tirx.stmt.For) assert isinstance(rt_func.body.block.body.body, tirx.stmt.For) assert isinstance(rt_func.body.block.body.body.body, tirx.stmt.SeqStmt) - assert isinstance(rt_func.body.block.body.body.body[0].block, tirx.stmt.SBlock) + assert isinstance(rt_func.body.block.body.body.body[0].block, s_tir.SBlock) assert isinstance(rt_func.body.block.body.body.body[1], tirx.stmt.For) - assert isinstance(rt_func.body.block.body.body.body[1].body.block, tirx.stmt.SBlock) + assert isinstance(rt_func.body.block.body.body.body[1].body.block, s_tir.SBlock) def test_element_wise(): @@ -1790,15 +1790,15 @@ def test_element_wise(): rt_func = tvm.script.from_source(func.script()) tvm.ir.assert_structural_equal(func, rt_func) - assert isinstance(rt_func.body.block, tirx.stmt.SBlock) + assert isinstance(rt_func.body.block, s_tir.SBlock) assert isinstance(rt_func.body.block.body, tirx.stmt.SeqStmt) assert isinstance(rt_func.body.block.body[0], tirx.stmt.For) assert isinstance(rt_func.body.block.body[0].body, tirx.stmt.For) - assert isinstance(rt_func.body.block.body[0].body.body.block, tirx.stmt.SBlock) + assert isinstance(rt_func.body.block.body[0].body.body.block, s_tir.SBlock) assert isinstance(rt_func.body.block.body[1], tirx.stmt.For) assert isinstance(rt_func.body.block.body[1].body, tirx.stmt.For) - assert isinstance(rt_func.body.block.body[1].body.body.block, tirx.stmt.SBlock) + assert isinstance(rt_func.body.block.body[1].body.body.block, s_tir.SBlock) def test_predicate(): @@ -1806,11 +1806,11 @@ def test_predicate(): rt_func = tvm.script.from_source(func.script()) tvm.ir.assert_structural_equal(func, rt_func) - assert isinstance(rt_func.body.block, tirx.stmt.SBlock) + assert isinstance(rt_func.body.block, s_tir.SBlock) assert isinstance(rt_func.body.block.body, tirx.stmt.For) assert isinstance(rt_func.body.block.body.body, tirx.stmt.For) assert isinstance(rt_func.body.block.body.body.body, tirx.stmt.For) - assert isinstance(rt_func.body.block.body.body.body.body.block, tirx.stmt.SBlock) + assert isinstance(rt_func.body.block.body.body.body.body.block, s_tir.SBlock) def for_thread_binding(): @@ -1867,19 +1867,19 @@ def test_match_buffer_region(): rt_func = tvm.script.from_source(func.script()) tvm.ir.assert_structural_equal(func, rt_func) - assert isinstance(rt_func.body, tirx.stmt.SBlockRealize) + assert isinstance(rt_func.body, s_tir.SBlockRealize) root = rt_func.body.block assert isinstance(root.body, tirx.stmt.For) assert isinstance(root.body.body, tirx.stmt.For) - assert isinstance(root.body.body.body, tirx.stmt.SBlockRealize) + assert isinstance(root.body.body.body, s_tir.SBlockRealize) outer_block = root.body.body.body.block assert len(outer_block.match_buffers) == 1 buffer_C = outer_block.match_buffers[0].buffer tvm.ir.assert_structural_equal(buffer_C.shape, [T.int32(16), T.int32(1), T.int32(4)]) assert isinstance(outer_block.body, tirx.stmt.For) - assert isinstance(outer_block.body.body, tirx.stmt.SBlockRealize) + assert isinstance(outer_block.body.body, s_tir.SBlockRealize) inner_block = outer_block.body.body.block assert len(inner_block.match_buffers) == 1 buffer_D = inner_block.match_buffers[0].buffer @@ -1912,9 +1912,9 @@ def test_block_elements(): rt_func = tvm.script.from_source(func.script()) tvm.ir.assert_structural_equal(func, rt_func) - assert isinstance(rt_func.body.block, tirx.stmt.SBlock) - assert isinstance(rt_func.body.block.body, tirx.stmt.SBlockRealize) - assert isinstance(rt_func.body.block.body.block, tirx.stmt.SBlock) + assert isinstance(rt_func.body.block, s_tir.SBlock) + assert isinstance(rt_func.body.block.body, s_tir.SBlockRealize) + assert isinstance(rt_func.body.block.body.block, s_tir.SBlock) block = rt_func.body.block.body.block assert isinstance(block.body, tirx.stmt.BufferStore) assert isinstance(block.init, tirx.stmt.BufferStore) @@ -1949,14 +1949,14 @@ def test_opaque_block(): tvm.ir.assert_structural_equal(func, rt_func) root_block = rt_func.body.block - assert isinstance(root_block, tirx.stmt.SBlock) + assert isinstance(root_block, s_tir.SBlock) assert isinstance(root_block.body, tirx.stmt.For) assert isinstance(root_block.body.body[0], tirx.stmt.For) - assert isinstance(root_block.body.body[0].body, tirx.stmt.SBlockRealize) - assert isinstance(root_block.body.body[0].body.block, tirx.stmt.SBlock) + assert isinstance(root_block.body.body[0].body, s_tir.SBlockRealize) + assert isinstance(root_block.body.body[0].body.block, s_tir.SBlock) assert len(root_block.body.body[0].body.block.iter_vars) == 0 - assert isinstance(root_block.body.body[1], tirx.stmt.SBlockRealize) - assert isinstance(root_block.body.body[1].block, tirx.stmt.SBlock) + assert isinstance(root_block.body.body[1], s_tir.SBlockRealize) + assert isinstance(root_block.body.body[1].block, s_tir.SBlock) assert len(root_block.body.body[1].block.iter_vars) == 0 From e799184e83e4677cacfc7059833ced0c944c0356 Mon Sep 17 00:00:00 2001 From: Tianqi Chen Date: Thu, 17 Sep 2026 02:02:22 +0000 Subject: [PATCH 03/18] [S-TIR] Name TensorIntrin modules for their owned API Use tensor_intrin names for the native declaration and registry implementation, with a dedicated Python class module that preserves the existing backend package and public TensorIntrin export. --- include/tvm/s_tir/{function.h => tensor_intrin.h} | 8 ++++---- python/tvm/s_tir/__init__.py | 2 +- python/tvm/s_tir/{function.py => _tensor_intrin.py} | 0 .../schedule_rule/multi_level_tiling_tensor_core.cc | 2 +- .../schedule_rule/multi_level_tiling_with_intrin.cc | 2 +- src/s_tir/meta_schedule/schedule_rule/schedule_rule.cc | 2 +- src/s_tir/schedule/concrete_schedule.cc | 2 +- src/s_tir/schedule/primitive.h | 2 +- src/s_tir/schedule/primitive/blockize_tensorize.cc | 2 +- src/s_tir/schedule/transform.cc | 2 +- src/s_tir/{function.cc => tensor_intrin.cc} | 4 ++-- 11 files changed, 14 insertions(+), 14 deletions(-) rename include/tvm/s_tir/{function.h => tensor_intrin.h} (95%) rename python/tvm/s_tir/{function.py => _tensor_intrin.py} (100%) rename src/s_tir/{function.cc => tensor_intrin.cc} (97%) diff --git a/include/tvm/s_tir/function.h b/include/tvm/s_tir/tensor_intrin.h similarity index 95% rename from include/tvm/s_tir/function.h rename to include/tvm/s_tir/tensor_intrin.h index 73ae70084fcb..cff27d56a1a9 100644 --- a/include/tvm/s_tir/function.h +++ b/include/tvm/s_tir/tensor_intrin.h @@ -18,11 +18,11 @@ */ /*! - * \file tvm/s_tir/function.h + * \file tvm/s_tir/tensor_intrin.h * \brief Tensor intrinsics for schedulable TIR. */ -#ifndef TVM_S_TIR_FUNCTION_H_ -#define TVM_S_TIR_FUNCTION_H_ +#ifndef TVM_S_TIR_TENSOR_INTRIN_H_ +#define TVM_S_TIR_TENSOR_INTRIN_H_ #include @@ -87,4 +87,4 @@ class TensorIntrin : public ffi::ObjectRef { } // namespace s_tir } // namespace tvm -#endif // TVM_S_TIR_FUNCTION_H_ +#endif // TVM_S_TIR_TENSOR_INTRIN_H_ diff --git a/python/tvm/s_tir/__init__.py b/python/tvm/s_tir/__init__.py index 6586b90f4d35..7bb80b3419d4 100644 --- a/python/tvm/s_tir/__init__.py +++ b/python/tvm/s_tir/__init__.py @@ -18,7 +18,7 @@ # pylint: disable=invalid-name """S-TIR namespace for scheduable TensorIR""" -from .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), diff --git a/python/tvm/s_tir/function.py b/python/tvm/s_tir/_tensor_intrin.py similarity index 100% rename from python/tvm/s_tir/function.py rename to python/tvm/s_tir/_tensor_intrin.py 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 ccf35589d727..c6258227e456 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 @@ -18,9 +18,9 @@ */ #include #include -#include #include #include +#include #include #include 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 571820508a75..c499820395e6 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 @@ -18,8 +18,8 @@ */ #include -#include #include +#include #include "../../schedule/analysis.h" #include "../../schedule/transform.h" 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 0a34cd417795..4a5467247aad 100644 --- a/src/s_tir/meta_schedule/schedule_rule/schedule_rule.cc +++ b/src/s_tir/meta_schedule/schedule_rule/schedule_rule.cc @@ -18,7 +18,7 @@ */ #include #include -#include +#include #include "../utils.h" diff --git a/src/s_tir/schedule/concrete_schedule.cc b/src/s_tir/schedule/concrete_schedule.cc index ac94bd5c1ee3..903203e5a8e6 100644 --- a/src/s_tir/schedule/concrete_schedule.cc +++ b/src/s_tir/schedule/concrete_schedule.cc @@ -20,8 +20,8 @@ #include #include -#include #include +#include #include diff --git a/src/s_tir/schedule/primitive.h b/src/s_tir/schedule/primitive.h index d7f0a77c0f1d..89c058f9ba70 100644 --- a/src/s_tir/schedule/primitive.h +++ b/src/s_tir/schedule/primitive.h @@ -20,10 +20,10 @@ #define TVM_S_TIR_SCHEDULE_PRIMITIVE_H_ #include -#include #include #include #include +#include #include diff --git a/src/s_tir/schedule/primitive/blockize_tensorize.cc b/src/s_tir/schedule/primitive/blockize_tensorize.cc index f8343288d26a..6a5c9ea83c28 100644 --- a/src/s_tir/schedule/primitive/blockize_tensorize.cc +++ b/src/s_tir/schedule/primitive/blockize_tensorize.cc @@ -21,8 +21,8 @@ #include #include #include -#include #include +#include #include diff --git a/src/s_tir/schedule/transform.cc b/src/s_tir/schedule/transform.cc index 9de1345711cb..bb0e1565c34a 100644 --- a/src/s_tir/schedule/transform.cc +++ b/src/s_tir/schedule/transform.cc @@ -20,8 +20,8 @@ #include #include #include -#include #include +#include #include #include "../../tirx/transform/ir_utils.h" diff --git a/src/s_tir/function.cc b/src/s_tir/tensor_intrin.cc similarity index 97% rename from src/s_tir/function.cc rename to src/s_tir/tensor_intrin.cc index c33a538790bf..e381448198d5 100644 --- a/src/s_tir/function.cc +++ b/src/s_tir/tensor_intrin.cc @@ -18,11 +18,11 @@ */ /*! - * \file src/s_tir/function.cc + * \file src/s_tir/tensor_intrin.cc * \brief Tensor intrinsic registry and construction. */ #include -#include +#include namespace tvm { namespace s_tir { From b8f114264c0ab77e41138f3158fe47d5289030d3 Mon Sep 17 00:00:00 2001 From: Tianqi Chen Date: Thu, 17 Sep 2026 02:16:53 +0000 Subject: [PATCH 04/18] [S-TIR] Own storage alignment collection directly Keep alignment annotations and their collector with the S-TIR compaction and lowering consumers. Visit block and allocation annotations directly without a TIRX registration bridge. --- .../schedule/primitive/block_annotate.cc | 2 +- src/s_tir/transform/compact_buffer_region.cc | 2 +- src/s_tir/transform/ir_utils.cc | 51 +++++++++++++++++++ src/s_tir/transform/ir_utils.h | 15 ++++++ src/s_tir/transform/lower_opaque_block.cc | 2 +- src/s_tir/transform/stmt_extension.cc | 15 ------ src/tirx/transform/ir_utils.cc | 42 --------------- src/tirx/transform/ir_utils.h | 11 ---- src/tirx/transform/stmt_extension.h | 30 ----------- 9 files changed, 69 insertions(+), 101 deletions(-) diff --git a/src/s_tir/schedule/primitive/block_annotate.cc b/src/s_tir/schedule/primitive/block_annotate.cc index 97288e803ef9..6e6c0f480433 100644 --- a/src/s_tir/schedule/primitive/block_annotate.cc +++ b/src/s_tir/schedule/primitive/block_annotate.cc @@ -21,7 +21,7 @@ #include #include -#include "../../../tirx/transform/ir_utils.h" +#include "../../transform/ir_utils.h" #include "../utils.h" namespace tvm { diff --git a/src/s_tir/transform/compact_buffer_region.cc b/src/s_tir/transform/compact_buffer_region.cc index f4062b3d236e..4d3535fd4f56 100644 --- a/src/s_tir/transform/compact_buffer_region.cc +++ b/src/s_tir/transform/compact_buffer_region.cc @@ -36,10 +36,10 @@ #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/ir_utils.cc b/src/s_tir/transform/ir_utils.cc index e7bb0d9adb77..55ed50e9916f 100644 --- a/src/s_tir/transform/ir_utils.cc +++ b/src/s_tir/transform/ir_utils.cc @@ -20,6 +20,7 @@ #include "ir_utils.h" #include +#include #include namespace tvm { @@ -73,5 +74,55 @@ Region ConvertRegion(const MatchBufferRegion& match_buffer, const Region& region return result; } +/*! \brief Collect storage alignment information from annotations. */ +class StorageAlignCollector : public StmtExprVisitor { + public: + ffi::Optional Visit(ffi::AnyView value) override { + if (value.as()) return std::nullopt; + return StmtExprVisitor::Visit(value); + } + + private: + friend std::unordered_map CollectStorageAlignAnnotation( + const Stmt& body); + + /*! \brief SBlock: resolve each annotation's buffer index through the write regions. */ + ffi::Optional 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(); + for (const auto& item : annotation) { + storage_align_[op->writes[item.get<0>()]->buffer.var()].push_back(item); + } + } + return StmtExprVisitor::Visit_(op); + } + + /*! \brief AllocBuffer: check for buffer_dim_align annotations. */ + ffi::Optional 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(); + 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 storage_align_; +}; + +std::unordered_map CollectStorageAlignAnnotation(const Stmt& body) { + auto collector = ffi::make_object(); + 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 index dcc0b0bcf73e..eaf925ac6f21 100644 --- a/src/s_tir/transform/ir_utils.h +++ b/src/s_tir/transform/ir_utils.h @@ -20,8 +20,11 @@ #ifndef TVM_S_TIR_TRANSFORM_IR_UTILS_H_ #define TVM_S_TIR_TRANSFORM_IR_UTILS_H_ +#include #include +#include + #include "../../tirx/transform/ir_utils.h" namespace tvm { @@ -42,6 +45,18 @@ ffi::Array ConvertIndices(const MatchBufferRegion& match_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; +/*! \brief A list of StorageAlignTuple, used by StorageAlign */ +using StorageAlignAnnotation = ffi::Array; +/*! + * \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 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/lower_opaque_block.cc b/src/s_tir/transform/lower_opaque_block.cc index fabb258aed00..1ac40d183247 100644 --- a/src/s_tir/transform/lower_opaque_block.cc +++ b/src/s_tir/transform/lower_opaque_block.cc @@ -27,7 +27,7 @@ #include #include -#include "../../tirx/transform/ir_utils.h" +#include "ir_utils.h" namespace tvm { namespace s_tir { diff --git a/src/s_tir/transform/stmt_extension.cc b/src/s_tir/transform/stmt_extension.cc index bcd902270523..64a03e4bc09a 100644 --- a/src/s_tir/transform/stmt_extension.cc +++ b/src/s_tir/transform/stmt_extension.cc @@ -110,21 +110,6 @@ TVM_FFI_STATIC_INIT_BLOCK() { return StmtExprVisitor::VisitBlock(self, block); }); }); - tirx::StorageAlignVisitor::RegisterExtension([](tirx::StorageAlignVisitor::VTable* table) { - table->ClearDispatch(); - table->SetDispatch([](const ffi::Object* node, ObjectVisitor* base) { - auto* self = static_cast(base); - auto* block = static_cast(node); - auto it = block->annotations.find(attr::buffer_dim_align); - if (it != block->annotations.end()) { - auto annotation = (*it).second.as_or_throw(); - for (const auto& item : annotation) { - self->RecordAlignment(block->writes[item.get<0>()]->buffer.var(), item); - } - } - return StmtExprVisitor::VisitBlock(self, block); - }); - }); } } // namespace diff --git a/src/tirx/transform/ir_utils.cc b/src/tirx/transform/ir_utils.cc index 81f7d000c06d..b8601952990b 100644 --- a/src/tirx/transform/ir_utils.cc +++ b/src/tirx/transform/ir_utils.cc @@ -770,48 +770,6 @@ std::pair GetAsyncWaitAttributes(const AttrStmtNode* op) { return std::make_pair(op->value, inner->value); } -/*! \brief Collect storage alignment information from annotations. */ -class StorageAlignCollector : public StorageAlignVisitor { - public: - ffi::Optional Visit(ffi::AnyView value) override { - if (value.as()) return std::nullopt; - return StmtExprVisitor::Visit(value); - } - - private: - friend std::unordered_map CollectStorageAlignAnnotation( - const Stmt& body); - - void RecordAlignment(const Var& buffer, const StorageAlignTuple& annotation) final { - storage_align_[buffer].push_back(annotation); - } - - /*! \brief AllocBuffer: check for buffer_dim_align annotations. */ - ffi::Optional Visit_(const AllocBufferNode* op) final { - auto it = op->annotations.find("buffer_dim_align"); - if (it != op->annotations.end()) { - auto storage_align_annotation = (*it).second.as_or_throw(); - 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 storage_align_; -}; - -std::unordered_map CollectStorageAlignAnnotation(const Stmt& body) { - auto collector = ffi::make_object(); - collector->Visit(body); - return std::move(collector->storage_align_); -} - int Stoi(const std::string& str) { try { return std::stoi(str); diff --git a/src/tirx/transform/ir_utils.h b/src/tirx/transform/ir_utils.h index 89895f63c9d4..05aa9e4e307c 100644 --- a/src/tirx/transform/ir_utils.h +++ b/src/tirx/transform/ir_utils.h @@ -25,7 +25,6 @@ #define TVM_TIR_TRANSFORM_IR_UTILS_H_ #include -#include #include #include #include @@ -275,16 +274,6 @@ std::unordered_map GetTensorCoreFragmentInfo(const // s_tir::attr::async_wait_queue_scope annotation. std::pair GetAsyncWaitAttributes(const AttrStmtNode* op); -/*! \brief The quad used by StorageAlign for (buffer_idx, axis, factor, offset) */ -using StorageAlignTuple = ffi::Tuple; -/*! \brief A list of StorageAlignTuple, used by StorageAlign */ -using StorageAlignAnnotation = ffi::Array; -/*! - * \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 CollectStorageAlignAnnotation(const Stmt& body); /*! * \brief Split string separated by "," to get wmma fragment dimension size. * \param shape_str The string to split. diff --git a/src/tirx/transform/stmt_extension.h b/src/tirx/transform/stmt_extension.h index 75beacf81b61..2cf709e82b71 100644 --- a/src/tirx/transform/stmt_extension.h +++ b/src/tirx/transform/stmt_extension.h @@ -27,7 +27,6 @@ #include #include "../ir_mutator_with_analyzer.h" -#include "ir_utils.h" namespace tvm { namespace tirx { @@ -125,35 +124,6 @@ class IndexDomainVisitor : public StmtExprVisitor { } }; -class StorageAlignVisitor : public StmtExprVisitor { - public: - using StmtExprVisitor::VTable; - static void RegisterExtension(void (*init)(VTable*)) { Extensions().push_back(init); } - StorageAlignVisitor() : StmtExprVisitor(GlobalVTable()) {} - virtual void RecordAlignment(const Var& buffer, const StorageAlignTuple& annotation) = 0; - - protected: - static void InitVTable(VTable* table) { - StmtExprVisitor::InitVTable(table); - for (auto init : Extensions()) init(table); - } - - private: - static std::vector& Extensions() { - static std::vector extensions; - return extensions; - } - static const VTable* GlobalVTable() { - static const VTable table = [] { - VTable table; - InitVTable(&table); - table.Finalize(); - return table; - }(); - return &table; - } -}; - } // namespace tirx } // namespace tvm #endif // TVM_TIRX_TRANSFORM_STMT_EXTENSION_H_ From ab471118be492b920b95e3311b45e782bcec502c Mon Sep 17 00:00:00 2001 From: Tianqi Chen Date: Thu, 17 Sep 2026 02:47:49 +0000 Subject: [PATCH 05/18] [S-TIR] Route block semantics through explicit native helpers Replace dialect callback registries with inherited S-TIR analyzer, SSA, path verification and simplification paths. Keep generic structural traversal definition-aware and preserve generic specialization through buffer remapping. Normalize schedulable TE and tensor intrinsic indices locally, keep ordinary dtype transforms after block lowering, and preserve scalar index temporaries when narrowing allocations. --- include/tvm/s_tir/analysis.h | 5 + include/tvm/s_tir/stmt_functor.h | 13 +- include/tvm/s_tir/transform.h | 7 + include/tvm/tirx/analysis.h | 7 +- include/tvm/tirx/stmt_functor.h | 10 +- python/tvm/s_tir/analysis/__init__.py | 10 + python/tvm/s_tir/backend/adreno/pipeline.py | 2 +- python/tvm/s_tir/pipeline.py | 2 +- python/tvm/s_tir/transform/transform.py | 10 + python/tvm/script/parser/core/entry.py | 10 +- python/tvm/tirx/analysis/analysis.py | 5 +- python/tvm/tirx/function.py | 14 +- src/relax/transform/legalize_ops.cc | 3 +- src/s_tir/analysis/is_pure_function.cc | 2 +- src/s_tir/analysis/verify_well_formed.cc | 32 +- src/s_tir/ir/data_type_rewriter.cc | 243 ++-- src/s_tir/ir/data_type_rewriter.h | 69 ++ src/s_tir/ir/ir_mutator_with_analyzer.cc | 28 +- src/s_tir/ir/ir_mutator_with_analyzer.h | 15 +- src/s_tir/ir/ir_visitor_with_analyzer.cc | 24 +- src/s_tir/ir/ir_visitor_with_analyzer.h | 15 +- src/s_tir/ir/specialize.cc | 70 -- src/s_tir/ir/tir_visitor_with_path.cc | 73 +- src/s_tir/ir/tir_visitor_with_path.h | 52 + .../feature_extractor/per_store_feature.cc | 4 +- .../meta_schedule/postproc/verify_gpu_code.cc | 2 +- .../schedule/primitive/blockize_tensorize.cc | 6 +- src/s_tir/schedule/state.cc | 2 +- src/s_tir/stmt.cc | 29 +- src/s_tir/stmt_functor.cc | 52 +- src/s_tir/transform/hoist_expression.cc | 10 +- src/s_tir/transform/inject_double_buffer.cc | 4 +- .../transform/inject_software_pipeline.cc | 4 +- src/s_tir/transform/inject_virtual_thread.cc | 4 +- src/s_tir/transform/ir_utils.cc | 70 ++ src/s_tir/transform/ir_utils.h | 4 + src/s_tir/transform/loop_partition.cc | 4 +- src/s_tir/transform/stmt_extension.cc | 117 -- src/s_tir/transform/stmt_simplify.cc | 93 ++ .../transform/stmt_simplify.h} | 26 +- src/te/operation/create_primfunc.cc | 6 +- src/tirx/analysis/verify_tirx_well_formed.cc | 17 +- src/tirx/analysis/verify_well_formed.cc | 272 +---- src/tirx/analysis/verify_well_formed.h | 261 +++- src/tirx/ir/data_type_rewriter.cc | 32 +- src/tirx/ir/data_type_rewriter.h | 17 +- src/tirx/ir/ir_mutator_with_analyzer.cc | 16 +- src/tirx/ir/ir_mutator_with_analyzer.h | 6 - src/tirx/ir/ir_visitor_with_analyzer.cc | 16 +- src/tirx/ir/ir_visitor_with_analyzer.h | 7 - src/tirx/ir/specialize.cc | 41 +- src/tirx/ir/stmt_functor.cc | 24 - src/tirx/ir/tir_visitor_with_path.cc | 21 - src/tirx/ir/tir_visitor_with_path.h | 24 +- src/tirx/transform/flatten_buffer.cc | 32 +- .../transform/force_narrow_index_to_i32.cc | 14 +- .../transform/inline_private_functions.cc | 20 +- src/tirx/transform/ir_utils.cc | 1079 ++++++++--------- src/tirx/transform/ir_utils.h | 118 ++ src/tirx/transform/narrow_datatype.cc | 8 +- src/tirx/transform/stmt_extension.h | 129 -- src/tirx/transform/stmt_simplify.cc | 304 +++-- src/tirx/transform/stmt_simplify.h | 68 ++ src/tirx/transform/tvm_ffi_binder.cc | 1 + tests/cpp/s_tir_functor_test.cc | 134 +- tests/python/relax/test_transform_fuse_ops.py | 10 +- tests/python/relax/test_transform_fuse_tir.py | 4 +- ...m_specialize_primfunc_based_on_callsite.py | 7 +- ...tproc_rewrite_parallel_vectorize_unroll.py | 4 +- ...t_s_tir_transform_compact_buffer_region.py | 4 +- ..._tir_transform_convert_blocks_to_opaque.py | 2 +- .../test_s_tir_transform_convert_ssa.py | 63 + ..._tir_transform_inject_software_pipeline.py | 2 +- ...test_s_tir_transform_lower_match_buffer.py | 2 +- ...st_s_tir_transform_unify_thread_binding.py | 2 +- .../test_tir_analysis_verify_well_formed.py | 109 +- tests/python/tirx-base/test_tir_specialize.py | 44 + .../test_tir_inline_private_functions.py | 4 +- .../test_tir_transform_flatten_buffer.py | 35 +- ...tir_transform_force_narrow_index_to_i32.py | 98 +- .../test_tir_transform_narrow_datatype.py | 63 +- .../test_tir_transform_simplify.py | 21 + 82 files changed, 2201 insertions(+), 2092 deletions(-) create mode 100644 src/s_tir/ir/data_type_rewriter.h delete mode 100644 src/s_tir/ir/specialize.cc create mode 100644 src/s_tir/ir/tir_visitor_with_path.h delete mode 100644 src/s_tir/transform/stmt_extension.cc create mode 100644 src/s_tir/transform/stmt_simplify.cc rename src/{tirx/ir/specialize.h => s_tir/transform/stmt_simplify.h} (57%) delete mode 100644 src/tirx/transform/stmt_extension.h create mode 100644 tests/python/s_tir/transform/test_s_tir_transform_convert_ssa.py diff --git a/include/tvm/s_tir/analysis.h b/include/tvm/s_tir/analysis.h index 48fdcded6705..f291d2c9b6c4 100644 --- a/include/tvm/s_tir/analysis.h +++ b/include/tvm/s_tir/analysis.h @@ -98,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/stmt_functor.h b/include/tvm/s_tir/stmt_functor.h index 80d6da180bd9..9b48e4ef0135 100644 --- a/include/tvm/s_tir/stmt_functor.h +++ b/include/tvm/s_tir/stmt_functor.h @@ -68,9 +68,10 @@ class StmtFunctor * 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. - * S-TIR registers this same native policy into generic TIRX visitors. This - * subclass additionally exposes virtual block hooks for block-aware passes. - * Other foreign nodes without a native policy still use structural fallback. + * 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: @@ -97,9 +98,9 @@ class TVM_DLL StmtExprVisitor : public tirx::StmtExprVisitor { * * Reuses inherited remapping and ownership checks. Block annotations are left * intact; structural mutation separately provides the full field rewrite. - * S-TIR registers the same policy into generic TIRX mutators so allocation - * remaps precede region uses and block binders retain their existing identity. - * Foreign nodes without a registered policy still use structural fallback. + * 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: diff --git a/include/tvm/s_tir/transform.h b/include/tvm/s_tir/transform.h index f898be2e467e..8de8210ff73f 100644 --- a/include/tvm/s_tir/transform.h +++ b/include/tvm/s_tir/transform.h @@ -48,6 +48,13 @@ namespace transform { using tirx::transform::CreatePrimFuncPass; using tvm::transform::Pass; + +/*! \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(); using tvm::transform::PassContext; /*! diff --git a/include/tvm/tirx/analysis.h b/include/tvm/tirx/analysis.h index ee163d5a26e6..11fa65f0b0cf 100644 --- a/include/tvm/tirx/analysis.h +++ b/include/tvm/tirx/analysis.h @@ -130,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 s_tir::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/stmt_functor.h b/include/tvm/tirx/stmt_functor.h index 6d3a8ef3ffbc..637a65a71621 100644 --- a/include/tvm/tirx/stmt_functor.h +++ b/include/tvm/tirx/stmt_functor.h @@ -118,6 +118,10 @@ class StmtFunctor { // Register inherited hooks in a fresh table before adding dialect nodes. static void InitVTable(VTable* vtable) { + vtable->template SetDispatch( + [](const ffi::ObjectRef& node, TSelf* self, Args... args) { + return self->VisitStmtDefault_(node.get(), std::forward(args)...); + }); IR_STMT_FUNCTOR_DISPATCH(BindNode); IR_STMT_FUNCTOR_DISPATCH(AttrStmtNode); IR_STMT_FUNCTOR_DISPATCH(IfThenElseNode); @@ -185,9 +189,6 @@ class StmtFunctor { */ class TVM_DLL StmtExprVisitor : public tvm::ExprVisitor { public: - using tvm::ExprVisitor::VTable; - // Register dialect hooks during library initialization, before first table use. - static void RegisterExtension(void (*init)(VTable*)); TVM_DEFINE_OBJECT_FUNCTOR_DEFAULT_CONSTRUCTOR(StmtExprVisitor, tvm::ExprVisitor) using tvm::ExprVisitor::Visit; @@ -240,9 +241,6 @@ class TVM_DLL StmtExprVisitor : public tvm::ExprVisitor { */ class TVM_DLL StmtExprMutator : public tvm::ExprMutator { public: - using tvm::ExprMutator::VTable; - // Register dialect hooks during library initialization, before first table use. - static void RegisterExtension(void (*init)(VTable*)); TVM_DEFINE_OBJECT_FUNCTOR_DEFAULT_CONSTRUCTOR(StmtExprMutator, tvm::ExprMutator) using tvm::ExprMutator::Mutate; using tvm::ExprMutator::Mutate_; diff --git a/python/tvm/s_tir/analysis/__init__.py b/python/tvm/s_tir/analysis/__init__.py index 0961ebcafad4..817e41472b3e 100644 --- a/python/tvm/s_tir/analysis/__init__.py +++ b/python/tvm/s_tir/analysis/__init__.py @@ -215,3 +215,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 7dd65a4b3016..1845f2f16ea5 100644 --- a/python/tvm/s_tir/backend/adreno/pipeline.py +++ b/python/tvm/s_tir/backend/adreno/pipeline.py @@ -44,7 +44,7 @@ def _pipeline(mod: tvm.ir.IRModule, _ctx: tvm.transform.PassContext) -> tvm.ir.I 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/pipeline.py b/python/tvm/s_tir/pipeline.py index 81a9fad21f69..283ec8fc9245 100644 --- a/python/tvm/s_tir/pipeline.py +++ b/python/tvm/s_tir/pipeline.py @@ -45,7 +45,7 @@ def _pipeline(mod: tvm.ir.IRModule, _ctx: tvm.transform.PassContext) -> tvm.ir.I 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/transform/transform.py b/python/tvm/s_tir/transform/transform.py index 0a56fe5bea85..6bbaeb0f33dc 100644 --- a/python/tvm/s_tir/transform/transform.py +++ b/python/tvm/s_tir/transform/transform.py @@ -22,6 +22,16 @@ 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 ec5b5d7ab114..cf3641d3b6fd 100644 --- a/python/tvm/script/parser/core/entry.py +++ b/python/tvm/script/parser/core/entry.py @@ -142,10 +142,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/tirx/analysis/analysis.py b/python/tvm/tirx/analysis/analysis.py index f48710a68312..faf2bf0b476a 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 708f55864976..d6d58c2d7909 100644 --- a/python/tvm/tirx/function.py +++ b/python/tvm/tirx/function.py @@ -119,19 +119,17 @@ def specialize(self, param_map: Mapping[Var, Expr | Buffer]): 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 @@ def mem_copy(a: T.handle, b: T.handle, m: T.int32, n: T.int32) -> None: .. 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 ------- diff --git a/src/relax/transform/legalize_ops.cc b/src/relax/transform/legalize_ops.cc index 12a51827926e..ecc1f0e0dbb7 100644 --- a/src/relax/transform/legalize_ops.cc +++ b/src/relax/transform/legalize_ops.cc @@ -31,6 +31,7 @@ #include #include #include +#include #include #include @@ -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/s_tir/analysis/is_pure_function.cc b/src/s_tir/analysis/is_pure_function.cc index 47a7f3deec2a..a5fb6a692cbc 100644 --- a/src/s_tir/analysis/is_pure_function.cc +++ b/src/s_tir/analysis/is_pure_function.cc @@ -28,7 +28,7 @@ #include #include -#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/verify_well_formed.cc b/src/s_tir/analysis/verify_well_formed.cc index f98b45eb0387..24381365d8b7 100644 --- a/src/s_tir/analysis/verify_well_formed.cc +++ b/src/s_tir/analysis/verify_well_formed.cc @@ -19,8 +19,12 @@ #include "../../tirx/analysis/verify_well_formed.h" +#include +#include #include +#include "../ir/tir_visitor_with_path.h" + namespace tvm { namespace s_tir { using tirx::BufferRegion; @@ -133,6 +137,32 @@ class BlockVarAccessVerifier : public StmtExprVisitor { bool has_error_{false}; }; -TVM_FFI_STATIC_INIT_BLOCK() { tirx::RegisterWellFormedExtension(BlockVarAccessVerifier::Verify); } +bool VerifyWellFormed(const tirx::PrimFunc& func, bool assert_mode) { + return BlockVarAccessVerifier::Verify(func, assert_mode) && + tirx::VerifyWellFormedCommon(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()) { + if (!BlockVarAccessVerifier::Verify(func.value(), assert_mode)) return false; + } + } + return tirx::VerifyWellFormedCommon(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()) { + return s_tir::VerifyWellFormed(func.value(), assert_mode); + } + if (auto mod = obj.as()) { + return s_tir::VerifyWellFormed(mod.value(), assert_mode); + } + TVM_FFI_THROW(TypeError) << "Expected a PrimFunc or IRModule, but received " + << obj->GetTypeKey(); + }); +} } // namespace s_tir } // namespace tvm diff --git a/src/s_tir/ir/data_type_rewriter.cc b/src/s_tir/ir/data_type_rewriter.cc index 4c1cefdcc939..d70e392489d8 100644 --- a/src/s_tir/ir/data_type_rewriter.cc +++ b/src/s_tir/ir/data_type_rewriter.cc @@ -17,128 +17,82 @@ * under the License. */ -#include "../../tirx/ir/data_type_rewriter.h" +#include "data_type_rewriter.h" -#include #include #include namespace tvm { -namespace tirx { +namespace s_tir { +using namespace tvm::tirx; using namespace tvm::prim; -using s_tir::MatchBufferRegion; -using s_tir::SBlock; -using s_tir::SBlockNode; -using s_tir::SBlockRealize; -using s_tir::SBlockRealizeNode; -class DataTypeLegalizer::Extension { - public: - static void InitVTable(VTable* vtable); - static UnchangedOr MutateBlockRealize(DataTypeLegalizer* self, - const s_tir::SBlockRealizeNode* op, - InplaceMode inplace_mode); - static UnchangedOr MutateBlock(DataTypeLegalizer* self, const s_tir::SBlockNode* op, - InplaceMode inplace_mode); -}; -class IndexDataTypeRewriter::Extension { - public: - static void InitVTable(VTable* vtable); - static UnchangedOr MutateBlockRealize(IndexDataTypeRewriter* self, - const s_tir::SBlockRealizeNode* op, - InplaceMode inplace_mode); - static UnchangedOr MutateBlock(IndexDataTypeRewriter* self, const s_tir::SBlockNode* op, - InplaceMode inplace_mode); - static ffi::Map VisitBlockAnnotations( - IndexDataTypeRewriter* self, const ffi::Map& annotations); - static IterVar VisitIterVar(IndexDataTypeRewriter* self, const IterVar& iter_var); - static BufferRegion VisitBufferRegion(IndexDataTypeRewriter* self, - const BufferRegion& buffer_region); -}; - -UnchangedOr DataTypeLegalizer::Extension::MutateBlockRealize(DataTypeLegalizer* self, - const SBlockRealizeNode* op, - InplaceMode inplace_mode) { - SBlockRealize realize = s_tir::StmtExprMutator::MutateBlockRealize(self, op, inplace_mode) - .ValueOrUnchanged(ffi::GetRef(op)) - .as_or_throw(); - ffi::Array new_iter_values; - bool changed = false; - for (int i = 0; i < static_cast(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]); +PrimFunc IndexDataTypeNormalizer::Rewrite(PrimFunc func) { + // 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 collect) + : IndexDataTypeNormalizer(PrimType::Int(64)), collect_(std::move(collect)) {} + using IndexDataTypeNormalizer::Mutate; + using IndexDataTypeNormalizer::Mutate_; + UnchangedOr Mutate_(const VarNode* op, InplaceMode mode) final { + if (def_region_kind() == kTVMFFIDefRegionKindNone && is_enabled_) collect_(op); + return IndexDataTypeNormalizer::Mutate_(op, mode); } - } - if (changed) { - realize.CopyOnWrite()->iter_values = std::move(new_iter_values); - } - return realize; -} -UnchangedOr DataTypeLegalizer::Extension::MutateBlock(DataTypeLegalizer* self, - const SBlockNode* op, - InplaceMode inplace_mode) { - SBlock new_block = s_tir::StmtExprMutator::MutateBlock(self, op, inplace_mode) - .ValueOrUnchanged(ffi::GetRef(op)) - .as_or_throw(); - ffi::Array 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; + protected: + bool CanRewriteDType(PrimType dtype) const final { return false; } + + private: + std::function collect_; + }; + auto seed = [this](const VarNode* var) { + auto dtype = var->ty.as(); + if (dtype && CanRewriteDType(dtype.value()) && dtype.value() != target_data_type_ && + VarRemapGet(ffi::AnyView(var)) == nullptr) { + VarRemapSet(ffi::AnyView(var), ffi::GetRef(var).CopyWithDType(target_data_type_)); + } + }; + auto collector = ffi::make_object(seed); + collector->Mutate(func->body); + for (const Var& param : func->params) { + if (param.as()) { + collector->WithDefRegionKind(kTVMFFIDefRegionKindSimple, + [&] { return collector->Mutate(param); }); } else { - return iter; + seed(param.get()); } - }); - if (!op->iter_vars.same_as(new_iter_vars)) { - new_block.CopyOnWrite()->iter_vars = std::move(new_iter_vars); } - return new_block; -} - -void DataTypeLegalizer::Extension::InitVTable(VTable* vtable) { - vtable->ClearDispatch(); - vtable->ClearDispatch(); - vtable->SetDispatch( - [](const ffi::Object* node, ObjectMutator* base, InplaceMode mode) -> UnchangedOr { - return MutateBlock(static_cast(base), - static_cast(node), mode); - }); - vtable->SetDispatch( - [](const ffi::Object* node, ObjectMutator* base, InplaceMode mode) -> UnchangedOr { - return MutateBlockRealize(static_cast(base), - static_cast(node), mode); - }); -} -TVM_FFI_STATIC_INIT_BLOCK() { - DataTypeLegalizer::RegisterExtension(DataTypeLegalizer::Extension::InitVTable); + ffi::Array params = func->params.Map([this](const Var& param) { + return WithDefRegionKind(kTVMFFIDefRegionKindSimple, [&] { + return Mutate(param).ValueOrUnchanged(param).as_or_throw(); + }); + }); + 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 IndexDataTypeRewriter::Extension::MutateBlockRealize(IndexDataTypeRewriter* self, - const SBlockRealizeNode* op, - InplaceMode inplace_mode) { - bool is_condition = self->is_condition_; - self->is_condition_ = true; - auto new_predicate_result = self->Mutate(op->predicate, inplace_mode); +UnchangedOr 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); - self->is_condition_ = is_condition; + this->is_condition_ = is_condition; - bool is_enabled = self->is_enabled_; - self->is_enabled_ = true; - auto new_iter_values = self->Mutate(op->iter_values, inplace_mode) + bool is_enabled = this->is_enabled_; + this->is_enabled_ = true; + auto new_iter_values = this->Mutate(op->iter_values, inplace_mode) .as_or_throw>>() .ValueOrUnchanged(op->iter_values); - self->is_enabled_ = is_enabled; + this->is_enabled_ = is_enabled; SBlock new_body = - self->Mutate(op->block, inplace_mode).ValueOrUnchanged(op->block).as_or_throw(); + this->Mutate(op->block, inplace_mode).ValueOrUnchanged(op->block).as_or_throw(); if (!new_predicate_unchanged || !new_iter_values.same_as(op->iter_values) || !new_body.same_as(op->block)) { SBlockRealize new_block_realize = ffi::GetRef(op); @@ -153,36 +107,34 @@ UnchangedOr IndexDataTypeRewriter::Extension::MutateBlockRealize(IndexData } } -UnchangedOr IndexDataTypeRewriter::Extension::MutateBlock(IndexDataTypeRewriter* self, - const SBlockNode* op, - InplaceMode inplace_mode) { - auto new_alloc_buffers = self->WithDefRegionKind(kTVMFFIDefRegionKindSimple, [&] { - return self->Mutate(op->alloc_buffers, inplace_mode) +UnchangedOr 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>>() .ValueOrUnchanged(op->alloc_buffers); }); - auto new_match_buffers = op->match_buffers.Map([self](const MatchBufferRegion& match) { - BufferVar buffer = self->WithDefRegionKind(kTVMFFIDefRegionKindSimple, [&] { - return self->Mutate(match->buffer, InplaceMode::kDisallow) + 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>() .ValueOrUnchanged(match->buffer); }); - BufferRegion source = VisitBufferRegion(self, match->source); + BufferRegion source = VisitBufferRegion(match->source); if (buffer.same_as(match->buffer) && source.same_as(match->source)) return match; return MatchBufferRegion(buffer, source); }); ffi::Array new_reads = op->reads.Map( - [self](const BufferRegion& buffer_region) { return VisitBufferRegion(self, buffer_region); }); + [this](const BufferRegion& buffer_region) { return VisitBufferRegion(buffer_region); }); ffi::Array new_writes = op->writes.Map( - [self](const BufferRegion& buffer_region) { return VisitBufferRegion(self, buffer_region); }); + [this](const BufferRegion& buffer_region) { return VisitBufferRegion(buffer_region); }); ffi::Array new_iter_vars = - op->iter_vars.Map([self](const IterVar& iter_var) { return VisitIterVar(self, iter_var); }); + op->iter_vars.Map([this](const IterVar& iter_var) { return VisitIterVar(iter_var); }); ffi::Optional new_init = std::nullopt; if (op->init.has_value()) { - new_init = self->Mutate(op->init.value(), inplace_mode).ValueOrUnchanged(op->init.value()); + new_init = this->Mutate(op->init.value(), inplace_mode).ValueOrUnchanged(op->init.value()); } - ffi::Map new_annotations = VisitBlockAnnotations(self, op->annotations); - auto new_body_result = self->Mutate(op->body, inplace_mode); + ffi::Map 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); @@ -201,24 +153,22 @@ UnchangedOr IndexDataTypeRewriter::Extension::MutateBlock(IndexDataTypeRew n->init = std::move(new_init); n->annotations = std::move(new_annotations); n->body = std::move(new_body); - for (const auto& buffer : new_block->alloc_buffers) self->ValidateAllocation(buffer); return new_block; } - for (const auto& buffer : op->alloc_buffers) self->ValidateAllocation(buffer); return ffi::Unchanged(); } -ffi::Map IndexDataTypeRewriter::Extension::VisitBlockAnnotations( - IndexDataTypeRewriter* self, const ffi::Map& annotations) { +ffi::Map IndexDataTypeNormalizer::VisitBlockAnnotations( + const ffi::Map& annotations) { auto new_annotations = annotations; - std::function f_mutate_obj = [self, &f_mutate_obj](const Any& obj) -> Any { + std::function f_mutate_obj = [this, &f_mutate_obj](const Any& obj) -> Any { if (obj == nullptr) { return obj; } if (auto var = obj.as(); var && var.value()->ty.as()) { BufferVar buffer(var.value()); - if (BufferVar new_buffer = self->Mutate(buffer, InplaceMode::kDisallow) + if (BufferVar new_buffer = this->Mutate(buffer, InplaceMode::kDisallow) .as_or_throw>() .ValueOrUnchanged(buffer); !new_buffer.same_as(buffer)) { @@ -240,18 +190,17 @@ ffi::Map IndexDataTypeRewriter::Extension::VisitBlockAnno return new_annotations; } -IterVar IndexDataTypeRewriter::Extension::VisitIterVar(IndexDataTypeRewriter* self, - const IterVar& iter_var) { - bool is_enabled = self->is_enabled_; - self->is_enabled_ = true; - PrimVar new_var = self->Mutate(iter_var->var, InplaceMode::kDisallow) +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(); PrimExpr min = - self->Mutate(iter_var->dom->min, InplaceMode::kDisallow).ValueOrUnchanged(iter_var->dom->min); - PrimExpr extent = self->Mutate(iter_var->dom->extent, InplaceMode::kDisallow) + 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); - self->is_enabled_ = is_enabled; + 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; @@ -263,20 +212,19 @@ IterVar IndexDataTypeRewriter::Extension::VisitIterVar(IndexDataTypeRewriter* se return iter_var; } -BufferRegion IndexDataTypeRewriter::Extension::VisitBufferRegion( - IndexDataTypeRewriter* self, const BufferRegion& buffer_region) { - BufferVar remapped_buffer = self->Mutate(buffer_region->buffer, InplaceMode::kDisallow) +BufferRegion IndexDataTypeNormalizer::VisitBufferRegion(const BufferRegion& buffer_region) { + BufferVar remapped_buffer = this->Mutate(buffer_region->buffer, InplaceMode::kDisallow) .as_or_throw>() .ValueOrUnchanged(buffer_region->buffer); - bool is_enabled = self->is_enabled_; - self->is_enabled_ = true; + bool is_enabled = this->is_enabled_; + this->is_enabled_ = true; auto new_region = buffer_region->region.Map([&](const Range& range) { return Range::FromMinExtent( - self->Mutate(range->min, InplaceMode::kDisallow).ValueOrUnchanged(range->min), - self->Mutate(range->extent, InplaceMode::kDisallow).ValueOrUnchanged(range->extent)); + this->Mutate(range->min, InplaceMode::kDisallow).ValueOrUnchanged(range->min), + this->Mutate(range->extent, InplaceMode::kDisallow).ValueOrUnchanged(range->extent)); }); - self->is_enabled_ = is_enabled; + this->is_enabled_ = is_enabled; if (!remapped_buffer.same_as(buffer_region->buffer) || !new_region.same_as(buffer_region->region)) { @@ -286,22 +234,5 @@ BufferRegion IndexDataTypeRewriter::Extension::VisitBufferRegion( } } -void IndexDataTypeRewriter::Extension::InitVTable(VTable* vtable) { - vtable->ClearDispatch(); - vtable->SetDispatch( - [](const ffi::Object* node, ObjectMutator* base, InplaceMode mode) -> UnchangedOr { - return MutateBlock(static_cast(base), - static_cast(node), mode); - }); - vtable->ClearDispatch(); - vtable->SetDispatch( - [](const ffi::Object* node, ObjectMutator* base, InplaceMode mode) -> UnchangedOr { - return MutateBlockRealize(static_cast(base), - static_cast(node), mode); - }); -} -TVM_FFI_STATIC_INIT_BLOCK() { - IndexDataTypeRewriter::RegisterExtension(IndexDataTypeRewriter::Extension::InitVTable); -} -} // namespace tirx +} // 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 000000000000..a5a74feaf0c8 --- /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 + +#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); + + UnchangedOr Mutate_(const SBlockNode* op, InplaceMode inplace_mode); + UnchangedOr Mutate_(const SBlockRealizeNode* op, InplaceMode inplace_mode); + + protected: + static void InitVTable(VTable* table) { + Parent::InitVTable(table); + SetDispatch(table); + SetDispatch(table); + } + static const VTable* GlobalVTable() { + static const VTable table = [] { + VTable table; + InitVTable(&table); + table.Finalize(); + return table; + }(); + return &table; + } + + private: + ffi::Map VisitBlockAnnotations( + const ffi::Map& annotations); + tirx::IterVar VisitIterVar(const tirx::IterVar& iter_var); + tirx::BufferRegion VisitBufferRegion(const tirx::BufferRegion& buffer_region); +}; + +} // namespace s_tir +} // namespace tvm +#endif // TVM_S_TIR_IR_DATA_TYPE_REWRITER_H_ diff --git a/src/s_tir/ir/ir_mutator_with_analyzer.cc b/src/s_tir/ir/ir_mutator_with_analyzer.cc index f97f340ece73..bd68223887e7 100644 --- a/src/s_tir/ir/ir_mutator_with_analyzer.cc +++ b/src/s_tir/ir/ir_mutator_with_analyzer.cc @@ -20,29 +20,17 @@ #include "ir_mutator_with_analyzer.h" namespace tvm { -namespace tirx { -UnchangedOr IRMutatorWithAnalyzer::Extension::MutateBlock(IRMutatorWithAnalyzer* self, - const s_tir::SBlockNode* op, - InplaceMode inplace_mode) { - return self->constraint_scope_.WithNewScope([&]() -> UnchangedOr { +namespace s_tir { +using namespace tirx; +UnchangedOr IRMutatorWithAnalyzer::Mutate_(const SBlockNode* op, InplaceMode inplace_mode) { + return constraint_scope_.WithNewScope([&]() -> UnchangedOr { for (const auto& iter_var : op->iter_vars) { - self->analyzer_->Bind(iter_var->var, iter_var->dom); - self->iter_vars_.Set(iter_var->var, iter_var->dom); + analyzer_->Bind(iter_var->var, iter_var->dom); + iter_vars_.Set(iter_var->var, iter_var->dom); } - return s_tir::StmtExprMutator::MutateBlock(self, op, inplace_mode); + return s_tir::StmtExprMutator::MutateBlock(this, op, inplace_mode); }); } -void IRMutatorWithAnalyzer::Extension::InitVTable(VTable* vtable) { - vtable->ClearDispatch(); - vtable->SetDispatch( - [](const ffi::Object* node, ObjectMutator* base, InplaceMode mode) -> UnchangedOr { - return MutateBlock(static_cast(base), - static_cast(node), mode); - }); -} -TVM_FFI_STATIC_INIT_BLOCK() { - IRMutatorWithAnalyzer::RegisterExtension(IRMutatorWithAnalyzer::Extension::InitVTable); -} -} // namespace tirx +} // namespace s_tir } // namespace tvm diff --git a/src/s_tir/ir/ir_mutator_with_analyzer.h b/src/s_tir/ir/ir_mutator_with_analyzer.h index 12a0a4672c31..1ca5e52bed59 100644 --- a/src/s_tir/ir/ir_mutator_with_analyzer.h +++ b/src/s_tir/ir/ir_mutator_with_analyzer.h @@ -25,14 +25,6 @@ #include "../../tirx/ir_mutator_with_analyzer.h" namespace tvm { -namespace tirx { -class IRMutatorWithAnalyzer::Extension { - public: - static UnchangedOr MutateBlock(IRMutatorWithAnalyzer* self, const s_tir::SBlockNode* op, - InplaceMode inplace_mode); - static void InitVTable(VTable* vtable); -}; -} // namespace tirx namespace s_tir { class IRMutatorWithAnalyzer : public tirx::IRMutatorWithAnalyzer { public: @@ -42,15 +34,16 @@ class IRMutatorWithAnalyzer : public tirx::IRMutatorWithAnalyzer { explicit IRMutatorWithAnalyzer(const arith::Analyzer& analyzer) : IRMutatorWithAnalyzer(analyzer.get()) {} explicit IRMutatorWithAnalyzer(arith::AnalyzerObj* analyzer) : Parent(analyzer, GlobalVTable()) {} - virtual UnchangedOr Mutate_(const SBlockNode* op, InplaceMode inplace_mode) { - return Parent::Extension::MutateBlock(this, op, inplace_mode); + virtual UnchangedOr Mutate_(const SBlockNode* op, InplaceMode inplace_mode); + virtual UnchangedOr 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); - vtable->ClearDispatch(); SetDispatch(vtable); + SetDispatch(vtable); } static const VTable* GlobalVTable() { static const VTable table = [] { diff --git a/src/s_tir/ir/ir_visitor_with_analyzer.cc b/src/s_tir/ir/ir_visitor_with_analyzer.cc index de5e01fe7095..56834a8b98cb 100644 --- a/src/s_tir/ir/ir_visitor_with_analyzer.cc +++ b/src/s_tir/ir/ir_visitor_with_analyzer.cc @@ -20,26 +20,16 @@ #include "ir_visitor_with_analyzer.h" namespace tvm { -namespace tirx { -ffi::Optional IRVisitorWithAnalyzer::Extension::VisitBlock( - IRVisitorWithAnalyzer* self, const s_tir::SBlockNode* op) { - return self->constraint_scope_.WithNewScope([&]() -> ffi::Optional { +namespace s_tir { +using namespace tirx; +ffi::Optional IRVisitorWithAnalyzer::Visit_(const SBlockNode* op) { + return constraint_scope_.WithNewScope([&]() -> ffi::Optional { for (const auto& iter_var : op->iter_vars) { - self->analyzer_->Bind(iter_var->var, iter_var->dom); + analyzer_->Bind(iter_var->var, iter_var->dom); } - return s_tir::StmtExprVisitor::VisitBlock(self, op); + return s_tir::StmtExprVisitor::VisitBlock(this, op); }); } -void IRVisitorWithAnalyzer::Extension::InitVTable(VTable* vtable) { - vtable->ClearDispatch(); - vtable->SetDispatch([](const ffi::Object* node, ObjectVisitor* base) { - return VisitBlock(static_cast(base), - static_cast(node)); - }); -} -TVM_FFI_STATIC_INIT_BLOCK() { - IRVisitorWithAnalyzer::RegisterExtension(IRVisitorWithAnalyzer::Extension::InitVTable); -} -} // namespace tirx +} // namespace s_tir } // namespace tvm diff --git a/src/s_tir/ir/ir_visitor_with_analyzer.h b/src/s_tir/ir/ir_visitor_with_analyzer.h index 87d20b82cb3f..4339077f84d3 100644 --- a/src/s_tir/ir/ir_visitor_with_analyzer.h +++ b/src/s_tir/ir/ir_visitor_with_analyzer.h @@ -25,14 +25,6 @@ #include "../../tirx/ir_visitor_with_analyzer.h" namespace tvm { -namespace tirx { -class IRVisitorWithAnalyzer::Extension { - public: - static ffi::Optional VisitBlock(IRVisitorWithAnalyzer* self, - const s_tir::SBlockNode* op); - static void InitVTable(VTable* vtable); -}; -} // namespace tirx namespace s_tir { class IRVisitorWithAnalyzer : public tirx::IRVisitorWithAnalyzer { public: @@ -40,15 +32,16 @@ class IRVisitorWithAnalyzer : public tirx::IRVisitorWithAnalyzer { using Parent::Visit; using Parent::Visit_; TVM_DEFINE_OBJECT_FUNCTOR_DEFAULT_CONSTRUCTOR(IRVisitorWithAnalyzer, Parent) - virtual ffi::Optional Visit_(const SBlockNode* op) { - return Parent::Extension::VisitBlock(this, op); + virtual ffi::Optional Visit_(const SBlockNode* op); + virtual ffi::Optional Visit_(const SBlockRealizeNode* op) { + return s_tir::StmtExprVisitor::VisitBlockRealize(this, op); } protected: static void InitVTable(VTable* vtable) { Parent::InitVTable(vtable); - vtable->ClearDispatch(); SetDispatch(vtable); + SetDispatch(vtable); } }; } // namespace s_tir diff --git a/src/s_tir/ir/specialize.cc b/src/s_tir/ir/specialize.cc deleted file mode 100644 index a46d00af8d5a..000000000000 --- a/src/s_tir/ir/specialize.cc +++ /dev/null @@ -1,70 +0,0 @@ -/* - * 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/ir/specialize.h" - -#include -#include - -namespace tvm { -namespace s_tir { -namespace { -ffi::Optional PlanBlockBuffers(tirx::StmtExprVisitor* planner, - const SBlockNode* op) { - // Block allocations were planned before all other block children by the specializer. - for (const tirx::BufferVar& buffer : op->alloc_buffers) { - TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(planner->WithDefRegionKind( - kTVMFFIDefRegionKindSimple, [&]() { return planner->Visit(buffer); })); - TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(planner->VisitBufferMetadata(buffer)); - } - for (const tirx::IterVar& iter : op->iter_vars) { - TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(planner->Visit(iter->dom->min)); - TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(planner->Visit(iter->dom->extent)); - } - for (const tirx::BufferRegion& region : op->reads) { - TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(planner->Visit(region)); - } - for (const tirx::BufferRegion& region : op->writes) { - TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(planner->Visit(region)); - } - for (const MatchBufferRegion& match : op->match_buffers) { - TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(planner->WithDefRegionKind( - kTVMFFIDefRegionKindSimple, [&]() { return planner->Visit(match->buffer); })); - TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(planner->VisitBufferMetadata(match->buffer)); - TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(planner->Visit(match->source)); - } - if (op->init.has_value()) { - TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(planner->Visit(op->init.value())); - } - return planner->Visit(op->body); -} - -void InitBufferPlanner(tirx::SpecializeVisitorVTable* vtable) { - vtable->ClearDispatch(); - vtable->SetDispatch([](const ffi::Object* node, ObjectVisitor* visitor) { - return PlanBlockBuffers(static_cast(visitor), - static_cast(node)); - }); -} - -} // namespace - -TVM_FFI_STATIC_INIT_BLOCK() { tirx::RegisterSpecializeBufferPlannerExtension(InitBufferPlanner); } -} // namespace s_tir -} // namespace tvm diff --git a/src/s_tir/ir/tir_visitor_with_path.cc b/src/s_tir/ir/tir_visitor_with_path.cc index 6aca37d59ab4..0a27c52f272c 100644 --- a/src/s_tir/ir/tir_visitor_with_path.cc +++ b/src/s_tir/ir/tir_visitor_with_path.cc @@ -17,32 +17,22 @@ * under the License. */ -#include "../../tirx/ir/tir_visitor_with_path.h" +#include "tir_visitor_with_path.h" #include namespace tvm { -namespace tirx { +namespace s_tir { +using namespace tirx; using AccessPath = ffi::reflection::AccessPath; -using s_tir::SBlockNode; -using s_tir::SBlockRealizeNode; -class TIRVisitorWithPath::Extension { - public: - static void InitVTable(VTable* vtable); - static void VisitBlock(TIRVisitorWithPath* self, const SBlockNode* op, AccessPath path); - static void VisitBlockRealize(TIRVisitorWithPath* self, const SBlockRealizeNode* op, - AccessPath path); -}; - -void TIRVisitorWithPath::Extension::VisitBlock(TIRVisitorWithPath* self, const SBlockNode* op, - AccessPath path) { +void TIRVisitorWithPath::VisitStmt_(const SBlockNode* op, AccessPath path) { std::vector, DefContext, DefContext>> context; { auto iter_path = path->Attr("iter_vars"); for (size_t i = 0; i < op->iter_vars.size(); i++) { - context.push_back(self->WithDef(op->iter_vars[i], iter_path->ArrayItem(i))); + context.push_back(WithDef(op->iter_vars[i], iter_path->ArrayItem(i))); } } @@ -53,62 +43,39 @@ void TIRVisitorWithPath::Extension::VisitBlock(TIRVisitorWithPath* self, const S 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(self->WithDef(buf, buffer_path)); + context.push_back(WithDef(buf, buffer_path)); } } - self->Visit(op->reads, path->Attr("reads")); - self->Visit(op->writes, path->Attr("writes")); - { auto match_path = path->Attr("match_buffers"); - for (size_t i = 0; i < op->match_buffers.size(); ++i) { - self->Visit(op->match_buffers[i]->source, match_path->ArrayItem(i)->Attr("source")); - } - 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 : self->WithMatchBufferDefs(buf, buffer_path)) { + for (auto& def : WithMatchBufferDefs(buf, buffer_path)) { context.push_back(std::move(def)); } - context.push_back(self->WithDef(buf, buffer_path)); + context.push_back(WithDef(buf, buffer_path)); } } - self->bind_scope_.WithNewScope([&]() { self->Visit(op->init, path->Attr("init")); }); - self->bind_scope_.WithNewScope([&]() { self->Visit(op->body, path->Attr("body")); }); + // 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::Extension::VisitBlockRealize(TIRVisitorWithPath* self, - const SBlockRealizeNode* op, - AccessPath path) { - self->Visit(op->iter_values, path->Attr("iter_values")); - self->Visit(op->predicate, path->Attr("predicate")); - self->Visit(op->block, path->Attr("block")); +void TIRVisitorWithPath::VisitStmt_(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")); } -void TIRVisitorWithPath::Extension::InitVTable(VTable* vtable) { - vtable->SetDispatch( - [](const ffi::ObjectRef& node, StmtVisitor* base, AccessPath path) { - auto* self = static_cast(base); - if (self->EnterExtensionStmt(node.get(), path)) { - VisitBlock(self, static_cast(node.get()), path); - } - }); - vtable->SetDispatch( - [](const ffi::ObjectRef& node, StmtVisitor* base, AccessPath path) { - auto* self = static_cast(base); - if (self->EnterExtensionStmt(node.get(), path)) { - VisitBlockRealize(self, static_cast(node.get()), path); - } - }); -} -TVM_FFI_STATIC_INIT_BLOCK() { - TIRVisitorWithPath::RegisterExtension(TIRVisitorWithPath::Extension::InitVTable); -} -} // namespace tirx +} // 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 000000000000..25afba3078e0 --- /dev/null +++ b/src/s_tir/ir/tir_visitor_with_path.h @@ -0,0 +1,52 @@ +/* + * 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 + +#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::VisitStmt_; + virtual void VisitStmt_(const SBlockNode* op, AccessPath path); + virtual void VisitStmt_(const SBlockRealizeNode* op, AccessPath path); + static void InitVTable(VTable* vtable) { + Parent::InitVTable(vtable); + vtable->SetDispatch( + [](const ffi::ObjectRef& node, StmtVisitor* self, AccessPath path) { + static_cast(self)->VisitStmt_( + static_cast(node.get()), path); + }); + vtable->SetDispatch( + [](const ffi::ObjectRef& node, StmtVisitor* self, AccessPath path) { + static_cast(self)->VisitStmt_( + static_cast(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 90fad9b3d3b2..123abd0d84f3 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 @@ -327,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/postproc/verify_gpu_code.cc b/src/s_tir/meta_schedule/postproc/verify_gpu_code.cc index a40d18bde7ef..cd707c1d72e8 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/schedule/primitive/blockize_tensorize.cc b/src/s_tir/schedule/primitive/blockize_tensorize.cc index 6a5c9ea83c28..fc8811822407 100644 --- a/src/s_tir/schedule/primitive/blockize_tensorize.cc +++ b/src/s_tir/schedule/primitive/blockize_tensorize.cc @@ -26,8 +26,8 @@ #include -#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" @@ -819,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/state.cc b/src/s_tir/schedule/state.cc index 86d51bbca4ce..f07a88c7ba50 100644 --- a/src/s_tir/schedule/state.cc +++ b/src/s_tir/schedule/state.cc @@ -432,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()) { auto func = opt.value(); - VerifyWellFormed(func); + s_tir::VerifyWellFormed(func); SBlockInfoCollector::Collect(self, func->body); } } diff --git a/src/s_tir/stmt.cc b/src/s_tir/stmt.cc index 2f4a78886a0c..2ec97dd3f361 100644 --- a/src/s_tir/stmt.cc +++ b/src/s_tir/stmt.cc @@ -88,15 +88,17 @@ TVMFFIAny MatchBufferRegionMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator } 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(value); TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->iter_vars)); - 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->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)); @@ -104,15 +106,13 @@ TVMFFIAny SBlockVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value) no } 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(value); TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr>, mapped_iter_vars, mutator->MutateExpected(self->iter_vars)); - TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr>, mapped_reads, - mutator->MutateExpected(self->reads)); - TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr>, mapped_writes, - mutator->MutateExpected(self->writes)); TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr>, mapped_alloc_buffers, mutator->WithDefRegionKind(kTVMFFIDefRegionKindSimple, [&]() { return mutator->MutateExpected(self->alloc_buffers); @@ -120,6 +120,10 @@ TVMFFIAny SBlockMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) n TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr>, mapped_match_buffers, mutator->MutateExpected(self->match_buffers)); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr>, mapped_reads, + mutator->MutateExpected(self->reads)); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr>, mapped_writes, + mutator->MutateExpected(self->writes)); using AnnotationMap = ffi::Map; TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, mapped_annotations, mutator->MutateExpected(self->annotations)); @@ -152,17 +156,14 @@ TVMFFIAny SBlockMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) n 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( ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value)); TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN( ffi::UnchangedOr>, mapped_iter_vars, mutator->MutateExpected(self->iter_vars, ffi::InplaceMode::kAllow)); - TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr>, mapped_reads, - mutator->MutateExpected(self->reads, ffi::InplaceMode::kAllow)); - TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN( - ffi::UnchangedOr>, mapped_writes, - mutator->MutateExpected(self->writes, ffi::InplaceMode::kAllow)); TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr>, mapped_alloc_buffers, mutator->WithDefRegionKind(kTVMFFIDefRegionKindSimple, [&]() { return mutator->MutateExpected(self->alloc_buffers, @@ -171,6 +172,11 @@ TVMFFIAny SBlockMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator, TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN( ffi::UnchangedOr>, mapped_match_buffers, mutator->MutateExpected(self->match_buffers, ffi::InplaceMode::kAllow)); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr>, mapped_reads, + mutator->MutateExpected(self->reads, ffi::InplaceMode::kAllow)); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN( + ffi::UnchangedOr>, mapped_writes, + mutator->MutateExpected(self->writes, ffi::InplaceMode::kAllow)); using AnnotationMap = ffi::Map; TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN( ffi::UnchangedOr, mapped_annotations, @@ -396,7 +402,6 @@ TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; SBlockRealizeNode::RegisterReflection(); refl::TypeAttrDef() - .def("tirx.prevent_inline", true) .attr(refl::type_attr::kStructuralVisit, reinterpret_cast(&SBlockRealizeVisit)) .attr(refl::type_attr::kStructuralMutate, reinterpret_cast(&SBlockRealizeMutate)) .attr(refl::type_attr::kStructuralMaybeInplaceMutate, diff --git a/src/s_tir/stmt_functor.cc b/src/s_tir/stmt_functor.cc index 6c8d66a65f8d..7846c595430d 100644 --- a/src/s_tir/stmt_functor.cc +++ b/src/s_tir/stmt_functor.cc @@ -20,7 +20,6 @@ * \file stmt_functor.cc * \brief Native traversal of schedulable TIR nodes. */ -#include #include #include @@ -31,49 +30,14 @@ namespace s_tir { using namespace tirx; -// Generic TIRX passes keep native dialect traversal without owning dialect nodes. -// StructuralVisitor/ObjectVisitor and structural mutation retain the full field walk. -TVM_FFI_STATIC_INIT_BLOCK() { - tirx::StmtExprVisitor::RegisterExtension([](tirx::StmtExprVisitor::VTable* vtable) { - vtable->SetDispatch([](const ffi::Object* node, ObjectVisitor* visitor) { - return StmtExprVisitor::VisitBlock(static_cast(visitor), - static_cast(node)); - }); - vtable->SetDispatch([](const ffi::Object* node, ObjectVisitor* visitor) { - return StmtExprVisitor::VisitBlockRealize(static_cast(visitor), - static_cast(node)); - }); - }); - tirx::StmtExprMutator::RegisterExtension([](tirx::StmtExprMutator::VTable* vtable) { - vtable->SetDispatch( - [](const ffi::Object* node, ObjectMutator* mutator, InplaceMode mode) { - return ffi::details::UnchangedOrUnsafe::MoveFromTVMFFIAny( - ffi::details::UnchangedOrUnsafe::MoveToTVMFFIAny( - StmtExprMutator::MutateBlock(static_cast(mutator), - static_cast(node), mode))); - }); - vtable->SetDispatch( - [](const ffi::Object* node, ObjectMutator* mutator, InplaceMode mode) { - return ffi::details::UnchangedOrUnsafe::MoveFromTVMFFIAny( - ffi::details::UnchangedOrUnsafe::MoveToTVMFFIAny(StmtExprMutator::MutateBlockRealize( - static_cast(mutator), - static_cast(node), mode))); - }); - }); -} - void StmtExprVisitor::InitVTable(VTable* vtable) { tirx::StmtExprVisitor::InitVTable(vtable); - vtable->ClearDispatch(); - vtable->ClearDispatch(); SetDispatch(vtable); SetDispatch(vtable); } void StmtExprMutator::InitVTable(VTable* vtable) { tirx::StmtExprMutator::InitVTable(vtable); - vtable->ClearDispatch(); - vtable->ClearDispatch(); SetDispatch(vtable); SetDispatch(vtable); } @@ -93,18 +57,18 @@ ffi::Optional StmtExprVisitor::VisitBlock(tirx::StmtExprVisitor* kTVMFFIDefRegionKindSimple, [&]() { return visitor->Visit(buf); })); TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitBufferMetadata(buf)); } - for (const BufferRegion& region : op->reads) { - TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->Visit(region)); - } - for (const BufferRegion& region : op->writes) { - TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->Visit(region)); - } 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 BufferRegion& region : op->reads) { + TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->Visit(region)); + } + for (const BufferRegion& 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())); } @@ -164,12 +128,12 @@ UnchangedOr StmtExprMutator::MutateBlock(tirx::StmtExprMutator* mutator, c ->WithDefRegionKind(kTVMFFIDefRegionKindSimple, [&] { return mutator->Mutate(op->alloc_buffers, inplace_mode); }) .as_or_throw>>(); + auto match_buffers = mutator->Mutate(op->match_buffers, inplace_mode) + .as_or_throw>>(); auto reads = mutator->Mutate(op->reads, inplace_mode).as_or_throw>>(); auto writes = mutator->Mutate(op->writes, inplace_mode) .as_or_throw>>(); - auto match_buffers = mutator->Mutate(op->match_buffers, inplace_mode) - .as_or_throw>>(); auto init = mutator->Mutate(op->init, inplace_mode).as_or_throw>>(); auto body = mutator->Mutate(op->body, inplace_mode); diff --git a/src/s_tir/transform/hoist_expression.cc b/src/s_tir/transform/hoist_expression.cc index 4420def0ae86..3d4c52df1fe4 100644 --- a/src/s_tir/transform/hoist_expression.cc +++ b/src/s_tir/transform/hoist_expression.cc @@ -39,7 +39,7 @@ #include "../../arith/interval_set.h" #include "../../runtime/thread_storage_scope.h" #include "../../s_tir/ir/ir_mutator_with_analyzer.h" -#include "../../tirx/transform/ir_utils.h" +#include "ir_utils.h" namespace tvm { namespace s_tir { @@ -467,7 +467,7 @@ class ExpressionHoister : public s_tir::IRMutatorWithAnalyzer { arith::Analyzer analyzer; auto hoister = ffi::make_object(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; } @@ -593,7 +593,7 @@ Pass HoistExpression() { return tvm::transform::Sequential( { insertion_pass, - tirx::transform::StmtSimplify(), + s_tir::transform::StmtSimplify(), tirx::transform::RemoveNoOp(), }, "s_tir.HoistExpression"); @@ -631,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"); @@ -649,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 f95dbd7b0c5a..7474c552ca76 100644 --- a/src/s_tir/transform/inject_double_buffer.cc +++ b/src/s_tir/transform/inject_double_buffer.cc @@ -31,7 +31,7 @@ #include #include -#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 Mutate_(const AttrStmtNode* op, InplaceMode inplace_mode) final { diff --git a/src/s_tir/transform/inject_software_pipeline.cc b/src/s_tir/transform/inject_software_pipeline.cc index 80fe163f3db5..0dd4c8a0a63e 100644 --- a/src/s_tir/transform/inject_software_pipeline.cc +++ b/src/s_tir/transform/inject_software_pipeline.cc @@ -36,8 +36,8 @@ #include #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 65ec0942ed77..b80cee1f49c0 100644 --- a/src/s_tir/transform/inject_virtual_thread.cc +++ b/src/s_tir/transform/inject_virtual_thread.cc @@ -34,7 +34,7 @@ #include #include "../../s_tir/ir/ir_mutator_with_analyzer.h" -#include "../../tirx/transform/ir_utils.h" +#include "ir_utils.h" namespace tvm { namespace s_tir { @@ -722,7 +722,7 @@ Pass InjectVirtualThread() { n->body = ffi::make_object(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 index 55ed50e9916f..f9da7ddab09d 100644 --- a/src/s_tir/transform/ir_utils.cc +++ b/src/s_tir/transform/ir_utils.cc @@ -21,6 +21,7 @@ #include #include +#include #include namespace tvm { @@ -28,6 +29,75 @@ 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(table); + SetDispatch(table); + } + + public: + UnchangedOr Mutate_(const SBlockNode* op, InplaceMode mode) { + SBlock block = ffi::GetRef(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(); + return iter; + }); + auto remap_region = [&](BufferRegion region) { + BufferVar buffer = GetRemappedBuffer(region->buffer); + if (!buffer.same_as(region->buffer)) region.CopyOnWrite()->buffer = buffer; + 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 Mutate_(const SBlockRealizeNode* op, InplaceMode mode) { + return s_tir::StmtExprMutator::MutateBlockRealize(this, op, mode); + } +}; +} // namespace + +Stmt ConvertSSA(Stmt stmt) { + return ffi::make_object() + ->Mutate(stmt, InplaceMode::kAllow) + .ValueOrUnchanged(stmt); +} + +IRModule ConvertSSA(IRModule mod) { + return ffi::make_object()->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 ConvertIndices(const MatchBufferRegion& match_buffer, const ffi::Array& indices) { const BufferVar& target = match_buffer->buffer; diff --git a/src/s_tir/transform/ir_utils.h b/src/s_tir/transform/ir_utils.h index eaf925ac6f21..e9d6c6cc2daa 100644 --- a/src/s_tir/transform/ir_utils.h +++ b/src/s_tir/transform/ir_utils.h @@ -30,6 +30,10 @@ 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 diff --git a/src/s_tir/transform/loop_partition.cc b/src/s_tir/transform/loop_partition.cc index ad32131734dd..392d821b71a8 100644 --- a/src/s_tir/transform/loop_partition.cc +++ b/src/s_tir/transform/loop_partition.cc @@ -43,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 { @@ -844,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/stmt_extension.cc b/src/s_tir/transform/stmt_extension.cc deleted file mode 100644 index 64a03e4bc09a..000000000000 --- a/src/s_tir/transform/stmt_extension.cc +++ /dev/null @@ -1,117 +0,0 @@ -/* - * 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/transform/stmt_extension.h" - -#include -#include - -namespace tvm { -namespace s_tir { -using namespace tirx; -namespace { - -UnchangedOr ConvertSSABlock(tirx::SSAStmtMutator* self, const SBlockNode* op, - InplaceMode mode) { - SBlock block = ffi::GetRef(op); - return self->WithScope([&]() -> Stmt { - auto iter_vars = op->iter_vars.Map([&](IterVar iter) { - Var var = self->DefineVar(iter->var); - if (!var.same_as(iter->var)) iter.CopyOnWrite()->var = var.as_or_throw(); - return iter; - }); - auto remap_region = [&](BufferRegion region) { - BufferVar buffer = self->RemapBuffer(region->buffer); - if (!buffer.same_as(region->buffer)) region.CopyOnWrite()->buffer = buffer; - 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 StmtExprMutator::MutateBlock(self, block.get(), - block.unique() ? mode : InplaceMode::kDisallow) - .ValueOrUnchanged(block); - }); -} - -UnchangedOr FlattenBlock(tirx::FlattenStmtMutator* self, const SBlockNode* op, - InplaceMode mode) { - TVM_FFI_ICHECK_EQ(op->match_buffers.size(), 0) - << "Unexpected MatchBufferRegion found during tirx.transform.FlattenBuffer. " - << "All MatchBufferRegion should be removed in tirx.transform.LowerMatchBuffer."; - SBlock block = ffi::GetRef(op); - auto alloc_buffers = op->alloc_buffers; - alloc_buffers.MutateByApply([&](BufferVar buffer) { return self->DefineBuffer(buffer); }); - if (!alloc_buffers.same_as(op->alloc_buffers)) block.CopyOnWrite()->alloc_buffers = alloc_buffers; - auto reads = op->reads; - reads.MutateByApply([&](BufferRegion region) { return self->RewriteRegion(region); }); - if (!reads.same_as(op->reads)) block.CopyOnWrite()->reads = reads; - auto writes = op->writes; - writes.MutateByApply([&](BufferRegion region) { return self->RewriteRegion(region); }); - if (!writes.same_as(op->writes)) block.CopyOnWrite()->writes = writes; - return StmtExprMutator::MutateBlock(self, block.get(), - block.unique() ? mode : InplaceMode::kDisallow) - .ValueOrUnchanged(block); -} - -// These hooks extend each pass's existing native table before it is finalized. -// All ordinary statements and expression remapping remain in the TIRX pass. -TVM_FFI_STATIC_INIT_BLOCK() { - tirx::SSAStmtMutator::RegisterExtension([](tirx::SSAStmtMutator::VTable* table) { - table->ClearDispatch(); - table->SetDispatch( - [](const ffi::Object* node, ObjectMutator* self, InplaceMode mode) { - return ffi::details::UnchangedOrUnsafe::MoveFromTVMFFIAny( - ffi::details::UnchangedOrUnsafe::MoveToTVMFFIAny( - ConvertSSABlock(static_cast(self), - static_cast(node), mode))); - }); - }); - tirx::FlattenStmtMutator::RegisterExtension([](tirx::FlattenStmtMutator::VTable* table) { - table->ClearDispatch(); - table->SetDispatch( - [](const ffi::Object* node, ObjectMutator* self, InplaceMode mode) { - return ffi::details::UnchangedOrUnsafe::MoveFromTVMFFIAny( - ffi::details::UnchangedOrUnsafe::MoveToTVMFFIAny( - FlattenBlock(static_cast(self), - static_cast(node), mode))); - }); - }); - tirx::IndexDomainVisitor::RegisterExtension([](tirx::IndexDomainVisitor::VTable* table) { - table->ClearDispatch(); - table->SetDispatch([](const ffi::Object* node, ObjectVisitor* base) { - auto* self = static_cast(base); - auto* block = static_cast(node); - for (const auto& iter : block->iter_vars) { - self->BindDomain(iter->var, Range::FromMinExtent(iter->dom->min, iter->dom->extent)); - } - return StmtExprVisitor::VisitBlock(self, block); - }); - }); -} - -} // namespace -} // namespace s_tir -} // namespace tvm diff --git a/src/s_tir/transform/stmt_simplify.cc b/src/s_tir/transform/stmt_simplify.cc new file mode 100644 index 000000000000..458abbd634a0 --- /dev/null +++ b/src/s_tir/transform/stmt_simplify.cc @@ -0,0 +1,93 @@ +/* + * 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 +#include +#include + +#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 : public arith::StmtSimplifier { + public: + using Parent = arith::StmtSimplifier; + StmtSimplifier(const arith::Analyzer& analyzer, arith::StmtSimplifyConfig config) + : Parent(GlobalVTable(), analyzer, config) {} + using Parent::Mutate_; + using Parent::Run; + + public: + UnchangedOr Mutate_(const SBlockNode* op, InplaceMode inplace_mode) { + return constraint_scope_.WithNewScope([&]() -> UnchangedOr { + 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 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(vtable); + SetDispatch(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(); + return ffi::make_object(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.StmtSimplify") + .value_or(tvm::transform::PassConfigWithDefaults()); + return ffi::make_object(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/ir/specialize.h b/src/s_tir/transform/stmt_simplify.h similarity index 57% rename from src/tirx/ir/specialize.h rename to src/s_tir/transform/stmt_simplify.h index 4e104f61c3ee..8dcff5bcc6d3 100644 --- a/src/tirx/ir/specialize.h +++ b/src/s_tir/transform/stmt_simplify.h @@ -17,23 +17,13 @@ * under the License. */ -#ifndef TVM_TIRX_IR_SPECIALIZE_H_ -#define TVM_TIRX_IR_SPECIALIZE_H_ - -#include - +#ifndef TVM_S_TIR_TRANSFORM_STMT_SIMPLIFY_H_ +#define TVM_S_TIR_TRANSFORM_STMT_SIMPLIFY_H_ +#include +#include namespace tvm { -namespace tirx { - -using SpecializeVisitorVTable = - ObjectFunctor(const ffi::Object*, ObjectVisitor*)>; - -// Register dialect-specific buffer planning before any specialization is run. -// Initializers extend the inherited native traversal table without exposing the -// specializer's private buffer remapping and declaration state. -void RegisterSpecializeBufferPlannerExtension(void (*init)(SpecializeVisitorVTable*)); - -} // namespace tirx +namespace s_tir { +tirx::PrimFunc StmtSimplify(tirx::PrimFunc func, const arith::Analyzer& analyzer); +} // namespace s_tir } // namespace tvm - -#endif // TVM_TIRX_IR_SPECIALIZE_H_ +#endif // TVM_S_TIR_TRANSFORM_STMT_SIMPLIFY_H_ diff --git a/src/te/operation/create_primfunc.cc b/src/te/operation/create_primfunc.cc index bd28eddf5659..ca0b2045d65f 100644 --- a/src/te/operation/create_primfunc.cc +++ b/src/te/operation/create_primfunc.cc @@ -40,7 +40,7 @@ #include #include -#include "../../tirx/ir/data_type_rewriter.h" +#include "../../s_tir/ir/data_type_rewriter.h" #include "graph.h" namespace tvm { @@ -874,7 +874,7 @@ PrimFunc CreatePrimFunc(const ffi::Array& 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(index_dtype_override.value()) + func = ffi::make_object(index_dtype_override.value()) ->Rewrite(std::move(func)); } auto result = ffi::make_object()->Process(std::move(func)); @@ -950,7 +950,7 @@ PrimFunc CreatePrimFunc(const ffi::Array& arg_list, } auto func = GenerateAndCompletePrimFunc(arg_list, root_stmts, &info); if (index_dtype_override.has_value()) { - func = ffi::make_object(index_dtype_override.value()) + func = ffi::make_object(index_dtype_override.value()) ->Rewrite(std::move(func)); } auto result = ffi::make_object()->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 d1c9c0d58573..42437656e849 100644 --- a/src/tirx/analysis/verify_tirx_well_formed.cc +++ b/src/tirx/analysis/verify_tirx_well_formed.cc @@ -50,10 +50,9 @@ class ExecScopeVerifier : public Verifier { private: using Verifier::Visit; - bool EnterExtensionStmt(const ffi::Object* op, ffi::reflection::AccessPath path) override { + void VisitStmtDefault_(const ffi::Object* op, ffi::reflection::AccessPath path) override { Verify(false) << "TIRxError: " << op->GetTypeKey() << " is not allowed in tirx=True mode at " << path; - return false; } void Dispatch_(const tirx::TilePrimitiveCallNode* op, ffi::reflection::AccessPath path) override { @@ -130,10 +129,9 @@ class LayoutVerifier : public Verifier { private: using Verifier::Visit; - bool EnterExtensionStmt(const ffi::Object* op, ffi::reflection::AccessPath path) override { + void VisitStmtDefault_(const ffi::Object* op, ffi::reflection::AccessPath path) override { Verify(false) << "TIRxError: " << op->GetTypeKey() << " is not allowed in tirx=True mode at " << path; - return false; } }; @@ -144,10 +142,9 @@ class AsyncStructsVerifier : public Verifier { private: using Verifier::Visit; - bool EnterExtensionStmt(const ffi::Object* op, ffi::reflection::AccessPath path) override { + void VisitStmtDefault_(const ffi::Object* op, ffi::reflection::AccessPath path) override { Verify(false) << "TIRxError: " << op->GetTypeKey() << " is not allowed in tirx=True mode at " << path; - return false; } }; @@ -158,10 +155,9 @@ class DeviceFuncVerifier : public Verifier { private: using Verifier::Visit; - bool EnterExtensionStmt(const ffi::Object* op, ffi::reflection::AccessPath path) override { + void VisitStmtDefault_(const ffi::Object* op, ffi::reflection::AccessPath path) override { Verify(false) << "TIRxError: " << op->GetTypeKey() << " is not allowed in tirx=True mode at " << path; - return false; } }; @@ -189,11 +185,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()) { - // 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 ea49da4c3984..adef48e1c111 100644 --- a/src/tirx/analysis/verify_well_formed.cc +++ b/src/tirx/analysis/verify_well_formed.cc @@ -43,248 +43,6 @@ namespace tirx { using AccessPath = ffi::reflection::AccessPath; -namespace { -std::vector& WellFormedExtensions() { - static std::vector extensions; - return extensions; -} -} // namespace -void RegisterWellFormedExtension(bool (*verify)(const PrimFunc&, bool)) { - WellFormedExtensions().push_back(verify); -} - -class UndefinedVarVerifier : public Verifier { - 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(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 currently_defined_; - - // Variables that were previously defined, and are now out of scope. - std::unordered_map 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 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 { - 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 - currently_defined_; - // Buffers that were previously defined and are now out of scope. - std::unordered_map - previously_defined_; -}; - -/*! \brief Verify the asserted type of each tirx buffer load. */ -class TensorLoadTypeVerifier : public Verifier { - public: - using Verifier::Verifier; - - private: - void Dispatch_(const TensorLoadNode* op, AccessPath path) override { - auto buffer = op->source.as(); - 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 index_ty = op->indices.empty() - ? ffi::Optional(buffer.value()->dtype) - : op->indices.back().ty().as(); - 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 asserted_ty = op->ty.as(); - ffi::Optional expected_ty = expected->ty.as(); - 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 @@ -327,37 +85,11 @@ class SingleEnvThreadVerifier : public Verifier { }; bool VerifyWellFormed(const PrimFunc& func, bool assert_mode) { - for (auto verify : WellFormedExtensions()) { - if (!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(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()) { - 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(mod, assert_mode); } TVM_FFI_STATIC_INIT_BLOCK() { diff --git a/src/tirx/analysis/verify_well_formed.h b/src/tirx/analysis/verify_well_formed.h index 0c19522d2621..16b57274ff90 100644 --- a/src/tirx/analysis/verify_well_formed.h +++ b/src/tirx/analysis/verify_well_formed.h @@ -19,11 +19,264 @@ #ifndef TVM_TIRX_ANALYSIS_VERIFY_WELL_FORMED_H_ #define TVM_TIRX_ANALYSIS_VERIFY_WELL_FORMED_H_ -#include + +#include + +#include "../ir/tir_visitor_with_path.h" namespace tvm { namespace tirx { -// Register supplementary dialect checks before verification is first invoked. -void RegisterWellFormedExtension(bool (*verify)(const PrimFunc&, bool)); +using AccessPath = ffi::reflection::AccessPath; + +template +class UndefinedVarVerifier : public Verifier, PathVisitor> { + using Verifier = tirx::Verifier, PathVisitor>; + + 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; + 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(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 currently_defined_; + + // Variables that were previously defined, and are now out of scope. + std::unordered_map 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 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 + * - Dialect-specific definitions exposed by PathVisitor + * + * 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. + */ +template +class UndefinedBufferVerifier : public Verifier, PathVisitor> { + using Verifier = tirx::Verifier, PathVisitor>; + + public: + using Verifier::Verifier; + using Verifier::Verify; + + 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 + currently_defined_; + // Buffers that were previously defined and are now out of scope. + std::unordered_map + previously_defined_; +}; + +/*! \brief Verify the asserted type of each tirx buffer load. */ +template +class TensorLoadTypeVerifier : public Verifier, PathVisitor> { + using Verifier = tirx::Verifier, 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(); + 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 index_ty = op->indices.empty() + ? ffi::Optional(buffer.value()->dtype) + : op->indices.back().ty().as(); + 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 asserted_ty = op->ty.as(); + ffi::Optional expected_ty = expected->ty.as(); + 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 << "."; + PathVisitor::Dispatch_(op, path); + } +}; + +template +bool VerifyWellFormedCommon(const NodeRef& node, bool assert_mode) { + return UndefinedVarVerifier::Verify(node, assert_mode) && + UndefinedBufferVerifier::Verify(node, assert_mode) && + TensorLoadTypeVerifier::Verify(node, assert_mode); +} } // namespace tirx } // namespace tvm -#endif +#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 3f6ed8956315..eaa54a388bb5 100644 --- a/src/tirx/ir/data_type_rewriter.cc +++ b/src/tirx/ir/data_type_rewriter.cc @@ -33,7 +33,6 @@ #include #include #include -#include #include "tvm/ir/expr.h" #include "tvm/ir/prim/expr.h" @@ -43,34 +42,6 @@ namespace tvm { namespace tirx { using namespace tvm::prim; -namespace { -std::vector& IndexDataTypeRewriterExtensions() { - static std::vector extensions; - return extensions; -} -} // namespace -void IndexDataTypeRewriter::RegisterExtension(void (*init)(VTable*)) { - IndexDataTypeRewriterExtensions().push_back(init); -} -void IndexDataTypeRewriter::InitVTable(VTable* vtable) { - DataTypeLegalizer::InitVTable(vtable); - for (auto init : IndexDataTypeRewriterExtensions()) init(vtable); -} - -namespace { -std::vector& DataTypeLegalizerExtensions() { - static std::vector extensions; - return extensions; -} -} // namespace -void DataTypeLegalizer::RegisterExtension(void (*init)(VTable*)) { - DataTypeLegalizerExtensions().push_back(init); -} -void DataTypeLegalizer::InitVTable(VTable* vtable) { - StmtExprMutator::InitVTable(vtable); - for (auto init : DataTypeLegalizerExtensions()) init(vtable); -} - UnchangedOr DataTypeLegalizer::Mutate_(const ForNode* op, InplaceMode inplace_mode) { auto result = StmtExprMutator::Mutate_(op, inplace_mode); if (!result.IsUnchanged()) { @@ -579,6 +550,9 @@ UnchangedOr 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 584539e02910..7d1e09d85b27 100644 --- a/src/tirx/ir/data_type_rewriter.h +++ b/src/tirx/ir/data_type_rewriter.h @@ -47,18 +47,12 @@ namespace tirx { */ class DataTypeLegalizer : public StmtExprMutator { public: - using StmtExprMutator::VTable; - // Dialect-owned handlers use nested access to the active traversal context. - class Extension; - // Register during library initialization, before the finalized table is first used. - static void RegisterExtension(void (*init)(VTable*)); TVM_DEFINE_OBJECT_FUNCTOR_DEFAULT_CONSTRUCTOR(DataTypeLegalizer, StmtExprMutator) using StmtExprMutator::Mutate; using StmtExprMutator::Mutate_; protected: explicit DataTypeLegalizer(const VTable* vtable) : StmtExprMutator(vtable) {} - static void InitVTable(VTable* vtable); UnchangedOr Mutate_(const ForNode* op, InplaceMode inplace_mode) override; UnchangedOr Mutate_(const AttrStmtNode* op, InplaceMode inplace_mode) override; UnchangedOr Mutate_(const BindNode* op, InplaceMode inplace_mode) override; @@ -103,17 +97,12 @@ class DataTypeLegalizer : public StmtExprMutator { */ class IndexDataTypeRewriter : public DataTypeLegalizer { public: - using DataTypeLegalizer::VTable; - // Dialect-owned handlers use nested access to the active traversal context. - class Extension; - // Register during library initialization, before the finalized table is first used. - static void RegisterExtension(void (*init)(VTable*)); TVM_DEFINE_OBJECT_FUNCTOR_DEFAULT_CONSTRUCTOR(IndexDataTypeRewriter, DataTypeLegalizer) using DataTypeLegalizer::Mutate; using DataTypeLegalizer::Mutate_; protected: - static void InitVTable(VTable* vtable); + explicit IndexDataTypeRewriter(const VTable* vtable) : DataTypeLegalizer(vtable) {} using Parent = DataTypeLegalizer; UnchangedOr Mutate(ffi::AnyView value, InplaceMode inplace_mode) override; UnchangedOr Mutate_(const BufferStoreNode* op, InplaceMode inplace_mode) override; @@ -133,9 +122,6 @@ class IndexDataTypeRewriter : public DataTypeLegalizer { UnchangedOr Mutate_(const ForNode* op, InplaceMode inplace_mode) override; - // Dialect allocation hooks retain common buffer validation in the active rewriter. - virtual void ValidateAllocation(const BufferVar& buffer) {} - // indicator of index expr to rewrite bool is_enabled_{false}; // indicator of condition @@ -157,6 +143,7 @@ class IndexDataTypeNormalizer : public IndexDataTypeRewriter { PrimFunc Rewrite(PrimFunc func); protected: + IndexDataTypeNormalizer(PrimType target_data_type, const VTable* vtable); using Parent = IndexDataTypeRewriter; UnchangedOr Mutate_(const IntImmNode* op, InplaceMode inplace_mode) override; diff --git a/src/tirx/ir/ir_mutator_with_analyzer.cc b/src/tirx/ir/ir_mutator_with_analyzer.cc index d11c6ee224df..ac5d6f6381d8 100644 --- a/src/tirx/ir/ir_mutator_with_analyzer.cc +++ b/src/tirx/ir/ir_mutator_with_analyzer.cc @@ -33,21 +33,7 @@ namespace tvm { namespace tirx { -namespace { -std::vector& IRMutatorWithAnalyzerExtensions() { - static std::vector extensions; - return extensions; -} -} // namespace - -void IRMutatorWithAnalyzer::RegisterExtension(void (*init)(VTable*)) { - IRMutatorWithAnalyzerExtensions().push_back(init); -} - -void IRMutatorWithAnalyzer::InitVTable(VTable* vtable) { - StmtExprMutator::InitVTable(vtable); - for (auto init : IRMutatorWithAnalyzerExtensions()) init(vtable); -} +void IRMutatorWithAnalyzer::InitVTable(VTable* vtable) { StmtExprMutator::InitVTable(vtable); } const IRMutatorWithAnalyzer::VTable* IRMutatorWithAnalyzer::GlobalVTable() { static const VTable table = [] { diff --git a/src/tirx/ir/ir_mutator_with_analyzer.h b/src/tirx/ir/ir_mutator_with_analyzer.h index 9cb079f4a3b9..bcc976c57ac5 100644 --- a/src/tirx/ir/ir_mutator_with_analyzer.h +++ b/src/tirx/ir/ir_mutator_with_analyzer.h @@ -34,7 +34,6 @@ #include #include -#include namespace tvm { namespace tirx { @@ -50,11 +49,6 @@ namespace tirx { */ class IRMutatorWithAnalyzer : public StmtExprMutator { public: - using StmtExprMutator::VTable; - // Dialect-owned handlers use nested access to the active traversal context. - class Extension; - // Extensions register during library initialization, before constructing a visitor. - static void RegisterExtension(void (*init)(VTable*)); using StmtExprMutator::Mutate; using StmtExprMutator::Mutate_; explicit IRMutatorWithAnalyzer(const arith::Analyzer& analyzer) diff --git a/src/tirx/ir/ir_visitor_with_analyzer.cc b/src/tirx/ir/ir_visitor_with_analyzer.cc index c944a02ba5f0..3324e6d2194f 100644 --- a/src/tirx/ir/ir_visitor_with_analyzer.cc +++ b/src/tirx/ir/ir_visitor_with_analyzer.cc @@ -30,21 +30,7 @@ namespace tvm { namespace tirx { -namespace { -std::vector& IRVisitorWithAnalyzerExtensions() { - static std::vector extensions; - return extensions; -} -} // namespace - -void IRVisitorWithAnalyzer::RegisterExtension(void (*init)(VTable*)) { - IRVisitorWithAnalyzerExtensions().push_back(init); -} - -void IRVisitorWithAnalyzer::InitVTable(VTable* vtable) { - StmtExprVisitor::InitVTable(vtable); - for (auto init : IRVisitorWithAnalyzerExtensions()) init(vtable); -} +void IRVisitorWithAnalyzer::InitVTable(VTable* vtable) { StmtExprVisitor::InitVTable(vtable); } ffi::Optional IRVisitorWithAnalyzer::Visit_(const ForNode* op) { return constraint_scope_.WithNewScope([&]() -> ffi::Optional { diff --git a/src/tirx/ir/ir_visitor_with_analyzer.h b/src/tirx/ir/ir_visitor_with_analyzer.h index 619d388d4d7f..63db9d9fc0cf 100644 --- a/src/tirx/ir/ir_visitor_with_analyzer.h +++ b/src/tirx/ir/ir_visitor_with_analyzer.h @@ -31,18 +31,11 @@ #include #include -#include - namespace tvm { namespace tirx { class IRVisitorWithAnalyzer : public StmtExprVisitor { public: - using StmtExprVisitor::VTable; - // Dialect-owned handlers use nested access to the active traversal context. - class Extension; - // Extensions register during library initialization, before constructing a visitor. - static void RegisterExtension(void (*init)(VTable*)); TVM_DEFINE_OBJECT_FUNCTOR_DEFAULT_CONSTRUCTOR(IRVisitorWithAnalyzer, StmtExprVisitor) PrimExpr Simplify(const PrimExpr& expr) { return analyzer_->Simplify(expr); } diff --git a/src/tirx/ir/specialize.cc b/src/tirx/ir/specialize.cc index 9f5d6ee61718..55195e3dc6e9 100644 --- a/src/tirx/ir/specialize.cc +++ b/src/tirx/ir/specialize.cc @@ -21,8 +21,6 @@ * \file src/tirx/ir/specialize.cc * \brief Specialize parameters of PrimFunc. */ -#include "specialize.h" - #include #include #include @@ -36,7 +34,6 @@ #include #include -#include #include "../transform/ir_utils.h" @@ -45,17 +42,6 @@ namespace tirx { using VarMap = std::unordered_map; -namespace { -std::vector& BufferPlannerExtensions() { - static std::vector extensions; - return extensions; -} -} // namespace - -void RegisterSpecializeBufferPlannerExtension(void (*init)(SpecializeVisitorVTable*)) { - BufferPlannerExtensions().push_back(init); -} - /**************** Helper functions ****************/ /*! \brief Helper function to check whether the given var is in function parameter list. */ @@ -168,25 +154,18 @@ class PrimFuncSpecializer : public StmtExprMutator { public: using StmtExprVisitor::Visit_; - explicit BufferPlanner(PrimFuncSpecializer* specializer) - : StmtExprVisitor(GlobalVTable()), specializer_(specializer) {} + explicit BufferPlanner(PrimFuncSpecializer* specializer) : specializer_(specializer) {} private: - static const VTable* GlobalVTable() { - static const VTable table = [] { - VTable table; - StmtExprVisitor::InitVTable(&table); - for (auto init : BufferPlannerExtensions()) init(&table); - table.Finalize(); - return table; - }(); - return &table; - } - ffi::Optional Visit_(const VarNode* op) final { if (op->ty.as()) { if (def_region_kind() == kTVMFFIDefRegionKindSimple) { - specializer_->MutateAllocBuffer(GetBufferVar(op)); + const BufferVar buffer = GetBufferVar(op); + specializer_->MutateAllocBuffer(buffer); + // Structural extension nodes expose buffer definitions without a native + // statement hook. Plan their metadata as uses after defining the buffer. + return this->WithDefRegionKind(kTVMFFIDefRegionKindNone, + [&]() { return VisitBufferMetadata(buffer); }); } else { specializer_->ValidateBufferUse(GetBufferVar(op)); } @@ -194,11 +173,15 @@ class PrimFuncSpecializer : public StmtExprMutator { return StmtExprVisitor::Visit_(op); } + ffi::Optional Visit_(const AllocBufferNode* op) final { + return this->WithDefRegionKind(kTVMFFIDefRegionKindSimple, + [&]() { return this->Visit(op->buffer); }); + } + ffi::Optional Visit_(const DeclBufferNode* op) final { // The declaration establishes the buffer before visiting its data expression. TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->WithDefRegionKind( kTVMFFIDefRegionKindSimple, [&]() { return this->Visit(op->buffer); })); - TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(VisitBufferMetadata(op->buffer)); return Visit(op->data); } diff --git a/src/tirx/ir/stmt_functor.cc b/src/tirx/ir/stmt_functor.cc index 5e19a6348c42..c72d63b50d57 100644 --- a/src/tirx/ir/stmt_functor.cc +++ b/src/tirx/ir/stmt_functor.cc @@ -40,17 +40,6 @@ namespace tvm { namespace tirx { -namespace { -std::vector& StmtExprVisitorExtensions() { - static std::vector extensions; - return extensions; -} -} // namespace - -void StmtExprVisitor::RegisterExtension(void (*init)(VTable*)) { - StmtExprVisitorExtensions().push_back(init); -} - void StmtExprVisitor::InitVTable(VTable* vtable) { tvm::ExprVisitor::InitVTable(vtable); SetDispatch(vtable); @@ -70,7 +59,6 @@ void StmtExprVisitor::InitVTable(VTable* vtable) { SetDispatch(vtable); SetDispatch(vtable); SetDispatch(vtable); - for (auto init : StmtExprVisitorExtensions()) init(vtable); } ffi::Optional StmtExprVisitor::Visit_(const VarNode* op) { return std::nullopt; } @@ -284,17 +272,6 @@ ffi::Optional StmtExprVisitor::Visit_(const TilePrimitiveCallNod return std::nullopt; } -namespace { -std::vector& StmtExprMutatorExtensions() { - static std::vector extensions; - return extensions; -} -} // namespace - -void StmtExprMutator::RegisterExtension(void (*init)(VTable*)) { - StmtExprMutatorExtensions().push_back(init); -} - void StmtExprMutator::InitVTable(VTable* vtable) { tvm::ExprMutator::InitVTable(vtable); SetDispatch(vtable); @@ -314,7 +291,6 @@ void StmtExprMutator::InitVTable(VTable* vtable) { SetDispatch(vtable); SetDispatch(vtable); SetDispatch(vtable); - for (auto init : StmtExprMutatorExtensions()) init(vtable); } UnchangedOr StmtExprMutator::Mutate_(const BindNode* op, InplaceMode inplace_mode) { diff --git a/src/tirx/ir/tir_visitor_with_path.cc b/src/tirx/ir/tir_visitor_with_path.cc index d7061f93cc52..9416ea44ac97 100644 --- a/src/tirx/ir/tir_visitor_with_path.cc +++ b/src/tirx/ir/tir_visitor_with_path.cc @@ -33,27 +33,6 @@ namespace tvm { namespace tirx { -namespace { -std::vector& PathVisitorExtensions() { - static std::vector extensions; - return extensions; -} -} // namespace -void TIRVisitorWithPath::RegisterExtension(void (*init)(VTable*)) { - PathVisitorExtensions().push_back(init); -} -TIRVisitorWithPath::TIRVisitorWithPath() - : StmtVisitor([] { - static const VTable table = [] { - VTable table; - StmtVisitor::InitVTable(&table); - for (auto init : PathVisitorExtensions()) init(&table); - table.Finalize(); - return table; - }(); - return &table; - }()) {} - using AccessPath = ffi::reflection::AccessPath; void TIRVisitorWithPath::Visit(const IRModule& mod, AccessPath path) { diff --git a/src/tirx/ir/tir_visitor_with_path.h b/src/tirx/ir/tir_visitor_with_path.h index aea29aac0c9f..9f5597128a01 100644 --- a/src/tirx/ir/tir_visitor_with_path.h +++ b/src/tirx/ir/tir_visitor_with_path.h @@ -45,23 +45,17 @@ namespace tirx { class TIRVisitorWithPath : protected ExprFunctor, protected StmtFunctor { public: - using StmtVisitor = StmtFunctor; - using VTable = StmtVisitor::VTable; - // Dialect-owned handlers use nested access to the active traversal context. - class Extension; - TIRVisitorWithPath(); - // Extensions register at library initialization, before first visitor construction. - static void RegisterExtension(void (*init)(VTable*)); + TIRVisitorWithPath() = default; template void operator()(TObjectRef&& obj) { Visit(std::forward(obj), ffi::reflection::AccessPath::Root()); } protected: - // A dialect restriction can reject extension statements without knowing concrete node types. - virtual bool EnterExtensionStmt(const ffi::Object* op, ffi::reflection::AccessPath path) { - return true; - } + using StmtVisitor = StmtFunctor; + using VTable = StmtVisitor::VTable; + explicit TIRVisitorWithPath(const VTable* vtable) : StmtVisitor(vtable) {} + static void InitVTable(VTable* vtable) { StmtVisitor::InitVTable(vtable); } // Delegate to ExprFunctor::Dispatch for PrimExpr, and any subclasses virtual inline void Visit(const PrimExpr& obj, ffi::reflection::AccessPath path) { Dispatch(obj, path); @@ -314,8 +308,8 @@ class TIRVisitorWithPath : protected ExprFunctor -class Verifier : protected TIRVisitorWithPath { +template +class Verifier : protected PathVisitor { public: template static bool Verify(const TirNodeRef& node, bool assert_on_error) { @@ -326,6 +320,10 @@ class Verifier : protected TIRVisitorWithPath { protected: explicit Verifier(bool assert_on_error) : assert_on_error_(assert_on_error) {} + void VisitStmtDefault_(const ffi::Object* op, ffi::reflection::AccessPath path) override { + Verify(false) << "TIR verifier does not support statement " << op->GetTypeKey() << " at " + << path; + } /* \brief Helper class to handle the bool-or-assert handles * diff --git a/src/tirx/transform/flatten_buffer.cc b/src/tirx/transform/flatten_buffer.cc index f41cc14fe0ba..461431b5abc4 100644 --- a/src/tirx/transform/flatten_buffer.cc +++ b/src/tirx/transform/flatten_buffer.cc @@ -34,7 +34,6 @@ #include "../ir/ir_mutator_with_analyzer.h" #include "ir_utils.h" -#include "stmt_extension.h" namespace tvm { namespace tirx { @@ -59,7 +58,7 @@ using namespace tvm::prim; * Every use site then only looks the pair up; a use before its definition is * a hard error instead of a silently stale reference. */ -class BufferFlattener : public FlattenStmtMutator { +class BufferFlattener : public IRMutatorWithAnalyzer { public: using IRMutatorWithAnalyzer::Mutate; using IRMutatorWithAnalyzer::Mutate_; @@ -99,7 +98,7 @@ class BufferFlattener : public FlattenStmtMutator { } public: - explicit BufferFlattener(const arith::Analyzer& ana) : FlattenStmtMutator(ana) {} + explicit BufferFlattener(const arith::Analyzer& ana) : IRMutatorWithAnalyzer(ana) {} private: struct FlatInfo { @@ -179,8 +178,6 @@ class BufferFlattener : public FlattenStmtMutator { return it->second; } - BufferVar DefineBuffer(BufferVar buffer) final { return Define(buffer).flattened; } - UnchangedOr Mutate_(const AllocBufferNode* op, InplaceMode inplace_mode) final { const FlatInfo& info = Define(op->buffer); if (info.flattened.same_as(op->buffer)) { @@ -295,31 +292,6 @@ class BufferFlattener : public FlattenStmtMutator { return BufferLoad(info.flattened, FoldIndices(info, node->indices), node->span); } - BufferRegion RewriteRegion(BufferRegion region) final { - const FlatInfo& info = Lookup(region->buffer); - if (info.flattened.same_as(region->buffer)) { - return region; - } - - ffi::Array min_values; - ffi::Array max_values; - for (const auto& range : region->region) { - min_values.push_back(range->min); - max_values.push_back(range->min + range->extent - 1); - } - - ffi::Array flattened_min = FoldIndices(info, min_values); - ffi::Array flattened_max = FoldIndices(info, max_values); - - ffi::Array flattened_ranges; - TVM_FFI_ICHECK_EQ(flattened_min.size(), flattened_max.size()); - for (size_t i = 0; i < flattened_min.size(); i++) { - flattened_ranges.push_back(Range(flattened_min[i], flattened_max[i] + 1)); - } - - return BufferRegion(info.flattened, flattened_ranges); - } - /*! \brief Set of buffers accessed during visitation (used to emit DeclBuffer for param buffers). */ std::unordered_set buffers_used_; diff --git a/src/tirx/transform/force_narrow_index_to_i32.cc b/src/tirx/transform/force_narrow_index_to_i32.cc index 956ba262baf8..5692dd7872e1 100644 --- a/src/tirx/transform/force_narrow_index_to_i32.cc +++ b/src/tirx/transform/force_narrow_index_to_i32.cc @@ -69,12 +69,22 @@ class Int32DTypeNarrower : public IndexDataTypeNormalizer { return ffi::Unchanged(); } - void ValidateAllocation(const BufferVar& buf) final { - if (buf->dtype.MatchesCode(DLDataTypeCode::kDLInt) && buf->dtype.bits() > 32) { + UnchangedOr Mutate_(const AllocBufferNode* op, InplaceMode inplace_mode) final { + auto result = IndexDataTypeNormalizer::Mutate_(op, inplace_mode); + auto alloc = + std::move(result).ValueOrUnchanged(ffi::GetRef(op)).as_or_throw(); + const BufferVar& buf = alloc->buffer; + // Scalar assignments in TVMScript use local scalar storage. Keep its explicit + // dtype (e.g. an int64 opaque call result) and cast at narrowed index uses. + // IsScalar checks the scalar layout contract, not merely the allocation size. + bool is_local_scalar = buf.scope() == "local" && buf.IsScalar(); + if (!is_local_scalar && buf->dtype.MatchesCode(DLDataTypeCode::kDLInt) && + buf->dtype.bits() > 32) { TVM_FFI_THROW(InternalError) << "The buffer " << buf << " allocated in the function has dtype " << buf->dtype << ". The function is " << func_; } + return alloc; } PrimFunc func_; diff --git a/src/tirx/transform/inline_private_functions.cc b/src/tirx/transform/inline_private_functions.cc index af52564a71f5..6e7278ae7ad8 100644 --- a/src/tirx/transform/inline_private_functions.cc +++ b/src/tirx/transform/inline_private_functions.cc @@ -121,13 +121,18 @@ bool IsInlinablePrimFunc(const GlobalVar& gvar, const PrimFunc& prim_func, if (param->ty.as()) return false; } - // We do not currently support inlining of schedulable TIR - // functions. To support this use case, repeated names in - // schedulable block nodes resulting from multiple calls to the same - // inlined function will need to be de-duplicated. - static ffi::reflection::TypeAttrColumn prevent_inline("tirx.prevent_inline"); - ffi::AnyView value = prevent_inline[prim_func->body->type_index()]; - if (value != nullptr && value.cast()) return false; + // Extension statement roots may introduce binder or naming rules that this + // pass cannot preserve. Only inline roots supported by native TIRX traversal. + struct NativeStmtTable : StmtExprVisitor { + static VTable Make() { + VTable table; + InitVTable(&table); + table.Finalize(); + return table; + } + }; + static const auto native_stmts = NativeStmtTable::Make(); + if (!native_stmts.CanDispatch(prim_func->body.get())) return false; return true; } @@ -304,7 +309,6 @@ Pass InlinePrivateFunctions() { } TVM_FFI_STATIC_INIT_BLOCK() { - ffi::reflection::EnsureTypeAttrColumn("tirx.prevent_inline"); namespace refl = tvm::ffi::reflection; refl::GlobalDef().def("tirx.transform.InlinePrivateFunctions", InlinePrivateFunctions); } diff --git a/src/tirx/transform/ir_utils.cc b/src/tirx/transform/ir_utils.cc index b8601952990b..1d4978167ffe 100644 --- a/src/tirx/transform/ir_utils.cc +++ b/src/tirx/transform/ir_utils.cc @@ -37,8 +37,6 @@ #include #include -#include "stmt_extension.h" - namespace tvm { namespace tirx { using namespace tvm::prim; @@ -89,651 +87,553 @@ Stmt MergeNest(const std::vector>& nest, Stmt body) { return body; } -class IRConvertSSA final : public SSAStmtMutator { - public: - using StmtExprMutator::Mutate; - using StmtExprMutator::Mutate_; - PrimFunc VisitPrimFunc(PrimFunc func) { - // Remap parameters, if they were used in another function. - // Function-scope remaps use function_scope_var_remap_ (not the scope stack), - // because they persist across the entire function body. - auto params = func->params.Map([&](const tirx::Var& var) -> tirx::Var { - if (defined_.count(var.get())) { - Var new_var = MakeNewVar(var); - PushVarRemap(var, new_var); - return new_var; - } else { - defined_.insert(var.get()); - return var; - } - }); - - // Remap implicitly defined buffer parameters - { - std::unordered_set defined_params; - for (const auto& var : func->params) { - defined_params.insert(var.get()); - } - std::unordered_set defined_match_vars; - for (const Var& param : func->params) { - auto buffer = param.as(); - if (!buffer) continue; - auto check_var = [&](const Var& var) { - const VarNode* var_ptr = var.get(); - if (defined_params.count(var_ptr)) return; - if (!defined_match_vars.insert(var_ptr).second) return; - - // Buffer-parameter shape vars use "match" semantics: first occurrence - // defines the var, subsequent occurrences (in other buffers) are - // just consistent uses of the same var -- not redefinitions. - if (defined_.count(var_ptr)) { - Var new_var = MakeNewVar(var); - PushVarRemap(var, new_var); - } else { - defined_.insert(var_ptr); - } - }; - auto walk_fn = [&](const Var& var) -> ffi::Expected { - check_var(var); - return ffi::WalkResult::Advance(); - }; - for (const auto& dim : buffer.value()->shape) { - ffi::StructuralWalk(dim, walk_fn); - } - for (const auto& stride : buffer.value()->strides) { - if (auto var = stride.as()) check_var(var.value()); +PrimFunc IRConvertSSA::VisitPrimFunc(PrimFunc func) { + // Remap parameters, if they were used in another function. + // Function-scope remaps use function_scope_var_remap_ (not the scope stack), + // because they persist across the entire function body. + auto params = func->params.Map([&](const tirx::Var& var) -> tirx::Var { + if (defined_.count(var.get())) { + Var new_var = MakeNewVar(var); + PushVarRemap(var, new_var); + return new_var; + } else { + defined_.insert(var.get()); + return var; + } + }); + + // Remap implicitly defined buffer parameters + { + std::unordered_set defined_params; + for (const auto& var : func->params) { + defined_params.insert(var.get()); + } + std::unordered_set defined_match_vars; + for (const Var& param : func->params) { + auto buffer = param.as(); + if (!buffer) continue; + auto check_var = [&](const Var& var) { + const VarNode* var_ptr = var.get(); + if (defined_params.count(var_ptr)) return; + if (!defined_match_vars.insert(var_ptr).second) return; + + // Buffer-parameter shape vars use "match" semantics: first occurrence + // defines the var, subsequent occurrences (in other buffers) are + // just consistent uses of the same var -- not redefinitions. + if (defined_.count(var_ptr)) { + Var new_var = MakeNewVar(var); + PushVarRemap(var, new_var); + } else { + defined_.insert(var_ptr); } - if (auto var = buffer.value()->elem_offset.as()) check_var(var.value()); + }; + auto walk_fn = [&](const Var& var) -> ffi::Expected { + check_var(var); + return ffi::WalkResult::Advance(); + }; + for (const auto& dim : buffer.value()->shape) { + ffi::StructuralWalk(dim, walk_fn); } + for (const auto& stride : buffer.value()->strides) { + if (auto var = stride.as()) check_var(var.value()); + } + if (auto var = buffer.value()->elem_offset.as()) check_var(var.value()); } + } - // Update the buffer parameters, based on the redefined parameters - bool buffer_params_changed = false; - for (size_t i = 0; i < func->params.size(); ++i) { - if (auto buffer = func->params[i].as()) { - BufferVar new_buffer = GetRemappedBuffer(buffer.value()); - if (!new_buffer.same_as(buffer.value()) || !params[i].same_as(new_buffer)) { - buffer_params_changed = true; - params.Set(i, new_buffer.var()); - } + // Update the buffer parameters, based on the redefined parameters + bool buffer_params_changed = false; + for (size_t i = 0; i < func->params.size(); ++i) { + if (auto buffer = func->params[i].as()) { + BufferVar new_buffer = GetRemappedBuffer(buffer.value()); + if (!new_buffer.same_as(buffer.value()) || !params[i].same_as(new_buffer)) { + buffer_params_changed = true; + params.Set(i, new_buffer.var()); } } + } - auto attrs = [&]() -> DictAttrs { - ffi::Map dict; - bool made_change = false; - - for (const auto& [key, old_value] : func->attrs->dict) { - auto value = old_value; - if (auto expr = value.as()) { - value = Mutate(expr.value(), InplaceMode::kDisallow).ValueOrUnchanged(expr.value()); - } else if (auto* stmt = value.as()) { - value = Mutate(ffi::GetRef(stmt), InplaceMode::kDisallow) - .ValueOrUnchanged(ffi::GetRef(stmt)); - } - - made_change = made_change || !value.same_as(old_value); - dict.Set(key, value); - } + auto attrs = [&]() -> DictAttrs { + ffi::Map dict; + bool made_change = false; - if (made_change) { - return DictAttrs(dict); - } else { - return func->attrs; + for (const auto& [key, old_value] : func->attrs->dict) { + auto value = old_value; + if (auto expr = value.as()) { + value = Mutate(expr.value(), InplaceMode::kDisallow).ValueOrUnchanged(expr.value()); + } else if (auto* stmt = value.as()) { + value = Mutate(ffi::GetRef(stmt), InplaceMode::kDisallow) + .ValueOrUnchanged(ffi::GetRef(stmt)); } - }(); - auto body_result = Mutate(func->body, InplaceMode::kDisallow); - bool body_unchanged = body_result.UnchangedOrSameAs(func->body); - auto body = std::move(body_result).ValueOrUnchanged(func->body); + made_change = made_change || !value.same_as(old_value); + dict.Set(key, value); + } - // If anything changed, update the returned function - if (!params.same_as(func->params) || buffer_params_changed || !attrs.same_as(func->attrs) || - !body_unchanged) { - func = PrimFunc(params, body, func->ret_type, attrs); + if (made_change) { + return DictAttrs(dict); + } else { + return func->attrs; } + }(); + + auto body_result = Mutate(func->body, InplaceMode::kDisallow); + bool body_unchanged = body_result.UnchangedOrSameAs(func->body); + auto body = std::move(body_result).ValueOrUnchanged(func->body); - // Pop function-scope remaps in reverse order - PopAllRemapsInCurrentScope(); - function_scope_var_remap_.clear(); - return func; + // If anything changed, update the returned function + if (!params.same_as(func->params) || buffer_params_changed || !attrs.same_as(func->attrs) || + !body_unchanged) { + func = PrimFunc(params, body, func->ret_type, attrs); } - UnchangedOr Mutate_(const VarNode* op, InplaceMode inplace_mode) final { - Var var = ffi::GetRef(op); - Var mapped = GetRemappedVar(var); - if (!mapped.same_as(var)) return mapped; + // Pop function-scope remaps in reverse order + PopAllRemapsInCurrentScope(); + function_scope_var_remap_.clear(); + return func; +} + +UnchangedOr IRConvertSSA::Mutate_(const VarNode* op, InplaceMode inplace_mode) { + Var var = ffi::GetRef(op); + Var mapped = GetRemappedVar(var); + if (!mapped.same_as(var)) return mapped; + return StmtExprMutator::Mutate_(op, inplace_mode); +} + +UnchangedOr IRConvertSSA::Mutate_(const prim::LetNode* op, InplaceMode inplace_mode) { + const Var& v = op->var; + if (defined_.count(v.get())) { + PrimExpr value = this->Mutate(op->value, inplace_mode).ValueOrUnchanged(op->value); + Var new_var = MakeNewVar(v); + PushVarRemap(v, new_var); + PrimExpr body = this->Mutate(op->body, inplace_mode).ValueOrUnchanged(op->body); + PopVarRemap(v, new_var); + return prim::Let(new_var, value, body); + } else { + defined_.insert(v.get()); return StmtExprMutator::Mutate_(op, inplace_mode); } - UnchangedOr Mutate_(const prim::LetNode* op, InplaceMode inplace_mode) final { - const Var& v = op->var; - if (defined_.count(v.get())) { - PrimExpr value = this->Mutate(op->value, inplace_mode).ValueOrUnchanged(op->value); - Var new_var = MakeNewVar(v); - PushVarRemap(v, new_var); - PrimExpr body = this->Mutate(op->body, inplace_mode).ValueOrUnchanged(op->body); - PopVarRemap(v, new_var); - return prim::Let(new_var, value, body); - } else { - defined_.insert(v.get()); - return StmtExprMutator::Mutate_(op, inplace_mode); - } - } +} - UnchangedOr Mutate_(const TensorLoadNode* op, InplaceMode inplace_mode) final { - auto node = StmtExprMutator::Mutate_(op, inplace_mode) - .ValueOrUnchanged(ffi::GetRef(op)) - .as_or_throw(); - auto output = VisitBufferAccess(std::move(node)); - return output; - } +UnchangedOr IRConvertSSA::Mutate_(const TensorLoadNode* op, InplaceMode inplace_mode) { + auto node = StmtExprMutator::Mutate_(op, inplace_mode) + .ValueOrUnchanged(ffi::GetRef(op)) + .as_or_throw(); + auto output = VisitBufferAccess(std::move(node)); + return output; +} - UnchangedOr Mutate_(const BufferStoreNode* op, InplaceMode inplace_mode) final { - auto node = StmtExprMutator::Mutate_(op, inplace_mode) - .ValueOrUnchanged(ffi::GetRef(op)) - .as_or_throw(); - auto output = VisitBufferAccess(std::move(node)); - return output; - } +UnchangedOr IRConvertSSA::Mutate_(const BufferStoreNode* op, InplaceMode inplace_mode) { + auto node = StmtExprMutator::Mutate_(op, inplace_mode) + .ValueOrUnchanged(ffi::GetRef(op)) + .as_or_throw(); + auto output = VisitBufferAccess(std::move(node)); + return output; +} - UnchangedOr Mutate_(const DeclBufferNode* op, InplaceMode inplace_mode) final { - Var v = op->buffer.var(); - if (defined_.count(v.get())) { - Var new_var = MakeNewVar(v); - PushVarRemap(v, new_var); - } else { - defined_.insert(v.get()); - } - DeclBuffer decl = StmtExprMutator::Mutate_(op, inplace_mode) - .ValueOrUnchanged(ffi::GetRef(op)) - .as_or_throw(); - BufferVar new_buffer = GetRemappedBuffer(decl->buffer); - if (!new_buffer.same_as(decl->buffer)) { - decl.CopyOnWrite()->buffer = std::move(new_buffer); - } - return decl; +UnchangedOr IRConvertSSA::Mutate_(const DeclBufferNode* op, InplaceMode inplace_mode) { + Var v = op->buffer.var(); + if (defined_.count(v.get())) { + Var new_var = MakeNewVar(v); + PushVarRemap(v, new_var); + } else { + defined_.insert(v.get()); } + DeclBuffer decl = StmtExprMutator::Mutate_(op, inplace_mode) + .ValueOrUnchanged(ffi::GetRef(op)) + .as_or_throw(); + BufferVar new_buffer = GetRemappedBuffer(decl->buffer); + if (!new_buffer.same_as(decl->buffer)) { + decl.CopyOnWrite()->buffer = std::move(new_buffer); + } + return decl; +} - Stmt WithScope(const std::function& body) final { return scope_.WithNewScope(body); } +Stmt IRConvertSSA::WithScope(const std::function& body) { + return scope_.WithNewScope(body); +} - Var DefineVar(Var var) final { - if (defined_.count(var.get())) { - Var new_var = MakeNewVar(var); - PushVarRemap(var, new_var); - return new_var; - } - defined_.insert(var.get()); - return var; +Var IRConvertSSA::DefineVar(Var var) { + if (defined_.count(var.get())) { + Var new_var = MakeNewVar(var); + PushVarRemap(var, new_var); + return new_var; } + defined_.insert(var.get()); + return var; +} - BufferVar RemapBuffer(BufferVar buffer) final { return GetRemappedBuffer(buffer); } +BufferStore IRConvertSSA::VisitBufferAccess(BufferStore node) { + BufferVar new_buf = GetRemappedBuffer(node->buffer); + if (!new_buf.same_as(node->buffer)) { + auto writer = node.CopyOnWrite(); + writer->buffer = new_buf; + } - template - Node VisitBufferAccess(Node node) { - BufferVar new_buf = GetRemappedBuffer(node->buffer); - if (!new_buf.same_as(node->buffer)) { - auto writer = node.CopyOnWrite(); - writer->buffer = new_buf; - } + return node; +} +TensorLoad IRConvertSSA::VisitBufferAccess(TensorLoad node) { + BufferVar buffer = node->source.as_or_throw(); + BufferVar new_buf = GetRemappedBuffer(buffer); + if (new_buf.same_as(buffer)) { return node; } + return BufferLoad(new_buf, node->indices, node->span); +} - TensorRegion VisitBufferAccess(TensorRegion node) { - BufferVar buffer = node->source.as_or_throw(); - BufferVar new_buf = GetRemappedBuffer(buffer); - if (!new_buf.same_as(buffer)) { - node.CopyOnWrite()->source = new_buf; - } - return node; +Var IRConvertSSA::GetRemappedVar(Var var) { + if (auto it = scoped_var_remap_.find(var.get()); + it != scoped_var_remap_.end() && it->second.size()) { + return it->second.back(); + } else if (auto it = function_scope_var_remap_.find(var.get()); + it != function_scope_var_remap_.end()) { + return it->second; + } else { + return var; } +} - TensorLoad VisitBufferAccess(TensorLoad node) { - BufferVar buffer = node->source.as_or_throw(); - BufferVar new_buf = GetRemappedBuffer(buffer); - if (new_buf.same_as(buffer)) { - return node; +BufferVar IRConvertSSA::GetRemappedBuffer(BufferVar buf) { + // Determine the buffer var that should be in the updated buffer, + // given the current scope. If no redefines are present, then the + // buffer var is unchanged. + Var new_buffer_var = GetRemappedVar(buf.var()); + PrimExpr elem_offset = + Mutate(buf->elem_offset, InplaceMode::kDisallow).ValueOrUnchanged(buf->elem_offset); + auto visit_expr = [this](const PrimExpr& expr) { + return Mutate(expr, InplaceMode::kDisallow).ValueOrUnchanged(expr); + }; + ffi::Array shape = buf->shape.Map(visit_expr); + ffi::Array strides = buf->strides.Map(visit_expr); + + // Rewrite the layout's per-iter extent/stride expressions in lockstep + // with the shape. If we don't, SSA-renamed shape vars end up as fresh + // Vars while the layout still references the original, producing + // structurally-unequal buffers whose shape and layout disagree (e.g., + // test_dynamic_launch_thread). + ffi::Optional new_layout = buf->layout; + bool layout_changed = false; + if (buf->layout.has_value()) { + if (auto opt_tile = buf->layout.value().as()) { + auto remap_iter = [&](const Iter& it) -> Iter { + PrimExpr new_extent = + Mutate(it->extent, InplaceMode::kDisallow).ValueOrUnchanged(it->extent); + PrimExpr new_stride = + Mutate(it->stride, InplaceMode::kDisallow).ValueOrUnchanged(it->stride); + if (new_extent.same_as(it->extent) && new_stride.same_as(it->stride)) { + return it; + } + return Iter(new_extent, new_stride, it->axis); + }; + auto new_shard = opt_tile->shard.Map(remap_iter); + auto new_replica = opt_tile->replica.Map(remap_iter); + if (!new_shard.same_as(opt_tile->shard) || !new_replica.same_as(opt_tile->replica)) { + new_layout = TileLayout(new_shard, new_replica, opt_tile->offset); + layout_changed = true; + } } - return BufferLoad(new_buf, node->indices, node->span); } - Var GetRemappedVar(Var var) { - if (auto it = scoped_var_remap_.find(var.get()); - it != scoped_var_remap_.end() && it->second.size()) { - return it->second.back(); - } else if (auto it = function_scope_var_remap_.find(var.get()); - it != function_scope_var_remap_.end()) { - return it->second; - } else { - return var; - } + // If no mapping is required, return the original buffer. + if (new_buffer_var.same_as(buf.var()) && elem_offset.same_as(buf->elem_offset) && + shape.same_as(buf->shape) && strides.same_as(buf->strides) && !layout_changed) { + return buf; } - BufferVar GetRemappedBuffer(BufferVar buf) { - // Determine the buffer var that should be in the updated buffer, - // given the current scope. If no redefines are present, then the - // buffer var is unchanged. - Var new_buffer_var = GetRemappedVar(buf.var()); - PrimExpr elem_offset = - Mutate(buf->elem_offset, InplaceMode::kDisallow).ValueOrUnchanged(buf->elem_offset); - auto visit_expr = [this](const PrimExpr& expr) { - return Mutate(expr, InplaceMode::kDisallow).ValueOrUnchanged(expr); - }; - ffi::Array shape = buf->shape.Map(visit_expr); - ffi::Array strides = buf->strides.Map(visit_expr); - - // Rewrite the layout's per-iter extent/stride expressions in lockstep - // with the shape. If we don't, SSA-renamed shape vars end up as fresh - // Vars while the layout still references the original, producing - // structurally-unequal buffers whose shape and layout disagree (e.g., - // test_dynamic_launch_thread). - ffi::Optional new_layout = buf->layout; - bool layout_changed = false; - if (buf->layout.has_value()) { - if (auto opt_tile = buf->layout.value().as()) { - auto remap_iter = [&](const Iter& it) -> Iter { - PrimExpr new_extent = - Mutate(it->extent, InplaceMode::kDisallow).ValueOrUnchanged(it->extent); - PrimExpr new_stride = - Mutate(it->stride, InplaceMode::kDisallow).ValueOrUnchanged(it->stride); - if (new_extent.same_as(it->extent) && new_stride.same_as(it->stride)) { - return it; - } - return Iter(new_extent, new_stride, it->axis); - }; - auto new_shard = opt_tile->shard.Map(remap_iter); - auto new_replica = opt_tile->replica.Map(remap_iter); - if (!new_shard.same_as(opt_tile->shard) || !new_replica.same_as(opt_tile->replica)) { - new_layout = TileLayout(new_shard, new_replica, opt_tile->offset); - layout_changed = true; - } - } - } + // If the current scope already has a mapping of this buffer, use + // the mapped buffer. + auto key = buf.get(); + std::vector& buffers = buf_remap_[key]; + if (buffers.size() && buffers.back().same_as(new_buffer_var)) { + return buffers.back(); + } - // If no mapping is required, return the original buffer. - if (new_buffer_var.same_as(buf.var()) && elem_offset.same_as(buf->elem_offset) && - shape.same_as(buf->shape) && strides.same_as(buf->strides) && !layout_changed) { - return buf; + // When only the buffer's identity changed, the remapped Var already has + // the desired BufferType. Reuse that exact Var so the definition and all + // subsequent uses remain in SSA. + if (const auto* type = new_buffer_var->ty.as()) { + BufferVar candidate(new_buffer_var); + if (shape.same_as(type->shape) && strides.same_as(type->strides) && + elem_offset.same_as(type->elem_offset) && !layout_changed) { + buffers.push_back(candidate); + return candidate; } + } - // If the current scope already has a mapping of this buffer, use - // the mapped buffer. - auto key = buf.get(); - std::vector& buffers = buf_remap_[key]; - if (buffers.size() && buffers.back().same_as(new_buffer_var)) { - return buffers.back(); - } + // Otherwise, make and return a new buffer object that uses the + // new buffer, pushing it onto the scoped stack of existing + // buffers. This will be popped when the new_buffer_var + // redefinition is popped. + auto type = CopyBufferType(buf); + type->shape = shape; + type->strides = strides; + type->elem_offset = elem_offset; + if (layout_changed) { + type->layout = std::move(new_layout); + } + BufferVar new_buf = RebuildBufferVar(buf, std::move(type), new_buffer_var->name); + + // A BufferVar's metadata lives in its Var type. If rewriting the + // metadata required a fresh Var, make it the active remap as well. This + // keeps BufferLoad/BufferStore and ordinary Var uses (such as + // buffer_data) on the same identity. + auto it = scoped_var_remap_.find(buf.get()); + if (it != scoped_var_remap_.end() && it->second.size() && + it->second.back().same_as(new_buffer_var)) { + it->second.back() = new_buf.var(); + } else if (auto function_it = function_scope_var_remap_.find(buf.get()); + function_it != function_scope_var_remap_.end() && + function_it->second.same_as(new_buffer_var)) { + function_it->second = new_buf.var(); + } else { + PushVarRemap(buf.var(), new_buf.var()); + } + buffers.push_back(new_buf); + return new_buf; +} - // When only the buffer's identity changed, the remapped Var already has - // the desired BufferType. Reuse that exact Var so the definition and all - // subsequent uses remain in SSA. - if (const auto* type = new_buffer_var->ty.as()) { - BufferVar candidate(new_buffer_var); - if (shape.same_as(type->shape) && strides.same_as(type->strides) && - elem_offset.same_as(type->elem_offset) && !layout_changed) { - buffers.push_back(candidate); - return candidate; - } - } +UnchangedOr IRConvertSSA::Mutate_(const BindNode* op, InplaceMode inplace_mode) { + // Bind var remaps are tracked in the current scope so they persist + // across SeqStmt siblings and are cleaned up when the enclosing + // body-carrying statement's scope exits. + const Var& v = op->var; + if (defined_.count(v.get())) { + Expr value = this->Mutate(op->value, inplace_mode).ValueOrUnchanged(op->value); + Var new_var = MakeNewVar(v); + PushVarRemap(v, new_var); + return Bind(new_var, value); + } else { + defined_.insert(v.get()); + return StmtExprMutator::Mutate_(op, inplace_mode); + } +} - // Otherwise, make and return a new buffer object that uses the - // new buffer, pushing it onto the scoped stack of existing - // buffers. This will be popped when the new_buffer_var - // redefinition is popped. - auto type = CopyBufferType(buf); - type->shape = shape; - type->strides = strides; - type->elem_offset = elem_offset; - if (layout_changed) { - type->layout = std::move(new_layout); - } - BufferVar new_buf = RebuildBufferVar(buf, std::move(type), new_buffer_var->name); - - // A BufferVar's metadata lives in its Var type. If rewriting the - // metadata required a fresh Var, make it the active remap as well. This - // keeps BufferLoad/BufferStore and ordinary Var uses (such as - // buffer_data) on the same identity. - auto it = scoped_var_remap_.find(buf.get()); - if (it != scoped_var_remap_.end() && it->second.size() && - it->second.back().same_as(new_buffer_var)) { - it->second.back() = new_buf.var(); - } else if (auto function_it = function_scope_var_remap_.find(buf.get()); - function_it != function_scope_var_remap_.end() && - function_it->second.same_as(new_buffer_var)) { - function_it->second = new_buf.var(); - } else { - PushVarRemap(buf.var(), new_buf.var()); - } - buffers.push_back(new_buf); - return new_buf; +UnchangedOr IRConvertSSA::Mutate_(const IfThenElseNode* op, InplaceMode inplace_mode) { + // Each branch gets its own scope so Bind remaps in one branch + // do not leak into the other. + auto condition_result = Mutate(op->condition, inplace_mode); + bool condition_unchanged = condition_result.UnchangedOrSameAs(op->condition); + PrimExpr condition = std::move(condition_result).ValueOrUnchanged(op->condition); + Stmt then_case = scope_.WithNewScope([&]() -> Stmt { + return Mutate(op->then_case, inplace_mode).ValueOrUnchanged(op->then_case); + }); + ffi::Optional else_case; + if (op->else_case) { + else_case = scope_.WithNewScope([&]() -> Stmt { + return Mutate(op->else_case.value(), inplace_mode).ValueOrUnchanged(op->else_case.value()); + }); } + if (condition_unchanged && then_case.same_as(op->then_case) && else_case.same_as(op->else_case)) { + return ffi::Unchanged(); + } + return IfThenElse(condition, then_case, else_case); +} - UnchangedOr Mutate_(const BindNode* op, InplaceMode inplace_mode) final { - // Bind var remaps are tracked in the current scope so they persist - // across SeqStmt siblings and are cleaned up when the enclosing - // body-carrying statement's scope exits. - const Var& v = op->var; - if (defined_.count(v.get())) { - Expr value = this->Mutate(op->value, inplace_mode).ValueOrUnchanged(op->value); +UnchangedOr IRConvertSSA::Mutate_(const ForNode* op, InplaceMode inplace_mode) { + const Var& v = op->loop_var; + if (defined_.count(v.get())) { + return scope_.WithNewScope([&]() -> Stmt { Var new_var = MakeNewVar(v); PushVarRemap(v, new_var); - return Bind(new_var, value); - } else { - defined_.insert(v.get()); - return StmtExprMutator::Mutate_(op, inplace_mode); - } - } - - UnchangedOr Mutate_(const IfThenElseNode* op, InplaceMode inplace_mode) final { - // Each branch gets its own scope so Bind remaps in one branch - // do not leak into the other. - auto condition_result = Mutate(op->condition, inplace_mode); - bool condition_unchanged = condition_result.UnchangedOrSameAs(op->condition); - PrimExpr condition = std::move(condition_result).ValueOrUnchanged(op->condition); - Stmt then_case = scope_.WithNewScope([&]() -> Stmt { - return Mutate(op->then_case, inplace_mode).ValueOrUnchanged(op->then_case); + Stmt stmt = + StmtExprMutator::Mutate_(op, inplace_mode).ValueOrUnchanged(ffi::GetRef(op)); + auto n = ffi::make_object(*stmt.as()); + n->loop_var = new_var.as_or_throw(); + return For(n); }); - ffi::Optional else_case; - if (op->else_case) { - else_case = scope_.WithNewScope([&]() -> Stmt { - return Mutate(op->else_case.value(), inplace_mode).ValueOrUnchanged(op->else_case.value()); - }); - } - if (condition_unchanged && then_case.same_as(op->then_case) && - else_case.same_as(op->else_case)) { - return ffi::Unchanged(); - } - return IfThenElse(condition, then_case, else_case); - } - - UnchangedOr Mutate_(const ForNode* op, InplaceMode inplace_mode) final { - const Var& v = op->loop_var; - if (defined_.count(v.get())) { - return scope_.WithNewScope([&]() -> Stmt { - Var new_var = MakeNewVar(v); - PushVarRemap(v, new_var); - Stmt stmt = - StmtExprMutator::Mutate_(op, inplace_mode).ValueOrUnchanged(ffi::GetRef(op)); - auto n = ffi::make_object(*stmt.as()); - n->loop_var = new_var.as_or_throw(); - return For(n); - }); - } else { - defined_.insert(v.get()); - return scope_.WithNewScope([&]() -> Stmt { - return StmtExprMutator::Mutate_(op, inplace_mode).ValueOrUnchanged(ffi::GetRef(op)); - }); - } - } - UnchangedOr Mutate_(const WhileNode* op, InplaceMode inplace_mode) final { + } else { + defined_.insert(v.get()); return scope_.WithNewScope([&]() -> Stmt { return StmtExprMutator::Mutate_(op, inplace_mode).ValueOrUnchanged(ffi::GetRef(op)); }); } - UnchangedOr Mutate_(const AllocBufferNode* op, InplaceMode inplace_mode) final { - Var v = op->buffer.var(); - if (defined_.count(v.get())) { - Var new_var = MakeNewVar(v); - PushVarRemap(v, new_var); - } else { - defined_.insert(v.get()); - } - Stmt stmt = StmtExprMutator::Mutate_(op, inplace_mode).ValueOrUnchanged(ffi::GetRef(op)); - op = stmt.as(); - // Use GetRemappedBuffer so that the AllocBuffer's buffer is the same - // object as the one used by BufferStore/TensorLoad in subsequent siblings. - BufferVar new_buf = GetRemappedBuffer(op->buffer); - if (!new_buf.same_as(op->buffer)) { - auto node = stmt.as_or_throw(); - node.CopyOnWrite()->buffer = std::move(new_buf); - return node; - } - return stmt; - } - UnchangedOr Mutate_(const AttrStmtNode* op, InplaceMode inplace_mode) final { - if (const IterVarNode* iter_var = op->node.as()) { - Range dom = iter_var->dom; - if (dom.defined()) { - // Retain the original domain while comparing and rebuilding its replacement. - auto min = Mutate(dom->min, InplaceMode::kDisallow).ValueOrUnchanged(dom->min); - auto extent = Mutate(dom->extent, InplaceMode::kDisallow).ValueOrUnchanged(dom->extent); - if (!min.same_as(iter_var->dom->min) || !extent.same_as(iter_var->dom->extent)) { - dom = Range::FromMinExtent(min, extent); - } - } +} - Var var = iter_var->var; - bool delayed_define = false; - if (auto it = function_scope_var_remap_.find(var.get()); - it != function_scope_var_remap_.end()) { - var = it->second; - } else if (defined_.count(var.get())) { - Var new_var(var->name, var->ty); - - function_scope_var_remap_.insert({var.get(), new_var}); - var = new_var; - } else { - // The AttrStmt refers to an undefined variable. This is - // allowed for some attributes, such as - // "pragma_parallel_launch_point", which annotates a variable - // that is about to occur in a ForNode. In these cases, the - // ForNode and the AttrStmt must continue using the same - // variable defintion. - // - // However, other AttrStmt, such as "thread_extent", act as - // points of definition for the variable they annotate. If - // the variable has not been defined after visiting the body, - // we should mark it as defined before exiting. This ensures - // correct de-duplication between multiple functions. - // - // This implementation may be simplified in the future by - // moving "pragma_parallel_launch_point" to be an annotation - // on the `ForNode`, rather than an `AttrStmt`. - delayed_define = true; - } +UnchangedOr IRConvertSSA::Mutate_(const WhileNode* op, InplaceMode inplace_mode) { + return scope_.WithNewScope([&]() -> Stmt { + return StmtExprMutator::Mutate_(op, inplace_mode).ValueOrUnchanged(ffi::GetRef(op)); + }); +} - IterVar new_iter_var; - if (dom.same_as(iter_var->dom) && var.same_as(iter_var->var)) { - new_iter_var = ffi::GetRef(iter_var); - } else { - new_iter_var = IterVar(dom, var.as_or_throw(), iter_var->iter_type, - iter_var->thread_tag, iter_var->span); - } - auto value_result = Mutate(op->value, inplace_mode); - bool value_unchanged = value_result.UnchangedOrSameAs(op->value); - auto value = std::move(value_result).ValueOrUnchanged(op->value); - auto body = scope_.WithNewScope( - [&]() -> Stmt { return Mutate(op->body, inplace_mode).ValueOrUnchanged(op->body); }); - - Stmt output; - if (new_iter_var.get() == iter_var && body.same_as(op->body) && value_unchanged) { - output = ffi::GetRef(op); - } else { - output = AttrStmt(new_iter_var, op->attr_key, value, body, iter_var->span); - } +UnchangedOr IRConvertSSA::Mutate_(const AllocBufferNode* op, InplaceMode inplace_mode) { + Var v = op->buffer.var(); + if (defined_.count(v.get())) { + Var new_var = MakeNewVar(v); + PushVarRemap(v, new_var); + } else { + defined_.insert(v.get()); + } + Stmt stmt = StmtExprMutator::Mutate_(op, inplace_mode).ValueOrUnchanged(ffi::GetRef(op)); + op = stmt.as(); + // Use GetRemappedBuffer so that the AllocBuffer's buffer is the same + // object as the one used by BufferStore/TensorLoad in subsequent siblings. + BufferVar new_buf = GetRemappedBuffer(op->buffer); + if (!new_buf.same_as(op->buffer)) { + auto node = stmt.as_or_throw(); + node.CopyOnWrite()->buffer = std::move(new_buf); + return node; + } + return stmt; +} - if (delayed_define) { - if (!defined_.count(var.get())) { - function_scope_var_remap_.insert({var.get(), var}); - defined_.insert(var.get()); - } +UnchangedOr IRConvertSSA::Mutate_(const AttrStmtNode* op, InplaceMode inplace_mode) { + if (const IterVarNode* iter_var = op->node.as()) { + Range dom = iter_var->dom; + if (dom.defined()) { + // Retain the original domain while comparing and rebuilding its replacement. + auto min = Mutate(dom->min, InplaceMode::kDisallow).ValueOrUnchanged(dom->min); + auto extent = Mutate(dom->extent, InplaceMode::kDisallow).ValueOrUnchanged(dom->extent); + if (!min.same_as(iter_var->dom->min) || !extent.same_as(iter_var->dom->extent)) { + dom = Range::FromMinExtent(min, extent); } + } + + Var var = iter_var->var; + bool delayed_define = false; + if (auto it = function_scope_var_remap_.find(var.get()); + it != function_scope_var_remap_.end()) { + var = it->second; + } else if (defined_.count(var.get())) { + Var new_var(var->name, var->ty); + + function_scope_var_remap_.insert({var.get(), new_var}); + var = new_var; + } else { + // The AttrStmt refers to an undefined variable. This is + // allowed for some attributes, such as + // "pragma_parallel_launch_point", which annotates a variable + // that is about to occur in a ForNode. In these cases, the + // ForNode and the AttrStmt must continue using the same + // variable defintion. + // + // However, other AttrStmt, such as "thread_extent", act as + // points of definition for the variable they annotate. If + // the variable has not been defined after visiting the body, + // we should mark it as defined before exiting. This ensures + // correct de-duplication between multiple functions. + // + // This implementation may be simplified in the future by + // moving "pragma_parallel_launch_point" to be an annotation + // on the `ForNode`, rather than an `AttrStmt`. + delayed_define = true; + } + + IterVar new_iter_var; + if (dom.same_as(iter_var->dom) && var.same_as(iter_var->var)) { + new_iter_var = ffi::GetRef(iter_var); + } else { + new_iter_var = IterVar(dom, var.as_or_throw(), iter_var->iter_type, + iter_var->thread_tag, iter_var->span); + } + auto value_result = Mutate(op->value, inplace_mode); + bool value_unchanged = value_result.UnchangedOrSameAs(op->value); + auto value = std::move(value_result).ValueOrUnchanged(op->value); + auto body = scope_.WithNewScope( + [&]() -> Stmt { return Mutate(op->body, inplace_mode).ValueOrUnchanged(op->body); }); + + Stmt output; + if (new_iter_var.get() == iter_var && body.same_as(op->body) && value_unchanged) { + output = ffi::GetRef(op); + } else { + output = AttrStmt(new_iter_var, op->attr_key, value, body, iter_var->span); + } - return output; - - } else if (const VarNode* v = op->node.as()) { - Stmt stmt = scope_.WithNewScope([&]() -> Stmt { - return StmtExprMutator::Mutate_(op, inplace_mode).ValueOrUnchanged(ffi::GetRef(op)); - }); - op = stmt.as(); - if (scoped_var_remap_.count(v) && scoped_var_remap_[v].size() != 0) { - return AttrStmt(scoped_var_remap_[v].back(), op->attr_key, op->value, op->body); - } else { - return stmt; + if (delayed_define) { + if (!defined_.count(var.get())) { + function_scope_var_remap_.insert({var.get(), var}); + defined_.insert(var.get()); } + } + + return output; + + } else if (const VarNode* v = op->node.as()) { + Stmt stmt = scope_.WithNewScope([&]() -> Stmt { + return StmtExprMutator::Mutate_(op, inplace_mode).ValueOrUnchanged(ffi::GetRef(op)); + }); + op = stmt.as(); + if (scoped_var_remap_.count(v) && scoped_var_remap_[v].size() != 0) { + return AttrStmt(scoped_var_remap_[v].back(), op->attr_key, op->value, op->body); } else { - return scope_.WithNewScope([&]() -> Stmt { - return StmtExprMutator::Mutate_(op, inplace_mode).ValueOrUnchanged(ffi::GetRef(op)); - }); + return stmt; } + } else { + return scope_.WithNewScope([&]() -> Stmt { + return StmtExprMutator::Mutate_(op, inplace_mode).ValueOrUnchanged(ffi::GetRef(op)); + }); } +} - private: - /*! \brief Record of a variable remap pushed to the current scope. */ - struct VarRemap { - Var old_var; - Var new_var; - }; - - /*! \brief Check whether a buffer uses a variable in any remapped field. */ - static bool BufferDependsOnVar(const BufferVar& buffer, const VarNode* var) { - if (buffer.get() == var) return true; +bool IRConvertSSA::BufferDependsOnVar(const BufferVar& buffer, const VarNode* var) { + if (buffer.get() == var) return true; - auto uses_var = [var](const PrimExpr& expr) { - auto walkfn = [var](const Var& candidate) -> ffi::Expected { - return candidate.get() == var ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(candidate)) - : ffi::WalkResult::Advance(); - }; - return expr.defined() && - ffi::StructuralWalk(expr, walkfn).has_value(); + auto uses_var = [var](const PrimExpr& expr) { + auto walkfn = [var](const Var& candidate) -> ffi::Expected { + return candidate.get() == var ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(candidate)) + : ffi::WalkResult::Advance(); }; - if (uses_var(buffer->elem_offset)) return true; - for (const PrimExpr& dim : buffer->shape) { - if (uses_var(dim)) return true; - } - for (const PrimExpr& stride : buffer->strides) { - if (uses_var(stride)) return true; - } - if (buffer->layout.has_value()) { - if (const auto* tile_layout = buffer->layout.value().as()) { - for (const Iter& iter : tile_layout->shard) { - if (uses_var(iter->extent) || uses_var(iter->stride)) return true; - } - for (const Iter& iter : tile_layout->replica) { - if (uses_var(iter->extent) || uses_var(iter->stride)) return true; - } + return expr.defined() && + ffi::StructuralWalk(expr, walkfn).has_value(); + }; + if (uses_var(buffer->elem_offset)) return true; + for (const PrimExpr& dim : buffer->shape) { + if (uses_var(dim)) return true; + } + for (const PrimExpr& stride : buffer->strides) { + if (uses_var(stride)) return true; + } + if (buffer->layout.has_value()) { + if (const auto* tile_layout = buffer->layout.value().as()) { + for (const Iter& iter : tile_layout->shard) { + if (uses_var(iter->extent) || uses_var(iter->stride)) return true; + } + for (const Iter& iter : tile_layout->replica) { + if (uses_var(iter->extent) || uses_var(iter->stride)) return true; } } - return false; } + return false; +} + +Var IRConvertSSA::MakeNewVar(const Var& old_var) { return Var(old_var->name, old_var->ty); } - /*! \brief Create a new variable with the same name and type as the original. */ - static Var MakeNewVar(const Var& old_var) { return Var(old_var->name, old_var->ty); } +void IRConvertSSA::PushVarRemap(const Var& old_var, const Var& new_var) { + scoped_var_remap_[old_var.get()].push_back(new_var); + auto& level = scope_.Current(); + level.parent = this; + level.push_back({old_var, new_var}); +} - /*! \brief Push a variable remap to the current scope and the scoped_var_remap_ stack. */ - void PushVarRemap(const Var& old_var, const Var& new_var) { - scoped_var_remap_[old_var.get()].push_back(new_var); - auto& level = scope_.Current(); - level.parent = this; - level.push_back({old_var, new_var}); +void IRConvertSSA::PopVarRemap(const Var& old_var, const Var& new_var) { + scoped_var_remap_[old_var.get()].pop_back(); + for (auto& kv : buf_remap_) { + std::vector& buffers = kv.second; + if (buffers.size() && BufferDependsOnVar(buffers.back(), new_var.get())) { + buffers.pop_back(); + } + } + // Also remove from the current scope's tracking vector + auto& current = scope_.Current(); + if (current.size() && current.back().new_var.same_as(new_var)) { + current.pop_back(); } +} - /*! \brief Pop a single variable remap (used for expression-level Let scoping). */ - void PopVarRemap(const Var& old_var, const Var& new_var) { - scoped_var_remap_[old_var.get()].pop_back(); +void IRConvertSSA::PopAllRemapsInCurrentScope() { + auto& current = scope_.Current(); + while (current.size()) { + auto& remap = current.back(); + scoped_var_remap_[remap.old_var.get()].pop_back(); for (auto& kv : buf_remap_) { std::vector& buffers = kv.second; - if (buffers.size() && BufferDependsOnVar(buffers.back(), new_var.get())) { + if (buffers.size() && BufferDependsOnVar(buffers.back(), remap.new_var.get())) { buffers.pop_back(); } } - // Also remove from the current scope's tracking vector - auto& current = scope_.Current(); - if (current.size() && current.back().new_var.same_as(new_var)) { - current.pop_back(); - } - } - - /*! \brief Pop all remaps in the current scope level (used for function-scope cleanup). */ - void PopAllRemapsInCurrentScope() { - auto& current = scope_.Current(); - while (current.size()) { - auto& remap = current.back(); - scoped_var_remap_[remap.old_var.get()].pop_back(); - for (auto& kv : buf_remap_) { - std::vector& buffers = kv.second; - if (buffers.size() && BufferDependsOnVar(buffers.back(), remap.new_var.get())) { - buffers.pop_back(); - } - } - current.pop_back(); - } + current.pop_back(); } - - /*! \brief Scope stack: each scope level holds the remaps introduced in that scope. - * - * When a body-carrying statement (For, SBlock, Allocate) calls - * scope_.WithNewScope([&]{...}), a new scope level is pushed. - * Bind statements push their remaps to the current scope. - * On scope exit, the destructor of std::vector triggers, - * and we undo all remaps in that level. - * - * Note: ScopeStack::WithNewScope calls T's destructor on exit. - * std::vector's destructor destroys elements but does NOT call custom - * cleanup. So we wrap the vector in ScopeLevel which handles cleanup. - */ - struct ScopeLevel { - std::vector remaps; - IRConvertSSA* parent{nullptr}; - - void push_back(VarRemap remap) { remaps.push_back(std::move(remap)); } - size_t size() const { return remaps.size(); } - VarRemap& back() { return remaps.back(); } - void pop_back() { remaps.pop_back(); } - - ~ScopeLevel() { - if (!parent) return; - // Pop remaps in reverse order - while (remaps.size()) { - auto& remap = remaps.back(); - parent->scoped_var_remap_[remap.old_var.get()].pop_back(); - for (auto& kv : parent->buf_remap_) { - std::vector& buffers = kv.second; - if (buffers.size() && BufferDependsOnVar(buffers.back(), remap.new_var.get())) { - buffers.pop_back(); - } - } - remaps.pop_back(); - } - } - - ScopeLevel() = default; - ScopeLevel(const ScopeLevel&) = delete; - ScopeLevel& operator=(const ScopeLevel&) = delete; - ScopeLevel(ScopeLevel&& other) noexcept - : remaps(std::move(other.remaps)), parent(other.parent) { - other.parent = nullptr; // prevent other's destructor from popping - } - ScopeLevel& operator=(ScopeLevel&& other) noexcept { - if (this != &other) { - // Run our destructor logic first - if (parent) { - while (remaps.size()) { - auto& remap = remaps.back(); - parent->scoped_var_remap_[remap.old_var.get()].pop_back(); - for (auto& kv : parent->buf_remap_) { - std::vector& buffers = kv.second; - if (buffers.size() && BufferDependsOnVar(buffers.back(), remap.new_var.get())) { - buffers.pop_back(); - } - } - remaps.pop_back(); - } - } - remaps = std::move(other.remaps); - parent = other.parent; - other.parent = nullptr; - } - return *this; - } - }; - - std::unordered_map> scoped_var_remap_; - std::unordered_set defined_; - std::unordered_map> buf_remap_; - std::unordered_map function_scope_var_remap_; - ScopeStack scope_; -}; +} Stmt ConvertSSA(Stmt stmt) { return ffi::make_object()->Mutate(stmt, InplaceMode::kAllow).ValueOrUnchanged(stmt); @@ -810,26 +710,29 @@ std::optional IsHostFunc(const PrimFunc& func) { } } +IRModule IRConvertSSA::VisitIRModule(IRModule mod) { + ffi::Map functions; + bool made_change = false; + for (auto [gvar, base_func] : mod->functions) { + if (auto* ptr = base_func.as()) { + auto updated = VisitPrimFunc(ffi::GetRef(ptr)); + if (!updated.same_as(base_func)) { + made_change = true; + base_func = updated; + } + } + functions.Set(gvar, base_func); + } + if (made_change) { + mod.CopyOnWrite()->functions = std::move(functions); + } + return mod; +} + namespace transform { Pass ConvertSSA() { auto pass_func = [](IRModule mod, PassContext ctx) { - auto converter = ffi::make_object(); - ffi::Map functions; - bool made_change = false; - for (auto [gvar, base_func] : mod->functions) { - if (auto* ptr = base_func.as()) { - auto updated = converter->VisitPrimFunc(ffi::GetRef(ptr)); - if (!updated.same_as(base_func)) { - made_change = true; - base_func = updated; - } - } - functions.Set(gvar, base_func); - } - if (made_change) { - mod.CopyOnWrite()->functions = std::move(functions); - } - return mod; + return ffi::make_object()->VisitIRModule(std::move(mod)); }; return tvm::transform::CreateModulePass(pass_func, 0, "tirx.ConvertSSA", {}); } diff --git a/src/tirx/transform/ir_utils.h b/src/tirx/transform/ir_utils.h index 05aa9e4e307c..0ee2c13028b0 100644 --- a/src/tirx/transform/ir_utils.h +++ b/src/tirx/transform/ir_utils.h @@ -27,17 +27,21 @@ #include #include #include +#include #include #include #include #include #include #include +#include +#include #include #include #include #include +#include #include #include @@ -223,6 +227,120 @@ inline Call StackAlloca(Type ret_type, std::string type, size_t num) { */ Stmt ConvertSSA(Stmt stmt); +/*! \brief Shared SSA renaming algorithm; dialects extend statement dispatch explicitly. */ +class IRConvertSSA : public StmtExprMutator { + public: + TVM_DEFINE_OBJECT_FUNCTOR_DEFAULT_CONSTRUCTOR(IRConvertSSA, StmtExprMutator) + using StmtExprMutator::Mutate; + using StmtExprMutator::Mutate_; + PrimFunc VisitPrimFunc(PrimFunc func); + IRModule VisitIRModule(IRModule mod); + + protected: + explicit IRConvertSSA(const VTable* table) : StmtExprMutator(table) {} + UnchangedOr Mutate_(const VarNode* op, InplaceMode inplace_mode) final; + UnchangedOr Mutate_(const prim::LetNode* op, InplaceMode inplace_mode) final; + UnchangedOr Mutate_(const TensorLoadNode* op, InplaceMode inplace_mode) final; + UnchangedOr Mutate_(const BufferStoreNode* op, InplaceMode inplace_mode) final; + UnchangedOr Mutate_(const DeclBufferNode* op, InplaceMode inplace_mode) final; + Stmt WithScope(const std::function& body); + Var DefineVar(Var var); + BufferStore VisitBufferAccess(BufferStore node); + TensorLoad VisitBufferAccess(TensorLoad node); + Var GetRemappedVar(Var var); + BufferVar GetRemappedBuffer(BufferVar buf); + UnchangedOr Mutate_(const BindNode* op, InplaceMode inplace_mode) final; + UnchangedOr Mutate_(const IfThenElseNode* op, InplaceMode inplace_mode) final; + UnchangedOr Mutate_(const ForNode* op, InplaceMode inplace_mode) final; + UnchangedOr Mutate_(const WhileNode* op, InplaceMode inplace_mode) final; + UnchangedOr Mutate_(const AllocBufferNode* op, InplaceMode inplace_mode) final; + UnchangedOr Mutate_(const AttrStmtNode* op, InplaceMode inplace_mode) final; + static bool BufferDependsOnVar(const BufferVar& buffer, const VarNode* var); + static Var MakeNewVar(const Var& old_var); + void PushVarRemap(const Var& old_var, const Var& new_var); + void PopVarRemap(const Var& old_var, const Var& new_var); + void PopAllRemapsInCurrentScope(); + + private: + struct VarRemap { + Var old_var; + Var new_var; + }; + /*! \brief Scope stack: each scope level holds the remaps introduced in that scope. + * + * When a body-carrying statement (For, Allocate, or a dialect statement) calls + * scope_.WithNewScope([&]{...}), a new scope level is pushed. + * Bind statements push their remaps to the current scope. + * On scope exit, the destructor of std::vector triggers, + * and we undo all remaps in that level. + * + * Note: ScopeStack::WithNewScope calls T's destructor on exit. + * std::vector's destructor destroys elements but does NOT call custom + * cleanup. So we wrap the vector in ScopeLevel which handles cleanup. + */ + struct ScopeLevel { + std::vector remaps; + IRConvertSSA* parent{nullptr}; + + void push_back(VarRemap remap) { remaps.push_back(std::move(remap)); } + size_t size() const { return remaps.size(); } + VarRemap& back() { return remaps.back(); } + void pop_back() { remaps.pop_back(); } + + ~ScopeLevel() { + if (!parent) return; + // Pop remaps in reverse order + while (remaps.size()) { + auto& remap = remaps.back(); + parent->scoped_var_remap_[remap.old_var.get()].pop_back(); + for (auto& kv : parent->buf_remap_) { + std::vector& buffers = kv.second; + if (buffers.size() && BufferDependsOnVar(buffers.back(), remap.new_var.get())) { + buffers.pop_back(); + } + } + remaps.pop_back(); + } + } + + ScopeLevel() = default; + ScopeLevel(const ScopeLevel&) = delete; + ScopeLevel& operator=(const ScopeLevel&) = delete; + ScopeLevel(ScopeLevel&& other) noexcept + : remaps(std::move(other.remaps)), parent(other.parent) { + other.parent = nullptr; // prevent other's destructor from popping + } + ScopeLevel& operator=(ScopeLevel&& other) noexcept { + if (this != &other) { + // Run our destructor logic first + if (parent) { + while (remaps.size()) { + auto& remap = remaps.back(); + parent->scoped_var_remap_[remap.old_var.get()].pop_back(); + for (auto& kv : parent->buf_remap_) { + std::vector& buffers = kv.second; + if (buffers.size() && BufferDependsOnVar(buffers.back(), remap.new_var.get())) { + buffers.pop_back(); + } + } + remaps.pop_back(); + } + } + remaps = std::move(other.remaps); + parent = other.parent; + other.parent = nullptr; + } + return *this; + } + }; + + std::unordered_map> scoped_var_remap_; + std::unordered_set defined_; + std::unordered_map> buf_remap_; + std::unordered_map function_scope_var_remap_; + ScopeStack scope_; +}; + /*! * \brief Return the storage scope associated with a buffer variable. * \param buffer_var The input buffer variable. diff --git a/src/tirx/transform/narrow_datatype.cc b/src/tirx/transform/narrow_datatype.cc index 99bc06b62ce9..a57919daa508 100644 --- a/src/tirx/transform/narrow_datatype.cc +++ b/src/tirx/transform/narrow_datatype.cc @@ -33,7 +33,6 @@ #include #include "../ir/data_type_rewriter.h" -#include "stmt_extension.h" namespace tvm { namespace tirx { @@ -75,7 +74,7 @@ using arith::ConstIntBound; // then we narrow `var` into `target_bits_`. That is, // `vmap[var] = min(target_bits_, var.dtype.bits())` // Otherwise, `var` is not narrowed, that is, `vmap[var] = var.dtype.bits()` -class DataTypeVisitor final : public IndexDomainVisitor { +class DataTypeVisitor final : public StmtExprVisitor { public: explicit DataTypeVisitor(int target_bits) : bits_(target_bits), target_bits_(target_bits) {} @@ -124,11 +123,6 @@ class DataTypeVisitor final : public IndexDomainVisitor { return StmtExprVisitor::Visit_(op); } - void BindDomain(const Var& var, const Range& domain) final { - analyzer_->Bind(var, domain); - vextent_.insert_or_assign(var.as(), domain->extent.ty()); - } - ffi::Optional Visit_(const AttrStmtNode* op) { if (op->attr_key == attr::thread_extent || op->attr_key == "virtual_thread") { IterVar iv = op->node.as_or_throw(); diff --git a/src/tirx/transform/stmt_extension.h b/src/tirx/transform/stmt_extension.h deleted file mode 100644 index 2cf709e82b71..000000000000 --- a/src/tirx/transform/stmt_extension.h +++ /dev/null @@ -1,129 +0,0 @@ -/* - * 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_extension.h - * \brief Narrow operation contexts for dialect-owned statement transform hooks. - */ -#ifndef TVM_TIRX_TRANSFORM_STMT_EXTENSION_H_ -#define TVM_TIRX_TRANSFORM_STMT_EXTENSION_H_ - -#include -#include - -#include "../ir_mutator_with_analyzer.h" - -namespace tvm { -namespace tirx { - -// Expose only the scope and remap operations needed at extension definition sites. -class SSAStmtMutator : public StmtExprMutator { - public: - using StmtExprMutator::VTable; - static void RegisterExtension(void (*init)(VTable*)) { Extensions().push_back(init); } - SSAStmtMutator() : StmtExprMutator(GlobalVTable()) {} - virtual Stmt WithScope(const std::function& body) = 0; - virtual Var DefineVar(Var var) = 0; - virtual BufferVar RemapBuffer(BufferVar buffer) = 0; - - protected: - static void InitVTable(VTable* table) { - StmtExprMutator::InitVTable(table); - for (auto init : Extensions()) init(table); - } - - private: - static std::vector& Extensions() { - static std::vector extensions; - return extensions; - } - static const VTable* GlobalVTable() { - static const VTable table = [] { - VTable table; - InitVTable(&table); - table.Finalize(); - return table; - }(); - return &table; - } -}; - -// Buffer geometry stays in the flattener; extensions identify definitions and regions. -class FlattenStmtMutator : public IRMutatorWithAnalyzer { - public: - using IRMutatorWithAnalyzer::VTable; - static void RegisterExtension(void (*init)(VTable*)) { Extensions().push_back(init); } - explicit FlattenStmtMutator(const arith::Analyzer& analyzer) - : IRMutatorWithAnalyzer(analyzer.get(), GlobalVTable()) {} - virtual BufferVar DefineBuffer(BufferVar buffer) = 0; - virtual BufferRegion RewriteRegion(BufferRegion region) = 0; - - protected: - static void InitVTable(VTable* table) { - IRMutatorWithAnalyzer::InitVTable(table); - for (auto init : Extensions()) init(table); - } - - private: - static std::vector& Extensions() { - static std::vector extensions; - return extensions; - } - static const VTable* GlobalVTable() { - static const VTable table = [] { - VTable table; - InitVTable(&table); - table.Finalize(); - return table; - }(); - return &table; - } -}; - -class IndexDomainVisitor : public StmtExprVisitor { - public: - using StmtExprVisitor::VTable; - static void RegisterExtension(void (*init)(VTable*)) { Extensions().push_back(init); } - IndexDomainVisitor() : StmtExprVisitor(GlobalVTable()) {} - virtual void BindDomain(const Var& var, const Range& domain) = 0; - - protected: - static void InitVTable(VTable* table) { - StmtExprVisitor::InitVTable(table); - for (auto init : Extensions()) init(table); - } - - private: - static std::vector& Extensions() { - static std::vector extensions; - return extensions; - } - static const VTable* GlobalVTable() { - static const VTable table = [] { - VTable table; - InitVTable(&table); - table.Finalize(); - return table; - }(); - return &table; - } -}; - -} // namespace tirx -} // namespace tvm -#endif // TVM_TIRX_TRANSFORM_STMT_EXTENSION_H_ diff --git a/src/tirx/transform/stmt_simplify.cc b/src/tirx/transform/stmt_simplify.cc index fbd51222a570..af2c5da6dd43 100644 --- a/src/tirx/transform/stmt_simplify.cc +++ b/src/tirx/transform/stmt_simplify.cc @@ -22,7 +22,7 @@ * \brief Statement simplifier based on analyzer */ -#include "../../tirx/transform/stmt_simplify.h" +#include "stmt_simplify.h" #include #include @@ -44,52 +44,37 @@ using namespace tvm::prim; using namespace tirx; -struct StmtSimplifyConfigNode : public ffi::Object { - bool transitively_prove_inequalities; - bool convert_boolean_to_and_of_ors; - bool apply_constraints_to_boolean_branches; +void StmtSimplifyConfigNode::RegisterReflection() { + namespace refl = tvm::ffi::reflection; + refl::ObjectDef() + .def_ro("transitively_prove_inequalities", + &StmtSimplifyConfigNode::transitively_prove_inequalities, + "If true, simplify conditionals with transitive combinations of scoped constraints", + refl::DefaultValue(false)) + .def_ro("convert_boolean_to_and_of_ors", + &StmtSimplifyConfigNode::convert_boolean_to_and_of_ors, + "If true, simplify conditionals into an AND of ORs", refl::DefaultValue(false)) + .def_ro("apply_constraints_to_boolean_branches", + &StmtSimplifyConfigNode::apply_constraints_to_boolean_branches, + "If true, simplify each branch of AND/OR under constraints provided by the other " + "branch", + refl::DefaultValue(false)); +} - static void RegisterReflection() { - namespace refl = tvm::ffi::reflection; - refl::ObjectDef() - .def_ro("transitively_prove_inequalities", - &StmtSimplifyConfigNode::transitively_prove_inequalities, - "If true, simplify conditionals with transitive combinations of scoped constraints", - refl::DefaultValue(false)) - .def_ro("convert_boolean_to_and_of_ors", - &StmtSimplifyConfigNode::convert_boolean_to_and_of_ors, - "If true, simplify conditionals into an AND of ORs", refl::DefaultValue(false)) - .def_ro("apply_constraints_to_boolean_branches", - &StmtSimplifyConfigNode::apply_constraints_to_boolean_branches, - "If true, simplify each branch of AND/OR under constraints provided by the other " - "branch", - refl::DefaultValue(false)); +RewriteSimplifier::Extension StmtSimplifyConfigNode::GetEnabledExtensions() const { + RewriteSimplifier::Extension flags = RewriteSimplifier::kNone; + if (transitively_prove_inequalities) { + flags = RewriteSimplifier::Extension(flags | RewriteSimplifier::kTransitivelyProveInequalities); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.transform.StmtSimplifyConfig", StmtSimplifyConfigNode, - ffi::Object); - - RewriteSimplifier::Extension GetEnabledExtensions() const { - RewriteSimplifier::Extension flags = RewriteSimplifier::kNone; - if (transitively_prove_inequalities) { - flags = - RewriteSimplifier::Extension(flags | RewriteSimplifier::kTransitivelyProveInequalities); - } - if (convert_boolean_to_and_of_ors) { - flags = RewriteSimplifier::Extension(flags | RewriteSimplifier::kConvertBooleanToAndOfOrs); - } - if (apply_constraints_to_boolean_branches) { - flags = RewriteSimplifier::Extension(flags | - RewriteSimplifier::kApplyConstraintsToBooleanBranches); - } - return flags; + if (convert_boolean_to_and_of_ors) { + flags = RewriteSimplifier::Extension(flags | RewriteSimplifier::kConvertBooleanToAndOfOrs); } -}; - -class StmtSimplifyConfig : public ffi::ObjectRef { - public: - TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(StmtSimplifyConfig, ffi::ObjectRef, - StmtSimplifyConfigNode); -}; + if (apply_constraints_to_boolean_branches) { + flags = + RewriteSimplifier::Extension(flags | RewriteSimplifier::kApplyConstraintsToBooleanBranches); + } + return flags; +} static StmtSimplifyConfig MakeDefaultStmtSimplifyConfig() { return tvm::transform::PassConfigWithDefaults(); @@ -99,155 +84,136 @@ TVM_FFI_STATIC_INIT_BLOCK() { StmtSimplifyConfigNode::RegisterReflection(); } TVM_REGISTER_PASS_CONFIG_OPTION("tirx.StmtSimplify", StmtSimplifyConfig); -class StmtSimplifier : public IRMutatorWithAnalyzer { - public: - using IRMutatorWithAnalyzer::Mutate; - using IRMutatorWithAnalyzer::Mutate_; - static PrimFunc Apply(PrimFunc func, const Analyzer& analyzer, - ffi::Optional config_opt = std::nullopt) { - auto config = config_opt.value_or(MakeDefaultStmtSimplifyConfig()); - analyzer->rewrite_simplify.SetEnabledExtensions(config->GetEnabledExtensions()); - - auto simplifier = ffi::make_object(analyzer, config); - simplifier->MarkBufferParamShapes(func); - auto* n = func.CopyOnWrite(); - n->body = simplifier->Mutate(n->body, InplaceMode::kAllow).ValueOrUnchanged(n->body); - return func; - } +PrimFunc StmtSimplifier::Apply(PrimFunc func, const Analyzer& analyzer, + ffi::Optional config_opt) { + auto config = config_opt.value_or(MakeDefaultStmtSimplifyConfig()); - public: - explicit StmtSimplifier(const Analyzer& analyzer, StmtSimplifyConfig config) - : IRMutatorWithAnalyzer(analyzer), config_(config) {} + auto simplifier = ffi::make_object(analyzer, config); + return simplifier->Run(std::move(func)); +} - private: - using Parent = IRMutatorWithAnalyzer; +PrimFunc StmtSimplifier::Run(PrimFunc func) { + analyzer_->rewrite_simplify.SetEnabledExtensions(config_->GetEnabledExtensions()); + MarkBufferParamShapes(func); + auto* n = func.CopyOnWrite(); + n->body = Mutate(n->body, InplaceMode::kAllow).ValueOrUnchanged(n->body); + return func; +} - UnchangedOr Mutate(ffi::AnyView input, InplaceMode inplace_mode) final { - if (input.as()) { - return ffi::Unchanged(); - } - if (auto expr = input.as()) { - PrimExpr simplified = analyzer_->Simplify(*expr); - if (simplified.same_as(*expr)) return ffi::Unchanged(); - return simplified; - } - return Parent::Mutate(input, inplace_mode); +UnchangedOr StmtSimplifier::Mutate(ffi::AnyView input, InplaceMode inplace_mode) { + if (input.as()) { + return ffi::Unchanged(); } + if (auto expr = input.as()) { + PrimExpr simplified = analyzer_->Simplify(*expr); + if (simplified.same_as(*expr)) return ffi::Unchanged(); + return simplified; + } + return Parent::Mutate(input, inplace_mode); +} - Stmt Simplify(Stmt stmt) { return Mutate(stmt, InplaceMode::kAllow).ValueOrUnchanged(stmt); } +UnchangedOr StmtSimplifier::Mutate_(const ForNode* op, InplaceMode inplace_mode) { + analyzer_->Bind(op->loop_var, Range::FromMinExtent(op->min, op->extent)); + With ctx1(analyzer_, op->loop_var >= op->min); + With ctx2(analyzer_, + static_cast(op->loop_var) < op->min + op->extent); + return Parent::Mutate_(op, inplace_mode); +} - UnchangedOr Mutate_(const ForNode* op, InplaceMode inplace_mode) final { - analyzer_->Bind(op->loop_var, Range::FromMinExtent(op->min, op->extent)); - With ctx1(analyzer_, op->loop_var >= op->min); - With ctx2(analyzer_, - static_cast(op->loop_var) < op->min + op->extent); +UnchangedOr StmtSimplifier::Mutate_(const BindNode* op, InplaceMode inplace_mode) { + auto prim_value = op->value.as(); + if (!prim_value) { return Parent::Mutate_(op, inplace_mode); } + PrimExpr value = + this->Mutate(prim_value.value(), inplace_mode).ValueOrUnchanged(prim_value.value()); + // Bind in analyzer for constraint proving and simplification of + // subsequent expressions. Don't remove the Bind statement -- + // with flat Bind there's no body to inspect for usage patterns, + // so we always keep the Bind. + if (SideEffect(value) <= CallEffectKind::kPure) { + analyzer_->Bind(op->var, value); + // Record the binding so we can substitute it into assert conditions + // (see Mutate_(const AssertStmtNode*, InplaceMode)). Under SSA each var is + // bound exactly once, so the map grows monotonically without key + // conflicts. No scope-based cleanup is needed because vars bound + // in inner scopes are only referenced within those scopes; stale + // entries are harmless and never consulted again. + non_inlined_bindings_.Set(op->var, value); + } - UnchangedOr Mutate_(const BindNode* op, InplaceMode inplace_mode) override { - auto prim_value = op->value.as(); - if (!prim_value) { - return Parent::Mutate_(op, inplace_mode); - } - PrimExpr value = - this->Mutate(prim_value.value(), inplace_mode).ValueOrUnchanged(prim_value.value()); - // Bind in analyzer for constraint proving and simplification of - // subsequent expressions. Don't remove the Bind statement -- - // with flat Bind there's no body to inspect for usage patterns, - // so we always keep the Bind. - if (SideEffect(value) <= CallEffectKind::kPure) { - analyzer_->Bind(op->var, value); - // Record the binding so we can substitute it into assert conditions - // (see Mutate_(const AssertStmtNode*, InplaceMode)). Under SSA each var is - // bound exactly once, so the map grows monotonically without key - // conflicts. No scope-based cleanup is needed because vars bound - // in inner scopes are only referenced within those scopes; stale - // entries are harmless and never consulted again. - non_inlined_bindings_.Set(op->var, value); - } - - if (value.same_as(op->value)) { - return ffi::Unchanged(); - } else { - if (inplace_mode == InplaceMode::kAllow) { - auto* n = const_cast(op); - n->value = std::move(value); - return ffi::Unchanged(); - } - auto n = ffi::make_object(*op); + if (value.same_as(op->value)) { + return ffi::Unchanged(); + } else { + if (inplace_mode == InplaceMode::kAllow) { + auto* n = const_cast(op); n->value = std::move(value); - return Stmt(n); + return ffi::Unchanged(); } + auto n = ffi::make_object(*op); + n->value = std::move(value); + return Stmt(n); } +} - UnchangedOr Mutate_(const IfThenElseNode* op, InplaceMode inplace_mode) override { - if (ffi::Optional cond = ProveCondition(op->condition)) { - if (cond.value()) { - return this->Mutate(op->then_case, inplace_mode).ValueOrUnchanged(op->then_case); - } else if (op->else_case) { - return this->Mutate(op->else_case.value(), inplace_mode) - .ValueOrUnchanged(op->else_case.value()); - } else { - return Evaluate(0); - } +UnchangedOr StmtSimplifier::Mutate_(const IfThenElseNode* op, InplaceMode inplace_mode) { + if (ffi::Optional cond = ProveCondition(op->condition)) { + if (cond.value()) { + return this->Mutate(op->then_case, inplace_mode).ValueOrUnchanged(op->then_case); + } else if (op->else_case) { + return this->Mutate(op->else_case.value(), inplace_mode) + .ValueOrUnchanged(op->else_case.value()); } else { - return Parent::Mutate_(op, inplace_mode); + return Evaluate(0); } + } else { + return Parent::Mutate_(op, inplace_mode); } +} - // eliminate useless stores - UnchangedOr Mutate_(const BufferStoreNode* op, InplaceMode inplace_mode) override { - BufferStore store = Parent::Mutate_(op, inplace_mode) - .ValueOrUnchanged(ffi::GetRef(op)) - .as_or_throw(); - if (const TensorLoadNode* load = store->value.as()) { - BufferVar buffer = load->source.as_or_throw(); - if (buffer.same_as(store->buffer) && ArrayDeepEqual(load->indices, store->indices) && - prim::ExprDeepEqual()(buffer->elem_offset, store->buffer->elem_offset) && - ArrayDeepEqual(buffer->shape, store->buffer->shape) && - ArrayDeepEqual(buffer->strides, store->buffer->strides)) { - return Evaluate(0); - } +UnchangedOr StmtSimplifier::Mutate_(const BufferStoreNode* op, InplaceMode inplace_mode) { + BufferStore store = Parent::Mutate_(op, inplace_mode) + .ValueOrUnchanged(ffi::GetRef(op)) + .as_or_throw(); + if (const TensorLoadNode* load = store->value.as()) { + BufferVar buffer = load->source.as_or_throw(); + if (buffer.same_as(store->buffer) && ArrayDeepEqual(load->indices, store->indices) && + prim::ExprDeepEqual()(buffer->elem_offset, store->buffer->elem_offset) && + ArrayDeepEqual(buffer->shape, store->buffer->shape) && + ArrayDeepEqual(buffer->strides, store->buffer->strides)) { + return Evaluate(0); } - return store; } - bool ArrayDeepEqual(const ffi::Array& lhs, const ffi::Array& rhs) { - if (lhs.size() != rhs.size()) { + return store; +} + +bool StmtSimplifier::ArrayDeepEqual(const ffi::Array& lhs, + const ffi::Array& rhs) { + if (lhs.size() != rhs.size()) { + return false; + } + for (size_t i = 0; i < lhs.size(); i++) { + if (!prim::ExprDeepEqual()(lhs[i], rhs[i])) { return false; } - for (size_t i = 0; i < lhs.size(); i++) { - if (!prim::ExprDeepEqual()(lhs[i], rhs[i])) { - return false; - } - } - return true; } + return true; +} - /* \brief Internal utility for checking conditionals - * - * Substitutes any known Bind values and then simplifies with the analyzer. - */ - ffi::Optional ProveCondition(PrimExpr condition) const { - auto f_substitute = [this](const Var& var) -> ffi::Expected> { - if (auto repl = non_inlined_bindings_.Get(var)) return ffi::Any(*std::move(repl)); - return ffi::Unchanged(); - }; - condition = ffi::StructuralMap(condition, f_substitute) - .as_or_throw(); - condition = analyzer_->Simplify(condition); - if (const auto* as_int = condition.as()) { - return as_int->value != 0; - } else { - return std::nullopt; - } +ffi::Optional StmtSimplifier::ProveCondition(PrimExpr condition) const { + auto f_substitute = [this](const Var& var) -> ffi::Expected> { + if (auto repl = non_inlined_bindings_.Get(var)) return ffi::Any(*std::move(repl)); + return ffi::Unchanged(); + }; + condition = ffi::StructuralMap(condition, f_substitute) + .as_or_throw(); + condition = analyzer_->Simplify(condition); + if (const auto* as_int = condition.as()) { + return as_int->value != 0; + } else { + return std::nullopt; } - - StmtSimplifyConfig config_; - - // Pure Bind values kept for substitution into assert conditions. - // Grows monotonically under SSA — no scope-based cleanup required. - ffi::Map non_inlined_bindings_; -}; +} } // namespace arith diff --git a/src/tirx/transform/stmt_simplify.h b/src/tirx/transform/stmt_simplify.h index 224df0ed8bfc..83e04ff97ff1 100644 --- a/src/tirx/transform/stmt_simplify.h +++ b/src/tirx/transform/stmt_simplify.h @@ -27,7 +27,75 @@ #include #include +#include "../ir_mutator_with_analyzer.h" + namespace tvm { +namespace arith { +using namespace tirx; + +struct StmtSimplifyConfigNode : public ffi::Object { + bool transitively_prove_inequalities; + bool convert_boolean_to_and_of_ors; + bool apply_constraints_to_boolean_branches; + + static void RegisterReflection(); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.transform.StmtSimplifyConfig", StmtSimplifyConfigNode, + ffi::Object); + + RewriteSimplifier::Extension GetEnabledExtensions() const; +}; + +class StmtSimplifyConfig : public ffi::ObjectRef { + public: + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(StmtSimplifyConfig, ffi::ObjectRef, + StmtSimplifyConfigNode); +}; + +class StmtSimplifier : public IRMutatorWithAnalyzer { + public: + using IRMutatorWithAnalyzer::Mutate; + using IRMutatorWithAnalyzer::Mutate_; + static PrimFunc Apply(PrimFunc func, const Analyzer& analyzer, + ffi::Optional config_opt = std::nullopt); + + explicit StmtSimplifier(const Analyzer& analyzer, StmtSimplifyConfig config) + : IRMutatorWithAnalyzer(analyzer), config_(config) {} + + protected: + using Parent = IRMutatorWithAnalyzer; + StmtSimplifier(const VTable* vtable, const Analyzer& analyzer, StmtSimplifyConfig config) + : Parent(analyzer.get(), vtable), config_(config) {} + PrimFunc Run(PrimFunc func); + + UnchangedOr Mutate(ffi::AnyView input, InplaceMode inplace_mode) final; + + UnchangedOr Mutate_(const ForNode* op, InplaceMode inplace_mode) final; + + UnchangedOr Mutate_(const BindNode* op, InplaceMode inplace_mode) override; + + UnchangedOr Mutate_(const IfThenElseNode* op, InplaceMode inplace_mode) override; + + // eliminate useless stores + UnchangedOr Mutate_(const BufferStoreNode* op, InplaceMode inplace_mode) override; + + private: + bool ArrayDeepEqual(const ffi::Array& lhs, const ffi::Array& rhs); + + /* \brief Internal utility for checking conditionals + * + * Substitutes any known Bind values and then simplifies with the analyzer. + */ + ffi::Optional ProveCondition(PrimExpr condition) const; + + StmtSimplifyConfig config_; + + // Pure Bind values kept for substitution into assert conditions. + // Grows monotonically under SSA — no scope-based cleanup required. + ffi::Map non_inlined_bindings_; +}; + +} // namespace arith + namespace tirx { /* \brief Simplify statements in the prim func diff --git a/src/tirx/transform/tvm_ffi_binder.cc b/src/tirx/transform/tvm_ffi_binder.cc index 837e34af7302..10a05d4ccab5 100644 --- a/src/tirx/transform/tvm_ffi_binder.cc +++ b/src/tirx/transform/tvm_ffi_binder.cc @@ -37,6 +37,7 @@ namespace tvm { namespace tirx { using namespace tvm::prim; +using tvm::prim::cast; using ffi::reflection::AccessPath; using ffi::reflection::AccessStep; diff --git a/tests/cpp/s_tir_functor_test.cc b/tests/cpp/s_tir_functor_test.cc index 9f9e5734e3db..653975ada1d9 100644 --- a/tests/cpp/s_tir_functor_test.cc +++ b/tests/cpp/s_tir_functor_test.cc @@ -50,6 +50,19 @@ TEST(STIRFunctor, LegacyInheritedDispatchAndContainsNode) { EXPECT_FALSE(ContainsNode(realize)); } +TEST(STIRFunctor, CoreLegacyDispatchReachesDefaultForDialectNodes) { + class Dispatch : public tirx::StmtFunctor { + public: + using tirx::StmtFunctor::VisitStmt_; + bool VisitStmt_(const EvaluateNode*) final { return true; } + bool VisitStmtDefault_(const ffi::Object*) final { return false; } + } dispatch; + SBlock block({}, {}, {}, "block", Evaluate(0)); + EXPECT_FALSE(dispatch(block)); + EXPECT_FALSE(dispatch(SBlockRealize({}, IntImm::Bool(true), block))); + EXPECT_TRUE(dispatch(block->body)); +} + TEST(STIRFunctor, NativeBlockOverrideReusesInheritedCoreHooks) { class Visitor : public StmtExprVisitor { public: @@ -154,8 +167,10 @@ void CheckMutationRemapsBufferDefinitionsAndUses() { BufferVar matched = decl_buffer({extent + 1}, PrimType::Int(32)); BufferRegion region(allocated, {Range::FromMinExtent(0, extent + 1)}); MatchBufferRegion match(matched, region); + BufferRegion matched_region(matched, {Range::FromMinExtent(0, extent + 1)}); Stmt body = SeqStmt({BufferStore(allocated, 0, {0}), BufferStore(matched, 0, {0})}); - SBlock block({}, {region}, {region}, "block", body, std::nullopt, {allocated}, {match}); + SBlock block({}, {region, matched_region}, {region, matched_region}, "block", body, std::nullopt, + {allocated}, {match}); class Mutator : public Base { public: using Base::Mutate_; @@ -173,6 +188,8 @@ void CheckMutationRemapsBufferDefinitionsAndUses() { EXPECT_TRUE(new_matched->shape[0].same_as(extent)); EXPECT_TRUE(changed->reads[0]->buffer.same_as(new_allocated)); EXPECT_TRUE(changed->writes[0]->buffer.same_as(new_allocated)); + EXPECT_TRUE(changed->reads[1]->buffer.same_as(new_matched)); + EXPECT_TRUE(changed->writes[1]->buffer.same_as(new_matched)); EXPECT_TRUE(changed->match_buffers[0]->source->buffer.same_as(new_allocated)); const auto* statements = changed->body.as(); ASSERT_NE(statements, nullptr); @@ -190,7 +207,104 @@ TEST(STIRFunctor, GenericTIRXMutationRemapsBufferDefinitionsAndUses) { CheckMutationRemapsBufferDefinitionsAndUses(); } -TEST(STIRFunctor, GenericTIRXVisitorUsesRegisteredNativePolicy) { +TEST(STIRFunctor, StructuralAndGenericSubstitutionPreserveDefinitionUses) { + PrimVar extent("extent"), new_extent("new_extent"), index("index"), new_index("new_index"); + BufferVar allocated = decl_buffer({extent}, PrimType::Int(32)); + BufferVar matched = decl_buffer({extent}, PrimType::Int(32)); + BufferRegion region(allocated, {Range::FromMinExtent(0, extent)}); + BufferRegion matched_region(matched, {Range::FromMinExtent(0, extent)}); + MatchBufferRegion match(matched, region); + IterVar iter(Range::FromMinExtent(0, extent), index, IterVarType::kDataPar); + Stmt body = + SeqStmt({BufferStore(allocated, extent, {index}), BufferStore(matched, extent, {index})}); + SBlock block({iter}, {region, matched_region}, {region, matched_region}, "block", body, + Evaluate(extent), {allocated}, {match}, {{"annotation", extent}}); + Stmt original = SBlockRealize({extent}, extent > 0, block); + + auto check = [&](const Stmt& result) { + const auto* realize = result.as(); + ASSERT_NE(realize, nullptr); + EXPECT_TRUE(realize->iter_values[0].same_as(new_extent)); + EXPECT_TRUE(realize->predicate.as()->a.same_as(new_extent)); + const auto* changed = realize->block.get(); + EXPECT_TRUE(changed->iter_vars[0]->var.same_as(new_index)); + EXPECT_TRUE(changed->iter_vars[0]->dom->extent.same_as(new_extent)); + EXPECT_TRUE(changed->alloc_buffers[0]->shape[0].same_as(new_extent)); + EXPECT_TRUE(changed->match_buffers[0]->buffer->shape[0].same_as(new_extent)); + EXPECT_TRUE(changed->match_buffers[0]->source->buffer.same_as(changed->alloc_buffers[0])); + EXPECT_TRUE(changed->reads[0]->buffer.same_as(changed->alloc_buffers[0])); + EXPECT_TRUE(changed->reads[1]->buffer.same_as(changed->match_buffers[0]->buffer)); + EXPECT_TRUE(changed->writes[0]->buffer.same_as(changed->alloc_buffers[0])); + EXPECT_TRUE(changed->writes[1]->buffer.same_as(changed->match_buffers[0]->buffer)); + EXPECT_TRUE(changed->reads[0]->region[0]->extent.same_as(new_extent)); + EXPECT_TRUE(changed->reads[1]->region[0]->extent.same_as(new_extent)); + const auto* statements = changed->body.as(); + ASSERT_NE(statements, nullptr); + const auto* store = statements->seq[0].as(); + EXPECT_TRUE(store->buffer.same_as(changed->alloc_buffers[0])); + EXPECT_TRUE(store->value.same_as(new_extent)); + EXPECT_TRUE(store->indices[0].same_as(new_index)); + EXPECT_TRUE(statements->seq[1].as()->buffer.same_as( + changed->match_buffers[0]->buffer)); + EXPECT_TRUE(changed->init.value().as()->value.same_as(new_extent)); + EXPECT_TRUE(changed->annotations.at("annotation").cast().same_as(new_extent)); + EXPECT_TRUE(block->iter_vars[0]->var.same_as(index)); + EXPECT_TRUE(block->iter_vars[0]->dom->extent.same_as(extent)); + EXPECT_TRUE(block->alloc_buffers[0].same_as(allocated)); + EXPECT_TRUE(block->match_buffers[0]->buffer.same_as(matched)); + EXPECT_TRUE(block->annotations.at("annotation").cast().same_as(extent)); + }; + + Stmt structural = ffi::StructuralMap( + original, + [&](const Var& var) -> ffi::Expected> { + if (var.same_as(extent)) return ffi::Any(new_extent); + if (var.same_as(index)) return ffi::Any(new_index); + return ffi::Unchanged(); + }) + .as_or_throw(); + check(structural); + Stmt generic = + SubstituteWithDataTypeLegalization(original, [&](const Var& var) -> ffi::Optional { + if (var.same_as(extent)) return new_extent; + if (var.same_as(index)) return new_index; + return std::nullopt; + }); + check(generic); + + // Unique outer nodes may update in place, but their shared child arrays and + // regions must not modify the retained block used to construct them. + for (bool use_generic : {false, true}) { + SBlock local(block->iter_vars, block->reads, block->writes, "unique", block->body, block->init, + block->alloc_buffers, block->match_buffers, block->annotations); + const auto* block_identity = local.get(); + Stmt input = SBlockRealize({extent}, extent > 0, std::move(local)); + const auto* realize_identity = input.get(); + Stmt result; + if (use_generic) { + result = SubstituteWithDataTypeLegalization(std::move(input), + [&](const Var& var) -> ffi::Optional { + if (var.same_as(extent)) return new_extent; + if (var.same_as(index)) return new_index; + return std::nullopt; + }); + } else { + result = ffi::StructuralMap( + std::move(input), + [&](const Var& var) -> ffi::Expected> { + if (var.same_as(extent)) return ffi::Any(new_extent); + if (var.same_as(index)) return ffi::Any(new_index); + return ffi::Unchanged(); + }) + .as_or_throw(); + } + EXPECT_EQ(result.get(), realize_identity); + EXPECT_EQ(result.as()->block.get(), block_identity); + check(result); + } +} + +TEST(STIRFunctor, GenericTIRXVisitorUsesFullStructuralTraversal) { PrimVar index("index"), annotation("annotation"); BufferVar buffer = decl_buffer({16}); BufferRegion region(buffer, {Range::FromMinExtent(0, 16)}); @@ -212,14 +326,14 @@ TEST(STIRFunctor, GenericTIRXVisitorUsesRegisteredNativePolicy) { }; auto visitor = ffi::make_object(); visitor->Visit(block); - EXPECT_EQ(std::count(visitor->vars.begin(), visitor->vars.end(), index.get()), 1); - EXPECT_EQ(std::count(visitor->vars.begin(), visitor->vars.end(), annotation.get()), 0); + EXPECT_EQ(std::count(visitor->vars.begin(), visitor->vars.end(), index.get()), 2); + EXPECT_EQ(std::count(visitor->vars.begin(), visitor->vars.end(), annotation.get()), 1); ASSERT_EQ(visitor->buffer_regions.size(), 2); EXPECT_EQ(visitor->buffer_regions[0], kTVMFFIDefRegionKindSimple); EXPECT_EQ(visitor->buffer_regions[1], kTVMFFIDefRegionKindNone); } -TEST(STIRFunctor, GenericTIRXMutationPreservesBindersAndAnnotations) { +TEST(STIRFunctor, GenericTIRXMutationRemapsBindersAndAnnotations) { PrimVar index("index"), replacement("replacement"), extent("extent"); PrimExpr expression = extent + 1; IterVar iter(Range::FromMinExtent(0, expression), index, IterVarType::kDataPar); @@ -237,9 +351,9 @@ TEST(STIRFunctor, GenericTIRXMutationPreservesBindersAndAnnotations) { Stmt result = mutator->Mutate(block, InplaceMode::kAllow).ValueOrUnchanged(block); const auto* changed = result.as(); ASSERT_NE(changed, nullptr); - EXPECT_TRUE(changed->iter_vars[0]->var.same_as(index)); + EXPECT_TRUE(changed->iter_vars[0]->var.same_as(replacement)); EXPECT_TRUE(changed->iter_vars[0]->dom->extent.same_as(extent)); - EXPECT_TRUE(changed->annotations.at("annotation").cast().same_as(expression)); + EXPECT_TRUE(changed->annotations.at("annotation").cast().same_as(extent)); EXPECT_TRUE(changed->body.as()->value.same_as(replacement)); EXPECT_TRUE(changed->body.as()->buffer.same_as(changed->alloc_buffers[0])); EXPECT_TRUE(changed->alloc_buffers[0]->shape[0].same_as(extent)); @@ -253,10 +367,10 @@ TEST(STIRFunctor, GenericTIRXMutationPreservesBindersAndAnnotations) { EXPECT_TRUE(update.IsUnchanged()); EXPECT_EQ(unique.get(), original); EXPECT_TRUE(unique->body.as()->value.same_as(extent)); - EXPECT_TRUE(unique->annotations.at("annotation").cast().same_as(expression)); + EXPECT_TRUE(unique->annotations.at("annotation").cast().same_as(extent)); } -TEST(STIRFunctor, GenericTIRXPolicyPreservesInterruptAndErrorIdentity) { +TEST(STIRFunctor, GenericTIRXFallbackPreservesInterruptAndErrorIdentity) { PrimVar annotation("annotation"), body("body"); SBlock block({}, {}, {}, "block", Evaluate(body), std::nullopt, {}, {}, {{"annotation", annotation}}); @@ -272,7 +386,7 @@ TEST(STIRFunctor, GenericTIRXPolicyPreservesInterruptAndErrorIdentity) { auto visitor = ffi::make_object(); auto interrupt = visitor->Visit(block); ASSERT_TRUE(interrupt.has_value()); - EXPECT_TRUE(interrupt.value()->value.cast().same_as(body)); + EXPECT_TRUE(interrupt.value()->value.cast().same_as(annotation)); EXPECT_EQ(visitor->count, 1); class Mutator : public tirx::StmtExprMutator { diff --git a/tests/python/relax/test_transform_fuse_ops.py b/tests/python/relax/test_transform_fuse_ops.py index bee21e0686f6..d9d0d128c221 100644 --- a/tests/python/relax/test_transform_fuse_ops.py +++ b/tests/python/relax/test_transform_fuse_ops.py @@ -1415,7 +1415,7 @@ def main( for global_var in mod.get_global_vars() if global_var.name_hint.startswith("fused_") ) - assert tvm.tirx.analysis.verify_well_formed(fused_tir) + assert tvm.s_tir.analysis.verify_well_formed(fused_tir) def test_symbolic_prim_arg_before_tensor_arg(): @@ -1481,7 +1481,7 @@ def main( for global_var in mod.get_global_vars() if global_var.name_hint.startswith("fused_") ) - assert tvm.tirx.analysis.verify_well_formed(fused_tir) + assert tvm.s_tir.analysis.verify_well_formed(fused_tir) def test_symbolic_prim_arg_reused_from_derived_tensor_shape(): @@ -1562,7 +1562,7 @@ def main( for global_var in mod.get_global_vars() if global_var.name_hint.startswith("fused_") ) - assert tvm.tirx.analysis.verify_well_formed(fused_tir) + assert tvm.s_tir.analysis.verify_well_formed(fused_tir) def test_symbolic_prim_arg_not_bound_by_derived_tensor_shape(): @@ -1683,7 +1683,7 @@ def main(x: R.Tensor((4,), dtype="int64")): for global_var in mod.get_global_vars() if global_var.name_hint.startswith("fused_") ) - assert tvm.tirx.analysis.verify_well_formed(fused_tir) + assert tvm.s_tir.analysis.verify_well_formed(fused_tir) def test_primitive_call_arg_used_by_output_shape_not_inlined(): @@ -1821,7 +1821,7 @@ def main( for global_var in mod.get_global_vars() if global_var.name_hint.startswith("fused_") ) - assert tvm.tirx.analysis.verify_well_formed(fused_tir) + assert tvm.s_tir.analysis.verify_well_formed(fused_tir) def test_shape_expr_arg(): diff --git a/tests/python/relax/test_transform_fuse_tir.py b/tests/python/relax/test_transform_fuse_tir.py index 5108116abca9..85200527515b 100644 --- a/tests/python/relax/test_transform_fuse_tir.py +++ b/tests/python/relax/test_transform_fuse_tir.py @@ -2415,7 +2415,7 @@ def main(x: R.Tensor((4,), "int64"), p: R.Prim("int64")) -> R.Tensor((1,), "int6 after = relax.transform.FuseTIR()(Before) assert relax.analysis.check_well_formed(after) - assert tvm.tirx.analysis.verify_well_formed(after["fused"]) + assert tvm.s_tir.analysis.verify_well_formed(after["fused"]) def test_inplace_argument_after_primitive_scalar(): @@ -2448,7 +2448,7 @@ def main(p: R.Prim("int64"), x: R.Tensor((4,), "int64")) -> R.Tensor((4,), "int6 after = relax.transform.FuseTIR()(Before) assert relax.analysis.check_well_formed(after) - assert tvm.tirx.analysis.verify_well_formed(after["fused"]) + assert tvm.s_tir.analysis.verify_well_formed(after["fused"]) if __name__ == "__main__": diff --git a/tests/python/relax/test_transform_specialize_primfunc_based_on_callsite.py b/tests/python/relax/test_transform_specialize_primfunc_based_on_callsite.py index b9c9b16e1ea9..4b2867d0714c 100644 --- a/tests/python/relax/test_transform_specialize_primfunc_based_on_callsite.py +++ b/tests/python/relax/test_transform_specialize_primfunc_based_on_callsite.py @@ -80,6 +80,7 @@ def verify(input): ValidateBufferScopes(False).visit(input) mod = tvm.relax.transform.SpecializePrimFuncBasedOnCallSite()(input) ValidateBufferScopes(True).visit(mod) + return mod def test_single_arg_return(): @@ -208,7 +209,11 @@ def main( R.output(gv2) return gv2 - verify(Input) + specialized = verify(Input) + # This pass runs before DLight: specialized blocks must remain schedulable. + schedule = tvm.s_tir.Schedule(specialized, debug_mask="all") + block = schedule.get_sblock("pool_max", func_name="max_pool2d_opencl") + assert len(schedule.get_loops(block)) == 7 def test_multi_arg_return(): diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_rewrite_parallel_vectorize_unroll.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_rewrite_parallel_vectorize_unroll.py index 5152e2fa15f0..bef78431ef07 100644 --- a/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_rewrite_parallel_vectorize_unroll.py +++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_rewrite_parallel_vectorize_unroll.py @@ -207,7 +207,7 @@ def test_meta_schedule_postproc_rewrite_parallel_unroll_vectorize(): postproc = RewriteParallelVectorizeUnroll() sch = Schedule(Move_PUV) assert postproc.apply(sch) - mod = tvm.tirx.transform.StmtSimplify()(sch.mod) + mod = tvm.s_tir.transform.StmtSimplify()(sch.mod) tvm.ir.assert_structural_equal(mod["main"], Move_PUV0) @@ -283,7 +283,7 @@ def expected(A: T.Buffer((1, 4, 4, 32), "float32"), B: T.Buffer((4, 4, 32), "flo postproc = RewriteParallelVectorizeUnroll() sch = Schedule(layer_norm) assert postproc.apply(sch) - mod = tvm.tirx.transform.StmtSimplify()(sch.mod) + mod = tvm.s_tir.transform.StmtSimplify()(sch.mod) assert_structural_equal_ignore_global_symbol(mod["main"], expected) diff --git a/tests/python/s_tir/transform/test_s_tir_transform_compact_buffer_region.py b/tests/python/s_tir/transform/test_s_tir_transform_compact_buffer_region.py index f0000581e1f1..d42ab45ea397 100644 --- a/tests/python/s_tir/transform/test_s_tir_transform_compact_buffer_region.py +++ b/tests/python/s_tir/transform/test_s_tir_transform_compact_buffer_region.py @@ -37,7 +37,7 @@ def test_compact(self): before = tvm.IRModule.from_expr(self.before.with_attr("global_symbol", "main")) expected = tvm.IRModule.from_expr(self.expected.with_attr("global_symbol", "main")) simplify = tvm.transform.Sequential( - [tirx.transform.StmtSimplify(), tirx.transform.RemoveNoOp()] + [s_tir.transform.StmtSimplify(), tirx.transform.RemoveNoOp()] ) after = simplify(s_tir.transform.CompactBufferAllocation(is_strict=is_strict)(before)) expected = simplify(expected) @@ -1300,7 +1300,7 @@ def before(a: T.handle): tmp[j] = A[j] after = s_tir.transform.CompactBufferAllocation()(tvm.IRModule.from_expr(before)) - assert tirx.analysis.verify_well_formed(after) + assert s_tir.analysis.verify_well_formed(after) class TestCompactSymbolicBound0: diff --git a/tests/python/s_tir/transform/test_s_tir_transform_convert_blocks_to_opaque.py b/tests/python/s_tir/transform/test_s_tir_transform_convert_blocks_to_opaque.py index ff93ad3319b1..b4e2e0e1e0d1 100644 --- a/tests/python/s_tir/transform/test_s_tir_transform_convert_blocks_to_opaque.py +++ b/tests/python/s_tir/transform/test_s_tir_transform_convert_blocks_to_opaque.py @@ -28,7 +28,7 @@ def _check(original, transformed): func = original mod = tvm.IRModule.from_expr(func.with_attr("global_symbol", "main")) mod = tvm.s_tir.transform.ConvertBlocksToOpaque()(mod) - mod = tvm.tirx.transform.StmtSimplify()(mod) + mod = tvm.s_tir.transform.StmtSimplify()(mod) tvm.ir.assert_structural_equal(mod["main"], transformed.with_attr("global_symbol", "main")) diff --git a/tests/python/s_tir/transform/test_s_tir_transform_convert_ssa.py b/tests/python/s_tir/transform/test_s_tir_transform_convert_ssa.py new file mode 100644 index 000000000000..86aed688eb66 --- /dev/null +++ b/tests/python/s_tir/transform/test_s_tir_transform_convert_ssa.py @@ -0,0 +1,63 @@ +# 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. + +import tvm +import tvm.testing +from tvm import s_tir, tirx + + +def test_reused_block_iterator(): + """A shared block defines a fresh iterator at each realization.""" + var = tirx.Var("v", "int32") + iterator = tirx.IterVar(tvm.ir.Range(0, 4), var, tirx.IterVar.DataPar) + block = s_tir.SBlock([iterator], [], [], "block", tirx.Evaluate(var)) + realize = s_tir.SBlockRealize([0], True, block) + before = tirx.PrimFunc([], tirx.SeqStmt([realize, realize])) + + after = s_tir.transform.ConvertSSA()(tvm.IRModule.from_expr(before))["main"] + + first, second = [realize.block for realize in after.body.seq] + assert not first.iter_vars[0].var.same_as(second.iter_vars[0].var) + assert first.body.value.same_as(first.iter_vars[0].var) + assert second.body.value.same_as(second.iter_vars[0].var) + # Shared input ownership must protect both original occurrences. + assert block.iter_vars[0].var.same_as(var) + assert block.body.value.same_as(var) + + +def test_shared_buffer_parameter_regions_across_functions(): + """Parameter renaming reaches both block regions and buffer accesses.""" + n = tirx.Var("n", "int32") + buffer = tirx.decl_buffer((n,), "float32", "buffer") + region = tirx.BufferRegion(buffer, [tvm.ir.Range(0, n)]) + block = s_tir.SBlock([], [region], [], "root", tirx.Evaluate(tirx.BufferLoad(buffer, [0]))) + func = tirx.PrimFunc([buffer], s_tir.SBlockRealize([], True, block)) + before = tvm.IRModule({"first": func, "second": func}) + + after = s_tir.transform.ConvertSSA()(before) + + first, second = after["first"], after["second"] + assert not first.params[0].same_as(second.params[0]) + for updated in [first, second]: + updated_block = updated.body.block + assert updated_block.reads[0].buffer.same_as(updated.params[0]) + assert updated_block.body.value.source.same_as(updated.params[0]) + assert updated_block.reads[0].region[0].extent.same_as(updated.params[0].ty.shape[0]) + + +if __name__ == "__main__": + tvm.testing.main() diff --git a/tests/python/s_tir/transform/test_s_tir_transform_inject_software_pipeline.py b/tests/python/s_tir/transform/test_s_tir_transform_inject_software_pipeline.py index 56633f3eb144..5898899a35ff 100644 --- a/tests/python/s_tir/transform/test_s_tir_transform_inject_software_pipeline.py +++ b/tests/python/s_tir/transform/test_s_tir_transform_inject_software_pipeline.py @@ -42,7 +42,7 @@ def _check(original, transformed): func = original mod = tvm.IRModule.from_expr(func.with_attr("global_symbol", "main")) mod = tvm.s_tir.transform.InjectSoftwarePipeline()(mod) - mod = tvm.tirx.transform.StmtSimplify()(mod) + mod = tvm.s_tir.transform.StmtSimplify()(mod) tvm.ir.assert_structural_equal( mod["main"], transformed.with_attr("global_symbol", "main"), True ) diff --git a/tests/python/s_tir/transform/test_s_tir_transform_lower_match_buffer.py b/tests/python/s_tir/transform/test_s_tir_transform_lower_match_buffer.py index 2f5184c69f57..5152a51dcf37 100644 --- a/tests/python/s_tir/transform/test_s_tir_transform_lower_match_buffer.py +++ b/tests/python/s_tir/transform/test_s_tir_transform_lower_match_buffer.py @@ -26,7 +26,7 @@ def _check(original, transformed): mod = tvm.IRModule.from_expr(original.with_attr("global_symbol", "main")) mod = tvm.s_tir.transform.LowerMatchBuffer()(mod) - mod = tvm.tirx.transform.StmtSimplify()(mod) + mod = tvm.s_tir.transform.StmtSimplify()(mod) tvm.ir.assert_structural_equal(mod["main"], transformed.with_attr("global_symbol", "main")) diff --git a/tests/python/s_tir/transform/test_s_tir_transform_unify_thread_binding.py b/tests/python/s_tir/transform/test_s_tir_transform_unify_thread_binding.py index c5e421774698..9bbbc5cb9977 100644 --- a/tests/python/s_tir/transform/test_s_tir_transform_unify_thread_binding.py +++ b/tests/python/s_tir/transform/test_s_tir_transform_unify_thread_binding.py @@ -28,7 +28,7 @@ def _check(original, transformed): mod = tvm.IRModule.from_expr(original.with_attr("global_symbol", "main")) mod = tvm.s_tir.transform.UnifyThreadBinding()(mod) - mod = tvm.tirx.transform.StmtSimplify()(mod) + mod = tvm.s_tir.transform.StmtSimplify()(mod) tvm.ir.assert_structural_equal( mod["main"], transformed.with_attr("global_symbol", "main"), True ) diff --git a/tests/python/tirx-analysis/test_tir_analysis_verify_well_formed.py b/tests/python/tirx-analysis/test_tir_analysis_verify_well_formed.py index 39d0e732cc12..20af6b0f231d 100644 --- a/tests/python/tirx-analysis/test_tir_analysis_verify_well_formed.py +++ b/tests/python/tirx-analysis/test_tir_analysis_verify_well_formed.py @@ -43,8 +43,8 @@ def element_wise( # It's a opaque block , so it can use outside variables C[i, j] = B[i, j] * 2.0 - assert tvm.tirx.analysis.verify_well_formed(element_wise) - assert tvm.tirx.analysis.verify_well_formed(tvm.IRModule.from_expr(element_wise)) + assert tvm.s_tir.analysis.verify_well_formed(element_wise) + assert tvm.s_tir.analysis.verify_well_formed(tvm.IRModule.from_expr(element_wise)) def test_buffer_region_bounds_are_visited(): @@ -56,7 +56,7 @@ def test_buffer_region_bounds_are_visited(): region = tvm.tirx.BufferRegion(buffer, [tvm.ir.Range.from_min_extent(undefined, 4)]) block = tvm.s_tir.SBlock([], [region], [], "region", tvm.tirx.Evaluate(0)) func = tvm.tirx.PrimFunc([buffer], block) - assert not tvm.tirx.analysis.verify_well_formed(func, assert_mode=False) + assert not tvm.s_tir.analysis.verify_well_formed(func, assert_mode=False) def test_fail_use_out_loop_var(): @@ -71,7 +71,7 @@ def element_wise( # we cannot use `i` since it's defined outside the block B[vi, vj] = A[i, vj] * 2.0 - assert not tvm.tirx.analysis.verify_well_formed(element_wise, assert_mode=False) + assert not tvm.s_tir.analysis.verify_well_formed(element_wise, assert_mode=False) def test_error_for_out_of_scope_usage(): @@ -99,7 +99,7 @@ def test_error_for_out_of_scope_usage(): (ValueError, tvm.error.InternalError), match="Invalid use of undefined variable i at .* no longer in-scope.", ): - tvm.tirx.analysis.verify_well_formed(func) + tvm.s_tir.analysis.verify_well_formed(func) def test_error_for_nested_rebind_usage(): @@ -116,7 +116,7 @@ def func(): (ValueError, tvm.error.InternalError), match="ill-formed, due to multiple nested definitions of variable i", ): - tvm.tirx.analysis.verify_well_formed(func) + tvm.s_tir.analysis.verify_well_formed(func) def test_error_for_repeated_binding(): @@ -138,7 +138,7 @@ def func(): with pytest.raises( (ValueError, tvm.error.InternalError), match="multiple nested definitions of variable i" ): - tvm.tirx.analysis.verify_well_formed(func) + tvm.s_tir.analysis.verify_well_formed(func) def test_error_for_cross_function_reuse(): @@ -161,7 +161,7 @@ def func2(): with pytest.raises( (ValueError, tvm.error.InternalError), match="multiple definitions of variable i" ): - tvm.tirx.analysis.verify_well_formed(mod) + tvm.s_tir.analysis.verify_well_formed(mod) def test_reuse_of_env_thread_in_function_is_well_formed(): @@ -180,7 +180,7 @@ def func(A: T.Buffer([256], "float32")): with T.launch_thread(threadIdx_x, 256): A[threadIdx_x] = A[threadIdx_x] + 2.0 - tvm.tirx.analysis.verify_well_formed(func) + tvm.s_tir.analysis.verify_well_formed(func) def test_reuse_of_env_thread_in_function_is_mandatory(): @@ -201,7 +201,7 @@ def func(A: T.Buffer([256], "float32")): with T.launch_thread("threadIdx.x", 256) as threadIdx_x: A[threadIdx_x] = A[threadIdx_x] + 2.0 - tvm.tirx.analysis.verify_well_formed(func) + tvm.s_tir.analysis.verify_well_formed(func) def test_reuse_of_env_thread_across_functions_is_ill_formed(): @@ -237,7 +237,7 @@ def kernel_2(A: T.Buffer([256], "float32")): with pytest.raises( (ValueError, tvm.error.InternalError), match="multiple definitions of variable threadIdx_x" ): - tvm.tirx.analysis.verify_well_formed(mod) + tvm.s_tir.analysis.verify_well_formed(mod) def test_multiple_buffer_arguments_may_share_allocation(): @@ -257,7 +257,7 @@ def func(A_handle: T.handle, B_handle: T.handle): pass - tvm.tirx.analysis.verify_well_formed(mod) + tvm.s_tir.analysis.verify_well_formed(mod) def test_block_match_buffer_defines_buffer_obj(): @@ -276,7 +276,7 @@ def func(A: T.Buffer([256, 256], "float32")): ) B[i, j] = 0.0 - tvm.tirx.analysis.verify_well_formed(mod) + tvm.s_tir.analysis.verify_well_formed(mod) def test_block_match_buffer_defines_symbolic_variables(): @@ -299,7 +299,7 @@ def func(A: T.Buffer([256, 256], "int32")): B[i, j] = elem_offset - tvm.tirx.analysis.verify_well_formed(mod) + tvm.s_tir.analysis.verify_well_formed(mod) def test_error_message_without_previous_definition_location(): @@ -325,7 +325,7 @@ def func(): T.evaluate(x) with pytest.raises((ValueError, tvm.error.InternalError)) as exc_info: - tvm.tirx.analysis.verify_well_formed(func, assert_mode=True) + tvm.s_tir.analysis.verify_well_formed(func, assert_mode=True) error_msg = str(exc_info.value) @@ -350,7 +350,7 @@ def func(): T.evaluate(x) with pytest.raises((ValueError, tvm.error.InternalError)) as exc_info: - tvm.tirx.analysis.verify_well_formed(func, assert_mode=True) + tvm.s_tir.analysis.verify_well_formed(func, assert_mode=True) error_msg = str(exc_info.value) @@ -381,7 +381,7 @@ def func(): T.evaluate(x) with pytest.raises((ValueError, tvm.error.InternalError)) as exc_info: - tvm.tirx.analysis.verify_well_formed(func, assert_mode=True) + tvm.s_tir.analysis.verify_well_formed(func, assert_mode=True) error_msg = str(exc_info.value) @@ -399,7 +399,7 @@ def func(A: T.Buffer((128,), "float32"), B: T.Buffer((128,), "float32")): for i in T.grid(128): B[i] = A[i] * 2.0 - tvm.tirx.analysis.verify_well_formed(func) + tvm.s_tir.analysis.verify_well_formed(func) def test_decl_buffer_is_well_formed(): @@ -411,7 +411,7 @@ def func(A: T.Buffer((128,), "float32")): for i in T.grid(128): B[i] = A[i] * 2.0 - tvm.tirx.analysis.verify_well_formed(func) + tvm.s_tir.analysis.verify_well_formed(func) def test_alloc_buffer_in_block_is_well_formed(): @@ -428,7 +428,7 @@ def func(A: T.Buffer((128,), "float32")): vi = T.axis.remap("S", [i]) B[vi] = A[vi] * 2.0 - tvm.tirx.analysis.verify_well_formed(mod) + tvm.s_tir.analysis.verify_well_formed(mod) def test_match_buffer_in_block_is_well_formed(): @@ -447,7 +447,7 @@ def func(A: T.Buffer((128, 128), "float32")): ) A_tile[i, j] = A_tile[i, j] * 2.0 - tvm.tirx.analysis.verify_well_formed(mod) + tvm.s_tir.analysis.verify_well_formed(mod) def test_error_undeclared_buffer_in_schedulable_tir(): @@ -489,7 +489,7 @@ def test_error_undeclared_buffer_in_schedulable_tir(): with pytest.raises( (ValueError, tvm.error.InternalError), match="buffer B.*without a prior DeclBuffer" ): - tvm.tirx.analysis.verify_well_formed(prim_func) + tvm.s_tir.analysis.verify_well_formed(prim_func) def test_tensor_load_asserted_type_matches_source_and_indices(): @@ -544,5 +544,70 @@ def test_tensor_load_malformed_indices_return_false_without_asserting(): tvm.tirx.analysis.verify_well_formed(non_final_vector) +@pytest.mark.parametrize("as_module", [False, True]) +def test_core_verifiers_reject_blocks(as_module): + block = tvm.s_tir.SBlock([], [], [], "block", tvm.tirx.Evaluate(0)) + func = tvm.tirx.PrimFunc([], block).with_attr("s_tir", True) + obj = tvm.IRModule.from_expr(func) if as_module else func + assert tvm.s_tir.analysis.verify_well_formed(obj) + for verify in [ + tvm.tirx.analysis.verify_well_formed, + tvm.tirx.analysis.verify_tirx_well_formed, + ]: + assert not verify(obj, assert_mode=False) + with pytest.raises(tvm.error.InternalError, match="(does not support|not allowed)"): + verify(obj) + + +def test_mixed_module_parser_checks_both_dialects(): + @I.ir_module + class Mixed: + @T.prim_func(s_tir=True) + def scheduled(A: T.Buffer((4,), "int32")): + for i in range(4): + with T.sblock("write"): + vi = T.axis.spatial(4, i) + A[vi] = vi + + @T.prim_func + def lowered(): + T.evaluate(0) + + assert tvm.s_tir.analysis.verify_well_formed(Mixed) + assert tvm.tirx.analysis.verify_tirx_well_formed(Mixed["lowered"]) + assert not tvm.tirx.analysis.verify_tirx_well_formed(Mixed, assert_mode=False) + + +def test_s_tir_verifier_preserves_shared_definition_check(): + shared = tvm.tirx.Var("shared", "int32") + core = tvm.tirx.PrimFunc([shared], tvm.tirx.Evaluate(shared)) + block = tvm.s_tir.SBlock([], [], [], "block", tvm.tirx.Evaluate(shared)) + scheduled = tvm.tirx.PrimFunc([shared], block).with_attr("s_tir", True) + mod = tvm.IRModule({"core": core, "scheduled": scheduled}) + assert not tvm.s_tir.analysis.verify_well_formed(mod, assert_mode=False) + with pytest.raises(tvm.error.InternalError, match="multiple definitions"): + tvm.s_tir.analysis.verify_well_formed(mod) + + +def test_matched_buffer_is_defined_before_block_regions(): + allocated = tvm.tirx.decl_buffer((4,), "float32", name="allocated") + matched = tvm.tirx.decl_buffer((4,), "float32", name="matched") + source = tvm.tirx.BufferRegion(allocated, [tvm.ir.Range(4)]) + region = tvm.tirx.BufferRegion(matched, [tvm.ir.Range(4)]) + block = tvm.s_tir.SBlock( + [], + [region], + [], + "use", + tvm.tirx.Evaluate(matched[0]), + alloc_buffers=[allocated], + match_buffers=[tvm.s_tir.MatchBufferRegion(matched, source)], + ) + func = tvm.tirx.PrimFunc([], block) + assert tvm.s_tir.analysis.verify_well_formed(func) + out_of_scope = tvm.tirx.PrimFunc([], tvm.tirx.SeqStmt([block, tvm.tirx.Evaluate(matched[0])])) + assert not tvm.s_tir.analysis.verify_well_formed(out_of_scope, assert_mode=False) + + if __name__ == "__main__": tvm.testing.main() diff --git a/tests/python/tirx-base/test_tir_specialize.py b/tests/python/tirx-base/test_tir_specialize.py index 4dffc8dc11e9..8dc8832c41bf 100644 --- a/tests/python/tirx-base/test_tir_specialize.py +++ b/tests/python/tirx-base/test_tir_specialize.py @@ -368,5 +368,49 @@ def expected() -> T.int32: tvm.ir.assert_structural_equal(after.ty, ty_expected) +def test_specialize_structural_buffer_definitions(): + """Extension definitions remain consistent across metadata, regions, and uses.""" + n = tvm.tirx.Var("n", "int32") + allocated = tvm.tirx.decl_buffer((n,), "float32", name="allocated") + matched = tvm.tirx.decl_buffer((n,), "float32", name="matched") + source = tvm.tirx.BufferRegion(allocated, [tvm.ir.Range(n)]) + read = tvm.tirx.BufferRegion(matched, [tvm.ir.Range(n)]) + match = tvm.s_tir.MatchBufferRegion(matched, source) + block = tvm.s_tir.SBlock( + [], + [read], + [], + "use", + tvm.tirx.Evaluate(matched[n - 1]), + alloc_buffers=[allocated], + match_buffers=[match], + annotations={"extent": n}, + ) + before = tvm.tirx.PrimFunc([n], tvm.s_tir.SBlockRealize([], True, block)) + assert tvm.s_tir.analysis.verify_well_formed(before) + after = before.specialize({n: 8}) + assert tvm.s_tir.analysis.verify_well_formed(after) + result = after.body.block + assert not after.params + assert result.alloc_buffers[0].shape[0] == 8 + assert result.match_buffers[0].buffer.shape[0] == 8 + assert result.match_buffers[0].source.buffer.same_as(result.alloc_buffers[0]) + assert result.reads[0].buffer.same_as(result.match_buffers[0].buffer) + assert result.body.value.source.same_as(result.match_buffers[0].buffer) + assert result.body.value.indices[0] == 7 + assert result.annotations["extent"] == 8 + # Specialization must not rewrite another owner's unspecialized definition. + assert block.alloc_buffers[0].shape[0].same_as(n) + assert block.match_buffers[0].buffer.shape[0].same_as(n) + assert block.annotations["extent"].same_as(n) + + +def test_specialize_plain_tirx(): + n = tvm.tirx.Var("n", "int32") + before = tvm.tirx.PrimFunc([n], tvm.tirx.Evaluate(n + 1)) + expected = tvm.tirx.PrimFunc([], tvm.tirx.Evaluate(9)) + tvm.ir.assert_structural_equal(before.specialize({n: 8}), expected) + + if __name__ == "__main__": tvm.testing.main() diff --git a/tests/python/tirx-transform/test_tir_inline_private_functions.py b/tests/python/tirx-transform/test_tir_inline_private_functions.py index 903c84913606..16637aa4d5eb 100644 --- a/tests/python/tirx-transform/test_tir_inline_private_functions.py +++ b/tests/python/tirx-transform/test_tir_inline_private_functions.py @@ -26,7 +26,7 @@ class BaseTestCase: def test_well_formed(self): After = tvm.tirx.transform.InlinePrivateFunctions()(self.Before) - tvm.tirx.analysis.verify_well_formed(After) + tvm.s_tir.analysis.verify_well_formed(After) def test_produces_expected(self): After = tvm.tirx.transform.InlinePrivateFunctions()(self.Before) @@ -85,7 +85,7 @@ def main(A: T.Buffer([80, 16], "float32"), B: T.Buffer([64, 16], "float32")): @T.prim_func(private=True, s_tir=True) def subroutine(A_data: T.handle("float32"), B_data: T.handle("float32")): - T.func_attr({"target": T.target("cuda")}) + T.func_attr({"target": T.target({"kind": "cuda", "arch": "sm_80"})}) A = T.decl_buffer([16, 16], "float32", data=A_data) B = T.decl_buffer([16], "float32", data=B_data) for i in range(16): diff --git a/tests/python/tirx-transform/test_tir_transform_flatten_buffer.py b/tests/python/tirx-transform/test_tir_transform_flatten_buffer.py index 7c2943879cb8..75900822bca0 100644 --- a/tests/python/tirx-transform/test_tir_transform_flatten_buffer.py +++ b/tests/python/tirx-transform/test_tir_transform_flatten_buffer.py @@ -20,9 +20,20 @@ from tvm.script import tirx as T +def _lower_blocks(value): + """Use the same block-to-statement boundary as the S-TIR pipeline.""" + is_func = isinstance(value, tvm.tirx.PrimFunc) + mod = tvm.IRModule.from_expr(value) if is_func else value + mod = tvm.s_tir.transform.ConvertBlocksToOpaque()(mod) + mod = tvm.s_tir.transform.LowerOpaqueBlock()(mod) + return mod["main"] if is_func else mod + + def _transform(): return tvm.transform.Sequential( [ + tvm.s_tir.transform.ConvertBlocksToOpaque(), + tvm.s_tir.transform.LowerOpaqueBlock(), tvm.tirx.transform.FlattenBuffer(), tvm.tirx.transform.StmtSimplify(), ] @@ -57,7 +68,7 @@ def main(A: T.Buffer((16, 16), "float32"), C: T.Buffer((16, 16), "float32")): C_1[((i * 16) + j)] = B_new[j] * 2.0 After = _transform()(Before) - tvm.ir.assert_structural_equal(After, Expected) + tvm.ir.assert_structural_equal(After, _lower_blocks(Expected)) def test_elementwise_without_decl_buffer(): @@ -97,7 +108,7 @@ def main(input_A: T.Buffer((16, 16), "float32"), input_C: T.Buffer((16, 16), "fl C[((i * 16) + j)] = B_new[j] * 2.0 After = _transform()(Before) - tvm.ir.assert_structural_equal(After, Expected) + tvm.ir.assert_structural_equal(After, _lower_blocks(Expected)) def test_gpu(): @@ -141,7 +152,7 @@ def main(A: T.Buffer((16, 16), "float32"), C: T.Buffer((16, 16), "float32")): C_1[i0 * 64 + i1 * 32 + i2 * 16 + j] = B[j] * 2.0 After = _transform()(Before) - tvm.ir.assert_structural_equal(After, Expected) + tvm.ir.assert_structural_equal(After, _lower_blocks(Expected)) def test_symbolic(): @@ -178,7 +189,7 @@ def main(a: T.handle, c: T.handle, n: T.int32, m: T.int32) -> None: C_1[i * m + j] = B[j] * 2.0 After = _transform()(Before) - tvm.ir.assert_structural_equal(After, Expected) + tvm.ir.assert_structural_equal(After, _lower_blocks(Expected)) def test_fused_symbolic(): @@ -209,7 +220,7 @@ def main(a: T.handle, b: T.handle, n: T.int32) -> None: B[i] = A[i] After = _transform()(Before) - tvm.ir.assert_structural_equal(After, Expected) + tvm.ir.assert_structural_equal(After, _lower_blocks(Expected)) def test_fused_symbolic_with_predicate(): @@ -247,7 +258,7 @@ def main(a: T.handle, b: T.handle, n: T.int32) -> None: B[bx * 64 + tx] = A[bx * 64 + tx] After = _transform()(Before) - tvm.ir.assert_structural_equal(After, Expected) + tvm.ir.assert_structural_equal(After, _lower_blocks(Expected)) def test_multi_alloc(): @@ -279,7 +290,7 @@ def main(A: T.Buffer((4, 32), "float32"), D: T.Buffer((4, 32), "float32")): D_1[i * 32 + j] = C[i * 32 + j] * 2.0 After = _transform()(Before) - tvm.ir.assert_structural_equal(After, Expected) + tvm.ir.assert_structural_equal(After, _lower_blocks(Expected)) def test_strided(): @@ -314,7 +325,7 @@ def main(A: T.Buffer((16, 16), "float32"), C: T.Buffer((16, 16), "float32")): C_1[i0 * 64 + i1 * 16 + j] = B_1[i1 * 17 + j] * 2.0 After = _transform()(Before) - tvm.ir.assert_structural_equal(After, Expected) + tvm.ir.assert_structural_equal(After, _lower_blocks(Expected)) def test_boolean(): @@ -338,11 +349,11 @@ def main(input_A: T.Buffer(10, "bool"), input_B: T.Buffer(10, "bool")) -> None: B[i0] = A[i0] After = _transform()(Before) - tvm.ir.assert_structural_equal(After, Expected) + tvm.ir.assert_structural_equal(After, _lower_blocks(Expected)) -def test_flatten_inside_block(): - """Flattening access inside a block flattens the accessed region.""" +def test_flatten_lowered_block(): + """Flatten allocations and accesses after lowering a schedulable block.""" @I.ir_module(s_tir=True) class Before: @@ -365,7 +376,7 @@ def main(): T.evaluate(A[i * 32 + j]) After = _transform()(Before) - tvm.ir.assert_structural_equal(After, Expected) + tvm.ir.assert_structural_equal(After, _lower_blocks(Expected)) def test_build_with_optional_pragma_unroll_explicit(): diff --git a/tests/python/tirx-transform/test_tir_transform_force_narrow_index_to_i32.py b/tests/python/tirx-transform/test_tir_transform_force_narrow_index_to_i32.py index 10169667296b..9ce2bd83f7c0 100644 --- a/tests/python/tirx-transform/test_tir_transform_force_narrow_index_to_i32.py +++ b/tests/python/tirx-transform/test_tir_transform_force_narrow_index_to_i32.py @@ -22,6 +22,25 @@ from tvm.script import tirx as T +def _lower_blocks(value): + """Use the same block-to-statement boundary as the S-TIR pipeline.""" + is_func = isinstance(value, tvm.tirx.PrimFunc) + mod = tvm.IRModule.from_expr(value) if is_func else value + mod = tvm.s_tir.transform.ConvertBlocksToOpaque()(mod) + mod = tvm.s_tir.transform.LowerOpaqueBlock()(mod) + return mod["main"] if is_func else mod + + +def _transform(): + return tvm.transform.Sequential( + [ + tvm.s_tir.transform.ConvertBlocksToOpaque(), + tvm.s_tir.transform.LowerOpaqueBlock(), + tvm.tirx.transform.ForceNarrowIndexToInt32(), + ] + ) + + def test_thread_axis1(): @T.prim_func(private=True, s_tir=True) def before(A: T.Buffer((T.int64(64),), "float32"), B: T.Buffer((T.int64(64),), "float32")): @@ -42,8 +61,8 @@ def expected(A: T.Buffer((64,), "float32"), B: T.Buffer((64,), "float32")): B[blockIdx_x * 32 + threadIdx_x] = A[blockIdx_x * 32 + threadIdx_x] + T.float32(1) mod = tvm.IRModule.from_expr(before) - func = tvm.tirx.transform.ForceNarrowIndexToInt32()(mod)["main"] - tvm.ir.assert_structural_equal(func, expected) + func = _transform()(mod)["main"] + tvm.ir.assert_structural_equal(func, _lower_blocks(expected)) def test_thread_axis2(): @@ -157,8 +176,8 @@ def expected( ) mod = tvm.IRModule.from_expr(before) - func = tvm.tirx.transform.ForceNarrowIndexToInt32()(mod)["main"] - tvm.ir.assert_structural_equal(func, expected) + func = _transform()(mod)["main"] + tvm.ir.assert_structural_equal(func, _lower_blocks(expected)) def test_block(): @@ -179,8 +198,8 @@ def expected(A: T.Buffer((128,), "float32"), B: T.Buffer((128,), "float32")): B[vi] = A[vi] + T.float32(1) mod = tvm.IRModule.from_expr(before) - func = tvm.tirx.transform.ForceNarrowIndexToInt32()(mod)["main"] - tvm.ir.assert_structural_equal(func, expected) + func = _transform()(mod)["main"] + tvm.ir.assert_structural_equal(func, _lower_blocks(expected)) def test_i16_buffer(): @@ -201,8 +220,8 @@ def expected(A: T.Buffer((128,), "int16"), B: T.Buffer((128,), "int16")): B[vi] = A[vi] + T.int16(1) mod = tvm.IRModule.from_expr(before) - after = tvm.tirx.transform.ForceNarrowIndexToInt32()(mod)["main"] - tvm.ir.assert_structural_equal(after, expected) + after = _transform()(mod)["main"] + tvm.ir.assert_structural_equal(after, _lower_blocks(expected)) def test_fail_on_buffer_param(): @@ -216,7 +235,7 @@ def func(A: T.Buffer((128,), "int64"), B: T.Buffer((128,), "int64")): mod = tvm.IRModule.from_expr(func) with pytest.raises(RuntimeError): - tvm.tirx.transform.ForceNarrowIndexToInt32()(mod)["main"] + _transform()(mod)["main"] def test_fail_on_internal_buffer(): @@ -236,7 +255,7 @@ def func(A: T.Buffer((128,), "int32"), B: T.Buffer((128,), "int32")): mod = tvm.IRModule.from_expr(func) with pytest.raises(RuntimeError): - tvm.tirx.transform.ForceNarrowIndexToInt32()(mod)["main"] + _transform()(mod)["main"] def test_pod_params_and_select(): @@ -256,8 +275,8 @@ def main(A: T.Buffer((4,), "float32"), B: T.Buffer((4,), "float32"), n: T.int32) for i in range(4): B[i] = T.Select(1 <= i, A[i + n], T.Cast("float32", i)) - after = tvm.tirx.transform.ForceNarrowIndexToInt32()(Before) - tvm.ir.assert_structural_equal(Expected, after) + after = _transform()(Before) + tvm.ir.assert_structural_equal(_lower_blocks(Expected), after) def test_if_then_else_index(): @@ -273,8 +292,8 @@ class Expected: def main(A: T.Buffer((4,), "float32"), B: T.Buffer((1,), "float32"), n: T.int32): B[0] = A[T.if_then_else(n < 0, n + 1, n)] - after = tvm.tirx.transform.ForceNarrowIndexToInt32()(Before) - tvm.ir.assert_structural_equal(Expected, after) + after = _transform()(Before) + tvm.ir.assert_structural_equal(_lower_blocks(Expected), after) def test_conditional_index_mixed_width_branches(): @@ -298,8 +317,8 @@ def main(A: T.Buffer((4,), "float32"), B: T.Buffer((4,), "float32"), n: T.int32) B[2] = A[T.Select(n < 0, opaque_index, T.Cast("int64", n))] B[3] = A[T.Select(n < 0, T.Cast("int64", n), opaque_index)] - after = tvm.tirx.transform.ForceNarrowIndexToInt32()(Before) - tvm.ir.assert_structural_equal(Expected, after) + after = _transform()(Before) + tvm.ir.assert_structural_equal(_lower_blocks(Expected), after) def test_clz(): @@ -317,8 +336,8 @@ def main(B: T.Buffer((4,), "int32")): for i in range(4): B[i] = T.clz(i) - 32 + 64 - after = tvm.tirx.transform.ForceNarrowIndexToInt32()(Before) - tvm.ir.assert_structural_equal(Expected, after) + after = _transform()(Before) + tvm.ir.assert_structural_equal(_lower_blocks(Expected), after) def test_right_shift_preserves_sign_extension_after_narrowing(): @@ -337,8 +356,8 @@ def main(A: T.Buffer((6,), "float32"), B: T.Buffer((1,), "float32"), n: T.int32) # ForceNarrowIndexToInt32 assumes that index values fit in int32. Under # that precondition, shifting the original int64 value by 63 and shifting # the narrowed value by its sign-bit position have the same result. - after = tvm.tirx.transform.ForceNarrowIndexToInt32()(Before) - tvm.ir.assert_structural_equal(Expected, after) + after = _transform()(Before) + tvm.ir.assert_structural_equal(_lower_blocks(Expected), after) def test_right_shift_dynamic_and_vector_amounts(): @@ -376,8 +395,8 @@ def main( ) ] - after = tvm.tirx.transform.ForceNarrowIndexToInt32()(Before) - tvm.ir.assert_structural_equal(Expected, after) + after = _transform()(Before) + tvm.ir.assert_structural_equal(_lower_blocks(Expected), after) def test_left_shift_dynamic_and_vector_amounts_remain_valid(): @@ -415,8 +434,8 @@ def main( ) ] - after = tvm.tirx.transform.ForceNarrowIndexToInt32()(Before) - tvm.ir.assert_structural_equal(Expected, after) + after = _transform()(Before) + tvm.ir.assert_structural_equal(_lower_blocks(Expected), after) def test_let_binding(): @@ -443,8 +462,35 @@ def main(buf: T.handle): for i in range(T.Cast("int32", ceil_log2)): T.evaluate(0) - after = tvm.tirx.transform.ForceNarrowIndexToInt32()(Before) - tvm.ir.assert_structural_equal(Expected, after) + after = _transform()(Before) + tvm.ir.assert_structural_equal(_lower_blocks(Expected), after) + + +def test_preserve_local_scalar_storage(): + @T.prim_func(private=True) + def before(n: T.int64): + index = T.call_extern("opaque_index", n, dtype="int64") + T.evaluate(index) + + @T.prim_func(private=True) + def expected(n: T.int32): + index = T.call_extern("opaque_index", n, dtype="int64") + T.evaluate(index) + + after = tvm.tirx.transform.ForceNarrowIndexToInt32()(tvm.IRModule.from_expr(before))["main"] + tvm.ir.assert_structural_equal(after, expected) + + +@pytest.mark.parametrize("shape", [(), (1,), (8,)]) +@pytest.mark.parametrize("scope", ["local", "shared", "global"]) +def test_reject_int64_array_storage(shape, scope): + # Rank or element count alone does not make an allocation scalar storage. + buffer = tvm.tirx.decl_buffer(shape, "int64", "array", scope=scope, layout=None) + before = tvm.tirx.PrimFunc( + [], tvm.tirx.SeqStmt([tvm.tirx.AllocBuffer(buffer), tvm.tirx.Evaluate(0)]) + ) + with pytest.raises(tvm.error.InternalError, match="allocated in the function has dtype"): + tvm.tirx.transform.ForceNarrowIndexToInt32()(tvm.IRModule.from_expr(before)) if __name__ == "__main__": diff --git a/tests/python/tirx-transform/test_tir_transform_narrow_datatype.py b/tests/python/tirx-transform/test_tir_transform_narrow_datatype.py index 7c1d86a7e8b8..64728864f43b 100644 --- a/tests/python/tirx-transform/test_tir_transform_narrow_datatype.py +++ b/tests/python/tirx-transform/test_tir_transform_narrow_datatype.py @@ -20,9 +20,28 @@ from tvm.tirx import const +def _lower_blocks(value): + """Use the same block-to-statement boundary as the S-TIR pipeline.""" + is_func = isinstance(value, tvm.tirx.PrimFunc) + mod = tvm.IRModule.from_expr(value) if is_func else value + mod = tvm.s_tir.transform.ConvertBlocksToOpaque()(mod) + mod = tvm.s_tir.transform.LowerOpaqueBlock()(mod) + return mod["main"] if is_func else mod + + +def _transform(target_bits): + return tvm.transform.Sequential( + [ + tvm.s_tir.transform.ConvertBlocksToOpaque(), + tvm.s_tir.transform.LowerOpaqueBlock(), + tvm.tirx.transform.NarrowDataType(target_bits), + ] + ) + + def lower_stmt(params, stmt, target_bits): func = tvm.tirx.PrimFunc(params, stmt) - func = tvm.tirx.transform.NarrowDataType(target_bits)(tvm.IRModule.from_expr(func))["main"] + func = _transform(target_bits)(tvm.IRModule.from_expr(func))["main"] stmt = func.body return stmt @@ -31,7 +50,7 @@ def lower_func_body(func, target_bits): """Lower a TVMScript function and return the first For loop in the body.""" mod = tvm.IRModule.from_expr(func) gvar = next(iter(mod.functions.keys())) - func = tvm.tirx.transform.NarrowDataType(target_bits)(mod)[gvar] + func = _transform(target_bits)(mod)[gvar] body = func.body # With flat buffer semantics, navigate to the first For node if isinstance(body, tvm.tirx.SeqStmt): @@ -110,7 +129,7 @@ def func(A: T.Buffer((m * n,), "float32"), B: T.Buffer((m * n,), "float32")): mod = tvm.IRModule.from_expr(func) gvar = next(iter(mod.functions.keys())) - func_narrowed = tvm.tirx.transform.NarrowDataType(target_bits)(mod)[gvar] + func_narrowed = _transform(target_bits)(mod)[gvar] stmt = func_narrowed.body assert stmt.node.var.ty.dtype == target_dtype assert stmt.body.node.var.ty.dtype == target_dtype @@ -148,7 +167,7 @@ def func( mod = tvm.IRModule.from_expr(func) gvar = next(iter(mod.functions.keys())) - func_narrowed = tvm.tirx.transform.NarrowDataType(target_bits)(mod)[gvar] + func_narrowed = _transform(target_bits)(mod)[gvar] stmt = func_narrowed.body assert stmt.seq[0].loop_var.ty.dtype == target_dtype @@ -209,10 +228,12 @@ def expected_after(A: T.Buffer(128, "float32"), B: T.Buffer(130, "float32")): i * 65 + j >= 0 and i * 65 + j < 128, A[i * 65 + j], T.float32(0), dtype="float32" ) - after = tvm.tirx.transform.NarrowDataType(32)( - tvm.IRModule.from_expr(before.with_attr("global_symbol", "main")) - )["main"] - tvm.ir.assert_structural_equal(after, expected_after.with_attr("global_symbol", "main")) + after = _transform(32)(tvm.IRModule.from_expr(before.with_attr("global_symbol", "main")))[ + "main" + ] + tvm.ir.assert_structural_equal( + after, _lower_blocks(expected_after.with_attr("global_symbol", "main")) + ) def test_block(): @@ -232,10 +253,12 @@ def expected_after(A: T.Buffer((128,), "float32"), B: T.Buffer((128,), "float32" vi = T.axis.spatial(T.int32(128), i * T.int32(8) + j) B[vi] = A[vi] + T.float32(1) - after = tvm.tirx.transform.NarrowDataType(32)( - tvm.IRModule.from_expr(before.with_attr("global_symbol", "main")) - )["main"] - tvm.ir.assert_structural_equal(after, expected_after.with_attr("global_symbol", "main")) + after = _transform(32)(tvm.IRModule.from_expr(before.with_attr("global_symbol", "main")))[ + "main" + ] + tvm.ir.assert_structural_equal( + after, _lower_blocks(expected_after.with_attr("global_symbol", "main")) + ) def test_avg_pool2d(): @@ -294,11 +317,11 @@ def expected_after(PSUM: T.Buffer((313600,), "int32"), PAVG: T.Buffer((313600,), ), ) - after = tvm.tirx.transform.NarrowDataType(32)( - tvm.IRModule.from_expr(before.with_attr("global_symbol", "main")) - ) + after = _transform(32)(tvm.IRModule.from_expr(before.with_attr("global_symbol", "main"))) after = tvm.tirx.transform.StmtSimplify()(after) - tvm.ir.assert_structural_equal(after["main"], expected_after.with_attr("global_symbol", "main")) + tvm.ir.assert_structural_equal( + after["main"], _lower_blocks(expected_after.with_attr("global_symbol", "main")) + ) def test_narrow_i64_valued_bufferload_index_to_i32(): @@ -312,10 +335,10 @@ def expect(A: T.Buffer((16,), "int64")): for i in range(15): A[i + 1] = A[i] + T.int64(1) - after = tvm.tirx.transform.NarrowDataType(32)( - tvm.IRModule.from_expr(before.with_attr("global_symbol", "main")) - )["main"] - tvm.ir.assert_structural_equal(after, expect.with_attr("global_symbol", "main")) + after = _transform(32)(tvm.IRModule.from_expr(before.with_attr("global_symbol", "main")))[ + "main" + ] + tvm.ir.assert_structural_equal(after, _lower_blocks(expect.with_attr("global_symbol", "main"))) if __name__ == "__main__": diff --git a/tests/python/tirx-transform/test_tir_transform_simplify.py b/tests/python/tirx-transform/test_tir_transform_simplify.py index f62dcf77bac7..44f9eeaa4d11 100644 --- a/tests/python/tirx-transform/test_tir_transform_simplify.py +++ b/tests/python/tirx-transform/test_tir_transform_simplify.py @@ -1313,5 +1313,26 @@ def expected(a: T.Buffer((2, 8), "int32"), b: T.Buffer((2, 8), "int32")): tvm.ir.assert_structural_equal(after, expected) +def test_s_tir_block_iterator_constraints(): + @T.prim_func(private=True, s_tir=True) + def before(A: T.Buffer((4,), "int32")): + for i in range(4): + with T.sblock("write"): + vi = T.axis.spatial(4, i) + if vi < 4: + A[vi] = vi + + @T.prim_func(private=True, s_tir=True) + def expected(A: T.Buffer((4,), "int32")): + for i in range(4): + with T.sblock("write"): + vi = T.axis.spatial(4, i) + A[vi] = vi + + result = tvm.s_tir.transform.StmtSimplify()(tvm.IRModule.from_expr(before))["main"] + tvm.ir.assert_structural_equal(result, expected) + assert tvm.s_tir.analysis.verify_well_formed(result) + + if __name__ == "__main__": tvm.testing.main() From 7080969efe830eada76128779307a8289bf0662f Mon Sep 17 00:00:00 2001 From: Tianqi Chen Date: Thu, 17 Sep 2026 11:39:02 +0000 Subject: [PATCH 06/18] Preserve reference results through statement fallback dispatch --- include/tvm/tirx/stmt_functor.h | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/include/tvm/tirx/stmt_functor.h b/include/tvm/tirx/stmt_functor.h index 637a65a71621..614a7e6f8762 100644 --- a/include/tvm/tirx/stmt_functor.h +++ b/include/tvm/tirx/stmt_functor.h @@ -119,8 +119,8 @@ class StmtFunctor { // Register inherited hooks in a fresh table before adding dialect nodes. static void InitVTable(VTable* vtable) { vtable->template SetDispatch( - [](const ffi::ObjectRef& node, TSelf* self, Args... args) { - return self->VisitStmtDefault_(node.get(), std::forward(args)...); + [](const ffi::ObjectRef& node, TSelf* self, Args... args) -> R { + return self->DispatchDefault_(node.get(), std::forward(args)...); }); IR_STMT_FUNCTOR_DISPATCH(BindNode); IR_STMT_FUNCTOR_DISPATCH(AttrStmtNode); From b8ed1c7614bf675ef23601a947a4da3d9d0a23d9 Mon Sep 17 00:00:00 2001 From: Tianqi Chen Date: Thu, 17 Sep 2026 11:41:54 +0000 Subject: [PATCH 07/18] Reconcile S-TIR ownership with shared regions and native dispatch --- include/tvm/s_tir/analysis.h | 4 +- include/tvm/s_tir/stmt.h | 12 +- include/tvm/s_tir/stmt_functor.h | 12 +- include/tvm/tirx/stmt_functor.h | 174 +++++++++--------- python/tvm/ir/json_compact.py | 6 + python/tvm/s_tir/analysis/__init__.py | 1 - python/tvm/s_tir/stmt.py | 22 +-- python/tvm/tirx/stmt.py | 2 +- src/relax/analysis/tir_op_pattern_kind.cc | 4 +- src/relax/transform/fuse_tir.cc | 4 +- .../transform/split_call_tir_by_pattern.cc | 4 +- .../analysis/sblock_access_region_detector.cc | 6 +- .../sblock_buffer_access_lca_detector.cc | 5 +- src/s_tir/analysis/verify_well_formed.cc | 3 +- src/s_tir/ir/data_type_rewriter.cc | 21 ++- src/s_tir/ir/data_type_rewriter.h | 2 +- src/s_tir/ir/ir_mutator_with_analyzer.h | 2 +- src/s_tir/ir/ir_visitor_with_analyzer.h | 2 +- src/s_tir/ir/tir_visitor_with_path.cc | 4 +- src/s_tir/ir/tir_visitor_with_path.h | 10 +- .../multi_level_tiling_tensor_core.cc | 2 +- src/s_tir/stmt.cc | 35 ++-- src/s_tir/stmt_functor.cc | 8 +- src/s_tir/transform/default_gpu_schedule.cc | 4 +- src/s_tir/transform/ir_utils.cc | 15 +- src/tirx/analysis/verify_tirx_well_formed.cc | 8 +- src/tirx/ir/stmt.cc | 44 ----- src/tirx/ir/stmt_functor.cc | 20 -- src/tirx/ir/tir_visitor_with_path.cc | 2 +- src/tirx/ir/tir_visitor_with_path.h | 38 ++-- src/tirx/script/builder/frame.cc | 7 +- src/tirx/script/builder/ir.cc | 2 +- src/tirx/transform/stmt_simplify.h | 2 +- tests/cpp/ir_functor_test.cc | 4 +- tests/cpp/s_tir_functor_test.cc | 52 +++--- tests/python/s_tir/test_stmt.py | 53 +++++- .../test_s_tir_transform_convert_ssa.py | 2 +- tests/python/tirx-base/test_tir_specialize.py | 4 +- 38 files changed, 307 insertions(+), 295 deletions(-) diff --git a/include/tvm/s_tir/analysis.h b/include/tvm/s_tir/analysis.h index f291d2c9b6c4..acaf2a4ba676 100644 --- a/include/tvm/s_tir/analysis.h +++ b/include/tvm/s_tir/analysis.h @@ -48,7 +48,7 @@ namespace tirx { * - second: write regions * - third: opaque regions */ -TVM_DLL ffi::Array> GetSBlockAccessRegion( +TVM_DLL ffi::Array> GetSBlockAccessRegion( const s_tir::SBlock& block, const ffi::Map& buffer_var_map); /*! @@ -59,7 +59,7 @@ TVM_DLL ffi::Array> GetSBlockAccessRegion( * It is a map from buffer var to the buffer * \return An array only consisting of the read regions and write regions of the input block */ -TVM_DLL ffi::Array> GetSBlockReadWriteRegion( +TVM_DLL ffi::Array> GetSBlockReadWriteRegion( const s_tir::SBlock& block, const ffi::Map& buffer_var_map); /*! diff --git a/include/tvm/s_tir/stmt.h b/include/tvm/s_tir/stmt.h index 818776242030..9feff932144f 100644 --- a/include/tvm/s_tir/stmt.h +++ b/include/tvm/s_tir/stmt.h @@ -46,7 +46,7 @@ class MatchBufferRegionNode : public ffi::Object { /*! \brief The target buffer. */ tirx::BufferVar buffer; /*! \brief The source buffer region. */ - tirx::BufferRegion source; + TensorRegion source; static void RegisterReflection() { namespace refl = tvm::ffi::reflection; @@ -65,7 +65,7 @@ class MatchBufferRegionNode : public ffi::Object { */ class MatchBufferRegion : public ffi::ObjectRef { public: - TVM_DLL explicit MatchBufferRegion(tirx::BufferVar buffer, tirx::BufferRegion source); + TVM_DLL explicit MatchBufferRegion(tirx::BufferVar buffer, TensorRegion source); TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(MatchBufferRegion, ffi::ObjectRef, MatchBufferRegionNode); @@ -98,9 +98,9 @@ class SBlockNode : public tirx::StmtNode { /*! \brief The variables of the block. */ ffi::Array iter_vars; /*! \brief The read buffer regions of the block. */ - ffi::Array reads; + ffi::Array reads; /*! \brief The write buffer regions of the block. */ - ffi::Array writes; + ffi::Array writes; /*! \brief The name_hint of the block. */ ffi::String name_hint; /*! \brief The buffer allocated in the block. */ @@ -144,8 +144,8 @@ class SBlockNode : public tirx::StmtNode { class SBlock : public tirx::Stmt { public: TVM_DLL explicit SBlock( - ffi::Array iter_vars, ffi::Array reads, - ffi::Array writes, ffi::String name_hint, tirx::Stmt body, + ffi::Array iter_vars, ffi::Array reads, + ffi::Array writes, ffi::String name_hint, tirx::Stmt body, ffi::Optional init = std::nullopt, ffi::Array alloc_buffers = ffi::Array(), ffi::Array match_buffers = ffi::Array(), diff --git a/include/tvm/s_tir/stmt_functor.h b/include/tvm/s_tir/stmt_functor.h index 9b48e4ef0135..8d4734c07c5f 100644 --- a/include/tvm/s_tir/stmt_functor.h +++ b/include/tvm/s_tir/stmt_functor.h @@ -31,7 +31,7 @@ namespace s_tir { /*! * \brief Extend TIRX statement dispatch with schedulable blocks. - * \tparam FType The statement signature, retaining the VisitStmt API. + * \tparam FType The statement signature, using the native Dispatch API. */ template class StmtFunctor; @@ -43,13 +43,13 @@ class StmtFunctor public: TVM_DEFINE_OBJECT_FUNCTOR_DEFAULT_CONSTRUCTOR(StmtFunctor, Parent) - using Parent::VisitStmt_; + using Parent::Dispatch_; - virtual R VisitStmt_(const SBlockNode* op, Args... args) { - return this->VisitStmtDefault_(op, std::forward(args)...); + virtual R Dispatch_(const SBlockNode* op, Args... args) { + return this->DispatchDefault_(op, std::forward(args)...); } - virtual R VisitStmt_(const SBlockRealizeNode* op, Args... args) { - return this->VisitStmtDefault_(op, std::forward(args)...); + virtual R Dispatch_(const SBlockRealizeNode* op, Args... args) { + return this->DispatchDefault_(op, std::forward(args)...); } protected: diff --git a/include/tvm/tirx/stmt_functor.h b/include/tvm/tirx/stmt_functor.h index 614a7e6f8762..66d2b2d68543 100644 --- a/include/tvm/tirx/stmt_functor.h +++ b/include/tvm/tirx/stmt_functor.h @@ -52,116 +52,120 @@ namespace tirx { template class StmtFunctor; -#define STMT_FUNCTOR_DEFAULT \ - { \ - return VisitStmtDefault_(op, std::forward(args)...); \ - } - -#define IR_STMT_FUNCTOR_DISPATCH(OP) \ - vtable->template SetDispatch([](const ffi::ObjectRef& n, TSelf* self, Args... args) { \ - return self->VisitStmt_(static_cast(n.get()), std::forward(args)...); \ - }); - template class StmtFunctor { private: - using TSelf = StmtFunctor; + using TSelf = StmtFunctor; public: /*! \brief The result type of this functor. */ using result_type = R; + /*! \brief Construct a functor with the TIRx statement hooks. */ StmtFunctor() : StmtFunctor(GlobalVTable()) {} - /*! \brief virtual destructor */ - virtual ~StmtFunctor() {} - /*! - * \brief Same as call. - * \param n The stmt node. - * \param args Additional arguments. - * \return The result of the call - */ - R operator()(const Stmt& n, Args... args) { return VisitStmt(n, std::forward(args)...); } - /*! - * \brief The functor call. - * \param n The stmt node. - * \param args Additional arguments. - * \return The result of the call - */ - virtual R VisitStmt(const Stmt& n, Args... args) { - return (*vtable_)(n, this, std::forward(args)...); + /*! \brief Destroy through the statement functor base. */ + virtual ~StmtFunctor() = default; + /*! \brief Dispatch a statement, forwarding additional arguments to its hook. */ + TVM_FFI_INLINE R operator()(const Stmt& node, Args... args) { + return Dispatch(node, std::forward(args)...); + } + /*! \brief Dispatch to a node hook, including registered ancestor hooks. */ + TVM_FFI_INLINE virtual R Dispatch(const Stmt& node, Args... args) { + TVM_FFI_ICHECK(node.defined()) << "Cannot dispatch a null statement"; + return (*vtable_)(node, this, std::forward(args)...); + } + + virtual R Dispatch_(const BindNode* node, Args... args) { + return DispatchDefault_(node, std::forward(args)...); + } + virtual R Dispatch_(const AttrStmtNode* node, Args... args) { + return DispatchDefault_(node, std::forward(args)...); + } + virtual R Dispatch_(const IfThenElseNode* node, Args... args) { + return DispatchDefault_(node, std::forward(args)...); + } + virtual R Dispatch_(const ForNode* node, Args... args) { + return DispatchDefault_(node, std::forward(args)...); + } + virtual R Dispatch_(const WhileNode* node, Args... args) { + return DispatchDefault_(node, std::forward(args)...); + } + virtual R Dispatch_(const ReturnNode* node, Args... args) { + return DispatchDefault_(node, std::forward(args)...); + } + virtual R Dispatch_(const BreakNode* node, Args... args) { + return DispatchDefault_(node, std::forward(args)...); + } + virtual R Dispatch_(const ContinueNode* node, Args... args) { + return DispatchDefault_(node, std::forward(args)...); + } + virtual R Dispatch_(const AllocBufferNode* node, Args... args) { + return DispatchDefault_(node, std::forward(args)...); + } + virtual R Dispatch_(const DeclBufferNode* node, Args... args) { + return DispatchDefault_(node, std::forward(args)...); + } + virtual R Dispatch_(const BufferStoreNode* node, Args... args) { + return DispatchDefault_(node, std::forward(args)...); + } + virtual R Dispatch_(const AssertStmtNode* node, Args... args) { + return DispatchDefault_(node, std::forward(args)...); + } + virtual R Dispatch_(const SeqStmtNode* node, Args... args) { + return DispatchDefault_(node, std::forward(args)...); + } + virtual R Dispatch_(const EvaluateNode* node, Args... args) { + return DispatchDefault_(node, std::forward(args)...); } - // Functions that can be overriden by subclass - virtual R VisitStmt_(const BindNode* op, Args... args) STMT_FUNCTOR_DEFAULT; - virtual R VisitStmt_(const AttrStmtNode* op, Args... args) STMT_FUNCTOR_DEFAULT; - virtual R VisitStmt_(const IfThenElseNode* op, Args... args) STMT_FUNCTOR_DEFAULT; - virtual R VisitStmt_(const ForNode* op, Args... args) STMT_FUNCTOR_DEFAULT; - virtual R VisitStmt_(const WhileNode* op, Args... args) STMT_FUNCTOR_DEFAULT; - virtual R VisitStmt_(const ReturnNode* op, Args... args) STMT_FUNCTOR_DEFAULT; - virtual R VisitStmt_(const BreakNode* op, Args... args) STMT_FUNCTOR_DEFAULT; - virtual R VisitStmt_(const ContinueNode* op, Args... args) STMT_FUNCTOR_DEFAULT; - virtual R VisitStmt_(const AllocBufferNode* op, Args... args) STMT_FUNCTOR_DEFAULT; - virtual R VisitStmt_(const DeclBufferNode* op, Args... args) STMT_FUNCTOR_DEFAULT; - virtual R VisitStmt_(const BufferStoreNode* op, Args... args) STMT_FUNCTOR_DEFAULT; - virtual R VisitStmt_(const AssertStmtNode* op, Args... args) STMT_FUNCTOR_DEFAULT; - virtual R VisitStmt_(const SeqStmtNode* op, Args... args) STMT_FUNCTOR_DEFAULT; - virtual R VisitStmt_(const EvaluateNode* op, Args... args) STMT_FUNCTOR_DEFAULT; - virtual R VisitStmt_(const ScopeIdDefStmtNode* op, Args... args) STMT_FUNCTOR_DEFAULT; - virtual R VisitStmt_(const tirx::TilePrimitiveCallNode* op, Args... args) STMT_FUNCTOR_DEFAULT; - virtual R VisitStmtDefault_(const ffi::Object* op, Args...) { - TVM_FFI_THROW(InternalError) << "Do not have a default for " << op->GetTypeKey(); + virtual R Dispatch_(const ScopeIdDefStmtNode* node, Args... args) { + return DispatchDefault_(node, std::forward(args)...); + } + virtual R Dispatch_(const tirx::TilePrimitiveCallNode* node, Args... args) { + return DispatchDefault_(node, std::forward(args)...); + } + /*! \brief Default behavior for statement hooks not overridden by a subclass. */ + virtual R DispatchDefault_(const ffi::Object* node, Args...) { + TVM_FFI_THROW(InternalError) << "Do not have a default for " << node->GetTypeKey(); TVM_FFI_UNREACHABLE(); } protected: + /*! \brief Dispatch table shared by this signature and its subclasses. */ using VTable = ObjectFunctor; - + /*! \brief Construct with a finalized table that outlives the functor. */ explicit StmtFunctor(const VTable* vtable) : vtable_(vtable) {} - - // Register inherited hooks in a fresh table before adding dialect nodes. + /*! \brief Register statement hooks in a fresh mutable table. */ static void InitVTable(VTable* vtable) { vtable->template SetDispatch( [](const ffi::ObjectRef& node, TSelf* self, Args... args) -> R { return self->DispatchDefault_(node.get(), std::forward(args)...); }); - IR_STMT_FUNCTOR_DISPATCH(BindNode); - IR_STMT_FUNCTOR_DISPATCH(AttrStmtNode); - IR_STMT_FUNCTOR_DISPATCH(IfThenElseNode); - IR_STMT_FUNCTOR_DISPATCH(ForNode); - IR_STMT_FUNCTOR_DISPATCH(WhileNode); - IR_STMT_FUNCTOR_DISPATCH(ReturnNode); - IR_STMT_FUNCTOR_DISPATCH(BreakNode); - IR_STMT_FUNCTOR_DISPATCH(ContinueNode); - IR_STMT_FUNCTOR_DISPATCH(AllocBufferNode); - IR_STMT_FUNCTOR_DISPATCH(DeclBufferNode); - IR_STMT_FUNCTOR_DISPATCH(AssertStmtNode); - IR_STMT_FUNCTOR_DISPATCH(SeqStmtNode); - IR_STMT_FUNCTOR_DISPATCH(EvaluateNode); - IR_STMT_FUNCTOR_DISPATCH(BufferStoreNode); - IR_STMT_FUNCTOR_DISPATCH(ScopeIdDefStmtNode); - IR_STMT_FUNCTOR_DISPATCH(tirx::TilePrimitiveCallNode); + SetDispatch(vtable); + SetDispatch(vtable); + SetDispatch(vtable); + SetDispatch(vtable); + SetDispatch(vtable); + SetDispatch(vtable); + SetDispatch(vtable); + SetDispatch(vtable); + SetDispatch(vtable); + SetDispatch(vtable); + SetDispatch(vtable); + SetDispatch(vtable); + SetDispatch(vtable); + SetDispatch(vtable); + SetDispatch(vtable); + SetDispatch(vtable); } - + /*! \brief Register an additional node hook implemented by Self. */ template static void SetDispatch(VTable* vtable) { - vtable->template SetDispatch([](const ffi::ObjectRef& node, TSelf* self, Args... args) { - return static_cast(self)->VisitStmt_(static_cast(node.get()), - std::forward(args)...); - }); - } - - private: - static const VTable* GlobalVTable() { - static const VTable table = [] { - VTable table; - InitVTable(&table); - table.Finalize(); - return table; - }(); - return &table; + vtable->template SetDispatch( + [](const ffi::ObjectRef& node, TSelf* self, Args... args) -> R { + return static_cast(self)->Dispatch_(static_cast(node.get()), + std::forward(args)...); + }); } - const VTable* const vtable_; -}; - private: static const VTable* GlobalVTable() { static const VTable table = [] { diff --git a/python/tvm/ir/json_compact.py b/python/tvm/ir/json_compact.py index 63d882d8b39f..b5de5b52f64b 100644 --- a/python/tvm/ir/json_compact.py +++ b/python/tvm/ir/json_compact.py @@ -19,6 +19,7 @@ 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", @@ -124,6 +125,11 @@ def upgrade_json(json_str): # 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", []): + if node.get("type") == "tirx.BufferRegion": + # TensorRegion keeps the inherited type/span and range references; + # only the buffer field became the shared expression source. + fields = node.get("data", {}) + fields["source"] = fields.pop("buffer") 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/s_tir/analysis/__init__.py b/python/tvm/s_tir/analysis/__init__.py index 817e41472b3e..24732261a5aa 100644 --- a/python/tvm/s_tir/analysis/__init__.py +++ b/python/tvm/s_tir/analysis/__init__.py @@ -23,7 +23,6 @@ import tvm from tvm.ir import IRModule, TensorRegion from tvm.tirx.expr import Var -from tvm.tirx.stmt import BufferRegion from tvm.s_tir import SBlock from tvm.tirx import Buffer, Stmt diff --git a/python/tvm/s_tir/stmt.py b/python/tvm/s_tir/stmt.py index 329dc3c6e219..3c84f5a59c87 100644 --- a/python/tvm/s_tir/stmt.py +++ b/python/tvm/s_tir/stmt.py @@ -20,11 +20,11 @@ import tvm_ffi -from tvm.ir import Expr, Span +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 BufferRegion, Stmt, _normalize_legacy_stmt +from tvm.tirx.stmt import Stmt, _normalize_legacy_stmt from . import _ffi_api @@ -38,14 +38,14 @@ class MatchBufferRegion(Object, Scriptable): buffer : Buffer The target buffer - source : BufferRegion + source : TensorRegion The region of source buffer """ buffer: Buffer - source: BufferRegion + source: TensorRegion - def __init__(self, buffer: Buffer, source: BufferRegion) -> None: + def __init__(self, buffer: Buffer, source: TensorRegion) -> None: self.__init_handle_by_constructor__( _ffi_api.MatchBufferRegion, buffer, @@ -62,10 +62,10 @@ class SBlock(Stmt): iter_vars : List[IterVar] The block Variable. - reads : List[BufferRegion] + reads : List[TensorRegion] The read buffer regions of the block. - writes: List[BufferRegion] + writes: List[TensorRegion] The write buffer regions of the block. name_hint: str @@ -91,8 +91,8 @@ class SBlock(Stmt): """ iter_vars: list[IterVar] - reads: list[BufferRegion] - writes: list[BufferRegion] + reads: list[TensorRegion] + writes: list[TensorRegion] name_hint: str body: Stmt init: Stmt | None @@ -104,8 +104,8 @@ class SBlock(Stmt): def __init__( self, iter_vars: list[IterVar], - reads: list[BufferRegion], - writes: list[BufferRegion], + reads: list[TensorRegion], + writes: list[TensorRegion], name_hint: str, body: Stmt, init: Stmt | None = None, diff --git a/python/tvm/tirx/stmt.py b/python/tvm/tirx/stmt.py index 2787a7b2b897..f609d4439fce 100644 --- a/python/tvm/tirx/stmt.py +++ b/python/tvm/tirx/stmt.py @@ -33,7 +33,7 @@ import tvm_ffi -from tvm.ir import Expr, Range, Span, Type +from tvm.ir import Expr, Range, Span, TensorRegion, Type from tvm.runtime import Object, Scriptable from . import _ffi_api diff --git a/src/relax/analysis/tir_op_pattern_kind.cc b/src/relax/analysis/tir_op_pattern_kind.cc index f49a800eacd6..f9f428c33ab7 100644 --- a/src/relax/analysis/tir_op_pattern_kind.cc +++ b/src/relax/analysis/tir_op_pattern_kind.cc @@ -52,8 +52,8 @@ class PatternKindAnalyzer : public s_tir::StmtExprVisitor { private: bool IsOutputBlock(const s_tir::SBlockNode* block) { - for (const BufferRegion& write_region : block->writes) { - if (param_buffers_.count(write_region->buffer)) { + for (const TensorRegion& write_region : block->writes) { + if (param_buffers_.count(write_region->source.as_or_throw())) { return true; } } diff --git a/src/relax/transform/fuse_tir.cc b/src/relax/transform/fuse_tir.cc index e2e9a4287d1f..925bbda6932a 100644 --- a/src/relax/transform/fuse_tir.cc +++ b/src/relax/transform/fuse_tir.cc @@ -195,8 +195,8 @@ class FuseTIRBufferSubstitutor : public s_tir::StmtExprMutator { s_tir::SBlock block = s_tir::StmtExprMutator::Mutate_(op, inplace_mode) .ValueOrUnchanged(ffi::GetRef(op)) .as_or_throw(); - ffi::Array reads = UnionAccessRegion(block->reads); - ffi::Array writes = UnionAccessRegion(block->writes); + ffi::Array reads = UnionAccessRegion(block->reads); + ffi::Array writes = UnionAccessRegion(block->writes); if (!reads.same_as(block->reads) || !writes.same_as(block->writes)) { auto* n = block.CopyOnWrite(); n->reads = std::move(reads); diff --git a/src/relax/transform/split_call_tir_by_pattern.cc b/src/relax/transform/split_call_tir_by_pattern.cc index 4d29bd47267f..54bf928e7121 100644 --- a/src/relax/transform/split_call_tir_by_pattern.cc +++ b/src/relax/transform/split_call_tir_by_pattern.cc @@ -274,7 +274,7 @@ class ForMatcher : public TensorizeComparator { return Dispatch(op->body, rhs->body); } - bool VisitStmt_(const s_tir::SBlockNode* op, const Stmt& other) final { + bool Dispatch_(const s_tir::SBlockNode* op, const Stmt& other) final { const auto* rhs = other.as(); // Check block equality. // All iter vars and buffer regions including the order should match. @@ -303,7 +303,7 @@ class ForMatcher : public TensorizeComparator { return Dispatch(op->body, rhs->body); } - bool VisitStmt_(const s_tir::SBlockRealizeNode* op, const Stmt& other) final { + bool Dispatch_(const s_tir::SBlockRealizeNode* op, const Stmt& other) final { const auto* rhs = other.as(); // Only allow trivial bindings for (size_t i = 0; i < op->iter_values.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 2331078005cd..b293b0d41dea 100644 --- a/src/s_tir/analysis/sblock_access_region_detector.cc +++ b/src/s_tir/analysis/sblock_access_region_detector.cc @@ -436,7 +436,7 @@ void BlockReadWriteDetector::Update(std::vector* buffers, auto it = match_buffers_.find(buffer.get()); if (it != match_buffers_.end()) { const s_tir::MatchBufferRegion& match_buffer = it->second; - buffer = match_buffer->source->buffer; + buffer = match_buffer->source->source.as_or_throw(); region = ConvertMatchedRegion(match_buffer, std::move(region)); } TVM_FFI_ICHECK_EQ(buffers->size(), regions->size()) @@ -497,7 +497,7 @@ void BlockReadWriteDetector::UpdateOpaque(const Var& buffer_var) { } } -ffi::Array> GetSBlockAccessRegion( +ffi::Array> GetSBlockAccessRegion( const s_tir::SBlock& block, const ffi::Map& buffer_var_map) { auto detector = ffi::make_object(buffer_var_map); detector->operator()(block); @@ -514,7 +514,7 @@ ffi::Array> GetSBlockAccessRegion( return {reads, writes, opaques}; } -ffi::Array> GetSBlockReadWriteRegion( +ffi::Array> GetSBlockReadWriteRegion( const s_tir::SBlock& block, const ffi::Map& buffer_var_map) { auto detector = ffi::make_object(buffer_var_map); detector->operator()(block); 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 9ec400dbc7f0..0d73fb07655e 100644 --- a/src/s_tir/analysis/sblock_buffer_access_lca_detector.cc +++ b/src/s_tir/analysis/sblock_buffer_access_lca_detector.cc @@ -137,7 +137,8 @@ class LCADetector : public s_tir::StmtExprVisitor { // Update match_buffers for (const s_tir::MatchBufferRegion& match_buffer : block->match_buffers) { - UpdateBufferLCA(match_buffer->source->buffer.get(), ancestor_scopes_.back()); + UpdateBufferLCA(match_buffer->source->source.as_or_throw().get(), + ancestor_scopes_.back()); match_buffers_.insert(match_buffer->buffer.get()); } @@ -269,7 +270,7 @@ class LCADetector : public s_tir::StmtExprVisitor { // Declared regions carry bounds, not opaque runtime accesses. ffi::Optional Visit_(const TensorRegionNode* op) final { - if (!op->source.as()) return StmtExprVisitor::Visit_(op); + if (!op->source.as()) 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)); diff --git a/src/s_tir/analysis/verify_well_formed.cc b/src/s_tir/analysis/verify_well_formed.cc index 24381365d8b7..3e00f8e9e1ba 100644 --- a/src/s_tir/analysis/verify_well_formed.cc +++ b/src/s_tir/analysis/verify_well_formed.cc @@ -27,7 +27,6 @@ namespace tvm { namespace s_tir { -using tirx::BufferRegion; using tirx::ForNode; using tirx::PrimFunc; @@ -95,7 +94,7 @@ class BlockVarAccessVerifier : public StmtExprVisitor { // Step 0. Skip block iter var's domain // Step 1. Visit read/write regions - auto fvisit_buffer_region = [this](const BufferRegion& s) -> ffi::Optional { + auto fvisit_buffer_region = [this](const TensorRegion& s) -> ffi::Optional { 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)); diff --git a/src/s_tir/ir/data_type_rewriter.cc b/src/s_tir/ir/data_type_rewriter.cc index d70e392489d8..b92fe1b65ab9 100644 --- a/src/s_tir/ir/data_type_rewriter.cc +++ b/src/s_tir/ir/data_type_rewriter.cc @@ -119,14 +119,14 @@ UnchangedOr IndexDataTypeNormalizer::Mutate_(const SBlockNode* op, Inplace .as_or_throw>() .ValueOrUnchanged(match->buffer); }); - BufferRegion source = VisitBufferRegion(match->source); + TensorRegion source = VisitBufferRegion(match->source); if (buffer.same_as(match->buffer) && source.same_as(match->source)) return match; return MatchBufferRegion(buffer, source); }); - ffi::Array new_reads = op->reads.Map( - [this](const BufferRegion& buffer_region) { return VisitBufferRegion(buffer_region); }); - ffi::Array new_writes = op->writes.Map( - [this](const BufferRegion& buffer_region) { return VisitBufferRegion(buffer_region); }); + ffi::Array new_reads = op->reads.Map( + [this](const TensorRegion& buffer_region) { return VisitBufferRegion(buffer_region); }); + ffi::Array new_writes = op->writes.Map( + [this](const TensorRegion& buffer_region) { return VisitBufferRegion(buffer_region); }); ffi::Array new_iter_vars = op->iter_vars.Map([this](const IterVar& iter_var) { return VisitIterVar(iter_var); }); ffi::Optional new_init = std::nullopt; @@ -212,10 +212,11 @@ IterVar IndexDataTypeNormalizer::VisitIterVar(const IterVar& iter_var) { return iter_var; } -BufferRegion IndexDataTypeNormalizer::VisitBufferRegion(const BufferRegion& buffer_region) { - BufferVar remapped_buffer = this->Mutate(buffer_region->buffer, InplaceMode::kDisallow) - .as_or_throw>() - .ValueOrUnchanged(buffer_region->buffer); +TensorRegion IndexDataTypeNormalizer::VisitBufferRegion(const TensorRegion& buffer_region) { + BufferVar remapped_buffer = + this->Mutate(buffer_region->source.as_or_throw(), InplaceMode::kDisallow) + .as_or_throw>() + .ValueOrUnchanged(buffer_region->source.as_or_throw()); bool is_enabled = this->is_enabled_; this->is_enabled_ = true; @@ -226,7 +227,7 @@ BufferRegion IndexDataTypeNormalizer::VisitBufferRegion(const BufferRegion& buff }); this->is_enabled_ = is_enabled; - if (!remapped_buffer.same_as(buffer_region->buffer) || + if (!remapped_buffer.same_as(buffer_region->source.as_or_throw()) || !new_region.same_as(buffer_region->region)) { return BufferRegion(remapped_buffer, new_region); } else { diff --git a/src/s_tir/ir/data_type_rewriter.h b/src/s_tir/ir/data_type_rewriter.h index a5a74feaf0c8..5ccfb2cb7c67 100644 --- a/src/s_tir/ir/data_type_rewriter.h +++ b/src/s_tir/ir/data_type_rewriter.h @@ -61,7 +61,7 @@ class IndexDataTypeNormalizer : public tirx::IndexDataTypeNormalizer { ffi::Map VisitBlockAnnotations( const ffi::Map& annotations); tirx::IterVar VisitIterVar(const tirx::IterVar& iter_var); - tirx::BufferRegion VisitBufferRegion(const tirx::BufferRegion& buffer_region); + TensorRegion VisitBufferRegion(const TensorRegion& buffer_region); }; } // namespace s_tir diff --git a/src/s_tir/ir/ir_mutator_with_analyzer.h b/src/s_tir/ir/ir_mutator_with_analyzer.h index 1ca5e52bed59..46a7648258a5 100644 --- a/src/s_tir/ir/ir_mutator_with_analyzer.h +++ b/src/s_tir/ir/ir_mutator_with_analyzer.h @@ -22,7 +22,7 @@ #include -#include "../../tirx/ir_mutator_with_analyzer.h" +#include "../../tirx/ir/ir_mutator_with_analyzer.h" namespace tvm { namespace s_tir { diff --git a/src/s_tir/ir/ir_visitor_with_analyzer.h b/src/s_tir/ir/ir_visitor_with_analyzer.h index 4339077f84d3..cf0d11813a22 100644 --- a/src/s_tir/ir/ir_visitor_with_analyzer.h +++ b/src/s_tir/ir/ir_visitor_with_analyzer.h @@ -22,7 +22,7 @@ #include -#include "../../tirx/ir_visitor_with_analyzer.h" +#include "../../tirx/ir/ir_visitor_with_analyzer.h" namespace tvm { namespace s_tir { diff --git a/src/s_tir/ir/tir_visitor_with_path.cc b/src/s_tir/ir/tir_visitor_with_path.cc index 0a27c52f272c..af69a4b4feaa 100644 --- a/src/s_tir/ir/tir_visitor_with_path.cc +++ b/src/s_tir/ir/tir_visitor_with_path.cc @@ -26,7 +26,7 @@ namespace s_tir { using namespace tirx; using AccessPath = ffi::reflection::AccessPath; -void TIRVisitorWithPath::VisitStmt_(const SBlockNode* op, AccessPath path) { +void TIRVisitorWithPath::Dispatch_(const SBlockNode* op, AccessPath path) { std::vector, DefContext, DefContext>> context; { @@ -71,7 +71,7 @@ void TIRVisitorWithPath::VisitStmt_(const SBlockNode* op, AccessPath path) { while (context.size()) context.pop_back(); } -void TIRVisitorWithPath::VisitStmt_(const SBlockRealizeNode* op, AccessPath path) { +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")); diff --git a/src/s_tir/ir/tir_visitor_with_path.h b/src/s_tir/ir/tir_visitor_with_path.h index 25afba3078e0..3da346860772 100644 --- a/src/s_tir/ir/tir_visitor_with_path.h +++ b/src/s_tir/ir/tir_visitor_with_path.h @@ -30,19 +30,19 @@ class TIRVisitorWithPath : public tirx::TIRVisitorWithPath { TVM_DEFINE_OBJECT_FUNCTOR_DEFAULT_CONSTRUCTOR(TIRVisitorWithPath, Parent) protected: using AccessPath = ffi::reflection::AccessPath; - using Parent::VisitStmt_; - virtual void VisitStmt_(const SBlockNode* op, AccessPath path); - virtual void VisitStmt_(const SBlockRealizeNode* op, AccessPath path); + 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); vtable->SetDispatch( [](const ffi::ObjectRef& node, StmtVisitor* self, AccessPath path) { - static_cast(self)->VisitStmt_( + static_cast(self)->Dispatch_( static_cast(node.get()), path); }); vtable->SetDispatch( [](const ffi::ObjectRef& node, StmtVisitor* self, AccessPath path) { - static_cast(self)->VisitStmt_( + static_cast(self)->Dispatch_( static_cast(node.get()), path); }); } 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 c6258227e456..bf5e1d4f1a87 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 @@ -868,7 +868,7 @@ ffi::Optional MultiLevelTilingTensorCoreNode::TransformWithTensorIntrin( visited_buffers.insert(lhs_buffer); // Refresh block pointer (block sref is not invalidated) block = TVM_SREF_TO_SBLOCK(block_sref); - const tirx::BufferRegion& reindexed_buffer_region = s_tir::GetNthAccessBufferRegion( + const tvm::TensorRegion& reindexed_buffer_region = s_tir::GetNthAccessBufferRegion( state->sch->state(), ffi::GetRef(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); diff --git a/src/s_tir/stmt.cc b/src/s_tir/stmt.cc index 2ec97dd3f361..98753c031b2b 100644 --- a/src/s_tir/stmt.cc +++ b/src/s_tir/stmt.cc @@ -53,7 +53,7 @@ TVMFFIAny MatchBufferRegionMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyVi mutator->WithDefRegionKind(kTVMFFIDefRegionKindSimple, [&]() { return mutator->MutateExpected(self->buffer); })); - TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, mapped_source, + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, mapped_source, mutator->MutateExpected(self->source)); if (mapped_buffer.UnchangedOrSameAs(self->buffer) && mapped_source.UnchangedOrSameAs(self->source)) { @@ -76,7 +76,7 @@ TVMFFIAny MatchBufferRegionMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator ffi::InplaceMode::kAllow); })); TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN( - ffi::UnchangedOr, mapped_source, + ffi::UnchangedOr, mapped_source, mutator->MutateExpected(self->source, ffi::InplaceMode::kAllow)); if (mapped_buffer.UnchangedOrSameAs(self->buffer) && mapped_source.UnchangedOrSameAs(self->source)) { @@ -120,9 +120,9 @@ TVMFFIAny SBlockMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) n TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr>, mapped_match_buffers, mutator->MutateExpected(self->match_buffers)); - TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr>, mapped_reads, + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr>, mapped_reads, mutator->MutateExpected(self->reads)); - TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr>, mapped_writes, + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr>, mapped_writes, mutator->MutateExpected(self->writes)); using AnnotationMap = ffi::Map; TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, mapped_annotations, @@ -172,10 +172,10 @@ TVMFFIAny SBlockMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator, TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN( ffi::UnchangedOr>, mapped_match_buffers, mutator->MutateExpected(self->match_buffers, ffi::InplaceMode::kAllow)); - TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr>, mapped_reads, + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr>, mapped_reads, mutator->MutateExpected(self->reads, ffi::InplaceMode::kAllow)); TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN( - ffi::UnchangedOr>, mapped_writes, + ffi::UnchangedOr>, mapped_writes, mutator->MutateExpected(self->writes, ffi::InplaceMode::kAllow)); using AnnotationMap = ffi::Map; TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN( @@ -269,8 +269,10 @@ TVMFFIAny SBlockRealizeMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator, } // namespace // MatchBufferRegion -MatchBufferRegion::MatchBufferRegion(BufferVar buffer, BufferRegion source) { - const BufferVar& source_buffer = source->buffer; +MatchBufferRegion::MatchBufferRegion(BufferVar buffer, TensorRegion source) { + const BufferVar& source_buffer = source->source.as_or_throw(); + 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()) @@ -322,17 +324,24 @@ TVM_FFI_STATIC_INIT_BLOCK() { .attr(refl::type_attr::kStructuralMaybeInplaceMutate, reinterpret_cast(&MatchBufferRegionMaybeInplaceMutate)); - refl::GlobalDef().def("s_tir.MatchBufferRegion", [](BufferVar buffer, BufferRegion source) { + refl::GlobalDef().def("s_tir.MatchBufferRegion", [](BufferVar buffer, TensorRegion source) { return MatchBufferRegion(buffer, source); }); } // Block -SBlock::SBlock(ffi::Array iter_vars, ffi::Array reads, - ffi::Array writes, ffi::String name_hint, Stmt body, +SBlock::SBlock(ffi::Array iter_vars, ffi::Array reads, + ffi::Array writes, ffi::String name_hint, Stmt body, ffi::Optional init, ffi::Array alloc_buffers, ffi::Array match_buffers, ffi::Map annotations, Span span) { + for (const auto& regions : {reads, writes}) { + for (const TensorRegion& region : regions) { + const auto buffer = region->source.as_or_throw(); + TVM_FFI_ICHECK_EQ(buffer->shape.size(), region->region.size()) + << "SBlock region must match its buffer rank"; + } + } ffi::ObjectPtr node = ffi::make_object(); node->iter_vars = std::move(iter_vars); node->reads = std::move(reads); @@ -372,8 +381,8 @@ TVM_FFI_STATIC_INIT_BLOCK() { reinterpret_cast(&SBlockMaybeInplaceMutate)); refl::GlobalDef().def("s_tir.SBlock", - [](ffi::Array iter_vars, ffi::Array reads, - ffi::Array writes, ffi::String name_hint, Stmt body, + [](ffi::Array iter_vars, ffi::Array reads, + ffi::Array writes, ffi::String name_hint, Stmt body, ffi::Optional init, ffi::Array alloc_buffers, ffi::Array match_buffers, ffi::Map annotations, Span span) { diff --git a/src/s_tir/stmt_functor.cc b/src/s_tir/stmt_functor.cc index 7846c595430d..91a20e902425 100644 --- a/src/s_tir/stmt_functor.cc +++ b/src/s_tir/stmt_functor.cc @@ -63,10 +63,10 @@ ffi::Optional StmtExprVisitor::VisitBlock(tirx::StmtExprVisitor* 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 BufferRegion& region : op->reads) { + for (const TensorRegion& region : op->reads) { TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->Visit(region)); } - for (const BufferRegion& region : op->writes) { + for (const TensorRegion& region : op->writes) { TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->Visit(region)); } if (op->init.has_value()) { @@ -131,9 +131,9 @@ UnchangedOr StmtExprMutator::MutateBlock(tirx::StmtExprMutator* mutator, c auto match_buffers = mutator->Mutate(op->match_buffers, inplace_mode) .as_or_throw>>(); auto reads = - mutator->Mutate(op->reads, inplace_mode).as_or_throw>>(); + mutator->Mutate(op->reads, inplace_mode).as_or_throw>>(); auto writes = mutator->Mutate(op->writes, inplace_mode) - .as_or_throw>>(); + .as_or_throw>>(); auto init = mutator->Mutate(op->init, inplace_mode).as_or_throw>>(); auto body = mutator->Mutate(op->body, inplace_mode); diff --git a/src/s_tir/transform/default_gpu_schedule.cc b/src/s_tir/transform/default_gpu_schedule.cc index 353326c21d67..35686953f328 100644 --- a/src/s_tir/transform/default_gpu_schedule.cc +++ b/src/s_tir/transform/default_gpu_schedule.cc @@ -146,8 +146,8 @@ tirx::PrimFunc WrapBareSBlockBody(const tirx::PrimFunc& func) { tirx::Stmt for_stmt = tirx::For(loop_var.as_or_throw(), zero, one, tirx::ForKind::kSerial, inner_realize); s_tir::SBlock root_block(/*iter_vars=*/ffi::Array{}, - /*reads=*/ffi::Array{}, - /*writes=*/ffi::Array{}, + /*reads=*/ffi::Array{}, + /*writes=*/ffi::Array{}, /*name_hint=*/"root", /*body=*/for_stmt); s_tir::SBlockRealize root_realize(/*iter_values=*/ffi::Array{}, /*predicate=*/IntImm::Bool(true), root_block); diff --git a/src/s_tir/transform/ir_utils.cc b/src/s_tir/transform/ir_utils.cc index f9da7ddab09d..adb3e54c5f11 100644 --- a/src/s_tir/transform/ir_utils.cc +++ b/src/s_tir/transform/ir_utils.cc @@ -51,9 +51,11 @@ class SIRConvertSSA final : public tirx::IRConvertSSA { if (!var.same_as(iter->var)) iter.CopyOnWrite()->var = var.as_or_throw(); return iter; }); - auto remap_region = [&](BufferRegion region) { - BufferVar buffer = GetRemappedBuffer(region->buffer); - if (!buffer.same_as(region->buffer)) region.CopyOnWrite()->buffer = buffer; + auto remap_region = [&](TensorRegion region) { + BufferVar buffer = GetRemappedBuffer(region->source.as_or_throw()); + if (!buffer.same_as(region->source.as_or_throw())) { + region.CopyOnWrite()->source = buffer.var(); + } return region; }; auto reads = block->reads.Map(remap_region); @@ -101,7 +103,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { ffi::Array ConvertIndices(const MatchBufferRegion& match_buffer, const ffi::Array& indices) { const BufferVar& target = match_buffer->buffer; - const BufferRegion& source = match_buffer->source; + const TensorRegion& source = match_buffer->source; TVM_FFI_ICHECK_EQ(indices.size(), target->shape.size()); arith::Analyzer analyzer; @@ -123,7 +125,7 @@ ffi::Array ConvertIndices(const MatchBufferRegion& match_buffer, Region ConvertRegion(const MatchBufferRegion& match_buffer, const Region& region) { const BufferVar& target = match_buffer->buffer; - const BufferRegion& source = match_buffer->source; + const TensorRegion& source = match_buffer->source; TVM_FFI_ICHECK_EQ(region.size(), target->shape.size()); arith::Analyzer analyzer; @@ -162,7 +164,8 @@ class StorageAlignCollector : public StmtExprVisitor { if (it != op->annotations.end()) { auto annotation = (*it).second.as_or_throw(); for (const auto& item : annotation) { - storage_align_[op->writes[item.get<0>()]->buffer.var()].push_back(item); + storage_align_[op->writes[item.get<0>()]->source.as_or_throw().var()].push_back( + item); } } return StmtExprVisitor::Visit_(op); diff --git a/src/tirx/analysis/verify_tirx_well_formed.cc b/src/tirx/analysis/verify_tirx_well_formed.cc index 42437656e849..6a95a2976a57 100644 --- a/src/tirx/analysis/verify_tirx_well_formed.cc +++ b/src/tirx/analysis/verify_tirx_well_formed.cc @@ -50,7 +50,7 @@ class ExecScopeVerifier : public Verifier { private: using Verifier::Visit; - void VisitStmtDefault_(const ffi::Object* op, ffi::reflection::AccessPath path) override { + void DispatchDefault_(const ffi::Object* op, ffi::reflection::AccessPath path) override { Verify(false) << "TIRxError: " << op->GetTypeKey() << " is not allowed in tirx=True mode at " << path; } @@ -129,7 +129,7 @@ class LayoutVerifier : public Verifier { private: using Verifier::Visit; - void VisitStmtDefault_(const ffi::Object* op, ffi::reflection::AccessPath path) override { + void DispatchDefault_(const ffi::Object* op, ffi::reflection::AccessPath path) override { Verify(false) << "TIRxError: " << op->GetTypeKey() << " is not allowed in tirx=True mode at " << path; } @@ -142,7 +142,7 @@ class AsyncStructsVerifier : public Verifier { private: using Verifier::Visit; - void VisitStmtDefault_(const ffi::Object* op, ffi::reflection::AccessPath path) override { + void DispatchDefault_(const ffi::Object* op, ffi::reflection::AccessPath path) override { Verify(false) << "TIRxError: " << op->GetTypeKey() << " is not allowed in tirx=True mode at " << path; } @@ -155,7 +155,7 @@ class DeviceFuncVerifier : public Verifier { private: using Verifier::Visit; - void VisitStmtDefault_(const ffi::Object* op, ffi::reflection::AccessPath path) override { + void DispatchDefault_(const ffi::Object* op, ffi::reflection::AccessPath path) override { Verify(false) << "TIRxError: " << op->GetTypeKey() << " is not allowed in tirx=True mode at " << path; } diff --git a/src/tirx/ir/stmt.cc b/src/tirx/ir/stmt.cc index 8d7823594a85..179f546cd234 100644 --- a/src/tirx/ir/stmt.cc +++ b/src/tirx/ir/stmt.cc @@ -721,50 +721,6 @@ TVMFFIAny BufferRegionTypeMaybeInplaceMutate(ffi::StructuralMutatorObj*, ffi::An return ffi::Unchanged().CopyToTVMFFIAny(); } -TVMFFIAny BufferRegionVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept { - const BufferRegionNode* self = - ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value); - TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->buffer)); - TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->region)); - return ffi::AnyView(nullptr).CopyToTVMFFIAny(); -} - -TVMFFIAny BufferRegionMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept { - const BufferRegionNode* self = - ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value); - TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, mapped_buffer, - mutator->MutateExpected(self->buffer)); - TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr>, mapped_region, - mutator->MutateExpected(self->region)); - if (mapped_buffer.UnchangedOrSameAs(self->buffer) && - mapped_region.UnchangedOrSameAs(self->region)) { - return ffi::Unchanged().CopyToTVMFFIAny(); - } - ffi::ObjectPtr copy = ffi::make_object(*self); - copy->buffer = std::move(mapped_buffer).ValueOrUnchanged(std::move(copy->buffer)); - copy->region = std::move(mapped_region).ValueOrUnchanged(std::move(copy->region)); - return ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(copy))); -} - -TVMFFIAny BufferRegionMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator, - ffi::AnyView value) noexcept { - BufferRegionNode* self = const_cast( - ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value)); - TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN( - ffi::UnchangedOr, mapped_buffer, - mutator->MutateExpected(self->buffer, ffi::InplaceMode::kAllow)); - TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN( - ffi::UnchangedOr>, mapped_region, - mutator->MutateExpected(self->region, ffi::InplaceMode::kAllow)); - if (mapped_buffer.UnchangedOrSameAs(self->buffer) && - mapped_region.UnchangedOrSameAs(self->region)) { - return ffi::Unchanged().CopyToTVMFFIAny(); - } - if (!mapped_buffer.IsUnchanged()) self->buffer = std::move(mapped_buffer).ValueUnchecked(); - if (!mapped_region.IsUnchanged()) self->region = std::move(mapped_region).ValueUnchecked(); - return ffi::Unchanged().CopyToTVMFFIAny(); -} - TVMFFIAny ScopeIdDefStmtVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept { const ScopeIdDefStmtNode* self = ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value); diff --git a/src/tirx/ir/stmt_functor.cc b/src/tirx/ir/stmt_functor.cc index c72d63b50d57..b1d7884012cd 100644 --- a/src/tirx/ir/stmt_functor.cc +++ b/src/tirx/ir/stmt_functor.cc @@ -58,7 +58,6 @@ void StmtExprVisitor::InitVTable(VTable* vtable) { SetDispatch(vtable); SetDispatch(vtable); SetDispatch(vtable); - SetDispatch(vtable); } ffi::Optional StmtExprVisitor::Visit_(const VarNode* op) { return std::nullopt; } @@ -290,7 +289,6 @@ void StmtExprMutator::InitVTable(VTable* vtable) { SetDispatch(vtable); SetDispatch(vtable); SetDispatch(vtable); - SetDispatch(vtable); } UnchangedOr StmtExprMutator::Mutate_(const BindNode* op, InplaceMode inplace_mode) { @@ -493,24 +491,6 @@ UnchangedOr StmtExprMutator::Mutate_(const BufferStoreNode* op, InplaceMod return Stmt(std::move(copy)); } -UnchangedOr StmtExprMutator::Mutate_(const BufferRegionNode* op, InplaceMode inplace_mode) { - auto buffer = Mutate(op->buffer, inplace_mode).as_or_throw>(); - auto region = Mutate(op->region, inplace_mode).as_or_throw>>(); - if (buffer.UnchangedOrSameAs(op->buffer) && region.UnchangedOrSameAs(op->region)) { - return ffi::Unchanged(); - } - if (inplace_mode == InplaceMode::kAllow) { - auto* writable = const_cast(op); - if (!buffer.IsUnchanged()) writable->buffer = std::move(buffer).ValueUnchecked(); - if (!region.IsUnchanged()) writable->region = std::move(region).ValueUnchecked(); - return ffi::Unchanged(); - } - auto copy = ffi::make_object(*op); - if (!buffer.IsUnchanged()) copy->buffer = std::move(buffer).ValueUnchecked(); - if (!region.IsUnchanged()) copy->region = std::move(region).ValueUnchecked(); - return Expr(std::move(copy)); -} - UnchangedOr StmtExprMutator::Mutate_(const SeqStmtNode* op, InplaceMode inplace_mode) { return detail::MutateSeqStmt(op, inplace_mode, [this](ffi::AnyView element, InplaceMode mode) { return Mutate(element, mode).as_or_throw>(); diff --git a/src/tirx/ir/tir_visitor_with_path.cc b/src/tirx/ir/tir_visitor_with_path.cc index 9416ea44ac97..c37c443ed374 100644 --- a/src/tirx/ir/tir_visitor_with_path.cc +++ b/src/tirx/ir/tir_visitor_with_path.cc @@ -257,7 +257,7 @@ void TIRVisitorWithPath::Dispatch_(const EvaluateNode* op, AccessPath path) { Visit(op->value, path->Attr("value")); } -void TIRVisitorWithPath::VisitStmt_(const tirx::TilePrimitiveCallNode* op, AccessPath path) { +void TIRVisitorWithPath::Dispatch_(const tirx::TilePrimitiveCallNode* op, AccessPath path) { for (size_t i = 0; i < op->args.size(); i++) { if (op->args[i] == nullptr) { continue; diff --git a/src/tirx/ir/tir_visitor_with_path.h b/src/tirx/ir/tir_visitor_with_path.h index 9f5597128a01..674a3e6784cb 100644 --- a/src/tirx/ir/tir_visitor_with_path.h +++ b/src/tirx/ir/tir_visitor_with_path.h @@ -97,7 +97,7 @@ class TIRVisitorWithPath : protected ExprFunctorGetTypeKey() << " at " << path; } diff --git a/src/tirx/script/builder/frame.cc b/src/tirx/script/builder/frame.cc index 6ed127670902..9ae7e1915477 100644 --- a/src/tirx/script/builder/frame.cc +++ b/src/tirx/script/builder/frame.cc @@ -214,10 +214,9 @@ void SBlockFrameNode::ExitWithScope() { if (int detect_access = (!reads.has_value()) | (!writes.has_value() << 1)) { attrs.Set("tirx.script_parsing_detect_access", tvm::IntImm::Int64(detect_access)); } - tvm::s_tir::SBlock block(iter_vars, reads.value_or(ffi::Array()), - writes.value_or(ffi::Array()), name, - AsStmt(stmts), init, tir_alloc_buffers, match_buffers, attrs, - tvm::Span()); + tvm::s_tir::SBlock block(iter_vars, reads.value_or(ffi::Array()), + writes.value_or(ffi::Array()), name, AsStmt(stmts), + init, tir_alloc_buffers, match_buffers, attrs, tvm::Span()); if (no_realize) { TVM_FFI_CHECK(iter_values.empty(), ValueError) << "Block bindings are not allowed when `no_realize=True`"; diff --git a/src/tirx/script/builder/ir.cc b/src/tirx/script/builder/ir.cc index b776e48c6b20..97623a658596 100644 --- a/src/tirx/script/builder/ir.cc +++ b/src/tirx/script/builder/ir.cc @@ -163,7 +163,7 @@ BufferVar MatchBuffer(ffi::ObjectRef param, ffi::Array shape, PrimType } else if (const auto* buffer_region = param.as()) { SBlockFrame frame = FindSBlockFrame("T.match_buffer"); frame->match_buffers.push_back( - tvm::s_tir::MatchBufferRegion(buffer, ffi::GetRef(buffer_region))); + tvm::s_tir::MatchBufferRegion(buffer, ffi::GetRef(buffer_region))); } else { TVM_FFI_THROW(InternalError) << "ValueError: Unexpected type for TIR MatchBuffer."; } diff --git a/src/tirx/transform/stmt_simplify.h b/src/tirx/transform/stmt_simplify.h index 83e04ff97ff1..78ead6accfb2 100644 --- a/src/tirx/transform/stmt_simplify.h +++ b/src/tirx/transform/stmt_simplify.h @@ -27,7 +27,7 @@ #include #include -#include "../ir_mutator_with_analyzer.h" +#include "../ir/ir_mutator_with_analyzer.h" namespace tvm { namespace arith { diff --git a/tests/cpp/ir_functor_test.cc b/tests/cpp/ir_functor_test.cc index d21225ee3014..3ba10ebe20b3 100644 --- a/tests/cpp/ir_functor_test.cc +++ b/tests/cpp/ir_functor_test.cc @@ -234,7 +234,7 @@ TEST(IRF, StmtVisitor) { tirx::Var buf_var("b", PointerType(dtype)); BufferVar buffer = decl_buffer({16}); body = SeqStmt({DeclBuffer(buffer, buf_var), std::move(body)}); - BufferRegion buffer_region(buffer, {Range::FromMinExtent(x + 1, 1)}); + TensorRegion buffer_region = BufferRegion(buffer, {Range::FromMinExtent(x + 1, 1)}); s_tir::MatchBufferRegion match_buffer_region(decl_buffer({1}), buffer_region); // construct block and block_realize @@ -362,7 +362,7 @@ TEST(IRF, StmtExprMutator) { Stmt alloc = fmakealloc(); // body is: DeclBuffer, AllocBuffer, Evaluate Stmt body = SeqStmt({decl, alloc, eval_body}); - BufferRegion buffer_region(buffer, {Range::FromMinExtent(x + 1, 1)}); + TensorRegion buffer_region = BufferRegion(buffer, {Range::FromMinExtent(x + 1, 1)}); s_tir::MatchBufferRegion match_buffer_region(decl_buffer({1}), buffer_region); // construct block and block_realize s_tir::SBlock block = s_tir::SBlock({}, {buffer_region}, {buffer_region}, "block", body, body, diff --git a/tests/cpp/s_tir_functor_test.cc b/tests/cpp/s_tir_functor_test.cc index 653975ada1d9..3e756b96e13b 100644 --- a/tests/cpp/s_tir_functor_test.cc +++ b/tests/cpp/s_tir_functor_test.cc @@ -34,10 +34,10 @@ using namespace tirx; TEST(STIRFunctor, LegacyInheritedDispatchAndContainsNode) { class Dispatch : public StmtFunctor { public: - using StmtFunctor::VisitStmt_; - int VisitStmt_(const SBlockNode*, int value) final { return value + 1; } - int VisitStmt_(const SBlockRealizeNode*, int value) final { return value + 2; } - int VisitStmt_(const EvaluateNode*, int value) final { return value + 3; } + using StmtFunctor::Dispatch_; + int Dispatch_(const SBlockNode*, int value) final { return value + 1; } + int Dispatch_(const SBlockRealizeNode*, int value) final { return value + 2; } + int Dispatch_(const EvaluateNode*, int value) final { return value + 3; } } dispatch; Stmt body = Evaluate(0); SBlock block({}, {}, {}, "block", body); @@ -53,9 +53,9 @@ TEST(STIRFunctor, LegacyInheritedDispatchAndContainsNode) { TEST(STIRFunctor, CoreLegacyDispatchReachesDefaultForDialectNodes) { class Dispatch : public tirx::StmtFunctor { public: - using tirx::StmtFunctor::VisitStmt_; - bool VisitStmt_(const EvaluateNode*) final { return true; } - bool VisitStmtDefault_(const ffi::Object*) final { return false; } + using tirx::StmtFunctor::Dispatch_; + bool Dispatch_(const EvaluateNode*) final { return true; } + bool DispatchDefault_(const ffi::Object*) final { return false; } } dispatch; SBlock block({}, {}, {}, "block", Evaluate(0)); EXPECT_FALSE(dispatch(block)); @@ -88,7 +88,7 @@ TEST(STIRFunctor, NativeBlockOverrideReusesInheritedCoreHooks) { TEST(STIRFunctor, NativeVisitPreservesBlockOrderAndBinders) { PrimVar index("index"), extent("extent"), annotation("annotation"); BufferVar buffer = decl_buffer({16}); - BufferRegion region(buffer, {Range::FromMinExtent(0, 16)}); + TensorRegion region = BufferRegion(buffer, {Range::FromMinExtent(0, 16)}); IterVar iter(Range::FromMinExtent(0, extent), index, IterVarType::kDataPar); SBlock block({iter}, {region}, {}, "block", Evaluate(index), std::nullopt, {buffer}, {}, {{"annotation", annotation}}); @@ -165,9 +165,9 @@ void CheckMutationRemapsBufferDefinitionsAndUses() { PrimVar extent("extent"); BufferVar allocated = decl_buffer({extent + 1}, PrimType::Int(32)); BufferVar matched = decl_buffer({extent + 1}, PrimType::Int(32)); - BufferRegion region(allocated, {Range::FromMinExtent(0, extent + 1)}); + TensorRegion region = BufferRegion(allocated, {Range::FromMinExtent(0, extent + 1)}); MatchBufferRegion match(matched, region); - BufferRegion matched_region(matched, {Range::FromMinExtent(0, extent + 1)}); + TensorRegion matched_region = BufferRegion(matched, {Range::FromMinExtent(0, extent + 1)}); Stmt body = SeqStmt({BufferStore(allocated, 0, {0}), BufferStore(matched, 0, {0})}); SBlock block({}, {region, matched_region}, {region, matched_region}, "block", body, std::nullopt, {allocated}, {match}); @@ -186,11 +186,12 @@ void CheckMutationRemapsBufferDefinitionsAndUses() { EXPECT_FALSE(new_matched.same_as(matched)); EXPECT_TRUE(new_allocated->shape[0].same_as(extent)); EXPECT_TRUE(new_matched->shape[0].same_as(extent)); - EXPECT_TRUE(changed->reads[0]->buffer.same_as(new_allocated)); - EXPECT_TRUE(changed->writes[0]->buffer.same_as(new_allocated)); - EXPECT_TRUE(changed->reads[1]->buffer.same_as(new_matched)); - EXPECT_TRUE(changed->writes[1]->buffer.same_as(new_matched)); - EXPECT_TRUE(changed->match_buffers[0]->source->buffer.same_as(new_allocated)); + EXPECT_TRUE(changed->reads[0]->source.as_or_throw().same_as(new_allocated)); + EXPECT_TRUE(changed->writes[0]->source.as_or_throw().same_as(new_allocated)); + EXPECT_TRUE(changed->reads[1]->source.as_or_throw().same_as(new_matched)); + EXPECT_TRUE(changed->writes[1]->source.as_or_throw().same_as(new_matched)); + EXPECT_TRUE( + changed->match_buffers[0]->source->source.as_or_throw().same_as(new_allocated)); const auto* statements = changed->body.as(); ASSERT_NE(statements, nullptr); EXPECT_TRUE(statements->seq[0].as()->buffer.same_as(new_allocated)); @@ -211,8 +212,8 @@ TEST(STIRFunctor, StructuralAndGenericSubstitutionPreserveDefinitionUses) { PrimVar extent("extent"), new_extent("new_extent"), index("index"), new_index("new_index"); BufferVar allocated = decl_buffer({extent}, PrimType::Int(32)); BufferVar matched = decl_buffer({extent}, PrimType::Int(32)); - BufferRegion region(allocated, {Range::FromMinExtent(0, extent)}); - BufferRegion matched_region(matched, {Range::FromMinExtent(0, extent)}); + TensorRegion region = BufferRegion(allocated, {Range::FromMinExtent(0, extent)}); + TensorRegion matched_region = BufferRegion(matched, {Range::FromMinExtent(0, extent)}); MatchBufferRegion match(matched, region); IterVar iter(Range::FromMinExtent(0, extent), index, IterVarType::kDataPar); Stmt body = @@ -231,11 +232,16 @@ TEST(STIRFunctor, StructuralAndGenericSubstitutionPreserveDefinitionUses) { EXPECT_TRUE(changed->iter_vars[0]->dom->extent.same_as(new_extent)); EXPECT_TRUE(changed->alloc_buffers[0]->shape[0].same_as(new_extent)); EXPECT_TRUE(changed->match_buffers[0]->buffer->shape[0].same_as(new_extent)); - EXPECT_TRUE(changed->match_buffers[0]->source->buffer.same_as(changed->alloc_buffers[0])); - EXPECT_TRUE(changed->reads[0]->buffer.same_as(changed->alloc_buffers[0])); - EXPECT_TRUE(changed->reads[1]->buffer.same_as(changed->match_buffers[0]->buffer)); - EXPECT_TRUE(changed->writes[0]->buffer.same_as(changed->alloc_buffers[0])); - EXPECT_TRUE(changed->writes[1]->buffer.same_as(changed->match_buffers[0]->buffer)); + EXPECT_TRUE(changed->match_buffers[0]->source->source.as_or_throw().same_as( + changed->alloc_buffers[0])); + EXPECT_TRUE( + changed->reads[0]->source.as_or_throw().same_as(changed->alloc_buffers[0])); + EXPECT_TRUE(changed->reads[1]->source.as_or_throw().same_as( + changed->match_buffers[0]->buffer)); + EXPECT_TRUE( + changed->writes[0]->source.as_or_throw().same_as(changed->alloc_buffers[0])); + EXPECT_TRUE(changed->writes[1]->source.as_or_throw().same_as( + changed->match_buffers[0]->buffer)); EXPECT_TRUE(changed->reads[0]->region[0]->extent.same_as(new_extent)); EXPECT_TRUE(changed->reads[1]->region[0]->extent.same_as(new_extent)); const auto* statements = changed->body.as(); @@ -307,7 +313,7 @@ TEST(STIRFunctor, StructuralAndGenericSubstitutionPreserveDefinitionUses) { TEST(STIRFunctor, GenericTIRXVisitorUsesFullStructuralTraversal) { PrimVar index("index"), annotation("annotation"); BufferVar buffer = decl_buffer({16}); - BufferRegion region(buffer, {Range::FromMinExtent(0, 16)}); + TensorRegion region = BufferRegion(buffer, {Range::FromMinExtent(0, 16)}); IterVar iter(Range::FromMinExtent(0, 16), index, IterVarType::kDataPar); SBlock block({iter}, {region}, {}, "block", Evaluate(index), std::nullopt, {buffer}, {}, {{"annotation", annotation}}); diff --git a/tests/python/s_tir/test_stmt.py b/tests/python/s_tir/test_stmt.py index 76b13b4501b2..11767939be90 100644 --- a/tests/python/s_tir/test_stmt.py +++ b/tests/python/s_tir/test_stmt.py @@ -26,10 +26,12 @@ @pytest.mark.parametrize("legacy", [False, True]) -def test_sblock_serialization(legacy): +@pytest.mark.parametrize("legacy_region", [False, True]) +def test_sblock_serialization(legacy, legacy_region): source = tirx.decl_buffer((4,), "float32", name="source") target = tirx.decl_buffer((4,), "float32", name="target") - region = tirx.BufferRegion(source, [tvm.ir.Range(0, 4)]) + span = tvm.ir.Span(tvm.ir.SourceName("region"), 2, 3, 4, 5) + region = tvm.ir.TensorRegion(source, [tvm.ir.Range(0, 4)], tirx.BufferRegionType(), span) match = s_tir.MatchBufferRegion(target, region) block = s_tir.SBlock([], [region], [region], "copy", tirx.Evaluate(0), match_buffers=[match]) realize = s_tir.SBlockRealize([], True, block) @@ -41,6 +43,13 @@ def test_sblock_serialization(legacy): assert getattr(s_tir, name).__module__ == "tvm.s_tir.stmt" assert not hasattr(tirx, name) assert not hasattr(tirx.stmt, name) + assert "ir.TensorRegion" in type_keys + assert "tirx.BufferRegion" not in type_keys + if legacy_region: + for node in graph["nodes"]: + if node.get("type") == "ir.TensorRegion": + node["type"] = "tirx.BufferRegion" + node["data"]["buffer"] = node["data"].pop("source") if legacy: for node in graph["nodes"]: if node.get("type") in { @@ -59,7 +68,47 @@ def test_sblock_serialization(legacy): assert restored_block.match_buffers[0].same_as(restored_match) assert restored_block.reads[0].same_as(restored_block.writes[0]) assert restored_match.source.same_as(restored_block.reads[0]) + restored_region = restored_match.source + assert isinstance(restored_region, tvm.ir.TensorRegion) + assert isinstance(restored_region.ty, tirx.BufferRegionType) + assert restored_region.span.source_name.name == "region" + assert restored_region.span.line == 2 + assert restored_region.span.end_line == 3 + assert restored_region.span.column == 4 + assert restored_region.span.end_column == 5 tvm.ir.assert_structural_equal(restored_realize, realize, map_free_vars=True) + canonical = json.loads(tvm.ir.save_json([restored_block, restored_realize, restored_match])) + assert not any( + node.get("type") + in { + "tirx.BufferRegion", + "tirx.SBlock", + "tirx.SBlockRealize", + "tirx.MatchBufferRegion", + } + for node in canonical["nodes"] + ) + + +@pytest.mark.parametrize("field", ["reads", "writes", "match_source"]) +@pytest.mark.parametrize("invalid", ["source", "rank"]) +def test_region_requires_buffer_source_and_rank(field, invalid): + source = ( + tirx.Var("source", "int32") if invalid == "source" else tirx.decl_buffer((4, 4), "float32") + ) + region = tvm.ir.TensorRegion(source, [tvm.ir.Range(0, 4)], tirx.BufferRegionType()) + message = None if invalid == "source" else "must match its buffer rank" + with pytest.raises((TypeError, tvm.error.InternalError), match=message): + if field == "match_source": + s_tir.MatchBufferRegion(tirx.decl_buffer((4,), "float32"), region) + else: + s_tir.SBlock( + [], + [region] if field == "reads" else [], + [region] if field == "writes" else [], + "invalid", + tirx.Evaluate(0), + ) if __name__ == "__main__": diff --git a/tests/python/s_tir/transform/test_s_tir_transform_convert_ssa.py b/tests/python/s_tir/transform/test_s_tir_transform_convert_ssa.py index 86aed688eb66..d881981f0918 100644 --- a/tests/python/s_tir/transform/test_s_tir_transform_convert_ssa.py +++ b/tests/python/s_tir/transform/test_s_tir_transform_convert_ssa.py @@ -54,7 +54,7 @@ def test_shared_buffer_parameter_regions_across_functions(): assert not first.params[0].same_as(second.params[0]) for updated in [first, second]: updated_block = updated.body.block - assert updated_block.reads[0].buffer.same_as(updated.params[0]) + assert updated_block.reads[0].source.same_as(updated.params[0]) assert updated_block.body.value.source.same_as(updated.params[0]) assert updated_block.reads[0].region[0].extent.same_as(updated.params[0].ty.shape[0]) diff --git a/tests/python/tirx-base/test_tir_specialize.py b/tests/python/tirx-base/test_tir_specialize.py index 8dc8832c41bf..4be5ed121dc6 100644 --- a/tests/python/tirx-base/test_tir_specialize.py +++ b/tests/python/tirx-base/test_tir_specialize.py @@ -394,8 +394,8 @@ def test_specialize_structural_buffer_definitions(): assert not after.params assert result.alloc_buffers[0].shape[0] == 8 assert result.match_buffers[0].buffer.shape[0] == 8 - assert result.match_buffers[0].source.buffer.same_as(result.alloc_buffers[0]) - assert result.reads[0].buffer.same_as(result.match_buffers[0].buffer) + assert result.match_buffers[0].source.source.same_as(result.alloc_buffers[0]) + assert result.reads[0].source.same_as(result.match_buffers[0].buffer) assert result.body.value.source.same_as(result.match_buffers[0].buffer) assert result.body.value.indices[0] == 7 assert result.annotations["extent"] == 8 From 38e8a9bb1a1360132513577705e570f91b87a1ab Mon Sep 17 00:00:00 2001 From: Tianqi Chen Date: Thu, 17 Sep 2026 11:45:26 +0000 Subject: [PATCH 08/18] Bind the call argument in the constructor roundtrip fixture --- tests/python/tirx-base/test_tir_constructor.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/python/tirx-base/test_tir_constructor.py b/tests/python/tirx-base/test_tir_constructor.py index cca9089609d9..079f2558fd65 100644 --- a/tests/python/tirx-base/test_tir_constructor.py +++ b/tests/python/tirx-base/test_tir_constructor.py @@ -140,7 +140,7 @@ def test_expr_constructor(): script = tvm.tirx.Evaluate(x_with_attrs).script() assert "attrs" in script assert "disable_tma" in script - func = tvm.tirx.PrimFunc([], tvm.tirx.Evaluate(x_with_attrs)) + func = tvm.tirx.PrimFunc([attr_arg], tvm.tirx.Evaluate(x_with_attrs)) assert tvm.script.from_source(func.script()).script() == func.script() y = tvm.tirx.Var("y", "float32") From a8985b2a1ce6f12b611848e7b7af7cf6ca900049 Mon Sep 17 00:00:00 2001 From: Tianqi Chen Date: Thu, 17 Sep 2026 11:57:02 +0000 Subject: [PATCH 09/18] Own statement simplifier implementation in TIRX Place StmtSimplifier and its configuration implementation in the tirx namespace and use that shared base from S-TIR. Keep arithmetic analyzer dependencies explicit and preserve existing runtime and pass configuration identities. --- src/s_tir/transform/stmt_simplify.cc | 13 +++++----- src/tirx/transform/stmt_simplify.cc | 36 +++++++++++++--------------- src/tirx/transform/stmt_simplify.h | 15 ++++-------- 3 files changed, 27 insertions(+), 37 deletions(-) diff --git a/src/s_tir/transform/stmt_simplify.cc b/src/s_tir/transform/stmt_simplify.cc index 458abbd634a0..8f5d9aa5cd65 100644 --- a/src/s_tir/transform/stmt_simplify.cc +++ b/src/s_tir/transform/stmt_simplify.cc @@ -30,10 +30,10 @@ namespace s_tir { using namespace tirx; // Reuse ordinary TIRX simplification, adding scoped constraints for S-TIR blocks. -class StmtSimplifier : public arith::StmtSimplifier { +class StmtSimplifier : public tirx::StmtSimplifier { public: - using Parent = arith::StmtSimplifier; - StmtSimplifier(const arith::Analyzer& analyzer, arith::StmtSimplifyConfig config) + using Parent = tirx::StmtSimplifier; + StmtSimplifier(const arith::Analyzer& analyzer, tirx::StmtSimplifyConfig config) : Parent(GlobalVTable(), analyzer, config) {} using Parent::Mutate_; using Parent::Run; @@ -70,7 +70,7 @@ class StmtSimplifier : public arith::StmtSimplifier { }; PrimFunc StmtSimplify(PrimFunc func, const arith::Analyzer& analyzer) { - auto config = tvm::transform::PassConfigWithDefaults(); + auto config = tvm::transform::PassConfigWithDefaults(); return ffi::make_object(analyzer, config)->Run(std::move(func)); } @@ -78,9 +78,8 @@ namespace transform { Pass StmtSimplify() { auto pass_func = [](PrimFunc func, IRModule, tvm::transform::PassContext ctx) { arith::Analyzer analyzer; - auto config = - ctx->GetConfig("tirx.StmtSimplify") - .value_or(tvm::transform::PassConfigWithDefaults()); + auto config = ctx->GetConfig("tirx.StmtSimplify") + .value_or(tvm::transform::PassConfigWithDefaults()); return ffi::make_object(analyzer, config)->Run(std::move(func)); }; return tirx::transform::CreatePrimFuncPass(pass_func, 0, "s_tir.StmtSimplify", {}); diff --git a/src/tirx/transform/stmt_simplify.cc b/src/tirx/transform/stmt_simplify.cc index af2c5da6dd43..6af72187811b 100644 --- a/src/tirx/transform/stmt_simplify.cc +++ b/src/tirx/transform/stmt_simplify.cc @@ -39,11 +39,9 @@ #include "../ir/ir_mutator_with_analyzer.h" namespace tvm { -namespace arith { +namespace tirx { using namespace tvm::prim; -using namespace tirx; - void StmtSimplifyConfigNode::RegisterReflection() { namespace refl = tvm::ffi::reflection; refl::ObjectDef() @@ -61,17 +59,19 @@ void StmtSimplifyConfigNode::RegisterReflection() { refl::DefaultValue(false)); } -RewriteSimplifier::Extension StmtSimplifyConfigNode::GetEnabledExtensions() const { - RewriteSimplifier::Extension flags = RewriteSimplifier::kNone; +arith::RewriteSimplifier::Extension StmtSimplifyConfigNode::GetEnabledExtensions() const { + arith::RewriteSimplifier::Extension flags = arith::RewriteSimplifier::kNone; if (transitively_prove_inequalities) { - flags = RewriteSimplifier::Extension(flags | RewriteSimplifier::kTransitivelyProveInequalities); + flags = arith::RewriteSimplifier::Extension( + flags | arith::RewriteSimplifier::kTransitivelyProveInequalities); } if (convert_boolean_to_and_of_ors) { - flags = RewriteSimplifier::Extension(flags | RewriteSimplifier::kConvertBooleanToAndOfOrs); + flags = arith::RewriteSimplifier::Extension( + flags | arith::RewriteSimplifier::kConvertBooleanToAndOfOrs); } if (apply_constraints_to_boolean_branches) { - flags = - RewriteSimplifier::Extension(flags | RewriteSimplifier::kApplyConstraintsToBooleanBranches); + flags = arith::RewriteSimplifier::Extension( + flags | arith::RewriteSimplifier::kApplyConstraintsToBooleanBranches); } return flags; } @@ -84,7 +84,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { StmtSimplifyConfigNode::RegisterReflection(); } TVM_REGISTER_PASS_CONFIG_OPTION("tirx.StmtSimplify", StmtSimplifyConfig); -PrimFunc StmtSimplifier::Apply(PrimFunc func, const Analyzer& analyzer, +PrimFunc StmtSimplifier::Apply(PrimFunc func, const arith::Analyzer& analyzer, ffi::Optional config_opt) { auto config = config_opt.value_or(MakeDefaultStmtSimplifyConfig()); @@ -114,9 +114,9 @@ UnchangedOr StmtSimplifier::Mutate(ffi::AnyView input, InplaceMode inp UnchangedOr StmtSimplifier::Mutate_(const ForNode* op, InplaceMode inplace_mode) { analyzer_->Bind(op->loop_var, Range::FromMinExtent(op->min, op->extent)); - With ctx1(analyzer_, op->loop_var >= op->min); - With ctx2(analyzer_, - static_cast(op->loop_var) < op->min + op->extent); + With ctx1(analyzer_, op->loop_var >= op->min); + With ctx2(analyzer_, + static_cast(op->loop_var) < op->min + op->extent); return Parent::Mutate_(op, inplace_mode); } @@ -215,12 +215,8 @@ ffi::Optional StmtSimplifier::ProveCondition(PrimExpr condition) const { } } -} // namespace arith - -namespace tirx { - PrimFunc StmtSimplify(PrimFunc func, const arith::Analyzer& analyzer) { - return arith::StmtSimplifier::Apply(std::move(func), analyzer); + return StmtSimplifier::Apply(std::move(func), analyzer); } namespace transform { @@ -228,9 +224,9 @@ namespace transform { Pass StmtSimplify() { auto pass_func = [](PrimFunc f, IRModule m, PassContext ctx) { arith::Analyzer analyzer; - auto cfg = ctx->GetConfig("tirx.StmtSimplify"); + auto cfg = ctx->GetConfig("tirx.StmtSimplify"); - return arith::StmtSimplifier::Apply(f, analyzer, cfg); + return StmtSimplifier::Apply(f, analyzer, cfg); }; return CreatePrimFuncPass(pass_func, 0, "tirx.StmtSimplify", {}); } diff --git a/src/tirx/transform/stmt_simplify.h b/src/tirx/transform/stmt_simplify.h index 78ead6accfb2..4b48c61f4a99 100644 --- a/src/tirx/transform/stmt_simplify.h +++ b/src/tirx/transform/stmt_simplify.h @@ -30,8 +30,7 @@ #include "../ir/ir_mutator_with_analyzer.h" namespace tvm { -namespace arith { -using namespace tirx; +namespace tirx { struct StmtSimplifyConfigNode : public ffi::Object { bool transitively_prove_inequalities; @@ -42,7 +41,7 @@ struct StmtSimplifyConfigNode : public ffi::Object { TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.transform.StmtSimplifyConfig", StmtSimplifyConfigNode, ffi::Object); - RewriteSimplifier::Extension GetEnabledExtensions() const; + arith::RewriteSimplifier::Extension GetEnabledExtensions() const; }; class StmtSimplifyConfig : public ffi::ObjectRef { @@ -55,15 +54,15 @@ class StmtSimplifier : public IRMutatorWithAnalyzer { public: using IRMutatorWithAnalyzer::Mutate; using IRMutatorWithAnalyzer::Mutate_; - static PrimFunc Apply(PrimFunc func, const Analyzer& analyzer, + static PrimFunc Apply(PrimFunc func, const arith::Analyzer& analyzer, ffi::Optional config_opt = std::nullopt); - explicit StmtSimplifier(const Analyzer& analyzer, StmtSimplifyConfig config) + explicit StmtSimplifier(const arith::Analyzer& analyzer, StmtSimplifyConfig config) : IRMutatorWithAnalyzer(analyzer), config_(config) {} protected: using Parent = IRMutatorWithAnalyzer; - StmtSimplifier(const VTable* vtable, const Analyzer& analyzer, StmtSimplifyConfig config) + StmtSimplifier(const VTable* vtable, const arith::Analyzer& analyzer, StmtSimplifyConfig config) : Parent(analyzer.get(), vtable), config_(config) {} PrimFunc Run(PrimFunc func); @@ -94,10 +93,6 @@ class StmtSimplifier : public IRMutatorWithAnalyzer { ffi::Map non_inlined_bindings_; }; -} // namespace arith - -namespace tirx { - /* \brief Simplify statements in the prim func * * Applies the same behavior as the tirx.transform.StmtSimplify pass. From cadd0353d50aa0d1e80acf56aa4eb3dab84d1367 Mon Sep 17 00:00:00 2001 From: Tianqi Chen Date: Thu, 17 Sep 2026 11:57:02 +0000 Subject: [PATCH 10/18] Document block traversal order and native-root inline eligibility Visit match-buffer definitions before read/write regions, changing the former native read/write-before-match order to agree with structural definition-before-use traversal. Generalize the private-function inline exclusion from SBlockRealize roots to every root absent from native TIRX statement dispatch. This also excludes future non-native roots whose binder and naming rules the inliner cannot preserve. --- src/s_tir/stmt_functor.cc | 3 +++ src/tirx/transform/inline_private_functions.cc | 5 +++-- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/src/s_tir/stmt_functor.cc b/src/s_tir/stmt_functor.cc index 91a20e902425..924a8efb8737 100644 --- a/src/s_tir/stmt_functor.cc +++ b/src/s_tir/stmt_functor.cc @@ -57,6 +57,9 @@ ffi::Optional StmtExprVisitor::VisitBlock(tirx::StmtExprVisitor* 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); })); diff --git a/src/tirx/transform/inline_private_functions.cc b/src/tirx/transform/inline_private_functions.cc index 6e7278ae7ad8..9644c30495d2 100644 --- a/src/tirx/transform/inline_private_functions.cc +++ b/src/tirx/transform/inline_private_functions.cc @@ -121,8 +121,9 @@ bool IsInlinablePrimFunc(const GlobalVar& gvar, const PrimFunc& prim_func, if (param->ty.as()) return false; } - // Extension statement roots may introduce binder or naming rules that this - // pass cannot preserve. Only inline roots supported by native TIRX traversal. + // Generalize the old SBlockRealize exclusion to all non-native statement roots: + // they may introduce binder or naming rules that this pass cannot preserve. + // Only inline roots supported by native TIRX traversal. struct NativeStmtTable : StmtExprVisitor { static VTable Make() { VTable table; From 7c1e6f1d25e3e7c3ddf036bc8b5f3f52dc07d2de Mon Sep 17 00:00:00 2001 From: Tianqi Chen Date: Thu, 17 Sep 2026 13:04:27 +0000 Subject: [PATCH 11/18] Share lowered execution attribute keys with TIRX Centralize virtual-thread and asynchronous-wait keys in TIRX and retain S-TIR aliases so both schedulable and lowered consumers use the same protocol names. --- include/tvm/s_tir/stmt.h | 6 +++--- include/tvm/tirx/stmt.h | 5 +++++ src/tirx/ir/data_type_rewriter.cc | 4 ++-- src/tirx/ir/ir_mutator_with_analyzer.cc | 3 ++- src/tirx/ir/ir_visitor_with_analyzer.cc | 3 ++- src/tirx/ir/tir_visitor_with_path.cc | 3 ++- src/tirx/script/builder/ir.cc | 3 ++- src/tirx/script/printer/stmt.cc | 3 ++- src/tirx/transform/bind_target.cc | 4 ++-- src/tirx/transform/ir_utils.cc | 4 ++-- src/tirx/transform/ir_utils.h | 2 +- src/tirx/transform/lower_tirx_opaque.cc | 2 +- src/tirx/transform/lower_warp_memory.cc | 2 +- src/tirx/transform/narrow_datatype.cc | 2 +- src/tirx/transform/remove_no_op.cc | 2 +- src/tirx/transform/storage_rewrite.cc | 8 ++++---- 16 files changed, 33 insertions(+), 23 deletions(-) diff --git a/include/tvm/s_tir/stmt.h b/include/tvm/s_tir/stmt.h index 9feff932144f..eefe47d8602f 100644 --- a/include/tvm/s_tir/stmt.h +++ b/include/tvm/s_tir/stmt.h @@ -204,8 +204,8 @@ 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. @@ -244,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/tirx/stmt.h b/include/tvm/tirx/stmt.h index 912088a5d022..5c3f4f2cd99d 100644 --- a/include/tvm/tirx/stmt.h +++ b/include/tvm/tirx/stmt.h @@ -829,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/src/tirx/ir/data_type_rewriter.cc b/src/tirx/ir/data_type_rewriter.cc index eaa54a388bb5..8f79792d1826 100644 --- a/src/tirx/ir/data_type_rewriter.cc +++ b/src/tirx/ir/data_type_rewriter.cc @@ -72,7 +72,7 @@ UnchangedOr DataTypeLegalizer::Mutate_(const ForNode* op, InplaceMode inpl } UnchangedOr DataTypeLegalizer::Mutate_(const AttrStmtNode* op, InplaceMode inplace_mode) { - if (op->attr_key == attr::thread_extent || op->attr_key == "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(op)); op = s.as(); TVM_FFI_ICHECK(op != nullptr) << "Expected type to be AttrStmtNode" @@ -328,7 +328,7 @@ UnchangedOr DataTypeLegalizer::Mutate_(const CallNode* op, InplaceMode inp } UnchangedOr IndexDataTypeRewriter::Mutate_(const AttrStmtNode* op, InplaceMode inplace_mode) { - if (op->attr_key == attr::thread_extent || op->attr_key == "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); diff --git a/src/tirx/ir/ir_mutator_with_analyzer.cc b/src/tirx/ir/ir_mutator_with_analyzer.cc index ac5d6f6381d8..1c65f94874a3 100644 --- a/src/tirx/ir/ir_mutator_with_analyzer.cc +++ b/src/tirx/ir/ir_mutator_with_analyzer.cc @@ -202,7 +202,8 @@ UnchangedOr IRMutatorWithAnalyzer::Mutate_(const IfThenElseNode* op, UnchangedOr IRMutatorWithAnalyzer::Mutate_(const AttrStmtNode* op, InplaceMode inplace_mode) { return constraint_scope_.WithNewScope([&]() -> UnchangedOr { - if (op->attr_key == tirx::attr::thread_extent || op->attr_key == "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(); 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_visitor_with_analyzer.cc b/src/tirx/ir/ir_visitor_with_analyzer.cc index 3324e6d2194f..ae823b3ddedc 100644 --- a/src/tirx/ir/ir_visitor_with_analyzer.cc +++ b/src/tirx/ir/ir_visitor_with_analyzer.cc @@ -80,7 +80,8 @@ ffi::Optional IRVisitorWithAnalyzer::Visit_(const IfThenElseNode ffi::Optional IRVisitorWithAnalyzer::Visit_(const AttrStmtNode* op) { return constraint_scope_.WithNewScope([&]() -> ffi::Optional { - if (op->attr_key == tirx::attr::thread_extent || op->attr_key == "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(); TVM_FFI_ICHECK_NE(iv->thread_tag.length(), 0U); analyzer_->Bind(iv->var, Range::FromMinExtent(IntImm(op->value.ty(), 0), op->value)); diff --git a/src/tirx/ir/tir_visitor_with_path.cc b/src/tirx/ir/tir_visitor_with_path.cc index c37c443ed374..37eaa20d396a 100644 --- a/src/tirx/ir/tir_visitor_with_path.cc +++ b/src/tirx/ir/tir_visitor_with_path.cc @@ -181,7 +181,8 @@ void TIRVisitorWithPath::Dispatch_(const AttrStmtNode* op, AccessPath path) { std::vector, DefContext, DefContext>> context; if (auto iter_var = op->node.as(); - iter_var && (op->attr_key == attr::thread_extent || op->attr_key == "virtual_thread")) { + iter_var && + (op->attr_key == attr::thread_extent || op->attr_key == tvm::tirx::attr::virtual_thread)) { // Some attributes serve as a source of definition for the // tirx::Var they annotate. context.push_back(WithDef(iter_var.value(), path->Attr("node"))); diff --git a/src/tirx/script/builder/ir.cc b/src/tirx/script/builder/ir.cc index 97623a658596..4fa7629cd02d 100644 --- a/src/tirx/script/builder/ir.cc +++ b/src/tirx/script/builder/ir.cc @@ -686,7 +686,8 @@ LaunchThreadFrame LaunchThread(Var var, PrimExpr extent) { } n->iter_var = iter_var; n->extent = extent; - n->attr_key = iter_var->thread_tag == "vthread" ? "virtual_thread" : "thread_extent"; + n->attr_key = + iter_var->thread_tag == "vthread" ? tvm::tirx::attr::virtual_thread : "thread_extent"; return LaunchThreadFrame(n); } diff --git a/src/tirx/script/printer/stmt.cc b/src/tirx/script/printer/stmt.cc index f13136512b8b..2bc677cf1e80 100644 --- a/src/tirx/script/printer/stmt.cc +++ b/src/tirx/script/printer/stmt.cc @@ -809,7 +809,8 @@ TVM_FFI_STATIC_INIT_BLOCK() { ffi::Optional define_var = std::nullopt; tirx::Stmt body = stmt->body; AccessPath body_p = stmt_p->Attr("body"); - if (stmt->attr_key == "thread_extent" || stmt->attr_key == "virtual_thread") { + if (stmt->attr_key == "thread_extent" || + stmt->attr_key == tvm::tirx::attr::virtual_thread) { if (stmt->node.as()) { rhs = DocsifyLaunchThread(stmt, stmt_p, &define_var, d); } diff --git a/src/tirx/transform/bind_target.cc b/src/tirx/transform/bind_target.cc index 91ae81f98daf..75fd157dacb5 100644 --- a/src/tirx/transform/bind_target.cc +++ b/src/tirx/transform/bind_target.cc @@ -115,7 +115,7 @@ class FunctionClassifierVisitor : public StmtExprVisitor { } ffi::Optional Visit_(const AttrStmtNode* op) final { - 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 || op->attr_key == attr::kDeviceEntry) { // Enter GPU scope for thread extent and virtual thread attributes bool last_is_under_gpu_scope = is_under_gpu_scope_; @@ -204,7 +204,7 @@ class CallSubstitutor : public StmtExprMutator { } UnchangedOr Mutate_(const AttrStmtNode* op, InplaceMode inplace_mode) final { - 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 || op->attr_key == attr::kDeviceEntry) { // Enter GPU scope for thread extent and virtual thread attributes bool last_is_under_gpu_scope = is_under_gpu_scope_; diff --git a/src/tirx/transform/ir_utils.cc b/src/tirx/transform/ir_utils.cc index 1d4978167ffe..be6d6cba09db 100644 --- a/src/tirx/transform/ir_utils.cc +++ b/src/tirx/transform/ir_utils.cc @@ -664,9 +664,9 @@ ffi::Array GetBufferAllocationShape(const BufferVar& buffer) { // Attribute strings are the metadata protocol shared by lowered and schedulable statements. std::pair GetAsyncWaitAttributes(const AttrStmtNode* op) { - TVM_FFI_ICHECK(op && op->attr_key == "async_wait_queue_scope"); + TVM_FFI_ICHECK(op && op->attr_key == tvm::tirx::attr::async_wait_queue_scope); auto inner = op->body.as(); - TVM_FFI_ICHECK(inner && inner->attr_key == "async_wait_inflight_count"); + TVM_FFI_ICHECK(inner && inner->attr_key == tvm::tirx::attr::async_wait_inflight_count); return std::make_pair(op->value, inner->value); } diff --git a/src/tirx/transform/ir_utils.h b/src/tirx/transform/ir_utils.h index 0ee2c13028b0..004bfa6f9d2b 100644 --- a/src/tirx/transform/ir_utils.h +++ b/src/tirx/transform/ir_utils.h @@ -389,7 +389,7 @@ struct FragmentInfo { std::unordered_map GetTensorCoreFragmentInfo(const Stmt& stmt); // Return the queue id and the in-flight count associated with the given -// s_tir::attr::async_wait_queue_scope annotation. +// tvm::tirx::attr::async_wait_queue_scope annotation. std::pair GetAsyncWaitAttributes(const AttrStmtNode* op); /*! diff --git a/src/tirx/transform/lower_tirx_opaque.cc b/src/tirx/transform/lower_tirx_opaque.cc index ec245fa25ccf..7c2c51f9541a 100644 --- a/src/tirx/transform/lower_tirx_opaque.cc +++ b/src/tirx/transform/lower_tirx_opaque.cc @@ -100,7 +100,7 @@ class TIRxOpaqueLower : public StmtExprMutator { /*thread_tag=*/thread_tag); ffi::String attr_key = (thread_tag == "vthread" || thread_tag == "vthread.x" || thread_tag == "vthread.y" || thread_tag == "vthread.z") - ? s_tir::attr::virtual_thread + ? tvm::tirx::attr::virtual_thread : tirx::attr::thread_extent; return AttrStmt(/*node=*/std::move(iter_var), /*attr_key=*/std::move(attr_key), diff --git a/src/tirx/transform/lower_warp_memory.cc b/src/tirx/transform/lower_warp_memory.cc index 871b9a35594d..9aab893b2d32 100644 --- a/src/tirx/transform/lower_warp_memory.cc +++ b/src/tirx/transform/lower_warp_memory.cc @@ -499,7 +499,7 @@ class BindVarBoundInfo : public StmtExprVisitor { } ffi::Optional Visit_(const AttrStmtNode* op) { - 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) { IterVar iv = op->node.as_or_throw(); TVM_FFI_ICHECK_NE(iv->thread_tag.length(), 0U); if (!var_dom_.count(iv->var.get())) { diff --git a/src/tirx/transform/narrow_datatype.cc b/src/tirx/transform/narrow_datatype.cc index a57919daa508..9fa93f78d881 100644 --- a/src/tirx/transform/narrow_datatype.cc +++ b/src/tirx/transform/narrow_datatype.cc @@ -124,7 +124,7 @@ class DataTypeVisitor final : public StmtExprVisitor { } ffi::Optional Visit_(const AttrStmtNode* op) { - if (op->attr_key == attr::thread_extent || op->attr_key == "virtual_thread") { + if (op->attr_key == attr::thread_extent || op->attr_key == tvm::tirx::attr::virtual_thread) { IterVar iv = op->node.as_or_throw(); TVM_FFI_ICHECK_NE(iv->thread_tag.length(), 0U); analyzer_->Bind(iv->var, Range::FromMinExtent(0, op->value)); diff --git a/src/tirx/transform/remove_no_op.cc b/src/tirx/transform/remove_no_op.cc index a688acd074b1..5c6d90d0397a 100644 --- a/src/tirx/transform/remove_no_op.cc +++ b/src/tirx/transform/remove_no_op.cc @@ -94,7 +94,7 @@ class NoOpRemover : public IRMutatorWithAnalyzer { UnchangedOr Mutate_(const AttrStmtNode* op, InplaceMode inplace_mode) final { if (op->attr_key == "pragma_debug_skip_region") { return MakeEvaluate(0); - } else if (op->attr_key == s_tir::attr::async_wait_queue_scope) { + } else if (op->attr_key == tvm::tirx::attr::async_wait_queue_scope) { auto wait_attrs = GetAsyncWaitAttributes(op); auto wait_cnt = wait_attrs.second; arith::Analyzer ana; diff --git a/src/tirx/transform/storage_rewrite.cc b/src/tirx/transform/storage_rewrite.cc index df1911179f59..44cdd6491ac5 100644 --- a/src/tirx/transform/storage_rewrite.cc +++ b/src/tirx/transform/storage_rewrite.cc @@ -242,7 +242,7 @@ class LinearAccessPatternFinder final : public StmtExprVisitor { in_thread_env_ = false; } else if (op->attr_key == attr::extern_scope) { TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(VisitNewScope(op)); - } else if (op->attr_key == s_tir::attr::virtual_thread) { + } else if (op->attr_key == tvm::tirx::attr::virtual_thread) { TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(VisitNewScope(op)); } else { TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(op)); @@ -668,7 +668,7 @@ class StoragePlanRewriter : public StmtExprMutator { } UnchangedOr Mutate_(const AttrStmtNode* op, InplaceMode inplace_mode) final { - 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 || attr::IsPragmaKey(op->attr_key)) { // remake all the allocation at the attach scope. if (attach_map_.count(op)) { @@ -1088,8 +1088,8 @@ class StoragePlanRewriter : public StmtExprMutator { // enter/exit new scope if (s.stmt->IsInstance()) { const auto* op = static_cast(s.stmt); - if (op->attr_key == attr::thread_extent || op->attr_key == s_tir::attr::virtual_thread || - attr::IsPragmaKey(op->attr_key)) { + if (op->attr_key == attr::thread_extent || + op->attr_key == tvm::tirx::attr::virtual_thread || attr::IsPragmaKey(op->attr_key)) { PlanNewScope(op); } else { TVM_FFI_ICHECK(op->attr_key == attr::extern_scope); From 057af8a05044ee0a469b45c6e440698885abe132 Mon Sep 17 00:00:00 2001 From: Tianqi Chen Date: Thu, 17 Sep 2026 13:04:27 +0000 Subject: [PATCH 12/18] Preserve object-form buffer region JSON compatibility Supply default region type and span metadata for the older two-field BufferRegion schema without shifting graph indices or replacing existing typed metadata. --- python/tvm/ir/json_compact.py | 19 ++++++++--- tests/python/s_tir/test_stmt.py | 58 +++++++++++++++++++++++++++++++++ 2 files changed, 73 insertions(+), 4 deletions(-) diff --git a/python/tvm/ir/json_compact.py b/python/tvm/ir/json_compact.py index b5de5b52f64b..e26aab947a12 100644 --- a/python/tvm/ir/json_compact.py +++ b/python/tvm/ir/json_compact.py @@ -124,12 +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": - # TensorRegion keeps the inherited type/span and range references; - # only the buffer field became the shared expression source. - fields = node.get("data", {}) + 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/tests/python/s_tir/test_stmt.py b/tests/python/s_tir/test_stmt.py index 11767939be90..66f887ff9ba5 100644 --- a/tests/python/s_tir/test_stmt.py +++ b/tests/python/s_tir/test_stmt.py @@ -23,6 +23,7 @@ import tvm import tvm.testing from tvm import s_tir, tirx +from tvm.ir.json_compact import upgrade_json @pytest.mark.parametrize("legacy", [False, True]) @@ -90,6 +91,63 @@ def test_sblock_serialization(legacy, legacy_region): ) +def test_untyped_buffer_region_serialization(): + source = tirx.decl_buffer((4,), "float32", name="source") + region = [tvm.ir.Range(0, 4)] + graph = json.loads(tvm.ir.save_json([source, region])) + nodes = graph["nodes"] + source_index, region_index = nodes[graph["root_index"]]["data"] + # Before c836e8c942, BufferRegion inherited PrimExprConvertible (an Object), + # and reflection registered exactly buffer/region, with no type or span. + # Construct that historical schema directly, independently of TensorRegion + # serialization, while using current buffer/range schemas for dependencies. + legacy_index = len(nodes) + for _ in range(2): + nodes.append( + { + "type": "tirx.BufferRegion", + "data": {"buffer": source_index, "region": region_index}, + } + ) + graph["root_index"] = len(nodes) + nodes.append({"type": "ffi.Array", "data": [legacy_index, legacy_index, legacy_index + 1]}) + legacy_json = json.dumps(graph) + upgraded = json.loads(upgrade_json(legacy_json)) + assert upgraded["root_index"] == graph["root_index"] + assert len(upgraded["nodes"]) == len(nodes) + 1 + assert upgraded["nodes"][:legacy_index] == nodes[:legacy_index] + assert upgraded["nodes"][graph["root_index"]] == nodes[graph["root_index"]] + for index in (legacy_index, legacy_index + 1): + assert upgraded["nodes"][index] == { + "type": "ir.TensorRegion", + "data": { + "source": source_index, + "region": region_index, + "ty": len(nodes), + "span": 0, + }, + } + first, repeated, second = tvm.ir.load_json(legacy_json) + assert first.same_as(repeated) + assert not first.same_as(second) + assert first.source.same_as(second.source) + assert first.region.same_as(second.region) + assert first.ty.same_as(second.ty) + assert isinstance(first.ty, tirx.BufferRegionType) + assert first.span is None + expected = tvm.ir.TensorRegion(source, region, tirx.BufferRegionType()) + tvm.ir.assert_structural_equal(first, expected, map_free_vars=True) + + +@pytest.mark.parametrize("data", [None, {"region": 0}]) +def test_malformed_legacy_buffer_region(data): + node = {"type": "tirx.BufferRegion"} + if data is not None: + node["data"] = data + with pytest.raises(ValueError, match="requires a buffer field"): + upgrade_json(json.dumps({"nodes": [{"type": "None"}, node], "root_index": 1})) + + @pytest.mark.parametrize("field", ["reads", "writes", "match_source"]) @pytest.mark.parametrize("invalid", ["source", "rank"]) def test_region_requires_buffer_source_and_rank(field, invalid): From ba6409286f233aa06e92b6ece231a5fcc817048e Mon Sep 17 00:00:00 2001 From: Tianqi Chen Date: Thu, 17 Sep 2026 13:04:27 +0000 Subject: [PATCH 13/18] Clarify legalization boundaries and simplify dialect helpers Document dtype-changing statement substitution as a post-lowering TIRX operation while preserving structural type-preserving mappings. Share verifier fallback diagnostics, retain protected path-dispatch thunks, and make the normalizer hooks overridable without adding dialect collector or binding callbacks. --- include/tvm/s_tir/transform.h | 2 +- include/tvm/tirx/stmt_functor.h | 6 ++++- src/s_tir/analysis/verify_well_formed.cc | 1 + src/s_tir/ir/data_type_rewriter.cc | 2 ++ src/s_tir/ir/data_type_rewriter.h | 4 ++-- src/s_tir/ir/ir_mutator_with_analyzer.h | 2 +- src/s_tir/ir/ir_visitor_with_analyzer.h | 2 +- src/s_tir/ir/tir_visitor_with_path.h | 2 ++ src/s_tir/transform/stmt_simplify.cc | 4 +++- src/tirx/analysis/verify_tirx_well_formed.cc | 20 ---------------- tests/cpp/ir_functor_test.cc | 24 ++++++++++++++++++++ 11 files changed, 42 insertions(+), 27 deletions(-) diff --git a/include/tvm/s_tir/transform.h b/include/tvm/s_tir/transform.h index 8de8210ff73f..0a92264c4c6f 100644 --- a/include/tvm/s_tir/transform.h +++ b/include/tvm/s_tir/transform.h @@ -48,6 +48,7 @@ namespace transform { 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(); @@ -55,7 +56,6 @@ TVM_DLL Pass ConvertSSA(); /*! \brief Simplify schedulable TIR using block iteration constraints and shared simplifier options. */ TVM_DLL Pass StmtSimplify(); -using tvm::transform::PassContext; /*! * \brief Canonicalize loop to start from zero . diff --git a/include/tvm/tirx/stmt_functor.h b/include/tvm/tirx/stmt_functor.h index 66d2b2d68543..17c93ece7263 100644 --- a/include/tvm/tirx/stmt_functor.h +++ b/include/tvm/tirx/stmt_functor.h @@ -289,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. */ diff --git a/src/s_tir/analysis/verify_well_formed.cc b/src/s_tir/analysis/verify_well_formed.cc index 3e00f8e9e1ba..b1a6439fae44 100644 --- a/src/s_tir/analysis/verify_well_formed.cc +++ b/src/s_tir/analysis/verify_well_formed.cc @@ -161,6 +161,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { } TVM_FFI_THROW(TypeError) << "Expected a PrimFunc or IRModule, but received " << obj->GetTypeKey(); + TVM_FFI_UNREACHABLE(); }); } } // namespace s_tir diff --git a/src/s_tir/ir/data_type_rewriter.cc b/src/s_tir/ir/data_type_rewriter.cc index b92fe1b65ab9..edf356e1e5cb 100644 --- a/src/s_tir/ir/data_type_rewriter.cc +++ b/src/s_tir/ir/data_type_rewriter.cc @@ -29,6 +29,8 @@ 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 { diff --git a/src/s_tir/ir/data_type_rewriter.h b/src/s_tir/ir/data_type_rewriter.h index 5ccfb2cb7c67..cf0c21e79657 100644 --- a/src/s_tir/ir/data_type_rewriter.h +++ b/src/s_tir/ir/data_type_rewriter.h @@ -38,8 +38,8 @@ class IndexDataTypeNormalizer : public tirx::IndexDataTypeNormalizer { : Parent(std::move(target_data_type), GlobalVTable()) {} tirx::PrimFunc Rewrite(tirx::PrimFunc func); - UnchangedOr Mutate_(const SBlockNode* op, InplaceMode inplace_mode); - UnchangedOr Mutate_(const SBlockRealizeNode* op, InplaceMode inplace_mode); + virtual UnchangedOr Mutate_(const SBlockNode* op, InplaceMode inplace_mode); + virtual UnchangedOr Mutate_(const SBlockRealizeNode* op, InplaceMode inplace_mode); protected: static void InitVTable(VTable* table) { diff --git a/src/s_tir/ir/ir_mutator_with_analyzer.h b/src/s_tir/ir/ir_mutator_with_analyzer.h index 46a7648258a5..49179ab43cb9 100644 --- a/src/s_tir/ir/ir_mutator_with_analyzer.h +++ b/src/s_tir/ir/ir_mutator_with_analyzer.h @@ -57,4 +57,4 @@ class IRMutatorWithAnalyzer : public tirx::IRMutatorWithAnalyzer { }; } // namespace s_tir } // namespace tvm -#endif +#endif // TVM_S_TIR_IR_MUTATOR_WITH_ANALYZER_H_ diff --git a/src/s_tir/ir/ir_visitor_with_analyzer.h b/src/s_tir/ir/ir_visitor_with_analyzer.h index cf0d11813a22..ea52dba33fb9 100644 --- a/src/s_tir/ir/ir_visitor_with_analyzer.h +++ b/src/s_tir/ir/ir_visitor_with_analyzer.h @@ -46,4 +46,4 @@ class IRVisitorWithAnalyzer : public tirx::IRVisitorWithAnalyzer { }; } // namespace s_tir } // namespace tvm -#endif +#endif // TVM_S_TIR_IR_VISITOR_WITH_ANALYZER_H_ diff --git a/src/s_tir/ir/tir_visitor_with_path.h b/src/s_tir/ir/tir_visitor_with_path.h index 3da346860772..6556d2eec7f8 100644 --- a/src/s_tir/ir/tir_visitor_with_path.h +++ b/src/s_tir/ir/tir_visitor_with_path.h @@ -35,6 +35,8 @@ class TIRVisitorWithPath : public tirx::TIRVisitorWithPath { 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( [](const ffi::ObjectRef& node, StmtVisitor* self, AccessPath path) { static_cast(self)->Dispatch_( diff --git a/src/s_tir/transform/stmt_simplify.cc b/src/s_tir/transform/stmt_simplify.cc index 8f5d9aa5cd65..36d8cf432806 100644 --- a/src/s_tir/transform/stmt_simplify.cc +++ b/src/s_tir/transform/stmt_simplify.cc @@ -30,7 +30,7 @@ namespace s_tir { using namespace tirx; // Reuse ordinary TIRX simplification, adding scoped constraints for S-TIR blocks. -class StmtSimplifier : public tirx::StmtSimplifier { +class StmtSimplifier final : public tirx::StmtSimplifier { public: using Parent = tirx::StmtSimplifier; StmtSimplifier(const arith::Analyzer& analyzer, tirx::StmtSimplifyConfig config) @@ -40,6 +40,8 @@ class StmtSimplifier : public tirx::StmtSimplifier { public: UnchangedOr 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 { for (const auto& iter_var : op->iter_vars) { analyzer_->Bind(iter_var->var, iter_var->dom); diff --git a/src/tirx/analysis/verify_tirx_well_formed.cc b/src/tirx/analysis/verify_tirx_well_formed.cc index 6a95a2976a57..8407bdb67122 100644 --- a/src/tirx/analysis/verify_tirx_well_formed.cc +++ b/src/tirx/analysis/verify_tirx_well_formed.cc @@ -50,11 +50,6 @@ class ExecScopeVerifier : public Verifier { private: using Verifier::Visit; - void DispatchDefault_(const ffi::Object* op, ffi::reflection::AccessPath path) override { - Verify(false) << "TIRxError: " << op->GetTypeKey() << " 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("TIRxOpCategory"); Verify(category_map.get(op->op, ffi::String("")) == "tile_primitive") @@ -128,11 +123,6 @@ class LayoutVerifier : public Verifier { private: using Verifier::Visit; - - void DispatchDefault_(const ffi::Object* op, ffi::reflection::AccessPath path) override { - Verify(false) << "TIRxError: " << op->GetTypeKey() << " is not allowed in tirx=True mode at " - << path; - } }; class AsyncStructsVerifier : public Verifier { @@ -141,11 +131,6 @@ class AsyncStructsVerifier : public Verifier { private: using Verifier::Visit; - - void DispatchDefault_(const ffi::Object* op, ffi::reflection::AccessPath path) override { - Verify(false) << "TIRxError: " << op->GetTypeKey() << " is not allowed in tirx=True mode at " - << path; - } }; class DeviceFuncVerifier : public Verifier { @@ -154,11 +139,6 @@ class DeviceFuncVerifier : public Verifier { private: using Verifier::Visit; - - void DispatchDefault_(const ffi::Object* op, ffi::reflection::AccessPath path) override { - Verify(false) << "TIRxError: " << op->GetTypeKey() << " is not allowed in tirx=True mode at " - << path; - } }; bool VerifyTIRxWellFormed(const PrimFunc& func, bool assert_mode, bool device_func) { diff --git a/tests/cpp/ir_functor_test.cc b/tests/cpp/ir_functor_test.cc index 3ba10ebe20b3..9050b1c3c882 100644 --- a/tests/cpp/ir_functor_test.cc +++ b/tests/cpp/ir_functor_test.cc @@ -742,3 +742,27 @@ TEST(IRF, SubstituteWithDataTypeLegalizationPreservesShiftAmounts) { EXPECT_TRUE(structural_equal(actual_left, widened_y << shift_amount)); EXPECT_TRUE(structural_equal(actual_right, widened_y >> shift_amount)); } + +TEST(IRF, SubstituteWithDataTypeLegalizationCastsCoreLoopBounds) { + using namespace tvm::prim; + using namespace tvm; + using namespace tvm::tirx; + + PrimVar index("i", PrimType::Int(32)); + PrimVar extent("n", PrimType::Int(32)); + PrimVar wide_extent("n64", PrimType::Int(64)); + Stmt original = For(index, 0, extent, ForKind::kSerial, Evaluate(index)); + Stmt actual = SubstituteWithDataTypeLegalization( + original, [&](const tirx::Var& var) -> ffi::Optional { + if (var.same_as(extent)) return PrimExpr(wide_extent); + return std::nullopt; + }); + + auto* loop = actual.as(); + ASSERT_NE(loop, nullptr); + EXPECT_TRUE(loop->loop_var.same_as(index)); + EXPECT_EQ(loop->min.ty(), index.ty()); + EXPECT_EQ(loop->extent.ty(), index.ty()); + EXPECT_TRUE(ffi::StructuralEqual()(loop->extent, cast(index.ty(), wide_extent))); + EXPECT_TRUE(original.as()->extent.same_as(extent)); +} From d36d691a2dc222299190130f1578d581e3076c2f Mon Sep 17 00:00:00 2001 From: Tianqi Chen Date: Thu, 17 Sep 2026 13:04:27 +0000 Subject: [PATCH 14/18] Clarify parser verification selection and S-TIR test ownership Keep the legacy parse dialect argument as documented compatibility syntax while internal callers rely on per-function attributes. Organize block-aware verification and simplification coverage under S-TIR while retaining core and mixed-dialect checks. --- python/tvm/script/parser/core/entry.py | 5 + python/tvm/script/parser/ir/entry.py | 2 +- python/tvm/tirx/script/parser/entry.py | 3 +- .../test_s_tir_analysis_verify_well_formed.py | 225 +++++++++++++++++ .../test_s_tir_transform_simplify.py | 45 ++++ .../test_tir_analysis_verify_well_formed.py | 231 ++---------------- .../test_tir_transform_simplify.py | 21 -- 7 files changed, 295 insertions(+), 237 deletions(-) create mode 100644 tests/python/s_tir/analysis/test_s_tir_analysis_verify_well_formed.py create mode 100644 tests/python/s_tir/transform/test_s_tir_transform_simplify.py diff --git a/python/tvm/script/parser/core/entry.py b/python/tvm/script/parser/core/entry.py index cf3641d3b6fd..afd155f3be1a 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. diff --git a/python/tvm/script/parser/ir/entry.py b/python/tvm/script/parser/ir/entry.py index 4cfe60b77cac..5191ce0f87ed 100644 --- a/python/tvm/script/parser/ir/entry.py +++ b/python/tvm/script/parser/ir/entry.py @@ -61,7 +61,7 @@ def decorator_wrapper(mod): 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/script/parser/entry.py b/python/tvm/tirx/script/parser/entry.py index d708aa621a5a..03bcc9d6431c 100644 --- a/python/tvm/tirx/script/parser/entry.py +++ b/python/tvm/tirx/script/parser/entry.py @@ -70,7 +70,7 @@ def decorator_wrapper(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 @@ def specialize(self, **specialization_kwargs) -> PrimFunc: 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/tests/python/s_tir/analysis/test_s_tir_analysis_verify_well_formed.py b/tests/python/s_tir/analysis/test_s_tir_analysis_verify_well_formed.py new file mode 100644 index 000000000000..357e2d288e40 --- /dev/null +++ b/tests/python/s_tir/analysis/test_s_tir_analysis_verify_well_formed.py @@ -0,0 +1,225 @@ +# 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. + +import pytest + +import tvm +import tvm.testing +from tvm.script import ir as I +from tvm.script import tirx as T + + +def test_pass_simple(): + @T.prim_func(s_tir=True) + def element_wise( + A: T.Buffer((128, 128), "float32"), + C: T.Buffer((128, 128), "float32"), + ): + B = T.sblock_alloc_buffer((128, 128), "float32") + for i, j in T.grid(128, 128): + with T.sblock("B"): + vi, vj = T.axis.remap("SS", [i, j]) + B[vi, vj] = A[vi, vj] * 2.0 + for i, j in T.grid(128, 128): + with T.sblock("C"): + # It's a opaque block , so it can use outside variables + C[i, j] = B[i, j] * 2.0 + + assert tvm.s_tir.analysis.verify_well_formed(element_wise) + assert tvm.s_tir.analysis.verify_well_formed(tvm.IRModule.from_expr(element_wise)) + + +def test_buffer_region_bounds_are_visited(): + data = tvm.tirx.Var( + "data", tvm.ir.PointerType(tvm.ir.PrimType("int32"), storage_scope="global") + ) + buffer = tvm.tirx.decl_buffer([4], "int32", data=data) + undefined = tvm.tirx.Var("undefined", "int32") + region = tvm.tirx.BufferRegion(buffer, [tvm.ir.Range.from_min_extent(undefined, 4)]) + block = tvm.s_tir.SBlock([], [region], [], "region", tvm.tirx.Evaluate(0)) + func = tvm.tirx.PrimFunc([buffer], block) + assert not tvm.s_tir.analysis.verify_well_formed(func, assert_mode=False) + + +def test_fail_use_out_loop_var(): + @T.prim_func(check_well_formed=False, s_tir=True) + def element_wise( + A: T.Buffer((128, 128), "float32"), + B: T.Buffer((128, 128), "float32"), + ): + for i, j in T.grid(128, 128): + with T.sblock("B"): + vi, vj = T.axis.remap("SS", [i, j]) + # we cannot use `i` since it's defined outside the block + B[vi, vj] = A[i, vj] * 2.0 + + assert not tvm.s_tir.analysis.verify_well_formed(element_wise, assert_mode=False) + + +def test_block_match_buffer_defines_buffer_obj(): + """In a block, T.match_buffer defines a buffer view""" + + @I.ir_module(s_tir=True) + class mod: + @T.prim_func(s_tir=True) + def func(A: T.Buffer([256, 256], "float32")): + for iters in T.grid(16, 16, 16, 16): + with T.sblock("compute"): + tile_i, tile_j, i, j = T.axis.remap("SSSS", iters) + B = T.match_buffer( + A[tile_i * 16 : (tile_i + 1) * 16, tile_j * 16 : (tile_j + 1) * 16], + dtype="float32", + ) + B[i, j] = 0.0 + + tvm.s_tir.analysis.verify_well_formed(mod) + + +def test_block_match_buffer_defines_symbolic_variables(): + """In a block, T.match_buffer may define symbolic variables""" + + @I.ir_module(s_tir=True) + class mod: + @T.prim_func(s_tir=True) + def func(A: T.Buffer([256, 256], "int32")): + for iters in T.grid(16, 16, 16, 16): + with T.sblock("compute"): + tile_i, tile_j, i, j = T.axis.remap("SSSS", iters) + + elem_offset = T.int32() + B = T.match_buffer( + A[tile_i * 16 : (tile_i + 1) * 16, tile_j * 16 : (tile_j + 1) * 16], + dtype="float32", + elem_offset=elem_offset, + ) + + B[i, j] = elem_offset + + tvm.s_tir.analysis.verify_well_formed(mod) + + +def test_alloc_buffer_in_block_is_well_formed(): + """SBlock::alloc_buffers introduces a buffer into scope for the block body.""" + + @I.ir_module(s_tir=True) + class mod: + @T.prim_func(s_tir=True) + def func(A: T.Buffer((128,), "float32")): + with T.sblock("root"): + B = T.sblock_alloc_buffer([128], "float32") + for i in T.grid(128): + with T.sblock("write_B"): + vi = T.axis.remap("S", [i]) + B[vi] = A[vi] * 2.0 + + tvm.s_tir.analysis.verify_well_formed(mod) + + +def test_match_buffer_in_block_is_well_formed(): + """SBlock::match_buffers introduces a buffer into scope for the block body.""" + + @I.ir_module(s_tir=True) + class mod: + @T.prim_func(s_tir=True) + def func(A: T.Buffer((128, 128), "float32")): + for iters in T.grid(8, 8, 16, 16): + with T.sblock("compute"): + ti, tj, i, j = T.axis.remap("SSSS", iters) + A_tile = T.match_buffer( + A[ti * 16 : (ti + 1) * 16, tj * 16 : (tj + 1) * 16], + dtype="float32", + ) + A_tile[i, j] = A_tile[i, j] * 2.0 + + tvm.s_tir.analysis.verify_well_formed(mod) + + +def test_error_undeclared_buffer_in_schedulable_tir(): + """In schedule-level TIR (with SBlock nodes), all buffers must be declared.""" + # Manually construct a BufferStore that uses a buffer without any declaration + # inside a block context. + n = tvm.tirx.Var("n", "int32") + A = tvm.tirx.decl_buffer([n], "float32", name="A") + i = tvm.tirx.Var("i", "int32") + + # Create an undeclared buffer using an explicit data pointer that is NOT + # a function parameter and NOT wrapped with DeclBuffer. + B_data = tvm.tirx.Var("B_data", tvm.ir.PointerType(tvm.ir.PrimType("float32"))) + B = tvm.tirx.decl_buffer([n], "float32", name="B", data=B_data) + + # Build a block that writes to B without any declaration of B. + bi = tvm.tirx.Var("bi", "int32") + block = tvm.s_tir.SBlock( + iter_vars=[tvm.tirx.IterVar(tvm.ir.Range(0, n), bi, 0)], # 0 = kDataPar + reads=[tvm.tirx.BufferRegion(A, [tvm.ir.Range(bi, bi + 1)])], + writes=[tvm.tirx.BufferRegion(B, [tvm.ir.Range(bi, bi + 1)])], + body=tvm.tirx.BufferStore(B, tvm.tirx.BufferLoad(A, [bi]), [bi]), + name_hint="write_B", + ) + block_realize = tvm.s_tir.SBlockRealize( + iter_values=[i], + predicate=tvm.tirx.const(True), + block=block, + ) + + prim_func = tvm.tirx.PrimFunc( + params=[A, B_data], + body=tvm.tirx.For(i, 0, n, tvm.tirx.ForKind.SERIAL, block_realize), + # Note: B is NOT a function parameter, so its declaration scope is only + # within a DeclBuffer node (which we intentionally omit here). + ) + + # B is used in the block but was never declared — should fail. + with pytest.raises( + (ValueError, tvm.error.InternalError), match="buffer B.*without a prior DeclBuffer" + ): + tvm.s_tir.analysis.verify_well_formed(prim_func) + + +def test_s_tir_verifier_preserves_shared_definition_check(): + shared = tvm.tirx.Var("shared", "int32") + core = tvm.tirx.PrimFunc([shared], tvm.tirx.Evaluate(shared)) + block = tvm.s_tir.SBlock([], [], [], "block", tvm.tirx.Evaluate(shared)) + scheduled = tvm.tirx.PrimFunc([shared], block).with_attr("s_tir", True) + mod = tvm.IRModule({"core": core, "scheduled": scheduled}) + assert not tvm.s_tir.analysis.verify_well_formed(mod, assert_mode=False) + with pytest.raises(tvm.error.InternalError, match="multiple definitions"): + tvm.s_tir.analysis.verify_well_formed(mod) + + +def test_matched_buffer_is_defined_before_block_regions(): + allocated = tvm.tirx.decl_buffer((4,), "float32", name="allocated") + matched = tvm.tirx.decl_buffer((4,), "float32", name="matched") + source = tvm.tirx.BufferRegion(allocated, [tvm.ir.Range(4)]) + region = tvm.tirx.BufferRegion(matched, [tvm.ir.Range(4)]) + block = tvm.s_tir.SBlock( + [], + [region], + [], + "use", + tvm.tirx.Evaluate(matched[0]), + alloc_buffers=[allocated], + match_buffers=[tvm.s_tir.MatchBufferRegion(matched, source)], + ) + func = tvm.tirx.PrimFunc([], block) + assert tvm.s_tir.analysis.verify_well_formed(func) + out_of_scope = tvm.tirx.PrimFunc([], tvm.tirx.SeqStmt([block, tvm.tirx.Evaluate(matched[0])])) + assert not tvm.s_tir.analysis.verify_well_formed(out_of_scope, assert_mode=False) + + +if __name__ == "__main__": + tvm.testing.main() diff --git a/tests/python/s_tir/transform/test_s_tir_transform_simplify.py b/tests/python/s_tir/transform/test_s_tir_transform_simplify.py new file mode 100644 index 000000000000..e6f979840d54 --- /dev/null +++ b/tests/python/s_tir/transform/test_s_tir_transform_simplify.py @@ -0,0 +1,45 @@ +# 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. + +import tvm +import tvm.testing +from tvm.script import tirx as T + + +def test_s_tir_block_iterator_constraints(): + @T.prim_func(private=True, s_tir=True) + def before(A: T.Buffer((4,), "int32")): + for i in range(4): + with T.sblock("write"): + vi = T.axis.spatial(4, i) + if vi < 4: + A[vi] = vi + + @T.prim_func(private=True, s_tir=True) + def expected(A: T.Buffer((4,), "int32")): + for i in range(4): + with T.sblock("write"): + vi = T.axis.spatial(4, i) + A[vi] = vi + + result = tvm.s_tir.transform.StmtSimplify()(tvm.IRModule.from_expr(before))["main"] + tvm.ir.assert_structural_equal(result, expected) + assert tvm.s_tir.analysis.verify_well_formed(result) + + +if __name__ == "__main__": + tvm.testing.main() diff --git a/tests/python/tirx-analysis/test_tir_analysis_verify_well_formed.py b/tests/python/tirx-analysis/test_tir_analysis_verify_well_formed.py index 20af6b0f231d..ce7fd285992e 100644 --- a/tests/python/tirx-analysis/test_tir_analysis_verify_well_formed.py +++ b/tests/python/tirx-analysis/test_tir_analysis_verify_well_formed.py @@ -27,53 +27,6 @@ from tvm.script import tirx as T -def test_pass_simple(): - @T.prim_func(s_tir=True) - def element_wise( - A: T.Buffer((128, 128), "float32"), - C: T.Buffer((128, 128), "float32"), - ): - B = T.sblock_alloc_buffer((128, 128), "float32") - for i, j in T.grid(128, 128): - with T.sblock("B"): - vi, vj = T.axis.remap("SS", [i, j]) - B[vi, vj] = A[vi, vj] * 2.0 - for i, j in T.grid(128, 128): - with T.sblock("C"): - # It's a opaque block , so it can use outside variables - C[i, j] = B[i, j] * 2.0 - - assert tvm.s_tir.analysis.verify_well_formed(element_wise) - assert tvm.s_tir.analysis.verify_well_formed(tvm.IRModule.from_expr(element_wise)) - - -def test_buffer_region_bounds_are_visited(): - data = tvm.tirx.Var( - "data", tvm.ir.PointerType(tvm.ir.PrimType("int32"), storage_scope="global") - ) - buffer = tvm.tirx.decl_buffer([4], "int32", data=data) - undefined = tvm.tirx.Var("undefined", "int32") - region = tvm.tirx.BufferRegion(buffer, [tvm.ir.Range.from_min_extent(undefined, 4)]) - block = tvm.s_tir.SBlock([], [region], [], "region", tvm.tirx.Evaluate(0)) - func = tvm.tirx.PrimFunc([buffer], block) - assert not tvm.s_tir.analysis.verify_well_formed(func, assert_mode=False) - - -def test_fail_use_out_loop_var(): - @T.prim_func(check_well_formed=False, s_tir=True) - def element_wise( - A: T.Buffer((128, 128), "float32"), - B: T.Buffer((128, 128), "float32"), - ): - for i, j in T.grid(128, 128): - with T.sblock("B"): - vi, vj = T.axis.remap("SS", [i, j]) - # we cannot use `i` since it's defined outside the block - B[vi, vj] = A[i, vj] * 2.0 - - assert not tvm.s_tir.analysis.verify_well_formed(element_wise, assert_mode=False) - - def test_error_for_out_of_scope_usage(): """A variable may not be used after its scope ends. @@ -99,7 +52,7 @@ def test_error_for_out_of_scope_usage(): (ValueError, tvm.error.InternalError), match="Invalid use of undefined variable i at .* no longer in-scope.", ): - tvm.s_tir.analysis.verify_well_formed(func) + tvm.tirx.analysis.verify_well_formed(func) def test_error_for_nested_rebind_usage(): @@ -116,7 +69,7 @@ def func(): (ValueError, tvm.error.InternalError), match="ill-formed, due to multiple nested definitions of variable i", ): - tvm.s_tir.analysis.verify_well_formed(func) + tvm.tirx.analysis.verify_well_formed(func) def test_error_for_repeated_binding(): @@ -138,7 +91,7 @@ def func(): with pytest.raises( (ValueError, tvm.error.InternalError), match="multiple nested definitions of variable i" ): - tvm.s_tir.analysis.verify_well_formed(func) + tvm.tirx.analysis.verify_well_formed(func) def test_error_for_cross_function_reuse(): @@ -161,7 +114,7 @@ def func2(): with pytest.raises( (ValueError, tvm.error.InternalError), match="multiple definitions of variable i" ): - tvm.s_tir.analysis.verify_well_formed(mod) + tvm.tirx.analysis.verify_well_formed(mod) def test_reuse_of_env_thread_in_function_is_well_formed(): @@ -180,7 +133,7 @@ def func(A: T.Buffer([256], "float32")): with T.launch_thread(threadIdx_x, 256): A[threadIdx_x] = A[threadIdx_x] + 2.0 - tvm.s_tir.analysis.verify_well_formed(func) + tvm.tirx.analysis.verify_well_formed(func) def test_reuse_of_env_thread_in_function_is_mandatory(): @@ -201,7 +154,7 @@ def func(A: T.Buffer([256], "float32")): with T.launch_thread("threadIdx.x", 256) as threadIdx_x: A[threadIdx_x] = A[threadIdx_x] + 2.0 - tvm.s_tir.analysis.verify_well_formed(func) + tvm.tirx.analysis.verify_well_formed(func) def test_reuse_of_env_thread_across_functions_is_ill_formed(): @@ -237,7 +190,7 @@ def kernel_2(A: T.Buffer([256], "float32")): with pytest.raises( (ValueError, tvm.error.InternalError), match="multiple definitions of variable threadIdx_x" ): - tvm.s_tir.analysis.verify_well_formed(mod) + tvm.tirx.analysis.verify_well_formed(mod) def test_multiple_buffer_arguments_may_share_allocation(): @@ -257,49 +210,7 @@ def func(A_handle: T.handle, B_handle: T.handle): pass - tvm.s_tir.analysis.verify_well_formed(mod) - - -def test_block_match_buffer_defines_buffer_obj(): - """In a block, T.match_buffer defines a buffer view""" - - @I.ir_module(s_tir=True) - class mod: - @T.prim_func(s_tir=True) - def func(A: T.Buffer([256, 256], "float32")): - for iters in T.grid(16, 16, 16, 16): - with T.sblock("compute"): - tile_i, tile_j, i, j = T.axis.remap("SSSS", iters) - B = T.match_buffer( - A[tile_i * 16 : (tile_i + 1) * 16, tile_j * 16 : (tile_j + 1) * 16], - dtype="float32", - ) - B[i, j] = 0.0 - - tvm.s_tir.analysis.verify_well_formed(mod) - - -def test_block_match_buffer_defines_symbolic_variables(): - """In a block, T.match_buffer may define symbolic variables""" - - @I.ir_module(s_tir=True) - class mod: - @T.prim_func(s_tir=True) - def func(A: T.Buffer([256, 256], "int32")): - for iters in T.grid(16, 16, 16, 16): - with T.sblock("compute"): - tile_i, tile_j, i, j = T.axis.remap("SSSS", iters) - - elem_offset = T.int32() - B = T.match_buffer( - A[tile_i * 16 : (tile_i + 1) * 16, tile_j * 16 : (tile_j + 1) * 16], - dtype="float32", - elem_offset=elem_offset, - ) - - B[i, j] = elem_offset - - tvm.s_tir.analysis.verify_well_formed(mod) + tvm.tirx.analysis.verify_well_formed(mod) def test_error_message_without_previous_definition_location(): @@ -325,7 +236,7 @@ def func(): T.evaluate(x) with pytest.raises((ValueError, tvm.error.InternalError)) as exc_info: - tvm.s_tir.analysis.verify_well_formed(func, assert_mode=True) + tvm.tirx.analysis.verify_well_formed(func, assert_mode=True) error_msg = str(exc_info.value) @@ -350,7 +261,7 @@ def func(): T.evaluate(x) with pytest.raises((ValueError, tvm.error.InternalError)) as exc_info: - tvm.s_tir.analysis.verify_well_formed(func, assert_mode=True) + tvm.tirx.analysis.verify_well_formed(func, assert_mode=True) error_msg = str(exc_info.value) @@ -381,7 +292,7 @@ def func(): T.evaluate(x) with pytest.raises((ValueError, tvm.error.InternalError)) as exc_info: - tvm.s_tir.analysis.verify_well_formed(func, assert_mode=True) + tvm.tirx.analysis.verify_well_formed(func, assert_mode=True) error_msg = str(exc_info.value) @@ -399,7 +310,7 @@ def func(A: T.Buffer((128,), "float32"), B: T.Buffer((128,), "float32")): for i in T.grid(128): B[i] = A[i] * 2.0 - tvm.s_tir.analysis.verify_well_formed(func) + tvm.tirx.analysis.verify_well_formed(func) def test_decl_buffer_is_well_formed(): @@ -411,85 +322,7 @@ def func(A: T.Buffer((128,), "float32")): for i in T.grid(128): B[i] = A[i] * 2.0 - tvm.s_tir.analysis.verify_well_formed(func) - - -def test_alloc_buffer_in_block_is_well_formed(): - """SBlock::alloc_buffers introduces a buffer into scope for the block body.""" - - @I.ir_module(s_tir=True) - class mod: - @T.prim_func(s_tir=True) - def func(A: T.Buffer((128,), "float32")): - with T.sblock("root"): - B = T.sblock_alloc_buffer([128], "float32") - for i in T.grid(128): - with T.sblock("write_B"): - vi = T.axis.remap("S", [i]) - B[vi] = A[vi] * 2.0 - - tvm.s_tir.analysis.verify_well_formed(mod) - - -def test_match_buffer_in_block_is_well_formed(): - """SBlock::match_buffers introduces a buffer into scope for the block body.""" - - @I.ir_module(s_tir=True) - class mod: - @T.prim_func(s_tir=True) - def func(A: T.Buffer((128, 128), "float32")): - for iters in T.grid(8, 8, 16, 16): - with T.sblock("compute"): - ti, tj, i, j = T.axis.remap("SSSS", iters) - A_tile = T.match_buffer( - A[ti * 16 : (ti + 1) * 16, tj * 16 : (tj + 1) * 16], - dtype="float32", - ) - A_tile[i, j] = A_tile[i, j] * 2.0 - - tvm.s_tir.analysis.verify_well_formed(mod) - - -def test_error_undeclared_buffer_in_schedulable_tir(): - """In schedule-level TIR (with SBlock nodes), all buffers must be declared.""" - # Manually construct a BufferStore that uses a buffer without any declaration - # inside a block context. - n = tvm.tirx.Var("n", "int32") - A = tvm.tirx.decl_buffer([n], "float32", name="A") - i = tvm.tirx.Var("i", "int32") - - # Create an undeclared buffer using an explicit data pointer that is NOT - # a function parameter and NOT wrapped with DeclBuffer. - B_data = tvm.tirx.Var("B_data", tvm.ir.PointerType(tvm.ir.PrimType("float32"))) - B = tvm.tirx.decl_buffer([n], "float32", name="B", data=B_data) - - # Build a block that writes to B without any declaration of B. - bi = tvm.tirx.Var("bi", "int32") - block = tvm.s_tir.SBlock( - iter_vars=[tvm.tirx.IterVar(tvm.ir.Range(0, n), bi, 0)], # 0 = kDataPar - reads=[tvm.tirx.BufferRegion(A, [tvm.ir.Range(bi, bi + 1)])], - writes=[tvm.tirx.BufferRegion(B, [tvm.ir.Range(bi, bi + 1)])], - body=tvm.tirx.BufferStore(B, tvm.tirx.BufferLoad(A, [bi]), [bi]), - name_hint="write_B", - ) - block_realize = tvm.s_tir.SBlockRealize( - iter_values=[i], - predicate=tvm.tirx.const(True), - block=block, - ) - - prim_func = tvm.tirx.PrimFunc( - params=[A, B_data], - body=tvm.tirx.For(i, 0, n, tvm.tirx.ForKind.SERIAL, block_realize), - # Note: B is NOT a function parameter, so its declaration scope is only - # within a DeclBuffer node (which we intentionally omit here). - ) - - # B is used in the block but was never declared — should fail. - with pytest.raises( - (ValueError, tvm.error.InternalError), match="buffer B.*without a prior DeclBuffer" - ): - tvm.s_tir.analysis.verify_well_formed(prim_func) + tvm.tirx.analysis.verify_well_formed(func) def test_tensor_load_asserted_type_matches_source_and_indices(): @@ -559,7 +392,8 @@ def test_core_verifiers_reject_blocks(as_module): verify(obj) -def test_mixed_module_parser_checks_both_dialects(): +@pytest.mark.parametrize("legacy_s_tir", [False, True]) +def test_mixed_module_parser_checks_both_dialects(legacy_s_tir): @I.ir_module class Mixed: @T.prim_func(s_tir=True) @@ -573,41 +407,12 @@ def scheduled(A: T.Buffer((4,), "int32")): def lowered(): T.evaluate(0) - assert tvm.s_tir.analysis.verify_well_formed(Mixed) + restored = tvm.script.from_source(Mixed.script(), s_tir=legacy_s_tir) + tvm.ir.assert_structural_equal(restored, Mixed) + assert tvm.s_tir.analysis.verify_well_formed(restored) assert tvm.tirx.analysis.verify_tirx_well_formed(Mixed["lowered"]) assert not tvm.tirx.analysis.verify_tirx_well_formed(Mixed, assert_mode=False) -def test_s_tir_verifier_preserves_shared_definition_check(): - shared = tvm.tirx.Var("shared", "int32") - core = tvm.tirx.PrimFunc([shared], tvm.tirx.Evaluate(shared)) - block = tvm.s_tir.SBlock([], [], [], "block", tvm.tirx.Evaluate(shared)) - scheduled = tvm.tirx.PrimFunc([shared], block).with_attr("s_tir", True) - mod = tvm.IRModule({"core": core, "scheduled": scheduled}) - assert not tvm.s_tir.analysis.verify_well_formed(mod, assert_mode=False) - with pytest.raises(tvm.error.InternalError, match="multiple definitions"): - tvm.s_tir.analysis.verify_well_formed(mod) - - -def test_matched_buffer_is_defined_before_block_regions(): - allocated = tvm.tirx.decl_buffer((4,), "float32", name="allocated") - matched = tvm.tirx.decl_buffer((4,), "float32", name="matched") - source = tvm.tirx.BufferRegion(allocated, [tvm.ir.Range(4)]) - region = tvm.tirx.BufferRegion(matched, [tvm.ir.Range(4)]) - block = tvm.s_tir.SBlock( - [], - [region], - [], - "use", - tvm.tirx.Evaluate(matched[0]), - alloc_buffers=[allocated], - match_buffers=[tvm.s_tir.MatchBufferRegion(matched, source)], - ) - func = tvm.tirx.PrimFunc([], block) - assert tvm.s_tir.analysis.verify_well_formed(func) - out_of_scope = tvm.tirx.PrimFunc([], tvm.tirx.SeqStmt([block, tvm.tirx.Evaluate(matched[0])])) - assert not tvm.s_tir.analysis.verify_well_formed(out_of_scope, assert_mode=False) - - if __name__ == "__main__": tvm.testing.main() diff --git a/tests/python/tirx-transform/test_tir_transform_simplify.py b/tests/python/tirx-transform/test_tir_transform_simplify.py index 44f9eeaa4d11..f62dcf77bac7 100644 --- a/tests/python/tirx-transform/test_tir_transform_simplify.py +++ b/tests/python/tirx-transform/test_tir_transform_simplify.py @@ -1313,26 +1313,5 @@ def expected(a: T.Buffer((2, 8), "int32"), b: T.Buffer((2, 8), "int32")): tvm.ir.assert_structural_equal(after, expected) -def test_s_tir_block_iterator_constraints(): - @T.prim_func(private=True, s_tir=True) - def before(A: T.Buffer((4,), "int32")): - for i in range(4): - with T.sblock("write"): - vi = T.axis.spatial(4, i) - if vi < 4: - A[vi] = vi - - @T.prim_func(private=True, s_tir=True) - def expected(A: T.Buffer((4,), "int32")): - for i in range(4): - with T.sblock("write"): - vi = T.axis.spatial(4, i) - A[vi] = vi - - result = tvm.s_tir.transform.StmtSimplify()(tvm.IRModule.from_expr(before))["main"] - tvm.ir.assert_structural_equal(result, expected) - assert tvm.s_tir.analysis.verify_well_formed(result) - - if __name__ == "__main__": tvm.testing.main() From 56a1b01074de5d8d04acfda8b59e599781e50862 Mon Sep 17 00:00:00 2001 From: Tianqi Chen Date: Thu, 17 Sep 2026 14:40:04 +0000 Subject: [PATCH 15/18] Treat TRN tile implementations as captured fragments TRN tile implementations and initialization bodies capture buffers and indices from the scope where dispatch inserts them. Defer standalone well-formedness checks for these partial functions while retaining verification of complete parsed functions. --- .../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 ++++-- .../tvm/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 ++++-- 10 files changed, 38 insertions(+), 20 deletions(-) diff --git a/python/tvm/backend/trn/tile_primitive/binary/default.py b/python/tvm/backend/trn/tile_primitive/binary/default.py index 3a432c3faddf..85f19d11c591 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 func(*args): 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 802227c17e5d..7c29db37943f 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 get_srcs(inst_gen): # 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 165ce1a2ce95..a3a83bb3a400 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 impl(): 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 9f622cf4045c..2cc80e57c24d 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 impl(): 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 e22b3800b6a2..561fca20c982 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 identity_init(): 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_psum_output(): 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 fa9a67a7c1a0..75467dffc80e 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_inst_macro(lhs_b_loop, rhs_b_loop, reduction_b_loop, acc, C_as_output 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 impl_C_psum(): 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 81ab5827e54a..582b49275678 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 2d791908b216..78cbe25de18b 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 impl(): 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 9106f3c8a65a..8e281c8e9f03 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 f2589a7a7ddc..40eabd560277 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): From 7a8f835c31943af840472d01ee5b45fe0603be4d Mon Sep 17 00:00:00 2001 From: Tianqi Chen Date: Thu, 17 Sep 2026 14:40:16 +0000 Subject: [PATCH 16/18] Bind intrinsic printer operands as function parameters Represent scalar inputs as explicit parameters so intrinsic namespace roundtrips form complete functions under shared definition checks. --- tests/python/tirx/test_printer_tir_namespaces.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/tests/python/tirx/test_printer_tir_namespaces.py b/tests/python/tirx/test_printer_tir_namespaces.py index 369aff9d6378..f7fe49c1b615 100644 --- a/tests/python/tirx/test_printer_tir_namespaces.py +++ b/tests/python/tirx/test_printer_tir_namespaces.py @@ -135,9 +135,8 @@ def test_printer_cuda_more(): def test_printer_cuda_low_level_warp_intrinsics_roundtrip(): - @T.prim_func(check_well_formed=False) - def kernel(): - x = T.int32() + @T.prim_func + def kernel(x: T.int32): mask = T.cuda.__activemask() T.evaluate(T.cuda.__shfl_sync(mask, x, 0, 32)) T.evaluate(T.cuda.__shfl_up_sync(mask, x, 1, 32)) @@ -155,9 +154,8 @@ def kernel(): def test_printer_webgpu_namespace_roundtrip(): - @T.prim_func(check_well_formed=False) - def kernel(): - x = T.int32() + @T.prim_func + def kernel(x: T.int32): T.evaluate(T.webgpu.subgroup_shuffle(x, 0)) T.evaluate(T.webgpu.subgroup_shuffle_up(x, 1)) T.evaluate(T.webgpu.subgroup_shuffle_down(x, 1)) From f90a5feab58309875a8c3f9661e99f675f1844ee Mon Sep 17 00:00:00 2001 From: Tianqi Chen Date: Thu, 17 Sep 2026 14:40:17 +0000 Subject: [PATCH 17/18] Lower opaque blocks before flattening auto-copy buffers Follow the production lowering order when checking auto-copy padding and inspect the resulting core allocation shape. --- ...test_s_tir_transform_memhammer_lower_auto_copy.py | 12 ++++-------- 1 file changed, 4 insertions(+), 8 deletions(-) diff --git a/tests/python/s_tir/transform/test_s_tir_transform_memhammer_lower_auto_copy.py b/tests/python/s_tir/transform/test_s_tir_transform_memhammer_lower_auto_copy.py index 0a5935c86a8c..d537cb2ffee2 100644 --- a/tests/python/s_tir/transform/test_s_tir_transform_memhammer_lower_auto_copy.py +++ b/tests/python/s_tir/transform/test_s_tir_transform_memhammer_lower_auto_copy.py @@ -1138,14 +1138,9 @@ def verify_single_allocation(stmt, alloc_size=None): alloc_extents = [] def verify(n): - if ( - isinstance(n, tvm.s_tir.SBlock) - and n.alloc_buffers is not None - and (True in ((buf.scope() == "shared.dyn") for buf in n.alloc_buffers)) - ): - num_alloc[0] += len(n.alloc_buffers) - for buf in n.alloc_buffers: - alloc_extents.append(buf.shape) + if isinstance(n, tvm.tirx.AllocBuffer) and n.buffer.scope() == "shared.dyn": + num_alloc[0] += 1 + alloc_extents.append(n.buffer.shape) tvm_ffi.structural_walk(stmt, verify) assert num_alloc[0] == 1 @@ -1163,6 +1158,7 @@ def prod(arr): def test_auto_padding(): mod = tvm.s_tir.transform.LowerAutoCopy()(Transpose) + mod = tvm.s_tir.transform.LowerOpaqueBlock()(mod) mod = tvm.tirx.transform.FlattenBuffer()(mod) verify_single_allocation(mod["main"].body, 16 * 130) From 92328fe495e6764807a270598755b302de5c8140 Mon Sep 17 00:00:00 2001 From: Tianqi Chen Date: Thu, 17 Sep 2026 16:41:08 +0000 Subject: [PATCH 18/18] Keep S-TIR coverage in existing suites Remove standalone coverage added for the ownership split and keep existing verification cases in their original suite with S-TIR entry points. Retain the existing API and lowering-phase adaptations. --- tests/cpp/ir_functor_test.cc | 24 - tests/cpp/s_tir_functor_test.cc | 436 ------------------ ...m_specialize_primfunc_based_on_callsite.py | 7 +- .../test_s_tir_analysis_verify_well_formed.py | 225 --------- tests/python/s_tir/test_stmt.py | 173 ------- tests/python/s_tir/test_tensor_intrin.py | 71 --- .../test_s_tir_transform_convert_ssa.py | 63 --- .../test_s_tir_transform_simplify.py | 45 -- .../test_tir_analysis_verify_well_formed.py | 204 ++++++-- tests/python/tirx-base/test_tir_specialize.py | 44 -- .../test_tir_transform_flatten_buffer.py | 2 +- ...tir_transform_force_narrow_index_to_i32.py | 27 -- 12 files changed, 169 insertions(+), 1152 deletions(-) delete mode 100644 tests/cpp/s_tir_functor_test.cc delete mode 100644 tests/python/s_tir/analysis/test_s_tir_analysis_verify_well_formed.py delete mode 100644 tests/python/s_tir/test_stmt.py delete mode 100644 tests/python/s_tir/test_tensor_intrin.py delete mode 100644 tests/python/s_tir/transform/test_s_tir_transform_convert_ssa.py delete mode 100644 tests/python/s_tir/transform/test_s_tir_transform_simplify.py diff --git a/tests/cpp/ir_functor_test.cc b/tests/cpp/ir_functor_test.cc index 9050b1c3c882..3ba10ebe20b3 100644 --- a/tests/cpp/ir_functor_test.cc +++ b/tests/cpp/ir_functor_test.cc @@ -742,27 +742,3 @@ TEST(IRF, SubstituteWithDataTypeLegalizationPreservesShiftAmounts) { EXPECT_TRUE(structural_equal(actual_left, widened_y << shift_amount)); EXPECT_TRUE(structural_equal(actual_right, widened_y >> shift_amount)); } - -TEST(IRF, SubstituteWithDataTypeLegalizationCastsCoreLoopBounds) { - using namespace tvm::prim; - using namespace tvm; - using namespace tvm::tirx; - - PrimVar index("i", PrimType::Int(32)); - PrimVar extent("n", PrimType::Int(32)); - PrimVar wide_extent("n64", PrimType::Int(64)); - Stmt original = For(index, 0, extent, ForKind::kSerial, Evaluate(index)); - Stmt actual = SubstituteWithDataTypeLegalization( - original, [&](const tirx::Var& var) -> ffi::Optional { - if (var.same_as(extent)) return PrimExpr(wide_extent); - return std::nullopt; - }); - - auto* loop = actual.as(); - ASSERT_NE(loop, nullptr); - EXPECT_TRUE(loop->loop_var.same_as(index)); - EXPECT_EQ(loop->min.ty(), index.ty()); - EXPECT_EQ(loop->extent.ty(), index.ty()); - EXPECT_TRUE(ffi::StructuralEqual()(loop->extent, cast(index.ty(), wide_extent))); - EXPECT_TRUE(original.as()->extent.same_as(extent)); -} diff --git a/tests/cpp/s_tir_functor_test.cc b/tests/cpp/s_tir_functor_test.cc deleted file mode 100644 index 3e756b96e13b..000000000000 --- a/tests/cpp/s_tir_functor_test.cc +++ /dev/null @@ -1,436 +0,0 @@ -/* - * 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 -#include -#include -#include - -#include -#include -#include - -namespace tvm { -namespace s_tir { -namespace { - -using namespace tirx; - -TEST(STIRFunctor, LegacyInheritedDispatchAndContainsNode) { - class Dispatch : public StmtFunctor { - public: - using StmtFunctor::Dispatch_; - int Dispatch_(const SBlockNode*, int value) final { return value + 1; } - int Dispatch_(const SBlockRealizeNode*, int value) final { return value + 2; } - int Dispatch_(const EvaluateNode*, int value) final { return value + 3; } - } dispatch; - Stmt body = Evaluate(0); - SBlock block({}, {}, {}, "block", body); - Stmt realize = SBlockRealize({}, IntImm::Bool(true), block); - EXPECT_EQ(dispatch(block, 10), 11); - EXPECT_EQ(dispatch(realize, 10), 12); - EXPECT_EQ(dispatch(body, 10), 13); - EXPECT_TRUE(ContainsNode(realize)); - EXPECT_TRUE(ContainsNode(realize)); - EXPECT_FALSE(ContainsNode(realize)); -} - -TEST(STIRFunctor, CoreLegacyDispatchReachesDefaultForDialectNodes) { - class Dispatch : public tirx::StmtFunctor { - public: - using tirx::StmtFunctor::Dispatch_; - bool Dispatch_(const EvaluateNode*) final { return true; } - bool DispatchDefault_(const ffi::Object*) final { return false; } - } dispatch; - SBlock block({}, {}, {}, "block", Evaluate(0)); - EXPECT_FALSE(dispatch(block)); - EXPECT_FALSE(dispatch(SBlockRealize({}, IntImm::Bool(true), block))); - EXPECT_TRUE(dispatch(block->body)); -} - -TEST(STIRFunctor, NativeBlockOverrideReusesInheritedCoreHooks) { - class Visitor : public StmtExprVisitor { - public: - using StmtExprVisitor::Visit_; - ffi::Optional Visit_(const SBlockNode* op) final { - ++blocks; - return StmtExprVisitor::Visit_(op); - } - ffi::Optional Visit_(const EvaluateNode* op) final { - ++evaluates; - return StmtExprVisitor::Visit_(op); - } - int blocks = 0; - int evaluates = 0; - }; - SBlock block({}, {}, {}, "block", Evaluate(0)); - auto visitor = ffi::make_object(); - visitor->Visit(SBlockRealize({}, IntImm::Bool(true), block)); - EXPECT_EQ(visitor->blocks, 1); - EXPECT_EQ(visitor->evaluates, 1); -} - -TEST(STIRFunctor, NativeVisitPreservesBlockOrderAndBinders) { - PrimVar index("index"), extent("extent"), annotation("annotation"); - BufferVar buffer = decl_buffer({16}); - TensorRegion region = BufferRegion(buffer, {Range::FromMinExtent(0, 16)}); - IterVar iter(Range::FromMinExtent(0, extent), index, IterVarType::kDataPar); - SBlock block({iter}, {region}, {}, "block", Evaluate(index), std::nullopt, {buffer}, {}, - {{"annotation", annotation}}); - class Visitor : public StmtExprVisitor { - public: - using StmtExprVisitor::Visit_; - ffi::Optional Visit_(const VarNode* var) final { - vars.push_back(var); - if (var->ty.as()) { - buffer_regions.push_back(def_region_kind()); - } - return std::nullopt; - } - std::vector vars; - std::vector buffer_regions; - }; - auto visitor = ffi::make_object(); - visitor->Visit(block); - EXPECT_EQ(std::count(visitor->vars.begin(), visitor->vars.end(), index.get()), 1); - EXPECT_EQ(std::count(visitor->vars.begin(), visitor->vars.end(), annotation.get()), 0); - ASSERT_EQ(visitor->buffer_regions.size(), 2); - EXPECT_EQ(visitor->buffer_regions[0], kTVMFFIDefRegionKindSimple); - EXPECT_EQ(visitor->buffer_regions[1], kTVMFFIDefRegionKindNone); - - // A full structural walk retains its distinct binder/annotation traversal. - int structural_index = 0; - int structural_annotation = 0; - ffi::StructuralWalk( - block, [&](const Var& var) -> ffi::Expected { - structural_index += var.same_as(index); - structural_annotation += var.same_as(annotation); - return ffi::WalkResult::Advance(); - }); - EXPECT_EQ(structural_index, 2); - EXPECT_EQ(structural_annotation, 1); -} - -TEST(STIRFunctor, NativeMutationKeepsAnnotationsAndSharedIteratorBinders) { - PrimVar index("index"), extent("extent"); - PrimExpr expression = extent + 1; - IterVar iter(Range::FromMinExtent(0, expression), index, IterVarType::kDataPar); - SBlock block({iter}, {}, {}, "block", Evaluate(expression), std::nullopt, {}, {}, - {{"annotation", expression}}); - SBlock retained = block; - class Mutator : public StmtExprMutator { - public: - using StmtExprMutator::Mutate_; - UnchangedOr Mutate_(const prim::AddNode* op, InplaceMode) final { return op->a; } - }; - auto mutator = ffi::make_object(); - Stmt result = mutator->Mutate(block, InplaceMode::kAllow).ValueOrUnchanged(block); - const auto* changed = result.as(); - ASSERT_NE(changed, nullptr); - EXPECT_NE(changed, retained.get()); - EXPECT_TRUE(changed->iter_vars[0]->var.same_as(index)); - EXPECT_TRUE(changed->iter_vars[0]->dom->extent.same_as(extent)); - EXPECT_TRUE(changed->body.as()->value.same_as(extent)); - EXPECT_TRUE(changed->annotations.at("annotation").cast().same_as(expression)); - EXPECT_TRUE(retained->iter_vars[0]->dom->extent.same_as(expression)); - EXPECT_TRUE(iter->dom->extent.same_as(expression)); - - // A sole owning block can update in place while its shared iterator is copied. - SBlock unique({iter}, {}, {}, "unique", Evaluate(expression)); - const auto* original = unique.get(); - auto update = mutator->Mutate(unique, InplaceMode::kAllow); - EXPECT_TRUE(update.IsUnchanged()); - EXPECT_EQ(unique.get(), original); - EXPECT_TRUE(unique->iter_vars[0]->dom->extent.same_as(extent)); - EXPECT_TRUE(iter->dom->extent.same_as(expression)); -} - -template -void CheckMutationRemapsBufferDefinitionsAndUses() { - PrimVar extent("extent"); - BufferVar allocated = decl_buffer({extent + 1}, PrimType::Int(32)); - BufferVar matched = decl_buffer({extent + 1}, PrimType::Int(32)); - TensorRegion region = BufferRegion(allocated, {Range::FromMinExtent(0, extent + 1)}); - MatchBufferRegion match(matched, region); - TensorRegion matched_region = BufferRegion(matched, {Range::FromMinExtent(0, extent + 1)}); - Stmt body = SeqStmt({BufferStore(allocated, 0, {0}), BufferStore(matched, 0, {0})}); - SBlock block({}, {region, matched_region}, {region, matched_region}, "block", body, std::nullopt, - {allocated}, {match}); - class Mutator : public Base { - public: - using Base::Mutate_; - UnchangedOr Mutate_(const prim::AddNode* op, InplaceMode) final { return op->a; } - }; - auto mutator = ffi::make_object(); - Stmt result = mutator->Mutate(block).ValueOrUnchanged(block); - const auto* changed = result.as(); - ASSERT_NE(changed, nullptr); - BufferVar new_allocated = changed->alloc_buffers[0]; - BufferVar new_matched = changed->match_buffers[0]->buffer; - EXPECT_FALSE(new_allocated.same_as(allocated)); - EXPECT_FALSE(new_matched.same_as(matched)); - EXPECT_TRUE(new_allocated->shape[0].same_as(extent)); - EXPECT_TRUE(new_matched->shape[0].same_as(extent)); - EXPECT_TRUE(changed->reads[0]->source.as_or_throw().same_as(new_allocated)); - EXPECT_TRUE(changed->writes[0]->source.as_or_throw().same_as(new_allocated)); - EXPECT_TRUE(changed->reads[1]->source.as_or_throw().same_as(new_matched)); - EXPECT_TRUE(changed->writes[1]->source.as_or_throw().same_as(new_matched)); - EXPECT_TRUE( - changed->match_buffers[0]->source->source.as_or_throw().same_as(new_allocated)); - const auto* statements = changed->body.as(); - ASSERT_NE(statements, nullptr); - EXPECT_TRUE(statements->seq[0].as()->buffer.same_as(new_allocated)); - EXPECT_TRUE(statements->seq[1].as()->buffer.same_as(new_matched)); - EXPECT_TRUE(block->alloc_buffers[0].same_as(allocated)); - EXPECT_TRUE(block->match_buffers[0]->buffer.same_as(matched)); -} - -TEST(STIRFunctor, NativeMutationRemapsBufferDefinitionsAndUses) { - CheckMutationRemapsBufferDefinitionsAndUses(); -} - -TEST(STIRFunctor, GenericTIRXMutationRemapsBufferDefinitionsAndUses) { - CheckMutationRemapsBufferDefinitionsAndUses(); -} - -TEST(STIRFunctor, StructuralAndGenericSubstitutionPreserveDefinitionUses) { - PrimVar extent("extent"), new_extent("new_extent"), index("index"), new_index("new_index"); - BufferVar allocated = decl_buffer({extent}, PrimType::Int(32)); - BufferVar matched = decl_buffer({extent}, PrimType::Int(32)); - TensorRegion region = BufferRegion(allocated, {Range::FromMinExtent(0, extent)}); - TensorRegion matched_region = BufferRegion(matched, {Range::FromMinExtent(0, extent)}); - MatchBufferRegion match(matched, region); - IterVar iter(Range::FromMinExtent(0, extent), index, IterVarType::kDataPar); - Stmt body = - SeqStmt({BufferStore(allocated, extent, {index}), BufferStore(matched, extent, {index})}); - SBlock block({iter}, {region, matched_region}, {region, matched_region}, "block", body, - Evaluate(extent), {allocated}, {match}, {{"annotation", extent}}); - Stmt original = SBlockRealize({extent}, extent > 0, block); - - auto check = [&](const Stmt& result) { - const auto* realize = result.as(); - ASSERT_NE(realize, nullptr); - EXPECT_TRUE(realize->iter_values[0].same_as(new_extent)); - EXPECT_TRUE(realize->predicate.as()->a.same_as(new_extent)); - const auto* changed = realize->block.get(); - EXPECT_TRUE(changed->iter_vars[0]->var.same_as(new_index)); - EXPECT_TRUE(changed->iter_vars[0]->dom->extent.same_as(new_extent)); - EXPECT_TRUE(changed->alloc_buffers[0]->shape[0].same_as(new_extent)); - EXPECT_TRUE(changed->match_buffers[0]->buffer->shape[0].same_as(new_extent)); - EXPECT_TRUE(changed->match_buffers[0]->source->source.as_or_throw().same_as( - changed->alloc_buffers[0])); - EXPECT_TRUE( - changed->reads[0]->source.as_or_throw().same_as(changed->alloc_buffers[0])); - EXPECT_TRUE(changed->reads[1]->source.as_or_throw().same_as( - changed->match_buffers[0]->buffer)); - EXPECT_TRUE( - changed->writes[0]->source.as_or_throw().same_as(changed->alloc_buffers[0])); - EXPECT_TRUE(changed->writes[1]->source.as_or_throw().same_as( - changed->match_buffers[0]->buffer)); - EXPECT_TRUE(changed->reads[0]->region[0]->extent.same_as(new_extent)); - EXPECT_TRUE(changed->reads[1]->region[0]->extent.same_as(new_extent)); - const auto* statements = changed->body.as(); - ASSERT_NE(statements, nullptr); - const auto* store = statements->seq[0].as(); - EXPECT_TRUE(store->buffer.same_as(changed->alloc_buffers[0])); - EXPECT_TRUE(store->value.same_as(new_extent)); - EXPECT_TRUE(store->indices[0].same_as(new_index)); - EXPECT_TRUE(statements->seq[1].as()->buffer.same_as( - changed->match_buffers[0]->buffer)); - EXPECT_TRUE(changed->init.value().as()->value.same_as(new_extent)); - EXPECT_TRUE(changed->annotations.at("annotation").cast().same_as(new_extent)); - EXPECT_TRUE(block->iter_vars[0]->var.same_as(index)); - EXPECT_TRUE(block->iter_vars[0]->dom->extent.same_as(extent)); - EXPECT_TRUE(block->alloc_buffers[0].same_as(allocated)); - EXPECT_TRUE(block->match_buffers[0]->buffer.same_as(matched)); - EXPECT_TRUE(block->annotations.at("annotation").cast().same_as(extent)); - }; - - Stmt structural = ffi::StructuralMap( - original, - [&](const Var& var) -> ffi::Expected> { - if (var.same_as(extent)) return ffi::Any(new_extent); - if (var.same_as(index)) return ffi::Any(new_index); - return ffi::Unchanged(); - }) - .as_or_throw(); - check(structural); - Stmt generic = - SubstituteWithDataTypeLegalization(original, [&](const Var& var) -> ffi::Optional { - if (var.same_as(extent)) return new_extent; - if (var.same_as(index)) return new_index; - return std::nullopt; - }); - check(generic); - - // Unique outer nodes may update in place, but their shared child arrays and - // regions must not modify the retained block used to construct them. - for (bool use_generic : {false, true}) { - SBlock local(block->iter_vars, block->reads, block->writes, "unique", block->body, block->init, - block->alloc_buffers, block->match_buffers, block->annotations); - const auto* block_identity = local.get(); - Stmt input = SBlockRealize({extent}, extent > 0, std::move(local)); - const auto* realize_identity = input.get(); - Stmt result; - if (use_generic) { - result = SubstituteWithDataTypeLegalization(std::move(input), - [&](const Var& var) -> ffi::Optional { - if (var.same_as(extent)) return new_extent; - if (var.same_as(index)) return new_index; - return std::nullopt; - }); - } else { - result = ffi::StructuralMap( - std::move(input), - [&](const Var& var) -> ffi::Expected> { - if (var.same_as(extent)) return ffi::Any(new_extent); - if (var.same_as(index)) return ffi::Any(new_index); - return ffi::Unchanged(); - }) - .as_or_throw(); - } - EXPECT_EQ(result.get(), realize_identity); - EXPECT_EQ(result.as()->block.get(), block_identity); - check(result); - } -} - -TEST(STIRFunctor, GenericTIRXVisitorUsesFullStructuralTraversal) { - PrimVar index("index"), annotation("annotation"); - BufferVar buffer = decl_buffer({16}); - TensorRegion region = BufferRegion(buffer, {Range::FromMinExtent(0, 16)}); - IterVar iter(Range::FromMinExtent(0, 16), index, IterVarType::kDataPar); - SBlock block({iter}, {region}, {}, "block", Evaluate(index), std::nullopt, {buffer}, {}, - {{"annotation", annotation}}); - class Visitor : public tirx::StmtExprVisitor { - public: - using tirx::StmtExprVisitor::Visit_; - ffi::Optional Visit_(const VarNode* var) final { - vars.push_back(var); - if (var->ty.as()) { - buffer_regions.push_back(def_region_kind()); - } - return std::nullopt; - } - std::vector vars; - std::vector buffer_regions; - }; - auto visitor = ffi::make_object(); - visitor->Visit(block); - EXPECT_EQ(std::count(visitor->vars.begin(), visitor->vars.end(), index.get()), 2); - EXPECT_EQ(std::count(visitor->vars.begin(), visitor->vars.end(), annotation.get()), 1); - ASSERT_EQ(visitor->buffer_regions.size(), 2); - EXPECT_EQ(visitor->buffer_regions[0], kTVMFFIDefRegionKindSimple); - EXPECT_EQ(visitor->buffer_regions[1], kTVMFFIDefRegionKindNone); -} - -TEST(STIRFunctor, GenericTIRXMutationRemapsBindersAndAnnotations) { - PrimVar index("index"), replacement("replacement"), extent("extent"); - PrimExpr expression = extent + 1; - IterVar iter(Range::FromMinExtent(0, expression), index, IterVarType::kDataPar); - BufferVar buffer = decl_buffer({expression}, PrimType::Int(32)); - SBlock block({iter}, {}, {}, "block", BufferStore(buffer, index, {0}), std::nullopt, {buffer}, {}, - {{"annotation", expression}}); - class Mutator : public tirx::StmtExprMutator { - public: - using tirx::StmtExprMutator::Mutate_; - UnchangedOr Mutate_(const prim::AddNode* op, InplaceMode) final { return op->a; } - }; - auto mutator = ffi::make_object(); - mutator->VarRemapSet(index, replacement); - SBlock retained = block; - Stmt result = mutator->Mutate(block, InplaceMode::kAllow).ValueOrUnchanged(block); - const auto* changed = result.as(); - ASSERT_NE(changed, nullptr); - EXPECT_TRUE(changed->iter_vars[0]->var.same_as(replacement)); - EXPECT_TRUE(changed->iter_vars[0]->dom->extent.same_as(extent)); - EXPECT_TRUE(changed->annotations.at("annotation").cast().same_as(extent)); - EXPECT_TRUE(changed->body.as()->value.same_as(replacement)); - EXPECT_TRUE(changed->body.as()->buffer.same_as(changed->alloc_buffers[0])); - EXPECT_TRUE(changed->alloc_buffers[0]->shape[0].same_as(extent)); - EXPECT_TRUE(block->iter_vars[0]->var.same_as(index)); - EXPECT_TRUE(block->annotations.at("annotation").cast().same_as(expression)); - - SBlock unique({}, {}, {}, "unique", Evaluate(expression), std::nullopt, {}, {}, - {{"annotation", expression}}); - const auto* original = unique.get(); - auto update = mutator->Mutate(unique, InplaceMode::kAllow); - EXPECT_TRUE(update.IsUnchanged()); - EXPECT_EQ(unique.get(), original); - EXPECT_TRUE(unique->body.as()->value.same_as(extent)); - EXPECT_TRUE(unique->annotations.at("annotation").cast().same_as(extent)); -} - -TEST(STIRFunctor, GenericTIRXFallbackPreservesInterruptAndErrorIdentity) { - PrimVar annotation("annotation"), body("body"); - SBlock block({}, {}, {}, "block", Evaluate(body), std::nullopt, {}, {}, - {{"annotation", annotation}}); - class Visitor : public tirx::StmtExprVisitor { - public: - using tirx::StmtExprVisitor::Visit_; - ffi::Optional Visit_(const VarNode* op) final { - ++count; - return VisitInterrupt(ffi::GetRef(op)); - } - int count = 0; - }; - auto visitor = ffi::make_object(); - auto interrupt = visitor->Visit(block); - ASSERT_TRUE(interrupt.has_value()); - EXPECT_TRUE(interrupt.value()->value.cast().same_as(annotation)); - EXPECT_EQ(visitor->count, 1); - - class Mutator : public tirx::StmtExprMutator { - public: - using tirx::StmtExprMutator::Mutate_; - ffi::Error error{"ValueError", "block child mutation error", ""}; - UnchangedOr Mutate_(const VarNode*, InplaceMode) final { throw error; } - }; - auto mutator = ffi::make_object(); - auto result = mutator->MutateExpected(block); - ASSERT_TRUE(result.is_err()); - EXPECT_TRUE(result.error().same_as(mutator->error)); - auto context = ffi::VisitErrorContext::TryGetFromError(result.error()); - ASSERT_TRUE(context.has_value()); - EXPECT_TRUE(context.value()->reverse_visit_pattern.back().same_as(block)); -} - -TEST(STIRFunctor, NativeInterruptStopsBeforeBlockBody) { - PrimVar stop("stop"), body("body"); - IterVar iter(Range::FromMinExtent(0, 16), PrimVar("index"), IterVarType::kDataPar); - SBlock block({iter}, {}, {}, "block", Evaluate(body)); - Stmt realize = SBlockRealize({stop}, IntImm::Bool(true), block); - class Visitor : public StmtExprVisitor { - public: - using StmtExprVisitor::Visit_; - ffi::Optional Visit_(const VarNode* op) final { - ++count; - return VisitInterrupt(ffi::GetRef(op)); - } - int count = 0; - }; - auto visitor = ffi::make_object(); - auto interrupt = visitor->Visit(realize); - ASSERT_TRUE(interrupt.has_value()); - EXPECT_TRUE(interrupt.value()->value.cast().same_as(stop)); - EXPECT_EQ(visitor->count, 1); -} - -} // namespace -} // namespace s_tir -} // namespace tvm diff --git a/tests/python/relax/test_transform_specialize_primfunc_based_on_callsite.py b/tests/python/relax/test_transform_specialize_primfunc_based_on_callsite.py index 4b2867d0714c..b9c9b16e1ea9 100644 --- a/tests/python/relax/test_transform_specialize_primfunc_based_on_callsite.py +++ b/tests/python/relax/test_transform_specialize_primfunc_based_on_callsite.py @@ -80,7 +80,6 @@ def verify(input): ValidateBufferScopes(False).visit(input) mod = tvm.relax.transform.SpecializePrimFuncBasedOnCallSite()(input) ValidateBufferScopes(True).visit(mod) - return mod def test_single_arg_return(): @@ -209,11 +208,7 @@ def main( R.output(gv2) return gv2 - specialized = verify(Input) - # This pass runs before DLight: specialized blocks must remain schedulable. - schedule = tvm.s_tir.Schedule(specialized, debug_mask="all") - block = schedule.get_sblock("pool_max", func_name="max_pool2d_opencl") - assert len(schedule.get_loops(block)) == 7 + verify(Input) def test_multi_arg_return(): diff --git a/tests/python/s_tir/analysis/test_s_tir_analysis_verify_well_formed.py b/tests/python/s_tir/analysis/test_s_tir_analysis_verify_well_formed.py deleted file mode 100644 index 357e2d288e40..000000000000 --- a/tests/python/s_tir/analysis/test_s_tir_analysis_verify_well_formed.py +++ /dev/null @@ -1,225 +0,0 @@ -# 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. - -import pytest - -import tvm -import tvm.testing -from tvm.script import ir as I -from tvm.script import tirx as T - - -def test_pass_simple(): - @T.prim_func(s_tir=True) - def element_wise( - A: T.Buffer((128, 128), "float32"), - C: T.Buffer((128, 128), "float32"), - ): - B = T.sblock_alloc_buffer((128, 128), "float32") - for i, j in T.grid(128, 128): - with T.sblock("B"): - vi, vj = T.axis.remap("SS", [i, j]) - B[vi, vj] = A[vi, vj] * 2.0 - for i, j in T.grid(128, 128): - with T.sblock("C"): - # It's a opaque block , so it can use outside variables - C[i, j] = B[i, j] * 2.0 - - assert tvm.s_tir.analysis.verify_well_formed(element_wise) - assert tvm.s_tir.analysis.verify_well_formed(tvm.IRModule.from_expr(element_wise)) - - -def test_buffer_region_bounds_are_visited(): - data = tvm.tirx.Var( - "data", tvm.ir.PointerType(tvm.ir.PrimType("int32"), storage_scope="global") - ) - buffer = tvm.tirx.decl_buffer([4], "int32", data=data) - undefined = tvm.tirx.Var("undefined", "int32") - region = tvm.tirx.BufferRegion(buffer, [tvm.ir.Range.from_min_extent(undefined, 4)]) - block = tvm.s_tir.SBlock([], [region], [], "region", tvm.tirx.Evaluate(0)) - func = tvm.tirx.PrimFunc([buffer], block) - assert not tvm.s_tir.analysis.verify_well_formed(func, assert_mode=False) - - -def test_fail_use_out_loop_var(): - @T.prim_func(check_well_formed=False, s_tir=True) - def element_wise( - A: T.Buffer((128, 128), "float32"), - B: T.Buffer((128, 128), "float32"), - ): - for i, j in T.grid(128, 128): - with T.sblock("B"): - vi, vj = T.axis.remap("SS", [i, j]) - # we cannot use `i` since it's defined outside the block - B[vi, vj] = A[i, vj] * 2.0 - - assert not tvm.s_tir.analysis.verify_well_formed(element_wise, assert_mode=False) - - -def test_block_match_buffer_defines_buffer_obj(): - """In a block, T.match_buffer defines a buffer view""" - - @I.ir_module(s_tir=True) - class mod: - @T.prim_func(s_tir=True) - def func(A: T.Buffer([256, 256], "float32")): - for iters in T.grid(16, 16, 16, 16): - with T.sblock("compute"): - tile_i, tile_j, i, j = T.axis.remap("SSSS", iters) - B = T.match_buffer( - A[tile_i * 16 : (tile_i + 1) * 16, tile_j * 16 : (tile_j + 1) * 16], - dtype="float32", - ) - B[i, j] = 0.0 - - tvm.s_tir.analysis.verify_well_formed(mod) - - -def test_block_match_buffer_defines_symbolic_variables(): - """In a block, T.match_buffer may define symbolic variables""" - - @I.ir_module(s_tir=True) - class mod: - @T.prim_func(s_tir=True) - def func(A: T.Buffer([256, 256], "int32")): - for iters in T.grid(16, 16, 16, 16): - with T.sblock("compute"): - tile_i, tile_j, i, j = T.axis.remap("SSSS", iters) - - elem_offset = T.int32() - B = T.match_buffer( - A[tile_i * 16 : (tile_i + 1) * 16, tile_j * 16 : (tile_j + 1) * 16], - dtype="float32", - elem_offset=elem_offset, - ) - - B[i, j] = elem_offset - - tvm.s_tir.analysis.verify_well_formed(mod) - - -def test_alloc_buffer_in_block_is_well_formed(): - """SBlock::alloc_buffers introduces a buffer into scope for the block body.""" - - @I.ir_module(s_tir=True) - class mod: - @T.prim_func(s_tir=True) - def func(A: T.Buffer((128,), "float32")): - with T.sblock("root"): - B = T.sblock_alloc_buffer([128], "float32") - for i in T.grid(128): - with T.sblock("write_B"): - vi = T.axis.remap("S", [i]) - B[vi] = A[vi] * 2.0 - - tvm.s_tir.analysis.verify_well_formed(mod) - - -def test_match_buffer_in_block_is_well_formed(): - """SBlock::match_buffers introduces a buffer into scope for the block body.""" - - @I.ir_module(s_tir=True) - class mod: - @T.prim_func(s_tir=True) - def func(A: T.Buffer((128, 128), "float32")): - for iters in T.grid(8, 8, 16, 16): - with T.sblock("compute"): - ti, tj, i, j = T.axis.remap("SSSS", iters) - A_tile = T.match_buffer( - A[ti * 16 : (ti + 1) * 16, tj * 16 : (tj + 1) * 16], - dtype="float32", - ) - A_tile[i, j] = A_tile[i, j] * 2.0 - - tvm.s_tir.analysis.verify_well_formed(mod) - - -def test_error_undeclared_buffer_in_schedulable_tir(): - """In schedule-level TIR (with SBlock nodes), all buffers must be declared.""" - # Manually construct a BufferStore that uses a buffer without any declaration - # inside a block context. - n = tvm.tirx.Var("n", "int32") - A = tvm.tirx.decl_buffer([n], "float32", name="A") - i = tvm.tirx.Var("i", "int32") - - # Create an undeclared buffer using an explicit data pointer that is NOT - # a function parameter and NOT wrapped with DeclBuffer. - B_data = tvm.tirx.Var("B_data", tvm.ir.PointerType(tvm.ir.PrimType("float32"))) - B = tvm.tirx.decl_buffer([n], "float32", name="B", data=B_data) - - # Build a block that writes to B without any declaration of B. - bi = tvm.tirx.Var("bi", "int32") - block = tvm.s_tir.SBlock( - iter_vars=[tvm.tirx.IterVar(tvm.ir.Range(0, n), bi, 0)], # 0 = kDataPar - reads=[tvm.tirx.BufferRegion(A, [tvm.ir.Range(bi, bi + 1)])], - writes=[tvm.tirx.BufferRegion(B, [tvm.ir.Range(bi, bi + 1)])], - body=tvm.tirx.BufferStore(B, tvm.tirx.BufferLoad(A, [bi]), [bi]), - name_hint="write_B", - ) - block_realize = tvm.s_tir.SBlockRealize( - iter_values=[i], - predicate=tvm.tirx.const(True), - block=block, - ) - - prim_func = tvm.tirx.PrimFunc( - params=[A, B_data], - body=tvm.tirx.For(i, 0, n, tvm.tirx.ForKind.SERIAL, block_realize), - # Note: B is NOT a function parameter, so its declaration scope is only - # within a DeclBuffer node (which we intentionally omit here). - ) - - # B is used in the block but was never declared — should fail. - with pytest.raises( - (ValueError, tvm.error.InternalError), match="buffer B.*without a prior DeclBuffer" - ): - tvm.s_tir.analysis.verify_well_formed(prim_func) - - -def test_s_tir_verifier_preserves_shared_definition_check(): - shared = tvm.tirx.Var("shared", "int32") - core = tvm.tirx.PrimFunc([shared], tvm.tirx.Evaluate(shared)) - block = tvm.s_tir.SBlock([], [], [], "block", tvm.tirx.Evaluate(shared)) - scheduled = tvm.tirx.PrimFunc([shared], block).with_attr("s_tir", True) - mod = tvm.IRModule({"core": core, "scheduled": scheduled}) - assert not tvm.s_tir.analysis.verify_well_formed(mod, assert_mode=False) - with pytest.raises(tvm.error.InternalError, match="multiple definitions"): - tvm.s_tir.analysis.verify_well_formed(mod) - - -def test_matched_buffer_is_defined_before_block_regions(): - allocated = tvm.tirx.decl_buffer((4,), "float32", name="allocated") - matched = tvm.tirx.decl_buffer((4,), "float32", name="matched") - source = tvm.tirx.BufferRegion(allocated, [tvm.ir.Range(4)]) - region = tvm.tirx.BufferRegion(matched, [tvm.ir.Range(4)]) - block = tvm.s_tir.SBlock( - [], - [region], - [], - "use", - tvm.tirx.Evaluate(matched[0]), - alloc_buffers=[allocated], - match_buffers=[tvm.s_tir.MatchBufferRegion(matched, source)], - ) - func = tvm.tirx.PrimFunc([], block) - assert tvm.s_tir.analysis.verify_well_formed(func) - out_of_scope = tvm.tirx.PrimFunc([], tvm.tirx.SeqStmt([block, tvm.tirx.Evaluate(matched[0])])) - assert not tvm.s_tir.analysis.verify_well_formed(out_of_scope, assert_mode=False) - - -if __name__ == "__main__": - tvm.testing.main() diff --git a/tests/python/s_tir/test_stmt.py b/tests/python/s_tir/test_stmt.py deleted file mode 100644 index 66f887ff9ba5..000000000000 --- a/tests/python/s_tir/test_stmt.py +++ /dev/null @@ -1,173 +0,0 @@ -# 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. -"""S-TIR node ownership and serialization compatibility.""" - -import json - -import pytest - -import tvm -import tvm.testing -from tvm import s_tir, tirx -from tvm.ir.json_compact import upgrade_json - - -@pytest.mark.parametrize("legacy", [False, True]) -@pytest.mark.parametrize("legacy_region", [False, True]) -def test_sblock_serialization(legacy, legacy_region): - source = tirx.decl_buffer((4,), "float32", name="source") - target = tirx.decl_buffer((4,), "float32", name="target") - span = tvm.ir.Span(tvm.ir.SourceName("region"), 2, 3, 4, 5) - region = tvm.ir.TensorRegion(source, [tvm.ir.Range(0, 4)], tirx.BufferRegionType(), span) - match = s_tir.MatchBufferRegion(target, region) - block = s_tir.SBlock([], [region], [region], "copy", tirx.Evaluate(0), match_buffers=[match]) - realize = s_tir.SBlockRealize([], True, block) - graph = json.loads(tvm.ir.save_json([block, realize, match])) - type_keys = {node.get("type") for node in graph["nodes"]} - for name in ("SBlock", "SBlockRealize", "MatchBufferRegion"): - assert f"s_tir.{name}" in type_keys - assert f"tirx.{name}" not in type_keys - assert getattr(s_tir, name).__module__ == "tvm.s_tir.stmt" - assert not hasattr(tirx, name) - assert not hasattr(tirx.stmt, name) - assert "ir.TensorRegion" in type_keys - assert "tirx.BufferRegion" not in type_keys - if legacy_region: - for node in graph["nodes"]: - if node.get("type") == "ir.TensorRegion": - node["type"] = "tirx.BufferRegion" - node["data"]["buffer"] = node["data"].pop("source") - if legacy: - for node in graph["nodes"]: - if node.get("type") in { - "s_tir.SBlock", - "s_tir.SBlockRealize", - "s_tir.MatchBufferRegion", - }: - node["type"] = node["type"].replace("s_tir.", "tirx.") - restored_block, restored_realize, restored_match = tvm.ir.load_json(json.dumps(graph)) - assert isinstance(restored_block, s_tir.SBlock) - assert isinstance(restored_realize, s_tir.SBlockRealize) - assert isinstance(restored_match, s_tir.MatchBufferRegion) - assert isinstance(restored_block, tirx.Stmt) - assert isinstance(restored_realize, tirx.Stmt) - assert restored_realize.block.same_as(restored_block) - assert restored_block.match_buffers[0].same_as(restored_match) - assert restored_block.reads[0].same_as(restored_block.writes[0]) - assert restored_match.source.same_as(restored_block.reads[0]) - restored_region = restored_match.source - assert isinstance(restored_region, tvm.ir.TensorRegion) - assert isinstance(restored_region.ty, tirx.BufferRegionType) - assert restored_region.span.source_name.name == "region" - assert restored_region.span.line == 2 - assert restored_region.span.end_line == 3 - assert restored_region.span.column == 4 - assert restored_region.span.end_column == 5 - tvm.ir.assert_structural_equal(restored_realize, realize, map_free_vars=True) - canonical = json.loads(tvm.ir.save_json([restored_block, restored_realize, restored_match])) - assert not any( - node.get("type") - in { - "tirx.BufferRegion", - "tirx.SBlock", - "tirx.SBlockRealize", - "tirx.MatchBufferRegion", - } - for node in canonical["nodes"] - ) - - -def test_untyped_buffer_region_serialization(): - source = tirx.decl_buffer((4,), "float32", name="source") - region = [tvm.ir.Range(0, 4)] - graph = json.loads(tvm.ir.save_json([source, region])) - nodes = graph["nodes"] - source_index, region_index = nodes[graph["root_index"]]["data"] - # Before c836e8c942, BufferRegion inherited PrimExprConvertible (an Object), - # and reflection registered exactly buffer/region, with no type or span. - # Construct that historical schema directly, independently of TensorRegion - # serialization, while using current buffer/range schemas for dependencies. - legacy_index = len(nodes) - for _ in range(2): - nodes.append( - { - "type": "tirx.BufferRegion", - "data": {"buffer": source_index, "region": region_index}, - } - ) - graph["root_index"] = len(nodes) - nodes.append({"type": "ffi.Array", "data": [legacy_index, legacy_index, legacy_index + 1]}) - legacy_json = json.dumps(graph) - upgraded = json.loads(upgrade_json(legacy_json)) - assert upgraded["root_index"] == graph["root_index"] - assert len(upgraded["nodes"]) == len(nodes) + 1 - assert upgraded["nodes"][:legacy_index] == nodes[:legacy_index] - assert upgraded["nodes"][graph["root_index"]] == nodes[graph["root_index"]] - for index in (legacy_index, legacy_index + 1): - assert upgraded["nodes"][index] == { - "type": "ir.TensorRegion", - "data": { - "source": source_index, - "region": region_index, - "ty": len(nodes), - "span": 0, - }, - } - first, repeated, second = tvm.ir.load_json(legacy_json) - assert first.same_as(repeated) - assert not first.same_as(second) - assert first.source.same_as(second.source) - assert first.region.same_as(second.region) - assert first.ty.same_as(second.ty) - assert isinstance(first.ty, tirx.BufferRegionType) - assert first.span is None - expected = tvm.ir.TensorRegion(source, region, tirx.BufferRegionType()) - tvm.ir.assert_structural_equal(first, expected, map_free_vars=True) - - -@pytest.mark.parametrize("data", [None, {"region": 0}]) -def test_malformed_legacy_buffer_region(data): - node = {"type": "tirx.BufferRegion"} - if data is not None: - node["data"] = data - with pytest.raises(ValueError, match="requires a buffer field"): - upgrade_json(json.dumps({"nodes": [{"type": "None"}, node], "root_index": 1})) - - -@pytest.mark.parametrize("field", ["reads", "writes", "match_source"]) -@pytest.mark.parametrize("invalid", ["source", "rank"]) -def test_region_requires_buffer_source_and_rank(field, invalid): - source = ( - tirx.Var("source", "int32") if invalid == "source" else tirx.decl_buffer((4, 4), "float32") - ) - region = tvm.ir.TensorRegion(source, [tvm.ir.Range(0, 4)], tirx.BufferRegionType()) - message = None if invalid == "source" else "must match its buffer rank" - with pytest.raises((TypeError, tvm.error.InternalError), match=message): - if field == "match_source": - s_tir.MatchBufferRegion(tirx.decl_buffer((4,), "float32"), region) - else: - s_tir.SBlock( - [], - [region] if field == "reads" else [], - [region] if field == "writes" else [], - "invalid", - tirx.Evaluate(0), - ) - - -if __name__ == "__main__": - tvm.testing.main() diff --git a/tests/python/s_tir/test_tensor_intrin.py b/tests/python/s_tir/test_tensor_intrin.py deleted file mode 100644 index 29b9824db849..000000000000 --- a/tests/python/s_tir/test_tensor_intrin.py +++ /dev/null @@ -1,71 +0,0 @@ -# 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 intrinsic ownership and runtime compatibility.""" - -import json - -import pytest - -import tvm -import tvm.testing -from tvm import s_tir, tirx - - -def test_tensor_intrin_serialization(): - func = tirx.PrimFunc([], tirx.Evaluate(0)) - intrin = s_tir.TensorIntrin(func, func) - graph = json.loads(tvm.ir.save_json(intrin)) - assert any(node.get("type") == "s_tir.TensorIntrin" for node in graph["nodes"]) - for legacy in (False, True): - if legacy: - for node in graph["nodes"]: - if node.get("type") == "s_tir.TensorIntrin": - node["type"] = "tirx.TensorIntrin" - restored = tvm.ir.load_json(json.dumps(graph)) - assert isinstance(restored, s_tir.TensorIntrin) - assert restored.desc.same_as(restored.impl) - tvm.ir.assert_structural_equal(restored.desc, func) - - -def test_tensor_intrin_registration(): - func = tirx.PrimFunc([], tirx.Evaluate(0)) - name = "test_s_tir_tensor_intrin_registration" - s_tir.TensorIntrin.register(name, func, func, override=True) - assert s_tir.TensorIntrin.get(name).desc.same_as(func) - with pytest.raises(ValueError, match="already been registered"): - s_tir.TensorIntrin.register(name, func, func) - replacement = tirx.PrimFunc([], tirx.Evaluate(1)) - s_tir.TensorIntrin.register(name, func, replacement, override=True) - assert s_tir.TensorIntrin.get(name).impl.same_as(replacement) - - -def test_tensor_intrin_constructor_constraints(): - empty = tirx.PrimFunc([], tirx.Evaluate(0)) - scalar = tirx.PrimFunc([tirx.Var("x", "int32")], tirx.Evaluate(0)) - with pytest.raises(ValueError, match="number of parameters"): - s_tir.TensorIntrin(empty, scalar) - with pytest.raises(ValueError, match="description.*handle only"): - s_tir.TensorIntrin(scalar, scalar) - pointer = tirx.PrimFunc( - [tirx.Var("p", tvm.ir.PointerType(tvm.ir.PrimType("float32")))], tirx.Evaluate(0) - ) - with pytest.raises(ValueError, match="implementation.*handle only"): - s_tir.TensorIntrin(pointer, scalar) - - -if __name__ == "__main__": - tvm.testing.main() diff --git a/tests/python/s_tir/transform/test_s_tir_transform_convert_ssa.py b/tests/python/s_tir/transform/test_s_tir_transform_convert_ssa.py deleted file mode 100644 index d881981f0918..000000000000 --- a/tests/python/s_tir/transform/test_s_tir_transform_convert_ssa.py +++ /dev/null @@ -1,63 +0,0 @@ -# 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. - -import tvm -import tvm.testing -from tvm import s_tir, tirx - - -def test_reused_block_iterator(): - """A shared block defines a fresh iterator at each realization.""" - var = tirx.Var("v", "int32") - iterator = tirx.IterVar(tvm.ir.Range(0, 4), var, tirx.IterVar.DataPar) - block = s_tir.SBlock([iterator], [], [], "block", tirx.Evaluate(var)) - realize = s_tir.SBlockRealize([0], True, block) - before = tirx.PrimFunc([], tirx.SeqStmt([realize, realize])) - - after = s_tir.transform.ConvertSSA()(tvm.IRModule.from_expr(before))["main"] - - first, second = [realize.block for realize in after.body.seq] - assert not first.iter_vars[0].var.same_as(second.iter_vars[0].var) - assert first.body.value.same_as(first.iter_vars[0].var) - assert second.body.value.same_as(second.iter_vars[0].var) - # Shared input ownership must protect both original occurrences. - assert block.iter_vars[0].var.same_as(var) - assert block.body.value.same_as(var) - - -def test_shared_buffer_parameter_regions_across_functions(): - """Parameter renaming reaches both block regions and buffer accesses.""" - n = tirx.Var("n", "int32") - buffer = tirx.decl_buffer((n,), "float32", "buffer") - region = tirx.BufferRegion(buffer, [tvm.ir.Range(0, n)]) - block = s_tir.SBlock([], [region], [], "root", tirx.Evaluate(tirx.BufferLoad(buffer, [0]))) - func = tirx.PrimFunc([buffer], s_tir.SBlockRealize([], True, block)) - before = tvm.IRModule({"first": func, "second": func}) - - after = s_tir.transform.ConvertSSA()(before) - - first, second = after["first"], after["second"] - assert not first.params[0].same_as(second.params[0]) - for updated in [first, second]: - updated_block = updated.body.block - assert updated_block.reads[0].source.same_as(updated.params[0]) - assert updated_block.body.value.source.same_as(updated.params[0]) - assert updated_block.reads[0].region[0].extent.same_as(updated.params[0].ty.shape[0]) - - -if __name__ == "__main__": - tvm.testing.main() diff --git a/tests/python/s_tir/transform/test_s_tir_transform_simplify.py b/tests/python/s_tir/transform/test_s_tir_transform_simplify.py deleted file mode 100644 index e6f979840d54..000000000000 --- a/tests/python/s_tir/transform/test_s_tir_transform_simplify.py +++ /dev/null @@ -1,45 +0,0 @@ -# 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. - -import tvm -import tvm.testing -from tvm.script import tirx as T - - -def test_s_tir_block_iterator_constraints(): - @T.prim_func(private=True, s_tir=True) - def before(A: T.Buffer((4,), "int32")): - for i in range(4): - with T.sblock("write"): - vi = T.axis.spatial(4, i) - if vi < 4: - A[vi] = vi - - @T.prim_func(private=True, s_tir=True) - def expected(A: T.Buffer((4,), "int32")): - for i in range(4): - with T.sblock("write"): - vi = T.axis.spatial(4, i) - A[vi] = vi - - result = tvm.s_tir.transform.StmtSimplify()(tvm.IRModule.from_expr(before))["main"] - tvm.ir.assert_structural_equal(result, expected) - assert tvm.s_tir.analysis.verify_well_formed(result) - - -if __name__ == "__main__": - tvm.testing.main() diff --git a/tests/python/tirx-analysis/test_tir_analysis_verify_well_formed.py b/tests/python/tirx-analysis/test_tir_analysis_verify_well_formed.py index ce7fd285992e..5c4acdae1d2c 100644 --- a/tests/python/tirx-analysis/test_tir_analysis_verify_well_formed.py +++ b/tests/python/tirx-analysis/test_tir_analysis_verify_well_formed.py @@ -27,6 +27,53 @@ from tvm.script import tirx as T +def test_pass_simple(): + @T.prim_func(s_tir=True) + def element_wise( + A: T.Buffer((128, 128), "float32"), + C: T.Buffer((128, 128), "float32"), + ): + B = T.sblock_alloc_buffer((128, 128), "float32") + for i, j in T.grid(128, 128): + with T.sblock("B"): + vi, vj = T.axis.remap("SS", [i, j]) + B[vi, vj] = A[vi, vj] * 2.0 + for i, j in T.grid(128, 128): + with T.sblock("C"): + # It's a opaque block , so it can use outside variables + C[i, j] = B[i, j] * 2.0 + + assert tvm.s_tir.analysis.verify_well_formed(element_wise) + assert tvm.s_tir.analysis.verify_well_formed(tvm.IRModule.from_expr(element_wise)) + + +def test_buffer_region_bounds_are_visited(): + data = tvm.tirx.Var( + "data", tvm.ir.PointerType(tvm.ir.PrimType("int32"), storage_scope="global") + ) + buffer = tvm.tirx.decl_buffer([4], "int32", data=data) + undefined = tvm.tirx.Var("undefined", "int32") + region = tvm.tirx.BufferRegion(buffer, [tvm.ir.Range.from_min_extent(undefined, 4)]) + block = tvm.s_tir.SBlock([], [region], [], "region", tvm.tirx.Evaluate(0)) + func = tvm.tirx.PrimFunc([buffer], block) + assert not tvm.s_tir.analysis.verify_well_formed(func, assert_mode=False) + + +def test_fail_use_out_loop_var(): + @T.prim_func(check_well_formed=False, s_tir=True) + def element_wise( + A: T.Buffer((128, 128), "float32"), + B: T.Buffer((128, 128), "float32"), + ): + for i, j in T.grid(128, 128): + with T.sblock("B"): + vi, vj = T.axis.remap("SS", [i, j]) + # we cannot use `i` since it's defined outside the block + B[vi, vj] = A[i, vj] * 2.0 + + assert not tvm.s_tir.analysis.verify_well_formed(element_wise, assert_mode=False) + + def test_error_for_out_of_scope_usage(): """A variable may not be used after its scope ends. @@ -213,6 +260,48 @@ def func(A_handle: T.handle, B_handle: T.handle): tvm.tirx.analysis.verify_well_formed(mod) +def test_block_match_buffer_defines_buffer_obj(): + """In a block, T.match_buffer defines a buffer view""" + + @I.ir_module(s_tir=True) + class mod: + @T.prim_func(s_tir=True) + def func(A: T.Buffer([256, 256], "float32")): + for iters in T.grid(16, 16, 16, 16): + with T.sblock("compute"): + tile_i, tile_j, i, j = T.axis.remap("SSSS", iters) + B = T.match_buffer( + A[tile_i * 16 : (tile_i + 1) * 16, tile_j * 16 : (tile_j + 1) * 16], + dtype="float32", + ) + B[i, j] = 0.0 + + tvm.s_tir.analysis.verify_well_formed(mod) + + +def test_block_match_buffer_defines_symbolic_variables(): + """In a block, T.match_buffer may define symbolic variables""" + + @I.ir_module(s_tir=True) + class mod: + @T.prim_func(s_tir=True) + def func(A: T.Buffer([256, 256], "int32")): + for iters in T.grid(16, 16, 16, 16): + with T.sblock("compute"): + tile_i, tile_j, i, j = T.axis.remap("SSSS", iters) + + elem_offset = T.int32() + B = T.match_buffer( + A[tile_i * 16 : (tile_i + 1) * 16, tile_j * 16 : (tile_j + 1) * 16], + dtype="float32", + elem_offset=elem_offset, + ) + + B[i, j] = elem_offset + + tvm.s_tir.analysis.verify_well_formed(mod) + + def test_error_message_without_previous_definition_location(): """Test case 1: Error message without 'It was first defined at' @@ -325,6 +414,84 @@ def func(A: T.Buffer((128,), "float32")): tvm.tirx.analysis.verify_well_formed(func) +def test_alloc_buffer_in_block_is_well_formed(): + """SBlock::alloc_buffers introduces a buffer into scope for the block body.""" + + @I.ir_module(s_tir=True) + class mod: + @T.prim_func(s_tir=True) + def func(A: T.Buffer((128,), "float32")): + with T.sblock("root"): + B = T.sblock_alloc_buffer([128], "float32") + for i in T.grid(128): + with T.sblock("write_B"): + vi = T.axis.remap("S", [i]) + B[vi] = A[vi] * 2.0 + + tvm.s_tir.analysis.verify_well_formed(mod) + + +def test_match_buffer_in_block_is_well_formed(): + """SBlock::match_buffers introduces a buffer into scope for the block body.""" + + @I.ir_module(s_tir=True) + class mod: + @T.prim_func(s_tir=True) + def func(A: T.Buffer((128, 128), "float32")): + for iters in T.grid(8, 8, 16, 16): + with T.sblock("compute"): + ti, tj, i, j = T.axis.remap("SSSS", iters) + A_tile = T.match_buffer( + A[ti * 16 : (ti + 1) * 16, tj * 16 : (tj + 1) * 16], + dtype="float32", + ) + A_tile[i, j] = A_tile[i, j] * 2.0 + + tvm.s_tir.analysis.verify_well_formed(mod) + + +def test_error_undeclared_buffer_in_schedulable_tir(): + """In schedule-level TIR (with SBlock nodes), all buffers must be declared.""" + # Manually construct a BufferStore that uses a buffer without any declaration + # inside a block context. + n = tvm.tirx.Var("n", "int32") + A = tvm.tirx.decl_buffer([n], "float32", name="A") + i = tvm.tirx.Var("i", "int32") + + # Create an undeclared buffer using an explicit data pointer that is NOT + # a function parameter and NOT wrapped with DeclBuffer. + B_data = tvm.tirx.Var("B_data", tvm.ir.PointerType(tvm.ir.PrimType("float32"))) + B = tvm.tirx.decl_buffer([n], "float32", name="B", data=B_data) + + # Build a block that writes to B without any declaration of B. + bi = tvm.tirx.Var("bi", "int32") + block = tvm.s_tir.SBlock( + iter_vars=[tvm.tirx.IterVar(tvm.ir.Range(0, n), bi, 0)], # 0 = kDataPar + reads=[tvm.tirx.BufferRegion(A, [tvm.ir.Range(bi, bi + 1)])], + writes=[tvm.tirx.BufferRegion(B, [tvm.ir.Range(bi, bi + 1)])], + body=tvm.tirx.BufferStore(B, tvm.tirx.BufferLoad(A, [bi]), [bi]), + name_hint="write_B", + ) + block_realize = tvm.s_tir.SBlockRealize( + iter_values=[i], + predicate=tvm.tirx.const(True), + block=block, + ) + + prim_func = tvm.tirx.PrimFunc( + params=[A, B_data], + body=tvm.tirx.For(i, 0, n, tvm.tirx.ForKind.SERIAL, block_realize), + # Note: B is NOT a function parameter, so its declaration scope is only + # within a DeclBuffer node (which we intentionally omit here). + ) + + # B is used in the block but was never declared — should fail. + with pytest.raises( + (ValueError, tvm.error.InternalError), match="buffer B.*without a prior DeclBuffer" + ): + tvm.s_tir.analysis.verify_well_formed(prim_func) + + def test_tensor_load_asserted_type_matches_source_and_indices(): @T.prim_func def func(): @@ -377,42 +544,5 @@ def test_tensor_load_malformed_indices_return_false_without_asserting(): tvm.tirx.analysis.verify_well_formed(non_final_vector) -@pytest.mark.parametrize("as_module", [False, True]) -def test_core_verifiers_reject_blocks(as_module): - block = tvm.s_tir.SBlock([], [], [], "block", tvm.tirx.Evaluate(0)) - func = tvm.tirx.PrimFunc([], block).with_attr("s_tir", True) - obj = tvm.IRModule.from_expr(func) if as_module else func - assert tvm.s_tir.analysis.verify_well_formed(obj) - for verify in [ - tvm.tirx.analysis.verify_well_formed, - tvm.tirx.analysis.verify_tirx_well_formed, - ]: - assert not verify(obj, assert_mode=False) - with pytest.raises(tvm.error.InternalError, match="(does not support|not allowed)"): - verify(obj) - - -@pytest.mark.parametrize("legacy_s_tir", [False, True]) -def test_mixed_module_parser_checks_both_dialects(legacy_s_tir): - @I.ir_module - class Mixed: - @T.prim_func(s_tir=True) - def scheduled(A: T.Buffer((4,), "int32")): - for i in range(4): - with T.sblock("write"): - vi = T.axis.spatial(4, i) - A[vi] = vi - - @T.prim_func - def lowered(): - T.evaluate(0) - - restored = tvm.script.from_source(Mixed.script(), s_tir=legacy_s_tir) - tvm.ir.assert_structural_equal(restored, Mixed) - assert tvm.s_tir.analysis.verify_well_formed(restored) - assert tvm.tirx.analysis.verify_tirx_well_formed(Mixed["lowered"]) - assert not tvm.tirx.analysis.verify_tirx_well_formed(Mixed, assert_mode=False) - - if __name__ == "__main__": tvm.testing.main() diff --git a/tests/python/tirx-base/test_tir_specialize.py b/tests/python/tirx-base/test_tir_specialize.py index 4be5ed121dc6..4dffc8dc11e9 100644 --- a/tests/python/tirx-base/test_tir_specialize.py +++ b/tests/python/tirx-base/test_tir_specialize.py @@ -368,49 +368,5 @@ def expected() -> T.int32: tvm.ir.assert_structural_equal(after.ty, ty_expected) -def test_specialize_structural_buffer_definitions(): - """Extension definitions remain consistent across metadata, regions, and uses.""" - n = tvm.tirx.Var("n", "int32") - allocated = tvm.tirx.decl_buffer((n,), "float32", name="allocated") - matched = tvm.tirx.decl_buffer((n,), "float32", name="matched") - source = tvm.tirx.BufferRegion(allocated, [tvm.ir.Range(n)]) - read = tvm.tirx.BufferRegion(matched, [tvm.ir.Range(n)]) - match = tvm.s_tir.MatchBufferRegion(matched, source) - block = tvm.s_tir.SBlock( - [], - [read], - [], - "use", - tvm.tirx.Evaluate(matched[n - 1]), - alloc_buffers=[allocated], - match_buffers=[match], - annotations={"extent": n}, - ) - before = tvm.tirx.PrimFunc([n], tvm.s_tir.SBlockRealize([], True, block)) - assert tvm.s_tir.analysis.verify_well_formed(before) - after = before.specialize({n: 8}) - assert tvm.s_tir.analysis.verify_well_formed(after) - result = after.body.block - assert not after.params - assert result.alloc_buffers[0].shape[0] == 8 - assert result.match_buffers[0].buffer.shape[0] == 8 - assert result.match_buffers[0].source.source.same_as(result.alloc_buffers[0]) - assert result.reads[0].source.same_as(result.match_buffers[0].buffer) - assert result.body.value.source.same_as(result.match_buffers[0].buffer) - assert result.body.value.indices[0] == 7 - assert result.annotations["extent"] == 8 - # Specialization must not rewrite another owner's unspecialized definition. - assert block.alloc_buffers[0].shape[0].same_as(n) - assert block.match_buffers[0].buffer.shape[0].same_as(n) - assert block.annotations["extent"].same_as(n) - - -def test_specialize_plain_tirx(): - n = tvm.tirx.Var("n", "int32") - before = tvm.tirx.PrimFunc([n], tvm.tirx.Evaluate(n + 1)) - expected = tvm.tirx.PrimFunc([], tvm.tirx.Evaluate(9)) - tvm.ir.assert_structural_equal(before.specialize({n: 8}), expected) - - if __name__ == "__main__": tvm.testing.main() diff --git a/tests/python/tirx-transform/test_tir_transform_flatten_buffer.py b/tests/python/tirx-transform/test_tir_transform_flatten_buffer.py index 75900822bca0..7d4688b747d0 100644 --- a/tests/python/tirx-transform/test_tir_transform_flatten_buffer.py +++ b/tests/python/tirx-transform/test_tir_transform_flatten_buffer.py @@ -352,7 +352,7 @@ def main(input_A: T.Buffer(10, "bool"), input_B: T.Buffer(10, "bool")) -> None: tvm.ir.assert_structural_equal(After, _lower_blocks(Expected)) -def test_flatten_lowered_block(): +def test_flatten_inside_block(): """Flatten allocations and accesses after lowering a schedulable block.""" @I.ir_module(s_tir=True) diff --git a/tests/python/tirx-transform/test_tir_transform_force_narrow_index_to_i32.py b/tests/python/tirx-transform/test_tir_transform_force_narrow_index_to_i32.py index 9ce2bd83f7c0..bb99a483635f 100644 --- a/tests/python/tirx-transform/test_tir_transform_force_narrow_index_to_i32.py +++ b/tests/python/tirx-transform/test_tir_transform_force_narrow_index_to_i32.py @@ -466,32 +466,5 @@ def main(buf: T.handle): tvm.ir.assert_structural_equal(_lower_blocks(Expected), after) -def test_preserve_local_scalar_storage(): - @T.prim_func(private=True) - def before(n: T.int64): - index = T.call_extern("opaque_index", n, dtype="int64") - T.evaluate(index) - - @T.prim_func(private=True) - def expected(n: T.int32): - index = T.call_extern("opaque_index", n, dtype="int64") - T.evaluate(index) - - after = tvm.tirx.transform.ForceNarrowIndexToInt32()(tvm.IRModule.from_expr(before))["main"] - tvm.ir.assert_structural_equal(after, expected) - - -@pytest.mark.parametrize("shape", [(), (1,), (8,)]) -@pytest.mark.parametrize("scope", ["local", "shared", "global"]) -def test_reject_int64_array_storage(shape, scope): - # Rank or element count alone does not make an allocation scalar storage. - buffer = tvm.tirx.decl_buffer(shape, "int64", "array", scope=scope, layout=None) - before = tvm.tirx.PrimFunc( - [], tvm.tirx.SeqStmt([tvm.tirx.AllocBuffer(buffer), tvm.tirx.Evaluate(0)]) - ) - with pytest.raises(tvm.error.InternalError, match="allocated in the function has dtype"): - tvm.tirx.transform.ForceNarrowIndexToInt32()(tvm.IRModule.from_expr(before)) - - if __name__ == "__main__": tvm.testing.main()