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..acaf2a4ba676 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 @@ -48,7 +49,7 @@ namespace tirx { * - third: opaque regions */ TVM_DLL ffi::Array> GetSBlockAccessRegion( - const SBlock& block, const ffi::Map& buffer_var_map); + 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 @@ -59,7 +60,7 @@ TVM_DLL ffi::Array> GetSBlockAccessRegion( * \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); + 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 @@ -97,6 +98,11 @@ class Analyzer; namespace s_tir { using namespace tvm::tirx; +/*! \brief Verify variable/buffer definitions, load types and schedulable block boundaries. */ +TVM_DLL bool VerifyWellFormed(const tirx::PrimFunc& func, bool assert_mode = true); +/*! \brief Verify S-TIR or mixed modules, including definitions shared across functions. */ +TVM_DLL bool VerifyWellFormed(const IRModule& mod, bool assert_mode = true); + /*! * \brief Estimate the FLOPs of a TIR fragment. * \param stmt The TIR fragment to be estimated. diff --git a/include/tvm/s_tir/sblock_scope.h b/include/tvm/s_tir/sblock_scope.h index 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..eefe47d8602f 100644 --- a/include/tvm/s_tir/stmt.h +++ b/include/tvm/s_tir/stmt.h @@ -18,7 +18,7 @@ */ /*! * \file tvm/s_tir/stmt.h - * \brief S-TIR (Schedulable TIR) statement attribute declarations. + * \brief S-TIR (Schedulable TIR) statements and attributes. * * This file contains attribute keys that are specific to the schedulable TIR * (S-TIR) layer, including meta_schedule annotations and schedule primitive / @@ -27,18 +27,185 @@ #ifndef TVM_S_TIR_STMT_H_ #define TVM_S_TIR_STMT_H_ -#include +#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. */ + 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("s_tir.MatchBufferRegion", MatchBufferRegionNode, ffi::Object); +}; + +/*! + * \brief Managed reference to MatchBufferRegionNode. + * \sa MatchBufferRegionNode + */ +class MatchBufferRegion : public ffi::ObjectRef { + public: + TVM_DLL explicit MatchBufferRegion(tirx::BufferVar buffer, TensorRegion source); + + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(MatchBufferRegion, ffi::ObjectRef, + MatchBufferRegionNode); + TVM_DEFINE_OBJECT_REF_COW_METHOD(MatchBufferRegionNode); +}; + +/*! + * \brief A block is a basic schedule unit in TIR. + * \note SBlock's body is parameterized by iter vars. + * \code + * + * with T.sblock(name): + * v0 = T.axis.S(domain, value0) + * v1 = T.axis.R(domain, value1) + * ... + * T.reads([buffer0[start:end, ...], ...]) + * T.writes([buffer1[start:end, ...], ...]) + * T.where(predicate) + * buffer2 = T.alloc_buffer(shape, dtype) + * buffer3 = T.match_buffer(source_buffer[start:end, ...]) + * T.attr({attr_key: attr_value, ...}) + * with T.init(): + * // init body + * // body + * + * \endcode + */ +class SBlockNode : public tirx::StmtNode { + public: + /*! \brief The variables of the block. */ + ffi::Array 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 { /*! * \brief Annotations for invoking and synchronizing asynchronous operations. */ constexpr const char* async_commit_queue_scope = "async_commit_queue_scope"; -constexpr const char* async_wait_queue_scope = "async_wait_queue_scope"; -constexpr const char* async_wait_inflight_count = "async_wait_inflight_count"; +constexpr const char* async_wait_queue_scope = tirx::attr::async_wait_queue_scope; +constexpr const char* async_wait_inflight_count = tirx::attr::async_wait_inflight_count; /*! * \brief Mark that the attached statement runs asynchronously. @@ -77,7 +244,7 @@ constexpr const char* pragma_loop_partition_hint = "pragma_loop_partition_hint"; constexpr const char* reduce_scope = "reduce_scope"; /*! \brief Mark launching of a virtual thread. */ -constexpr const char* virtual_thread = "virtual_thread"; +constexpr const char* virtual_thread = tirx::attr::virtual_thread; // ----------------------------------------------------------------------- // meta_schedule annotations diff --git a/include/tvm/s_tir/stmt_functor.h b/include/tvm/s_tir/stmt_functor.h new file mode 100644 index 000000000000..8d4734c07c5f --- /dev/null +++ b/include/tvm/s_tir/stmt_functor.h @@ -0,0 +1,128 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +/*! + * \file tvm/s_tir/stmt_functor.h + * \brief Native traversal and statement dispatch for schedulable TIR. + */ +#ifndef TVM_S_TIR_STMT_FUNCTOR_H_ +#define TVM_S_TIR_STMT_FUNCTOR_H_ + +#include +#include + +namespace tvm { +namespace s_tir { + +/*! + * \brief Extend TIRX statement dispatch with schedulable blocks. + * \tparam FType The statement signature, using the native Dispatch 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::Dispatch_; + + virtual R Dispatch_(const SBlockNode* op, Args... args) { + return this->DispatchDefault_(op, std::forward(args)...); + } + virtual R Dispatch_(const SBlockRealizeNode* op, Args... args) { + return this->DispatchDefault_(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. + * Generic TIRX visitors traverse these nodes structurally, including whole + * iterators and annotations. Both paths visit allocation and match-buffer + * definitions before region uses. Only this S-TIR subclass supplies native + * block hooks; the core TIRX table does not register dialect nodes. + */ +class TVM_DLL StmtExprVisitor : public tirx::StmtExprVisitor { + public: + TVM_DEFINE_OBJECT_FUNCTOR_DEFAULT_CONSTRUCTOR(StmtExprVisitor, tirx::StmtExprVisitor) + using tirx::StmtExprVisitor::Visit; + using tirx::StmtExprVisitor::Visit_; + + virtual ffi::Optional 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. + * Generic TIRX mutators instead use the full structural rewrite, including + * iterator definitions and annotations. Both paths establish allocation and + * match-buffer remaps before visiting region uses and preserve copy-on-write. + */ +class TVM_DLL StmtExprMutator : public tirx::StmtExprMutator { + public: + TVM_DEFINE_OBJECT_FUNCTOR_DEFAULT_CONSTRUCTOR(StmtExprMutator, tirx::StmtExprMutator) + using tirx::StmtExprMutator::Mutate; + using tirx::StmtExprMutator::Mutate_; + + virtual UnchangedOr 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/tensor_intrin.h b/include/tvm/s_tir/tensor_intrin.h new file mode 100644 index 000000000000..cff27d56a1a9 --- /dev/null +++ b/include/tvm/s_tir/tensor_intrin.h @@ -0,0 +1,90 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +/*! + * \file tvm/s_tir/tensor_intrin.h + * \brief Tensor intrinsics for schedulable TIR. + */ +#ifndef TVM_S_TIR_TENSOR_INTRIN_H_ +#define TVM_S_TIR_TENSOR_INTRIN_H_ + +#include + +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_TENSOR_INTRIN_H_ diff --git a/include/tvm/s_tir/transform.h b/include/tvm/s_tir/transform.h index 6bbc91edfbf3..0a92264c4c6f 100644 --- a/include/tvm/s_tir/transform.h +++ b/include/tvm/s_tir/transform.h @@ -26,6 +26,7 @@ #include #include +#include #include #include @@ -49,6 +50,13 @@ using tirx::transform::CreatePrimFuncPass; using tvm::transform::Pass; using tvm::transform::PassContext; +/*! \brief De-duplicate definitions, including schedulable block iterators, across PrimFuncs. */ +TVM_DLL Pass ConvertSSA(); + +/*! \brief Simplify schedulable TIR using block iteration constraints and shared simplifier options. + */ +TVM_DLL Pass StmtSimplify(); + /*! * \brief Canonicalize loop to start from zero . * \return The pass. diff --git a/include/tvm/s_tir/utils.h b/include/tvm/s_tir/utils.h index 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..11fa65f0b0cf 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,12 +130,7 @@ TVM_DLL size_t CalculateWorkspaceBytes(const PrimFunc& func, int64_t workspace_b * * - Each variable has a single point of definition. * - * - Expressions within a tirx::SBlock may not reference variables - * defined outside the block. For example, for a block with iter - * vars `vi, vj = T.axis.remap('SS', [i,j])`, the statement - * `B[i,j] = A[i,j]` would be ill-formed, because it uses the loop - * variables `i` and `j` instead of the block variables `vi` and - * `vj`. + * Dialect statements require their dialect-specific verifier. * * \param func The PrimFunc to be verified. * \param assert_mode The indicator if it raises an error when the function is not well-formed. diff --git a/include/tvm/tirx/function.h b/include/tvm/tirx/function.h index 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/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..5c3f4f2cd99d 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. @@ -995,6 +829,11 @@ constexpr const char* pragma_unroll_explicit = "pragma_unroll_explicit"; constexpr const char* storage_alignment = "storage_alignment"; /*! \brief Mark launching extent of thread, used by device API. */ constexpr const char* thread_extent = "thread_extent"; + +/*! \brief Shared execution attributes consumed before and after block lowering. */ +constexpr const char* virtual_thread = "virtual_thread"; +constexpr const char* async_wait_queue_scope = "async_wait_queue_scope"; +constexpr const char* async_wait_inflight_count = "async_wait_inflight_count"; /*! \brief Annotation key on AllocBuffer marking the allocation as volatile. */ constexpr const char* kVolatile = "tirx.volatile"; /*! \brief Mark buffer initial addr alignment in bytes */ diff --git a/include/tvm/tirx/stmt_functor.h b/include/tvm/tirx/stmt_functor.h index c1147f5f8010..17c93ece7263 100644 --- a/include/tvm/tirx/stmt_functor.h +++ b/include/tvm/tirx/stmt_functor.h @@ -116,12 +116,6 @@ class StmtFunctor { 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)...); } @@ -141,6 +135,10 @@ class StmtFunctor { explicit StmtFunctor(const VTable* vtable) : vtable_(vtable) {} /*! \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)...); + }); SetDispatch(vtable); SetDispatch(vtable); SetDispatch(vtable); @@ -155,8 +153,6 @@ class StmtFunctor { SetDispatch(vtable); SetDispatch(vtable); SetDispatch(vtable); - SetDispatch(vtable); - SetDispatch(vtable); SetDispatch(vtable); SetDispatch(vtable); } @@ -216,8 +212,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 +227,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); }; @@ -282,8 +276,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); @@ -297,7 +289,11 @@ class TVM_DLL StmtExprMutator : public tvm::ExprMutator { * \param stmt The source statement to be substituted * \param vmap returns a new value if re-mapping is needed, otherwise returns nullptr. * - * Substitution may change the data type of the expression. + * This statement overload legalizes only core TIRX nodes. Dtype-changing + * substitutions must be applied after lowering dialect blocks: structural + * traversal of extension statements does not legalize their iterator domains or + * bindings. Type-preserving mappings continue to traverse schedulable blocks + * structurally before lowering. * * \return The result. */ @@ -334,14 +330,13 @@ bool ContainsNode(const Stmt& stmt) { if (contains_node || value.as()) { 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/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): diff --git a/python/tvm/ir/json_compact.py b/python/tvm/ir/json_compact.py index a7986a1491e7..e26aab947a12 100644 --- a/python/tvm/ir/json_compact.py +++ b/python/tvm/ir/json_compact.py @@ -19,6 +19,11 @@ import json _PRIM_TYPE_KEY_RENAMES = { + "tirx.BufferRegion": "ir.TensorRegion", + "tirx.SBlock": "s_tir.SBlock", + "tirx.SBlockRealize": "s_tir.SBlockRealize", + "tirx.MatchBufferRegion": "s_tir.MatchBufferRegion", + "tirx.TensorIntrin": "s_tir.TensorIntrin", "tirx.StringImm": "ir.prim.StringImm", "tirx.Cast": "ir.prim.Cast", "tirx.Add": "ir.prim.Add", @@ -119,7 +124,23 @@ def upgrade_json(json_str): # compatible with the pre-unification Relax/TIRx schemas and with graphs # written before the canonical Var field was renamed to `name`. Rewriting # nodes in place preserves node indices and shared references. - for node in data.get("nodes", []): + nodes = data.get("nodes", []) + buffer_region_type = None + for node in nodes: + if node.get("type") == "tirx.BufferRegion": + fields = node.get("data") + if not isinstance(fields, dict) or "buffer" not in fields: + raise ValueError("Legacy tirx.BufferRegion requires a buffer field") + fields["source"] = fields.pop("buffer") + # Typed BufferRegion already carries type/span. Before it became + # an Expr, it had only buffer/region; supply that form's defaults + # by appending a type node so existing graph indices stay intact. + if "ty" not in fields: + if buffer_region_type is None: + buffer_region_type = len(nodes) + nodes.append({"type": "tirx.BufferRegionType", "data": {"span": 0}}) + fields["ty"] = buffer_region_type + fields.setdefault("span", 0) node["type"] = _PRIM_TYPE_KEY_RENAMES.get(node.get("type"), node.get("type")) if node.get("type") == "relax.expr.Var": node["type"] = "ir.Var" diff --git a/python/tvm/relax/analysis/analysis.py b/python/tvm/relax/analysis/analysis.py index 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 164dcc99019b..7bb80b3419d4 100644 --- a/python/tvm/s_tir/__init__.py +++ b/python/tvm/s_tir/__init__.py @@ -18,7 +18,8 @@ # pylint: disable=invalid-name """S-TIR namespace for scheduable TensorIR""" -from tvm.tirx.function import TensorIntrin +from ._tensor_intrin import TensorIntrin +from .stmt import MatchBufferRegion, SBlock, SBlockRealize # dlight depends on compiler-only C++ functions (e.g. s_tir.schedule.GetSBlockRealize), # so skip it in runtime-only builds. diff --git a/python/tvm/s_tir/_tensor_intrin.py b/python/tvm/s_tir/_tensor_intrin.py new file mode 100644 index 000000000000..00967c08eea6 --- /dev/null +++ b/python/tvm/s_tir/_tensor_intrin.py @@ -0,0 +1,80 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Tensor intrinsics for schedulable TIR.""" + +from typing import Optional + +import tvm_ffi + +from tvm.runtime import Object +from tvm.tirx.function import PrimFunc + +from . import _ffi_api + + +@tvm_ffi.register_object("s_tir.TensorIntrin") +class TensorIntrin(Object): + """A tensor intrinsic. + + Parameters + ---------- + desc : PrimFunc + The function to describe the computation. + + impl : PrimFunc + The function of the implementation for the execution. + """ + + def __init__(self, desc, impl): + self.__init_handle_by_constructor__(_ffi_api.TensorIntrin, desc, impl) + + @staticmethod + def register(name: str, desc: PrimFunc, impl: PrimFunc, override: bool = False): + """Register a tensor intrinsic with its name. + + Parameters + ---------- + name : str + The name of the TensorIntrin to register. + desc : PrimFunc + The function to describe the computation. + impl : PrimFunc + The function of the implementation for the execution. + override: bool + Whether override existing intrinsic. + """ + return _ffi_api.TensorIntrinRegister(name, TensorIntrin(desc, impl), override) # type: ignore + + @staticmethod + def get(name: str, allow_missing: bool = False) -> Optional["TensorIntrin"]: + """Look up a tensor intrinsic by its name. + + Parameters + ---------- + name : str + The name of the TensorIntrin to look up. + + allow_missing : bool + Whether to allow missing tensor intrin. If False, raise an error if the tensor intrin + doesn't exist. + + Returns + ------- + result : Optional[TensorIntrin] + The TensorIntrin with the specified name, or None if not found. + """ + return _ffi_api.TensorIntrinGet(name, allow_missing) # pylint: type: ignore diff --git a/python/tvm/s_tir/analysis/__init__.py b/python/tvm/s_tir/analysis/__init__.py index 07021f092d95..24732261a5aa 100644 --- a/python/tvm/s_tir/analysis/__init__.py +++ b/python/tvm/s_tir/analysis/__init__.py @@ -23,7 +23,7 @@ import tvm from tvm.ir import IRModule, TensorRegion from tvm.tirx.expr import Var -from tvm.tirx.stmt import SBlock +from tvm.s_tir import SBlock from tvm.tirx import Buffer, Stmt from tvm.tirx.function import PrimFunc @@ -38,7 +38,7 @@ def get_sblock_access_region( Parameters ---------- - block: tvm.tirx.SBlock + block: tvm.s_tir.SBlock The block in which we are detecting read/write regions. buffer_var_map : Dict[Var, Buffer] @@ -63,7 +63,7 @@ def get_sblock_read_write_region( Parameters ---------- - block: tvm.tirx.SBlock + block: tvm.s_tir.SBlock The block in which we are detecting read/write regions. buffer_var_map : Dict[Var, Buffer] @@ -214,3 +214,13 @@ def is_pure_function(func: PrimFunc) -> bool: def assert_pure_function(func: PrimFunc) -> bool: """Asserts that the function is a pure function""" return _ffi_api.is_pure_function(func, True) # type: ignore # pylint: disable=no-member + + +def verify_well_formed(obj: PrimFunc | IRModule, assert_mode: bool = True) -> bool: + """Verify definitions, buffer loads and S-TIR block boundaries. + + Modules may contain both S-TIR and ordinary PrimFuncs. Shared variable + identities are checked across function boundaries. Use the TIRX-specific + verifier separately for execution-scope restrictions on ordinary PrimFuncs. + """ + return _ffi_api.VerifyWellFormed(obj, assert_mode) diff --git a/python/tvm/s_tir/backend/adreno/pipeline.py b/python/tvm/s_tir/backend/adreno/pipeline.py index 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/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/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/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 e97195ad54f4..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 @@ -3043,7 +3044,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/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..3c84f5a59c87 --- /dev/null +++ b/python/tvm/s_tir/stmt.py @@ -0,0 +1,179 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Schedulable TensorIR statement nodes.""" + +from collections.abc import Mapping + +import tvm_ffi + +from tvm.ir import Expr, Span, TensorRegion +from tvm.runtime import Object, Scriptable, const +from tvm.tirx.buffer import Buffer +from tvm.tirx.expr import IterVar +from tvm.tirx.stmt import Stmt, _normalize_legacy_stmt + +from . import _ffi_api + + +@tvm_ffi.register_object("s_tir.MatchBufferRegion") +class MatchBufferRegion(Object, Scriptable): + """MatchBufferRegion node. + + Parameters + ---------- + buffer : Buffer + The target buffer + + source : TensorRegion + The region of source buffer + """ + + buffer: Buffer + source: TensorRegion + + def __init__(self, buffer: Buffer, source: TensorRegion) -> None: + self.__init_handle_by_constructor__( + _ffi_api.MatchBufferRegion, + buffer, + source, # type: ignore + ) + + +@tvm_ffi.register_object("s_tir.SBlock") +class SBlock(Stmt): + """SBlock node. + + Parameters + ---------- + iter_vars : List[IterVar] + The block Variable. + + reads : List[TensorRegion] + The read buffer regions of the block. + + writes: List[TensorRegion] + The write buffer regions of the block. + + name_hint: str + the name_hint of the block. + + body: Stmt + The body of the block. + + init: Optional[Stmt] + The init block of the reduction block + + alloc_buffers: Optional[list[Buffer]] + The buffer allocations + + match_buffers: Optional[List[MatchBufferRegion]] + The subregion buffer match + + annotations: Optional[Mapping[str, Object]] + Additional annotation hints. + + span : Optional[Span] + The location of this block in the source code. + """ + + iter_vars: list[IterVar] + reads: list[TensorRegion] + writes: list[TensorRegion] + name_hint: str + body: Stmt + init: Stmt | None + alloc_buffers: list[Buffer] + match_buffers: list[MatchBufferRegion] + annotations: Mapping[str, Object] + span: Span | None + + def __init__( + self, + iter_vars: list[IterVar], + reads: list[TensorRegion], + writes: list[TensorRegion], + name_hint: str, + body: Stmt, + init: Stmt | None = None, + alloc_buffers: list[Buffer] | None = None, + match_buffers: list[MatchBufferRegion] | None = None, + annotations: Mapping[str, Object] | None = None, + span: Span | None = None, + ) -> None: + if alloc_buffers is None: + alloc_buffers = [] + if match_buffers is None: + match_buffers = [] + if annotations is None: + annotations = {} + body = _normalize_legacy_stmt(body) + init = _normalize_legacy_stmt(init) + self.__init_handle_by_constructor__( + _ffi_api.SBlock, # type: ignore + iter_vars, + reads, + writes, + name_hint, + body, + init, + alloc_buffers, + match_buffers, + annotations, + span, + ) # type: ignore + + +@tvm_ffi.register_object("s_tir.SBlockRealize") +class SBlockRealize(Stmt): + """SBlockRealize node. + + Parameters + ---------- + iter_values : List[Expr] + The binding values of the block var. + + predicate : Union[Expr, bool] + The predicate of the block. + + block : SBlock + The block to realize + + span : Optional[Span] + The location of this block_realize in the source code. + """ + + iter_values: list[Expr] + predicate: Expr + block: SBlock + span: Span | None + + def __init__( + self, + iter_values: list[Expr], + predicate: Expr | bool, + block: SBlock, + span: Span | None = None, + ) -> None: + if isinstance(predicate, bool): + predicate = const(predicate, "bool") + self.__init_handle_by_constructor__( + _ffi_api.SBlockRealize, # type: ignore + iter_values, + predicate, + block, + span, + ) # type: ignore diff --git a/python/tvm/s_tir/tensor_intrin/arm_cpu.py b/python/tvm/s_tir/tensor_intrin/arm_cpu.py index 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/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..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. @@ -142,10 +147,12 @@ def parse( parser.report_error(source_ast, err=WELL_FORMED_ERROR_MESSAGE) try: - if s_tir: - tvm.tirx.analysis.verify_well_formed(check_ret) - else: - tvm.tirx.analysis.verify_tirx_well_formed(check_ret) + # Keep shared-definition checks across both dialects, then apply + # ordinary TIRX execution restrictions only to its own functions. + tvm.s_tir.analysis.verify_well_formed(check_ret) + for func in check_ret.functions.values(): + if isinstance(func, tvm.tirx.PrimFunc) and not func.attrs.get("s_tir", False): + tvm.tirx.analysis.verify_tirx_well_formed(func) except Exception as err: # pylint: disable=broad-exception-caught parser.report_error( source_ast, diff --git a/python/tvm/script/parser/ir/entry.py b/python/tvm/script/parser/ir/entry.py index 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/__init__.py b/python/tvm/tirx/__init__.py index af4b7161b0bf..e5a08aed0cf8 100644 --- a/python/tvm/tirx/__init__.py +++ b/python/tvm/tirx/__init__.py @@ -52,11 +52,11 @@ from .stmt import SeqStmt from .stmt import IfThenElse, Evaluate, stmt_seq, stmt_list -from .stmt import BufferRegion, BufferRegionType, MatchBufferRegion, SBlock, SBlockRealize +from .stmt import BufferRegion, BufferRegionType from .stmt import ScopeIdDefStmt from .tile_primitive import DispatchContext, LambdaExpr, TilePrimitiveCall -from .function import PrimFunc, TensorIntrin, IndexMap +from .function import PrimFunc, IndexMap from .op import call_packed_lowered, call_cpacked_lowered, register_intrin_lowering from .op import call_packed, call_cpacked, call_intrin, call_pure_extern, call_extern diff --git a/python/tvm/tirx/analysis/analysis.py b/python/tvm/tirx/analysis/analysis.py index 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 38c701075fde..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 ------- @@ -164,60 +160,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/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/python/tvm/tirx/stmt.py b/python/tvm/tirx/stmt.py index 5d15b227af16..f609d4439fce 100644 --- a/python/tvm/tirx/stmt.py +++ b/python/tvm/tirx/stmt.py @@ -34,7 +34,7 @@ import tvm_ffi from tvm.ir import Expr, Range, Span, TensorRegion, Type -from tvm.runtime import Object, Scriptable, const +from tvm.runtime import Object, Scriptable from . import _ffi_api from .buffer import Buffer @@ -628,156 +628,6 @@ def BufferRegion(buffer: Buffer, region: list[Range]) -> TensorRegion: return _ffi_api.BufferRegion(buffer, region) -@tvm_ffi.register_object("tirx.MatchBufferRegion") -class MatchBufferRegion(Object, Scriptable): - """MatchBufferRegion node. - - Parameters - ---------- - buffer : Buffer - The target buffer - - source : TensorRegion - The region of source buffer - """ - - buffer: Buffer - source: TensorRegion - - def __init__(self, buffer: Buffer, source: TensorRegion) -> None: - self.__init_handle_by_constructor__( - _ffi_api.MatchBufferRegion, - buffer, - source, # type: ignore - ) - - -@tvm_ffi.register_object("tirx.SBlock") -class SBlock(Stmt): - """SBlock node. - - Parameters - ---------- - iter_vars : List[IterVar] - The block Variable. - - reads : List[TensorRegion] - The read buffer regions of the block. - - writes: List[TensorRegion] - The write buffer regions of the block. - - name_hint: str - the name_hint of the block. - - body: Stmt - The body of the block. - - init: Optional[Stmt] - The init block of the reduction block - - alloc_buffers: Optional[list[Buffer]] - The buffer allocations - - match_buffers: Optional[List[MatchBufferRegion]] - The subregion buffer match - - annotations: Optional[Mapping[str, Object]] - Additional annotation hints. - - span : Optional[Span] - The location of this block in the source code. - """ - - iter_vars: list[IterVar] - reads: list[TensorRegion] - writes: list[TensorRegion] - name_hint: str - body: Stmt - init: Stmt | None - alloc_buffers: list[Buffer] - match_buffers: list[MatchBufferRegion] - annotations: Mapping[str, Object] - span: Span | None - - def __init__( - self, - iter_vars: list[IterVar], - reads: list[TensorRegion], - writes: list[TensorRegion], - name_hint: str, - body: Stmt, - init: Stmt | None = None, - alloc_buffers: list[Buffer] | None = None, - match_buffers: list[MatchBufferRegion] | None = None, - annotations: Mapping[str, Object] | None = None, - span: Span | None = None, - ) -> None: - if alloc_buffers is None: - alloc_buffers = [] - if match_buffers is None: - match_buffers = [] - if annotations is None: - annotations = {} - body = _normalize_legacy_stmt(body) - init = _normalize_legacy_stmt(init) - self.__init_handle_by_constructor__( - _ffi_api.SBlock, # type: ignore - iter_vars, - reads, - writes, - name_hint, - body, - init, - alloc_buffers, - match_buffers, - annotations, - span, - ) # type: ignore - - -@tvm_ffi.register_object("tirx.SBlockRealize") -class SBlockRealize(Stmt): - """SBlockRealize node. - - Parameters - ---------- - iter_values : List[Expr] - The binding values of the block var. - - predicate : Union[Expr, bool] - The predicate of the block. - - block : SBlock - The block to realize - - span : Optional[Span] - The location of this block_realize in the source code. - """ - - iter_values: list[Expr] - predicate: Expr - block: SBlock - span: Span | None - - def __init__( - self, - iter_values: list[Expr], - predicate: Expr | bool, - block: SBlock, - span: Span | None = None, - ) -> None: - if isinstance(predicate, bool): - predicate = const(predicate, "bool") - self.__init_handle_by_constructor__( - _ffi_api.SBlockRealize, # type: ignore - iter_values, - predicate, - block, - span, - ) # type: ignore - - @tvm_ffi.register_object("tirx.ScopeIdDefStmt") class ScopeIdDefStmt(Stmt): """ScopeIdDefStmt node. diff --git a/src/relax/analysis/layout_transformation.cc b/src/relax/analysis/layout_transformation.cc index 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..f9f428c33ab7 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,7 +51,7 @@ class PatternKindAnalyzer : public StmtExprVisitor { } private: - bool IsOutputBlock(const SBlockNode* block) { + bool IsOutputBlock(const s_tir::SBlockNode* block) { for (const TensorRegion& write_region : block->writes) { if (param_buffers_.count(write_region->source.as_or_throw())) { 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..925bbda6932a 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,10 +191,10 @@ 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(); + 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)) { @@ -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/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/relax/transform/split_call_tir_by_pattern.cc b/src/relax/transform/split_call_tir_by_pattern.cc index d179e1efe30c..54bf928e7121 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 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. // 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 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) { 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..a5fb6a692cbc 100644 --- a/src/s_tir/analysis/is_pure_function.cc +++ b/src/s_tir/analysis/is_pure_function.cc @@ -24,10 +24,11 @@ #include #include #include +#include +#include #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/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..b293b0d41dea 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,7 +435,7 @@ 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; + const s_tir::MatchBufferRegion& match_buffer = it->second; buffer = match_buffer->source->source.as_or_throw(); region = ConvertMatchedRegion(match_buffer, std::move(region)); } @@ -495,7 +498,7 @@ void BlockReadWriteDetector::UpdateOpaque(const Var& buffer_var) { } ffi::Array> GetSBlockAccessRegion( - const SBlock& block, const ffi::Map& buffer_var_map) { + 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(); @@ -512,7 +515,7 @@ ffi::Array> GetSBlockAccessRegion( } ffi::Array> GetSBlockReadWriteRegion( - const SBlock& block, const ffi::Map& buffer_var_map) { + 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..0d73fb07655e 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,18 @@ class LCADetector : public StmtExprVisitor { UpdateDominateScopeOfNonDataParIter(op); // Update match_buffers - for (const MatchBufferRegion& match_buffer : block->match_buffers) { + for (const s_tir::MatchBufferRegion& match_buffer : block->match_buffers) { UpdateBufferLCA(match_buffer->source->source.as_or_throw().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 +181,7 @@ class LCADetector : public StmtExprVisitor { // collect non-data-parallel block iteration's dominate scope. // for reduction iter type, we maintain the highest dominate scope for all reduce iters. // for other iter type, we maintain the dict for each individual iter. - const SBlock& block = block_realize->block; + const s_tir::SBlock& block = block_realize->block; bool is_reduce_block = false; for (size_t i = 0; i < block_realize->iter_values.size(); ++i) { const IterVar& iter_var = block->iter_vars[i]; @@ -263,12 +265,12 @@ class LCADetector : public StmtExprVisitor { blockidx_scopes_.push_back(ancestor_scopes_.back()); } } - return StmtExprVisitor::Visit_(op); + return s_tir::StmtExprVisitor::Visit_(op); } // Declared regions carry bounds, not opaque runtime accesses. ffi::Optional 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)); @@ -354,7 +356,7 @@ class LCADetector : public StmtExprVisitor { return lhs; } - /*! \brief The ancestor scope stacks info (SBlock and For). The + /*! \brief The ancestor scope stacks info (s_tir::SBlock and For). The * first element is initialized in LCADetector::Detect to represent * the root scope. */ diff --git a/src/s_tir/analysis/verify_gpu_code.cc b/src/s_tir/analysis/verify_gpu_code.cc index 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..b1a6439fae44 --- /dev/null +++ b/src/s_tir/analysis/verify_well_formed.cc @@ -0,0 +1,168 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +#include "../../tirx/analysis/verify_well_formed.h" + +#include +#include +#include + +#include "../ir/tir_visitor_with_path.h" + +namespace tvm { +namespace s_tir { +using tirx::ForNode; +using tirx::PrimFunc; + +/*! \brief Verify all Expr inside the block does not contain: + * 1. loop vars outside the current block. + * 2. block vars of parent blocks. + */ +class BlockVarAccessVerifier : public StmtExprVisitor { + public: + static bool Verify(const PrimFunc& func, bool assert_mode) { + auto verifier = ffi::make_object(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}; +}; + +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(); + TVM_FFI_UNREACHABLE(); + }); +} +} // namespace s_tir +} // namespace tvm diff --git a/src/s_tir/backend/adreno/inject_texture_alloc.cc b/src/s_tir/backend/adreno/inject_texture_alloc.cc index 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..edf356e1e5cb --- /dev/null +++ b/src/s_tir/ir/data_type_rewriter.cc @@ -0,0 +1,241 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +#include "data_type_rewriter.h" + +#include + +#include + +namespace tvm { +namespace s_tir { +using namespace tvm::tirx; +using namespace tvm::prim; + +PrimFunc IndexDataTypeNormalizer::Rewrite(PrimFunc func) { + // Keep this short setup local so its collector uses S-TIR block semantics + // without adding a dialect-specific collector hook to the TIRX normalizer. + // Collect scalar dtype requirements without changing types. Buffer definitions + // are rewritten only after every scalar replacement has been seeded. + class IndexVarCollector : public IndexDataTypeNormalizer { + public: + explicit IndexVarCollector(std::function 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); + } + + 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 { + seed(param.get()); + } + } + 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 IndexDataTypeNormalizer::Mutate_(const SBlockRealizeNode* op, + InplaceMode inplace_mode) { + bool is_condition = this->is_condition_; + this->is_condition_ = true; + auto new_predicate_result = this->Mutate(op->predicate, inplace_mode); + bool new_predicate_unchanged = new_predicate_result.UnchangedOrSameAs(op->predicate); + auto new_predicate = std::move(new_predicate_result).ValueOrUnchanged(op->predicate); + this->is_condition_ = is_condition; + + bool is_enabled = this->is_enabled_; + this->is_enabled_ = true; + auto new_iter_values = this->Mutate(op->iter_values, inplace_mode) + .as_or_throw>>() + .ValueOrUnchanged(op->iter_values); + this->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 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([this](const MatchBufferRegion& match) { + BufferVar buffer = this->WithDefRegionKind(kTVMFFIDefRegionKindSimple, [&] { + return this->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 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; + 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 IndexDataTypeNormalizer::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 = this->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 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 = + this->Mutate(iter_var->dom->min, InplaceMode::kDisallow).ValueOrUnchanged(iter_var->dom->min); + PrimExpr extent = this->Mutate(iter_var->dom->extent, InplaceMode::kDisallow) + .ValueOrUnchanged(iter_var->dom->extent); + this->is_enabled_ = is_enabled; + if (!new_var.same_as(iter_var->var) || !min.same_as(iter_var->dom->min) || + !extent.same_as(iter_var->dom->extent)) { + IterVar new_iter_var = iter_var; + IterVarNode* n = new_iter_var.CopyOnWrite(); + n->var = std::move(new_var); + n->dom = Range(min, extent); + return new_iter_var; + } + return iter_var; +} + +TensorRegion IndexDataTypeNormalizer::VisitBufferRegion(const TensorRegion& buffer_region) { + BufferVar remapped_buffer = + this->Mutate(buffer_region->source.as_or_throw(), InplaceMode::kDisallow) + .as_or_throw>() + .ValueOrUnchanged(buffer_region->source.as_or_throw()); + + bool is_enabled = this->is_enabled_; + this->is_enabled_ = true; + auto new_region = buffer_region->region.Map([&](const Range& range) { + return Range::FromMinExtent( + this->Mutate(range->min, InplaceMode::kDisallow).ValueOrUnchanged(range->min), + this->Mutate(range->extent, InplaceMode::kDisallow).ValueOrUnchanged(range->extent)); + }); + this->is_enabled_ = is_enabled; + + if (!remapped_buffer.same_as(buffer_region->source.as_or_throw()) || + !new_region.same_as(buffer_region->region)) { + return BufferRegion(remapped_buffer, new_region); + } else { + return buffer_region; + } +} + +} // namespace s_tir +} // namespace tvm diff --git a/src/s_tir/ir/data_type_rewriter.h b/src/s_tir/ir/data_type_rewriter.h new file mode 100644 index 000000000000..cf0c21e79657 --- /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); + + virtual UnchangedOr Mutate_(const SBlockNode* op, InplaceMode inplace_mode); + virtual 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); + TensorRegion VisitBufferRegion(const TensorRegion& 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 new file mode 100644 index 000000000000..bd68223887e7 --- /dev/null +++ b/src/s_tir/ir/ir_mutator_with_analyzer.cc @@ -0,0 +1,36 @@ +/* + * 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 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) { + analyzer_->Bind(iter_var->var, iter_var->dom); + iter_vars_.Set(iter_var->var, iter_var->dom); + } + return s_tir::StmtExprMutator::MutateBlock(this, op, inplace_mode); + }); +} + +} // namespace 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 new file mode 100644 index 000000000000..49179ab43cb9 --- /dev/null +++ b/src/s_tir/ir/ir_mutator_with_analyzer.h @@ -0,0 +1,60 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +#ifndef TVM_S_TIR_IR_MUTATOR_WITH_ANALYZER_H_ +#define TVM_S_TIR_IR_MUTATOR_WITH_ANALYZER_H_ + +#include + +#include "../../tirx/ir/ir_mutator_with_analyzer.h" + +namespace tvm { +namespace s_tir { +class IRMutatorWithAnalyzer : public tirx::IRMutatorWithAnalyzer { + public: + using Parent = tirx::IRMutatorWithAnalyzer; + using Parent::Mutate; + using Parent::Mutate_; + explicit IRMutatorWithAnalyzer(const arith::Analyzer& analyzer) + : IRMutatorWithAnalyzer(analyzer.get()) {} + explicit IRMutatorWithAnalyzer(arith::AnalyzerObj* analyzer) : Parent(analyzer, GlobalVTable()) {} + virtual UnchangedOr 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); + SetDispatch(vtable); + 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 // TVM_S_TIR_IR_MUTATOR_WITH_ANALYZER_H_ 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..56834a8b98cb --- /dev/null +++ b/src/s_tir/ir/ir_visitor_with_analyzer.cc @@ -0,0 +1,35 @@ +/* + * 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 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) { + analyzer_->Bind(iter_var->var, iter_var->dom); + } + return s_tir::StmtExprVisitor::VisitBlock(this, op); + }); +} + +} // 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 new file mode 100644 index 000000000000..ea52dba33fb9 --- /dev/null +++ b/src/s_tir/ir/ir_visitor_with_analyzer.h @@ -0,0 +1,49 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +#ifndef TVM_S_TIR_IR_VISITOR_WITH_ANALYZER_H_ +#define TVM_S_TIR_IR_VISITOR_WITH_ANALYZER_H_ + +#include + +#include "../../tirx/ir/ir_visitor_with_analyzer.h" + +namespace tvm { +namespace s_tir { +class IRVisitorWithAnalyzer : public tirx::IRVisitorWithAnalyzer { + public: + using Parent = tirx::IRVisitorWithAnalyzer; + using Parent::Visit; + using Parent::Visit_; + TVM_DEFINE_OBJECT_FUNCTOR_DEFAULT_CONSTRUCTOR(IRVisitorWithAnalyzer, Parent) + virtual ffi::Optional 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); + SetDispatch(vtable); + SetDispatch(vtable); + } +}; +} // namespace s_tir +} // namespace tvm +#endif // TVM_S_TIR_IR_VISITOR_WITH_ANALYZER_H_ diff --git a/src/s_tir/ir/tir_visitor_with_path.cc b/src/s_tir/ir/tir_visitor_with_path.cc new file mode 100644 index 000000000000..af69a4b4feaa --- /dev/null +++ b/src/s_tir/ir/tir_visitor_with_path.cc @@ -0,0 +1,81 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +#include "tir_visitor_with_path.h" + +#include + +namespace tvm { +namespace s_tir { +using namespace tirx; +using AccessPath = ffi::reflection::AccessPath; + +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)); + } + } + + { + auto match_path = path->Attr("match_buffers"); + for (size_t i = 0; i < op->match_buffers.size(); i++) { + Visit(op->match_buffers[i]->source, match_path->ArrayItem(i)->Attr("source")); + auto buf = op->match_buffers[i]->buffer; + auto buffer_path = match_path->ArrayItem(i)->Attr("buffer"); + + for (auto& def : WithMatchBufferDefs(buf, buffer_path)) { + context.push_back(std::move(def)); + } + context.push_back(WithDef(buf, buffer_path)); + } + } + + // Regions may use allocation and match-buffer definitions in this block. + Visit(op->reads, path->Attr("reads")); + Visit(op->writes, path->Attr("writes")); + + bind_scope_.WithNewScope([&]() { Visit(op->init, path->Attr("init")); }); + bind_scope_.WithNewScope([&]() { Visit(op->body, path->Attr("body")); }); + + while (context.size()) context.pop_back(); +} + +void TIRVisitorWithPath::Dispatch_(const SBlockRealizeNode* op, AccessPath path) { + Visit(op->iter_values, path->Attr("iter_values")); + Visit(op->predicate, path->Attr("predicate")); + Visit(op->block, path->Attr("block")); +} + +} // namespace s_tir +} // namespace tvm diff --git a/src/s_tir/ir/tir_visitor_with_path.h b/src/s_tir/ir/tir_visitor_with_path.h new file mode 100644 index 000000000000..6556d2eec7f8 --- /dev/null +++ b/src/s_tir/ir/tir_visitor_with_path.h @@ -0,0 +1,54 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +#ifndef TVM_S_TIR_IR_TIR_VISITOR_WITH_PATH_H_ +#define TVM_S_TIR_IR_TIR_VISITOR_WITH_PATH_H_ +#include + +#include "../../tirx/ir/tir_visitor_with_path.h" +namespace tvm { +namespace s_tir { +class TIRVisitorWithPath : public tirx::TIRVisitorWithPath { + public: + using Parent = tirx::TIRVisitorWithPath; + TVM_DEFINE_OBJECT_FUNCTOR_DEFAULT_CONSTRUCTOR(TIRVisitorWithPath, Parent) + protected: + using AccessPath = ffi::reflection::AccessPath; + using Parent::Dispatch_; + virtual void Dispatch_(const SBlockNode* op, AccessPath path); + virtual void Dispatch_(const SBlockRealizeNode* op, AccessPath path); + static void InitVTable(VTable* vtable) { + Parent::InitVTable(vtable); + // StmtVisitor is a protected base of the path visitor, so create the + // downcast thunks here instead of in StmtVisitor::SetDispatch. + vtable->SetDispatch( + [](const ffi::ObjectRef& node, StmtVisitor* self, AccessPath path) { + static_cast(self)->Dispatch_( + static_cast(node.get()), path); + }); + vtable->SetDispatch( + [](const ffi::ObjectRef& node, StmtVisitor* self, AccessPath path) { + static_cast(self)->Dispatch_( + 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 9ad9fc5a7144..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 @@ -19,6 +19,7 @@ #include #include #include +#include #include #include @@ -326,11 +327,11 @@ tvm::transform::Sequential PassListForPerStoreFeature() { s_tir::transform::PlanAndUpdateBufferAllocationLocation(), s_tir::transform::ConvertBlocksToOpaque(), s_tir::transform::CompactBufferAllocation(), - tirx::transform::StmtSimplify(), + s_tir::transform::StmtSimplify(), s_tir::transform::LowerAutoCopy(), s_tir::transform::UnifyThreadBinding(), s_tir::transform::LowerMatchBuffer(), - tirx::transform::StmtSimplify(), + s_tir::transform::StmtSimplify(), }); } diff --git a/src/s_tir/meta_schedule/module_equality.cc b/src/s_tir/meta_schedule/module_equality.cc index 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/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/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 7c4e9238a902..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 @@ -20,6 +20,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,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 = tirx::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); @@ -627,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 @@ -782,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,7 +869,7 @@ ffi::Optional MultiLevelTilingTensorCoreNode::TransformWithTensorIntrin( // Refresh block pointer (block sref is not invalidated) block = TVM_SREF_TO_SBLOCK(block_sref); const tvm::TensorRegion& reindexed_buffer_region = s_tir::GetNthAccessBufferRegion( - state->sch->state(), ffi::GetRef(block), buffer_index, index_type); + 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/schedule_rule/multi_level_tiling_with_intrin.cc b/src/s_tir/meta_schedule/schedule_rule/multi_level_tiling_with_intrin.cc index fab8937ae56c..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 @@ -19,6 +19,7 @@ #include #include +#include #include "../../schedule/analysis.h" #include "../../schedule/transform.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..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,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/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 bb3e1e77019f..903203e5a8e6 100644 --- a/src/s_tir/schedule/concrete_schedule.cc +++ b/src/s_tir/schedule/concrete_schedule.cc @@ -20,6 +20,8 @@ #include #include +#include +#include #include @@ -927,7 +929,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 +938,7 @@ void ConcreteScheduleNode::Tensorize(const LoopRV& loop_rv, const ffi::String& i void ConcreteScheduleNode::Tensorize(const SBlockRV& block_rv, const ffi::String& intrin, bool preserve_unit_iters) { TVM_TIR_SCHEDULE_BEGIN(); - s_tir::Tensorize(state_, this->GetSRef(block_rv), tirx::TensorIntrin::Get(intrin).value(), + s_tir::Tensorize(state_, this->GetSRef(block_rv), TensorIntrin::Get(intrin).value(), preserve_unit_iters); this->state_->DebugVerify(); TVM_TIR_SCHEDULE_END("tensorize", this->error_render_level_); diff --git a/src/s_tir/schedule/concrete_schedule.h b/src/s_tir/schedule/concrete_schedule.h index 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 83ab03ec5780..89c058f9ba70 100644 --- a/src/s_tir/schedule/primitive.h +++ b/src/s_tir/schedule/primitive.h @@ -22,6 +22,8 @@ #include #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/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/schedule/primitive/blockize_tensorize.cc b/src/s_tir/schedule/primitive/blockize_tensorize.cc index 572f8c48048e..fc8811822407 100644 --- a/src/s_tir/schedule/primitive/blockize_tensorize.cc +++ b/src/s_tir/schedule/primitive/blockize_tensorize.cc @@ -21,11 +21,13 @@ #include #include #include +#include +#include #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" @@ -817,7 +819,7 @@ void Tensorize(ScheduleState self, const StmtSRef& sref, const TensorIntrin& int } arith::Analyzer analyzer; - PrimFunc intrin_desc = StmtSimplify(intrin->desc, analyzer); + PrimFunc intrin_desc = s_tir::StmtSimplify(intrin->desc, analyzer); PrimFunc intrin_impl = DeepCopy(intrin->impl); int index_dtype_bits = -1; diff --git a/src/s_tir/schedule/primitive/cache_index.cc b/src/s_tir/schedule/primitive/cache_index.cc index 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..f07a88c7ba50 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 { @@ -431,7 +432,7 @@ ScheduleState::ScheduleState(IRModule mod, int debug_mask, bool enable_check) { const BaseFunc& base_func = kv.second; if (auto opt = base_func.as()) { auto func = opt.value(); - VerifyWellFormed(func); + s_tir::VerifyWellFormed(func); SBlockInfoCollector::Collect(self, func->body); } } diff --git a/src/s_tir/schedule/traced_schedule.cc b/src/s_tir/schedule/traced_schedule.cc index 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 c962e1f0f8fb..bb0e1565c34a 100644 --- a/src/s_tir/schedule/transform.cc +++ b/src/s_tir/schedule/transform.cc @@ -20,6 +20,8 @@ #include #include #include +#include +#include #include #include "../../tirx/transform/ir_utils.h" @@ -316,7 +318,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()) { @@ -468,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(); @@ -479,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); @@ -488,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..98753c031b2b --- /dev/null +++ b/src/s_tir/stmt.cc @@ -0,0 +1,426 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +/*! + * \file tvm/s_tir/stmt.cc + * \brief Schedulable block definitions and structural traversal. + */ +#include +#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 { + // 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->WithDefRegionKind( + kTVMFFIDefRegionKindSimple, [&]() { return visitor->VisitExpected(self->alloc_buffers); })); + TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->match_buffers)); + TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->reads)); + TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->writes)); + TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->annotations)); + TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->init)); + TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->body)); + return ffi::AnyView(nullptr).CopyToTVMFFIAny(); +} + +TVMFFIAny SBlockMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept { + // Establish allocation and match-buffer definitions before their region uses. + // Whole iterators and annotations remain part of structural traversal. + // skips: name_hint + const SBlockNode* self = + ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(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_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)); + 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)); + 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 { + // 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_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)); + 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, + 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, 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("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, + 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("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() + .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..924a8efb8737 --- /dev/null +++ b/src/s_tir/stmt_functor.cc @@ -0,0 +1,201 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +/*! + * \file stmt_functor.cc + * \brief Native traversal of schedulable TIR nodes. + */ +#include + +#include +#include + +namespace tvm { +namespace s_tir { + +using namespace tirx; + +void StmtExprVisitor::InitVTable(VTable* vtable) { + tirx::StmtExprVisitor::InitVTable(vtable); + SetDispatch(vtable); + SetDispatch(vtable); +} + +void StmtExprMutator::InitVTable(VTable* vtable) { + tirx::StmtExprMutator::InitVTable(vtable); + 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)); + } + // Define match-buffer targets before visiting reads/writes that may use them. + // This differs from the old TIRX native order (reads/writes before matches) + // and agrees with structural traversal's definition-before-use contract. + for (const MatchBufferRegion& match_buffer_region : op->match_buffers) { + TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->WithDefRegionKind( + kTVMFFIDefRegionKindSimple, [&]() { return visitor->Visit(match_buffer_region->buffer); })); + TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitBufferMetadata(match_buffer_region->buffer)); + TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->Visit(match_buffer_region->source)); + } + for (const TensorRegion& region : op->reads) { + TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->Visit(region)); + } + for (const TensorRegion& region : op->writes) { + TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->Visit(region)); + } + if (op->init.has_value()) { + TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->Visit(op->init.value())); + } + return visitor->Visit(op->body); +} + +ffi::Optional 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 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 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/tensor_intrin.cc b/src/s_tir/tensor_intrin.cc new file mode 100644 index 000000000000..e381448198d5 --- /dev/null +++ b/src/s_tir/tensor_intrin.cc @@ -0,0 +1,100 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +/*! + * \file src/s_tir/tensor_intrin.cc + * \brief Tensor intrinsic registry and construction. + */ +#include +#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/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..4d3535fd4f56 100644 --- a/src/s_tir/transform/compact_buffer_region.cc +++ b/src/s_tir/transform/compact_buffer_region.cc @@ -27,19 +27,19 @@ #include #include #include +#include #include #include -#include #include #include #include "../../support/arena.h" #include "../../support/utils.h" -#include "../../tirx/transform/ir_utils.h" #include "../analysis/conditional_bounds.h" #include "../schedule/utils.h" #include "../support/nd_int_set.h" +#include "ir_utils.h" namespace tvm { namespace s_tir { diff --git a/src/s_tir/transform/convert_blocks_to_opaque.cc b/src/s_tir/transform/convert_blocks_to_opaque.cc index 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..35686953f328 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..3d4c52df1fe4 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,8 +38,8 @@ #include "../../arith/interval_set.h" #include "../../runtime/thread_storage_scope.h" -#include "../../tirx/ir/ir_mutator_with_analyzer.h" -#include "../../tirx/transform/ir_utils.h" +#include "../../s_tir/ir/ir_mutator_with_analyzer.h" +#include "ir_utils.h" namespace tvm { namespace s_tir { @@ -455,10 +456,10 @@ class HoistInfoCollector : public StmtExprVisitor { std::unordered_set active_loop_vars; }; -class ExpressionHoister : public tirx::IRMutatorWithAnalyzer { +class ExpressionHoister : public s_tir::IRMutatorWithAnalyzer { public: - using tirx::IRMutatorWithAnalyzer::Mutate; - using tirx::IRMutatorWithAnalyzer::Mutate_; + using s_tir::IRMutatorWithAnalyzer::Mutate; + using s_tir::IRMutatorWithAnalyzer::Mutate_; static Stmt Hoist(Stmt stmt, HoistExpressionConfig config) { auto loop_info = HoistInfoCollector::Collect(stmt, config); @@ -466,12 +467,12 @@ class ExpressionHoister : public tirx::IRMutatorWithAnalyzer { arith::Analyzer analyzer; auto hoister = ffi::make_object(std::move(loop_info), config, analyzer); stmt = hoister->Mutate(stmt, InplaceMode::kAllow).ValueOrUnchanged(std::move(stmt)); - stmt = ConvertSSA(std::move(stmt)); + stmt = s_tir::ConvertSSA(std::move(stmt)); return stmt; } private: - using Parent = tirx::IRMutatorWithAnalyzer; + using Parent = s_tir::IRMutatorWithAnalyzer; public: explicit ExpressionHoister(std::vector loop_info, @@ -592,7 +593,7 @@ Pass HoistExpression() { return tvm::transform::Sequential( { insertion_pass, - tirx::transform::StmtSimplify(), + s_tir::transform::StmtSimplify(), tirx::transform::RemoveNoOp(), }, "s_tir.HoistExpression"); @@ -630,7 +631,7 @@ static Pass HoistIfThenElseImpl() { return tvm::transform::Sequential( { insertion_pass, - tirx::transform::StmtSimplify(), + s_tir::transform::StmtSimplify(), tirx::transform::RemoveNoOp(), }, "s_tir.HoistIfThenElse"); @@ -648,7 +649,7 @@ static Pass HoistIfThenElseBasicImpl() { return tvm::transform::Sequential( { insertion_pass, - tirx::transform::StmtSimplify(), + s_tir::transform::StmtSimplify(), tirx::transform::RemoveNoOp(), }, "s_tir.HoistIfThenElseBasic"); diff --git a/src/s_tir/transform/inject_double_buffer.cc b/src/s_tir/transform/inject_double_buffer.cc index 2100a82780bb..7474c552ca76 100644 --- a/src/s_tir/transform/inject_double_buffer.cc +++ b/src/s_tir/transform/inject_double_buffer.cc @@ -27,11 +27,11 @@ #include #include #include +#include #include #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_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_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 99ccdeea8bfb..b80cee1f49c0 100644 --- a/src/s_tir/transform/inject_virtual_thread.cc +++ b/src/s_tir/transform/inject_virtual_thread.cc @@ -27,14 +27,14 @@ #include #include #include +#include #include #include -#include #include -#include "../../tirx/ir/ir_mutator_with_analyzer.h" -#include "../../tirx/transform/ir_utils.h" +#include "../../s_tir/ir/ir_mutator_with_analyzer.h" +#include "ir_utils.h" namespace tvm { namespace s_tir { @@ -230,10 +230,10 @@ class VarTouchedAnalysis : public StmtExprVisitor { // Inject virtual thread loop // rewrite the buffer access pattern when necessary. -class VTInjector : public tirx::IRMutatorWithAnalyzer { +class VTInjector : public s_tir::IRMutatorWithAnalyzer { public: - using tirx::IRMutatorWithAnalyzer::Mutate; - using tirx::IRMutatorWithAnalyzer::Mutate_; + using s_tir::IRMutatorWithAnalyzer::Mutate; + using s_tir::IRMutatorWithAnalyzer::Mutate_; // constructor VTInjector(arith::AnalyzerObj* analyzer, Var var, int num_threads, @@ -686,10 +686,10 @@ class VTInjector : public tirx::IRMutatorWithAnalyzer { */ }; -class VirtualThreadInjector : public tirx::IRMutatorWithAnalyzer { +class VirtualThreadInjector : public s_tir::IRMutatorWithAnalyzer { public: - using tirx::IRMutatorWithAnalyzer::Mutate; - using tirx::IRMutatorWithAnalyzer::Mutate_; + using s_tir::IRMutatorWithAnalyzer::Mutate; + using s_tir::IRMutatorWithAnalyzer::Mutate_; using IRMutatorWithAnalyzer::IRMutatorWithAnalyzer; @@ -722,7 +722,7 @@ Pass InjectVirtualThread() { n->body = ffi::make_object(analyzer) ->Mutate(n->body, InplaceMode::kAllow) .ValueOrUnchanged(std::move(n->body)); - n->body = ConvertSSA(std::move(n->body)); + n->body = s_tir::ConvertSSA(std::move(n->body)); return f; }; return CreatePrimFuncPass(pass_func, 0, "s_tir.InjectVirtualThread", {}); diff --git a/src/s_tir/transform/ir_utils.cc b/src/s_tir/transform/ir_utils.cc new file mode 100644 index 000000000000..adb3e54c5f11 --- /dev/null +++ b/src/s_tir/transform/ir_utils.cc @@ -0,0 +1,201 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +#include "ir_utils.h" + +#include +#include +#include +#include + +namespace tvm { +namespace s_tir { +using namespace tirx; +using namespace tvm::prim; + +namespace { +class SIRConvertSSA final : public tirx::IRConvertSSA { + public: + TVM_DEFINE_OBJECT_FUNCTOR_DEFAULT_CONSTRUCTOR(SIRConvertSSA, tirx::IRConvertSSA) + using tirx::IRConvertSSA::Mutate_; + + protected: + static void InitVTable(VTable* table) { + tirx::IRConvertSSA::InitVTable(table); + SetDispatch(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 = [&](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); + 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; + 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; +} + +/*! \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>()]->source.as_or_throw().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 new file mode 100644 index 000000000000..e9d6c6cc2daa --- /dev/null +++ b/src/s_tir/transform/ir_utils.h @@ -0,0 +1,66 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +#ifndef TVM_S_TIR_TRANSFORM_IR_UTILS_H_ +#define TVM_S_TIR_TRANSFORM_IR_UTILS_H_ + +#include +#include + +#include + +#include "../../tirx/transform/ir_utils.h" + +namespace tvm { +namespace s_tir { + +/*! \brief Convert a schedulable statement or module to SSA form. */ +tirx::Stmt ConvertSSA(tirx::Stmt stmt); +IRModule ConvertSSA(IRModule mod); + +/*! + * \brief Convert match buffer target buffer access indices to original one. + * \param indices The indices of the target buffer + * \return The indices of source buffer. + */ +ffi::Array 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); + +/*! \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/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..392d821b71a8 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 @@ -42,7 +43,7 @@ #include "../../arith/interval_set.h" #include "../../runtime/thread_storage_scope.h" -#include "../../tirx/transform/ir_utils.h" +#include "ir_utils.h" namespace tvm { namespace s_tir { @@ -843,7 +844,7 @@ Stmt LoopPartitioner::TryPartition(const Stmt& stmt, Var var, PrimExpr min, Prim ->Mutate(stmt) .ValueOrUnchanged(stmt); } - s = ConvertSSA(s); + s = s_tir::ConvertSSA(s); return s; } diff --git a/src/s_tir/transform/lower_async_dma.cc b/src/s_tir/transform/lower_async_dma.cc index 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..1ac40d183247 100644 --- a/src/s_tir/transform/lower_opaque_block.cc +++ b/src/s_tir/transform/lower_opaque_block.cc @@ -24,10 +24,10 @@ #include #include #include +#include #include -#include -#include "../../tirx/transform/ir_utils.h" +#include "ir_utils.h" namespace tvm { namespace s_tir { diff --git a/src/s_tir/transform/lower_thread_allreduce.cc b/src/s_tir/transform/lower_thread_allreduce.cc index 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_simplify.cc b/src/s_tir/transform/stmt_simplify.cc new file mode 100644 index 000000000000..36d8cf432806 --- /dev/null +++ b/src/s_tir/transform/stmt_simplify.cc @@ -0,0 +1,94 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +#include "stmt_simplify.h" + +#include +#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 final : public tirx::StmtSimplifier { + public: + using Parent = tirx::StmtSimplifier; + StmtSimplifier(const arith::Analyzer& analyzer, tirx::StmtSimplifyConfig config) + : Parent(GlobalVTable(), analyzer, config) {} + using Parent::Mutate_; + using Parent::Run; + + public: + UnchangedOr 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); + 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/s_tir/transform/stmt_simplify.h b/src/s_tir/transform/stmt_simplify.h new file mode 100644 index 000000000000..8dcff5bcc6d3 --- /dev/null +++ b/src/s_tir/transform/stmt_simplify.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_S_TIR_TRANSFORM_STMT_SIMPLIFY_H_ +#define TVM_S_TIR_TRANSFORM_STMT_SIMPLIFY_H_ +#include +#include +namespace tvm { +namespace s_tir { +tirx::PrimFunc StmtSimplify(tirx::PrimFunc func, const arith::Analyzer& analyzer); +} // namespace s_tir +} // namespace tvm +#endif // TVM_S_TIR_TRANSFORM_STMT_SIMPLIFY_H_ diff --git a/src/s_tir/transform/storage_access.h b/src/s_tir/transform/storage_access.h index 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..ca0b2045d65f 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 @@ -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 { @@ -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) { @@ -872,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)); @@ -948,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 66adcebb0b3e..8407bdb67122 100644 --- a/src/tirx/analysis/verify_tirx_well_formed.cc +++ b/src/tirx/analysis/verify_tirx_well_formed.cc @@ -50,14 +50,6 @@ class ExecScopeVerifier : public Verifier { private: using Verifier::Visit; - void Dispatch_(const SBlockNode* op, ffi::reflection::AccessPath path) override { - Verify(false) << "TIRxError: SBlock is not allowed in tirx=True mode at " << path; - } - - void Dispatch_(const SBlockRealizeNode* op, ffi::reflection::AccessPath path) override { - Verify(false) << "TIRxError: SBlockRealize is not allowed in tirx=True mode at " << path; - } - void Dispatch_(const tirx::TilePrimitiveCallNode* op, ffi::reflection::AccessPath path) override { static const auto& category_map = Op::GetAttrMap("TIRxOpCategory"); Verify(category_map.get(op->op, ffi::String("")) == "tile_primitive") @@ -131,14 +123,6 @@ 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; - } }; class AsyncStructsVerifier : public Verifier { @@ -147,14 +131,6 @@ 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; - } }; class DeviceFuncVerifier : public Verifier { @@ -163,14 +139,6 @@ 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 VerifyTIRxWellFormed(const PrimFunc& func, bool assert_mode, bool device_func) { @@ -197,11 +165,6 @@ bool VerifyTIRxWellFormed(const PrimFunc& func, bool assert_mode, bool device_fu bool VerifyTIRxWellFormed(const IRModule& mod, bool assert_mode, bool device_func) { for (const auto& [gvar, base_func] : mod->functions) { if (auto prim_func = base_func.as()) { - // 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 0123c463b64c..adef48e1c111 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,344 +43,6 @@ namespace tirx { using AccessPath = ffi::reflection::AccessPath; -/*! \brief Verify all Expr inside the block does not contain: - * 1. loop vars outside the current block. - * 2. block vars of parent blocks. - */ -class BlockVarAccessVerifier : public StmtExprVisitor { - public: - static bool Verify(const PrimFunc& func, bool assert_mode) { - auto verifier = ffi::make_object(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}; -}; - -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 @@ -421,37 +85,11 @@ class SingleEnvThreadVerifier : public Verifier { }; bool VerifyWellFormed(const PrimFunc& func, bool assert_mode) { - if (!BlockVarAccessVerifier::Verify(func, assert_mode)) { - return false; - } - - if (!UndefinedVarVerifier::Verify(func, assert_mode)) return false; - - if (!UndefinedBufferVerifier::Verify(func, assert_mode)) return false; - - if (!TensorLoadTypeVerifier::Verify(func, assert_mode)) return false; - - // TODO(Siyuan): add more checks here. - return true; + return VerifyWellFormedCommon(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 new file mode 100644 index 000000000000..16b57274ff90 --- /dev/null +++ b/src/tirx/analysis/verify_well_formed.h @@ -0,0 +1,282 @@ +/* + * 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 + +#include "../ir/tir_visitor_with_path.h" +namespace tvm { +namespace tirx { +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 // 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 2b7504513de7..8f79792d1826 100644 --- a/src/tirx/ir/data_type_rewriter.cc +++ b/src/tirx/ir/data_type_rewriter.cc @@ -27,14 +27,12 @@ #include #include #include -#include #include #include #include #include #include -#include #include "tvm/ir/expr.h" #include "tvm/ir/prim/expr.h" @@ -44,7 +42,6 @@ namespace tvm { namespace tirx { using namespace tvm::prim; - UnchangedOr DataTypeLegalizer::Mutate_(const ForNode* op, InplaceMode inplace_mode) { auto result = StmtExprMutator::Mutate_(op, inplace_mode); if (!result.IsUnchanged()) { @@ -74,51 +71,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 == 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" @@ -374,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 == s_tir::attr::virtual_thread) { + if (op->attr_key == attr::thread_extent || op->attr_key == tvm::tirx::attr::virtual_thread) { bool is_enabled = is_enabled_; is_enabled_ = true; auto stmt = DataTypeLegalizer::Mutate_(op, inplace_mode); @@ -392,165 +346,6 @@ UnchangedOr 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); @@ -755,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 c2cca01ed8d6..7d1e09d85b27 100644 --- a/src/tirx/ir/data_type_rewriter.h +++ b/src/tirx/ir/data_type_rewriter.h @@ -40,8 +40,6 @@ namespace tirx { * bounds. * - The data type of the binary and ternary expressions must be consistent with the data types of * each of their operands. - * - The data type of the bounds and binding values of block iter vars must be consistent with the - * data type of the block iter vars. * * Usually we enforce the consistency of data types when constructing the IR nodes. However, such * inconsistency may happen as a result of IR mutation in some passes. This class can be used as @@ -49,14 +47,14 @@ namespace tirx { */ class DataTypeLegalizer : public StmtExprMutator { public: + TVM_DEFINE_OBJECT_FUNCTOR_DEFAULT_CONSTRUCTOR(DataTypeLegalizer, StmtExprMutator) using StmtExprMutator::Mutate; using StmtExprMutator::Mutate_; protected: + explicit DataTypeLegalizer(const VTable* vtable) : StmtExprMutator(vtable) {} UnchangedOr 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 +97,14 @@ class DataTypeLegalizer : public StmtExprMutator { */ class IndexDataTypeRewriter : public DataTypeLegalizer { public: + TVM_DEFINE_OBJECT_FUNCTOR_DEFAULT_CONSTRUCTOR(IndexDataTypeRewriter, DataTypeLegalizer) using DataTypeLegalizer::Mutate; using DataTypeLegalizer::Mutate_; protected: + explicit IndexDataTypeRewriter(const VTable* vtable) : DataTypeLegalizer(vtable) {} using Parent = DataTypeLegalizer; UnchangedOr 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 +122,6 @@ 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); // indicator of index expr to rewrite bool is_enabled_{false}; // indicator of condition @@ -149,6 +143,7 @@ class IndexDataTypeNormalizer : public IndexDataTypeRewriter { PrimFunc Rewrite(PrimFunc func); protected: + IndexDataTypeNormalizer(PrimType target_data_type, const VTable* vtable); using Parent = IndexDataTypeRewriter; UnchangedOr Mutate_(const IntImmNode* op, InplaceMode inplace_mode) override; 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/src/tirx/ir/ir_mutator_with_analyzer.cc b/src/tirx/ir/ir_mutator_with_analyzer.cc index e26fdf5282e1..1c65f94874a3 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,18 @@ namespace tvm { namespace tirx { +void IRMutatorWithAnalyzer::InitVTable(VTable* vtable) { StmtExprMutator::InitVTable(vtable); } + +const IRMutatorWithAnalyzer::VTable* IRMutatorWithAnalyzer::GlobalVTable() { + static const VTable table = [] { + VTable table; + InitVTable(&table); + table.Finalize(); + return table; + }(); + return &table; +} + using namespace tvm::prim; using arith::detail::EnterConstraintFacts; @@ -113,16 +124,6 @@ UnchangedOr 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 +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 == s_tir::attr::virtual_thread) { + if (op->attr_key == tirx::attr::thread_extent || + op->attr_key == tvm::tirx::attr::virtual_thread) { IterVar iv = op->node.as_or_throw(); 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..bcc976c57ac5 100644 --- a/src/tirx/ir/ir_mutator_with_analyzer.h +++ b/src/tirx/ir/ir_mutator_with_analyzer.h @@ -51,12 +51,13 @@ class IRMutatorWithAnalyzer : public StmtExprMutator { public: using StmtExprMutator::Mutate; using StmtExprMutator::Mutate_; - explicit IRMutatorWithAnalyzer(const arith::Analyzer& analyzer) : analyzer_(analyzer.get()) {} - explicit IRMutatorWithAnalyzer(arith::AnalyzerObj* analyzer) : analyzer_(analyzer) {} + explicit IRMutatorWithAnalyzer(const arith::Analyzer& analyzer) + : IRMutatorWithAnalyzer(analyzer.get()) {} + explicit IRMutatorWithAnalyzer(arith::AnalyzerObj* analyzer) + : IRMutatorWithAnalyzer(analyzer, GlobalVTable()) {} // override functions that need to populate the context information. UnchangedOr 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 +67,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..ae823b3ddedc 100644 --- a/src/tirx/ir/ir_visitor_with_analyzer.cc +++ b/src/tirx/ir/ir_visitor_with_analyzer.cc @@ -24,13 +24,13 @@ #include #include -#include #include #include #include namespace tvm { namespace tirx { +void IRVisitorWithAnalyzer::InitVTable(VTable* vtable) { StmtExprVisitor::InitVTable(vtable); } ffi::Optional IRVisitorWithAnalyzer::Visit_(const ForNode* op) { return constraint_scope_.WithNewScope([&]() -> ffi::Optional { @@ -47,15 +47,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 +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 == s_tir::attr::virtual_thread) { + if (op->attr_key == tirx::attr::thread_extent || + op->attr_key == tvm::tirx::attr::virtual_thread) { IterVar iv = op->node.as_or_throw(); 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..63db9d9fc0cf 100644 --- a/src/tirx/ir/ir_visitor_with_analyzer.h +++ b/src/tirx/ir/ir_visitor_with_analyzer.h @@ -36,12 +36,13 @@ namespace tirx { class IRVisitorWithAnalyzer : public StmtExprVisitor { public: + 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 +55,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..55195e3dc6e9 100644 --- a/src/tirx/ir/specialize.cc +++ b/src/tirx/ir/specialize.cc @@ -26,7 +26,6 @@ #include #include #include -#include #include #include #include @@ -161,7 +160,12 @@ class PrimFuncSpecializer : public StmtExprMutator { 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)); } @@ -169,43 +173,18 @@ 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); } - 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 +326,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/stmt.cc b/src/tirx/ir/stmt.cc index 1699a8b686b6..179f546cd234 100644 --- a/src/tirx/ir/stmt.cc +++ b/src/tirx/ir/stmt.cc @@ -721,176 +721,6 @@ 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)); - 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 ScopeIdDefStmtVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept { const ScopeIdDefStmtNode* self = ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value); @@ -925,61 +755,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 +1328,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 +1350,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..b1d7884012cd 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" @@ -55,8 +56,6 @@ void StmtExprVisitor::InitVTable(VTable* vtable) { SetDispatch(vtable); SetDispatch(vtable); SetDispatch(vtable); - SetDispatch(vtable); - SetDispatch(vtable); SetDispatch(vtable); SetDispatch(vtable); } @@ -228,42 +227,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. @@ -324,8 +287,6 @@ void StmtExprMutator::InitVTable(VTable* vtable) { SetDispatch(vtable); SetDispatch(vtable); SetDispatch(vtable); - SetDispatch(vtable); - SetDispatch(vtable); SetDispatch(vtable); SetDispatch(vtable); } @@ -467,28 +428,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,75 +491,6 @@ 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)) - 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 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 0f0f684a50fa..37eaa20d396a 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,7 +33,6 @@ namespace tvm { namespace tirx { - using AccessPath = ffi::reflection::AccessPath; void TIRVisitorWithPath::Visit(const IRModule& mod, AccessPath path) { @@ -159,14 +157,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")); @@ -192,7 +182,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)) { + (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"))); @@ -268,57 +258,6 @@ 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) { for (size_t i = 0; i < op->args.size(); i++) { if (op->args[i] == nullptr) { diff --git a/src/tirx/ir/tir_visitor_with_path.h b/src/tirx/ir/tir_visitor_with_path.h index f4547becd6cc..674a3e6784cb 100644 --- a/src/tirx/ir/tir_visitor_with_path.h +++ b/src/tirx/ir/tir_visitor_with_path.h @@ -45,12 +45,17 @@ namespace tirx { class TIRVisitorWithPath : protected ExprFunctor, protected StmtFunctor { public: + TIRVisitorWithPath() = default; template void operator()(TObjectRef&& obj) { Visit(std::forward(obj), ffi::reflection::AccessPath::Root()); } protected: + 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); @@ -93,7 +98,6 @@ 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) { @@ -318,6 +320,10 @@ class Verifier : protected TIRVisitorWithPath { protected: explicit Verifier(bool assert_on_error) : assert_on_error_(assert_on_error) {} + void DispatchDefault_(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/script/builder/frame.cc b/src/tirx/script/builder/frame.cc index 031984e83d1b..9ae7e1915477 100644 --- a/src/tirx/script/builder/frame.cc +++ b/src/tirx/script/builder/frame.cc @@ -20,13 +20,14 @@ #include #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,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::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 +225,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..4fa7629cd02d 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. @@ -685,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/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/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/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/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/flatten_buffer.cc b/src/tirx/transform/flatten_buffer.cc index 0bf7e3cb436c..461431b5abc4 100644 --- a/src/tirx/transform/flatten_buffer.cc +++ b/src/tirx/transform/flatten_buffer.cc @@ -178,37 +178,6 @@ 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); - } - UnchangedOr Mutate_(const AllocBufferNode* op, InplaceMode inplace_mode) final { const FlatInfo& info = Define(op->buffer); if (info.flattened.same_as(op->buffer)) { @@ -323,31 +292,6 @@ 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())) { - 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 91e0d35e5a6b..5692dd7872e1 100644 --- a/src/tirx/transform/force_narrow_index_to_i32.cc +++ b/src/tirx/transform/force_narrow_index_to_i32.cc @@ -69,19 +69,22 @@ 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_; - } + 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 block_; + return alloc; } PrimFunc func_; diff --git a/src/tirx/transform/inline_private_functions.cc b/src/tirx/transform/inline_private_functions.cc index e385963b2200..9644c30495d2 100644 --- a/src/tirx/transform/inline_private_functions.cc +++ b/src/tirx/transform/inline_private_functions.cc @@ -121,12 +121,19 @@ 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 - // `tirx::SBlock` 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; + // 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; + 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; } diff --git a/src/tirx/transform/ir_utils.cc b/src/tirx/transform/ir_utils.cc index 121b12ba8af1..be6d6cba09db 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 @@ -88,674 +87,553 @@ Stmt MergeNest(const std::vector>& nest, Stmt body) { return body; } -class IRConvertSSA final : public StmtExprMutator { - 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; } + }(); - // Pop function-scope remaps in reverse order - PopAllRemapsInCurrentScope(); - function_scope_var_remap_.clear(); - return func; + 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); + + // 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 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 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 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; +} - 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; +Stmt IRConvertSSA::WithScope(const std::function& body) { + return scope_.WithNewScope(body); +} + +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; +} - 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; +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; } - 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; + 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); +} - UnchangedOr Mutate_(const SBlockNode* op, InplaceMode inplace_mode) final { - SBlock block = ffi::GetRef(op); +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; + } +} - // 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()); +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_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; + 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 StmtExprMutator::Mutate_(block.get(), - block.unique() ? inplace_mode : InplaceMode::kDisallow) - .ValueOrUnchanged(block) - .as_or_throw(); - }); + } } - 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; - } + // 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; + } + + // 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(); + } + + // 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; + } + } + + // 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; +} - return node; +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); } +} - 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; +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); +} - TensorLoad 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); +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); + 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)); + }); } +} - 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; - } +UnchangedOr IRConvertSSA::Mutate_(const WhileNode* op, InplaceMode inplace_mode) { + return scope_.WithNewScope([&]() -> Stmt { + return StmtExprMutator::Mutate_(op, inplace_mode).ValueOrUnchanged(ffi::GetRef(op)); + }); +} + +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; +} - 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; - } +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); } } - // 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; - } + 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); - // 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(); + 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); } - // 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 (delayed_define) { + if (!defined_.count(var.get())) { + function_scope_var_remap_.insert({var.get(), var}); + defined_.insert(var.get()); } } - // 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 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); - 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); - } - } + return output; - 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); + } else if (const VarNode* v = op->node.as()) { + Stmt stmt = scope_.WithNewScope([&]() -> Stmt { + return StmtExprMutator::Mutate_(op, inplace_mode).ValueOrUnchanged(ffi::GetRef(op)); }); - 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); - }); + 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 { - defined_.insert(v.get()); - return scope_.WithNewScope([&]() -> Stmt { - return StmtExprMutator::Mutate_(op, inplace_mode).ValueOrUnchanged(ffi::GetRef(op)); - }); + return stmt; } - } - UnchangedOr Mutate_(const WhileNode* op, InplaceMode inplace_mode) final { + } else { 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; - } +bool IRConvertSSA::BufferDependsOnVar(const BufferVar& buffer, const VarNode* var) { + if (buffer.get() == var) return 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); - } - - if (delayed_define) { - if (!defined_.count(var.get())) { - function_scope_var_remap_.insert({var.get(), var}); - defined_.insert(var.get()); - } + 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(); + }; + 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; } - - 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; + for (const Iter& iter : tile_layout->replica) { + if (uses_var(iter->extent) || uses_var(iter->stride)) return true; } - } else { - return scope_.WithNewScope([&]() -> Stmt { - return StmtExprMutator::Mutate_(op, inplace_mode).ValueOrUnchanged(ffi::GetRef(op)); - }); } } + return false; +} - private: - /*! \brief Record of a variable remap pushed to the current scope. */ - struct VarRemap { - Var old_var; - Var new_var; - }; +Var IRConvertSSA::MakeNewVar(const Var& old_var) { return Var(old_var->name, old_var->ty); } - /*! \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; +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}); +} - 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(); - }; - 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; +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(); } - 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; } - - /*! \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); } - - /*! \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}); + // 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); @@ -784,112 +662,14 @@ 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 == tvm::tirx::attr::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 == tvm::tirx::attr::async_wait_inflight_count); return std::make_pair(op->value, inner->value); } -/*! \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 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); - } - - /*! \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); - 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); @@ -930,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 1b24e462ee9a..004bfa6f9d2b 100644 --- a/src/tirx/transform/ir_utils.h +++ b/src/tirx/transform/ir_utils.h @@ -25,21 +25,23 @@ #define TVM_TIR_TRANSFORM_IR_UTILS_H_ #include -#include #include #include +#include #include #include -#include #include #include #include #include +#include +#include #include #include #include #include +#include #include #include @@ -225,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. @@ -232,21 +348,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. @@ -288,19 +389,9 @@ 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); -/*! \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/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 6e0d65d7f440..9fa93f78d881 100644 --- a/src/tirx/transform/narrow_datatype.cc +++ b/src/tirx/transform/narrow_datatype.cc @@ -27,7 +27,6 @@ #include #include #include -#include #include #include #include @@ -124,16 +123,8 @@ 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); - } - 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); 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/stmt_simplify.cc b/src/tirx/transform/stmt_simplify.cc index fbd51222a570..6af72187811b 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 @@ -39,57 +39,42 @@ #include "../ir/ir_mutator_with_analyzer.h" namespace tvm { -namespace arith { +namespace tirx { 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)); +arith::RewriteSimplifier::Extension StmtSimplifyConfigNode::GetEnabledExtensions() const { + arith::RewriteSimplifier::Extension flags = arith::RewriteSimplifier::kNone; + if (transitively_prove_inequalities) { + flags = arith::RewriteSimplifier::Extension( + flags | arith::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 = arith::RewriteSimplifier::Extension( + flags | arith::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 = arith::RewriteSimplifier::Extension( + flags | arith::RewriteSimplifier::kApplyConstraintsToBooleanBranches); + } + return flags; +} static StmtSimplifyConfig MakeDefaultStmtSimplifyConfig() { return tvm::transform::PassConfigWithDefaults(); @@ -99,162 +84,139 @@ 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()); +PrimFunc StmtSimplifier::Apply(PrimFunc func, const arith::Analyzer& analyzer, + ffi::Optional config_opt) { + auto config = config_opt.value_or(MakeDefaultStmtSimplifyConfig()); - 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; - } - - 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 - -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 { @@ -262,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 224df0ed8bfc..4b48c61f4a99 100644 --- a/src/tirx/transform/stmt_simplify.h +++ b/src/tirx/transform/stmt_simplify.h @@ -27,9 +27,72 @@ #include #include +#include "../ir/ir_mutator_with_analyzer.h" + namespace tvm { 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); + + arith::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 arith::Analyzer& analyzer, + ffi::Optional config_opt = std::nullopt); + + explicit StmtSimplifier(const arith::Analyzer& analyzer, StmtSimplifyConfig config) + : IRMutatorWithAnalyzer(analyzer), config_(config) {} + + protected: + using Parent = IRMutatorWithAnalyzer; + StmtSimplifier(const VTable* vtable, const arith::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_; +}; + /* \brief Simplify statements in the prim func * * Applies the same behavior as the tirx.transform.StmtSimplify pass. 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); 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/ir_functor_test.cc b/tests/cpp/ir_functor_test.cc index b62b5aa134d6..3ba10ebe20b3 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 @@ -234,12 +235,12 @@ TEST(IRF, StmtVisitor) { 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); + 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 = [&]() { @@ -362,14 +363,14 @@ TEST(IRF, StmtExprMutator) { // 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); + 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/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/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/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_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/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_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_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/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_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/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_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_memhammer_lower_auto_copy.py b/tests/python/s_tir/transform/test_s_tir_transform_memhammer_lower_auto_copy.py index 5d050149f8ef..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.tirx.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) 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/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 5f1102ea6437..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 @@ -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(): @@ -54,9 +54,9 @@ 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) + 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(): @@ -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(): @@ -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(): @@ -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, @@ -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(): 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") 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..7d4688b747d0 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.""" + """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..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 @@ -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,8 @@ 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) 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/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)) 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