diff --git a/.github/CODEOWNERSHIP b/.github/CODEOWNERSHIP index 9e4c6f4fb074..51b30d810605 100644 --- a/.github/CODEOWNERSHIP +++ b/.github/CODEOWNERSHIP @@ -119,10 +119,10 @@ src/target/** @junrushao1994 @vinx13 @tqchen @kparzysz-quic @ZihengJiang @masah include/tvm/target/** @junrushao1994 @vinx13 @tqchen @kparzysz-quic @ZihengJiang @masahi python/tvm/target/** @junrushao1994 @vinx13 @tqchen @kparzysz-quic @ZihengJiang @masahi -# arith: Arithmetic module and simplifiers -src/arith/** @tqchen @junrushao1994 @vinx13 -include/tvm/arith/** @tqchen @junrushao1994 @vinx13 -python/tvm/arith/** @tqchen @junrushao1994 @vinx13 +# sym: Symbolic analysis and simplifiers +src/sym/** @tqchen @junrushao1994 @vinx13 +include/tvm/sym/** @tqchen @junrushao1994 @vinx13 +python/tvm/sym/** @tqchen @junrushao1994 @vinx13 # parser src/parser/** @jroesch @slyubomirsky diff --git a/CMakeLists.txt b/CMakeLists.txt index 8c92057f3777..482d356230e9 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -296,7 +296,7 @@ assign_source_group("Include" ${GROUP_INCLUDE}) # Source file lists tvm_file_glob(GLOB_RECURSE COMPILER_SRCS src/ir/*.cc - src/arith/*.cc + src/sym/*.cc src/te/*.cc src/tirx/*.cc src/s_tir/*.cc diff --git a/cmake/modules/contrib/Z3.cmake b/cmake/modules/contrib/Z3.cmake index ec4372bc322e..bf61ec13524b 100644 --- a/cmake/modules/contrib/Z3.cmake +++ b/cmake/modules/contrib/Z3.cmake @@ -15,8 +15,8 @@ # specific language governing permissions and limitations # under the License. -# src/arith/z3_prover.cc is always part of COMPILER_SRCS (picked up by the -# src/arith/*.cc glob). It compiles a conservative stub by default and switches +# src/sym/z3_prover.cc is always part of COMPILER_SRCS (picked up by the +# src/sym/*.cc glob). It compiles a conservative stub by default and switches # to the real Z3 implementation only when the TVM_USE_Z3 macro is defined below. if(${USE_Z3} MATCHES ${IS_FALSE_PATTERN}) return() @@ -92,6 +92,6 @@ else() return() endif() -# Enable the real Z3 implementation inside the single src/arith/z3_prover.cc file. +# Enable the real Z3 implementation inside the single src/sym/z3_prover.cc file. add_compile_definitions(TVM_USE_Z3) message(STATUS "Build with Z3 SMT solver support") diff --git a/docs/arch/index.rst b/docs/arch/index.rst index 381abd1f751e..d500341f7907 100644 --- a/docs/arch/index.rst +++ b/docs/arch/index.rst @@ -418,11 +418,11 @@ schedule primitives and auto-tuning tools that operate on ``tirx::PrimFunc``: Please refer to the :ref:`TensorIR Deep Dive ` for more details. -tvm/arith ---------- +tvm/sym +------- This module is closely tied to TensorIR. One of the key problems in the low-level code generation is the analysis of the indices' -arithmetic properties — the positiveness, variable bound, and the integer set that describes the iterator space. arith module provides +arithmetic properties — the positiveness, variable bound, and the integer set that describes the iterator space. sym module provides a collection of tools that do (primarily integer) analysis. A TensorIR pass can use these analyses to simplify and optimize the code. tvm/te and tvm/topi diff --git a/docs/reference/api/python/index.rst b/docs/reference/api/python/index.rst index 6fba876950e4..f81de243f902 100644 --- a/docs/reference/api/python/index.rst +++ b/docs/reference/api/python/index.rst @@ -24,7 +24,7 @@ Python API error ir - arith + sym instrument transform target diff --git a/docs/reference/api/python/arith.rst b/docs/reference/api/python/sym.rst similarity index 94% rename from docs/reference/api/python/arith.rst rename to docs/reference/api/python/sym.rst index 742fbf5aef68..0dbb6ee72b72 100644 --- a/docs/reference/api/python/arith.rst +++ b/docs/reference/api/python/sym.rst @@ -15,8 +15,8 @@ specific language governing permissions and limitations under the License. -tvm.arith ---------- -.. automodule:: tvm.arith +tvm.sym +------- +.. automodule:: tvm.sym :members: :imported-members: diff --git a/docs/tirx/arch/lowering_pipeline.rst b/docs/tirx/arch/lowering_pipeline.rst index 72f7cc3100ae..19836b1a3b90 100644 --- a/docs/tirx/arch/lowering_pipeline.rst +++ b/docs/tirx/arch/lowering_pipeline.rst @@ -62,7 +62,7 @@ The ``tirx_pipeline`` module pass applies this exact sequence (a few are gated b axis is declared once * - 3 - ``StmtSimplify`` - - statement-level arithmetic simplification (the arith analyzer) + - statement-level arithmetic simplification (the sym analyzer) * - 4 - ``LowerTIRxOpaque`` - lowers remaining opaque constructs to lower-level TIRx forms diff --git a/include/tvm/relax/analysis.h b/include/tvm/relax/analysis.h index 092745f1718e..bbae775cab23 100644 --- a/include/tvm/relax/analysis.h +++ b/include/tvm/relax/analysis.h @@ -24,12 +24,12 @@ #ifndef TVM_RELAX_ANALYSIS_H_ #define TVM_RELAX_ANALYSIS_H_ -#include #include #include #include #include #include +#include #include #include @@ -55,7 +55,7 @@ namespace relax { * two shapes equals to each other during runtime. */ TVM_DLL bool CanProveShapeEqual(const ffi::Array& lhs, const ffi::Array& rhs, - const arith::Analyzer& ana); + const sym::Analyzer& ana); /*! * \brief Can prove the two symbolic shape expressions equals to each other. @@ -68,7 +68,7 @@ TVM_DLL bool CanProveShapeEqual(const ffi::Array& lhs, const ffi::Arra * if result is false, there is still possibility that * two shapes equals to each other during runtime. */ -TVM_DLL bool CanProveShapeEqual(const Expr& lhs, const Expr& rhs, const arith::Analyzer& ana); +TVM_DLL bool CanProveShapeEqual(const Expr& lhs, const Expr& rhs, const sym::Analyzer& ana); //----------------------------------- // Foundational Type analysis @@ -106,7 +106,7 @@ TVM_DLL Type DeriveCallRetType(const FuncType& finfo, const Call& call, const Bl * \return The derived type of the call. */ TVM_DLL Type DeriveCallRetType(const FuncType& finfo, const Call& call, const BlockBuilder& ctx, - const arith::Analyzer& ana); + const sym::Analyzer& ana); /*! * \brief Erase the info to a corresponding more coarse grained @@ -173,7 +173,7 @@ TVM_DLL Type EraseToWellDefined( */ TVM_DLL Type EraseToWellDefined(const Type& info, std::function(const Var& var)> f_var_map, - const arith::Analyzer& ana); + const sym::Analyzer& ana); /*! * \brief EraseToWellDefined variant with map. @@ -195,7 +195,7 @@ TVM_DLL Type EraseToWellDefined(const Type& info, ffi::Map var_map); * \return the corresponding erased type. */ TVM_DLL Type EraseToWellDefined(const Type& info, ffi::Map var_map, - const arith::Analyzer& ana); + const sym::Analyzer& ana); /*! * \brief Fine grained result of base check. @@ -266,7 +266,7 @@ TVM_DLL BaseCheckResult TypeBaseCheck(const Type& base, const Type& derived); * \sa BaseCheckResult */ TVM_DLL BaseCheckResult TypeBaseCheck(const Type& base, const Type& derived, - const arith::Analyzer& ana); + const sym::Analyzer& ana); /*! * \brief Check the relation of two type to see if one subsumes another one. @@ -283,7 +283,7 @@ TVM_DLL bool IsBaseOf(const Type& base, const Type& derived); * \param ana Context analyzer to prove symbolic expression equality. * \return Whether the relation holds. */ -TVM_DLL bool IsBaseOf(const Type& base, const Type& derived, const arith::Analyzer& ana); +TVM_DLL bool IsBaseOf(const Type& base, const Type& derived, const sym::Analyzer& ana); /*! * \brief Return the condition for which base is a superset of derived @@ -322,7 +322,7 @@ TVM_DLL Type TypeLCA(const Type& lhs, const Type& rhs); * \param ana Context analyzer to prove symbolic expression equality. * \return The unified information. */ -TVM_DLL Type TypeLCA(const Type& lhs, const Type& rhs, const arith::Analyzer& ana); +TVM_DLL Type TypeLCA(const Type& lhs, const Type& rhs, const sym::Analyzer& ana); /*! * \brief Get the TIR variables that appear in the input type. diff --git a/include/tvm/relax/block_builder.h b/include/tvm/relax/block_builder.h index ac20e2ad8f9a..77813588625a 100644 --- a/include/tvm/relax/block_builder.h +++ b/include/tvm/relax/block_builder.h @@ -24,11 +24,11 @@ #ifndef TVM_RELAX_BLOCK_BUILDER_H_ #define TVM_RELAX_BLOCK_BUILDER_H_ -#include #include #include #include #include +#include namespace tvm { namespace relax { @@ -248,7 +248,7 @@ class BlockBuilderNode : public ffi::Object { * \brief Get the analyzer of the BlockBuilder. * \return The BlockBuilder's arithmetic analyzer. */ - virtual arith::Analyzer GetAnalyzer() = 0; + virtual sym::Analyzer GetAnalyzer() = 0; static constexpr const bool _type_mutable = true; TVM_FFI_DECLARE_OBJECT_INFO("relax.BlockBuilder", BlockBuilderNode, ffi::Object); diff --git a/include/tvm/relax/dataflow_pattern.h b/include/tvm/relax/dataflow_pattern.h index 0599506cce1f..d3893fe49e9e 100644 --- a/include/tvm/relax/dataflow_pattern.h +++ b/include/tvm/relax/dataflow_pattern.h @@ -43,10 +43,10 @@ namespace tvm { -namespace arith { +namespace sym { class AnalyzerObj; class Analyzer; -} // namespace arith +} // namespace sym namespace relax { diff --git a/include/tvm/relax/distributed/axis_group_graph.h b/include/tvm/relax/distributed/axis_group_graph.h index 9590394297f7..eb3570de142a 100644 --- a/include/tvm/relax/distributed/axis_group_graph.h +++ b/include/tvm/relax/distributed/axis_group_graph.h @@ -20,11 +20,11 @@ #ifndef TVM_RELAX_DISTRIBUTED_AXIS_GROUP_GRAPH_H_ #define TVM_RELAX_DISTRIBUTED_AXIS_GROUP_GRAPH_H_ -#include #include #include #include #include +#include #include #include @@ -60,7 +60,7 @@ class BufferAxisHash { * \return The iter var whose extent to be changed */ Var GetShardingVarFromIndex(PrimExpr index, ffi::Map var_range, - const arith::Analyzer& analyzer); + const sym::Analyzer& analyzer); /*! * \brief Construct an axis group graph from a PrimFunc. Two buffer axis are connected if they @@ -134,7 +134,7 @@ class BufferAxisGraphExtractor : public s_tir::StmtExprVisitor { } bool Match(PrimExpr a, PrimExpr buffer_shape_a, PrimExpr b, PrimExpr buffer_shape_b, - const arith::Analyzer& analyzer) { + const sym::Analyzer& analyzer) { if (b.as()) { std::swap(a, b); std::swap(buffer_shape_a, buffer_shape_b); @@ -147,7 +147,7 @@ class BufferAxisGraphExtractor : public s_tir::StmtExprVisitor { analyzer->Bind(iter_var_range_); b = analyzer->Simplify(b); // index var `a` must access whole range of a specific buffer dimension - arith::IntSet intset_b = arith::EvalSet(b, arith::AsIntSet(iter_var_range_)); + sym::IntSet intset_b = sym::EvalSet(b, sym::AsIntSet(iter_var_range_)); if (!analyzer->CanProveEqual(buffer_shape_a, iter_var_range_[var]->extent) || !intset_b.MatchRange(Range::FromMinExtent(0, buffer_shape_b))) { return false; @@ -169,7 +169,7 @@ class BufferAxisGraphExtractor : public s_tir::StmtExprVisitor { for (const auto& iter_var : op->iter_vars) { iter_var_range_.Set(iter_var->var, iter_var->dom); } - arith::Analyzer analyzer; + sym::Analyzer analyzer; for (const auto& access_pr : buffer_access_indices_) { BufferVar buffer = access_pr.first; ffi::Array indices = access_pr.second; diff --git a/include/tvm/relax/utils.h b/include/tvm/relax/utils.h index d081d991668e..e764f6097f66 100644 --- a/include/tvm/relax/utils.h +++ b/include/tvm/relax/utils.h @@ -24,10 +24,10 @@ #ifndef TVM_RELAX_UTILS_H_ #define TVM_RELAX_UTILS_H_ -#include #include #include #include +#include namespace tvm { namespace relax { @@ -71,7 +71,7 @@ TVM_DLL Type Bind(const Type& ty, const tvm::ffi::Map& binds); * \return The input binding map augmented with inferred symbolic bindings. */ TVM_DLL tvm::ffi::Map InferSymbolicVarMap( - const tvm::ffi::Map& binds, const arith::Analyzer& analyzer); + const tvm::ffi::Map& binds, const sym::Analyzer& analyzer); /*! * \brief Check if the given Type is for a boolean scalar (tensor of rank 0 with a boolean diff --git a/include/tvm/s_tir/analysis.h b/include/tvm/s_tir/analysis.h index acaf2a4ba676..9b7dea02e62c 100644 --- a/include/tvm/s_tir/analysis.h +++ b/include/tvm/s_tir/analysis.h @@ -90,10 +90,10 @@ const s_tir::SBlockNode* FindAnchorBlock(const IRModule& mod); } // namespace tirx -namespace arith { +namespace sym { class AnalyzerObj; class Analyzer; -} // namespace arith +} // namespace sym namespace s_tir { using namespace tvm::tirx; @@ -144,8 +144,7 @@ struct MemCpyDetails { * \param analyzer The analyzer with which to check any algebraic expressions * \returns The source and destination regions being copied, if the loop is equivalent to memcpy. */ -TVM_DLL std::optional IdentifyMemCpy(const For& loop, - const arith::Analyzer& analyzer); +TVM_DLL std::optional IdentifyMemCpy(const For& loop, const sym::Analyzer& analyzer); /*! * \brief Infer the domain touched by buffer accesses within a statement. diff --git a/include/tvm/arith/analyzer.h b/include/tvm/sym/analyzer.h similarity index 98% rename from include/tvm/arith/analyzer.h rename to include/tvm/sym/analyzer.h index 9b2de84d3a17..c1dd52f319af 100644 --- a/include/tvm/arith/analyzer.h +++ b/include/tvm/sym/analyzer.h @@ -18,19 +18,19 @@ */ /*! - * \file tvm/arith/analyzer.h + * \file tvm/sym/analyzer.h * \brief Algebra expression simplifications. */ -#ifndef TVM_ARITH_ANALYZER_H_ -#define TVM_ARITH_ANALYZER_H_ +#ifndef TVM_SYM_ANALYZER_H_ +#define TVM_SYM_ANALYZER_H_ -#include #include #include #include #include #include #include +#include #include #include @@ -39,8 +39,8 @@ #include namespace tvm { -/*! \brief namespace of arithmetic analysis. */ -namespace arith { +/*! \brief namespace of symbolic analysis. */ +namespace sym { //------------------------------------------------------- // Base integer analysis API. // @@ -68,7 +68,7 @@ enum DivMode { * \brief The strength used in top-level condition proves * \note The higher, the more time consuming it can be. * - * Do not use level beyond kDefault in internal recursive rewriting in arith + * Do not use level beyond kDefault in internal recursive rewriting in sym * analysis and only use it at top-level simplification to avoid speed issues. */ enum class ProofStrength : int { @@ -107,7 +107,7 @@ class ConstIntBoundNode : public ffi::Object { static const constexpr int64_t kNegInf = -kPosInf; static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode; - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("arith.ConstIntBound", ConstIntBoundNode, ffi::Object); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("sym.ConstIntBound", ConstIntBoundNode, ffi::Object); }; /*! @@ -221,7 +221,7 @@ class ModularSetNode : public ffi::Object { } static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode; - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("arith.ModularSet", ModularSetNode, ffi::Object); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("sym.ModularSet", ModularSetNode, ffi::Object); }; /*! @@ -911,7 +911,7 @@ class TVM_DLL AnalyzerObj : public ffi::Object { * not make the underlying AnalyzerObj immutable. */ static constexpr bool _type_mutable = true; - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("arith.Analyzer", AnalyzerObj, ffi::Object); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("sym.Analyzer", AnalyzerObj, ffi::Object); }; /*! @@ -943,9 +943,9 @@ class Analyzer : public ffi::ObjectRef { * \code * * Var x("x"); - * arith::Analyzer analyzer; + * sym::Analyzer analyzer; * { - * With scope(analyzer, tvm::floormod(x, 3) == 0); + * With scope(analyzer, tvm::floormod(x, 3) == 0); * TVM_FFI_ICHECK_EQ(analyzer->modular_set(x)->coeff, 3); * } * // constraint no longer in effect. @@ -1005,6 +1005,6 @@ class ConstraintContext { bool is_assume_; }; -} // namespace arith +} // namespace sym } // namespace tvm -#endif // TVM_ARITH_ANALYZER_H_ +#endif // TVM_SYM_ANALYZER_H_ diff --git a/include/tvm/arith/bound.h b/include/tvm/sym/bound.h similarity index 92% rename from include/tvm/arith/bound.h rename to include/tvm/sym/bound.h index 6004dff53533..2c880a7152a2 100644 --- a/include/tvm/arith/bound.h +++ b/include/tvm/sym/bound.h @@ -17,20 +17,20 @@ * under the License. */ /*! - * \file tvm/arith/bound.h + * \file tvm/sym/bound.h * \brief Bound deducers. */ -#ifndef TVM_ARITH_BOUND_H_ -#define TVM_ARITH_BOUND_H_ +#ifndef TVM_SYM_BOUND_H_ +#define TVM_SYM_BOUND_H_ -#include #include #include +#include #include namespace tvm { -namespace arith { +namespace sym { /*! * \brief Deduce the bound of the target variable in a expression, @@ -63,6 +63,6 @@ IntSet DeduceBound(PrimExpr v, PrimExpr cond, const std::unordered_map& hint_map, const std::unordered_map& relax_map); -} // namespace arith +} // namespace sym } // namespace tvm -#endif // TVM_ARITH_BOUND_H_ +#endif // TVM_SYM_BOUND_H_ diff --git a/include/tvm/arith/int_set.h b/include/tvm/sym/int_set.h similarity index 94% rename from include/tvm/arith/int_set.h rename to include/tvm/sym/int_set.h index cdb5c9af9e1f..197f7b69e5db 100644 --- a/include/tvm/arith/int_set.h +++ b/include/tvm/sym/int_set.h @@ -18,11 +18,11 @@ */ /*! - * \file tvm/arith/int_set.h + * \file tvm/sym/int_set.h * \brief Integer set */ -#ifndef TVM_ARITH_INT_SET_H_ -#define TVM_ARITH_INT_SET_H_ +#ifndef TVM_SYM_INT_SET_H_ +#define TVM_SYM_INT_SET_H_ #include #include @@ -30,7 +30,7 @@ #include namespace tvm { -namespace arith { +namespace sym { class AnalyzerObj; class Analyzer; @@ -267,7 +267,7 @@ IntSet Intersect(const ffi::Array& sets); * \param var_dom The ranges of variables * \return The integer sets of the variables */ -ffi::Map AsIntSet(const ffi::Map& var_dom); +ffi::Map AsIntSet(const ffi::Map& var_dom); /*! * \brief Analyze the region with affine map, given the domain of variables and their predicate. @@ -276,12 +276,12 @@ ffi::Map AsIntSet(const ffi::Map& var_dom); * \param var_dom The ranges of the variables * \param predicate The predicate for the affine map * \param analyzer The analyzer used - * \return std::nullopt if the detection fails, or an array of arith::IntSet as the result of + * \return std::nullopt if the detection fails, or an array of sym::IntSet as the result of * analysis */ TVM_DLL ffi::Optional> EstimateRegionStrictBound( const ffi::Array& region, const ffi::Map& var_dom, const PrimExpr& predicate, - const arith::Analyzer& analyzer); + const sym::Analyzer& analyzer); /*! * \brief Analyze the region with affine map, given the domain of variables and their predicate. @@ -290,12 +290,12 @@ TVM_DLL ffi::Optional> EstimateRegionStrictBound( * \param var_dom The ranges of the variables * \param predicate The predicate for the affine map * \param analyzer The analyzer used - * \return std::nullopt if the detection fails, or an array of arith::IntSet as the result of + * \return std::nullopt if the detection fails, or an array of sym::IntSet as the result of * analysis */ TVM_DLL ffi::Optional> EstimateRegionLowerBound( const ffi::Array& region, const ffi::Map& var_dom, const PrimExpr& predicate, - const arith::Analyzer& analyzer); + const sym::Analyzer& analyzer); /*! * \brief Analyze the region with affine map, given the domain of variables and their predicate @@ -305,13 +305,13 @@ TVM_DLL ffi::Optional> EstimateRegionLowerBound( * \param var_dom The ranges of the variables * \param predicate The predicate for the affine map * \param analyzer The analyzer used - * \return an array of arith::IntSet as the result of analysis + * \return an array of sym::IntSet as the result of analysis */ TVM_DLL ffi::Array EstimateRegionUpperBound(const ffi::Array& region, const ffi::Map& var_dom, const PrimExpr& predicate, - const arith::Analyzer& analyzer); + const sym::Analyzer& analyzer); -} // namespace arith +} // namespace sym } // namespace tvm -#endif // TVM_ARITH_INT_SET_H_ +#endif // TVM_SYM_INT_SET_H_ diff --git a/include/tvm/arith/iter_affine_map.h b/include/tvm/sym/iter_affine_map.h similarity index 92% rename from include/tvm/arith/iter_affine_map.h rename to include/tvm/sym/iter_affine_map.h index f47492f497bf..17e951f49688 100644 --- a/include/tvm/arith/iter_affine_map.h +++ b/include/tvm/sym/iter_affine_map.h @@ -18,7 +18,7 @@ */ /*! - * \file tvm/arith/iter_affine_map.h + * \file tvm/sym/iter_affine_map.h * \brief Iterator quasi-affine mapping patterns. * * This file defines a collection of mapping patterns @@ -45,16 +45,16 @@ * while split corresponds to additional floordiv/mod operations * that can appear in quasi-affine transformations. */ -#ifndef TVM_ARITH_ITER_AFFINE_MAP_H_ -#define TVM_ARITH_ITER_AFFINE_MAP_H_ +#ifndef TVM_SYM_ITER_AFFINE_MAP_H_ +#define TVM_SYM_ITER_AFFINE_MAP_H_ -#include #include #include #include +#include namespace tvm { -namespace arith { +namespace sym { /*! * \brief Base class of all iter map expressions. @@ -66,7 +66,7 @@ namespace arith { class IterMapExprNode : public ExprNode { public: static constexpr const uint32_t _type_child_slots = 2; - TVM_FFI_DECLARE_OBJECT_INFO("arith.IterMapExpr", IterMapExprNode, ExprNode); + TVM_FFI_DECLARE_OBJECT_INFO("sym.IterMapExpr", IterMapExprNode, ExprNode); }; /*! @@ -105,7 +105,7 @@ class IterMarkNode : public ffi::Object { } static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindDAGNode; - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("arith.IterMark", IterMarkNode, ffi::Object); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("sym.IterMark", IterMarkNode, ffi::Object); }; /*! @@ -151,7 +151,7 @@ class IterSplitExprNode : public IterMapExprNode { } static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode; - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("arith.IterSplitExpr", IterSplitExprNode, IterMapExprNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("sym.IterSplitExpr", IterSplitExprNode, IterMapExprNode); }; /*! @@ -205,7 +205,7 @@ class IterSumExprNode : public IterMapExprNode { } static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode; - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("arith.IterSumExpr", IterSumExprNode, IterMapExprNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("sym.IterSumExpr", IterSumExprNode, IterMapExprNode); }; /*! @@ -225,16 +225,16 @@ class IterSumExpr : public IterMapExpr { TVM_DEFINE_OBJECT_REF_COW_METHOD(IterSumExprNode); }; -} // namespace arith +} // namespace sym namespace ffi { template <> -inline constexpr bool object_ref_contains_v = true; +inline constexpr bool object_ref_contains_v = true; template <> -inline constexpr bool object_ref_contains_v = true; +inline constexpr bool object_ref_contains_v = true; } // namespace ffi -namespace arith { +namespace sym { /*! \brief Mapping level for iterators. */ enum IterMapLevel { @@ -275,7 +275,7 @@ class IterMapResultNode : public ffi::Object { .def_ro("errors", &IterMapResultNode::errors) .def_ro("padding_predicate", &IterMapResultNode::padding_predicate); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("arith.IterMapResult", IterMapResultNode, ffi::Object); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("sym.IterMapResult", IterMapResultNode, ffi::Object); }; /*! @@ -316,7 +316,7 @@ class IterMapResult : public ffi::ObjectRef { */ IterMapResult DetectIterMap(const ffi::Array& indices, const ffi::Map& input_iters, const PrimExpr& predicate, - IterMapLevel check_level, const arith::Analyzer& analyzer, + IterMapLevel check_level, const sym::Analyzer& analyzer, bool simplify_trivial_iterators = true); /*! @@ -333,7 +333,7 @@ IterMapResult DetectIterMap(const ffi::Array& indices, ffi::Array IterMapSimplify(const ffi::Array& indices, const ffi::Map& input_iters, const PrimExpr& input_pred, IterMapLevel check_level, - const arith::Analyzer& analyzer, + const sym::Analyzer& analyzer, bool simplify_trivial_iterators = true); /*! @@ -390,7 +390,7 @@ ffi::Array> SubspaceDivide(const ffi::Array& bind const ffi::Map& input_iters, const ffi::Array& sub_iters, const PrimExpr& predicate, IterMapLevel check_level, - const arith::Analyzer& analyzer, + const sym::Analyzer& analyzer, bool simplify_trivial_iterators = true); /*! @@ -417,8 +417,8 @@ PrimExpr NormalizeIterMapToExpr(const PrimExpr& expr); * \note This function is useful to detect iterator stride patterns. */ IterSumExpr NormalizeToIterSum(PrimExpr index, const ffi::Map& input_iters, - const arith::Analyzer& analyzer); + const sym::Analyzer& analyzer); -} // namespace arith +} // namespace sym } // namespace tvm -#endif // TVM_ARITH_ITER_AFFINE_MAP_H_ +#endif // TVM_SYM_ITER_AFFINE_MAP_H_ diff --git a/include/tvm/arith/pattern.h b/include/tvm/sym/pattern.h similarity index 91% rename from include/tvm/arith/pattern.h rename to include/tvm/sym/pattern.h index dd48abb763ea..f6258442e66b 100644 --- a/include/tvm/arith/pattern.h +++ b/include/tvm/sym/pattern.h @@ -18,16 +18,16 @@ */ /*! - * \file tvm/arith/pattern.h + * \file tvm/sym/pattern.h * \brief Expression pattern detectors. */ -#ifndef TVM_ARITH_PATTERN_H_ -#define TVM_ARITH_PATTERN_H_ +#ifndef TVM_SYM_PATTERN_H_ +#define TVM_SYM_PATTERN_H_ #include namespace tvm { -namespace arith { +namespace sym { /*! * \brief Detect if e can be rewritten as e = sum_{i=0}^{n-1} var[i] * coeff[i] + coeff[n] * Where coeff[i] and base are invariant of var[j] for all i and j. @@ -48,6 +48,6 @@ ffi::Array DetectLinearEquation(const PrimExpr& e, const ffi::Array DetectClipBound(const PrimExpr& e, const ffi::Array& vars); -} // namespace arith +} // namespace sym } // namespace tvm -#endif // TVM_ARITH_PATTERN_H_ +#endif // TVM_SYM_PATTERN_H_ diff --git a/include/tvm/te/operation.h b/include/tvm/te/operation.h index 42800875efdb..c59ec678273f 100644 --- a/include/tvm/te/operation.h +++ b/include/tvm/te/operation.h @@ -24,10 +24,10 @@ #ifndef TVM_TE_OPERATION_H_ #define TVM_TE_OPERATION_H_ -#include #include #include #include +#include #include #include #include diff --git a/include/tvm/te/tensor.h b/include/tvm/te/tensor.h index da66e034a6b5..4dfa5cdc3bfa 100644 --- a/include/tvm/te/tensor.h +++ b/include/tvm/te/tensor.h @@ -24,9 +24,9 @@ #ifndef TVM_TE_TENSOR_H_ #define TVM_TE_TENSOR_H_ -#include #include #include +#include #include #include @@ -36,7 +36,7 @@ namespace tvm { namespace te { -using arith::IntSet; +using sym::IntSet; using namespace tvm::tirx; // internal node container for Operation diff --git a/include/tvm/tirx/index_map.h b/include/tvm/tirx/index_map.h index 0214949fce60..1c420c480269 100644 --- a/include/tvm/tirx/index_map.h +++ b/include/tvm/tirx/index_map.h @@ -34,9 +34,9 @@ #include namespace tvm { -namespace arith { +namespace sym { class Analyzer; -} // namespace arith +} // namespace sym } // namespace tvm namespace tvm { @@ -109,7 +109,7 @@ class IndexMapNode : public ffi::Object { * each expression in `final_indices`. */ ffi::Array MapIndices(const ffi::Array& indices, - const arith::Analyzer& analyzer) const; + const sym::Analyzer& analyzer) const; /*! \brief Map a memory range to the output space using a fresh analyzer. * @@ -137,8 +137,7 @@ class IndexMapNode : public ffi::Object { * \returns The ranges in the output space. Contains one value for * each expression in `final_indices`. */ - ffi::Array MapRanges(const ffi::Array& ranges, - const arith::Analyzer& analyzer) const; + ffi::Array MapRanges(const ffi::Array& ranges, const sym::Analyzer& analyzer) const; /*! \brief Map a buffer shape to the output space using a fresh analyzer. * @@ -157,7 +156,7 @@ class IndexMapNode : public ffi::Object { * value for each expression in `final_indices`. */ ffi::Array MapShape(const ffi::Array& shape, - const arith::Analyzer& analyzer) const; + const sym::Analyzer& analyzer) const; /* \brief Map an Tensor according to this index map * @@ -231,7 +230,7 @@ class IndexMap : public ffi::ObjectRef { * \param analyzer An analyzer to be used while deriving and validating * the inverse. */ - IndexMap Inverse(ffi::Array initial_ranges, const arith::Analyzer& analyzer) const; + IndexMap Inverse(ffi::Array initial_ranges, const sym::Analyzer& analyzer) const; /*! \brief Rename the variables in the index map and ensure the names are unique. * @@ -268,7 +267,7 @@ class IndexMap : public ffi::ObjectRef { * which the inverse maps to a valid range. */ std::pair NonSurjectiveInverse(ffi::Array initial_ranges, - const arith::Analyzer& analyzer) const; + const sym::Analyzer& analyzer) const; TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(IndexMap, ffi::ObjectRef, IndexMapNode); }; diff --git a/include/tvm/topi/detail/constant_utils.h b/include/tvm/topi/detail/constant_utils.h index f56437de0f23..51f9e910ce9f 100644 --- a/include/tvm/topi/detail/constant_utils.h +++ b/include/tvm/topi/detail/constant_utils.h @@ -24,9 +24,9 @@ #ifndef TVM_TOPI_DETAIL_CONSTANT_UTILS_H_ #define TVM_TOPI_DETAIL_CONSTANT_UTILS_H_ -#include #include #include +#include #include #include @@ -132,7 +132,7 @@ inline bool EqualCheck(PrimExpr lhs, PrimExpr rhs) { tvm::prim::ExprDeepEqual expr_equal; bool result = expr_equal(lhs, rhs); if (!result) { - PrimExpr t = tvm::arith::Analyzer()->Simplify(lhs - rhs); + PrimExpr t = tvm::sym::Analyzer()->Simplify(lhs - rhs); if (const IntImmNode* i = t.as()) { result = i->value == 0; } diff --git a/include/tvm/topi/nn.h b/include/tvm/topi/nn.h index 5945960f0481..6289ebecbf8f 100644 --- a/include/tvm/topi/nn.h +++ b/include/tvm/topi/nn.h @@ -24,8 +24,8 @@ #ifndef TVM_TOPI_NN_H_ #define TVM_TOPI_NN_H_ -#include #include +#include #include #include #include @@ -164,7 +164,7 @@ inline tvm::te::Tensor pad( } } - arith::Analyzer analyzer; + sym::Analyzer analyzer; TVM_FFI_ICHECK_GE(pad_before.size(), 1); TVM_FFI_ICHECK_EQ(pad_before.size(), pad_after.size()); tvm::ffi::Array pad_before_int32; diff --git a/include/tvm/topi/nn/bnn.h b/include/tvm/topi/nn/bnn.h index f114c4dc9fb4..9b532fe963d1 100644 --- a/include/tvm/topi/nn/bnn.h +++ b/include/tvm/topi/nn/bnn.h @@ -24,7 +24,7 @@ #ifndef TVM_TOPI_NN_BNN_H_ #define TVM_TOPI_NN_BNN_H_ -#include +#include #include #include #include @@ -55,7 +55,7 @@ inline tvm::te::Tensor binarize_pack(const tvm::te::Tensor& data, int axis, TVM_FFI_ICHECK_EQ(GetConstInt(ishape[axis]) % 32, 0) << "binarize_pack: axis size must be a multiple of 32"; - arith::Analyzer analyzer; + sym::Analyzer analyzer; auto n = ishape.size(); ffi::Array oshape; for (size_t i = 0; i < n; ++i) { diff --git a/include/tvm/topi/nn/dilate.h b/include/tvm/topi/nn/dilate.h index a0cfe7f26bf1..c7e57fc337fd 100644 --- a/include/tvm/topi/nn/dilate.h +++ b/include/tvm/topi/nn/dilate.h @@ -24,7 +24,7 @@ #ifndef TVM_TOPI_NN_DILATE_H_ #define TVM_TOPI_NN_DILATE_H_ -#include +#include #include #include @@ -75,7 +75,7 @@ inline Tensor dilate(const Tensor& x, ffi::Array strides, double dilat << "strides size (" << strides.size() << ") must match dimension of x (" << n << ")"; ffi::Array out_shape; - arith::Analyzer analyzer; + sym::Analyzer analyzer; for (size_t i = 0; i < n; ++i) { out_shape.push_back(analyzer->Simplify((x->shape[i] - 1) * (strides[i] + 1))); } diff --git a/include/tvm/topi/nn/pooling.h b/include/tvm/topi/nn/pooling.h index 69de16660eef..829b0901c51d 100644 --- a/include/tvm/topi/nn/pooling.h +++ b/include/tvm/topi/nn/pooling.h @@ -24,7 +24,7 @@ #ifndef TVM_TOPI_NN_POOLING_H_ #define TVM_TOPI_NN_POOLING_H_ -#include +#include #include #include #include @@ -86,7 +86,7 @@ inline Tensor pool_grad_impl(const Tensor& out_grad, const Tensor& x, ffi::Array pad_after(std::vector(x->shape.size(), 0)); pad_after.Set(height_axis, pad_bottom); pad_after.Set(width_axis, pad_right); - arith::Analyzer analyzer; + sym::Analyzer analyzer; auto out_height = analyzer->Simplify((height - kernel_height + pad_top + pad_bottom) / stride_height + 1); auto out_width = @@ -572,7 +572,7 @@ inline Tensor pool_impl_nd(const Tensor& x, const ffi::Array& kernel_s pad_before.Set(ii, pad_head[i]); pad_after.Set(ii, pad_tail[i]); - arith::Analyzer analyzer; + sym::Analyzer analyzer; PrimExpr numerator = data_shape[ii] - (kernel[i] - 1) * dilation[i] - 1 + pad_head[i] + pad_tail[i]; diff --git a/include/tvm/topi/transform.h b/include/tvm/topi/transform.h index e90a24db8388..24aad7495709 100644 --- a/include/tvm/topi/transform.h +++ b/include/tvm/topi/transform.h @@ -24,9 +24,9 @@ #ifndef TVM_TOPI_TRANSFORM_H_ #define TVM_TOPI_TRANSFORM_H_ -#include #include #include +#include #include #include #include @@ -492,7 +492,7 @@ inline Tensor concatenate(const ffi::Array& inputs, int axis = 0, for (auto t : inputs) { axis_sizes.push_back(t->shape[axis]); } - arith::Analyzer analyzer; + sym::Analyzer analyzer; PrimExpr join_size = axis_sizes[0]; for (size_t i = 1; i < axis_sizes.size(); ++i) { join_size += axis_sizes[i]; @@ -725,7 +725,7 @@ inline te::Tensor dynamic_strided_slice_with_axes( TVM_FFI_ICHECK_LT(axis, src_tensor_dim); } - arith::Analyzer analyzer; + sym::Analyzer analyzer; ffi::Array out_shape = x->shape; for (size_t i = 0; i < begin.size(); i++) { @@ -785,7 +785,7 @@ inline Tensor dynamic_strided_slice(const Tensor& x, const ffi::Array& const size_t num_slice_axes = begin.size(); ffi::Array out_shape; - arith::Analyzer analyzer; + sym::Analyzer analyzer; for (size_t i = 0; i < num_slice_axes; ++i) { // Dynamic scalar tensor loads cannot be simplified while inferring shape. if (!te::IsTensorLoad(begin[i]) && !te::IsTensorLoad(end[i]) && !te::IsTensorLoad(strides[i])) { @@ -1762,7 +1762,7 @@ inline Tensor tensordot(const Tensor& A, const tvm::te::Tensor& B, ffi::Array iter_domain; iter_domain.reserve(src->shape.size()); for (const PrimExpr& e : src->shape) { diff --git a/python/tvm/__init__.py b/python/tvm/__init__.py index 16cef942ad68..ad21df744c90 100644 --- a/python/tvm/__init__.py +++ b/python/tvm/__init__.py @@ -65,7 +65,7 @@ from .driver import build, compile # others -from . import arith +from . import sym # support infra from . import support diff --git a/python/tvm/backend/cuda/tile_primitive/common.py b/python/tvm/backend/cuda/tile_primitive/common.py index 25af4e545300..a74481242b43 100644 --- a/python/tvm/backend/cuda/tile_primitive/common.py +++ b/python/tvm/backend/cuda/tile_primitive/common.py @@ -22,10 +22,10 @@ import re from enum import Enum -from tvm.arith.analyzer import Analyzer from tvm.ir import TensorRegion from tvm.runtime import DataType from tvm.script import tirx as T +from tvm.sym.analyzer import Analyzer from tvm.tirx import Buffer, PrimFunc from tvm.tirx.operator.tile_primitive import DispatchContext, fail from tvm.tirx.tile_primitive import TilePrimitiveCall diff --git a/python/tvm/backend/cuda/tile_primitive/copy/_common.py b/python/tvm/backend/cuda/tile_primitive/copy/_common.py index b7697716174e..42dae4debca0 100644 --- a/python/tvm/backend/cuda/tile_primitive/copy/_common.py +++ b/python/tvm/backend/cuda/tile_primitive/copy/_common.py @@ -25,7 +25,7 @@ to call, allowed vec widths). All the layout/partition logic lives here. """ -from tvm import arith +from tvm import sym from tvm.tirx.layout import ComposeLayout, Iter, TileLayout from tvm.tirx.operator.tile_primitive.registry import DispatchContext @@ -34,7 +34,7 @@ def _alignment_ok(vec_len: int, terms) -> bool: """Every term must be a multiple of ``vec_len``. Constants checked - directly; Expr / symbolic terms checked via ``arith.Analyzer``. + directly; Expr / symbolic terms checked via ``sym.Analyzer``. ``vec_len=1`` always passes (the scalar fallback). When a symbolic term can't be proved divisible, returns ``False`` conservatively — @@ -42,7 +42,7 @@ def _alignment_ok(vec_len: int, terms) -> bool: """ if vec_len <= 1: return True - analyzer = arith.Analyzer() + analyzer = sym.Analyzer() for t in terms: if isinstance(t, int): if t % vec_len != 0: @@ -322,7 +322,7 @@ def _extract_tile(layout, region): # Region bounds may be constant-valued but remain as unfolded expressions # after substitution. Simplify before converting to a Python integer; # genuinely symbolic extents fall back inside ``strip_swizzle_to_tile``. - analyzer = arith.Analyzer() + analyzer = sym.Analyzer() return strip_swizzle_to_tile( layout, lambda: [int(analyzer.simplify(end - start)) for (start, end) in region] ) diff --git a/python/tvm/backend/cuda/tile_primitive/copy/vec_auto_reg.py b/python/tvm/backend/cuda/tile_primitive/copy/vec_auto_reg.py index 5648045be054..4c35f346a303 100644 --- a/python/tvm/backend/cuda/tile_primitive/copy/vec_auto_reg.py +++ b/python/tvm/backend/cuda/tile_primitive/copy/vec_auto_reg.py @@ -30,9 +30,9 @@ import tvm_ffi import tvm -from tvm.arith import Analyzer, ConstIntBound from tvm.runtime import DataType from tvm.script import tirx as T +from tvm.sym import Analyzer, ConstIntBound from tvm.tirx import Buffer, PrimFunc from tvm.tirx import Var as _TirVar from tvm.tirx.expr import IntImm as _IntImm diff --git a/python/tvm/backend/cuda/tile_primitive/copy/vec_forced.py b/python/tvm/backend/cuda/tile_primitive/copy/vec_forced.py index fcb9f2f794d5..a0e43abdca1f 100644 --- a/python/tvm/backend/cuda/tile_primitive/copy/vec_forced.py +++ b/python/tvm/backend/cuda/tile_primitive/copy/vec_forced.py @@ -32,10 +32,10 @@ cache="nc", l1_evict="L1::no_allocate", prefetch_size="L2::256B") """ -from tvm.arith.analyzer import Analyzer from tvm.ir import TensorRegion from tvm.runtime import DataType from tvm.script import tirx as T +from tvm.sym.analyzer import Analyzer from tvm.tirx import Buffer, PrimFunc from tvm.tirx.operator.tile_primitive.dispatcher import predicate, register_dispatch from tvm.tirx.operator.tile_primitive.registry import DispatchContext diff --git a/python/tvm/backend/cuda/tile_primitive/copy_async/tcgen05_cp.py b/python/tvm/backend/cuda/tile_primitive/copy_async/tcgen05_cp.py index 578e2dcf41c5..c9edef825dce 100644 --- a/python/tvm/backend/cuda/tile_primitive/copy_async/tcgen05_cp.py +++ b/python/tvm/backend/cuda/tile_primitive/copy_async/tcgen05_cp.py @@ -109,9 +109,9 @@ import operator import tvm -from tvm.arith import Analyzer from tvm.runtime import DataType from tvm.script import tirx as T +from tvm.sym import Analyzer from tvm.tirx import Buffer, PrimFunc from tvm.tirx.layout import ComposeLayout, TCol, TileLayout, TLane from tvm.tirx.layout import m as m_axis diff --git a/python/tvm/backend/cuda/tile_primitive/copy_async/tcgen05_ldst.py b/python/tvm/backend/cuda/tile_primitive/copy_async/tcgen05_ldst.py index ff23c11806e5..c9909c663748 100644 --- a/python/tvm/backend/cuda/tile_primitive/copy_async/tcgen05_ldst.py +++ b/python/tvm/backend/cuda/tile_primitive/copy_async/tcgen05_ldst.py @@ -23,9 +23,9 @@ """ import tvm -from tvm.arith import Analyzer from tvm.runtime import DataType from tvm.script import tirx as T +from tvm.sym import Analyzer from tvm.tirx import Buffer, PrimFunc from tvm.tirx.layout import ( S, diff --git a/python/tvm/backend/cuda/tile_primitive/copy_async/tma.py b/python/tvm/backend/cuda/tile_primitive/copy_async/tma.py index e5b5aa677c75..1f75023fb297 100644 --- a/python/tvm/backend/cuda/tile_primitive/copy_async/tma.py +++ b/python/tvm/backend/cuda/tile_primitive/copy_async/tma.py @@ -33,8 +33,8 @@ from itertools import pairwise import tvm -from tvm.arith import Analyzer from tvm.script import tirx as T +from tvm.sym import Analyzer from tvm.tirx import Buffer, IntImm, PrimFunc, is_buffer_var from tvm.tirx.layout import Layout, TileLayout from tvm.tirx.operator.tile_primitive import ( diff --git a/python/tvm/backend/cuda/tile_primitive/copy_async/utils.py b/python/tvm/backend/cuda/tile_primitive/copy_async/utils.py index d6d4cf754d78..4bac26810594 100644 --- a/python/tvm/backend/cuda/tile_primitive/copy_async/utils.py +++ b/python/tvm/backend/cuda/tile_primitive/copy_async/utils.py @@ -21,7 +21,7 @@ layout helpers other variants (e.g. ``dsmem.py``) still import. """ -from tvm.arith import Analyzer +from tvm.sym import Analyzer from tvm.tirx.layout import Layout, TileLayout from ..layout_utils import strip_swizzle_to_tile diff --git a/python/tvm/backend/cuda/tile_primitive/elementwise/_common.py b/python/tvm/backend/cuda/tile_primitive/elementwise/_common.py index 5e668ffddb19..90aea6abc51c 100644 --- a/python/tvm/backend/cuda/tile_primitive/elementwise/_common.py +++ b/python/tvm/backend/cuda/tile_primitive/elementwise/_common.py @@ -30,10 +30,10 @@ import functools import operator -from tvm.arith.analyzer import Analyzer from tvm.ir import TensorRegion from tvm.runtime import DataType from tvm.script import tirx as T +from tvm.sym.analyzer import Analyzer from tvm.tirx.layout import Axis, Iter, TileLayout from ..common import get_indices, get_st_extent diff --git a/python/tvm/backend/cuda/tile_primitive/elementwise/reg.py b/python/tvm/backend/cuda/tile_primitive/elementwise/reg.py index cec6b6831c59..f4bbf688a471 100644 --- a/python/tvm/backend/cuda/tile_primitive/elementwise/reg.py +++ b/python/tvm/backend/cuda/tile_primitive/elementwise/reg.py @@ -34,8 +34,8 @@ import functools import operator -from tvm.arith import Analyzer from tvm.script import tirx as T +from tvm.sym import Analyzer from tvm.tirx import PrimFunc, TilePrimitiveCall from tvm.tirx.layout import TileLayout from tvm.tirx.operator.tile_primitive import DispatchContext diff --git a/python/tvm/backend/cuda/tile_primitive/gemm/mma_m16n8k_.py b/python/tvm/backend/cuda/tile_primitive/gemm/mma_m16n8k_.py index c76a6f41b740..c9794ef9020d 100644 --- a/python/tvm/backend/cuda/tile_primitive/gemm/mma_m16n8k_.py +++ b/python/tvm/backend/cuda/tile_primitive/gemm/mma_m16n8k_.py @@ -19,8 +19,8 @@ from dataclasses import dataclass -from tvm.arith.analyzer import Analyzer from tvm.script import tirx as T +from tvm.sym.analyzer import Analyzer from tvm.tirx import PrimFunc from tvm.tirx.layout import TileLayout from tvm.tirx.operator.tile_primitive import ( diff --git a/python/tvm/backend/cuda/tile_primitive/gemm_async/tcgen05.py b/python/tvm/backend/cuda/tile_primitive/gemm_async/tcgen05.py index 16a3c739adf8..b4c489191fef 100644 --- a/python/tvm/backend/cuda/tile_primitive/gemm_async/tcgen05.py +++ b/python/tvm/backend/cuda/tile_primitive/gemm_async/tcgen05.py @@ -26,10 +26,10 @@ import operator import tvm -from tvm.arith.analyzer import Analyzer from tvm.ir import TensorRegion from tvm.runtime import DataType from tvm.script import tirx as T +from tvm.sym.analyzer import Analyzer from tvm.tirx import PrimFunc from tvm.tirx import op as tirx_op from tvm.tirx.layout import ( diff --git a/python/tvm/backend/cuda/tile_primitive/gemm_utils.py b/python/tvm/backend/cuda/tile_primitive/gemm_utils.py index a7ada173d340..981ab6652661 100644 --- a/python/tvm/backend/cuda/tile_primitive/gemm_utils.py +++ b/python/tvm/backend/cuda/tile_primitive/gemm_utils.py @@ -17,7 +17,7 @@ """GEMM-related utilities for CUDA op dispatches.""" -from tvm.arith.analyzer import Analyzer +from tvm.sym.analyzer import Analyzer from tvm.tirx import Buffer from tvm.tirx.operator.tile_primitive import DispatchContext from tvm.tirx.tile_primitive import TilePrimitiveCall diff --git a/python/tvm/backend/cuda/tile_primitive/layout_utils.py b/python/tvm/backend/cuda/tile_primitive/layout_utils.py index d59045de7809..137327992d70 100644 --- a/python/tvm/backend/cuda/tile_primitive/layout_utils.py +++ b/python/tvm/backend/cuda/tile_primitive/layout_utils.py @@ -27,7 +27,7 @@ import operator from collections import defaultdict -from tvm.arith import Analyzer +from tvm.sym import Analyzer from tvm.tirx.layout import ComposeLayout, S, TileLayout diff --git a/python/tvm/backend/cuda/tile_primitive/reduction/local.py b/python/tvm/backend/cuda/tile_primitive/reduction/local.py index f5629235a408..c6fe3a471128 100644 --- a/python/tvm/backend/cuda/tile_primitive/reduction/local.py +++ b/python/tvm/backend/cuda/tile_primitive/reduction/local.py @@ -61,9 +61,9 @@ import operator from typing import Any -from tvm.arith.analyzer import Analyzer from tvm.ir import TensorRegion from tvm.script import tirx as T +from tvm.sym.analyzer import Analyzer from tvm.tirx import PrimFunc from tvm.tirx.layout import TileLayout, laneid from tvm.tirx.operator.tile_primitive import DispatchContext, fail diff --git a/python/tvm/backend/cuda/tile_primitive/reduction/shared.py b/python/tvm/backend/cuda/tile_primitive/reduction/shared.py index cee4f7e1cc8e..74334c2de00d 100644 --- a/python/tvm/backend/cuda/tile_primitive/reduction/shared.py +++ b/python/tvm/backend/cuda/tile_primitive/reduction/shared.py @@ -57,9 +57,9 @@ import math import operator -from tvm.arith.analyzer import Analyzer from tvm.ir import TensorRegion from tvm.script import tirx as T +from tvm.sym.analyzer import Analyzer from tvm.tirx import PrimFunc from tvm.tirx.operator.tile_primitive import DispatchContext, fail from tvm.tirx.operator.tile_primitive.common import ReduceOpType diff --git a/python/tvm/backend/cuda/tile_primitive/reduction/utils.py b/python/tvm/backend/cuda/tile_primitive/reduction/utils.py index ece1d2734abf..d00a85d6e15b 100644 --- a/python/tvm/backend/cuda/tile_primitive/reduction/utils.py +++ b/python/tvm/backend/cuda/tile_primitive/reduction/utils.py @@ -21,9 +21,9 @@ import math import operator -from tvm.arith.analyzer import Analyzer from tvm.ir import TensorRegion from tvm.script import tirx as T +from tvm.sym.analyzer import Analyzer from tvm.tirx.operator.tile_primitive import DispatchContext from tvm.tirx.operator.tile_primitive.common import ReduceOpType from tvm.tirx.tile_primitive import TilePrimitiveCall diff --git a/python/tvm/backend/cuda/tile_primitive/tma_utils.py b/python/tvm/backend/cuda/tile_primitive/tma_utils.py index 99886ff363ca..85be6673d913 100644 --- a/python/tvm/backend/cuda/tile_primitive/tma_utils.py +++ b/python/tvm/backend/cuda/tile_primitive/tma_utils.py @@ -21,7 +21,7 @@ from enum import Enum import tvm -from tvm.arith.analyzer import Analyzer +from tvm.sym.analyzer import Analyzer from tvm.tirx.layout import ComposeLayout, Layout, S, TileLayout diff --git a/python/tvm/backend/trn/layout.py b/python/tvm/backend/trn/layout.py index 2cdfecfa571b..2c4976b74a42 100644 --- a/python/tvm/backend/trn/layout.py +++ b/python/tvm/backend/trn/layout.py @@ -41,7 +41,7 @@ def is_trainium_layout(layout: Layout | None) -> bool: def trainium_layout(annotation: str, shape: tuple[Expr], is_psum: bool = False) -> TileLayout: """Create a Trainium tile layout from a PF annotation string and logical shape.""" - analyzer = tvm.arith.Analyzer() + analyzer = tvm.sym.Analyzer() assert re.fullmatch(r"[PF]*", annotation), ( f"annotation {annotation} must be a string of 'P' and 'F'" ) @@ -97,7 +97,7 @@ def trainium_layout(annotation: str, shape: tuple[Expr], is_psum: bool = False) def to_psum_layout(layout: TileLayout) -> TileLayout: """Convert a Trainium sbuf layout to its psum physical-bank layout.""" - analyzer = tvm.arith.Analyzer() + analyzer = tvm.sym.Analyzer() shard = [] for iter in layout.shard: if iter.axis.name == "F": diff --git a/python/tvm/backend/trn/tile_primitive/binary/utils.py b/python/tvm/backend/trn/tile_primitive/binary/utils.py index 66453e55ee7d..306c564d453b 100644 --- a/python/tvm/backend/trn/tile_primitive/binary/utils.py +++ b/python/tvm/backend/trn/tile_primitive/binary/utils.py @@ -19,9 +19,9 @@ from enum import Enum -from tvm.arith.analyzer import Analyzer from tvm.backend.trn.layout import is_trainium_layout from tvm.ir import TensorRegion +from tvm.sym.analyzer import Analyzer from tvm.tirx import FloatImm from tvm.tirx.operator.tile_primitive.common import MapOpType diff --git a/python/tvm/backend/trn/tile_primitive/common.py b/python/tvm/backend/trn/tile_primitive/common.py index 9a7bbaa2fc4e..7ca9a5815bb4 100644 --- a/python/tvm/backend/trn/tile_primitive/common.py +++ b/python/tvm/backend/trn/tile_primitive/common.py @@ -17,7 +17,7 @@ """Common utilities for TRN operator scheduling.""" -from tvm.arith.analyzer import Analyzer +from tvm.sym.analyzer import Analyzer from tvm.tirx.operator.tile_primitive import DispatchContext # Used to generate the correct [:, None] for mask/predicate diff --git a/python/tvm/backend/trn/tile_primitive/dim_utils.py b/python/tvm/backend/trn/tile_primitive/dim_utils.py index f1de58c3348e..968fcfaac342 100644 --- a/python/tvm/backend/trn/tile_primitive/dim_utils.py +++ b/python/tvm/backend/trn/tile_primitive/dim_utils.py @@ -19,9 +19,9 @@ from collections import namedtuple -from tvm.arith.analyzer import Analyzer from tvm.ir import TensorRegion from tvm.script import tirx as T +from tvm.sym.analyzer import Analyzer # Represents the part of data iter covered by the buffer region RangeInfo = namedtuple( diff --git a/python/tvm/backend/trn/tile_primitive/gemm/default.py b/python/tvm/backend/trn/tile_primitive/gemm/default.py index 75467dffc80e..9935a3897015 100644 --- a/python/tvm/backend/trn/tile_primitive/gemm/default.py +++ b/python/tvm/backend/trn/tile_primitive/gemm/default.py @@ -20,10 +20,10 @@ import functools import operator -from tvm.arith.analyzer import Analyzer from tvm.backend.trn.layout import is_trainium_layout from tvm.ir import TensorRegion, assert_structural_equal from tvm.script import tirx as T +from tvm.sym.analyzer import Analyzer from tvm.tirx import PrimFunc from tvm.tirx.operator.tile_primitive import ( DispatchContext, diff --git a/python/tvm/backend/trn/tile_primitive/instruction_generator.py b/python/tvm/backend/trn/tile_primitive/instruction_generator.py index 48591d886017..aa75170bee69 100644 --- a/python/tvm/backend/trn/tile_primitive/instruction_generator.py +++ b/python/tvm/backend/trn/tile_primitive/instruction_generator.py @@ -26,10 +26,10 @@ import tvm_ffi import tvm -from tvm.arith.analyzer import Analyzer from tvm.backend.trn.layout import is_trainium_layout from tvm.ir import Range, TensorRegion from tvm.script import tirx as T +from tvm.sym.analyzer import Analyzer from tvm.tirx import BufferRegion, Expr, Var, is_buffer_var from tvm.tirx.layout import Iter diff --git a/python/tvm/backend/trn/tile_primitive/unary/utils.py b/python/tvm/backend/trn/tile_primitive/unary/utils.py index 40eabd560277..106648b17d59 100644 --- a/python/tvm/backend/trn/tile_primitive/unary/utils.py +++ b/python/tvm/backend/trn/tile_primitive/unary/utils.py @@ -17,10 +17,10 @@ """Shared helpers, op tables, and validation functions for unary operator dispatches.""" -from tvm.arith.analyzer import Analyzer from tvm.backend.trn.layout import is_trainium_layout from tvm.ir import TensorRegion from tvm.script import tirx as T +from tvm.sym.analyzer import Analyzer from tvm.tirx import FloatImm from tvm.tirx.operator.tile_primitive.common import MapOpType diff --git a/python/tvm/contrib/cutlass/build.py b/python/tvm/contrib/cutlass/build.py index a7b4cd667c5d..ca60e4d80f4e 100644 --- a/python/tvm/contrib/cutlass/build.py +++ b/python/tvm/contrib/cutlass/build.py @@ -427,7 +427,7 @@ def is_shape_valid_for_cutlass_matmul( # This could be regular matmul or batch matmul with shape ND x 2D or 2D x ND return True - analyzer = tvm.arith.Analyzer() + analyzer = tvm.sym.Analyzer() # If one side has less dimensions, use 1 to fill the gap batch_dim_pairs = list( itertools.zip_longest( diff --git a/python/tvm/ir/json_compact.py b/python/tvm/ir/json_compact.py index e26aab947a12..2bdc5c52ab45 100644 --- a/python/tvm/ir/json_compact.py +++ b/python/tvm/ir/json_compact.py @@ -19,6 +19,20 @@ import json _PRIM_TYPE_KEY_RENAMES = { + "arith.Analyzer": "sym.Analyzer", + "arith.CanonicalExpr": "sym.CanonicalExpr", + "arith.ConstIntBound": "sym.ConstIntBound", + "arith.IntervalSet": "sym.IntervalSet", + "arith.IterMapExpr": "sym.IterMapExpr", + "arith.IterMapResult": "sym.IterMapResult", + "arith.IterMark": "sym.IterMark", + "arith.IterSplitExpr": "sym.IterSplitExpr", + "arith.IterSumExpr": "sym.IterSumExpr", + "arith.ModularSet": "sym.ModularSet", + "arith.PresburgerSet": "sym.PresburgerSet", + "arith.RewriteSimplifierStats": "sym.RewriteSimplifierStats", + "arith.SplitExpr": "sym.SplitExpr", + "arith.SumExpr": "sym.SumExpr", "tirx.BufferRegion": "ir.TensorRegion", "tirx.SBlock": "s_tir.SBlock", "tirx.SBlockRealize": "s_tir.SBlockRealize", diff --git a/python/tvm/relax/backend/cuda/cublas.py b/python/tvm/relax/backend/cuda/cublas.py index 54bc7679c6c3..6459faa945a0 100644 --- a/python/tvm/relax/backend/cuda/cublas.py +++ b/python/tvm/relax/backend/cuda/cublas.py @@ -22,9 +22,9 @@ import tvm from tvm import DataType -from tvm.arith import Analyzer from tvm.relax import transform from tvm.relax.transform import PatternCheckContext +from tvm.sym import Analyzer from ..pattern_registry import get_patterns_with_prefix, register_patterns from ..patterns import ( diff --git a/python/tvm/relax/backend/cuda/cutlass.py b/python/tvm/relax/backend/cuda/cutlass.py index 2ce5a3de60a9..566e737236d4 100644 --- a/python/tvm/relax/backend/cuda/cutlass.py +++ b/python/tvm/relax/backend/cuda/cutlass.py @@ -83,7 +83,7 @@ def _has_dependency(from_var: Var, to_var: Var, var_usages: Mapping[Var, Sequenc def _is_same_shape(shape1, shape2): - analyzer = tvm.arith.Analyzer() + analyzer = tvm.sym.Analyzer() return all([analyzer.can_prove_equal(s1, s2) for s1, s2 in zip(shape1, shape2)]) diff --git a/python/tvm/relax/frontend/torch/base_fx_graph_translator.py b/python/tvm/relax/frontend/torch/base_fx_graph_translator.py index 3ff3b596af3b..2b05d62693e5 100644 --- a/python/tvm/relax/frontend/torch/base_fx_graph_translator.py +++ b/python/tvm/relax/frontend/torch/base_fx_graph_translator.py @@ -3132,7 +3132,7 @@ def _item(self, node: fx.Node) -> relax.Var: x = self.env[node.args[0]] shape = self.shape_of(x) dtype = x.ty.dtype - analyzer = tvm.arith.Analyzer() + analyzer = tvm.sym.Analyzer() has_single_element = shape is not None and all( analyzer.can_prove_equal(dim, 1) for dim in shape ) diff --git a/python/tvm/relax/op/_op_gradient.py b/python/tvm/relax/op/_op_gradient.py index 4e9d8d516b6e..dd1baf908156 100644 --- a/python/tvm/relax/op/_op_gradient.py +++ b/python/tvm/relax/op/_op_gradient.py @@ -21,9 +21,9 @@ import operator from tvm import relax -from tvm.arith import Analyzer from tvm.ir import Call, PrimType from tvm.relax.type import ShapeType +from tvm.sym import Analyzer from ..block_builder import BlockBuilder from ..expr import Expr, ShapeExpr, Var diff --git a/python/tvm/relax/transform/legalize_ops/ccl.py b/python/tvm/relax/transform/legalize_ops/ccl.py index 659b1f7d4397..f5340c1e2140 100644 --- a/python/tvm/relax/transform/legalize_ops/ccl.py +++ b/python/tvm/relax/transform/legalize_ops/ccl.py @@ -18,7 +18,7 @@ # ruff: noqa: RUF005 """Default legalization function for ccl operators.""" -from tvm import arith, tirx, topi +from tvm import sym, tirx, topi from tvm.ir import Call from ...block_builder import BlockBuilder @@ -91,7 +91,7 @@ def _transpose_for_ccl(_bb: BlockBuilder, expr: Expr, axis: int, num_workers: in new_shape = [] for i, shape_value in enumerate(arg_shape.values): if i == axis: - modulo = arith.Analyzer().simplify(shape_value % num_workers) + modulo = sym.Analyzer().simplify(shape_value % num_workers) assert modulo == 0, ( f"scatter_from_worker0 expects the size of axis {axis} of input tensor " "to be divisible by num_workers. However, the axis 0 of input tensor " diff --git a/python/tvm/s_tir/dlight/analysis/gemv.py b/python/tvm/s_tir/dlight/analysis/gemv.py index 33a83b38c7ca..4211cea6f1b8 100644 --- a/python/tvm/s_tir/dlight/analysis/gemv.py +++ b/python/tvm/s_tir/dlight/analysis/gemv.py @@ -18,7 +18,7 @@ import tvm_ffi -from tvm import arith, s_tir, tirx +from tvm import s_tir, sym, tirx from .common_analysis import ( SBlockInfo, @@ -107,7 +107,7 @@ def normalize( ) -> bool | None: """Normalize the main block.""" block_stmt: s_tir.SBlock = sch.get(block_info.block_rv) - access = arith.normalize_to_iter_sum( + access = sym.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/general_reduction.py b/python/tvm/s_tir/dlight/gpu/general_reduction.py index 22f1fde590c7..1bc14609cdc2 100644 --- a/python/tvm/s_tir/dlight/gpu/general_reduction.py +++ b/python/tvm/s_tir/dlight/gpu/general_reduction.py @@ -19,7 +19,7 @@ import tvm_ffi -from tvm import arith, s_tir, tirx +from tvm import s_tir, sym, tirx from tvm.target import Target from ..analysis import get_root_block, normalize_prim_func @@ -73,7 +73,7 @@ def apply( # pylint: disable=too-many-locals # preserving global scope for buffers accessed by another block. blocks = [sch.get(info.block_rv) for info in block_infos] alloc_buffers = list(sch.get(get_root_block(sch)).alloc_buffers) - analyzer = arith.Analyzer() + analyzer = sym.Analyzer() for block_index, (info, block) in enumerate(zip(block_infos[:-1], blocks[:-1])): loops = sch.get_loops(info.block_rv) if not all(analyzer.can_prove_equal(sch.get(loop).extent, 1) for loop in loops): @@ -106,7 +106,7 @@ def apply( # pylint: disable=too-many-locals return sch def f_layout_mapping(*iters): - analyzer = arith.Analyzer() + analyzer = sym.Analyzer() # Try to match the iters of last block to the iters of the first block. # For matched positions, use the iter from the input `iters`. # For unmatched positions, use a new iter which is constant 0. 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 2732fef989f7..c378d64fb163 100644 --- a/python/tvm/s_tir/dlight/gpu/low_batch_gemv.py +++ b/python/tvm/s_tir/dlight/gpu/low_batch_gemv.py @@ -22,7 +22,7 @@ import tvm_ffi -from tvm import arith, s_tir, tirx +from tvm import s_tir, sym, tirx from tvm.target import Target from ..analysis import ( @@ -157,7 +157,7 @@ def normalize( dynamic_iter_vars = set( iter_var.var for iter_var in block_stmt.iter_vars if iter_var.var not in const_iter_vars ) - access = arith.normalize_to_iter_sum( + access = sym.normalize_to_iter_sum( detect_dominant_read(block_stmt, const_iter_vars), input_iters={i.var: i.dom for i in block_stmt.iter_vars}, ) diff --git a/python/tvm/s_tir/dlight/gpu/reduction.py b/python/tvm/s_tir/dlight/gpu/reduction.py index 6c010ade697c..18b958049369 100644 --- a/python/tvm/s_tir/dlight/gpu/reduction.py +++ b/python/tvm/s_tir/dlight/gpu/reduction.py @@ -21,7 +21,7 @@ import tvm_ffi -from tvm import arith, s_tir, tirx +from tvm import s_tir, sym, tirx from tvm.target import Target from ..analysis import ( @@ -65,7 +65,7 @@ def _suggest_inner_spatial_tx(s_factor: int | tirx.Expr) -> int: def _get_spatial_domains_in_access_order( - block_info: SBlockInfo, access: arith.IterSumExpr + block_info: SBlockInfo, access: sym.IterSumExpr ) -> list[int | tirx.Expr] | None: """Return normalized spatial extents in access order.""" iter_to_info = {info.var: info for info in block_info.iters} @@ -151,7 +151,7 @@ def apply( # pylint: disable=too-many-locals,too-many-branches,too-many-return- ): return None # Step 2. Normalize the block, merge spatial and reduction iters - access = arith.normalize_to_iter_sum( + access = sym.normalize_to_iter_sum( detect_dominant_read(block_stmt), input_iters={i.var: i.dom for i in block_stmt.iter_vars}, ) @@ -187,7 +187,7 @@ def _normalize( # pylint: disable=too-many-branches self, sch: s_tir.Schedule, block_info: SBlockInfo, - access: arith.IterSumExpr, + access: sym.IterSumExpr, ) -> tuple[bool | None, int | None, Mapping[int, int] | None, int | None]: if access.base != 0: return None, None, None, None diff --git a/python/tvm/s_tir/dlight/gpu/transpose.py b/python/tvm/s_tir/dlight/gpu/transpose.py index aebdde59e6bf..d57ec7df219e 100644 --- a/python/tvm/s_tir/dlight/gpu/transpose.py +++ b/python/tvm/s_tir/dlight/gpu/transpose.py @@ -16,7 +16,7 @@ # under the License. """Reduction rule for operators including softmax, layer norm, RMS norm, etc""" -from tvm import arith, s_tir, tirx +from tvm import s_tir, sym, tirx from tvm.ir import TensorLoad from tvm.s_tir import Schedule from tvm.s_tir.schedule import SBlockRV @@ -91,7 +91,7 @@ def apply( # pylint: disable=too-many-locals c_factor = 1 if prologue is not None: block_stmt = sch.get(prologue) - result = arith.normalize_to_iter_sum( + result = sym.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/arith/__init__.py b/python/tvm/sym/__init__.py similarity index 95% rename from python/tvm/arith/__init__.py rename to python/tvm/sym/__init__.py index b646a6bf69a6..dc449f9f6d11 100644 --- a/python/tvm/arith/__init__.py +++ b/python/tvm/sym/__init__.py @@ -15,7 +15,7 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -"""Integer bound analysis, simplification and pattern detection.""" +"""Symbolic analysis: integer bounds, simplification and pattern detection.""" from .int_set import ( IntSet, diff --git a/python/tvm/arith/_ffi_api.py b/python/tvm/sym/_ffi_api.py similarity index 92% rename from python/tvm/arith/_ffi_api.py rename to python/tvm/sym/_ffi_api.py index 5211096f4253..f675ab8e9170 100644 --- a/python/tvm/arith/_ffi_api.py +++ b/python/tvm/sym/_ffi_api.py @@ -14,8 +14,8 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -"""FFI APIs for tvm.arith""" +"""FFI APIs for tvm.sym""" import tvm_ffi -tvm_ffi.init_ffi_api("arith", __name__) +tvm_ffi.init_ffi_api("sym", __name__) diff --git a/python/tvm/arith/analyzer.py b/python/tvm/sym/analyzer.py similarity index 98% rename from python/tvm/arith/analyzer.py rename to python/tvm/sym/analyzer.py index 59f8ba4e5d26..9ac1cfc85dfe 100644 --- a/python/tvm/arith/analyzer.py +++ b/python/tvm/sym/analyzer.py @@ -23,8 +23,8 @@ import tvm_ffi from tvm import ir -from tvm.arith import IntSet from tvm.runtime import Object +from tvm.sym import IntSet from . import _ffi_api @@ -39,7 +39,7 @@ class ProofStrength(enum.IntEnum): class CompareResult(enum.IntEnum): """Result of a transitive comparison. - Values must match the C++ ``arith::CompareResult`` enum. + Values must match the C++ ``sym::CompareResult`` enum. """ INCONSISTENT = 0 @@ -65,7 +65,7 @@ class Extension(enum.Flag): ComparisonOfProductAndSum = 1 << 3 -@tvm_ffi.register_object("arith.ModularSet") +@tvm_ffi.register_object("sym.ModularSet") class ModularSet(Object): """Represent range of (coeff * x + base) for x in Z""" @@ -73,7 +73,7 @@ def __init__(self, coeff, base): self.__init_handle_by_constructor__(_ffi_api.ModularSet, coeff, base) -@tvm_ffi.register_object("arith.ConstIntBound") +@tvm_ffi.register_object("sym.ConstIntBound") class ConstIntBound(Object): """Represent constant integer bound @@ -132,7 +132,7 @@ def __exit__(self, ptype, value, trace): _ffi_api.ExitZ3ContextScope() -@tvm_ffi.register_object("arith.Analyzer") +@tvm_ffi.register_object("sym.Analyzer") class Analyzer(Object): """Integer arithmetic analyzer @@ -357,7 +357,7 @@ def int_set(self, expr: ir.Expr, dom_map: dict[ir.Var, IntSet] | None = None) -> expr : Expr The expression. - dom_map : Optional[Dict[tvm.ir.Var, tvm.arith.IntSet]] + dom_map : Optional[Dict[tvm.ir.Var, tvm.sym.IntSet]] The domain for variables to be relaxed. When omitted, the analyzer uses the domains of the variables already bound to it. @@ -442,7 +442,7 @@ def constraint_scope(self, constraint: ir.Expr) -> ConstraintScope: .. code-block:: python x = te.var("x") - analyzer = tvm.arith.Analyzer() + analyzer = tvm.sym.Analyzer() with analyzer.constraint_scope(x % 3 == 0): # constraint in effect assert analyzer.modular_set(x).coeff == 3 diff --git a/python/tvm/arith/bound.py b/python/tvm/sym/bound.py similarity index 100% rename from python/tvm/arith/bound.py rename to python/tvm/sym/bound.py diff --git a/python/tvm/arith/int_set.py b/python/tvm/sym/int_set.py similarity index 96% rename from python/tvm/arith/int_set.py rename to python/tvm/sym/int_set.py index c7e79f679591..2b2307cd1040 100644 --- a/python/tvm/arith/int_set.py +++ b/python/tvm/sym/int_set.py @@ -68,7 +68,7 @@ def single_point(point): return _ffi_api.intset_single_point(point) -@tvm_ffi.register_object("arith.IntervalSet") +@tvm_ffi.register_object("sym.IntervalSet") class IntervalSet(IntSet): """Represent set of continuous interval [min_value, max_value] @@ -85,7 +85,7 @@ def __init__(self, min_value, max_value): self.__init_handle_by_constructor__(_ffi_api.IntervalSet, min_value, max_value) -@tvm_ffi.register_object("arith.PresburgerSet") +@tvm_ffi.register_object("sym.PresburgerSet") class PresburgerSet(IntSet): """Represent of Presburger Set""" @@ -108,7 +108,7 @@ def estimate_region_lower_bound(region, var_dom, predicate, analyzer=None): predicate : Expr The predicate for the affine map - analyzer : Optional[tvm.arith.Analyzer] + analyzer : Optional[tvm.sym.Analyzer] The analyzer to use. When provided, its accumulated bindings and constraints are reused; otherwise a fresh analyzer is created. @@ -135,7 +135,7 @@ def estimate_region_strict_bound(region, var_dom, predicate, analyzer=None): predicate : Expr The predicate for the affine map - analyzer : Optional[tvm.arith.Analyzer] + analyzer : Optional[tvm.sym.Analyzer] The analyzer to use. When provided, its accumulated bindings and constraints are reused; otherwise a fresh analyzer is created. @@ -163,7 +163,7 @@ def estimate_region_upper_bound(region, var_dom, predicate, analyzer=None): predicate : Expr The predicate for the affine map - analyzer : Optional[tvm.arith.Analyzer] + analyzer : Optional[tvm.sym.Analyzer] The analyzer to use. When provided, its accumulated bindings and constraints are reused; otherwise a fresh analyzer is created. diff --git a/python/tvm/arith/iter_affine_map.py b/python/tvm/sym/iter_affine_map.py similarity index 96% rename from python/tvm/arith/iter_affine_map.py rename to python/tvm/sym/iter_affine_map.py index bac1b030b187..98da4c99fc39 100644 --- a/python/tvm/arith/iter_affine_map.py +++ b/python/tvm/sym/iter_affine_map.py @@ -26,12 +26,12 @@ from . import _ffi_api -@tvm_ffi.register_object("arith.IterMapExpr") +@tvm_ffi.register_object("sym.IterMapExpr") class IterMapExpr(Expr): """Base class of all IterMap expressions.""" -@tvm_ffi.register_object("arith.IterMark") +@tvm_ffi.register_object("sym.IterMark") class IterMark(Object): """Mark the source as an iterator in [0, extent). @@ -48,7 +48,7 @@ def __init__(self, source, extent): self.__init_handle_by_constructor__(_ffi_api.IterMark, source, extent) -@tvm_ffi.register_object("arith.IterSplitExpr") +@tvm_ffi.register_object("sym.IterSplitExpr") class IterSplitExpr(IterMapExpr): """Split of an iterator. @@ -75,7 +75,7 @@ def __init__(self, source, lower_factor, extent, scale): ) -@tvm_ffi.register_object("arith.IterSumExpr") +@tvm_ffi.register_object("sym.IterSumExpr") class IterSumExpr(IterMapExpr): """Fuse multiple iterators by summing them with scaling. @@ -94,7 +94,7 @@ def __init__(self, args, base): self.__init_handle_by_constructor__(_ffi_api.IterSumExpr, args, base) -@tvm_ffi.register_object("arith.IterMapResult") +@tvm_ffi.register_object("sym.IterMapResult") class IterMapResult(Object): """Result of iter map detection.""" @@ -151,7 +151,7 @@ def detect_iter_map( If true, iterators with extent of 1 will be replaced with a constant value. - analyzer : Optional[tvm.arith.Analyzer] + analyzer : Optional[tvm.sym.Analyzer] The analyzer to use. When provided, its accumulated bindings and constraints are reused; otherwise a fresh analyzer is created. @@ -186,7 +186,7 @@ def normalize_to_iter_sum(index, input_iters, analyzer=None): input_iters : Map[tvm.ir.Var, Range] The domain of each input iterators. - analyzer : Optional[tvm.arith.Analyzer] + analyzer : Optional[tvm.sym.Analyzer] The analyzer to use. When provided, its accumulated bindings and constraints are reused; otherwise a fresh analyzer is created. @@ -234,7 +234,7 @@ def iter_map_simplify( If true, iterators with extent of 1 will be replaced with a constant value. - analyzer : Optional[tvm.arith.Analyzer] + analyzer : Optional[tvm.sym.Analyzer] The analyzer to use. When provided, its accumulated bindings and constraints are reused; otherwise a fresh analyzer is created. @@ -320,7 +320,7 @@ def subspace_divide( If true, iterators with extent of 1 will be replaced with a constant value. - analyzer : Optional[tvm.arith.Analyzer] + analyzer : Optional[tvm.sym.Analyzer] The analyzer to use. When provided, its accumulated bindings and constraints are reused; otherwise a fresh analyzer is created. diff --git a/python/tvm/arith/pattern.py b/python/tvm/sym/pattern.py similarity index 100% rename from python/tvm/arith/pattern.py rename to python/tvm/sym/pattern.py diff --git a/python/tvm/te/operation.py b/python/tvm/te/operation.py index fe69b3e42b4f..cd8f4d88599f 100644 --- a/python/tvm/te/operation.py +++ b/python/tvm/te/operation.py @@ -23,7 +23,7 @@ from tvm_ffi import Array -import tvm.arith._ffi_api +import tvm.sym._ffi_api import tvm.tirx import tvm.tirx._ffi_api from tvm.ir import is_prim_expr diff --git a/python/tvm/testing/utils.py b/python/tvm/testing/utils.py index 05f1e97c24a5..eb117a11b6de 100644 --- a/python/tvm/testing/utils.py +++ b/python/tvm/testing/utils.py @@ -84,8 +84,8 @@ def test_cuda_vectorize_add(): import tvm_ffi import tvm -import tvm.arith import tvm.support.utils +import tvm.sym import tvm.te import tvm.tirx from tvm.contrib import cudnn @@ -279,7 +279,7 @@ def assert_prim_expr_equal(lhs, rhs): rhs : tvm.tirx.Expr The left operand. """ - ana = tvm.arith.Analyzer() + ana = tvm.sym.Analyzer() if not ana.can_prove_equal(lhs, rhs): raise ValueError(f"{lhs} and {rhs} are not equal") @@ -330,7 +330,7 @@ def _compute_body(*us): counterex = [(str(v), i + r.min) for (v, r), i in zip(vranges.items(), indices)] counterex = sorted(counterex, key=lambda x: x[0]) counterex = ", ".join([v + " = " + str(i) for v, i in counterex]) - ana = tvm.arith.Analyzer() + ana = tvm.sym.Analyzer() raise AssertionError( f"Expression {ana.simplify(bool_expr)}\nis not true on {vranges}\n" f"Counterexample: {counterex}" diff --git a/python/tvm/tirx/_buffer_view.py b/python/tvm/tirx/_buffer_view.py index 90c4a13a6126..5a384d3ee0f4 100644 --- a/python/tvm/tirx/_buffer_view.py +++ b/python/tvm/tirx/_buffer_view.py @@ -142,7 +142,7 @@ def local(buf: Buffer, *shape, layout=None) -> Buffer: ) local_extent = buf.layout.storage().span() shape_total = functools.reduce(lambda x, y: x * y, shape, 1) - if not tvm.arith.Analyzer().can_prove_equal(shape_total, local_extent): + if not tvm.sym.Analyzer().can_prove_equal(shape_total, local_extent): raise ValueError( f"Local view shape {shape} has {shape_total} elements, " f"but the buffer has physical storage span {local_extent} per thread" @@ -342,7 +342,7 @@ def _swizzle_offset_commutes(swizzle, extra_offset): offset_c = _concrete_int(extra_offset) if offset_c is not None: return offset_c % period == 0 - from ..arith import Analyzer # pylint: disable=import-outside-toplevel + from ..sym import Analyzer # pylint: disable=import-outside-toplevel return Analyzer().can_prove_equal(tvm.tirx.floormod(extra_offset, period), 0) @@ -389,7 +389,7 @@ def _tmem_element_offset_to_column_offset(buf: Buffer, element_offset): ) return bit_offset_c // 32 - from ..arith import Analyzer # pylint: disable=import-outside-toplevel + from ..sym import Analyzer # pylint: disable=import-outside-toplevel analyzer = Analyzer() if not analyzer.can_prove_equal(tvm.tirx.floormod(bit_offset, 32), 0): diff --git a/python/tvm/tirx/function.py b/python/tvm/tirx/function.py index d6d58c2d7909..ec1f1fc57da9 100644 --- a/python/tvm/tirx/function.py +++ b/python/tvm/tirx/function.py @@ -303,7 +303,7 @@ def is_equivalent_to(self, other_map: "IndexMap", analyzer=None) -> bool: The IndexMap to which the comparison should be made. - analyzer : Optional[tvm.arith.Analyzer] + analyzer : Optional[tvm.sym.Analyzer] The analyzer to use while comparing the mapped indices. When provided, its accumulated bindings and constraints are reused so @@ -323,7 +323,7 @@ def is_equivalent_to(self, other_map: "IndexMap", analyzer=None) -> bool: return False if analyzer is None: - analyzer = tvm.arith.Analyzer() + analyzer = tvm.sym.Analyzer() mapped_other_final_indices = other_map.map_indices(self.initial_indices, analyzer=analyzer) for self_index, other_index in zip(self.final_indices, mapped_other_final_indices): @@ -339,7 +339,7 @@ def map_indices(self, indices: list[Expr], analyzer=None) -> list[Expr]: ---------- indices : List[Expr] The indices to be mapped - analyzer : Optional[tvm.arith.Analyzer] + analyzer : Optional[tvm.sym.Analyzer] The analyzer to use while simplifying mapped indices. Returns @@ -356,7 +356,7 @@ def map_shape(self, shape: list[Expr], analyzer=None) -> list[Expr]: ---------- shape : List[Expr] The buffer shape to be mapped - analyzer : Optional[tvm.arith.Analyzer] + analyzer : Optional[tvm.sym.Analyzer] The analyzer to use while simplifying mapped shape expressions. Returns @@ -393,7 +393,7 @@ def inverse(self, shape: list[Range | Expr], analyzer=None) -> "IndexMap": The region over which the inverse should be determined. Used for validating that the mapping is bijective over this range. - analyzer : Optional[tvm.arith.Analyzer] + analyzer : Optional[tvm.sym.Analyzer] The analyzer to use while deriving and validating the inverse. Returns @@ -419,7 +419,7 @@ def non_surjective_inverse( The region over which the inverse should be determined. Used for determining the predicate. - analyzer : Optional[tvm.arith.Analyzer] + analyzer : Optional[tvm.sym.Analyzer] The analyzer to use while deriving the inverse and padding predicate. Returns diff --git a/python/tvm/tirx/layout.py b/python/tvm/tirx/layout.py index a5ca42370734..33c78b7ff614 100644 --- a/python/tvm/tirx/layout.py +++ b/python/tvm/tirx/layout.py @@ -1186,7 +1186,7 @@ def __init__(self, terms: dict[Axis, Expr] | None = None): def _add_term(self, axis: Axis, value: Expr): if axis in self.terms: - # Merge if both exist; rely on tvm arith for symbolic add + # Merge if both exist; rely on tvm sym for symbolic add self.terms[axis] = self.terms[axis] + value # type: ignore[operator] else: self.terms[axis] = value @@ -1429,7 +1429,7 @@ def get_scope(self) -> tuple[ExecScope, ExecScope] | None: @classmethod def trainium(cls, annotation: str, shape: tuple[Expr], is_psum: bool = False) -> "TileLayout": """Create a TileLayout from an annotation string and a shape.""" - analyzer = tvm.arith.Analyzer() + analyzer = tvm.sym.Analyzer() assert re.fullmatch(r"[PF]*", annotation), ( f"annotation {annotation} must be a string of 'P' and 'F'" ) @@ -1487,7 +1487,7 @@ def trainium(cls, annotation: str, shape: tuple[Expr], is_psum: bool = False) -> def to_psum(self) -> "TileLayout": """Convert the layout to a psum layout.""" - analyzer = tvm.arith.Analyzer() + analyzer = tvm.sym.Analyzer() shard = [] for i in self.shard: if i.axis.name == "F": diff --git a/python/tvm/tirx/script/builder/ir.py b/python/tvm/tirx/script/builder/ir.py index f9be5ec39d75..ec8cc5465290 100644 --- a/python/tvm/tirx/script/builder/ir.py +++ b/python/tvm/tirx/script/builder/ir.py @@ -1029,7 +1029,7 @@ def _as_range(dom: ir.Range | list[Expr]) -> ir.Range: if isinstance(dom, ir.Range): return dom if isinstance(dom, list | tuple): - from tvm.arith import Analyzer # pylint: disable=import-outside-toplevel + from tvm.sym import Analyzer # pylint: disable=import-outside-toplevel extent = Analyzer().simplify(dom[1] - dom[0]) if isinstance(extent, tir.IntImm): @@ -2357,7 +2357,7 @@ def buffer_store( The indices location to be stored. """ - from tvm.arith import Analyzer # pylint: disable=import-outside-toplevel + from tvm.sym import Analyzer # pylint: disable=import-outside-toplevel if not isinstance(indices, list | tuple | ir.Array): indices = [indices] diff --git a/python/tvm/tirx/script/builder/tirx.py b/python/tvm/tirx/script/builder/tirx.py index e8cba7d76b27..ec2c5eb3e9d6 100644 --- a/python/tvm/tirx/script/builder/tirx.py +++ b/python/tvm/tirx/script/builder/tirx.py @@ -521,7 +521,7 @@ def _payload_bits(region): for axis in region.region: value = value * axis.extent if analyzer is None: - from tvm.arith import Analyzer # pylint: disable=import-outside-toplevel + from tvm.sym import Analyzer # pylint: disable=import-outside-toplevel analyzer = Analyzer() return analyzer.simplify(value) @@ -575,7 +575,7 @@ def _fail(): _fail() continue if analyzer is None: - from tvm.arith import Analyzer # pylint: disable=import-outside-toplevel + from tvm.sym import Analyzer # pylint: disable=import-outside-toplevel analyzer = Analyzer() if not analyzer.can_prove_equal(d, s): diff --git a/python/tvm/topi/nn/dilate.py b/python/tvm/topi/nn/dilate.py index 536654a43f1e..5ae90cdd4d74 100644 --- a/python/tvm/topi/nn/dilate.py +++ b/python/tvm/topi/nn/dilate.py @@ -49,7 +49,7 @@ def dilate(data, strides, dilation_value=0.0, name="DilatedInput"): n = len(data.shape) if len(strides) != n: raise ValueError(f"data dimension and strides size dismatch : {n} vs {len(strides)}") - ana = tvm.arith.Analyzer() + ana = tvm.sym.Analyzer() out_shape = tuple(ana.simplify((data.shape[i] - 1) * strides[i] + 1) for i in range(n)) def _dilate(*indices): diff --git a/python/tvm/topi/nn/pad.py b/python/tvm/topi/nn/pad.py index 2319166ffbb7..0b7202071ea7 100644 --- a/python/tvm/topi/nn/pad.py +++ b/python/tvm/topi/nn/pad.py @@ -48,7 +48,7 @@ def get_padded_shape(data, pad_before, pad_after=None): if len(pad_after) != n: raise ValueError(f"pad_after length {len(pad_after)} != input dims {n}") - ana = tvm.arith.Analyzer() + ana = tvm.sym.Analyzer() out_shape = tuple(ana.simplify(data.shape[i] + pad_before[i] + pad_after[i]) for i in range(n)) return out_shape @@ -86,7 +86,7 @@ def pad(data, pad_before, pad_after=None, pad_value=0.0, name="PadInput", attrs= raise ValueError(f"Input dimension and pad_before dismatch : {n} vs {len(pad_before)}") if len(pad_after) != n: raise ValueError(f"Input dimension and pad_after dismatch : {n} vs {len(pad_after)}") - ana = tvm.arith.Analyzer() + ana = tvm.sym.Analyzer() dshape = [] for dim in data.shape: dshape.append(dim) @@ -145,7 +145,7 @@ def mirror_pad(data, pad_before, pad_after=None, mode="SYMMETRIC", name="MirrorP raise ValueError(f"Input dimension and pad_before dismatch : {n} vs {len(pad_before)}") if len(pad_after) != n: raise ValueError(f"Input dimension and pad_after dismatch : {n} vs {len(pad_after)}") - ana = tvm.arith.Analyzer() + ana = tvm.sym.Analyzer() out_shape = tuple(ana.simplify(data.shape[i] + pad_before[i] + pad_after[i]) for i in range(n)) assert mode in ("SYMMETRIC", "REFLECT") mode = int(mode == "SYMMETRIC") diff --git a/python/tvm/topi/scatter.py b/python/tvm/topi/scatter.py index 453cfbedf496..05c3dba215c9 100644 --- a/python/tvm/topi/scatter.py +++ b/python/tvm/topi/scatter.py @@ -19,9 +19,9 @@ """ScatterND operator""" from tvm import DataTypeCode, ir, te, tirx # hide redefinition of min and max -from tvm.arith.analyzer import Analyzer from tvm.script.ir_builder import IRBuilder from tvm.script.ir_builder import tirx as T +from tvm.sym.analyzer import Analyzer def _verify_scatter_nd_inputs(data, indices, updates): diff --git a/python/tvm/topi/utils.py b/python/tvm/topi/utils.py index b05bb8c2bb02..b1bbe29178f5 100644 --- a/python/tvm/topi/utils.py +++ b/python/tvm/topi/utils.py @@ -119,7 +119,7 @@ def get_const_int(expr): if isinstance(expr, Integral): return expr if not isinstance(expr, tvm.tirx.IntImm): - ana = tvm.arith.Analyzer() + ana = tvm.sym.Analyzer() expr = ana.simplify(expr) if not isinstance(expr, tvm.tirx.IntImm): raise ValueError("Expect value to be constant int") @@ -142,7 +142,7 @@ def get_const_float(expr): if isinstance(expr, float): return float(expr) if not isinstance(expr, tvm.tirx.FloatImm): - ana = tvm.arith.Analyzer() + ana = tvm.sym.Analyzer() expr = ana.simplify(expr) if not isinstance(expr, tvm.tirx.FloatImm): raise ValueError("Expect value to be constant float") @@ -165,7 +165,7 @@ def equal_const_int(expr, value): if isinstance(expr, Integral): return expr == value if not isinstance(expr, tvm.tirx.IntImm): - ana = tvm.arith.Analyzer() + ana = tvm.sym.Analyzer() expr = ana.simplify(expr) if not isinstance(expr, tvm.tirx.IntImm): return False @@ -196,7 +196,7 @@ def get_const_tuple(in_tuple): if tvm.ir.is_prim_var(elem): ret.append(elem) elif not isinstance(elem, tvm.tirx.IntImm | int): - ana = tvm.arith.Analyzer() if ana is None else ana + ana = tvm.sym.Analyzer() if ana is None else ana elem = ana.simplify(elem) if not isinstance(elem, tvm.tirx.IntImm): ret.append(elem) @@ -271,12 +271,12 @@ def simplify(expr): if isinstance(expr, te.Tensor): return te.compute( expr.shape, - lambda *indices: tvm.arith.Analyzer().simplify(expr[indices]), + lambda *indices: tvm.sym.Analyzer().simplify(expr[indices]), name="simplify_output", tag="simplify", ) elif tvm.ir.is_prim_expr(expr): - return tvm.arith.Analyzer().simplify(expr) + return tvm.sym.Analyzer().simplify(expr) else: return expr diff --git a/src/backend/cuda/codegen/codegen_cuda.cc b/src/backend/cuda/codegen/codegen_cuda.cc index 53d52c6a12b2..55e61927b156 100644 --- a/src/backend/cuda/codegen/codegen_cuda.cc +++ b/src/backend/cuda/codegen/codegen_cuda.cc @@ -23,10 +23,10 @@ #include "codegen_cuda.h" -#include #include #include #include +#include #include #include @@ -240,7 +240,7 @@ class ThreadIdxExtractor : public tirx::StmtExprVisitor { void CodeGenCUDA::PrintExtraAttrs(const PrimFunc& f, std::ostream& os) { auto extractor = ffi::make_object(); extractor->Visit(f->body); - arith::Analyzer analyzer; + sym::Analyzer analyzer; PrimExpr threadIdx_ext = analyzer->Simplify( extractor->threadIdx_x_ext * extractor->threadIdx_y_ext * extractor->threadIdx_z_ext); PrimExpr cluster_cta_yz_ext = @@ -343,7 +343,7 @@ void CodeGenCUDA::Dispatch_(const tirx::ForNode* op) { // those declarations are printed between the pragma and the for statement, // nvcc is free to unroll the loop despite disable_unroll. std::string begin_str = PrintExpr(op->min); - PrimExpr end = is_zero(op->min) ? op->extent : arith::Analyzer()->Simplify(op->min + op->extent); + PrimExpr end = is_zero(op->min) ? op->extent : sym::Analyzer()->Simplify(op->min + op->extent); std::string end_str = PrintExpr(end); std::string step_str = op->step.has_value() ? PrintExpr(*op->step) : ""; if (op->annotations.count("disable_unroll")) { @@ -1165,7 +1165,7 @@ void CodeGenCUDA::Dispatch_(const CallNode* op, std::ostream& os) { tvm::ffi::Function::GetGlobal("tirx.index_map.shared_16x16_to_ldmatrix_32x8_layout"); TVM_FFI_ICHECK(index_map_func.has_value()); - arith::Analyzer analyzer; + sym::Analyzer analyzer; auto inverse_index_map = IndexMap::FromFunc(2, *index_map_func).Inverse({Range(0, m), Range(0, n)}, analyzer); auto indices_16x16 = inverse_index_map->final_indices; @@ -1274,7 +1274,7 @@ void CodeGenCUDA::Dispatch_(const CallNode* op, std::ostream& os) { tvm::ffi::Function::GetGlobal("tirx.index_map.shared_16x16_to_ldmatrix_32x8_layout"); TVM_FFI_ICHECK(index_map_func.has_value()); - arith::Analyzer analyzer; + sym::Analyzer analyzer; auto inverse_index_map = IndexMap::FromFunc(2, *index_map_func).Inverse({Range(0, m), Range(0, n)}, analyzer); auto indices_16x16 = inverse_index_map->final_indices; diff --git a/src/backend/hexagon/codegen/llvm/codegen_hexagon.cc b/src/backend/hexagon/codegen/llvm/codegen_hexagon.cc index e290d63367c8..86f700fa25ea 100644 --- a/src/backend/hexagon/codegen/llvm/codegen_hexagon.cc +++ b/src/backend/hexagon/codegen/llvm/codegen_hexagon.cc @@ -338,7 +338,7 @@ llvm::Value* CodeGenHexagon::VectorLookupLoad(BufferVar buffer, PrimType buffer_ if (buffer_type.bits() != 8) return nullptr; int table_elem_count = - arith::Analyzer()->Simplify(buffer->shape[0]).as()->value.as().value(); + sym::Analyzer()->Simplify(buffer->shape[0]).as()->value.as().value(); if (table_elem_count <= 0 || table_elem_count > 256) return nullptr; auto int32 = PrimType::Int(32); diff --git a/src/backend/metal/codegen/codegen_metal.cc b/src/backend/metal/codegen/codegen_metal.cc index 40e87e4f4309..e20a8a45d8eb 100644 --- a/src/backend/metal/codegen/codegen_metal.cc +++ b/src/backend/metal/codegen/codegen_metal.cc @@ -22,12 +22,12 @@ */ #include "codegen_metal.h" -#include #include #include #include #include #include +#include #include #include @@ -362,14 +362,14 @@ void CodeGenMetal::Dispatch_(const AllocBufferNode* op) { this->PrintIndent(); // Compute a compile-time upper bound on the number of buffer elements. size_t constant_size = 1; - arith::Analyzer analyzer; + sym::Analyzer analyzer; for (const auto& dim : op->buffer->shape) { const auto* dim_imm = dim.as(); int64_t dim_size = dim_imm ? static_cast(dim_imm->value) : analyzer->const_int_bound(dim)->max_value; if (dim_imm == nullptr) { // An integer dtype's intrinsic maximum is not a program-derived allocation bound. - TVM_FFI_ICHECK(dim_size != arith::ConstIntBound::kPosInf) + TVM_FFI_ICHECK(dim_size != sym::ConstIntBound::kPosInf) << "Metal allocation extent requires a finite compile-time upper bound, but got " << dim; if (const auto* dtype_max = max_value(dim.ty()).as()) { TVM_FFI_ICHECK_LT(dim_size, dtype_max->value) diff --git a/src/backend/opencl/codegen/intrin_rule_opencl.cc b/src/backend/opencl/codegen/intrin_rule_opencl.cc index 0d61f6d8ea6f..c39d8c8dc6ae 100644 --- a/src/backend/opencl/codegen/intrin_rule_opencl.cc +++ b/src/backend/opencl/codegen/intrin_rule_opencl.cc @@ -21,7 +21,7 @@ * \file intrin_rule_opencl.cc * \brief OpenCL intrinsic rules. */ -#include +#include #include #include "../../../target/intrin_rule.h" @@ -37,7 +37,7 @@ static PrimExpr DispatchIntelShuffle(const PrimExpr& e) { const CallNode* call = e.as(); TVM_FFI_ICHECK(call != nullptr); TVM_FFI_ICHECK_EQ(call->args.size(), 5); // mask, value, warp_id, width, warp_size - arith::Analyzer analyzer; + sym::Analyzer analyzer; TVM_FFI_ICHECK(analyzer->CanProve(call->args[3].as_or_throw() == call->args[4].as_or_throw())) << "Intel warp shuffle dose not support width != warp_size"; diff --git a/src/backend/trn/transform/lower_trainium_layout.cc b/src/backend/trn/transform/lower_trainium_layout.cc index 6d44daf95f67..4aff399a0f96 100644 --- a/src/backend/trn/transform/lower_trainium_layout.cc +++ b/src/backend/trn/transform/lower_trainium_layout.cc @@ -22,9 +22,9 @@ * \brief Trainium-specific TIRx layout lowering. */ -#include #include #include +#include #include #include #include @@ -55,7 +55,7 @@ static bool IsTrainiumLayout(const TileLayoutNode* layout) { class TrainiumLayoutApplier : public tirx::IRMutatorWithAnalyzer { public: static std::pair> Lower(const Stmt& stmt, const ffi::Array& params) { - arith::Analyzer ana; + sym::Analyzer ana; auto storage_lower = ffi::make_object(ana); ffi::Array new_params; new_params.reserve(params.size()); @@ -84,7 +84,7 @@ class TrainiumLayoutApplier : public tirx::IRMutatorWithAnalyzer { return std::make_pair(new_stmt, new_params); } - explicit TrainiumLayoutApplier(const arith::Analyzer& analyzer) + explicit TrainiumLayoutApplier(const sym::Analyzer& analyzer) : tirx::IRMutatorWithAnalyzer(analyzer) {} protected: @@ -153,7 +153,7 @@ class TrainiumLayoutApplier : public tirx::IRMutatorWithAnalyzer { } else if (is_alloc) { if (auto tile_layout = buf->layout.as(); tile_layout && tile_layout->HasThreadAxis()) { - arith::Analyzer ana; + sym::Analyzer ana; PrimExpr mem_span = IntImm::Int32(1); for (const auto& iter : tile_layout->shard) { if (iter->axis->IsMemoryAxis()) { diff --git a/src/backend/vulkan/codegen/codegen_spirv.cc b/src/backend/vulkan/codegen/codegen_spirv.cc index 83ba13b2eeb4..63f2f577db84 100644 --- a/src/backend/vulkan/codegen/codegen_spirv.cc +++ b/src/backend/vulkan/codegen/codegen_spirv.cc @@ -160,7 +160,7 @@ void CodeGenSPIRV::InitFuncState() { std::fill(workgroup_size_, workgroup_size_ + 3, 1); var_map_.clear(); storage_info_.clear(); - analyzer_ = arith::Analyzer(); + analyzer_ = sym::Analyzer(); builder_.reset(new spirv::IRBuilder(spirv_support_)); builder_->InitHeader(); shared_memory_bytes_used_ = 0; diff --git a/src/backend/vulkan/codegen/codegen_spirv.h b/src/backend/vulkan/codegen/codegen_spirv.h index e5ee1321b2eb..df84e7e47088 100644 --- a/src/backend/vulkan/codegen/codegen_spirv.h +++ b/src/backend/vulkan/codegen/codegen_spirv.h @@ -24,8 +24,8 @@ #ifndef TVM_TARGET_VULKAN_CODEGEN_SPIRV_H_ #define TVM_TARGET_VULKAN_CODEGEN_SPIRV_H_ -#include #include +#include #include #include #include @@ -226,7 +226,7 @@ class CodeGenSPIRV : public tirx::ExprFunctor, std::unordered_map var_map_; // The analyzer. - arith::Analyzer analyzer_; + sym::Analyzer analyzer_; // deep comparison of PrimExpr prim::ExprDeepEqual deep_equal_; diff --git a/src/backend/webgpu/codegen/codegen_webgpu.cc b/src/backend/webgpu/codegen/codegen_webgpu.cc index 233c7b8f52d4..d3af342cd6e9 100644 --- a/src/backend/webgpu/codegen/codegen_webgpu.cc +++ b/src/backend/webgpu/codegen/codegen_webgpu.cc @@ -22,12 +22,12 @@ */ #include "codegen_webgpu.h" -#include #include #include #include #include #include +#include #include #include @@ -40,11 +40,11 @@ #include #include -#include "../../../arith/pattern_match.h" #include "../../../runtime/file_utils.h" #include "../../../runtime/metadata.h" #include "../../../runtime/thread_storage_scope.h" #include "../../../support/bytes_io.h" +#include "../../../sym/pattern_match.h" #include "../../../target/build_common.h" #include "webgpu_fallback_module.h" @@ -622,8 +622,8 @@ void CodeGenWebGPU::Dispatch_(const TensorLoadNode* op, std::ostream& os) { // TVM_FFI_ICHECK_EQ(element_ty.lanes(), 1) << "Can only vector load scalar array"; TVM_FFI_ICHECK(value_ty.WithLanes(1) == element_ty) << "WebGPU vector loading requires base type to match"; - arith::PVar base; - if (arith::ramp(base, 1, value_ty.lanes()).Match(index)) { + sym::PVar base; + if (sym::ramp(base, 1, value_ty.lanes()).Match(index)) { // vec3(buf[base + 0], buf[base + 1], buf[base + 2]); std::string base_vid = SSAGetID(PrintExpr(base.Eval()), base.Eval().ty()); PrintType(element_ty.WithLanes(value_ty.lanes()), os); @@ -699,8 +699,8 @@ void CodeGenWebGPU::Dispatch_(const BufferStoreNode* op) { TVM_FFI_ICHECK(value_ty.WithLanes(1) == element_ty) << "WebGPU vector stire requires base type to match"; std::string value_vid = PrintExpr(op->value); - arith::PVar base; - if (arith::ramp(base, 1, value_ty.lanes()).Match(index)) { + sym::PVar base; + if (sym::ramp(base, 1, value_ty.lanes()).Match(index)) { // buf[base + 0] = value[0] // buf[base + 1] = value[1] std::string base_vid = SSAGetID(PrintExpr(base.Eval()), base.Eval().ty()); @@ -726,7 +726,7 @@ void CodeGenWebGPU::Dispatch_(const AllocBufferNode* op) { TVM_FFI_ICHECK(op->buffer.defined()); std::string vid = AllocVarID(op->buffer.get()); size_t constant_size = 1; - arith::Analyzer analyzer; + sym::Analyzer analyzer; for (const auto& dim : op->buffer->shape) { const auto* dim_imm = dim.as(); int64_t dim_size = @@ -787,7 +787,7 @@ void CodeGenWebGPU::Dispatch_(const AllocBufferNode* op) { void CodeGenWebGPU::Dispatch_(const ForNode* op) { std::string begin_str = PrintExpr(op->min); - PrimExpr end = is_zero(op->min) ? op->extent : arith::Analyzer()->Simplify(op->min + op->extent); + PrimExpr end = is_zero(op->min) ? op->extent : sym::Analyzer()->Simplify(op->min + op->extent); std::string end_str = PrintExpr(end); std::string step_str = op->step.has_value() ? PrintExpr(*op->step) : ""; std::string vid = AllocVarID(op->loop_var.get()); diff --git a/src/backend/webgpu/codegen/intrin_rule_webgpu.cc b/src/backend/webgpu/codegen/intrin_rule_webgpu.cc index a275c4a98429..75ce8c6367c9 100644 --- a/src/backend/webgpu/codegen/intrin_rule_webgpu.cc +++ b/src/backend/webgpu/codegen/intrin_rule_webgpu.cc @@ -21,7 +21,7 @@ * \file intrin_rule_webgpu.cc * \brief WebGPU intrinsic rules. */ -#include +#include #include #include "../../../target/intrin_rule.h" diff --git a/src/relax/analysis/layout_transformation.cc b/src/relax/analysis/layout_transformation.cc index 38af950dd20b..b8e47e43b1ed 100644 --- a/src/relax/analysis/layout_transformation.cc +++ b/src/relax/analysis/layout_transformation.cc @@ -22,14 +22,14 @@ * \brief Analyze the PrimFunc and suggest layout transformation on it's blocks and buffers based on * the user provided layout transformations on it's outputs. */ -#include -#include #include #include #include #include #include #include +#include +#include #include #include @@ -47,9 +47,9 @@ static bool IsBijectiveAffine(const IndexMap& m, const ffi::Array& ranges for (size_t i = 0; i < ranges.size(); i++) { input_iters.Set(m->initial_indices[i], ranges[i]); } - arith::Analyzer analyzer; + sym::Analyzer analyzer; auto iter_map_result = DetectIterMap(m->final_indices, input_iters, /* predicate = */ 1, - /*check_level=*/arith::IterMapLevel::Bijective, analyzer, + /*check_level=*/sym::IterMapLevel::Bijective, analyzer, /*simplify_trivial_iterators=*/true); return !iter_map_result->indices.empty(); } @@ -62,7 +62,7 @@ static bool IsBijectiveAffine(const IndexMap& m, const ffi::Array& ranges */ class IndexAnalyzer : public s_tir::StmtExprVisitor { public: - ffi::Array Analyze(const arith::IterSumExpr& expr) { + ffi::Array Analyze(const sym::IterSumExpr& expr) { Visit(expr); return iterators_; } @@ -70,11 +70,11 @@ class IndexAnalyzer : public s_tir::StmtExprVisitor { private: /*! \brief Override Visit for iter expr type processing */ ffi::Optional Visit(ffi::AnyView value) override { - if (const auto* op = value.as()) { + if (const auto* op = value.as()) { for (const auto& arg : op->args) TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Visit(arg)); return Visit(op->base); } - if (const auto* op = value.as()) { + if (const auto* op = value.as()) { TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(VisitIterMark(op->source)); TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Visit(op->lower_factor)); TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Visit(op->extent)); @@ -83,7 +83,7 @@ class IndexAnalyzer : public s_tir::StmtExprVisitor { return s_tir::StmtExprVisitor::Visit(value); } - ffi::Optional VisitIterMark(const arith::IterMark& op) { + ffi::Optional VisitIterMark(const sym::IterMark& op) { if (auto var = op->source.as()) iterators_.push_back(var.value()); else @@ -111,15 +111,15 @@ class IndexAnalyzer : public s_tir::StmtExprVisitor { * SpatialLayout(A[s0 * c + s1]) = undefined */ using SpatialLayout = ffi::Array>; -static SpatialLayout GetSpatialLayout(const arith::IterMapResult& iter_map_result) { +static SpatialLayout GetSpatialLayout(const sym::IterMapResult& iter_map_result) { TVM_FFI_ICHECK(!iter_map_result->indices.empty()); SpatialLayout result; - for (const arith::IterSumExpr& index : iter_map_result->indices) { + for (const sym::IterSumExpr& index : iter_map_result->indices) { auto index_analyzer = ffi::make_object(); ffi::Array iter_vars = index_analyzer->Analyze(index); if (iter_vars.size() >= 2) { LOG(WARNING) << "[LayoutInference] Unable to get spatial layout of access: " - << arith::NormalizeIterMapToExpr(index); + << sym::NormalizeIterMapToExpr(index); return {}; } if (iter_vars.empty()) { @@ -174,7 +174,7 @@ static bool AreIdenticalTransforms(const IndexMap& t0, const IndexMap& t1) { // Create a new shape expression. ffi::Array t1_initial_indices = t1->initial_indices.Map([](tirx::Var i) { return i.as_or_throw(); }); - arith::Analyzer analyzer; + sym::Analyzer analyzer; auto t0_output = t0->MapIndices(t1_initial_indices, analyzer); for (size_t i = 0; i < t0_output.size(); ++i) { if (!analyzer->CanProveEqual(t0_output[i], t1->final_indices[i])) return false; @@ -456,9 +456,9 @@ class BlockAnalyzer : public s_tir::StmtExprVisitor { // Helper to break down the indices of buffer access. SpatialLayout DetectBufferAccessIterMap(ffi::Array indices) { - auto result = arith::DetectIterMap( + auto result = sym::DetectIterMap( /*indices=*/indices, /*input_iters*/ spatial_dom_, - /*predicate*/ 1, /*check_level*/ arith::IterMapLevel::NoCheck, arith_analyzer_); + /*predicate*/ 1, /*check_level*/ sym::IterMapLevel::NoCheck, sym_analyzer_); if (result->indices.empty()) { DLOG(INFO) << "[LayoutInference] Failed to analyze indices " << indices << ", error: " << result->errors; @@ -550,7 +550,7 @@ class BlockAnalyzer : public s_tir::StmtExprVisitor { bool can_transform_block_; IndexMap write_transformation_; ffi::Map spatial_dom_; - arith::Analyzer arith_analyzer_; + sym::Analyzer sym_analyzer_; s_tir::SBlock block_; IndexMap block_transformation_; diff --git a/src/relax/analysis/shape_analysis.cc b/src/relax/analysis/shape_analysis.cc index df4bd8376f01..126654cf1d37 100644 --- a/src/relax/analysis/shape_analysis.cc +++ b/src/relax/analysis/shape_analysis.cc @@ -23,14 +23,14 @@ * \brief Utilities for shape analysis. */ -#include #include +#include namespace tvm { namespace relax { bool CanProveShapeEqual(const ffi::Array& lhs, const ffi::Array& rhs, - const arith::Analyzer& ana) { + const sym::Analyzer& ana) { if (lhs.same_as(rhs)) return true; if (lhs.size() != rhs.size()) return false; for (size_t i = 0; i < lhs.size(); ++i) { @@ -39,7 +39,7 @@ bool CanProveShapeEqual(const ffi::Array& lhs, const ffi::Array(); auto* rhs_shape = rhs.as(); diff --git a/src/relax/analysis/tir_op_pattern_kind.cc b/src/relax/analysis/tir_op_pattern_kind.cc index f9f428c33ab7..ca19476bcc9a 100644 --- a/src/relax/analysis/tir_op_pattern_kind.cc +++ b/src/relax/analysis/tir_op_pattern_kind.cc @@ -17,7 +17,6 @@ * under the License. */ -#include #include #include #include @@ -26,6 +25,7 @@ #include #include #include +#include #include #include #include @@ -462,11 +462,11 @@ bool HasReshapePattern(const PrimFunc& func) { idx = idx * buffer->shape[i] + indices[i]; } idx = ana_->Simplify(idx); - return arith::IterMapSimplify( + return sym::IterMapSimplify( /*indices=*/{idx}, /*input_iters=*/var_range, /*input_pred=*/IntImm::Bool(true), - /*check_level=*/arith::IterMapLevel::Surjective, + /*check_level=*/sym::IterMapLevel::Surjective, /*analyzer=*/ana_, /*simplify_trivial_iterators=*/true)[0]; }; @@ -524,12 +524,12 @@ bool HasReshapePattern(const PrimFunc& func) { ffi::StructuralMap(std::move(flattened_idx), f_substitute) .as_or_throw(); - ffi::Array simplify_res = arith::IterMapSimplify( + ffi::Array simplify_res = sym::IterMapSimplify( /*indices=*/{flattened_idx}, /*input_iters=*/ ffi::Map{{fused_var, Range(IntImm(dtype, /*value=*/0), stride)}}, /*input_pred=*/IntImm::Bool(true), - /*check_level=*/arith::IterMapLevel::Surjective, + /*check_level=*/sym::IterMapLevel::Surjective, /*analyzer=*/this->ana_, /*simplify_trivial_iterators=*/true); TVM_FFI_ICHECK_EQ(simplify_res.size(), 1); @@ -556,7 +556,7 @@ bool HasReshapePattern(const PrimFunc& func) { bool is_reshape_; const BufferVar& src_buffer_; const BufferVar& dst_buffer_; - arith::Analyzer ana_; + sym::Analyzer ana_; }; ffi::Array buffer_args; diff --git a/src/relax/analysis/type_analysis.cc b/src/relax/analysis/type_analysis.cc index 12f39e4d1d09..c7750f000a59 100644 --- a/src/relax/analysis/type_analysis.cc +++ b/src/relax/analysis/type_analysis.cc @@ -125,7 +125,7 @@ Type TypeFromStaticType(const Type& type) { class WellDefinedEraser : public TypeMutator, public ExprMutatorBase { public: WellDefinedEraser(std::function(const Var& var)> f_var_map, - arith::AnalyzerObj* ana) + sym::AnalyzerObj* ana) : f_var_map_(f_var_map), ana_(ana) {} Type VisitType_(const PrimTypeNode* op) final { return ffi::GetRef(op); } @@ -247,27 +247,27 @@ class WellDefinedEraser : public TypeMutator, public ExprMutatorBase { private: bool has_undefined_ = false; std::function(const Var& var)> f_var_map_; - arith::AnalyzerObj* ana_; + sym::AnalyzerObj* ana_; }; Type EraseToWellDefined(const Type& info, std::function(const Var& var)> f_var_map) { - arith::Analyzer analyzer; + sym::Analyzer analyzer; return EraseToWellDefined(info, f_var_map, analyzer); } Type EraseToWellDefined(const Type& info, std::function(const Var& var)> f_var_map, - const arith::Analyzer& ana) { + const sym::Analyzer& ana) { return WellDefinedEraser(f_var_map, ana.get()).VisitType(info); } Type EraseToWellDefined(const Type& info, ffi::Map var_map) { - arith::Analyzer analyzer; + sym::Analyzer analyzer; return EraseToWellDefined(info, var_map, analyzer); } -Type EraseToWellDefined(const Type& info, ffi::Map var_map, const arith::Analyzer& ana) { +Type EraseToWellDefined(const Type& info, ffi::Map var_map, const sym::Analyzer& ana) { std::function(const Var& var)> f_var_map = nullptr; if (!var_map.empty()) { @@ -300,7 +300,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { //-------------------------- class TypeBaseChecker : public TypeFunctor { public: - explicit TypeBaseChecker(arith::AnalyzerObj* ana) : analyzer_(ana) {} + explicit TypeBaseChecker(sym::AnalyzerObj* ana) : analyzer_(ana) {} BaseCheckResult VisitType(const Type& lhs, const Type& other) override { // quick path @@ -474,7 +474,7 @@ class TypeBaseChecker : public TypeFunctor(analyzer_))) { + if (CanProveShapeEqual(mapped_value, rhs, ffi::GetRef(analyzer_))) { return BaseCheckResult::kPass; } return BaseCheckResult::kFailL2; @@ -952,12 +952,12 @@ class CallRetTypeDeriver : public TypeBaseChecker { }; Type DeriveCallRetType(const FuncType& finfo, const Call& call, const BlockBuilder& ctx) { - arith::Analyzer analyzer; + sym::Analyzer analyzer; return DeriveCallRetType(finfo, call, ctx, analyzer); } Type DeriveCallRetType(const FuncType& finfo, const Call& call, const BlockBuilder& ctx, - const arith::Analyzer& ana) { + const sym::Analyzer& ana) { // The deriver's TVM_FFI_VISIT_THROW seeds a VisitErrorContext on the error; // the outer pass wrapper catches it and enriches the message with the access // path. Nothing to do here but propagate. @@ -977,7 +977,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { //-------------------------- class TypeLCAFinder : public TypeFunctor { public: - explicit TypeLCAFinder(arith::AnalyzerObj* ana) : analyzer_(ana) {} + explicit TypeLCAFinder(sym::AnalyzerObj* ana) : analyzer_(ana) {} Type VisitType(const Type& lhs, const Type& other) final { // quick path @@ -1008,7 +1008,7 @@ class TypeLCAFinder : public TypeFunctor { int ndim = lhs->ndim == rhs->ndim ? lhs->ndim : kUnknownNDim; if (lhs->ndim != rhs->ndim || !lhs->values.has_value() || !rhs->values.has_value() || !CanProveShapeEqual(lhs->values.value(), rhs->values.value(), - ffi::GetRef(analyzer_))) { + ffi::GetRef(analyzer_))) { // prefers return same when possible if (!lhs->values.has_value() && lhs->ndim == ndim) { return ffi::GetRef(lhs); @@ -1039,7 +1039,7 @@ class TypeLCAFinder : public TypeFunctor { // then we cannot keep in symbolic shape if (lhs->ndim != rhs->ndim || !lhs->shape.has_value() || !rhs->shape.has_value() || !CanProveShapeEqual(lhs->shape.value(), rhs->shape.value(), - ffi::GetRef(analyzer_))) { + ffi::GetRef(analyzer_))) { // reuse lhs when possible if (!lhs->shape.has_value() && lhs->dtype == dtype && lhs->ndim == ndim && (!lhs->vdevice.has_value() || vdev.defined())) { @@ -1138,7 +1138,7 @@ class TypeLCAFinder : public TypeFunctor { private: // analyzer - arith::AnalyzerObj* analyzer_; + sym::AnalyzerObj* analyzer_; // struct equal checker ffi::StructuralEqual struct_equal_; @@ -1153,11 +1153,11 @@ class TypeLCAFinder : public TypeFunctor { }; Type TypeLCA(const Type& lhs, const Type& rhs) { - arith::Analyzer analyzer; + sym::Analyzer analyzer; return TypeLCA(lhs, rhs, analyzer); } -Type TypeLCA(const Type& lhs, const Type& rhs, const arith::Analyzer& ana) { +Type TypeLCA(const Type& lhs, const Type& rhs, const sym::Analyzer& ana) { return TypeLCAFinder(ana.get())(lhs, rhs); } diff --git a/src/relax/distributed/axis_group_graph.cc b/src/relax/distributed/axis_group_graph.cc index 468f8bb58984..80cc1d5b0690 100644 --- a/src/relax/distributed/axis_group_graph.cc +++ b/src/relax/distributed/axis_group_graph.cc @@ -33,7 +33,7 @@ namespace tirx { using namespace tvm::prim; Var GetShardingVarFromIndex(PrimExpr index, ffi::Map var_range, - const arith::Analyzer& analyzer) { + const sym::Analyzer& analyzer) { if (auto prim_var = index.as()) { return prim_var.value(); } @@ -41,7 +41,7 @@ Var GetShardingVarFromIndex(PrimExpr index, ffi::Map var_range, for (const auto& [var, range] : var_range) { primitive_var_range.Set(var.as_or_throw(), range); } - arith::IterSumExpr iter_sum = arith::NormalizeToIterSum(index, primitive_var_range, analyzer); + sym::IterSumExpr iter_sum = sym::NormalizeToIterSum(index, primitive_var_range, analyzer); if (!is_zero(iter_sum->base)) { return Var(); } @@ -49,7 +49,7 @@ Var GetShardingVarFromIndex(PrimExpr index, ffi::Map var_range, return Var(); } // floormod(floordiv(source, lower_factor), extent) * scale - arith::IterSplitExpr highest_iter_split = iter_sum->args[0]; + sym::IterSplitExpr highest_iter_split = iter_sum->args[0]; auto source_var = highest_iter_split->source->source.as(); if (!source_var) { return Var(); @@ -128,7 +128,7 @@ void BuildAxisGraphBinary(const Var& output_var, const Call& call, const auto* x1_shape = x1_ty->shape.as(); const auto* x2_shape = x2_ty->shape.as(); TVM_FFI_ICHECK(x1_shape && x2_shape); - arith::Analyzer analyzer; + sym::Analyzer analyzer; for (int i = 1; i <= std::min(x1_ndim, x2_ndim); ++i) { const PrimExpr& dim0 = x1_shape->values[x1_ndim - i]; const PrimExpr& dim1 = x2_shape->values[x2_ndim - i]; @@ -241,7 +241,7 @@ void BuildAxisGraphMatmul(const Var& output_var, const Call& call, int x1_prefix_ndim = x1_shape_prefix.size(); int x2_prefix_ndim = x2_shape_prefix.size(); - arith::Analyzer analyzer; + sym::Analyzer analyzer; for (int i = 1; i <= std::min(x1_prefix_ndim, x2_prefix_ndim); ++i) { const PrimExpr& dim0 = x1_shape_prefix[x1_prefix_ndim - i]; const PrimExpr& dim1 = x2_shape_prefix[x2_prefix_ndim - i]; @@ -322,7 +322,7 @@ void BuildAxisGraphReshape(const Var& output_var, const Call& call, int i = old_shape_values.size(); int j = new_shape_values.size(); PrimExpr old_shape_product = 1, new_shape_product = 1; - arith::Analyzer analyzer_; + sym::Analyzer analyzer_; while (i > 0 && j > 0) { if (analyzer_->CanProve(new_shape_product > old_shape_product)) { i--; 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 a7ab6cf621a6..02ec388acf89 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 @@ -198,7 +198,7 @@ class DistributedBufferCompactor : public s_tir::StmtExprMutator { for (const auto& iter_var : block->iter_vars) { iter_var_range.Set(iter_var->var, iter_var->dom); } - arith::Analyzer analyzer; + sym::Analyzer analyzer; for (const auto& buffer : buffers) { if (buffer_access_indices.count(buffer) == 0 || buffer_shards_.count(buffer) == 0) { continue; @@ -224,7 +224,7 @@ class DistributedBufferCompactor : public s_tir::StmtExprMutator { if (shard > 1) { Range dom = iter_var->dom; TVM_FFI_ICHECK(is_zero(dom->min)); - arith::Analyzer analyzer; + sym::Analyzer analyzer; TVM_FFI_ICHECK(analyzer->CanProve(floormod(dom->extent, shard) == 0)); new_iter_vars.push_back( IterVar(Range::FromMinExtent(dom->min, floordiv(dom->extent, shard)), iter_var->var, @@ -321,7 +321,7 @@ class DistributedBufferCompactor : public s_tir::StmtExprMutator { if (loop_var_shards_.count(op->loop_var)) { int shard = loop_var_shards_[op->loop_var]; if (shard > 1) { - arith::Analyzer analyzer; + sym::Analyzer analyzer; TVM_FFI_ICHECK(analyzer->CanProve(floormod(new_loop->extent, shard) == 0)); new_loop.CopyOnWrite()->extent = floordiv(new_loop->extent, shard); return new_loop; diff --git a/src/relax/ir/block_builder.cc b/src/relax/ir/block_builder.cc index a5f63ca03d97..cbdf74794127 100644 --- a/src/relax/ir/block_builder.cc +++ b/src/relax/ir/block_builder.cc @@ -20,7 +20,6 @@ /*! * \file src/relax/block_builder.cc */ -#include #include #include #include @@ -34,6 +33,7 @@ #include #include #include +#include #include #include @@ -296,7 +296,7 @@ class BlockBuilderImpl : public BlockBuilderNode { } } - arith::Analyzer GetAnalyzer() final { return analyzer_; } + sym::Analyzer GetAnalyzer() final { return analyzer_; } protected: /*! @@ -353,7 +353,7 @@ class BlockBuilderImpl : public BlockBuilderNode { IRModule context_mod_; /*! \brief Internal analzyer */ - arith::Analyzer analyzer_; + sym::Analyzer analyzer_; /*! * \return The current frame. diff --git a/src/relax/ir/dataflow_block_rewriter.cc b/src/relax/ir/dataflow_block_rewriter.cc index f5fda40e3dee..0d480c38593a 100644 --- a/src/relax/ir/dataflow_block_rewriter.cc +++ b/src/relax/ir/dataflow_block_rewriter.cc @@ -22,7 +22,6 @@ * \brief A transform to match a Relax DataflowBlock and rewrite */ -#include #include #include #include @@ -32,6 +31,7 @@ #include #include #include +#include #include #include @@ -190,7 +190,7 @@ static std::optional TryMatch(const PNode& p, const RNode& r, static std::optional TryValidate( const MatchState& current_match, const std::unordered_map& pattern2node, - const std::vector& validation_constraints, arith::AnalyzerObj* analyzer) { + const std::vector& validation_constraints, sym::AnalyzerObj* analyzer) { MatchState new_match; std::function(const DFPatternNode*)> query_match_state = @@ -244,7 +244,7 @@ static std::optional MatchTree( const std::unordered_map& pattern2node, const std::unordered_map& var2node, DFPatternMatcher* matcher, const std::vector& roots, const std::vector& validation_constraints, - const MatcherUseDefAnalysis& ud_analysis, arith::AnalyzerObj* analyzer) { + const MatcherUseDefAnalysis& ud_analysis, sym::AnalyzerObj* analyzer) { auto get_next_root = [&](size_t root_idx) -> const PNode* { // Look for the next unmatched root node. for (; root_idx < roots.size(); ++root_idx) { @@ -346,7 +346,7 @@ ffi::Optional> MatchGraph(const PatternContext& ctx, return std::nullopt; } - arith::Analyzer analyzer; + sym::Analyzer analyzer; auto match = MatchTree({}, 0, pattern2node, var2node, &matcher, roots, ctx->validation_constraints, ud_analysis, analyzer.get()); if (!match) { diff --git a/src/relax/ir/dataflow_matcher.cc b/src/relax/ir/dataflow_matcher.cc index be25544b860d..d5c3d4310a57 100644 --- a/src/relax/ir/dataflow_matcher.cc +++ b/src/relax/ir/dataflow_matcher.cc @@ -24,7 +24,6 @@ #include "dataflow_matcher.h" -#include #include #include #include @@ -34,6 +33,7 @@ #include #include #include +#include #include #include @@ -48,15 +48,15 @@ #include #include -#include "../../arith/constraint_extract.h" +#include "../../sym/constraint_extract.h" #include "../transform/utils.h" namespace tvm { namespace relax { using namespace tvm::prim; -using tvm::arith::Analyzer; -using tvm::arith::AnalyzerObj; +using tvm::sym::Analyzer; +using tvm::sym::AnalyzerObj; /*! * \brief Match the attributes of an object. @@ -456,7 +456,7 @@ PrimExpr DFPatternMatcher::SimplifyCondition(PrimExpr condition) { return condition; } - std::vector constraints = arith::ExtractConstraints(condition, false); + std::vector constraints = sym::ExtractConstraints(condition, false); if (constraints.size() == 1) { return condition; } diff --git a/src/relax/ir/dataflow_matcher.h b/src/relax/ir/dataflow_matcher.h index b02bf82e177c..636bc0d539c0 100644 --- a/src/relax/ir/dataflow_matcher.h +++ b/src/relax/ir/dataflow_matcher.h @@ -24,10 +24,10 @@ #ifndef TVM_RELAX_IR_DATAFLOW_MATCHER_H_ #define TVM_RELAX_IR_DATAFLOW_MATCHER_H_ -#include #include #include #include +#include #include #include @@ -95,7 +95,7 @@ class DFPatternMatcher : public DFPatternFunctor matched_nodes_; PrimExpr symbolic_expr_condition_{IntImm::Bool(true)}; - arith::Analyzer analyzer_; + sym::Analyzer analyzer_; bool memoize_ = true; }; diff --git a/src/relax/ir/expr_functor.cc b/src/relax/ir/expr_functor.cc index 37421363943a..9445690285f3 100644 --- a/src/relax/ir/expr_functor.cc +++ b/src/relax/ir/expr_functor.cc @@ -999,7 +999,7 @@ Expr ExprMutator::VisitWithNewScope(const Expr& expr, ffi::OptionalBeginScope(params); // Outer scope only includes TIR variables that can be inferred from // the function parameters. - With context(builder_->GetAnalyzer(), constraint); + With context(builder_->GetAnalyzer(), constraint); builder_->BeginInnerScope(); // Inner scope also includes any TIR variables that are defined by // MatchCast nodes, and are internal to the scope. diff --git a/src/relax/op/ccl/ccl.cc b/src/relax/op/ccl/ccl.cc index 39a2a5e778ba..f5b660d1bcf4 100644 --- a/src/relax/op/ccl/ccl.cc +++ b/src/relax/op/ccl/ccl.cc @@ -149,7 +149,7 @@ Type InferTypeScatter(const Call& call, const BlockBuilder& ctx) { const auto* attrs = call->attrs.as(); int num_workers = attrs->num_workers; - arith::Analyzer analyzer = ctx->GetAnalyzer(); + sym::Analyzer analyzer = ctx->GetAnalyzer(); auto input_shape = input_ty->GetShape(); TVM_FFI_ICHECK(input_shape.has_value()) << "input tensor of scatter_from_worker0 should have defined shape."; diff --git a/src/relax/op/distributed/distributed.cc b/src/relax/op/distributed/distributed.cc index 69eefd20d09f..0ae767aaf058 100644 --- a/src/relax/op/distributed/distributed.cc +++ b/src/relax/op/distributed/distributed.cc @@ -149,7 +149,7 @@ Type InferTypeRtoS(const Call& call, const BlockBuilder& ctx) { const auto* attrs = call->attrs.as(); int num_workers = attrs->num_workers; - arith::Analyzer analyzer = ctx->GetAnalyzer(); + sym::Analyzer analyzer = ctx->GetAnalyzer(); auto input_shape = input_ty->GetShape(); TVM_FFI_ICHECK(input_shape.has_value()) << "input tensor of redistribute_replica_to_shard should have defined shape."; @@ -177,7 +177,7 @@ Type InferDistTypeRtoS(const Call& call, const BlockBuilder& ctx) { TensorType tensor_ty = input_dtensor_ty->tensor_ty; const auto* attrs = call->attrs.as(); int num_workers = attrs->num_workers; - arith::Analyzer analyzer = ctx->GetAnalyzer(); + sym::Analyzer analyzer = ctx->GetAnalyzer(); auto input_shape = tensor_ty->GetShape(); TVM_FFI_ICHECK(input_shape.has_value()) << "input tensor of redistribute_replica_to_shard should have defined shape."; diff --git a/src/relax/op/distributed/linear_algebra.cc b/src/relax/op/distributed/linear_algebra.cc index 75b6a97eb50e..9fc5d295834d 100644 --- a/src/relax/op/distributed/linear_algebra.cc +++ b/src/relax/op/distributed/linear_algebra.cc @@ -75,7 +75,7 @@ Type InferDistTypeMatmul(const Call& call, const BlockBuilder& ctx) { ffi::Optional> output_shape_prefix = InferBinaryBroadcastShape(call, ctx, x1_shape_prefix, x2_shape_prefix); TVM_FFI_ICHECK(output_shape_prefix.has_value()) << "Failed to infer output shape of Matmul"; - arith::Analyzer analyzer = ctx->GetAnalyzer(); + sym::Analyzer analyzer = ctx->GetAnalyzer(); PrimExpr x1_reduction_length = x1_shape->values[x1_ty->ndim - 1]; PrimExpr x2_reduction_length = x2_shape->values[x2_ndim - 2]; if (analyzer->CanProve(x1_reduction_length != x2_reduction_length)) { diff --git a/src/relax/op/distributed/utils.h b/src/relax/op/distributed/utils.h index 78ac15755811..36ac64880b4c 100644 --- a/src/relax/op/distributed/utils.h +++ b/src/relax/op/distributed/utils.h @@ -25,9 +25,9 @@ #ifndef TVM_RELAX_OP_DISTRIBUTED_UTILS_H_ #define TVM_RELAX_OP_DISTRIBUTED_UTILS_H_ -#include #include #include +#include #include "../op_common.h" diff --git a/src/relax/op/nn/attention.cc b/src/relax/op/nn/attention.cc index ce2fa5113e41..f38a0d7a2d2a 100644 --- a/src/relax/op/nn/attention.cc +++ b/src/relax/op/nn/attention.cc @@ -90,7 +90,7 @@ Type InferTypeAttention(const Call& call, const BlockBuilder& ctx) { PrimExpr head_dim = q_shape->values[3]; PrimExpr num_keys = k_shape->values[1]; PrimExpr head_dim_value = v_shape->values[3]; - arith::Analyzer analyzer = ctx->GetAnalyzer(); + sym::Analyzer analyzer = ctx->GetAnalyzer(); auto diag_equal = [&](PrimExpr v1, PrimExpr v2, ffi::String m1, ffi::String m2, ffi::String dim) { if (analyzer->CanProve(v1 != v2)) { TVM_FFI_VISIT_THROW(ValueError, call) diff --git a/src/relax/op/nn/convolution.cc b/src/relax/op/nn/convolution.cc index b55bdf90bdfd..1927a7d5bccd 100644 --- a/src/relax/op/nn/convolution.cc +++ b/src/relax/op/nn/convolution.cc @@ -104,7 +104,7 @@ Type InferTypeConv1d(const Call& call, const BlockBuilder& ctx) { ffi::Array data_NCW_shape = data2NCW.ForwardShape(data_shape.value()->values); ffi::Array weight_OIW_shape = weight2OIW.ForwardShape(weight_shape.value()->values); - arith::Analyzer analyzer = ctx->GetAnalyzer(); + sym::Analyzer analyzer = ctx->GetAnalyzer(); PrimExpr input_channel_data = data_NCW_shape[1]; PrimExpr input_channel_kernel = weight_OIW_shape[1]; if (analyzer->CanProve(input_channel_data != input_channel_kernel * attrs->groups)) { @@ -274,7 +274,7 @@ Type InferTypeConv2d(const Call& call, const BlockBuilder& ctx) { ffi::Array data_NCHW_shape = data2NCHW.ForwardShape(data_shape.value()->values); ffi::Array weight_OIHW_shape = weight2OIHW.ForwardShape(weight_shape.value()->values); - arith::Analyzer analyzer = ctx->GetAnalyzer(); + sym::Analyzer analyzer = ctx->GetAnalyzer(); PrimExpr input_channel_data = data_NCHW_shape[1]; PrimExpr input_channel_kernel = weight_OIHW_shape[1]; if (analyzer->CanProve(input_channel_data != input_channel_kernel * attrs->groups)) { @@ -488,7 +488,7 @@ Type InferTypeConv3d(const Call& call, const BlockBuilder& ctx) { ffi::Array data_NCDHW_shape = data2NCDHW.ForwardShape(data_shape.value()->values); ffi::Array weight_OIDHW_shape = weight2OIDHW.ForwardShape(weight_shape.value()->values); - arith::Analyzer analyzer = ctx->GetAnalyzer(); + sym::Analyzer analyzer = ctx->GetAnalyzer(); PrimExpr input_channel_data = data_NCDHW_shape[1]; PrimExpr input_channel_kernel = weight_OIDHW_shape[1]; if (analyzer->CanProve(input_channel_data != input_channel_kernel * attrs->groups)) { @@ -676,7 +676,7 @@ Type InferTypeConv1dTranspose(const Call& call, const BlockBuilder& ctx) { ffi::Array data_NCW_shape = data2NCW.ForwardShape(data_shape.value()->values); ffi::Array weight_IOW_shape = weight2IOW.ForwardShape(weight_shape.value()->values); - arith::Analyzer analyzer = ctx->GetAnalyzer(); + sym::Analyzer analyzer = ctx->GetAnalyzer(); PrimExpr input_channel_data = data_NCW_shape[1]; PrimExpr input_channel_kernel = weight_IOW_shape[0]; if (analyzer->CanProve(input_channel_data != input_channel_kernel)) { @@ -868,7 +868,7 @@ Type InferTypeConv2dTranspose(const Call& call, const BlockBuilder& ctx) { ffi::Array data_NCHW_shape = data2NCHW.ForwardShape(data_shape.value()->values); ffi::Array weight_IOHW_shape = weight2IOHW.ForwardShape(weight_shape.value()->values); - arith::Analyzer analyzer = ctx->GetAnalyzer(); + sym::Analyzer analyzer = ctx->GetAnalyzer(); PrimExpr input_channel_data = data_NCHW_shape[1]; PrimExpr input_channel_kernel = weight_IOHW_shape[0]; if (analyzer->CanProve(input_channel_data != input_channel_kernel)) { @@ -1100,7 +1100,7 @@ Type InferTypeConv3dTranspose(const Call& call, const BlockBuilder& ctx) { ffi::Array data_NCDHW_shape = data2NCDHW.ForwardShape(data_shape.value()->values); ffi::Array weight_IODHW_shape = weight2IODHW.ForwardShape(weight_shape.value()->values); - arith::Analyzer analyzer = ctx->GetAnalyzer(); + sym::Analyzer analyzer = ctx->GetAnalyzer(); PrimExpr input_channel_data = data_NCDHW_shape[1]; PrimExpr input_channel_kernel = weight_IODHW_shape[0]; if (analyzer->CanProve(input_channel_data != input_channel_kernel)) { diff --git a/src/relax/op/nn/nn.cc b/src/relax/op/nn/nn.cc index cd6afe02b035..1da6baedd638 100644 --- a/src/relax/op/nn/nn.cc +++ b/src/relax/op/nn/nn.cc @@ -425,7 +425,7 @@ bool NormCheckDtypeAndShape(const Call& call, const BlockBuilder& ctx, } } - arith::Analyzer analyzer = ctx->GetAnalyzer(); + sym::Analyzer analyzer = ctx->GetAnalyzer(); for (int i = 1; i < static_cast(axis_lengths.size()); ++i) { for (int d = 0; d < n_axis; ++d) { if (analyzer->CanProve(axis_lengths[0][d] != axis_lengths[i][d])) { @@ -636,7 +636,7 @@ Type InferTypeGroupNorm(const Call& call, const BlockBuilder& ctx) { TVM_FFI_VISIT_THROW(TypeError, call) << op << " expects that data must be float, but got " << data_ty->dtype; } - arith::Analyzer analyzer = ctx->GetAnalyzer(); + sym::Analyzer analyzer = ctx->GetAnalyzer(); const auto* data_shape = data_ty->shape.as(); if (data_shape != nullptr && channel_axis != -1 && analyzer->CanProve(floormod(data_shape->values[channel_axis], attrs->num_groups) != 0)) { @@ -746,7 +746,7 @@ Type InferTypeInstanceNorm(const Call& call, const BlockBuilder& ctx) { } } const auto* data_shape = data_ty->shape.as(); - arith::Analyzer analyzer = ctx->GetAnalyzer(); + sym::Analyzer analyzer = ctx->GetAnalyzer(); for (int i = 1; i < static_cast(op->arguments.size()); ++i) { if (input_ty[i]->dtype != data_ty->dtype) { TVM_FFI_VISIT_THROW(TypeError, call) @@ -927,7 +927,7 @@ Type InferTypeCrossEntropy(const Call& call, const BlockBuilder& ctx) { } if (pred_shape_value.has_value() && label_shape_value.has_value()) { - arith::Analyzer analyzer = ctx->GetAnalyzer(); + sym::Analyzer analyzer = ctx->GetAnalyzer(); for (size_t i = 0; i < pred_shape_value.value().size(); ++i) { if (analyzer->CanProve(pred_shape_value.value()[i] != label_shape_value.value()[i])) { TVM_FFI_VISIT_THROW(ValueError, call) @@ -1062,7 +1062,7 @@ Type InferTypeNLLLoss(const Call& call, const BlockBuilder& ctx) { << wgt_ty->ndim; } - arith::Analyzer analyzer = ctx->GetAnalyzer(); + sym::Analyzer analyzer = ctx->GetAnalyzer(); ffi::Optional N; ffi::Optional C; ffi::Array output_shape; // N, d1, d2, ..., dk diff --git a/src/relax/op/nn/pooling.cc b/src/relax/op/nn/pooling.cc index 6c0da988e35a..50a45d53d0ee 100644 --- a/src/relax/op/nn/pooling.cc +++ b/src/relax/op/nn/pooling.cc @@ -103,7 +103,7 @@ Type InferTypePool1D(const Call& call, const BlockBuilder& ctx) { PrimExpr kernel_w = IntImm::Int32(attrs->pool_size[0]); PrimExpr padding_w = IntImm::Int32(attrs->padding[0]) + IntImm::Int32(attrs->padding[1]); - arith::Analyzer analyzer = ctx->GetAnalyzer(); + sym::Analyzer analyzer = ctx->GetAnalyzer(); std::vector out_NCW_shape; out_NCW_shape.resize(3); out_NCW_shape[0] = data_NCW_shape[0]; @@ -230,7 +230,7 @@ Type InferTypePool2D(const Call& call, const BlockBuilder& ctx) { PrimExpr padding_h = IntImm::Int32(attrs->padding[0]) + IntImm::Int32(attrs->padding[2]); PrimExpr padding_w = IntImm::Int32(attrs->padding[1]) + IntImm::Int32(attrs->padding[3]); - arith::Analyzer analyzer = ctx->GetAnalyzer(); + sym::Analyzer analyzer = ctx->GetAnalyzer(); std::vector out_NCHW_shape; out_NCHW_shape.resize(4); out_NCHW_shape[0] = data_NCHW_shape[0]; @@ -390,7 +390,7 @@ Type InferTypePool3D(const Call& call, const BlockBuilder& ctx) { PrimExpr padding_h = IntImm::Int32(attrs->padding[1]) + IntImm::Int32(attrs->padding[4]); PrimExpr padding_w = IntImm::Int32(attrs->padding[2]) + IntImm::Int32(attrs->padding[5]); - arith::Analyzer analyzer = ctx->GetAnalyzer(); + sym::Analyzer analyzer = ctx->GetAnalyzer(); std::vector out_NCDHW_shape; out_NCDHW_shape.resize(5); out_NCDHW_shape[0] = data_NCDHW_shape[0]; diff --git a/src/relax/op/op.cc b/src/relax/op/op.cc index 19505a617501..4c344c18cafb 100644 --- a/src/relax/op/op.cc +++ b/src/relax/op/op.cc @@ -52,7 +52,7 @@ bool EqualCheck(const PrimExpr& lhs, const PrimExpr& rhs) { if (const auto* pdiff = diff.as()) { return pdiff->value == 0; } - tvm::arith::Analyzer ana; + tvm::sym::Analyzer ana; diff = ana->Simplify(diff); if (const auto* pdiff = diff.as()) { return pdiff->value == 0; diff --git a/src/relax/op/op_common.cc b/src/relax/op/op_common.cc index ec219f17c737..0c656f9f4251 100644 --- a/src/relax/op/op_common.cc +++ b/src/relax/op/op_common.cc @@ -106,7 +106,7 @@ ffi::Array GetTensorTypeFromTuple(const Call& call, const BlockBuild return tensor_ty; } -BinaryBroadcastShapeInferResult InferBinaryBroadcastShape(arith::AnalyzerObj* analyzer, +BinaryBroadcastShapeInferResult InferBinaryBroadcastShape(sym::AnalyzerObj* analyzer, const ffi::Array& x1_shape, const ffi::Array& x2_shape) { BinaryBroadcastShapeInferResult result; @@ -217,7 +217,7 @@ bool CanProveLayoutTransform(const SLayout& input_layout, const SLayout& desired tirx::SBijectiveLayout todesired(input_layout, desired_layout); ffi::Array desired_shape = todesired.ForwardShape(shape); ffi::Array back_shape = todesired.BackwardShape(desired_shape); - arith::Analyzer analyzer; + sym::Analyzer analyzer; for (size_t i = 0; i < shape.size(); ++i) { if (tvm::prim::is_const_int(shape[i])) { if (!analyzer->CanProveEqual(shape[i], back_shape[i])) { diff --git a/src/relax/op/op_common.h b/src/relax/op/op_common.h index bec5ffcd4ab5..953892c1f112 100644 --- a/src/relax/op/op_common.h +++ b/src/relax/op/op_common.h @@ -25,11 +25,11 @@ #ifndef TVM_RELAX_OP_OP_COMMON_H_ #define TVM_RELAX_OP_OP_COMMON_H_ -#include #include #include #include #include +#include #include #include @@ -413,7 +413,7 @@ struct BinaryBroadcastShapeInferResult { * \param x2_shape The shape of the second operand. * \return Inference status and broadcasted shape, or a conflict message. */ -BinaryBroadcastShapeInferResult InferBinaryBroadcastShape(arith::AnalyzerObj* analyzer, +BinaryBroadcastShapeInferResult InferBinaryBroadcastShape(sym::AnalyzerObj* analyzer, const ffi::Array& x1_shape, const ffi::Array& x2_shape); diff --git a/src/relax/op/tensor/create.cc b/src/relax/op/tensor/create.cc index 1dd93cf3a501..d111f914ffeb 100644 --- a/src/relax/op/tensor/create.cc +++ b/src/relax/op/tensor/create.cc @@ -24,10 +24,10 @@ #include "create.h" -#include #include #include #include +#include #include #include @@ -383,7 +383,7 @@ Type InferTypeArange(const Call& call, const BlockBuilder& ctx) { tvm::prim::cast(tvm::PrimType::Int(64), tvm::ceil(tvm::prim::cast(tvm::PrimType::Float(32), end - start) / step)); } - arith::Analyzer analyzer; + sym::Analyzer analyzer; num_elem = analyzer->Simplify(num_elem); return TensorType(ShapeExpr({num_elem}), PrimType(dtype)); } @@ -432,7 +432,7 @@ Type InferTypeHammingWindow(const Call& call, const BlockBuilder& ctx) { }; PrimExpr window_size = get_prim_value(call->args[0], "window_size"); - arith::Analyzer analyzer; + sym::Analyzer analyzer; if (analyzer->CanProveLess(window_size, 1)) { TVM_FFI_VISIT_THROW(ValueError, call) << "Hamming_window expects the window_size must be greater than zero but got " diff --git a/src/relax/op/tensor/index.cc b/src/relax/op/tensor/index.cc index dcfe2b7f6b5f..d5aa64ab8687 100644 --- a/src/relax/op/tensor/index.cc +++ b/src/relax/op/tensor/index.cc @@ -429,8 +429,8 @@ Type InferTypeStridedSlice(const Call& call, const BlockBuilder& ctx) { PrimExpr output_dim = topi::GetLength(begin, end, strides_tuple[i], input_dim, attrs->assume_inbound); - arith::Analyzer analyzer = ctx->GetAnalyzer(); - std::optional> context; + sym::Analyzer analyzer = ctx->GetAnalyzer(); + std::optional> context; if (attrs->assume_inbound) { context.emplace(analyzer, 0 <= begin && begin <= input_dim && 0 <= end && end <= input_dim); } diff --git a/src/relax/op/tensor/linear_algebra.cc b/src/relax/op/tensor/linear_algebra.cc index de96cc2a8671..3a544a3ceac1 100644 --- a/src/relax/op/tensor/linear_algebra.cc +++ b/src/relax/op/tensor/linear_algebra.cc @@ -133,7 +133,7 @@ Type InferTypeMatmul(const Call& call, const BlockBuilder& ctx) { return TensorType(out_dtype, output_ndim); } - arith::Analyzer analyzer = ctx->GetAnalyzer(); + sym::Analyzer analyzer = ctx->GetAnalyzer(); PrimExpr x1_reduction_length = x1_shape->values[x1_ty->ndim - 1]; PrimExpr x2_reduction_length = x2_shape->values[x2_ndim - 2]; if (analyzer->CanProve(x1_reduction_length != x2_reduction_length)) { diff --git a/src/relax/op/tensor/manipulate.cc b/src/relax/op/tensor/manipulate.cc index 56b130fbbcc5..847cf1f3449d 100644 --- a/src/relax/op/tensor/manipulate.cc +++ b/src/relax/op/tensor/manipulate.cc @@ -107,7 +107,7 @@ Type InferTypeBroadcastTo(const Call& call, const BlockBuilder& ctx) { return TensorType(/*shape=*/call->args[1], data_ty->dtype, data_ty->vdevice); } - arith::Analyzer analyzer = ctx->GetAnalyzer(); + sym::Analyzer analyzer = ctx->GetAnalyzer(); ffi::Array old_shape_value = shape_ty->values.value(); ffi::Array tgt_shape_value = tgt_shape_ty->values.value(); int old_ndim = old_shape_value.size(); @@ -158,7 +158,7 @@ ffi::Optional> CheckConcatOutputShape( const Call& call, const BlockBuilder& ctx, const std::vector>& shape_values, int axis) { bool shape_unknown = false; - arith::Analyzer analyzer = ctx->GetAnalyzer(); + sym::Analyzer analyzer = ctx->GetAnalyzer(); PrimExpr concat_sum = [&]() { // For the specified axis, we compute the sum of shape value over each tensor. @@ -595,7 +595,7 @@ Type InferTypeIndexTensor(const Call& call, const BlockBuilder& ctx) { << data_ty->ndim << " dimensions"; } - arith::Analyzer analyzer = ctx->GetAnalyzer(); + sym::Analyzer analyzer = ctx->GetAnalyzer(); bool all_index_have_shape_value = true; std::vector> index_shapes; int max_index_ndim = 0; @@ -752,7 +752,7 @@ Type InferTypeLayoutTransform(const Call& call, const BlockBuilder& ctx) { return TensorType(data_ty->dtype, /*ndim=*/index_map->final_indices.size(), data_ty->vdevice); } - arith::Analyzer analyzer; + sym::Analyzer analyzer; ffi::Array output_shape = index_map->MapShape(shape_ty->values.value(), analyzer); return TensorType(ShapeExpr(output_shape), data_ty->dtype, data_ty->vdevice); } @@ -981,7 +981,7 @@ Expr ConvertNewShapeToExpr(const Expr& data, // Assign appropriate value to -1 dimension. if (dim_to_infer != -1) { - arith::Analyzer analyzer; + sym::Analyzer analyzer; PrimExpr old_shape_prod = ComputeShapeProduct(shape_ty->values.value()); array_ref.Set(dim_to_infer, analyzer->Simplify(floordiv(old_shape_prod, new_shape_prod))); } @@ -1386,7 +1386,7 @@ TVM_REGISTER_OP("relax.squeeze") void CheckCollapseShape(const Call& call, const BlockBuilder& ctx, const ffi::Array& data_shape, const ffi::Array& target_shape) { - arith::Analyzer analyzer = ctx->GetAnalyzer(); + sym::Analyzer analyzer = ctx->GetAnalyzer(); int data_ndim = data_shape.size(); int target_ndim = target_shape.size(); @@ -1441,7 +1441,7 @@ ffi::Optional> CheckStackOutputShape( const Call& call, const BlockBuilder& ctx, const std::vector>& shape_values, int axis) { bool shape_unknown = false; - arith::Analyzer analyzer = ctx->GetAnalyzer(); + sym::Analyzer analyzer = ctx->GetAnalyzer(); // Stack requires all input tensors to have identical shapes for (int d = 0; d < static_cast(shape_values[0].size()); ++d) { @@ -1749,7 +1749,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { } Type InferTypeRepeat(const Call& call, const BlockBuilder& ctx) { - arith::Analyzer analyzer = ctx->GetAnalyzer(); + sym::Analyzer analyzer = ctx->GetAnalyzer(); TensorType data_ty = GetUnaryInputTensorType(call, ctx); const auto* attrs = call->attrs.as(); const auto* data_shape = data_ty->shape.as(); @@ -1873,7 +1873,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { } Type InferTypeTile(const Call& call, const BlockBuilder& ctx) { - arith::Analyzer analyzer = ctx->GetAnalyzer(); + sym::Analyzer analyzer = ctx->GetAnalyzer(); TensorType data_ty = GetUnaryInputTensorType(call, ctx); const auto* attrs = call->attrs.as(); const auto* data_shape = data_ty->shape.as(); @@ -2633,7 +2633,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { } Type InferTypeScatterElements(const Call& call, const BlockBuilder& ctx) { - arith::Analyzer analyzer = ctx->GetAnalyzer(); + sym::Analyzer analyzer = ctx->GetAnalyzer(); const auto* data_ty = GetTypeAs(call->args[0]); const auto* indices_ty = GetTypeAs(call->args[1]); const auto* updates_ty = GetTypeAs(call->args[2]); @@ -2778,7 +2778,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { Type InferTypeScatterND(const Call& call, const BlockBuilder& ctx) { // `call->args` contains: [data, indices, updates] - arith::Analyzer analyzer = ctx->GetAnalyzer(); + sym::Analyzer analyzer = ctx->GetAnalyzer(); TVM_FFI_ICHECK_EQ(call->args.size(), 3); const auto* data_ty = GetTypeAs(call->args[0]); const auto* indices_ty = GetTypeAs(call->args[1]); @@ -2955,7 +2955,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { } Type InferTypeSliceScatter(const Call& call, const BlockBuilder& ctx) { - arith::Analyzer analyzer = ctx->GetAnalyzer(); + sym::Analyzer analyzer = ctx->GetAnalyzer(); const auto* data_ty = GetTypeAs(call->args[0]); const auto* src_ty = GetTypeAs(call->args[1]); auto* attrs = call->attrs.as(); diff --git a/src/relax/op/tensor/sampling.cc b/src/relax/op/tensor/sampling.cc index 0a8ba017709b..f497af5853ed 100644 --- a/src/relax/op/tensor/sampling.cc +++ b/src/relax/op/tensor/sampling.cc @@ -123,7 +123,7 @@ Type InferTypeMultinomialFromUniform(const Call& call, const BlockBuilder& ctx) PrimExpr batch = prob_shape->values[0]; PrimExpr n = uniform_sample_shape->values[0]; - arith::Analyzer ana; + sym::Analyzer ana; if (!ana->CanProveEqual(n, sample_indices_shape->values[0])) { TVM_FFI_VISIT_THROW(ValueError, call) << "Multinomial_from_uniform op requires the input uniform_sample and " diff --git a/src/relax/op/tensor/ternary.cc b/src/relax/op/tensor/ternary.cc index 46cd2eead73b..0672427a4c9d 100644 --- a/src/relax/op/tensor/ternary.cc +++ b/src/relax/op/tensor/ternary.cc @@ -85,7 +85,7 @@ Type InferTypeEwiseFMA(const Call& call, const BlockBuilder& ctx) { auto* s1 = t1->shape.as(); auto* s2 = t2->shape.as(); auto* s3 = t3->shape.as(); - arith::Analyzer analyzer = ctx->GetAnalyzer(); + sym::Analyzer analyzer = ctx->GetAnalyzer(); if (s1 && s2 && s3) { ffi::Array output_shape; for (int i = 0; i < ndim; ++i) { diff --git a/src/relax/op/vision/nms.cc b/src/relax/op/vision/nms.cc index e9ae9e9e119e..aaa097d899b5 100644 --- a/src/relax/op/vision/nms.cc +++ b/src/relax/op/vision/nms.cc @@ -18,7 +18,6 @@ */ #include "nms.h" -#include #include #include #include @@ -28,6 +27,7 @@ #include #include #include +#include #include #include @@ -266,7 +266,7 @@ Type InferTypeNMS(const Call& call, const BlockBuilder& ctx) { const auto* valid_count_shape = valid_count_ty->shape.as(); const auto* indices_shape = indices_ty->shape.as(); if (data_shape != nullptr) { - arith::Analyzer analyzer = ctx->GetAnalyzer(); + sym::Analyzer analyzer = ctx->GetAnalyzer(); PrimExpr batch = data_shape->values[0]; PrimExpr num_anchors = data_shape->values[1]; if (valid_count_shape != nullptr && diff --git a/src/relax/transform/adjust_matmul_order.cc b/src/relax/transform/adjust_matmul_order.cc index 3904b0dd2ec2..bdd28054d86d 100644 --- a/src/relax/transform/adjust_matmul_order.cc +++ b/src/relax/transform/adjust_matmul_order.cc @@ -85,7 +85,7 @@ bool IsLastTwoDimsSwap(const Expr& expr) { } ffi::Optional> InferBatchedMatmulBroadcastPrefix( - arith::AnalyzerObj* analyzer, const ffi::Array& x1, const ffi::Array& x2) { + sym::AnalyzerObj* analyzer, const ffi::Array& x1, const ffi::Array& x2) { auto infer_result = InferBinaryBroadcastShape(analyzer, x1, x2); if (infer_result.status == BinaryBroadcastShapeInferResult::Status::kSuccess) { return infer_result.shape; @@ -276,7 +276,7 @@ std::tuple)>> PrimExpr size_M = shape_c[shape_c.size() - 2]; // row of C and col of B PrimExpr size_B = shape_c[shape_c.size() - 1]; // col of C - arith::Analyzer analyzer; + sym::Analyzer analyzer; auto prefix_a = GetBatchPrefix(shape_a); auto prefix_b = GetBatchPrefix(shape_b); auto prefix_c = GetBatchPrefix(shape_c); @@ -312,12 +312,11 @@ std::tuple)>> PrimExpr ops_with_rhs_first = batch_bc * size_R * size_M * size_B + batch_outer_rhs * size_N * size_R * size_B; - analyzer->rewrite_simplify.SetEnabledExtensions( - static_cast( - analyzer->rewrite_simplify.GetEnabledExtensions() | - arith::RewriteSimplifier::Extension::kComparisonOfProductAndSum)); - With func_attr_constraint(analyzer, symbolic_var_constraints); - With analyzer_constraint( + analyzer->rewrite_simplify.SetEnabledExtensions(static_cast( + analyzer->rewrite_simplify.GetEnabledExtensions() | + sym::RewriteSimplifier::Extension::kComparisonOfProductAndSum)); + With func_attr_constraint(analyzer, symbolic_var_constraints); + With analyzer_constraint( analyzer, batch_ab > 0 && batch_bc > 0 && batch_outer_lhs > 0 && batch_outer_rhs > 0 && size_N > 0 && size_R > 0 && size_M > 0 && size_B > 0); diff --git a/src/relax/transform/alter_op_impl.cc b/src/relax/transform/alter_op_impl.cc index bb7f20b27894..bd1b4ca45541 100644 --- a/src/relax/transform/alter_op_impl.cc +++ b/src/relax/transform/alter_op_impl.cc @@ -23,7 +23,6 @@ * identify PrimFuncs to be replaced. Marks the new PrimFuncs with kFrozenLayout attribute set to * true. */ -#include #include #include #include @@ -32,6 +31,7 @@ #include #include #include +#include #include #include #include @@ -67,7 +67,7 @@ static IndexMap DeepCopyIndexMap(const IndexMap& index_map) { bool IsTransformBijective(const Expr& expr, const IndexMap& transform) { ffi::Array input_shape = GetShapeFromTensor(expr); ffi::Array initial_ranges = ConstructRangeFromShape(input_shape); - arith::Analyzer analyzer; + sym::Analyzer analyzer; auto [inverse, padding_predicate] = transform.NonSurjectiveInverse(initial_ranges, analyzer); (void)inverse; // to avoid unused variable warning; if (!analyzer->CanProve(!padding_predicate)) return false; @@ -234,7 +234,7 @@ class AlterOpImplMutator : public ExprMutator { } ffi::Array old_shape = GetShapeFromTensorType(old_tensor_ty); ffi::Array initial_ranges = ConstructRangeFromShape(old_shape); - arith::Analyzer analyzer; + sym::Analyzer analyzer; auto [inverse_index_map, padding_predicate] = index_map.NonSurjectiveInverse(initial_ranges, analyzer); @@ -317,7 +317,7 @@ class AlterOpImplMutator : public ExprMutator { Type UpdateOutputType(const TensorType& tensor_ty, const IndexMap& transform) { if (transform.get() == nullptr) return tensor_ty; auto shape = GetShapeFromTensorType(tensor_ty); - arith::Analyzer analyzer; + sym::Analyzer analyzer; auto new_shape = transform->MapShape(shape, analyzer); if (tensor_ty->vdevice.has_value()) { return TensorType(ShapeExpr(new_shape), tensor_ty->dtype, tensor_ty->vdevice.value()); diff --git a/src/relax/transform/bind_params.cc b/src/relax/transform/bind_params.cc index bbe15e7b5e98..78269ecce176 100644 --- a/src/relax/transform/bind_params.cc +++ b/src/relax/transform/bind_params.cc @@ -91,7 +91,7 @@ ffi::Map NormalizeBindings(const Function& func, relax_var_remap.Set(normalize_key(key), normalize_value(value)); } - arith::Analyzer analyzer; + sym::Analyzer analyzer; return InferSymbolicVarMap(relax_var_remap, analyzer); } diff --git a/src/relax/transform/combine_parallel_matmul.cc b/src/relax/transform/combine_parallel_matmul.cc index f892c6be0760..9d7110f5a94d 100644 --- a/src/relax/transform/combine_parallel_matmul.cc +++ b/src/relax/transform/combine_parallel_matmul.cc @@ -17,7 +17,6 @@ * under the License. */ -#include #include #include #include @@ -25,6 +24,7 @@ #include #include #include +#include #include #include @@ -120,7 +120,7 @@ ffi::TypedFunction(ffi::Map, ffi::Map& indices, const std::vector>& rhs_shapes) { - arith::Analyzer ana; + sym::Analyzer ana; for (auto ind : indices) { TVM_FFI_ICHECK_EQ(static_cast(rhs_shapes[ind].size()), rhs_dim); // -2 for reduction and concat axes diff --git a/src/relax/transform/fuse_tir.cc b/src/relax/transform/fuse_tir.cc index 925bbda6932a..387f830a397f 100644 --- a/src/relax/transform/fuse_tir.cc +++ b/src/relax/transform/fuse_tir.cc @@ -42,7 +42,7 @@ using namespace tvm::prim; */ class SymbolicMatcher : ExprFunctor { public: - explicit SymbolicMatcher(arith::AnalyzerObj* analyzer, ffi::Map* var_remap) + explicit SymbolicMatcher(sym::AnalyzerObj* analyzer, ffi::Map* var_remap) : analyzer_(analyzer), var_remap_(var_remap) {} void Match(const ffi::Array& params, const ffi::Array& args) { @@ -163,7 +163,7 @@ class SymbolicMatcher : ExprFunctor } } - arith::AnalyzerObj* analyzer_; + sym::AnalyzerObj* analyzer_; ffi::Map* var_remap_; PrimExpr must_prove_ = IntImm::Bool(true); }; @@ -951,7 +951,7 @@ class FusedTIRConstructor : public ExprVisitor { * `symbolic_var_matcher`, and must be before it in the struct * order. */ - arith::Analyzer analyzer; + sym::Analyzer analyzer; /*! \brief The map from symbolic var to its corresponding var in the fused function */ tirx::SymbolicMatcher symbolic_var_matcher = diff --git a/src/relax/transform/kill_after_last_use.cc b/src/relax/transform/kill_after_last_use.cc index 78f93f8bfacc..2222a5423a0d 100644 --- a/src/relax/transform/kill_after_last_use.cc +++ b/src/relax/transform/kill_after_last_use.cc @@ -20,13 +20,13 @@ * \file src/relax/transform/kill_after_last_use.cc * \brief Kill storage/tensor objects after last use, if not already killed */ -#include #include #include #include #include #include #include +#include #include #include diff --git a/src/relax/transform/remove_unused_parameters.cc b/src/relax/transform/remove_unused_parameters.cc index 866875d73cb6..c61f01a94258 100644 --- a/src/relax/transform/remove_unused_parameters.cc +++ b/src/relax/transform/remove_unused_parameters.cc @@ -129,7 +129,7 @@ std::optional AnalyzeCallee(Function func) { for (size_t i = 0; i < old_relax_params.size(); i++) { old_binding.Set(old_relax_params[i], old_args[i]); } - arith::Analyzer analyzer; + sym::Analyzer analyzer; auto tir_binding = InferSymbolicVarMap(old_binding, analyzer); for (const auto& tir_var : free_tir_vars) { diff --git a/src/relax/transform/rewrite_dataflow_reshape.cc b/src/relax/transform/rewrite_dataflow_reshape.cc index e99f10219d2a..b659487bec48 100644 --- a/src/relax/transform/rewrite_dataflow_reshape.cc +++ b/src/relax/transform/rewrite_dataflow_reshape.cc @@ -20,13 +20,13 @@ * \file src/relax/transform/rewrite_dataflow_reshape.cc * \brief Transform all reshape within dataflow block to a relax.reshape operator */ -#include #include #include #include #include #include #include +#include #include #include @@ -149,7 +149,7 @@ class DataflowReshapeRewriter : public ExprMutator { }; auto inp_count = product(inp_ty->GetShape().value()); auto res_count = product(res_ty->GetShape().value()); - if (!arith::Analyzer()->CanProveEqual(inp_count, res_count)) { + if (!sym::Analyzer()->CanProveEqual(inp_count, res_count)) { return false; } diff --git a/src/relax/transform/split_call_tir_by_pattern.cc b/src/relax/transform/split_call_tir_by_pattern.cc index 54bf928e7121..53a0812001aa 100644 --- a/src/relax/transform/split_call_tir_by_pattern.cc +++ b/src/relax/transform/split_call_tir_by_pattern.cc @@ -20,7 +20,6 @@ * \file src/relax/transform/to_non_dataflow.cc * \brief Transform all dataflow structure to non-dataflow version. */ -#include #include #include #include @@ -32,6 +31,7 @@ #include #include #include +#include #include #include @@ -385,7 +385,7 @@ class ForMatcher : public TensorizeComparator { return true; } - arith::Analyzer analyzer_; + sym::Analyzer analyzer_; std::vector loop_stack_lhs_, loop_stack_rhs_; tirx::PrimFunc pattern_; std::unordered_set pattern_vars_; diff --git a/src/relax/transform/static_plan_block_memory.cc b/src/relax/transform/static_plan_block_memory.cc index e5fe2c6c802f..53059da3e6af 100644 --- a/src/relax/transform/static_plan_block_memory.cc +++ b/src/relax/transform/static_plan_block_memory.cc @@ -65,13 +65,13 @@ * signature will have upper bound 1024. And we will use 1024 as its value * during memory planning. */ -#include #include #include #include #include #include #include +#include #include #include @@ -200,7 +200,7 @@ using Tokens = NestedMsg; */ class TokenAllocatorMixed { public: - explicit TokenAllocatorMixed(arith::AnalyzerObj* analyzer) : analyzer_(analyzer) {} + explicit TokenAllocatorMixed(sym::AnalyzerObj* analyzer) : analyzer_(analyzer) {} /*! * \brief Request a storage token from the available token pool for a @@ -322,7 +322,7 @@ class TokenAllocatorMixed { }; /*! \brief The arithmetic analyzer. */ - arith::AnalyzerObj* analyzer_; + sym::AnalyzerObj* analyzer_; /*! \brief A constant scale representing the token search range. */ const int match_range_{16}; /*! \brief The pool of available storage tokens for each storage scope and dtype. */ @@ -424,8 +424,8 @@ class StorageAllocatorBaseVisitor : public ExprVisitor { * \param ana The analyzer which contains the TIR var upper bounds. * \param dom_map The domain map of the TIR variables. */ -void SetTIRVarRangeConstraints(Function func, arith::AnalyzerObj* ana, - ffi::Map* dom_map) { +void SetTIRVarRangeConstraints(Function func, sym::AnalyzerObj* ana, + ffi::Map* dom_map) { // Use the attribute-annotated TIR var bounds as the TIR var values for // memory planning. // NOTE: we only apply the annotated bounds to the TIR variables that @@ -466,7 +466,7 @@ void SetTIRVarRangeConstraints(Function func, arith::AnalyzerObj* ana, tvm::Range range = tvm::Range::FromMinExtent(tvm::IntImm::Int64(lower), tvm::IntImm::Int64(upper - lower + 1)); ana->Bind(tir_var, range); - dom_map->Set(tir_var, arith::IntSet::FromRange(range)); + dom_map->Set(tir_var, sym::IntSet::FromRange(range)); } else if (it_lower != var_lower_bound_attr.end() && it_lower->second->value >= 0) { ana->MarkGlobalNonNegValue(tir_var.as_or_throw()); } else if (non_negative_var_attr.count(tir_var->name)) { @@ -484,15 +484,15 @@ void SetTIRVarRangeConstraints(Function func, arith::AnalyzerObj* ana, * \return The upper-bounded shape. When a dimension's upper bound * cannot be determined, we keep the dimension unchanged. */ -ffi::Array GetUpperBoundShape(ffi::Array shape, arith::AnalyzerObj* ana, - const ffi::Map& dom_map) { +ffi::Array GetUpperBoundShape(ffi::Array shape, sym::AnalyzerObj* ana, + const ffi::Map& dom_map) { // Use the upper bounds of TIR vars as their values. ffi::Array upper_bounded_shape; upper_bounded_shape.reserve(shape.size()); for (const PrimExpr& dim_len : shape) { int64_t max_bound = ana->const_int_bound(dim_len)->max_value; if (max_bound == std::numeric_limits::max()) { - arith::IntSet int_set = ana->int_set(dim_len, dom_map); + sym::IntSet int_set = ana->int_set(dim_len, dom_map); if (int_set.HasUpperBound()) { upper_bounded_shape.push_back(int_set.max()); } else { @@ -533,7 +533,7 @@ class StorageAllocatorInit : public StorageAllocatorBaseVisitor { * \return The mapping from each Expr to the token it uses. */ static std::unordered_map Initialize(const IRModule& mod, - arith::AnalyzerObj* analyzer) { + sym::AnalyzerObj* analyzer) { StorageAllocatorInit initializer(mod, analyzer); for (auto it : mod->functions) { @@ -549,7 +549,7 @@ class StorageAllocatorInit : public StorageAllocatorBaseVisitor { private: using ExprVisitor::VisitExpr_; - explicit StorageAllocatorInit(const IRModule& ctx_mod, arith::AnalyzerObj* analyzer) + explicit StorageAllocatorInit(const IRModule& ctx_mod, sym::AnalyzerObj* analyzer) : ctx_mod_(ctx_mod), analyzer_(analyzer) {} void VisitExpr_(const FunctionNode* func) final { @@ -742,9 +742,9 @@ class StorageAllocatorInit : public StorageAllocatorBaseVisitor { */ const IRModule& ctx_mod_; /*! \brief The arithmetic analyzer. */ - arith::AnalyzerObj* analyzer_; + sym::AnalyzerObj* analyzer_; /*! \brief The domain map of dynamic TIR variables for analysis. */ - ffi::Map dom_map_; + ffi::Map dom_map_; /*! \brief The mapping from each token to the binding block where it is created. */ std::unordered_map token2block_; /*! \brief The mapping from each token to the Exprs that are using this token. */ @@ -768,7 +768,7 @@ class StorageAllocatorInit : public StorageAllocatorBaseVisitor { class StorageAllocator : public StorageAllocatorBaseVisitor { public: explicit StorageAllocator(std::unordered_map token_map, - arith::AnalyzerObj* analyzer) + sym::AnalyzerObj* analyzer) : allocator_(analyzer) { this->token_map_ = std::move(token_map); } @@ -1018,9 +1018,9 @@ class StorageAllocationRewriter : public ExprMutator { } /*! \brief The arithmetic analyzer. */ - arith::Analyzer ana_; + sym::Analyzer ana_; /*! \brief The domain map of dynamic TIR variables for analysis. */ - ffi::Map dom_map_; + ffi::Map dom_map_; /*! \brief A boolean indicating whether to plan dynamic-shape function output tensors. */ bool plan_dynamic_output_; /*! @@ -1035,7 +1035,7 @@ class StorageAllocationRewriter : public ExprMutator { }; IRModule StaticPlanBlockMemory(IRModule mod) { - arith::Analyzer ana; + sym::Analyzer ana; // Step 1. Initialize. std::unordered_map token_map = diff --git a/src/relax/utils.cc b/src/relax/utils.cc index 4511ef8deac1..e9179d9f4261 100644 --- a/src/relax/utils.cc +++ b/src/relax/utils.cc @@ -101,7 +101,7 @@ class ExprBinder : public ExprMutator { } const tvm::ffi::Map& bindings_; - arith::Analyzer analyzer_; + sym::Analyzer analyzer_; }; /*! @@ -119,7 +119,7 @@ Type Bind(const Type& ty, const tvm::ffi::Map& binds) { } tvm::ffi::Map InferSymbolicVarMap( - const tvm::ffi::Map& relax_var_remap, const arith::Analyzer& analyzer) { + const tvm::ffi::Map& relax_var_remap, const sym::Analyzer& analyzer) { tvm::ffi::Map var_remap = relax_var_remap; for (const auto& [var, value] : relax_var_remap) { diff --git a/src/s_tir/analysis/calculate_allocated_memory.cc b/src/s_tir/analysis/calculate_allocated_memory.cc index c96a9c26667b..5ec2d5ed475b 100644 --- a/src/s_tir/analysis/calculate_allocated_memory.cc +++ b/src/s_tir/analysis/calculate_allocated_memory.cc @@ -21,13 +21,13 @@ * \file tirx/analysis/calculate_allocated_memory.cc * \brief Calculate allocated memory per memory scope required by PrimFuncs. */ -#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 448171b4da16..9221bb8414f2 100644 --- a/src/s_tir/analysis/conditional_bounds.cc +++ b/src/s_tir/analysis/conditional_bounds.cc @@ -23,12 +23,12 @@ */ #include "conditional_bounds.h" -#include -#include #include #include #include #include +#include +#include #include #include @@ -36,21 +36,21 @@ #include #include -#include "../../arith/int_operator.h" +#include "../../sym/int_operator.h" namespace tvm { namespace s_tir { using namespace tvm::prim; using namespace tvm::tirx; -using arith::Analyzer; -using arith::AnalyzerObj; -using arith::EvalSet; -using arith::IntSet; +using sym::Analyzer; +using sym::AnalyzerObj; +using sym::EvalSet; +using sym::IntSet; namespace { -using arith::ExtendedEuclidean; -using arith::LeastCommonMultiple; +using sym::ExtendedEuclidean; +using sym::LeastCommonMultiple; // The solver's intermediate representations remain local to this analysis. struct IntGroupBounds { @@ -268,7 +268,7 @@ class NormalizeComparisons : public tvm::ExprMutator { } return T(analyzer_->Simplify(a - b), IntImm(a.ty(), 0)); } - arith::Analyzer analyzer_; + sym::Analyzer analyzer_; }; void AddInequality(std::vector* inequality_set, const PrimExpr& new_ineq, @@ -306,7 +306,7 @@ void ClassifyByPolarity(const PrimVar& var, const std::vector& current // and store to coef_pos and coef_neg respectively. for (const PrimExpr& ineq : current_ineq_set) { if (const prim::LENode* le = ineq.as()) { - ffi::Array coef = arith::DetectLinearEquation(le->a, {var}); + ffi::Array coef = sym::DetectLinearEquation(le->a, {var}); const auto* imm = !coef.empty() ? coef[0].as() : nullptr; if (auto value = imm ? imm->value.as() : std::nullopt; value.has_value()) { int64_t coef0 = *value; @@ -321,7 +321,7 @@ void ClassifyByPolarity(const PrimVar& var, const std::vector& current continue; } } else if (const prim::EQNode* eq = ineq.as()) { - ffi::Array coef = arith::DetectLinearEquation(eq->a, {var}); + ffi::Array coef = sym::DetectLinearEquation(eq->a, {var}); const auto* imm = !coef.empty() ? coef[0].as() : nullptr; if (auto value = imm ? imm->value.as() : std::nullopt; value.has_value()) { int64_t coef0 = *value; @@ -362,7 +362,7 @@ void MoveEquality(std::vector* upper_bounds, std::vector* lo } PartialSolvedInequalities SolveLinearInequalities(const IntConstraints& system_to_solve) { - arith::Analyzer analyzer; + sym::Analyzer analyzer; analyzer->Bind(system_to_solve.ranges); // The algorithm consists in doing the following things for each variable v @@ -560,7 +560,7 @@ IntConstraints SolveInequalitiesToRange(const IntConstraints& inequalities) { // We process variables in the reverse direction to start with the most independent one. // This order is needed to compute new ranges. for (auto it = inequalities.variables.rbegin(); it != inequalities.variables.rend(); ++it) { - arith::Analyzer analyzer; + sym::Analyzer analyzer; analyzer->Bind(vranges); const PrimVar& var = *it; @@ -596,7 +596,7 @@ IntConstraints SolveInequalitiesToRange(const IntConstraints& inequalities) { } // Add the original conditions to the resulting conditions - arith::Analyzer analyzer; + sym::Analyzer analyzer; analyzer->Bind(vranges); for (const PrimExpr& old_cond : AsConditions(inequalities.variables, solved_bounds, solved_other_relations)) { @@ -619,7 +619,7 @@ IntConstraints SolveInequalitiesToRange(const IntConstraints& inequalities) { ffi::Optional> ConditionalBoundsContext::TrySolveCondition() { // extract equations and related vars from condition expression. // currently only extract simple integral equations which could be solvable. - arith::Analyzer analyzer; + sym::Analyzer analyzer; PrimExpr condition = analyzer->Simplify(condition_); if (is_const_int(condition)) { return std::nullopt; @@ -676,7 +676,7 @@ ffi::Optional> ConditionalBoundsContext::TrySolveCondition( // build dom ranges for related vars ffi::Map ranges; for (const Var& v : vars) { - arith::IntSet dom; + sym::IntSet dom; auto relax_it = relax_map_->find(v.get()); if (relax_it != relax_map_->end()) { dom = relax_it->second; @@ -700,8 +700,8 @@ ffi::Optional> ConditionalBoundsContext::TrySolveCondition( } ConditionalBoundsContext::ConditionalBoundsContext( - const PrimExpr& condition, std::unordered_map* relax_map, - std::unordered_map* hint_map, + const PrimExpr& condition, std::unordered_map* relax_map, + std::unordered_map* hint_map, std::vector* pending_conditions) : condition_(condition), relax_map_(relax_map), @@ -719,20 +719,20 @@ void ConditionalBoundsContext::EnterWithScope() { // update solved var ranges for (const auto& kv : constraints.value()) { const VarNode* var = kv.first.get(); - arith::IntSet new_dom = arith::IntSet::FromRange(kv.second); + sym::IntSet new_dom = sym::IntSet::FromRange(kv.second); auto relax_it = relax_map_->find(var); if (relax_it != relax_map_->end()) { // this is a bound for relaxed var origin_map_.emplace(var, relax_it->second); - relax_it->second = arith::Intersect({relax_it->second, new_dom}); + relax_it->second = sym::Intersect({relax_it->second, new_dom}); } else { // this is a bound for free var auto hint_it = hint_map_->find(var); if (hint_it != hint_map_->end()) { origin_map_.emplace(var, hint_it->second); - hint_it->second = arith::Intersect({hint_it->second, new_dom}); + hint_it->second = sym::Intersect({hint_it->second, new_dom}); } else { - origin_map_.emplace(var, arith::IntSet::Nothing()); + origin_map_.emplace(var, sym::IntSet::Nothing()); hint_map_->insert(hint_it, {var, new_dom}); } } diff --git a/src/s_tir/analysis/conditional_bounds.h b/src/s_tir/analysis/conditional_bounds.h index 779e89bb5ae7..27cbd751cbbe 100644 --- a/src/s_tir/analysis/conditional_bounds.h +++ b/src/s_tir/analysis/conditional_bounds.h @@ -24,9 +24,9 @@ #ifndef TVM_S_TIR_ANALYSIS_CONDITIONAL_BOUNDS_H_ #define TVM_S_TIR_ANALYSIS_CONDITIONAL_BOUNDS_H_ -#include #include #include +#include #include #include @@ -53,8 +53,8 @@ class ConditionalBoundsContext { * \param pending_conditions The stack of unresolved constraints. */ ConditionalBoundsContext(const PrimExpr& condition, - std::unordered_map* relax_map, - std::unordered_map* hint_map, + std::unordered_map* relax_map, + std::unordered_map* hint_map, std::vector* pending_constraints); void EnterWithScope(); void ExitWithScope(); @@ -65,13 +65,13 @@ class ConditionalBoundsContext { /*! \brief the condition holds on true branch. */ const PrimExpr& condition_; /*! \brief domain map for relaxed vars to update */ - std::unordered_map* relax_map_; + std::unordered_map* relax_map_; /*! \brief domain map for free vars to update */ - std::unordered_map* hint_map_; + std::unordered_map* hint_map_; /*! \brief unresolved condition stack */ std::vector* pending_conditions_; /*! \brief used to record and restore original var bounds */ - std::unordered_map origin_map_; + std::unordered_map origin_map_; /*! \brief used to record unresolved conditions num. */ size_t origin_pending_conditions_num_; }; diff --git a/src/s_tir/analysis/domain_touched.cc b/src/s_tir/analysis/domain_touched.cc index bbd8c4430da0..627acda46320 100644 --- a/src/s_tir/analysis/domain_touched.cc +++ b/src/s_tir/analysis/domain_touched.cc @@ -21,13 +21,13 @@ * \file domain_touched.cc * \brief Analyze buffer domains touched by a statement */ -#include #include #include #include #include #include #include +#include #include #include @@ -40,7 +40,7 @@ namespace tvm { namespace s_tir { using namespace tirx; -using arith::IntSet; +using sym::IntSet; namespace { @@ -93,7 +93,7 @@ class BufferTouchedDomain final : public s_tir::IRVisitorWithAnalyzer { << "Must consider at least on of either loads and stores, but both are false"; } for (size_t i = 0; i < bounds.size(); ++i) { - ret.push_back(arith::Union(bounds[i]).CoverRange(none)); + ret.push_back(sym::Union(bounds[i]).CoverRange(none)); } return ret; } diff --git a/src/s_tir/analysis/estimate_flops.cc b/src/s_tir/analysis/estimate_flops.cc index 3665f5006d48..d22793a7d0e0 100644 --- a/src/s_tir/analysis/estimate_flops.cc +++ b/src/s_tir/analysis/estimate_flops.cc @@ -22,7 +22,7 @@ #include #include -#include "tvm/arith/analyzer.h" +#include "tvm/sym/analyzer.h" namespace tvm { namespace s_tir { @@ -88,7 +88,7 @@ struct TResult { class FlopEstimator : private tirx::ExprFunctor, private StmtFunctor { - arith::Analyzer ana; + sym::Analyzer ana; public: using tirx::ExprFunctor::Dispatch; @@ -118,9 +118,9 @@ class FlopEstimator : private tirx::ExprFunctor, TResult Dispatch_(const prim::GTNode* op) override { return TResult(); } TResult Dispatch_(const prim::GENode* op) override { return TResult(); } - int64_t GetLoopExtent(const ForNode* node, const arith::Analyzer& ana) { + int64_t GetLoopExtent(const ForNode* node, const sym::Analyzer& ana) { int64_t bound = ana->const_int_bound(node->extent)->max_value; - if (bound == arith::ConstIntBound::kPosInf) { + if (bound == sym::ConstIntBound::kPosInf) { return 1; // Analyzer could not determine a valid bound, use 1 instead. } else { return bound; diff --git a/src/s_tir/analysis/identify_memcpy.cc b/src/s_tir/analysis/identify_memcpy.cc index 02cfdbe865e3..0e298d9375d9 100644 --- a/src/s_tir/analysis/identify_memcpy.cc +++ b/src/s_tir/analysis/identify_memcpy.cc @@ -22,12 +22,12 @@ * \brief Check if a loop nest is equivalent to memcpy */ -#include -#include #include #include #include #include +#include +#include #include #include #include @@ -45,8 +45,8 @@ namespace s_tir { using namespace tvm::tirx; std::variant IdentifyMemCpyImpl(const For& loop, - arith::AnalyzerObj* analyzer) { - ffi::Map loop_intervals; + sym::AnalyzerObj* analyzer) { + ffi::Map loop_intervals; ffi::Map loop_ranges; PrimExpr total_loop_iterations = 1; @@ -56,7 +56,7 @@ std::variant IdentifyMemCpyImpl(const For& loop, while (auto* for_node = stmt.as()) { loop_ranges.Set(for_node->loop_var, Range::FromMinExtent(for_node->min, for_node->extent)); loop_intervals.Set(for_node->loop_var, - arith::IntSet::FromMinExtent(for_node->min, for_node->extent)); + sym::IntSet::FromMinExtent(for_node->min, for_node->extent)); total_loop_iterations = total_loop_iterations * for_node->extent; stmt = for_node->body; @@ -108,22 +108,22 @@ std::variant IdentifyMemCpyImpl(const For& loop, // for i in T.serial(16): // B[i] = A[T.abs(i-8)] - arith::Analyzer analyzer_ref = ffi::GetRef(analyzer); - auto src_iter_map = arith::DetectIterMap({src_index}, loop_ranges, IntImm::Bool(true), - arith::IterMapLevel::Bijective, analyzer_ref); + sym::Analyzer analyzer_ref = ffi::GetRef(analyzer); + auto src_iter_map = sym::DetectIterMap({src_index}, loop_ranges, IntImm::Bool(true), + sym::IterMapLevel::Bijective, analyzer_ref); if (src_iter_map->errors.size()) { return static_cast(std::stringstream() - << "arith::DetectIterMap(src) returned " + << "sym::DetectIterMap(src) returned " << src_iter_map->errors.size() << " errors: [" << src_iter_map->errors << "]" << " for src_index = " << src_index) .str(); } - auto dst_iter_map = arith::DetectIterMap({dst_index}, loop_ranges, IntImm::Bool(true), - arith::IterMapLevel::Bijective, analyzer_ref); + auto dst_iter_map = sym::DetectIterMap({dst_index}, loop_ranges, IntImm::Bool(true), + sym::IterMapLevel::Bijective, analyzer_ref); if (dst_iter_map->errors.size()) { return static_cast(std::stringstream() - << "arith::DetectIterMap(dst) returned " + << "sym::DetectIterMap(dst) returned " << dst_iter_map->errors.size() << " errors: [" << dst_iter_map->errors << "]" << " for dst_index = " << dst_index) @@ -211,12 +211,12 @@ std::variant IdentifyMemCpyImpl(const For& loop, << "IterMaps were detected as src = " << src_iter_sum << ", dst = " << dst_iter_sum) .str(); } - std::vector src_iter_terms(src_iter_sum->args.begin(), - src_iter_sum->args.end()); - std::vector dst_iter_terms(dst_iter_sum->args.begin(), - dst_iter_sum->args.end()); + std::vector src_iter_terms(src_iter_sum->args.begin(), + src_iter_sum->args.end()); + std::vector dst_iter_terms(dst_iter_sum->args.begin(), + dst_iter_sum->args.end()); - auto make_comparison_tuple = [](const arith::IterSplitExpr& expr) { + auto make_comparison_tuple = [](const sym::IterSplitExpr& expr) { auto as_int_or_zero = [](auto& val) -> ffi::BigInt { if (auto* as_int = val.template as()) { return as_int->value; @@ -230,19 +230,19 @@ std::variant IdentifyMemCpyImpl(const For& loop, static_cast(expr->lower_factor.as()), as_int_or_zero(expr->lower_factor), }; }; - auto sorting_function = [&make_comparison_tuple](const arith::IterSplitExpr& lhs, - const arith::IterSplitExpr& rhs) -> bool { + auto sorting_function = [&make_comparison_tuple](const sym::IterSplitExpr& lhs, + const sym::IterSplitExpr& rhs) -> bool { return make_comparison_tuple(lhs) < make_comparison_tuple(rhs); }; std::sort(src_iter_terms.begin(), src_iter_terms.end(), sorting_function); std::sort(dst_iter_terms.begin(), dst_iter_terms.end(), sorting_function); for (size_t i = 0; i < src_iter_terms.size(); i++) { - const arith::IterSplitExpr& src_term = src_iter_terms[i]; - const arith::IterSplitExpr& dst_term = dst_iter_terms[i]; + const sym::IterSplitExpr& src_term = src_iter_terms[i]; + const sym::IterSplitExpr& dst_term = dst_iter_terms[i]; if (!analyzer->CanProve( - arith::NormalizeIterMapToExpr(src_term->source->source == dst_term->source->source))) { + sym::NormalizeIterMapToExpr(src_term->source->source == dst_term->source->source))) { return static_cast( std::stringstream() << "Term " << i << " had different source, src_term->source = " << src_term->source @@ -282,7 +282,7 @@ std::variant IdentifyMemCpyImpl(const For& loop, return MemCpyDetails{src_region, dst_region}; } -std::optional IdentifyMemCpy(const For& loop, const arith::Analyzer& analyzer) { +std::optional IdentifyMemCpy(const For& loop, const sym::Analyzer& analyzer) { auto result = IdentifyMemCpyImpl(loop, analyzer.get()); if (auto* ptr = std::get_if(&result)) { return *ptr; diff --git a/src/s_tir/analysis/oob_checker.cc b/src/s_tir/analysis/oob_checker.cc index 42e3161b5113..baa73f7e0d30 100644 --- a/src/s_tir/analysis/oob_checker.cc +++ b/src/s_tir/analysis/oob_checker.cc @@ -35,8 +35,8 @@ struct OOBLocation { BufferVar buf; size_t dimension; ffi::ObjectRef index; - arith::IntSet index_bounds; - arith::IntSet shape_bounds; + sym::IntSet index_bounds; + sym::IntSet shape_bounds; }; class OOBError : public s_tir::ScheduleErrorContextObj { diff --git a/src/s_tir/analysis/sblock_access_region_detector.cc b/src/s_tir/analysis/sblock_access_region_detector.cc index b293b0d41dea..b6dd57facb90 100644 --- a/src/s_tir/analysis/sblock_access_region_detector.cc +++ b/src/s_tir/analysis/sblock_access_region_detector.cc @@ -22,12 +22,12 @@ * \brief Detect sblock read/write regions by visiting its body */ -#include #include #include #include #include #include +#include #include #include @@ -74,9 +74,9 @@ class BlockReadWriteDetector : public s_tir::StmtExprVisitor { private: /*! \brief Iteration range for loop_vars */ - std::unordered_map dom_map_; + std::unordered_map dom_map_; /*! \brief Extra iteration range hint for free vars */ - std::unordered_map hint_map_; + std::unordered_map hint_map_; /*! \brief Unresolved conditions within current scope. */ std::vector pending_conditions_; /*! \brief The buffers that the current block reads */ @@ -86,11 +86,11 @@ class BlockReadWriteDetector : public s_tir::StmtExprVisitor { /*! \brief The opaque buffer which is access by buffer.data */ std::vector opaque_buffers_; /*! \brief The read regions of the current block */ - std::vector> read_regions_; + std::vector> read_regions_; /*! \brief The write regions of the current block */ - std::vector> write_regions_; + std::vector> write_regions_; /*! \brief The opaque regions of the current block */ - std::vector> opaque_regions_; + std::vector> opaque_regions_; /*! \brief The outside buffer data mapping to its buffer */ ffi::Map buffer_var_map_; /*! \brief The target buffer var mapping to its matching */ @@ -98,7 +98,7 @@ class BlockReadWriteDetector : public s_tir::StmtExprVisitor { /*! \brief let bindings inside the block */ std::unordered_map let_bindings_; /*!\ brief Internal analyzer. */ - arith::Analyzer ana_; + sym::Analyzer ana_; /*! * \brief Update read/write buffers and regions with provided buffer and region @@ -107,24 +107,24 @@ class BlockReadWriteDetector : public s_tir::StmtExprVisitor { * \param buffer The provided buffer * \param region The provided region */ - void Update(std::vector* buffers, std::vector>* regions, - BufferVar buffer, std::vector region); + void Update(std::vector* buffers, std::vector>* regions, + BufferVar buffer, std::vector region); /*! \brief Helper function to collect access regions. */ ffi::Array CollectRegions( const std::vector& buffers, - const std::vector>& regions, + const std::vector>& regions, const std::unordered_set* excluded_buffers = nullptr); /*! \brief Helper function to convert matched access region to source region. */ - std::vector ConvertMatchedRegion(const s_tir::MatchBufferRegion& match_buffer, - const std::vector& int_sets) const; + std::vector ConvertMatchedRegion(const s_tir::MatchBufferRegion& match_buffer, + const std::vector& int_sets) const; /*! \brief Helper function to update a opaque access. */ void UpdateOpaque(const Var& buffer_var); /*! \brief Helper function to relax the buffer indices */ - arith::IntSet RelaxAccessIndex(const PrimExpr& index); + sym::IntSet RelaxAccessIndex(const PrimExpr& index); // Declared regions carry bounds, not opaque runtime accesses. ffi::Optional Visit_(const TensorRegionNode* op) final { @@ -189,7 +189,7 @@ ffi::Optional BlockReadWriteDetector::Visit_(const TensorLoadNod } return ffi::Unchanged(); }; - std::vector relaxed_region; + std::vector relaxed_region; for (PrimExpr index : op->indices) { PrimExpr remapped_index = ffi::StructuralMap(index, f_substitute).as_or_throw(); @@ -198,7 +198,7 @@ ffi::Optional BlockReadWriteDetector::Visit_(const TensorLoadNod remapped_index = ffi::StructuralMap(index, f_substitute) .as_or_throw(); } - relaxed_region.push_back(arith::EvalSet(arith::IntSet::Vector(remapped_index), dom_map_)); + relaxed_region.push_back(sym::EvalSet(sym::IntSet::Vector(remapped_index), dom_map_)); } Update(&read_buffers_, &read_regions_, op->source.as_or_throw(), relaxed_region); @@ -210,7 +210,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); + dom_map_[op->loop_var.get()] = sym::IntSet::FromRange(range); TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(s_tir::StmtExprVisitor::Visit_(op)); dom_map_.erase(op->loop_var.get()); return std::nullopt; @@ -251,14 +251,14 @@ ffi::Optional BlockReadWriteDetector::Visit_(const BindNode* op) ffi::Optional BlockReadWriteDetector::Visit_(const CallNode* op) { auto update_masked_access = [this](const BufferVar& buffer, const ffi::Array& indices, std::vector* buffers, - std::vector>* regions) { + std::vector>* regions) { auto f_substitute = [this](const Var& var) -> ffi::Expected> { if (auto it = let_bindings_.find(var.get()); it != let_bindings_.end()) { return ffi::Any(it->second); } return ffi::Unchanged(); }; - std::vector relaxed_region; + std::vector relaxed_region; for (PrimExpr index : indices) { PrimExpr remapped_index = ffi::StructuralMap(index, f_substitute) .as_or_throw(); @@ -267,7 +267,7 @@ ffi::Optional BlockReadWriteDetector::Visit_(const CallNode* op) remapped_index = ffi::StructuralMap(index, f_substitute) .as_or_throw(); } - relaxed_region.push_back(arith::EvalSet(arith::IntSet::Vector(remapped_index), dom_map_)); + relaxed_region.push_back(sym::EvalSet(sym::IntSet::Vector(remapped_index), dom_map_)); } Update(buffers, regions, buffer, relaxed_region); }; @@ -299,10 +299,10 @@ ffi::Optional BlockReadWriteDetector::Visit_(const CallNode* op) const BufferVar& buffer = (*it).second; const TensorRegion buffer_region = FullBufferRegion(buffer); const ffi::Array& region = buffer_region->region; - std::vector int_set; + std::vector int_set; int_set.reserve(region.size()); for (const Range& range : region) { - int_set.push_back(arith::EvalSet(range, dom_map_)); + int_set.push_back(sym::EvalSet(range, dom_map_)); } // read access, write access or opaque access if ((access_mask->value & 1) && (access_mask->value & 2)) { @@ -347,7 +347,7 @@ ffi::Optional BlockReadWriteDetector::Visit_(const BufferStoreNo } return ffi::Unchanged(); }; - std::vector relaxed_region; + std::vector relaxed_region; for (PrimExpr index : op->indices) { PrimExpr remapped_index = ffi::StructuralMap(index, f_substitute).as_or_throw(); @@ -356,7 +356,7 @@ ffi::Optional BlockReadWriteDetector::Visit_(const BufferStoreNo remapped_index = ffi::StructuralMap(index, f_substitute) .as_or_throw(); } - relaxed_region.push_back(arith::EvalSet(arith::IntSet::Vector(remapped_index), dom_map_)); + relaxed_region.push_back(sym::EvalSet(sym::IntSet::Vector(remapped_index), dom_map_)); } Update(&writes_buffers_, &write_regions_, op->buffer, relaxed_region); TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Visit(op->value)); @@ -377,27 +377,27 @@ ffi::Optional BlockReadWriteDetector::Visit_(const s_tir::SBlock return ffi::Unchanged(); }; for (const auto& read : op->block->reads) { - std::vector relaxed_region; + std::vector relaxed_region; for (const auto& range : read->region) { PrimExpr min = ffi::StructuralMap(range->min, f_substitute) .as_or_throw(); PrimExpr extent = ffi::StructuralMap(range->extent, f_substitute) .as_or_throw(); relaxed_region.push_back( - arith::EvalSet(arith::IntSet::FromRange(Range::FromMinExtent(min, extent)), dom_map_)); + sym::EvalSet(sym::IntSet::FromRange(Range::FromMinExtent(min, extent)), dom_map_)); } Update(&read_buffers_, &read_regions_, read->source.as_or_throw(), relaxed_region); } for (const auto& write : op->block->writes) { - std::vector relaxed_region; + std::vector relaxed_region; for (const auto& range : write->region) { PrimExpr min = ffi::StructuralMap(range->min, f_substitute) .as_or_throw(); PrimExpr extent = ffi::StructuralMap(range->extent, f_substitute) .as_or_throw(); relaxed_region.push_back( - arith::EvalSet(arith::IntSet::FromRange(Range::FromMinExtent(min, extent)), dom_map_)); + sym::EvalSet(sym::IntSet::FromRange(Range::FromMinExtent(min, extent)), dom_map_)); } Update(&writes_buffers_, &write_regions_, write->source.as_or_throw(), relaxed_region); @@ -405,32 +405,31 @@ ffi::Optional BlockReadWriteDetector::Visit_(const s_tir::SBlock return std::nullopt; } -std::vector BlockReadWriteDetector::ConvertMatchedRegion( - const s_tir::MatchBufferRegion& match_buffer, - const std::vector& int_sets) const { +std::vector BlockReadWriteDetector::ConvertMatchedRegion( + const s_tir::MatchBufferRegion& match_buffer, const std::vector& int_sets) const { const BufferVar& buffer = match_buffer->buffer; ffi::Array region; region.reserve(int_sets.size()); TVM_FFI_ICHECK_EQ(buffer->shape.size(), int_sets.size()); for (size_t i = 0; i < int_sets.size(); ++i) { - const tvm::arith::IntSet& int_set = int_sets[i]; + const tvm::sym::IntSet& int_set = int_sets[i]; region.push_back(int_set.CoverRange(Range::FromMinExtent(0, buffer->shape[i]))); } region = ConvertRegion(match_buffer, region); - std::vector result; + std::vector result; result.reserve(region.size()); for (const Range& range : region) { - result.push_back(arith::EvalSet(range, dom_map_)); + result.push_back(sym::EvalSet(range, dom_map_)); } return result; } void BlockReadWriteDetector::Update(std::vector* buffers, - std::vector>* regions, - BufferVar buffer, std::vector region) { + std::vector>* regions, + BufferVar buffer, std::vector region) { if (buffer_var_map_.find(buffer.var()) == buffer_var_map_.end()) return; // Handle match_buffer remap auto it = match_buffers_.find(buffer.get()); @@ -445,7 +444,7 @@ void BlockReadWriteDetector::Update(std::vector* buffers, if ((*buffers)[i].same_as(buffer)) { TVM_FFI_ICHECK_EQ((*regions)[i].size(), region.size()) << "Inconsistent buffer dimension"; for (size_t j = 0; j < region.size(); ++j) { - (*regions)[i][j] = arith::Union({(*regions)[i][j], region[j]}); + (*regions)[i][j] = sym::Union({(*regions)[i][j], region[j]}); } return; } @@ -456,7 +455,7 @@ void BlockReadWriteDetector::Update(std::vector* buffers, ffi::Array BlockReadWriteDetector::CollectRegions( const std::vector& buffers, - const std::vector>& regions, + const std::vector>& regions, const std::unordered_set* excluded_buffers) { TVM_FFI_ICHECK_EQ(buffers.size(), regions.size()); ffi::Array res; @@ -469,7 +468,7 @@ ffi::Array BlockReadWriteDetector::CollectRegions( region.reserve(regions[i].size()); TVM_FFI_ICHECK_EQ(buffers[i]->shape.size(), regions[i].size()); for (size_t j = 0; j < regions[i].size(); j++) { - const tvm::arith::IntSet& range = regions[i][j]; + const tvm::sym::IntSet& range = regions[i][j]; if (range.CanProveSinglePoint(ana_)) { PrimExpr min = range.min(); region.push_back(Range::FromMinExtent(min, prim::MakeConst(min.ty(), 1))); @@ -488,10 +487,10 @@ void BlockReadWriteDetector::UpdateOpaque(const Var& buffer_var) { const BufferVar& buffer = (*it).second; const TensorRegion buffer_region = FullBufferRegion(buffer); const ffi::Array& region = buffer_region->region; - std::vector int_set; + std::vector int_set; int_set.reserve(region.size()); for (const Range& range : region) { - int_set.push_back(arith::EvalSet(range, dom_map_)); + int_set.push_back(sym::EvalSet(range, dom_map_)); } Update(&opaque_buffers_, &opaque_regions_, buffer, int_set); } diff --git a/src/s_tir/backend/adreno/inject_texture_alloc.cc b/src/s_tir/backend/adreno/inject_texture_alloc.cc index e9403acaebc7..19882d3bb22b 100644 --- a/src/s_tir/backend/adreno/inject_texture_alloc.cc +++ b/src/s_tir/backend/adreno/inject_texture_alloc.cc @@ -21,10 +21,10 @@ * \file inject_texture_alloc.cc */ -#include #include #include #include +#include #include #include "../../../backend/opencl/runtime/texture.h" @@ -49,7 +49,7 @@ class TextureAllocInjector : public s_tir::IRMutatorWithAnalyzer { using s_tir::IRMutatorWithAnalyzer::Mutate_; static PrimFunc Inject(PrimFunc func) { - arith::Analyzer ana; + sym::Analyzer ana; auto pass = ffi::make_object(ana); auto writer = func.CopyOnWrite(); pass->MarkBufferParamShapes(func); @@ -57,7 +57,7 @@ class TextureAllocInjector : public s_tir::IRMutatorWithAnalyzer { return func; } - explicit TextureAllocInjector(const arith::Analyzer& ana) : IRMutatorWithAnalyzer(ana) {} + explicit TextureAllocInjector(const sym::Analyzer& ana) : IRMutatorWithAnalyzer(ana) {} private: UnchangedOr Mutate_(const AllocBufferNode* op, InplaceMode inplace_mode) final { diff --git a/src/s_tir/data_layout.cc b/src/s_tir/data_layout.cc index 9000d298e586..8febe31d65e4 100644 --- a/src/s_tir/data_layout.cc +++ b/src/s_tir/data_layout.cc @@ -21,7 +21,6 @@ * \file src/lang/data_layout.cc * \brief Data SLayout expression. */ -#include #include #include #include @@ -32,6 +31,7 @@ #include #include #include +#include #include #include @@ -372,7 +372,7 @@ inline bool GetStoreRule(ffi::Array* index_rule, ffi::Array* } } - arith::Analyzer ana; + sym::Analyzer ana; for (size_t i = 0; i < dst_layout.ndim(); i++) { const auto dst_unpacked_axes = SLayout::UnpackIterVar(dst_layout.PackedAxisAt(i)); @@ -447,7 +447,7 @@ inline bool GetStoreRule(ffi::Array* index_rule, ffi::Array* inline ffi::Array TransformIndex(const ffi::Array& src_index, const ffi::Array& src_axis, const ffi::Array& transform_rule) { - arith::Analyzer ana; + sym::Analyzer ana; ffi::Array result; std::unordered_map bind_map; for (size_t i = 0; i < src_index.size(); ++i) { @@ -486,7 +486,7 @@ inline ffi::Array TransformShape(const ffi::Array& src_shape const ffi::Array& src_axis, const ffi::Array& target_axis, const ffi::Array& transform_rule) { - arith::Analyzer ana; + sym::Analyzer ana; TVM_FFI_ICHECK_EQ(src_shape.size(), src_axis.size()) << "Input shape size " << src_shape.size() << " mismatch with the expected shape size " << src_axis.size(); diff --git a/src/s_tir/ir/ir_mutator_with_analyzer.h b/src/s_tir/ir/ir_mutator_with_analyzer.h index 49179ab43cb9..204f9750c524 100644 --- a/src/s_tir/ir/ir_mutator_with_analyzer.h +++ b/src/s_tir/ir/ir_mutator_with_analyzer.h @@ -31,9 +31,9 @@ class IRMutatorWithAnalyzer : public tirx::IRMutatorWithAnalyzer { using Parent = tirx::IRMutatorWithAnalyzer; using Parent::Mutate; using Parent::Mutate_; - explicit IRMutatorWithAnalyzer(const arith::Analyzer& analyzer) + explicit IRMutatorWithAnalyzer(const sym::Analyzer& analyzer) : IRMutatorWithAnalyzer(analyzer.get()) {} - explicit IRMutatorWithAnalyzer(arith::AnalyzerObj* analyzer) : Parent(analyzer, GlobalVTable()) {} + explicit IRMutatorWithAnalyzer(sym::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); 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 123abd0d84f3..ee79c7ae3763 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 @@ -67,7 +67,7 @@ namespace utils { * \param analyzer The analyzer * \return The shape of the buffer */ -std::vector GetBufferShape(const BufferVar& buffer, arith::AnalyzerObj* analyzer) { +std::vector GetBufferShape(const BufferVar& buffer, sym::AnalyzerObj* analyzer) { int ndim = buffer->shape.size(); std::vector result; result.reserve(ndim); @@ -76,8 +76,8 @@ std::vector GetBufferShape(const BufferVar& buffer, arith::AnalyzerObj* result.push_back(static_cast(int_imm->value)); continue; } - arith::ConstIntBound bound = analyzer->const_int_bound(i); - if (0 <= bound->max_value && bound->max_value < arith::ConstIntBound::kPosInf) { + sym::ConstIntBound bound = analyzer->const_int_bound(i); + if (0 <= bound->max_value && bound->max_value < sym::ConstIntBound::kPosInf) { result.push_back(bound->max_value); } else { result.push_back(1); @@ -128,7 +128,7 @@ int64_t FirstLoopExtent(const ForVec& loops, int64_t default_value) { * \return The relaxed and unioned region */ IntVec RelaxAndUnion(const std::vector& multi_indices, int64_t* numel, - arith::AnalyzerObj* analyzer) { + sym::AnalyzerObj* analyzer) { *numel = 1; if (multi_indices.empty()) { return {}; @@ -137,10 +137,10 @@ IntVec RelaxAndUnion(const std::vector& multi_indices, int64_t* nume int ndim = multi_indices[0].size(); IntVec access_shape(ndim, 0); for (int i = 0; i < ndim; ++i) { - int64_t minimum = arith::ConstIntBound::kPosInf; - int64_t maximum = arith::ConstIntBound::kNegInf; + int64_t minimum = sym::ConstIntBound::kPosInf; + int64_t maximum = sym::ConstIntBound::kNegInf; for (int j = 0; j < n_indices; ++j) { - arith::ConstIntBound bound = analyzer->const_int_bound(multi_indices[j][i]); + sym::ConstIntBound bound = analyzer->const_int_bound(multi_indices[j][i]); minimum = std::min(minimum, bound->min_value); maximum = std::max(maximum, bound->max_value); } @@ -760,7 +760,7 @@ struct Feature { static void Pad(std::vector* v) { v->insert(v->end(), 18, 0.0); } - void SetStride(const LoopNest& loop_nest, arith::AnalyzerObj* analyzer); + void SetStride(const LoopNest& loop_nest, sym::AnalyzerObj* analyzer); void SetReuse(const LoopNest& loop_nest, // int64_t top_loop_touch_bytes, // @@ -789,14 +789,14 @@ struct Feature { explicit Feature(const BufferStoreNode* store, const LoopNest& loop_nest, int64_t cache_line_bytes, IntVec* for_touched_bytes, - ForBufferMap* buffer_touched_under_loop, arith::AnalyzerObj* analyzer); + ForBufferMap* buffer_touched_under_loop, sym::AnalyzerObj* analyzer); void Init(const BufferStoreNode* store, int n_loops); void SetRegion(const LoopNest& loop_nest, // IntVec* for_touched_bytes, // ForBufferMap* buffer_touched_under_loop, // - arith::AnalyzerObj* analyzer); + sym::AnalyzerObj* analyzer); std::vector sub_features; }; @@ -843,7 +843,7 @@ void Feature::Init(const BufferStoreNode* store, int n_loops) { void Feature::SetRegion(const LoopNest& loop_nest, IntVec* for_touched_bytes, ForBufferMap* buffer_touched_under_loop, - arith::AnalyzerObj* analyzer) { + sym::AnalyzerObj* analyzer) { int n_loops = loop_nest.loops.size(); const std::vector& loops = loop_nest.loops; // Step 1. Initialize and bind all the loop variables to a constant @@ -881,7 +881,7 @@ void Feature::SetRegion(const LoopNest& loop_nest, IntVec* for_touched_bytes, } } -void Feature::SubFeature::SetStride(const LoopNest& loop_nest, arith::AnalyzerObj* analyzer) { +void Feature::SubFeature::SetStride(const LoopNest& loop_nest, sym::AnalyzerObj* analyzer) { int n_loops = loop_nest.loops.size(); const std::vector& loops = loop_nest.loops; // For each buffer, we find the loop stride on it @@ -1041,7 +1041,7 @@ void Feature::SubFeature::SetFeature(const LoopNest& loop_nest, int64_t cache_li Feature::Feature(const BufferStoreNode* store, const LoopNest& loop_nest, int64_t cache_line_bytes, IntVec* for_touched_bytes, ForBufferMap* buffer_touched_under_loop, - arith::AnalyzerObj* analyzer) { + sym::AnalyzerObj* analyzer) { int n_loops = loop_nest.loops.size(); // Step 0. Initialize data structures this->Init(store, n_loops); @@ -1189,8 +1189,7 @@ struct Feature { Feature() = default; - explicit Feature(const LoopNest& loop_nest, const BufferVar& buffer, - arith::AnalyzerObj* analyzer) { + explicit Feature(const LoopNest& loop_nest, const BufferVar& buffer, sym::AnalyzerObj* analyzer) { std::vector shape = utils::GetBufferShape(buffer, analyzer); int64_t numel = 1; for (int64_t x : shape) { @@ -1410,7 +1409,7 @@ class PerStoreFeatureCollector : public StmtExprVisitor { bool is_gpu_; int64_t cache_line_bytes_; int64_t arith_intensity_curve_num_samples_; - arith::Analyzer analyzer_; + sym::Analyzer analyzer_; LoopNest loop_nest_ = {}; IntVec for_touched_bytes_ = {}; ForBufferMap buffer_touched_under_loop_ = {}; diff --git a/src/s_tir/meta_schedule/postproc/disallow_async_strided_mem_copy.cc b/src/s_tir/meta_schedule/postproc/disallow_async_strided_mem_copy.cc index 37073b2c1d50..94310d55df26 100644 --- a/src/s_tir/meta_schedule/postproc/disallow_async_strided_mem_copy.cc +++ b/src/s_tir/meta_schedule/postproc/disallow_async_strided_mem_copy.cc @@ -93,9 +93,9 @@ struct AsyncStridedMemCopyFinder : public StmtExprVisitor { ffi::Array store_index = bufferstorenode->indices; // Use DetectIterMap to detect whether store index is non-contiguous. - arith::Analyzer analyzer; + sym::Analyzer analyzer; auto store_iter_map = DetectIterMap(store_index, input_iters, 1, - arith::IterMapLevel::Surjective, analyzer, false); + sym::IterMapLevel::Surjective, analyzer, false); if (!store_iter_map->errors.empty()) { found_ = true; } @@ -105,7 +105,7 @@ struct AsyncStridedMemCopyFinder : public StmtExprVisitor { // Use DetectIterMap to detect whether load index is non-contiguous. auto load_iter_map = DetectIterMap(load_index, input_iters, 1, - arith::IterMapLevel::Surjective, analyzer, false); + sym::IterMapLevel::Surjective, analyzer, false); if (!load_iter_map->errors.empty()) { found_ = true; } diff --git a/src/s_tir/meta_schedule/postproc/rewrite_layout.cc b/src/s_tir/meta_schedule/postproc/rewrite_layout.cc index b47fc6c65c21..5c977ff175b3 100644 --- a/src/s_tir/meta_schedule/postproc/rewrite_layout.cc +++ b/src/s_tir/meta_schedule/postproc/rewrite_layout.cc @@ -112,7 +112,7 @@ class BufferReadPosCollector : public StmtExprVisitor { /*! \brief Loop stack for calculating IndexMap. */ ffi::Array loop_stack_; /*! \brief Arithmetic analyzer. */ - arith::Analyzer analyzer_; + sym::Analyzer analyzer_; /*! \brief Current BlockRealize scope, used in recursive visit */ SBlockRealize cur_realize_; }; diff --git a/src/s_tir/meta_schedule/postproc/rewrite_parallel_vectorize_unroll.cc b/src/s_tir/meta_schedule/postproc/rewrite_parallel_vectorize_unroll.cc index 3dcfb8f85ccf..fbfe65128c99 100644 --- a/src/s_tir/meta_schedule/postproc/rewrite_parallel_vectorize_unroll.cc +++ b/src/s_tir/meta_schedule/postproc/rewrite_parallel_vectorize_unroll.cc @@ -233,7 +233,7 @@ void AdjustParallelVectorize(const Schedule& sch, const SBlockRV& block_rv, for (const StmtSRef& loop_sref : loop_srefs) { int64_t stride = 0, buffer_stride = 1; const auto* var = loop_sref->StmtAs(); - arith::Analyzer analyzer; + sym::Analyzer analyzer; for (int i = access->region.size() - 1; i >= 0; i--) { PrimExpr idx = analyzer->Simplify( ffi::StructuralMap(access->region[i]->min, f_substitute) 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 5b8d16e28eea..d1bc43dc37ed 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 @@ -96,7 +96,7 @@ MultiLevelTilingWideVectorNode::SplitLoop(const Schedule& sch, SBlockRV block_rv const size_t innermost_axis = block_node->writes[0]->region.size() - 1; const PrimExpr innermost_iter_value = block_realize->iter_values[innermost_axis]; - if (!arith::Analyzer()->CanProve(static_cast(loop->loop_var) == innermost_iter_value)) { + if (!sym::Analyzer()->CanProve(static_cast(loop->loop_var) == innermost_iter_value)) { // If this is not the innermost spatial loop, split the loop in the normal way. return MultiLevelTilingNode::SplitLoop(sch, block_rv, loop_rv, n_tiles); } else { diff --git a/src/s_tir/meta_schedule/utils.h b/src/s_tir/meta_schedule/utils.h index d8d9e39a79ef..ca851ae63af6 100644 --- a/src/s_tir/meta_schedule/utils.h +++ b/src/s_tir/meta_schedule/utils.h @@ -19,7 +19,6 @@ #ifndef TVM_S_TIR_META_SCHEDULE_UTILS_H_ #define TVM_S_TIR_META_SCHEDULE_UTILS_H_ -#include #include #include #include @@ -44,6 +43,7 @@ #include #include #include +#include #include #include diff --git a/src/s_tir/schedule/analysis.h b/src/s_tir/schedule/analysis.h index 2ee6240b7268..f4c8541cd6d2 100644 --- a/src/s_tir/schedule/analysis.h +++ b/src/s_tir/schedule/analysis.h @@ -19,12 +19,12 @@ #ifndef TVM_S_TIR_SCHEDULE_ANALYSIS_H_ #define TVM_S_TIR_SCHEDULE_ANALYSIS_H_ -#include #include #include #include #include #include +#include #include #include @@ -84,7 +84,7 @@ StmtSRef GetSRefTreeRoot(const StmtSRef& sref); * \param analyzer The analyzer to be bound */ void AddShapeVarBounds(const ScheduleState& state, const StmtSRefNode* sref, - arith::AnalyzerObj* analyzer); + sym::AnalyzerObj* analyzer); /******** Scope ********/ /*! @@ -235,7 +235,7 @@ bool IsWriteCache(const StmtSRef& block_sref); * \return A boolean flag indicating if the binding is affine */ bool IsAffineBinding(const SBlockRealize& realize, const ffi::Map& loop_var_ranges, - arith::AnalyzerObj* analyzer); + sym::AnalyzerObj* analyzer); /*! * \brief Check whether a block has an affine binding using the cached flag, and throw an exception @@ -301,7 +301,7 @@ bool GetVarsTouchedByBlockIters(const SBlockRealize& block_realize, * \throw ScheduleError If the loop doesn't starts with zero. */ void CheckLoopStartsWithZero(const ScheduleState& self, const StmtSRef& loop_sref, - arith::AnalyzerObj* analyzer); + sym::AnalyzerObj* analyzer); /*! * \brief Check whether a block has a trivial binding, i.e. each block var is bound to a outer loop, @@ -606,7 +606,7 @@ bool CanReverseComputeAt(const ScheduleState& self, const StmtSRef& block_sref, ffi::Optional SuggestIndexMap(const BufferVar& buffer, const ffi::Array& indices, const ffi::Array& loops, const PrimExpr& predicate, - arith::AnalyzerObj* analyzer); + sym::AnalyzerObj* analyzer); /*! * \brief Checks if the given AST contains the specific operators @@ -706,11 +706,11 @@ bool NeedsRFactorOrCrossThreadReduction(const s_tir::ScheduleState& self, // * \param dom_high_exclusive The highest node in the sref tree path * \return An n-dimensional integer set */ -ffi::Array AnalyzeRegionUpperBound(const TensorRegion& region, - const PrimExpr& predicate, - const StmtSRef& dom_low_inclusive, - const StmtSRef& dom_high_exclusive, - arith::AnalyzerObj* analyzer); +ffi::Array AnalyzeRegionUpperBound(const TensorRegion& region, + const PrimExpr& predicate, + const StmtSRef& dom_low_inclusive, + const StmtSRef& dom_high_exclusive, + sym::AnalyzerObj* analyzer); /*! * \brief Analyze the buffer region under the sref tree path [dom_low_inclusive, dom_high_exclusive) @@ -722,11 +722,11 @@ ffi::Array AnalyzeRegionUpperBound(const TensorRegion& region, * \param analyzer The analyzer * \return An n-dimensional integer set */ -ffi::Array AnalyzeRegionLowerBound(const TensorRegion& region, - const PrimExpr& predicate, - const StmtSRef& dom_low_inclusive, - const StmtSRef& dom_high_exclusive, - arith::AnalyzerObj* analyzer); +ffi::Array AnalyzeRegionLowerBound(const TensorRegion& region, + const PrimExpr& predicate, + const StmtSRef& dom_low_inclusive, + const StmtSRef& dom_high_exclusive, + sym::AnalyzerObj* analyzer); /*! * \brief Simplify non-trivial expressions @@ -738,7 +738,7 @@ ffi::Array AnalyzeRegionLowerBound(const TensorRegion& region, * simplified to constant values for further scheduling and analysis because simplifing away the * block iters may result in loss of information for further analysis. */ -PrimExpr SimplifyNonTrivialExpr(const PrimExpr& expr, arith::AnalyzerObj* analyzer); +PrimExpr SimplifyNonTrivialExpr(const PrimExpr& expr, sym::AnalyzerObj* analyzer); /*! \brief Necessary information used for tensorization */ class TensorizeInfoNode : public ffi::Object { diff --git a/src/s_tir/schedule/analysis/analysis.cc b/src/s_tir/schedule/analysis/analysis.cc index 789dcbff7314..8c9194958b72 100644 --- a/src/s_tir/schedule/analysis/analysis.cc +++ b/src/s_tir/schedule/analysis/analysis.cc @@ -571,7 +571,7 @@ bool IsWriteCache(const StmtSRef& block_sref) { /******** Binding ********/ bool IsAffineBinding(const SBlockRealize& realize, const ffi::Map& loop_var_ranges, - arith::AnalyzerObj* analyzer) { + sym::AnalyzerObj* analyzer) { if (loop_var_ranges.empty()) { return true; } @@ -579,18 +579,18 @@ bool IsAffineBinding(const SBlockRealize& realize, const ffi::Map& l for (const auto& [var, range] : loop_var_ranges) { primitive_loop_var_ranges.Set(var.as_or_throw(), range); } - auto res = arith::DetectIterMap( + auto res = sym::DetectIterMap( /*indices=*/realize->iter_values, /*input_iters=*/primitive_loop_var_ranges, /*predicate=*/realize->predicate, - /*check_level=*/arith::IterMapLevel::Surjective, - /*analyzer=*/ffi::GetRef(analyzer), + /*check_level=*/sym::IterMapLevel::Surjective, + /*analyzer=*/ffi::GetRef(analyzer), /*simplify_trivial_iterators=*/false); if (res->indices.empty()) { return false; } - for (const arith::IterSumExpr& sum_expr : res->indices) { - const ffi::Array& args = sum_expr->args; + for (const sym::IterSumExpr& sum_expr : res->indices) { + const ffi::Array& args = sum_expr->args; if (!args.empty() && !is_one(args[0]->scale)) { return false; } @@ -643,7 +643,7 @@ void CheckPartialAffineBinding(const ScheduleState& self, SBlock block, } if (block_sref->parent && high_exclusive.has_value()) { // if it is not of global affine binding, check affineness under high_exclusive, - arith::Analyzer analyzer; + sym::Analyzer analyzer; ffi::Map dom_map = LoopDomainOfSRefTreePath(ffi::GetRef(block_sref->parent), high_exclusive); if (IsAffineBinding(GetSBlockRealize(self, block_sref), dom_map, analyzer.get())) { @@ -767,7 +767,7 @@ bool GetVarsTouchedByBlockIters(const SBlockRealize& block_realize, /******** Loop properties ********/ void CheckLoopStartsWithZero(const ScheduleState& self, const StmtSRef& loop_sref, - arith::AnalyzerObj* analyzer) { + sym::AnalyzerObj* analyzer) { class LoopNotStartWithZeroError : public ScheduleErrorContextObj { public: explicit LoopNotStartWithZeroError(IRModule mod, For loop) @@ -1353,7 +1353,7 @@ StmtSRef GetSRefTreeRoot(const StmtSRef& sref) { } void AddShapeVarBounds(const ScheduleState& state, const StmtSRefNode* sref, - arith::AnalyzerObj* analyzer) { + sym::AnalyzerObj* analyzer) { while (sref->parent != nullptr) { sref = sref->parent; } @@ -1733,7 +1733,7 @@ bool NeedsRFactorOrCrossThreadReduction(const s_tir::ScheduleState& self, // } } -PrimExpr SimplifyNonTrivialExpr(const PrimExpr& expr, arith::AnalyzerObj* analyzer) { +PrimExpr SimplifyNonTrivialExpr(const PrimExpr& expr, sym::AnalyzerObj* analyzer) { auto simplified = analyzer->Simplify(expr); if (simplified->IsInstance()) { return expr; @@ -1760,7 +1760,7 @@ struct TensorIntrinDescInfo { * \param desc_func The description PrimFunc * \return The auxilary information */ -TensorIntrinDescInfo ExtractTensorIntrinDescInfo(arith::AnalyzerObj* analyzer, +TensorIntrinDescInfo ExtractTensorIntrinDescInfo(sym::AnalyzerObj* analyzer, const PrimFunc& desc_func) { TensorIntrinDescInfo info; const auto* desc_scope_realize = desc_func->body.as(); @@ -1790,7 +1790,7 @@ ffi::Optional GetTensorizeLoopMapping(const s_tir::ScheduleState& const tirx::StmtSRef& block_sref, const tirx::PrimFunc& desc_func, bool allow_padding) { - arith::Analyzer analyzer; + sym::Analyzer analyzer; 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); @@ -1968,7 +1968,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { class AutoTensorizeMappingProposer { public: static ffi::Array ProposeMappings(const AutoTensorizeComparator* extractor, - arith::AnalyzerObj* analyzer) { + sym::AnalyzerObj* analyzer) { AutoTensorizeMappingProposer proposer(extractor, analyzer); proposer.CollectFeasibleSet(); return proposer.ProposeAllFuseMapping(); @@ -1976,7 +1976,7 @@ class AutoTensorizeMappingProposer { private: explicit AutoTensorizeMappingProposer(const AutoTensorizeComparator* extractor, - arith::AnalyzerObj* analyzer) + sym::AnalyzerObj* analyzer) : extractor_(extractor), analyzer_(analyzer) {} using VarSet = std::unordered_set; @@ -2142,7 +2142,7 @@ class AutoTensorizeMappingProposer { // tensor intrin. const AutoTensorizeComparator* extractor_; // The arithmetic analyzer. - arith::AnalyzerObj* analyzer_; + sym::AnalyzerObj* analyzer_; /*! \brief Potential mappings on RHS for each variable on LHS */ std::unordered_map lhs_feasible_vars_; }; @@ -2154,7 +2154,7 @@ bool CheckAutoTensorizeApplicable(const ScheduleState& state, const tirx::StmtSR // Step 2. Check if `desc_block` matches `block` // Ignore the scope of buffers when comparing, since we can do cache_read/write const SBlockRealize& block = GetSBlockRealize(state, block_sref); - arith::Analyzer analyzer; + sym::Analyzer analyzer; auto desc_info = ExtractTensorIntrinDescInfo(analyzer.get(), desc_func); return extractor->Dispatch(block->block, desc_info.desc_block->block); @@ -2173,7 +2173,7 @@ ffi::Optional GetAutoTensorizeMappingInfo( if (!CheckAutoTensorizeApplicable(self, block_sref, desc_func, &extractor)) { return std::nullopt; } - arith::Analyzer analyzer; + sym::Analyzer analyzer; ffi::Array mappings = AutoTensorizeMappingProposer::ProposeMappings(&extractor, analyzer.get()); if (mappings.empty()) { diff --git a/src/s_tir/schedule/analysis/layout.cc b/src/s_tir/schedule/analysis/layout.cc index 07781c85abfd..445e232d774e 100644 --- a/src/s_tir/schedule/analysis/layout.cc +++ b/src/s_tir/schedule/analysis/layout.cc @@ -80,11 +80,11 @@ class SplitExprCollector { static std::vector Collect(const PrimExpr& index, const ffi::Map& input_iters, // const PrimExpr& predicate, // - arith::IterMapLevel check_level, // - arith::AnalyzerObj* analyzer) { - arith::Analyzer analyzer_ref = ffi::GetRef(analyzer); - arith::IterMapResult res = arith::DetectIterMap({analyzer->Simplify(index)}, input_iters, - predicate, check_level, analyzer_ref); + sym::IterMapLevel check_level, // + sym::AnalyzerObj* analyzer) { + sym::Analyzer analyzer_ref = ffi::GetRef(analyzer); + sym::IterMapResult res = sym::DetectIterMap({analyzer->Simplify(index)}, input_iters, predicate, + check_level, analyzer_ref); const auto& iter_sum_exprs = res->indices; if (iter_sum_exprs.empty()) { return {}; @@ -102,7 +102,7 @@ class SplitExprCollector { } private: - void Visit(const arith::IterSplitExpr& expr) { + void Visit(const sym::IterSplitExpr& expr) { if (auto var = expr->source->source.as()) { const auto* lower_factor_imm = expr->lower_factor.as(); auto lower_factor = lower_factor_imm ? lower_factor_imm->value.as() : std::nullopt; @@ -113,15 +113,15 @@ class SplitExprCollector { return; } exprs_.push_back(SplitExpr{var.value(), *lower_factor, *extent}); - } else if (auto iter_sum_expr = expr->source->source.as()) { + } else if (auto iter_sum_expr = expr->source->source.as()) { Visit(iter_sum_expr.value()); } else { TVM_FFI_ICHECK(false) << "Unexpected type: " << expr->source->source->GetTypeKey(); } } - void Visit(const arith::IterSumExpr& expr) { - for (const arith::IterSplitExpr& arg : expr->args) { + void Visit(const sym::IterSumExpr& expr) { + for (const sym::IterSplitExpr& arg : expr->args) { Visit(arg); } } @@ -135,7 +135,7 @@ class SplitExprCollector { ffi::Optional SuggestIndexMap(const BufferVar& buffer, const ffi::Array& indices, const ffi::Array& loops, const PrimExpr& predicate, - arith::AnalyzerObj* analyzer) { + sym::AnalyzerObj* analyzer) { int ndim = buffer->shape.size(); int n_loops = loops.size(); // Step 1. Collect the domains and indices of loop variables @@ -159,7 +159,7 @@ ffi::Optional SuggestIndexMap(const BufferVar& buffer, // Step 3. Detect the IterSplitExpr of the indexing pattern std::vector split_exprs = SplitExprCollector::Collect( /*index=*/f_flatten_index(indices), input_iters, predicate, - /*check_level=*/arith::IterMapLevel::Surjective, analyzer); + /*check_level=*/sym::IterMapLevel::Surjective, analyzer); if (split_exprs.empty()) { return std::nullopt; } @@ -254,7 +254,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { refl::GlobalDef().def("s_tir.schedule.SuggestIndexMap", [](BufferVar buffer, ffi::Array indices, ffi::Array loops, PrimExpr predicate) { - arith::Analyzer analyzer; + sym::Analyzer analyzer; return SuggestIndexMap(buffer, indices, loops, predicate, analyzer.get()); }); } diff --git a/src/s_tir/schedule/analysis/reducer.cc b/src/s_tir/schedule/analysis/reducer.cc index 680d556df6f4..25b570f97143 100644 --- a/src/s_tir/schedule/analysis/reducer.cc +++ b/src/s_tir/schedule/analysis/reducer.cc @@ -33,7 +33,7 @@ using namespace tvm::tirx; /*! * \brief PrimExpr pattern matcher. * - * It is different from the pattern matcher in arith/pattern_match.h, which is dedicated + * It is different from the pattern matcher in sym/pattern_match.h, which is dedicated * for compile-time constant patterns. This pattern matcher can work on dynamic user-specific * patterns. * @@ -506,7 +506,7 @@ std::pair, ffi::Array> GetInitValuesAndUpdates const ffi::Array& expected_indices = updates[0]->indices; TVM_FFI_ICHECK_EQ(expected_shape.size(), expected_indices.size()); int n_dim = expected_indices.size(); - arith::Analyzer ana; + sym::Analyzer ana; for (int i = 0; i < n_buffers; ++i) { if (static_cast(updates[i]->buffer->shape.size()) != n_dim) { ErrorRFactorCrossThreadReductionNotApplicable(self, std::move(block), /*violated_cond=*/11); diff --git a/src/s_tir/schedule/concrete_schedule.cc b/src/s_tir/schedule/concrete_schedule.cc index 903203e5a8e6..71db5082b8ec 100644 --- a/src/s_tir/schedule/concrete_schedule.cc +++ b/src/s_tir/schedule/concrete_schedule.cc @@ -36,7 +36,7 @@ Schedule Schedule::Concrete(IRModule mod, LinearCongruentialEngine::TRandState s n->state_ = ScheduleState(mod, debug_mask, enable_check); n->error_render_level_ = error_render_level; n->symbol_table_ = {}; - n->analyzer_ = arith::Analyzer(); + n->analyzer_ = sym::Analyzer(); n->Seed(seed); GlobalVar gv; if (FindEntryFunc(mod, &gv) != nullptr) { @@ -204,7 +204,7 @@ Schedule ConcreteScheduleNode::Copy() { n->func_working_on_ = this->func_working_on_; n->error_render_level_ = this->error_render_level_; ConcreteScheduleNode::Copy(&n->state_, &n->symbol_table_); - n->analyzer_ = arith::Analyzer(); // new analyzer needed because it is stateful + n->analyzer_ = sym::Analyzer(); // new analyzer needed because it is stateful n->rand_state_ = ForkSeed(); return Schedule(std::move(n)); } diff --git a/src/s_tir/schedule/concrete_schedule.h b/src/s_tir/schedule/concrete_schedule.h index 205eda800a3e..952703944d5d 100644 --- a/src/s_tir/schedule/concrete_schedule.h +++ b/src/s_tir/schedule/concrete_schedule.h @@ -51,7 +51,7 @@ class ConcreteScheduleNode : public ScheduleNode { /*! \brief A symbol table that maps random variables to concrete StmtSRef/Integers */ TSymbolTable symbol_table_; /*! \brief A persistent stateless arithmetic analyzer. */ - arith::Analyzer analyzer_; + sym::Analyzer analyzer_; /*! \brief The value of random state for sampling. */ LinearCongruentialEngine::TRandState rand_state_; diff --git a/src/s_tir/schedule/ir_comparator.cc b/src/s_tir/schedule/ir_comparator.cc index 015ba132f661..e82df9bb13cf 100644 --- a/src/s_tir/schedule/ir_comparator.cc +++ b/src/s_tir/schedule/ir_comparator.cc @@ -735,7 +735,7 @@ bool AutoTensorizeComparator::Dispatch_(const SBlockNode* op, const Stmt& other) return false; } for (const IterVar& block_iter : op->iter_vars) { - inner_iter_dom_map_.Set(block_iter->var, arith::IntSet::FromRange(block_iter->dom)); + inner_iter_dom_map_.Set(block_iter->var, sym::IntSet::FromRange(block_iter->dom)); } } else { auto collect_iter = [&](const SBlockNode* op, std::vector& iters) -> bool { diff --git a/src/s_tir/schedule/ir_comparator.h b/src/s_tir/schedule/ir_comparator.h index 58db4f8ebdd2..2cf0444f2d81 100644 --- a/src/s_tir/schedule/ir_comparator.h +++ b/src/s_tir/schedule/ir_comparator.h @@ -111,12 +111,12 @@ class TensorizeComparator : public ExprComparator, public StmtComparator { /*! \brief Whether it is visiting the scope block (the outermost block). */ bool is_scope_block = true; /*! \brief The arithmetic analyzer for comparing LHS and RHS */ - arith::Analyzer analyzer_; + sym::Analyzer analyzer_; /*! * \brief The arithmetic analyzer for simplifying expressions on LHS. * This analyzer only contains the domains of the iterators on LHS. */ - arith::Analyzer lhs_analyzer_; + sym::Analyzer lhs_analyzer_; /*! \brief Additional error messages. Only used when assert_mode is true. */ std::vector error_messages_; // variable remap if any @@ -169,7 +169,7 @@ class AutoTensorizeComparator : public TensorizeComparator { private: /*! \brief The domain of the inner block iters. */ - ffi::Map inner_iter_dom_map_; + ffi::Map inner_iter_dom_map_; }; } // namespace s_tir diff --git a/src/s_tir/schedule/primitive/annotate_buffer_access.cc b/src/s_tir/schedule/primitive/annotate_buffer_access.cc index 41be161a8d83..ff24b191c0e1 100644 --- a/src/s_tir/schedule/primitive/annotate_buffer_access.cc +++ b/src/s_tir/schedule/primitive/annotate_buffer_access.cc @@ -96,7 +96,7 @@ void AnnotateBufferAccess(ScheduleState self, const StmtSRef& block_sref, int bu BufferVar buffer = GetNthAccessBuffer(self, ffi::GetRef(block), buffer_index, buffer_index_type); - arith::Analyzer analyzer; + sym::Analyzer analyzer; ffi::Array block_iter_vars; for (const IterVar& iter_var : block->iter_vars) { block_iter_vars.push_back(iter_var->var); diff --git a/src/s_tir/schedule/primitive/blockize_tensorize.cc b/src/s_tir/schedule/primitive/blockize_tensorize.cc index fc8811822407..6b815827878c 100644 --- a/src/s_tir/schedule/primitive/blockize_tensorize.cc +++ b/src/s_tir/schedule/primitive/blockize_tensorize.cc @@ -92,12 +92,13 @@ class SubspaceNotDivisibleError : public ScheduleErrorContextObj { * \param inner_iters The iters of the inner space * \return The result of the subspace division. */ -ffi::Array> TrivialSubspaceDivision( - const ffi::Array& iter_vars, const ffi::Array& bindings, - const PrimExpr& predicate, const ffi::Array& outer_iters, - const ffi::Array& inner_iters) { +ffi::Array> TrivialSubspaceDivision(const ffi::Array& iter_vars, + const ffi::Array& bindings, + const PrimExpr& predicate, + const ffi::Array& outer_iters, + const ffi::Array& inner_iters) { if (!is_one(predicate)) return {}; - ffi::Array> res; + ffi::Array> res; std::unordered_set outer_loop_vars; std::unordered_set inner_loop_vars; @@ -118,18 +119,18 @@ ffi::Array> TrivialSubspaceDivision( }; auto use_outer_loop_vars = make_uses_var(outer_iters); auto use_inner_loop_vars = make_uses_var(inner_iters); - arith::IterMark unit_iter_mark(arith::IterSumExpr({}, 0), 1); + sym::IterMark unit_iter_mark(sym::IterSumExpr({}, 0), 1); for (int i = 0, n = bindings.size(); i < n; ++i) { bool outer = use_outer_loop_vars(bindings[i]); bool inner = use_inner_loop_vars(bindings[i]); - arith::IterMark iter_mark; + sym::IterMark iter_mark; if (bindings[i].as()) { - iter_mark = arith::IterMark( - arith::IterSplitExpr(arith::IterMark(bindings[i], iter_vars[i]->dom->extent)), - iter_vars[i]->dom->extent); + iter_mark = + sym::IterMark(sym::IterSplitExpr(sym::IterMark(bindings[i], iter_vars[i]->dom->extent)), + iter_vars[i]->dom->extent); } else { - iter_mark = arith::IterMark(arith::IterSumExpr({}, bindings[i]), iter_vars[i]->dom->extent); + iter_mark = sym::IterMark(sym::IterSumExpr({}, bindings[i]), iter_vars[i]->dom->extent); } if (outer && !inner) { res.push_back({/*outer_iter=*/iter_mark, /*inner_iter=*/unit_iter_mark}); @@ -141,8 +142,8 @@ ffi::Array> TrivialSubspaceDivision( return {}; } } - res.push_back({arith::IterMark(arith::IterSumExpr({}, 0), IntImm::Bool(true)), - arith::IterMark(arith::IterSumExpr({}, 0), IntImm::Bool(true))}); + res.push_back({sym::IterMark(sym::IterSumExpr({}, 0), IntImm::Bool(true)), + sym::IterMark(sym::IterSumExpr({}, 0), IntImm::Bool(true))}); return res; } @@ -162,13 +163,13 @@ ffi::Array> TrivialSubspaceDivision( * \param preserve_unit_iters Whether or not to preserve unit iterators in block bindings * \param loop_sref_as_outer Whether loop_sref is divided into outer or inner */ -ffi::Array> SubspaceDivide(const SBlockRealize& realize, - const StmtSRef& block_sref, // - const StmtSRef& loop_sref, // - std::vector* loops, - arith::AnalyzerObj* analyzer, - bool preserve_unit_iters, - bool loop_sref_as_outer = false) { +ffi::Array> SubspaceDivide(const SBlockRealize& realize, + const StmtSRef& block_sref, // + const StmtSRef& loop_sref, // + std::vector* loops, + sym::AnalyzerObj* analyzer, + bool preserve_unit_iters, + bool loop_sref_as_outer = false) { ffi::Array inner_vars; ffi::Array outer_vars; ffi::Array primitive_inner_vars; @@ -190,9 +191,9 @@ ffi::Array> SubspaceDivide(const SBlockRealize& real inner = false; } } - ffi::Array> result = arith::SubspaceDivide( + ffi::Array> result = sym::SubspaceDivide( realize->iter_values, primitive_loop_var_domain, primitive_inner_vars, realize->predicate, - arith::IterMapLevel::Surjective, ffi::GetRef(analyzer), + sym::IterMapLevel::Surjective, ffi::GetRef(analyzer), /*simplify_trivial_iterators=*/!preserve_unit_iters); if (!result.empty()) { return result; @@ -215,23 +216,23 @@ ffi::Array> SubspaceDivide(const SBlockRealize& real * \return A substitution plan to the iterators in the original inner block. */ ffi::Map DeriveBlockBinding( - const ffi::Array& iter_vars, // - const ffi::Array>& division, // - ffi::Array* outer_iter_vars, // - ffi::Array* outer_bindings, // - ffi::Array* inner_iter_vars, // - ffi::Array* inner_bindings, // + const ffi::Array& iter_vars, // + const ffi::Array>& division, // + ffi::Array* outer_iter_vars, // + ffi::Array* outer_bindings, // + ffi::Array* inner_iter_vars, // + ffi::Array* inner_bindings, // bool preserve_unit_iters, bool reuse_outer = false) { - using arith::IterMapExpr; - using arith::IterMapExprNode; - using arith::NormalizeIterMapToExpr; + using sym::IterMapExpr; + using sym::IterMapExprNode; + using sym::NormalizeIterMapToExpr; ffi::Map block_var_subst; TVM_FFI_ICHECK_EQ(iter_vars.size() + 1, division.size()); - arith::Analyzer ana; + sym::Analyzer ana; for (int i = 0, n = iter_vars.size(); i < n; ++i) { const IterVar& iter_var = iter_vars[i]; - arith::IterMark outer_mark = division[i][0]; - arith::IterMark inner_mark = division[i][1]; + sym::IterMark outer_mark = division[i][0]; + sym::IterMark inner_mark = division[i][1]; IterMapExpr outer_binding = outer_mark->source.as_or_throw(); IterMapExpr inner_binding = inner_mark->source.as_or_throw(); // After computing the subspace division, bindings[i] can be written as @@ -403,14 +404,14 @@ Stmt GenerateOuterInit(const Stmt& block_init, const SBlockRealize& inner_realiz * \return The substituted stmt. */ Stmt ReplaceAndSimplify(const Stmt& stmt, const ffi::Map& sub, - ffi::Map* block_sref_reuse, arith::AnalyzerObj* analyzer) { + ffi::Map* block_sref_reuse, sym::AnalyzerObj* analyzer) { struct Replacer : public StmtExprMutator { public: using StmtExprMutator::Mutate; using StmtExprMutator::Mutate_; explicit Replacer(const ffi::Map& sub, - ffi::Map* block_sref_reuse, arith::AnalyzerObj* analyzer) + ffi::Map* block_sref_reuse, sym::AnalyzerObj* analyzer) : block_sref_reuse_(block_sref_reuse), analyzer_(analyzer) { for (const auto& [var, replacement] : sub) VarRemapSet(var, replacement); } @@ -437,7 +438,7 @@ Stmt ReplaceAndSimplify(const Stmt& stmt, const ffi::Map& sub, } ffi::Map* block_sref_reuse_; - arith::AnalyzerObj* analyzer_; + sym::AnalyzerObj* analyzer_; }; return ffi::make_object(sub, block_sref_reuse, analyzer) ->Mutate(stmt) @@ -451,12 +452,12 @@ Stmt ReplaceAndSimplify(const Stmt& stmt, const ffi::Map& sub, * \return The relaxed regions */ ffi::Array EvalSetRegions(const ffi::Array& regions, - const ffi::Map& dom_map) { + const ffi::Map& dom_map) { ffi::Array results; results.reserve(regions.size()); for (const TensorRegion& buffer_region : regions) { const BufferVar& buffer = buffer_region->source.as_or_throw(); - ffi::Array relaxed = arith::EvalSet(buffer_region->region, dom_map); + ffi::Array relaxed = sym::EvalSet(buffer_region->region, dom_map); TVM_FFI_ICHECK_EQ(relaxed.size(), buffer->shape.size()); int ndim = buffer->shape.size(); ffi::Array new_region; @@ -475,17 +476,16 @@ ffi::Array EvalSetRegions(const ffi::Array& regions, * \return The union regions */ ffi::Array UnionRegions(const ffi::Array& regions) { - typedef std::vector> ranges_t; + typedef std::vector> ranges_t; std::unordered_map intset_map; for (const TensorRegion& buffer_region : regions) { const BufferVar& buffer = buffer_region->source.as_or_throw(); if (intset_map.find(buffer) == intset_map.end()) { - intset_map[buffer] = {buffer->shape.size(), ffi::Array()}; + intset_map[buffer] = {buffer->shape.size(), ffi::Array()}; } - std::vector> dim_range(buffer->shape.size(), - ffi::Array()); + std::vector> dim_range(buffer->shape.size(), ffi::Array()); for (size_t dim = 0; dim < buffer->shape.size(); ++dim) { - intset_map[buffer][dim].push_back(arith::IntSet::FromRange(buffer_region->region[dim])); + intset_map[buffer][dim].push_back(sym::IntSet::FromRange(buffer_region->region[dim])); } } ffi::Array results; @@ -493,7 +493,7 @@ ffi::Array UnionRegions(const ffi::Array& regions) { const BufferVar& buffer = it.first; ffi::Array regions; for (size_t dim = 0; dim < buffer->shape.size(); ++dim) { - const arith::IntSet intset = arith::Union(it.second[dim]); + const sym::IntSet intset = sym::Union(it.second[dim]); regions.push_back({intset.min(), intset.max() + 1}); } results.push_back(BufferRegion(buffer, regions)); @@ -517,7 +517,7 @@ Stmt MakeLoopNest(Stmt stmt, const std::vector& loops) { } SBlockRealize BlockizeImpl(const ScheduleState& self, const StmtSRef& loop_sref, - ffi::Map* block_sref_reuse, arith::AnalyzerObj* analyzer, + ffi::Map* block_sref_reuse, sym::AnalyzerObj* analyzer, bool preserve_unit_iters) { TVM_SREF_TO_FOR(loop_sref); // Step 1: Check and get the only block under `loop`. @@ -526,7 +526,7 @@ SBlockRealize BlockizeImpl(const ScheduleState& self, const StmtSRef& loop_sref, StmtSRef block_sref = self->stmt2ref.at(block.get()); // Step 2: Derive subspace division std::vector loops; - ffi::Array> division = + ffi::Array> division = SubspaceDivide(block_realize, block_sref, loop_sref, &loops, analyzer, preserve_unit_iters); if (division.empty()) { throw MakeScheduleError(self->mod, ffi::GetRef(loops.back()), @@ -545,9 +545,9 @@ SBlockRealize BlockizeImpl(const ScheduleState& self, const StmtSRef& loop_sref, &inner_iter_vars, &inner_bindings, // preserve_unit_iters); // Step 4: Do var substitution to adjust to the new block bindings - ffi::Map inner_iter_dom; + ffi::Map inner_iter_dom; for (const IterVar& iter : inner_iter_vars) { - inner_iter_dom.Set(iter->var, arith::IntSet::FromRange(iter->dom)); + inner_iter_dom.Set(iter->var, sym::IntSet::FromRange(iter->dom)); analyzer->Bind(iter->var, iter->dom); } SBlock block_subst = @@ -588,7 +588,7 @@ SBlockRealize BlockizeImpl(const ScheduleState& self, const StmtSRef& loop_sref, } StmtSRef Blockize(ScheduleState self, const StmtSRef& loop_sref, bool preserve_unit_iters) { - arith::Analyzer analyzer; + sym::Analyzer analyzer; ffi::Map block_sref_reuse; SBlockRealize blockized = BlockizeImpl(self, loop_sref, &block_sref_reuse, analyzer.get(), preserve_unit_iters); @@ -612,13 +612,13 @@ SBlockRealize BlockizeBlocks(const ScheduleState& self, const ffi::Array write_regions; std::string outer_block_name = "outer_"; ffi::Map loop_var_subst; - arith::Analyzer analyzer; + sym::Analyzer analyzer; for (const auto& block_sref : block_srefs) { auto block_realize = GetSBlockRealize(self, block_sref); auto block = block_realize->block; // Step 1: Derive subspace division std::vector loops; - ffi::Array> division = SubspaceDivide( + ffi::Array> division = SubspaceDivide( block_realize, block_sref, lca, &loops, analyzer.get(), preserve_unit_iters, true); if (division.empty()) { throw MakeScheduleError(self->mod, ffi::GetRef(loops.back()), @@ -645,7 +645,7 @@ SBlockRealize BlockizeBlocks(const ScheduleState& self, const ffi::Array inner_iter_dom; + ffi::Map inner_iter_dom; for (const IterVar& iter : inner_iter_vars) { PrimExpr min = ffi::StructuralMap(iter->dom->min, f_substitute) .as_or_throw(); @@ -653,7 +653,7 @@ SBlockRealize BlockizeBlocks(const ScheduleState& self, const ffi::Array(iter->dom->extent, f_substitute) .as_or_throw(); Range dom = Range::FromMinExtent(min, extent); - inner_iter_dom.Set(iter->var, arith::IntSet::FromRange(dom)); + inner_iter_dom.Set(iter->var, sym::IntSet::FromRange(dom)); analyzer->Bind(iter->var, dom); } SBlock block_subst = @@ -808,7 +808,7 @@ void Tensorize(ScheduleState self, const StmtSRef& sref, const TensorIntrin& int block_realize = GetSBlockRealize(self, sref); old_block = block_realize->block; } else if (sref->stmt->IsInstance()) { - arith::Analyzer analyzer; + sym::Analyzer analyzer; ffi::Map block_sref_reuse; block_realize = BlockizeImpl(self, sref, &block_sref_reuse, analyzer.get(), preserve_unit_iters); @@ -818,7 +818,7 @@ void Tensorize(ScheduleState self, const StmtSRef& sref, const TensorIntrin& int throw; } - arith::Analyzer analyzer; + sym::Analyzer analyzer; PrimFunc intrin_desc = s_tir::StmtSimplify(intrin->desc, analyzer); PrimFunc intrin_impl = DeepCopy(intrin->impl); diff --git a/src/s_tir/schedule/primitive/cache_index.cc b/src/s_tir/schedule/primitive/cache_index.cc index 630c4a681482..51a1ba7fbfb5 100644 --- a/src/s_tir/schedule/primitive/cache_index.cc +++ b/src/s_tir/schedule/primitive/cache_index.cc @@ -16,11 +16,11 @@ * specific language governing permissions and limitations * under the License. */ -#include #include #include #include #include +#include #include "../../../tirx/transform/replace_selected_expr.h" #include "../utils.h" @@ -61,8 +61,8 @@ struct IndexInfo { * \param range The range of the integer. * \returns A data type that covers the input range. */ -PrimType DeterminePrimType(const arith::IntSet& range) { - arith::Analyzer ana; +PrimType DeterminePrimType(const sym::IntSet& range) { + sym::Analyzer ana; if (ana->CanProve(range.min() >= INT32_MIN && range.max() <= INT32_MAX)) { return PrimType::Int(32); } else { @@ -280,7 +280,7 @@ ffi::Array MakeIndexCacheStage(IndexInfo* info, const ffi::String& stora ffi::Array buffer_shape; for (const Var& it : info->origin_block_vars[expr_index]) { buffer_shape.push_back( - arith::EvalSet(info->var_binding.at(it), arith::AsIntSet(info->range_map)).max() + 1); + sym::EvalSet(info->var_binding.at(it), sym::AsIntSet(info->range_map)).max() + 1); } info->cache_buffer.push_back(BufferVar( index_buffer_name, BufferType(storage_scope, data_ty, buffer_shape, {1}, {0}, 0, 0))); @@ -289,7 +289,7 @@ ffi::Array MakeIndexCacheStage(IndexInfo* info, const ffi::String& stora std::vector loop_vars; ffi::Map replace_table; for (const Var& it : iter_vars) { - PrimType data_ty = DeterminePrimType(arith::IntSet::FromRange(info->range_map.at(it))); + PrimType data_ty = DeterminePrimType(sym::IntSet::FromRange(info->range_map.at(it))); PrimVar loop_var("ax" + std::to_string(replace_table.size()), data_ty); loop_vars.push_back(loop_var); replace_table.Set(it, loop_var); @@ -513,7 +513,7 @@ ffi::Array CacheIndex(ScheduleState self, const StmtSRef& block_sref, if (result_block_sref->parent == nullptr) { affine_binding = true; } else { - arith::Analyzer analyzer; + sym::Analyzer analyzer; StmtSRef parent_sref = ffi::GetRef(result_block_sref->parent); affine_binding = IsAffineBinding(/*realize=*/GetSBlockRealize(self, result_block_sref), /*loop_var_ranges=*/LoopDomainOfSRefTreePath(parent_sref), diff --git a/src/s_tir/schedule/primitive/cache_index_helpers.cc b/src/s_tir/schedule/primitive/cache_index_helpers.cc index c88a2b9b3284..28ddb38fd67c 100644 --- a/src/s_tir/schedule/primitive/cache_index_helpers.cc +++ b/src/s_tir/schedule/primitive/cache_index_helpers.cc @@ -25,11 +25,11 @@ #include "cache_index_helpers.h" -#include // For the arith::Analyzer::Simplify() method simplifying terms #include #include #include #include +#include // For the sym::Analyzer::Simplify() method simplifying terms #include #include #include @@ -409,7 +409,7 @@ bool EqualTerms(const PrimExpr& a, const PrimExpr& b) { */ PrimExpr NormalizeTerm(const PrimExpr& expr, bool do_normalization) { if (do_normalization) { - arith::Analyzer analyzer; + sym::Analyzer analyzer; return analyzer->Simplify(expr); } else { return expr; diff --git a/src/s_tir/schedule/primitive/cache_read_write.cc b/src/s_tir/schedule/primitive/cache_read_write.cc index d0d9ec36ef98..de27c902ed2c 100644 --- a/src/s_tir/schedule/primitive/cache_read_write.cc +++ b/src/s_tir/schedule/primitive/cache_read_write.cc @@ -473,7 +473,7 @@ bool CalculateAffineFlag(const ScheduleState& self, const StmtSRef& block_sref) if (block_sref->parent == nullptr) { return true; } - arith::Analyzer analyzer; + sym::Analyzer analyzer; StmtSRef parent_sref = ffi::GetRef(block_sref->parent); return IsAffineBinding(/*realize=*/GetSBlockRealize(self, block_sref), /*loop_var_ranges=*/LoopDomainOfSRefTreePath(parent_sref), @@ -666,7 +666,7 @@ TensorRegion RelaxBufferRegion(ScheduleState self, const TensorRegion& buffer_re SBlockRealize realize = GetSBlockRealize(self, block_sref); ffi::Map binding = GetBindings(realize); const BufferVar& buffer = buffer_region->source.as_or_throw(); - arith::Analyzer analyzer; + sym::Analyzer analyzer; auto f_substitute = [&binding](const Var& var) -> ffi::Expected> { if (auto repl = binding.Get(var)) return ffi::Any(*std::move(repl)); return ffi::Unchanged(); @@ -679,7 +679,7 @@ TensorRegion RelaxBufferRegion(ScheduleState self, const TensorRegion& buffer_re return Range::FromMinExtent(min, extent); }); TensorRegion subst_region = BufferRegion(buffer, mapped_region); - ffi::Array int_sets = AnalyzeRegionUpperBound( + ffi::Array int_sets = AnalyzeRegionUpperBound( /*region=*/subst_region, /*predicate=*/ ffi::StructuralMap(realize->predicate && extra_predicate, @@ -1179,7 +1179,7 @@ class CacheReadRewriter : public StmtExprMutator { */ bool cache_full_region_; /*! \brief Arithmetic analyzer. */ - arith::Analyzer ana_; + sym::Analyzer ana_; friend ReindexCacheReadRewriter; }; @@ -1524,7 +1524,7 @@ class CacheWriteRewriter : public StmtExprMutator { */ bool cache_full_region_; /*! \brief Arithmetic analyzer. */ - arith::Analyzer ana_; + sym::Analyzer ana_; friend ReindexCacheWriteRewriter; }; @@ -2219,7 +2219,7 @@ void CollectReindexCacheStageInfoAndCreateBuffer( ReindexCacheStageInfo* info, const IRModule& mod, const StmtSRef& block_sref, const ffi::String& storage_scope, const IndexMap& index_map, const SBlock& block, const SBlockRealize& realize, const BufferVar& old_buffer, const TensorRegion& cache_region) { - arith::Analyzer analyzer; + sym::Analyzer analyzer; ffi::Array block_iter_vars, block_shape; for (const IterVar& iter_var : block->iter_vars) { block_iter_vars.push_back(iter_var); @@ -2552,7 +2552,7 @@ StmtSRef ReIndex(ScheduleState self, const StmtSRef& block_sref, int buffer_inde SBlock block = ffi::GetRef(block_ptr); BufferVar buffer = GetNthAccessBuffer(self, block, buffer_index, buffer_index_type); StmtSRef scope_sref = GetScopeRoot(self, block_sref, /*require_stage_pipeline=*/true); - arith::Analyzer analyzer; + sym::Analyzer analyzer; // Step 1. Collect the original indices and check there's only single pattern of related // Load/Store and the buffer is not accessed opaquely diff --git a/src/s_tir/schedule/primitive/compute_at.cc b/src/s_tir/schedule/primitive/compute_at.cc index 9c73b9a41ed4..3790a26b44fd 100644 --- a/src/s_tir/schedule/primitive/compute_at.cc +++ b/src/s_tir/schedule/primitive/compute_at.cc @@ -85,7 +85,7 @@ class NotInSameScopeError : public ScheduleErrorContextObj { public: static void CheckAndBindLoopDomain(const ScheduleState& self, const StmtSRef& block_sref, const StmtSRef& loop_sref, const StmtSRef& scope_root_sref, - arith::AnalyzerObj* analyzer) { + sym::AnalyzerObj* analyzer) { for (const StmtSRefNode* p = loop_sref.get();; p = p->parent) { if (const ForNode* loop = p->StmtAs()) { analyzer->Bind(loop->loop_var, Range::FromMinExtent(loop->min, loop->extent)); @@ -193,22 +193,22 @@ int FindInsertionPoint( * extra predicates for non-trivial bound. The domain info class can also union with each other. */ struct BlockVarDomainInfo { - arith::IntSet dom{arith::IntSet::Nothing()}; // dom is ensured to be bounded - arith::IntSet bound{arith::IntSet::Nothing()}; + sym::IntSet dom{sym::IntSet::Nothing()}; // dom is ensured to be bounded + sym::IntSet bound{sym::IntSet::Nothing()}; /*! \brief Relaxed union operation */ void Union(const BlockVarDomainInfo& other) { // just relax (d0 ^ b0) v (d1 ^ b1) to (d0 v d1) ^ (b0 v b1) - dom = arith::Union({dom, other.dom}); - bound = arith::Union({bound, other.bound}); + dom = sym::Union({dom, other.dom}); + bound = sym::Union({bound, other.bound}); } /*! \brief Simplify domain info */ - void Simplify(arith::AnalyzerObj* analyzer) { - auto to_simplified = [analyzer](const arith::IntSet& set) { + void Simplify(sym::AnalyzerObj* analyzer) { + auto to_simplified = [analyzer](const sym::IntSet& set) { PrimExpr min = set.HasLowerBound() ? analyzer->Simplify(set.min()) : set.min(); PrimExpr max = set.HasUpperBound() ? analyzer->Simplify(set.max()) : set.max(); - return arith::IntSet::Interval(min, max); + return sym::IntSet::Interval(min, max); }; // if no dom specified, try use bound as dom if (dom.IsNothing()) { @@ -222,18 +222,18 @@ struct BlockVarDomainInfo { dom = to_simplified(dom); bound = to_simplified(bound); // if can proof the dom is within bound, remove bound - auto intersect = to_simplified(arith::Intersect({dom, bound})); + auto intersect = to_simplified(sym::Intersect({dom, bound})); if (analyzer->CanProveEqual(dom.min(), intersect.min()) && analyzer->CanProveEqual(dom.max(), intersect.max())) { - bound = arith::IntSet::Nothing(); + bound = sym::IntSet::Nothing(); } else if (analyzer->CanProveEqual(bound.min(), intersect.min()) && analyzer->CanProveEqual(bound.max(), intersect.max())) { dom = bound; - bound = arith::IntSet::Nothing(); + bound = sym::IntSet::Nothing(); } else if (is_const_int(intersect.min()) && is_const_int(intersect.max())) { // if the bound induce constant iter range, merge bound to loop domain dom = intersect; - bound = arith::IntSet::Nothing(); + bound = sym::IntSet::Nothing(); } } }; @@ -263,7 +263,7 @@ class ScopeReconstructor : public StmtExprMutator { * \param preserve_unit_loops Whether to generate unit loops where the loop extent is 1 */ void MakeNewLoop(int insert_position, std::vector iter_doms, - arith::AnalyzerObj* analyzer, bool preserve_unit_loops) { + sym::AnalyzerObj* analyzer, bool preserve_unit_loops) { int n_iters = iter_doms.size(); ffi::Array loop_vars; ffi::Array loop_extents; @@ -285,18 +285,18 @@ class ScopeReconstructor : public StmtExprMutator { } else { iter_values.push_back(iter_dom->min); } - const arith::IntSet& pred_bound = iter_doms[i].bound; + const sym::IntSet& pred_bound = iter_doms[i].bound; if (!pred_bound.IsNothing()) { // NOTE: Apply strong analyzer proofs to get rid of symbolic bound if (pred_bound.HasLowerBound()) { PrimExpr lower_bound = iter_values[i] >= pred_bound.min(); - if (!analyzer->CanProve(lower_bound, arith::ProofStrength::kSymbolicBound)) { + if (!analyzer->CanProve(lower_bound, sym::ProofStrength::kSymbolicBound)) { predicate = predicate && lower_bound; } } if (pred_bound.HasUpperBound()) { PrimExpr upper_bound = iter_values[i] < pred_bound.max() + 1; - if (!analyzer->CanProve(upper_bound, arith::ProofStrength::kSymbolicBound)) { + if (!analyzer->CanProve(upper_bound, sym::ProofStrength::kSymbolicBound)) { predicate = predicate && upper_bound; } } @@ -380,7 +380,7 @@ void RelaxBufferRegions(const ffi::Map& binding, runtime::StorageScope global_scope{runtime::StorageRank::kGlobal, ""}; // We cache the variable domains runtime::StorageRank previous_rank = runtime::StorageRank::kGlobal; - ffi::Optional> var_dom = std::nullopt; + ffi::Optional> var_dom = std::nullopt; auto f_substitute = [&binding](const Var& var) -> ffi::Expected> { if (auto repl = binding.Get(var)) return ffi::Any(*std::move(repl)); return ffi::Unchanged(); @@ -401,7 +401,7 @@ void RelaxBufferRegions(const ffi::Map& binding, runtime::StorageRank rank = scope.rank; if (rank != previous_rank || !var_dom.has_value()) { previous_rank = rank; - var_dom = arith::AsIntSet(LoopDomainOfSRefTreePath( + var_dom = sym::AsIntSet(LoopDomainOfSRefTreePath( /*low_inclusive=*/relax_path_low_inclusive, /*high_exclusive=*/relax_path_high_exclusive, /*extra_relax_scope=*/scope)); @@ -414,7 +414,7 @@ void RelaxBufferRegions(const ffi::Map& binding, .template as_or_throw(); return Range::FromMinExtent(min, extent); }); - ffi::Array relaxed_region = arith::EvalSet(mapped_region, var_dom.value()); + ffi::Array relaxed_region = sym::EvalSet(mapped_region, var_dom.value()); relaxed_regions.push_back({relaxed_region.begin(), relaxed_region.end()}); } } @@ -427,30 +427,30 @@ void RelaxBufferRegions(const ffi::Map& binding, * \param dim_max The maximum index bound by the buffer shape * \param analyzer The arithmetic analyzer */ -std::pair SolveBlockVarDomain(const arith::IntSet& provided, - const arith::IntSet& required, +std::pair SolveBlockVarDomain(const sym::IntSet& provided, + const sym::IntSet& required, PrimExpr dim_max, - arith::AnalyzerObj* analyzer) { + sym::AnalyzerObj* analyzer) { PrimExpr provided_min = analyzer->Simplify(provided.min()); PrimExpr provided_max = analyzer->Simplify(provided.max()); PrimExpr required_min = analyzer->Simplify(required.min()); PrimExpr required_max = analyzer->Simplify(required.max()); - arith::IntSet var_dom, var_bound; + sym::IntSet var_dom, var_bound; ffi::Optional var; - arith::PVar p_v; - arith::PVar p_e; + sym::PVar p_v; + sym::PVar p_e; if ((p_v * p_e).Match(provided_min) || (p_e * p_v).Match(provided_min)) { PrimExpr e = p_e.Eval(); var = p_v.Eval(); - var_dom = arith::IntSet::Interval(floordiv(required_min, e), floordiv(required_max, e)); - var_bound = arith::IntSet::Interval(0, floordiv(dim_max, e)); + var_dom = sym::IntSet::Interval(floordiv(required_min, e), floordiv(required_max, e)); + var_bound = sym::IntSet::Interval(0, floordiv(dim_max, e)); } else if (analyzer->CanProveEqual(provided_min, provided_max)) { if (p_v.Match(provided_min)) { var = p_v.Eval(); - var_dom = arith::IntSet::Interval(required_min, required_max); - var_bound = arith::IntSet::Interval(0, dim_max); + var_dom = sym::IntSet::Interval(required_min, required_max); + var_bound = sym::IntSet::Interval(0, dim_max); } else { - arith::PVar p_f1, p_f2; + sym::PVar p_f1, p_f2; if ((floordiv(p_f1, p_f2).Match(provided_min))) { PrimExpr var_expr = p_f1.Eval(); PrimExpr fac = p_f2.Eval(); @@ -458,12 +458,12 @@ std::pair SolveBlockVarDomain(const arith::IntSet& prov if (var_expr.as()) { // a <= (x // factor) <= b, fac > 0 ==> (a * fac) <= x <= (b * fac + fac - 1) var = var_expr.as_or_throw(); - var_dom = arith::IntSet::Interval(required_min * fac, - analyzer->Simplify(required_max * fac + fac - 1)); - var_bound = arith::IntSet::Interval(0, analyzer->Simplify(dim_max * fac + fac - 1)); + var_dom = sym::IntSet::Interval(required_min * fac, + analyzer->Simplify(required_max * fac + fac - 1)); + var_bound = sym::IntSet::Interval(0, analyzer->Simplify(dim_max * fac + fac - 1)); } else { - const arith::IntSet new_provided = arith::IntSet::SinglePoint(p_f1.Eval()); - const arith::IntSet new_required = arith::IntSet::Interval( + const sym::IntSet new_provided = sym::IntSet::SinglePoint(p_f1.Eval()); + const sym::IntSet new_required = sym::IntSet::Interval( required_min * fac, analyzer->Simplify(required_max * fac + fac - 1)); return SolveBlockVarDomain(new_provided, new_required, dim_max, analyzer); } @@ -479,10 +479,10 @@ std::pair SolveBlockVarDomain(const arith::IntSet& prov PrimExpr mod_2 = p_f2.Eval(); if (analyzer->CanProveGreaterEqual(mod_1, 1) && analyzer->CanProveGreaterEqual(mod_2, 1)) { - const arith::IntSet new_provided = arith::IntSet::SinglePoint(p_f1.Eval()); + const sym::IntSet new_provided = sym::IntSet::SinglePoint(p_f1.Eval()); if (analyzer->CanProveGreaterEqual(required_min, 0)) { - const arith::IntSet new_required = - arith::IntSet::Interval(required_min, arith::SymbolicLimits::pos_inf_); + const sym::IntSet new_required = + sym::IntSet::Interval(required_min, sym::SymbolicLimits::pos_inf_); return SolveBlockVarDomain(new_provided, new_required, dim_max, analyzer); } } @@ -506,14 +506,13 @@ std::pair SolveBlockVarDomain(const arith::IntSet& prov */ void UpdateBlockVarDomainDimwise( const VarNode* buffer, const NDIntSet& provided_region, const NDIntSet& required_region, - arith::AnalyzerObj* analyzer, - std::unordered_map* iter_doms) { + sym::AnalyzerObj* analyzer, std::unordered_map* iter_doms) { size_t ndim = GetBufferVar(buffer)->shape.size(); for (size_t i = 0; i < ndim; ++i) { - arith::IntSet provided = provided_region[i]; - arith::IntSet required = required_region[i]; + sym::IntSet provided = provided_region[i]; + sym::IntSet required = required_region[i]; PrimExpr dim_max = max(GetBufferVar(buffer)->shape[i] - 1, 0); - arith::Analyzer analyzer_ref = ffi::GetRef(analyzer); + sym::Analyzer analyzer_ref = ffi::GetRef(analyzer); if (provided.CanProveSinglePoint(analyzer_ref) && is_const_int(provided.min())) { TVM_FFI_ICHECK(required.CanProveSinglePoint(analyzer_ref) && @@ -533,9 +532,9 @@ void UpdateBlockVarDomainDimwise( } /*! \brief Helper function to implement intset version of `InverseAffineIterMap`. */ -ffi::Map InverseAffineIterMap(const ffi::Array& iter_map, - const NDIntSet& outputs, - arith::AnalyzerObj* analyzer) { +ffi::Map InverseAffineIterMap(const ffi::Array& iter_map, + const NDIntSet& outputs, + sym::AnalyzerObj* analyzer) { ffi::Array min_point, max_point; min_point.reserve(outputs.size()); max_point.reserve(outputs.size()); @@ -546,16 +545,15 @@ ffi::Map InverseAffineIterMap(const ffi::Array dom_map; + ffi::Map dom_map; for (const auto& kv : rev_min) { const Var& var = kv.first; auto it = rev_max.find(var); TVM_FFI_ICHECK(it != rev_max.end()); // InverseAffineIterMap's result vars are assumed stable const PrimExpr& rev_min_point = kv.second; const PrimExpr& rev_max_point = (*it).second; - dom_map.Set(var, - arith::IntSet::Interval(analyzer->Simplify(min(rev_min_point, rev_max_point)), - analyzer->Simplify(max(rev_min_point, rev_max_point)))); + dom_map.Set(var, sym::IntSet::Interval(analyzer->Simplify(min(rev_min_point, rev_max_point)), + analyzer->Simplify(max(rev_min_point, rev_max_point)))); } return dom_map; } @@ -573,10 +571,10 @@ ffi::Map InverseAffineIterMap(const ffi::Array& iter_vars, const NDIntSet& provided_region, const NDIntSet& required_region, - arith::AnalyzerObj* analyzer, + sym::AnalyzerObj* analyzer, std::unordered_map* iter_doms) { // we only support single point provided region now, which could cover most cases - arith::Analyzer analyzer_ref = ffi::GetRef(analyzer); + sym::Analyzer analyzer_ref = ffi::GetRef(analyzer); for (const auto& intset : provided_region) { if (!intset.CanProveSinglePoint(analyzer_ref)) return false; } @@ -591,20 +589,20 @@ bool UpdateBlockVarDomainAffine(const VarNode* buffer, const ffi::Array for (size_t i = 0; i < ndim; ++i) { provide_indices.push_back(provided_region[i].min()); } - auto res = arith::DetectIterMap(provide_indices, dom_map, IntImm::Bool(true), - arith::IterMapLevel::Bijective, analyzer_ref, false); + auto res = sym::DetectIterMap(provide_indices, dom_map, IntImm::Bool(true), + sym::IterMapLevel::Bijective, analyzer_ref, false); if (res->indices.empty()) { return false; } // calculate backward mapping (required region point -> block vars) NDIntSet required_bound; for (size_t i = 0; i < ndim; ++i) { - required_bound.push_back(arith::IntSet::Interval(IntImm(GetBufferVar(buffer)->shape[i].ty(), 0), - max(GetBufferVar(buffer)->shape[i] - 1, 0))); + required_bound.push_back(sym::IntSet::Interval(IntImm(GetBufferVar(buffer)->shape[i].ty(), 0), + max(GetBufferVar(buffer)->shape[i] - 1, 0))); } - ffi::Map var_dom = + ffi::Map var_dom = InverseAffineIterMap(res->indices, required_region, analyzer); - ffi::Map var_bound = + ffi::Map var_bound = InverseAffineIterMap(res->indices, required_bound, analyzer); for (const auto& kv : var_dom) { const Var& var = kv.first; @@ -627,7 +625,7 @@ std::vector CalculateBlockVarDomain( const ffi::Array& iter_vars, std::unordered_map> provided_regions, std::unordered_map> required_regions, - arith::AnalyzerObj* analyzer) { + sym::AnalyzerObj* analyzer) { int n_iters = iter_vars.size(); // Step 1. Construct the mapping from block var to their iteration domain (initialized to empty) std::unordered_map iter_doms; @@ -660,9 +658,9 @@ std::vector CalculateBlockVarDomain( for (const IterVar& iter_var : iter_vars) { BlockVarDomainInfo& info = iter_doms.at(iter_var->var.get()); if (info.bound.IsNothing()) { - info.bound = arith::IntSet::FromRange(iter_var->dom); + info.bound = sym::IntSet::FromRange(iter_var->dom); } else { - info.bound = arith::Intersect({info.bound, arith::IntSet::FromRange(iter_var->dom)}); + info.bound = sym::Intersect({info.bound, sym::IntSet::FromRange(iter_var->dom)}); } info.Simplify(analyzer); TVM_FFI_ICHECK(!info.dom.IsNothing()); @@ -719,7 +717,7 @@ void CalculateProvidedRequiredRegions( template void ComputeAtOrReverseComputeAtImpl(ScheduleState self, const StmtSRef& block_sref, const StmtSRef& loop_sref, bool preserve_unit_loops, - arith::AnalyzerObj* analyzer, bool check_only = false, + sym::AnalyzerObj* analyzer, bool check_only = false, int index = -1) { const SBlockNode* block = TVM_SREF_TO_SBLOCK(block_sref); const ForNode* loop = TVM_SREF_TO_FOR(loop_sref); @@ -797,21 +795,21 @@ void ComputeAtOrReverseComputeAtImpl(ScheduleState self, const StmtSRef& block_s void ComputeAt(ScheduleState self, const StmtSRef& block_sref, const StmtSRef& loop_sref, bool preserve_unit_loops, int index) { - arith::Analyzer analyzer; + sym::Analyzer analyzer; ComputeAtOrReverseComputeAtImpl(self, block_sref, loop_sref, preserve_unit_loops, analyzer.get(), false, index); } void ReverseComputeAt(ScheduleState self, const StmtSRef& block_sref, const StmtSRef& loop_sref, bool preserve_unit_loops, int index) { - arith::Analyzer analyzer; + sym::Analyzer analyzer; ComputeAtOrReverseComputeAtImpl(self, block_sref, loop_sref, preserve_unit_loops, analyzer.get(), false, index); } bool CanComputeAt(const ScheduleState& self, const StmtSRef& block_sref, const StmtSRef& loop_sref, bool preserve_unit_loops) { - arith::Analyzer analyzer; + sym::Analyzer analyzer; try { ComputeAtOrReverseComputeAtImpl(self, block_sref, loop_sref, preserve_unit_loops, analyzer.get(), true); @@ -823,7 +821,7 @@ bool CanComputeAt(const ScheduleState& self, const StmtSRef& block_sref, const S bool CanReverseComputeAt(const ScheduleState& self, const StmtSRef& block_sref, const StmtSRef& loop_sref, bool preserve_unit_loops) { - arith::Analyzer analyzer; + sym::Analyzer analyzer; try { ComputeAtOrReverseComputeAtImpl(self, block_sref, loop_sref, preserve_unit_loops, analyzer.get(), true); diff --git a/src/s_tir/schedule/primitive/compute_inline.cc b/src/s_tir/schedule/primitive/compute_inline.cc index 1340a49a78b6..b93e29f5b812 100644 --- a/src/s_tir/schedule/primitive/compute_inline.cc +++ b/src/s_tir/schedule/primitive/compute_inline.cc @@ -558,11 +558,11 @@ class ComputeInliner : public BaseInliner { for (const auto& iter : producer_block->iter_vars) { producer_iter_doms.Set(iter->var, iter->dom); } - arith::IterMapResult res = arith::DetectIterMap( + sym::IterMapResult res = sym::DetectIterMap( /*indices=*/inlined_store_->indices, /*input_iters=*/producer_iter_doms, /*predicate=*/true, - /*check_level=*/arith::IterMapLevel::Bijective, + /*check_level=*/sym::IterMapLevel::Bijective, /*analyzer=*/analyzer_, /*simplify_trivial_iterators=*/false); if (!res->errors.empty()) { @@ -576,7 +576,7 @@ class ComputeInliner : public BaseInliner { ffi::Array prim_idx_vars; prim_idx_vars.reserve(idx_vars_.size()); for (const Var& var : idx_vars_) prim_idx_vars.push_back(var.as_or_throw()); - auto inverse_iter_map = arith::InverseAffineIterMap(res->indices, prim_idx_vars); + auto inverse_iter_map = sym::InverseAffineIterMap(res->indices, prim_idx_vars); for (const auto& iter : producer_block->iter_vars) { if (is_const_int(iter->dom->min) && analyzer_->CanProveEqual(iter->dom->extent, 1)) { // fallback mapping for constant iters @@ -614,7 +614,7 @@ class ComputeInliner : public BaseInliner { } /*! \brief The arithmetic analyzer */ - arith::Analyzer analyzer_; + sym::Analyzer analyzer_; /*! \brief The store value for inlinement. If the producer store indices are trivial, it is wrt the producer block iter var, otherwise it is wrt to the placeholder vars of store indices. */ @@ -734,11 +734,11 @@ class ReverseComputeInliner : public BaseInliner { } } - arith::IterMapResult res = arith::DetectIterMap( + sym::IterMapResult res = sym::DetectIterMap( /*indices=*/buffer_load_indices_, /*input_iters=*/consumer_iter_doms, /*predicate=*/true, - /*check_level=*/arith::IterMapLevel::NoCheck, + /*check_level=*/sym::IterMapLevel::NoCheck, /*analyzer=*/analyzer_, /*simplify_trivial_iterators=*/false); buffer_load_iter_map_ = res->indices; @@ -858,15 +858,14 @@ class ReverseComputeInliner : public BaseInliner { * \return Whether the consumer block iter domains are covered */ bool CheckConsumerCovered() { - ffi::Map producer_iter_doms; + ffi::Map producer_iter_doms; for (const IterVar& iter_var : producer_block_->iter_vars) { - producer_iter_doms.Set(iter_var->var, arith::IntSet::FromRange(iter_var->dom)); + producer_iter_doms.Set(iter_var->var, sym::IntSet::FromRange(iter_var->dom)); } // For each block iter in the consumer block, find the corresponding expression in the producer for (const IterVar& iter : consumer_block_->iter_vars) { if (auto producer_iter = VarRemapGet(iter->var).as()) { - arith::IntSet producer_iter_range = - arith::EvalSet(producer_iter.value(), producer_iter_doms); + sym::IntSet producer_iter_range = sym::EvalSet(producer_iter.value(), producer_iter_doms); if (analyzer_->CanProve(producer_iter_range.min() > iter->dom->min) || analyzer_->CanProve(producer_iter_range.max() < iter->dom->min + iter->dom->extent - 1)) { @@ -886,7 +885,7 @@ class ReverseComputeInliner : public BaseInliner { * \param producer_indices The BufferStore indices of the producer. */ void CreateInverseMapping(const ffi::Array producer_indices) { - auto inverse_iter_map = arith::InverseAffineIterMap(buffer_load_iter_map_, producer_indices); + auto inverse_iter_map = sym::InverseAffineIterMap(buffer_load_iter_map_, producer_indices); for (const auto& pair : inverse_iter_map) { VarRemapSet(pair.first, pair.second); } @@ -955,7 +954,7 @@ class ReverseComputeInliner : public BaseInliner { /*! \brief The indices of the consumer's BufferLoad */ ffi::Array buffer_load_indices_; /*! \brief The IterMap representing the indices of the consumer's BufferLoad */ - ffi::Array buffer_load_iter_map_{nullptr}; + ffi::Array buffer_load_iter_map_{nullptr}; /*! \brief The producer block */ const SBlockNode* producer_block_{nullptr}; /* \brief The consumer block */ @@ -965,7 +964,7 @@ class ReverseComputeInliner : public BaseInliner { */ PrimExpr consumer_iter_in_bound_{nullptr}; /*! \brief The arithmetic analyzer */ - arith::Analyzer analyzer_; + sym::Analyzer analyzer_; }; void ComputeInlineImpl(ScheduleState self, const StmtSRef& producer_block_sref, @@ -1054,7 +1053,7 @@ void ReverseComputeInlineImpl(ScheduleState self, const StmtSRef& consumer_block } self->Replace(scope_root_sref, tgt_stmt, inliner->block_reuse); // Step 8. Update the cached flags - arith::Analyzer analyzer; + sym::Analyzer analyzer; SBlockInfo& block_info = self->block_info[producer_block_sref]; block_info.affine_binding = IsAffineBinding( /*realize=*/GetSBlockRealize(self, producer_block_sref), @@ -1411,7 +1410,7 @@ SBlock ReductionEpilogueFuser::CreateFusedReductionBlock( .as_or_throw(); // Simplify the expression (e.g., 0 + C[vi, vj] -> C[vi, vj]) - arith::Analyzer analyzer; + sym::Analyzer analyzer; init_epilogue = analyzer->Simplify(init_epilogue); ffi::Array init_indices = diff --git a/src/s_tir/schedule/primitive/decompose_padding.cc b/src/s_tir/schedule/primitive/decompose_padding.cc index 283a5609a350..0ad27fc00f73 100644 --- a/src/s_tir/schedule/primitive/decompose_padding.cc +++ b/src/s_tir/schedule/primitive/decompose_padding.cc @@ -74,7 +74,7 @@ class PaddingInfoAnalyzer { public: static PaddingSBlockInfo CheckAndGetPaddingInfo(IRModule mod, const SBlockRealizeNode* realize, const ffi::Map& dom_map, - arith::AnalyzerObj* analyzer) { + sym::AnalyzerObj* analyzer) { PaddingInfoAnalyzer padding_analyzer(analyzer); if (!padding_analyzer.MatchPadding(realize, dom_map)) { throw MakeScheduleError(mod, realize->block, @@ -84,7 +84,7 @@ class PaddingInfoAnalyzer { } private: - explicit PaddingInfoAnalyzer(arith::AnalyzerObj* analyzer) : analyzer_(analyzer) {} + explicit PaddingInfoAnalyzer(sym::AnalyzerObj* analyzer) : analyzer_(analyzer) {} /*! \brief Detect padding pattern and update result. */ bool MatchPadding(const SBlockRealizeNode* realize, const ffi::Map& dom_map) { @@ -154,7 +154,7 @@ class PaddingInfoAnalyzer { PrimExpr RewritePredicate(const PrimExpr& predicate) { PrimExpr res = IntImm::Bool(true); std::function update = [&res, &update](PrimExpr e) { - arith::PVar a, b; + sym::PVar a, b; if ((a && b).Match(e)) { update(a.Eval()); update(b.Eval()); @@ -177,14 +177,14 @@ class PaddingInfoAnalyzer { const PrimExpr& in_bound_predicate) { ffi::Array region; - arith::Analyzer analyzer_ref = ffi::GetRef(analyzer_); - auto res = arith::DetectIterMap(iter_values, dom_map, in_bound_predicate, - arith::IterMapLevel::Surjective, analyzer_ref); + sym::Analyzer analyzer_ref = ffi::GetRef(analyzer_); + auto res = sym::DetectIterMap(iter_values, dom_map, in_bound_predicate, + sym::IterMapLevel::Surjective, analyzer_ref); if (res->indices.empty()) { SetError("Block iters are not independent wrt padding condition"); return {}; } - for (const arith::IterSumExpr& sum : res->indices) { + for (const sym::IterSumExpr& sum : res->indices) { if (sum->args.empty()) { region.push_back(Range::FromMinExtent(sum->base, IntImm(sum->base.ty(), /* value */ 1))); } else { @@ -206,7 +206,7 @@ class PaddingInfoAnalyzer { /*! \brief current error message. */ std::string error_msg_; /*! \brief arithmetic analyzer. */ - arith::AnalyzerObj* analyzer_; + sym::AnalyzerObj* analyzer_; }; /*! \brief Create block to fill constant pad values into full region */ @@ -214,7 +214,7 @@ static std::pair CreateConstBlock(const SBlockRealizeNode* const PaddingSBlockInfo& info, const ffi::Array& loops, const Stmt& highest_pos_inclusive, - arith::AnalyzerObj* analyzer) { + sym::AnalyzerObj* analyzer) { const SBlock& block = realize->block; ffi::Array new_iter_vars; ffi::Map repl_dict; @@ -290,7 +290,7 @@ static std::pair CreateInBoundBlock(const SBlockRealizeNode const ffi::Array& loops, const Stmt& highest_pos_inclusive, - arith::AnalyzerObj* analyzer) { + sym::AnalyzerObj* analyzer) { const SBlock& block = realize->block; ffi::Array new_iter_vars; ffi::Map repl_dict; @@ -453,7 +453,7 @@ StmtSRef DecomposePaddingImpl(ScheduleState self, const StmtSRef& block_sref, const SBlockNode* block = TVM_SREF_TO_SBLOCK(block_sref); const SBlockRealizeNode* realize = GetSBlockRealize(self, block_sref).get(); ffi::Map dom_map; - arith::Analyzer analyzer; + sym::Analyzer analyzer; // Check 1. check the block is complete. StmtSRef scope_root_sref = GetScopeRoot(self, block_sref, /*require_stage_pipeline=*/false); diff --git a/src/s_tir/schedule/primitive/layout_transformation.cc b/src/s_tir/schedule/primitive/layout_transformation.cc index c6122710b257..bb56fa6bb7ce 100644 --- a/src/s_tir/schedule/primitive/layout_transformation.cc +++ b/src/s_tir/schedule/primitive/layout_transformation.cc @@ -17,13 +17,13 @@ * under the License. */ -#include #include #include #include #include #include #include +#include #include #include @@ -104,7 +104,7 @@ class TransformLayoutPlanner : public StmtExprVisitor { static TransformPlan Plan(SBlock block, BufferVar old_buffer, BufferVar new_buffer, IndexMap index_map, IndexMap inverse, PrimExpr padding_predicate, - ffi::Optional pad_value, arith::AnalyzerObj* analyzer) { + ffi::Optional pad_value, sym::AnalyzerObj* analyzer) { TVM_FFI_ICHECK(!pad_value.has_value() || pad_value.value()->final_indices.size() == 1) << "Internal error: Should be caught by ScheduleError checks prior to this point"; auto visitor = ffi::make_object(old_buffer); @@ -246,7 +246,7 @@ class TransformLayoutPlanner : public StmtExprVisitor { BufferStoreReplacer(const WriteInfo& info, const BufferVar& new_buffer, PrimExpr padding_predicate, const IndexMap& inverse, const ffi::Optional& pad_value, - ffi::Map* new_block_to_old, arith::AnalyzerObj* analyzer) + ffi::Map* new_block_to_old, sym::AnalyzerObj* analyzer) : info(info), new_buffer(new_buffer), new_indices( @@ -390,8 +390,8 @@ class TransformLayoutPlanner : public StmtExprVisitor { if (can_replace) { ffi::Array new_index_exprs = new_indices.Map([](const Var& var) { return var.as_or_throw(); }); - PrimExpr pad_value_at_index = pad_value.value()->MapIndices( - new_index_exprs, ffi::GetRef(analyzer))[0]; + PrimExpr pad_value_at_index = + pad_value.value()->MapIndices(new_index_exprs, ffi::GetRef(analyzer))[0]; store = BufferStore(new_buffer, if_then_else(padding_predicate, pad_value_at_index, op->value), new_index_exprs); @@ -470,12 +470,12 @@ class TransformLayoutPlanner : public StmtExprVisitor { const ffi::Optional& pad_value; ffi::Map& new_block_to_old; bool all_stores_replaced{true}; - arith::AnalyzerObj* analyzer; + sym::AnalyzerObj* analyzer; }; TransformPlan Finalize(BufferVar new_buffer, IndexMap index_map, IndexMap inverse, PrimExpr padding_predicate, ffi::Optional pad_value, - arith::AnalyzerObj* analyzer) const { + sym::AnalyzerObj* analyzer) const { if (auto prologue_plan = FinalizeProloguePlan(new_buffer, index_map, inverse, padding_predicate, pad_value, analyzer); prologue_plan.has_value()) { @@ -496,7 +496,7 @@ class TransformLayoutPlanner : public StmtExprVisitor { std::optional FinalizeProloguePlan(BufferVar new_buffer, IndexMap index_map, IndexMap inverse, PrimExpr padding_predicate, ffi::Optional pad_value, - arith::AnalyzerObj* analyzer) const { + sym::AnalyzerObj* analyzer) const { if (write_info_.size() || is_zero(padding_predicate) || !pad_value.has_value()) { return std::nullopt; } @@ -526,7 +526,7 @@ class TransformLayoutPlanner : public StmtExprVisitor { .as_or_throw(); PrimExpr pad_value_at_index = - pad_value.value()->MapIndices(indices, ffi::GetRef(analyzer))[0]; + pad_value.value()->MapIndices(indices, ffi::GetRef(analyzer))[0]; PrimExpr expr = (!padding_predicate) || (BufferLoad(new_buffer, indices) == pad_value_at_index); Stmt stmt = Evaluate(Call(PrimType::Bool(), tirx::builtin::assume(), {expr}).as_or_throw()); @@ -550,7 +550,7 @@ class TransformLayoutPlanner : public StmtExprVisitor { IndexMap inverse, PrimExpr padding_predicate, ffi::Optional pad_value, - arith::AnalyzerObj* analyzer) const { + sym::AnalyzerObj* analyzer) const { if (write_info_.empty() || is_zero(padding_predicate) || !pad_value.has_value()) { return std::nullopt; } @@ -601,7 +601,7 @@ class TransformLayoutPlanner : public StmtExprVisitor { std::optional FinalizeEpiloguePlan(BufferVar new_buffer, IndexMap index_map, IndexMap inverse, PrimExpr padding_predicate, ffi::Optional pad_value, - arith::AnalyzerObj* analyzer) const { + sym::AnalyzerObj* analyzer) const { if (write_info_.empty() || is_zero(padding_predicate) || !pad_value.has_value()) { return std::nullopt; } @@ -621,7 +621,7 @@ class TransformLayoutPlanner : public StmtExprVisitor { } PrimExpr pad_value_at_index = - pad_value.value()->MapIndices(indices, ffi::GetRef(analyzer))[0]; + pad_value.value()->MapIndices(indices, ffi::GetRef(analyzer))[0]; Stmt stmt = BufferStore(new_buffer, pad_value_at_index, indices); std::stringstream block_name; @@ -822,7 +822,7 @@ class TransformLayoutRewriter : public s_tir::IRMutatorWithAnalyzer { const SBlock& scope_stmt, const BufferVar& old_buffer, const BufferVar& new_buffer, const IndexMap& index_map, const ffi::Optional& opt_inverse, const PrimExpr& padding_predicate, const ffi::Optional& pad_value) { - arith::Analyzer analyzer; + sym::Analyzer analyzer; auto plan = pad_value.has_value() ? TransformLayoutPlanner::Plan(scope_stmt, old_buffer, new_buffer, index_map, opt_inverse.value(), padding_predicate, @@ -846,7 +846,7 @@ class TransformLayoutRewriter : public s_tir::IRMutatorWithAnalyzer { TransformLayoutRewriter(const BufferVar& old_buffer, const BufferVar& new_buffer, const IndexMap& index_map, const TransformLayoutPlanner::TransformPlan& plan, - const arith::Analyzer& analyzer) + const sym::Analyzer& analyzer) : IRMutatorWithAnalyzer(analyzer), old_buffer_(old_buffer), new_buffer_(new_buffer), @@ -1004,7 +1004,7 @@ class TransformLayoutRewriter : public s_tir::IRMutatorWithAnalyzer { const TransformLayoutPlanner::TransformPlan& plan_; ffi::Map buffer_data_to_buffer_; ffi::Map new_block_to_old_; - arith::Analyzer index_simplifier_; + sym::Analyzer index_simplifier_; }; class BufferIsSubregionError : public ScheduleErrorContextObj { @@ -1168,7 +1168,7 @@ class TransformationIntroducesPaddingError : public ScheduleErrorContextObj { } ffi::String DetailRenderTemplate() const final { - arith::Analyzer analyzer; + sym::Analyzer analyzer; auto new_shape = index_map_->MapShape(buffer_->shape, analyzer); std::ostringstream os; os << "The transformation " << index_map_ << " applied on buffer " << buffer_.name() @@ -1243,7 +1243,7 @@ IndexMap LegalizeIndexMapDType(const IndexMap& index_map, const ffi::Array& pad_value, bool assume_injective_transform) { - arith::Analyzer analyzer; + sym::Analyzer analyzer; AddShapeVarBounds(self, block_sref.get(), analyzer.get()); // Step 1: Input handling and error checking const SBlockNode* block_ptr = TVM_SREF_TO_SBLOCK(block_sref); @@ -1448,7 +1448,7 @@ void TransformBlockLayout(ScheduleState self, const StmtSRef& block_sref, const IndexMap& index_map) { const SBlockNode* block_ptr = TVM_SREF_TO_SBLOCK(block_sref); const SBlock& block = ffi::GetRef(block_ptr); - arith::Analyzer analyzer; + sym::Analyzer analyzer; AddShapeVarBounds(self, block_sref.get(), analyzer.get()); // Step 1: Collect outer loops and loop vars diff --git a/src/s_tir/schedule/primitive/loop_transformation.cc b/src/s_tir/schedule/primitive/loop_transformation.cc index f2987eae4706..57e98f5bf3cd 100644 --- a/src/s_tir/schedule/primitive/loop_transformation.cc +++ b/src/s_tir/schedule/primitive/loop_transformation.cc @@ -143,12 +143,12 @@ class IterMapSimplifyBlockBinding : public StmtExprMutator { return realize; } ffi::Array v = - arith::IterMapSimplify(/*indices=*/op->iter_values, - /*input_iters=*/loop_var2extent_, - /*input_pred=*/op->predicate, - /*check_level=*/arith::IterMapLevel::Surjective, - /*analyzer=*/analzyer_, - /*simplify_trivial_iterators=*/!preserve_unit_iters_); + sym::IterMapSimplify(/*indices=*/op->iter_values, + /*input_iters=*/loop_var2extent_, + /*input_pred=*/op->predicate, + /*check_level=*/sym::IterMapLevel::Surjective, + /*analyzer=*/analzyer_, + /*simplify_trivial_iterators=*/!preserve_unit_iters_); if (v.same_as(op->iter_values)) { return ffi::Unchanged(); } else { @@ -168,7 +168,7 @@ class IterMapSimplifyBlockBinding : public StmtExprMutator { /*! \brief The range of loops */ ffi::Map loop_var2extent_; /*! \brief Internal analyzer */ - arith::Analyzer analzyer_; + sym::Analyzer analzyer_; /*! \brief Whether or not to simplify unit iterators */ bool preserve_unit_iters_; }; @@ -441,7 +441,7 @@ ffi::Array Split(ScheduleState self, const StmtSRef& loop_sref, throw MakeScheduleError(self->mod, ffi::GetRef(loop)); } // Currently, loops not starting with 0 are not supported - arith::Analyzer analyzer; + sym::Analyzer analyzer; CheckLoopStartsWithZero(self, loop_sref, analyzer.get()); // Find the most common dtype @@ -472,8 +472,7 @@ ffi::Array Split(ScheduleState self, const StmtSRef& loop_sref, .ValueOrUnchanged(std::move(new_stmt)); // Step 3. Update predicate to guard the loop PrimExpr predicate = substitute_value < loop->extent; - if (!disable_predication && - !analyzer->CanProve(predicate, arith::ProofStrength::kSymbolicBound)) { + if (!disable_predication && !analyzer->CanProve(predicate, sym::ProofStrength::kSymbolicBound)) { new_stmt = ffi::make_object(/*predicate=*/predicate) ->Mutate(new_stmt, InplaceMode::kAllow) .ValueOrUnchanged(std::move(new_stmt)); @@ -717,7 +716,7 @@ ffi::Array LoopPartition(ScheduleState self, const StmtSRef& loop_sref throw MakeScheduleError(self->mod, ffi::GetRef(loop)); } - arith::Analyzer analyzer; + sym::Analyzer analyzer; // Find the most common dtype PrimType dtype = PrimType::Int(32); { @@ -882,7 +881,7 @@ StmtSRef Merge(ScheduleState self, const ffi::Array& loop_srefs) { // - The total repeat number has not changed for each direct child block. // - The execution order has not changed. (The block executes with the same // args and the same order with before.) - arith::Analyzer analyzer; + sym::Analyzer analyzer; StmtSRef scope_root_sref; StmtSRef lca = GetSRefLowestCommonAncestor(loop_srefs); std::vector> lca_nest_loops; @@ -958,7 +957,7 @@ StmtSRef Fuse(ScheduleState self, const ffi::Array& loop_srefs, loops.reserve(loop_srefs.size()); StmtSRef outer_loop_sref{nullptr}; const ForNode* outer_loop = nullptr; - arith::Analyzer analyzer; + sym::Analyzer analyzer; std::unordered_set outer_loop_vars; // Step 1. check correctness for (const StmtSRef& sref : loop_srefs) { diff --git a/src/s_tir/schedule/primitive/pad_einsum.cc b/src/s_tir/schedule/primitive/pad_einsum.cc index 7bfd67177913..c7758f334895 100644 --- a/src/s_tir/schedule/primitive/pad_einsum.cc +++ b/src/s_tir/schedule/primitive/pad_einsum.cc @@ -161,7 +161,7 @@ struct BufferPadding { return result; } - Stmt MakeCopyBlock(bool is_read, ffi::Array* blocks, arith::AnalyzerObj* analyzer) { + Stmt MakeCopyBlock(bool is_read, ffi::Array* blocks, sym::AnalyzerObj* analyzer) { ffi::Array loop_vars; ffi::Array loop_doms; ffi::Array iter_vars; @@ -392,7 +392,7 @@ class PadEinsumBufferReplacer : public StmtExprMutator { }; void PadEinsum(ScheduleState self, const StmtSRef& block_sref, const ffi::Array& padding) { - arith::Analyzer analyzer; + sym::Analyzer analyzer; // Step 1: Input checking and error handling const SBlockNode* block = TVM_SREF_TO_SBLOCK(block_sref); SBlockRealize realize = GetSBlockRealize(self, block_sref); diff --git a/src/s_tir/schedule/primitive/read_write_at.cc b/src/s_tir/schedule/primitive/read_write_at.cc index 92f8f35e94b0..7c8f2f134bd0 100644 --- a/src/s_tir/schedule/primitive/read_write_at.cc +++ b/src/s_tir/schedule/primitive/read_write_at.cc @@ -42,9 +42,9 @@ bool HasBuffer(const ffi::Array& buffer_regions, const BufferVar& } void RelaxBufferRegions(const ffi::Array& buffer_regions, - const BufferVar& buffer, // - const ffi::Map& var_dom, // - const ffi::Map& bindings, // + const BufferVar& buffer, // + const ffi::Map& var_dom, // + const ffi::Map& bindings, // std::vector* relaxed_regions) { auto f_substitute = [&bindings](const Var& var) -> ffi::Expected> { if (auto repl = bindings.Get(var)) return ffi::Any(*std::move(repl)); @@ -61,7 +61,7 @@ void RelaxBufferRegions(const ffi::Array& buffer_regions, .as_or_throw(); return Range::FromMinExtent(min, extent); }); - ffi::Array relaxed_region = arith::EvalSet(mapped_region, var_dom); + ffi::Array relaxed_region = sym::EvalSet(mapped_region, var_dom); relaxed_regions->push_back({relaxed_region.begin(), relaxed_region.end()}); } } @@ -227,7 +227,7 @@ struct ReadWriteAtImpl { /*buffer_regions=*/is_read ? block->reads : block->writes, /*buffer=*/src_, /*var_dom=*/ - arith::AsIntSet(LoopDomainOfSRefTreePath( + sym::AsIntSet(LoopDomainOfSRefTreePath( /*low_inclusive=*/ffi::GetRef(self_->stmt2ref.at(block)->parent), /*high_exclusive=*/loop_sref_, /*extra_relax_scope=*/scope)), @@ -270,7 +270,7 @@ struct ReadWriteAtImpl { ffi::Array domain; domain.reserve(ndim); for (int i = 0; i < ndim; ++i) { - const arith::IntSet& int_set = relaxed[i]; + const sym::IntSet& int_set = relaxed[i]; PrimExpr min = analyzer_->Simplify(int_set.min()); PrimExpr extent = analyzer_->Simplify(int_set.max() + 1 - min); domain.push_back(Range::FromMinExtent(min, extent)); @@ -362,7 +362,7 @@ struct ReadWriteAtImpl { dst_(dst), annotations_(annotations), block_sref_reuse_(), - analyzer_(arith::Analyzer()) { + analyzer_(sym::Analyzer()) { loop_ = TVM_SREF_TO_FOR(loop_sref); } @@ -373,7 +373,7 @@ struct ReadWriteAtImpl { const BufferVar& dst_; ffi::Map annotations_; ffi::Map block_sref_reuse_; - arith::Analyzer analyzer_; + sym::Analyzer analyzer_; }; StmtSRef ReadAt(ScheduleState self, const StmtSRef& loop_sref, const StmtSRef& block_sref, diff --git a/src/s_tir/schedule/primitive/rolling_buffer.cc b/src/s_tir/schedule/primitive/rolling_buffer.cc index 875a34b5f7a2..49aa13771f74 100644 --- a/src/s_tir/schedule/primitive/rolling_buffer.cc +++ b/src/s_tir/schedule/primitive/rolling_buffer.cc @@ -44,7 +44,7 @@ struct RollingBufferInfo { }; TensorRegion GetRelaxedBufferRegion(const SBlockRealize& realize, const TensorRegion& buffer_region, - const ffi::Map& dom_map) { + const ffi::Map& dom_map) { ffi::Map bindings = GetBindings(realize); auto f_substitute = [&bindings](const Var& var) -> ffi::Expected> { if (auto repl = bindings.Get(var)) return ffi::Any(*std::move(repl)); @@ -57,7 +57,7 @@ TensorRegion GetRelaxedBufferRegion(const SBlockRealize& realize, const TensorRe .as_or_throw(); return Range::FromMinExtent(min, extent); }); - ffi::Array relaxed_intsets = arith::EvalSet(mapped_region, dom_map); + ffi::Array relaxed_intsets = sym::EvalSet(mapped_region, dom_map); Region relaxed_region; relaxed_region.reserve(relaxed_intsets.size()); for (size_t i = 0; i < relaxed_intsets.size(); ++i) { @@ -187,8 +187,8 @@ class RollingBufferInfoCollector { std::vector> bound_iter_vars; std::vector bound_overlaps; - arith::PVar p_var; - arith::PVar p_stride, p_divisor; + sym::PVar p_var; + sym::PVar p_stride, p_divisor; for (auto bound : region) { auto stride = 0; auto divisor = 1; @@ -379,10 +379,10 @@ class RollingBufferRewriter : public StmtExprMutator { auto iter_var = info_->axis_iter_vars[i]; if (iter_var && info_->axis_overlaps[i] > 0) { Var var = iter_var.value(); - const ffi::Map dmap = { - std::make_pair(var, arith::IntSet::Interval(0, 0))}; + const ffi::Map dmap = { + std::make_pair(var, sym::IntSet::Interval(0, 0))}; auto iter_value = realize->iter_values[i]; - arith::Analyzer analyzer; + sym::Analyzer analyzer; auto term_2 = analyzer->int_set(iter_value, dmap).min(); condition = analyzer->Simplify(And(condition, Or(LT(var.as_or_throw(), 1), GE(term_2, info_->axis_overlaps[i])))); @@ -440,7 +440,7 @@ void RollingBuffer(ScheduleState self, const StmtSRef& block_sref, int write_buf * indices to circularize the buffer along the rolling dimension. * - Append block predicate to avoid recomputing overlapping elements. */ - ffi::Map dom_map; + ffi::Map dom_map; const SBlockRealize& realize = GetSBlockRealize(self, block_sref); const SBlock& block = realize->block; @@ -470,7 +470,7 @@ void RollingBuffer(ScheduleState self, const StmtSRef& block_sref, int write_buf } For cur_loop = ffi::GetRef(stmt->StmtAs()); Range range = Range::FromMinExtent(cur_loop->min, cur_loop->extent); - dom_map.Set(cur_loop->loop_var, arith::IntSet::FromRange(range)); + dom_map.Set(cur_loop->loop_var, sym::IntSet::FromRange(range)); } TensorRegion relaxed_region = GetRelaxedBufferRegion(realize, buffer_region, dom_map); diff --git a/src/s_tir/schedule/state.cc b/src/s_tir/schedule/state.cc index f07a88c7ba50..bc8fd8d01c18 100644 --- a/src/s_tir/schedule/state.cc +++ b/src/s_tir/schedule/state.cc @@ -16,11 +16,11 @@ * specific language governing permissions and limitations * under the License. */ -#include #include #include #include #include +#include #include "./utils.h" namespace tvm { @@ -43,17 +43,17 @@ using SMap = std::unordered_map; * \param dom_high_exclusive The highest node in the sref tree path * \return An n-dimensional integer set */ -ffi::Array AnalyzeRegionUpperBound(const TensorRegion& region, // - const PrimExpr& predicate, // - const StmtSRef& dom_low_inclusive, // - const StmtSRef& dom_high_exclusive, // - arith::AnalyzerObj* analyzer) { +ffi::Array AnalyzeRegionUpperBound(const TensorRegion& region, // + const PrimExpr& predicate, // + const StmtSRef& dom_low_inclusive, // + const StmtSRef& dom_high_exclusive, // + sym::AnalyzerObj* analyzer) { ffi::Map var_dom = LoopDomainOfSRefTreePath( /*low_inclusive=*/dom_low_inclusive, /*high_exclusive=*/dom_high_exclusive, /*extra_relax_scope=*/ runtime::StorageScope::Create(region->source.as_or_throw().scope())); - arith::Analyzer analyzer_ref = ffi::GetRef(analyzer); + sym::Analyzer analyzer_ref = ffi::GetRef(analyzer); return EstimateRegionUpperBound( /*region=*/region->region, /*var_dom=*/var_dom, @@ -70,25 +70,25 @@ ffi::Array AnalyzeRegionUpperBound(const TensorRegion& region, * \param analyzer The analyzer * \return An n-dimensional integer set */ -ffi::Array AnalyzeRegionLowerBound(const TensorRegion& region, // - const PrimExpr& predicate, // - const StmtSRef& dom_low_inclusive, // - const StmtSRef& dom_high_exclusive, // - arith::AnalyzerObj* analyzer) { +ffi::Array AnalyzeRegionLowerBound(const TensorRegion& region, // + const PrimExpr& predicate, // + const StmtSRef& dom_low_inclusive, // + const StmtSRef& dom_high_exclusive, // + sym::AnalyzerObj* analyzer) { ffi::Map var_dom = LoopDomainOfSRefTreePath( /*low_inclusive=*/dom_low_inclusive, /*high_exclusive=*/dom_high_exclusive, /*extra_relax_scope=*/ runtime::StorageScope::Create(region->source.as_or_throw().scope())); - arith::Analyzer analyzer_ref = ffi::GetRef(analyzer); - if (ffi::Optional> result = EstimateRegionLowerBound( + sym::Analyzer analyzer_ref = ffi::GetRef(analyzer); + if (ffi::Optional> result = EstimateRegionLowerBound( /*region=*/region->region, /*var_dom=*/var_dom, /*predicate=*/predicate, /*analyzer=*/analyzer_ref)) { return result.value(); } - return ffi::Array(region->source.as_or_throw()->shape.size(), - arith::IntSet::Nothing()); + return ffi::Array(region->source.as_or_throw()->shape.size(), + sym::IntSet::Nothing()); } /*! @@ -100,33 +100,33 @@ ffi::Array AnalyzeRegionLowerBound(const TensorRegion& region, * \return A boolean indicating if the produced region could cover the consumed region */ bool ProducerCoversConsumer(const ffi::Array& buffer_shape, - const ffi::Array& produced_region, - const ffi::Array& consumed_region, - arith::AnalyzerObj* analyzer) { + const ffi::Array& produced_region, + const ffi::Array& consumed_region, + sym::AnalyzerObj* analyzer) { TVM_FFI_ICHECK_EQ(buffer_shape.size(), consumed_region.size()); TVM_FFI_ICHECK_EQ(produced_region.size(), consumed_region.size()); int ndim = produced_region.size(); for (int i = 0; i < ndim; ++i) { - arith::IntSet buffer_size = arith::IntSet::FromMinExtent(0, buffer_shape[i]); + sym::IntSet buffer_size = sym::IntSet::FromMinExtent(0, buffer_shape[i]); if (produced_region[i].IsNothing()) { return false; } if (consumed_region[i].IsNothing()) { continue; } - arith::IntSet produced = - arith::IntSet::Interval(analyzer->canonical_simplify(produced_region[i].min()), - analyzer->canonical_simplify(produced_region[i].max())); - arith::IntSet consumed = - arith::IntSet::Interval(analyzer->canonical_simplify(consumed_region[i].min()), - analyzer->canonical_simplify(consumed_region[i].max())); - produced = arith::Intersect({produced, buffer_size}); - consumed = arith::Intersect({consumed, buffer_size}); - - produced = arith::IntSet::Interval(analyzer->Simplify(produced.min()), - analyzer->Simplify(produced.max())); - consumed = arith::IntSet::Interval(analyzer->Simplify(consumed.min()), - analyzer->Simplify(consumed.max())); + sym::IntSet produced = + sym::IntSet::Interval(analyzer->canonical_simplify(produced_region[i].min()), + analyzer->canonical_simplify(produced_region[i].max())); + sym::IntSet consumed = + sym::IntSet::Interval(analyzer->canonical_simplify(consumed_region[i].min()), + analyzer->canonical_simplify(consumed_region[i].max())); + produced = sym::Intersect({produced, buffer_size}); + consumed = sym::Intersect({consumed, buffer_size}); + + produced = sym::IntSet::Interval(analyzer->Simplify(produced.min()), + analyzer->Simplify(produced.max())); + consumed = sym::IntSet::Interval(analyzer->Simplify(consumed.min()), + analyzer->Simplify(consumed.max())); if (!analyzer->CanProve((analyzer->canonical_simplify(produced.min() - consumed.min()) <= 0) && (analyzer->canonical_simplify(consumed.max() - produced.max()) <= 0))) { @@ -306,7 +306,7 @@ class SBlockInfoCollector : public StmtExprVisitor { continue; } // For each buffer, record the regions generated under this loop - std::unordered_map>, ffi::ObjectPtrHash, + std::unordered_map>, ffi::ObjectPtrHash, ffi::ObjectPtrEqual> touched_regions; // Step 2.3.1. Find all the regions read by the consumer that we care about @@ -323,7 +323,7 @@ class SBlockInfoCollector : public StmtExprVisitor { auto it = touched_regions.find(buffer); // Skip the regions that is not read by the consumer if (it != touched_regions.end()) { - std::vector>& touched_region = it->second; + std::vector>& touched_region = it->second; // The analysis here is trying to be conservation to rule out false positive cases, // and to make sure region cover property must be satisfied once the flag is on // Therefore, we use lower-bound analysis for producers and upper-bound analysis for @@ -342,12 +342,11 @@ class SBlockInfoCollector : public StmtExprVisitor { StmtSRef parent_sref = ffi::GetRef(consumer_block_sref->parent); for (const TensorRegion& region : block_reads_unbound.at(consumer_block_sref.get())) { BufferVar buffer = region->source.as_or_throw(); - const std::vector>& touched_region = - touched_regions.at(buffer); + const std::vector>& touched_region = touched_regions.at(buffer); if (!touched_region.empty()) { - ffi::Array produced_region = - arith::UnionRegionLowerBound({touched_region.begin(), touched_region.end()}); - ffi::Array consumed_region = AnalyzeRegionUpperBound( + ffi::Array produced_region = + sym::UnionRegionLowerBound({touched_region.begin(), touched_region.end()}); + ffi::Array consumed_region = AnalyzeRegionUpperBound( /*region=*/region, /*predicate=*/consumer_realize->predicate, /*dom_low_inclusive=*/parent_sref, @@ -408,7 +407,7 @@ class SBlockInfoCollector : public StmtExprVisitor { /*! \brief The stack frames of blocks in the DFS visit. */ std::vector> block_frames_; /*! \brief The auxiliary analyzer */ - arith::Analyzer analyzer_; + sym::Analyzer analyzer_; }; /**************** Constructor ****************/ diff --git a/src/s_tir/schedule/traced_schedule.cc b/src/s_tir/schedule/traced_schedule.cc index 8cea5fc2bb50..e089d658c7ac 100644 --- a/src/s_tir/schedule/traced_schedule.cc +++ b/src/s_tir/schedule/traced_schedule.cc @@ -30,7 +30,7 @@ Schedule Schedule::Traced(IRModule mod, LinearCongruentialEngine::TRandState see n->state_ = ScheduleState(mod, debug_mask, enable_check); n->error_render_level_ = error_render_level; n->symbol_table_ = {}; - n->analyzer_ = arith::Analyzer(); + n->analyzer_ = sym::Analyzer(); n->trace_ = Trace(); n->Seed(seed); GlobalVar gv; @@ -47,7 +47,7 @@ Schedule TracedScheduleNode::Copy() { n->error_render_level_ = this->error_render_level_; ConcreteScheduleNode::Copy(&n->state_, &n->symbol_table_); n->func_working_on_ = this->func_working_on_; - n->analyzer_ = arith::Analyzer(); // new analyzer needed because it is stateful + n->analyzer_ = sym::Analyzer(); // new analyzer needed because it is stateful n->rand_state_ = ForkSeed(); n->trace_ = Trace(this->trace_->insts, this->trace_->decisions); return Schedule(std::move(n)); diff --git a/src/s_tir/schedule/transform.cc b/src/s_tir/schedule/transform.cc index bb0e1565c34a..6f7e27af083e 100644 --- a/src/s_tir/schedule/transform.cc +++ b/src/s_tir/schedule/transform.cc @@ -389,7 +389,7 @@ ffi::Optional TileWithTensorIntrin(const s_tir::Schedule& sch, } } // Split the loops - arith::Analyzer analyzer; + sym::Analyzer analyzer; std::unordered_set inner_loops; std::vector reorder_suffix; reorder_suffix.resize(info->loop_map.size()); diff --git a/src/s_tir/schedule/transform.h b/src/s_tir/schedule/transform.h index 5cf61192ce08..4893e749c209 100644 --- a/src/s_tir/schedule/transform.h +++ b/src/s_tir/schedule/transform.h @@ -222,12 +222,12 @@ class BlockBufferAccessSimplifier : public s_tir::IRMutatorWithAnalyzer { * \param analyzer The arithmetic analyzer * \return The simplified statement */ - static Stmt Simplify(const Stmt& stmt, const arith::Analyzer& analyzer) { + static Stmt Simplify(const Stmt& stmt, const sym::Analyzer& analyzer) { auto simplifier = ffi::make_object(analyzer); return simplifier->Mutate(stmt).ValueOrUnchanged(stmt); } - explicit BlockBufferAccessSimplifier(const arith::Analyzer& analyzer) + explicit BlockBufferAccessSimplifier(const sym::Analyzer& analyzer) : IRMutatorWithAnalyzer(analyzer) {} private: diff --git a/src/s_tir/schedule/utils.h b/src/s_tir/schedule/utils.h index 9ac1aaec06dc..cb6704c8a8e1 100644 --- a/src/s_tir/schedule/utils.h +++ b/src/s_tir/schedule/utils.h @@ -19,9 +19,6 @@ #ifndef TVM_S_TIR_SCHEDULE_UTILS_H_ #define TVM_S_TIR_SCHEDULE_UTILS_H_ -#include -#include -#include #include #include #include @@ -34,6 +31,9 @@ #include #include #include +#include +#include +#include #include #include #include @@ -43,9 +43,9 @@ #include #include -#include "../../arith/pattern_match.h" #include "../../ir/attr_registry.h" #include "../../runtime/thread_storage_scope.h" +#include "../../sym/pattern_match.h" #include "../support/array_utils.h" #include "../support/nd_int_set.h" #include "./analysis.h" @@ -216,8 +216,8 @@ inline ffi::Optional AnalyzeVarWithShift(const PrimExpr& expr, *constant = std::nullopt; return static_cast(var.value()); } - arith::PVar var; - arith::PVar shift; + sym::PVar var; + sym::PVar shift; // match: "var + shift" if ((var + shift).Match(expr) || (shift + var).Match(expr)) { *constant = shift.Eval(); diff --git a/src/s_tir/stmt.cc b/src/s_tir/stmt.cc index 98753c031b2b..45cbd47b0be7 100644 --- a/src/s_tir/stmt.cc +++ b/src/s_tir/stmt.cc @@ -21,12 +21,12 @@ * \file tvm/s_tir/stmt.cc * \brief Schedulable block definitions and structural traversal. */ -#include #include #include #include #include #include +#include namespace tvm { namespace s_tir { @@ -273,7 +273,7 @@ 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; + sym::Analyzer analyzer; // Check scope and dtype TVM_FFI_ICHECK_EQ(buffer.scope(), source_buffer.scope()) << "MatchBuffer " << buffer << " scope mismatch:" << buffer.scope() << " vs. " diff --git a/src/s_tir/support/nd_int_set.h b/src/s_tir/support/nd_int_set.h index 0422e333670a..2d4bfb612d44 100644 --- a/src/s_tir/support/nd_int_set.h +++ b/src/s_tir/support/nd_int_set.h @@ -19,8 +19,8 @@ #ifndef TVM_S_TIR_SUPPORT_ND_INT_SET_H_ #define TVM_S_TIR_SUPPORT_ND_INT_SET_H_ -#include #include +#include #include #include @@ -29,7 +29,7 @@ namespace tvm { namespace support { /*! \brief An N-dimensional integer set representing a rectangle region */ -using NDIntSet = std::vector; +using NDIntSet = std::vector; /*! * \brief Construct an N-dimensional integer set representing a region. @@ -40,7 +40,7 @@ inline NDIntSet NDIntSetFromRegion(const tirx::Region& region) { NDIntSet result; result.reserve(region.size()); for (const Range& range : region) { - result.push_back(arith::IntSet::FromRange(range)); + result.push_back(sym::IntSet::FromRange(range)); } return result; } @@ -55,7 +55,7 @@ inline NDIntSet NDIntSetFromShape(const ffi::Array& shape) { NDIntSet result; result.reserve(shape.size()); for (const PrimExpr& extent : shape) { - result.push_back(arith::IntSet::FromMinExtent(zero, extent)); + result.push_back(sym::IntSet::FromMinExtent(zero, extent)); } return result; } @@ -69,7 +69,7 @@ inline NDIntSet NDIntSetFromPoint(const ffi::Array& indices) { NDIntSet result; result.reserve(indices.size()); for (const PrimExpr& index : indices) { - result.push_back(arith::IntSet::SinglePoint(index)); + result.push_back(sym::IntSet::SinglePoint(index)); } return result; } @@ -84,8 +84,8 @@ inline void NDIntSetUnionWith(NDIntSet* lhs, const NDIntSet& rhs) { TVM_FFI_ICHECK_EQ(lhs->size(), rhs.size()); int ndim = rhs.size(); for (int i = 0; i < ndim; ++i) { - arith::IntSet& int_set = lhs->at(i); - int_set = arith::Union({int_set, rhs.at(i)}); + sym::IntSet& int_set = lhs->at(i); + int_set = sym::Union({int_set, rhs.at(i)}); } } @@ -106,12 +106,12 @@ inline NDIntSet NDIntSetUnion(const std::vector& nd_int_sets) { } NDIntSet result; result.reserve(ndim); - ffi::Array int_sets(n, arith::IntSet{nullptr}); + ffi::Array int_sets(n, sym::IntSet{nullptr}); for (int dim = 0; dim < ndim; ++dim) { for (int i = 0; i < n; ++i) { int_sets.Set(i, nd_int_sets[i][dim]); } - result.push_back(arith::Union(int_sets)); + result.push_back(sym::Union(int_sets)); } return result; } @@ -122,7 +122,7 @@ inline NDIntSet NDIntSetUnion(const std::vector& nd_int_sets) { * \return The constructed set. */ inline NDIntSet NDIntSetEmpty(int ndim) { - return std::vector(ndim, arith::IntSet::Nothing()); + return std::vector(ndim, sym::IntSet::Nothing()); } /*! @@ -133,12 +133,11 @@ inline NDIntSet NDIntSetEmpty(int ndim) { * integer set. * \sa EvalSet */ -inline NDIntSet NDIntSetEval( - const NDIntSet& nd_int_set, - const std::unordered_map& dom_map) { +inline NDIntSet NDIntSetEval(const NDIntSet& nd_int_set, + const std::unordered_map& dom_map) { NDIntSet ret; ret.reserve(nd_int_set.size()); - for (const arith::IntSet& s : nd_int_set) { + for (const sym::IntSet& s : nd_int_set) { ret.push_back(EvalSet(s, dom_map)); } return ret; diff --git a/src/s_tir/transform/bound_checker.cc b/src/s_tir/transform/bound_checker.cc index 00d8b1f8b1d0..b43d0b1750fd 100644 --- a/src/s_tir/transform/bound_checker.cc +++ b/src/s_tir/transform/bound_checker.cc @@ -22,7 +22,6 @@ */ // Instrument checkers for out of the bounds access. -#include #include #include #include @@ -31,6 +30,7 @@ #include #include #include +#include #include #include @@ -247,7 +247,7 @@ class BoundChecker : public StmtExprMutator { // Hashtable which maps buffer_var to shape. std::unordered_map> mem_to_shape_; // internal analyzer - arith::Analyzer analyzer_; + sym::Analyzer analyzer_; }; Stmt InstrumentBoundCheckers(Stmt stmt) { diff --git a/src/s_tir/transform/canonicalize_loop.cc b/src/s_tir/transform/canonicalize_loop.cc index 101acfd07f70..f7f7bdb8471c 100644 --- a/src/s_tir/transform/canonicalize_loop.cc +++ b/src/s_tir/transform/canonicalize_loop.cc @@ -21,12 +21,12 @@ * \file s_tir/transform/canonicalize_loop.cc * \brief Canonicalize all loops to start from zero and step one. */ -#include #include #include #include #include #include +#include #include #include @@ -82,7 +82,7 @@ class LoopCanonicalizer : public StmtExprMutator { } } - arith::Analyzer analyzer_; + sym::Analyzer analyzer_; }; namespace transform { diff --git a/src/s_tir/transform/compact_buffer_region.cc b/src/s_tir/transform/compact_buffer_region.cc index 4d3535fd4f56..de7c2d0623cb 100644 --- a/src/s_tir/transform/compact_buffer_region.cc +++ b/src/s_tir/transform/compact_buffer_region.cc @@ -22,13 +22,13 @@ * \brief Compact the buffer size into its exact need. */ -#include #include #include #include #include #include #include +#include #include #include @@ -49,15 +49,15 @@ using support::NDIntSet; /*! \brief a more constrained bound estimate for n-dimentional int set */ NDIntSet NDIntSetEval(Region region, PrimExpr predicate, - const std::unordered_map& dom_map, - arith::AnalyzerObj* analyzer) { + const std::unordered_map& dom_map, + sym::AnalyzerObj* analyzer) { std::unordered_map var_dom; for (const auto& it : dom_map) { var_dom[ffi::GetRef(it.first)] = it.second.CoverRange(Range::FromMinExtent(0, 0)); } - arith::Analyzer analyzer_ref = ffi::GetRef(analyzer); - ffi::Optional> eval_res = - arith::EstimateRegionUpperBound(region, var_dom, predicate, analyzer_ref); + sym::Analyzer analyzer_ref = ffi::GetRef(analyzer); + ffi::Optional> eval_res = + sym::EstimateRegionUpperBound(region, var_dom, predicate, analyzer_ref); if (eval_res.has_value()) { return NDIntSet(eval_res.value().begin(), eval_res.value().end()); @@ -192,7 +192,7 @@ class BufferAccessRegionCollector : public StmtExprVisitor { : IterVar(Range(), op->loop_var, IterVarType::kDataPar); ancestor_iters_.push_back(iter); dom_analyzer_->Bind(op->loop_var, loop_range); - dom_map_.emplace(op->loop_var.get(), arith::IntSet::FromRange(loop_range)); + dom_map_.emplace(op->loop_var.get(), sym::IntSet::FromRange(loop_range)); size_t n_pending_before = pending_flat_alloc_buffers_.size(); TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(op)); // Compact flat AllocBuffers defined inside this For scope @@ -204,21 +204,21 @@ class BufferAccessRegionCollector : public StmtExprVisitor { ffi::Optional Visit_(const BindNode* op) final { TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit(op->value)); - if (auto value = op->value.as(); value && arith::IsIndexTypedExpr(value.value())) { + if (auto value = op->value.as(); value && sym::IsIndexTypedExpr(value.value())) { dom_analyzer_->Bind(op->var, value.value()); - dom_map_.emplace(op->var.get(), arith::IntSet::SinglePoint(value.value())); + dom_map_.emplace(op->var.get(), sym::IntSet::SinglePoint(value.value())); } return std::nullopt; } ffi::Optional Visit_(const LetNode* op) final { TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit(op->value)); - if (arith::IsIndexTypedExpr(op->value)) { + if (sym::IsIndexTypedExpr(op->value)) { dom_analyzer_->Bind(op->var, op->value); - dom_map_.emplace(op->var.get(), arith::IntSet::SinglePoint(op->value)); + dom_map_.emplace(op->var.get(), sym::IntSet::SinglePoint(op->value)); } TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit(op->body)); - if (arith::IsIndexTypedExpr(op->value)) { + if (sym::IsIndexTypedExpr(op->value)) { dom_map_.erase(op->var.get()); } return std::nullopt; @@ -356,7 +356,7 @@ class BufferAccessRegionCollector : public StmtExprVisitor { dom = Range::FromMinExtent(IntImm(op->value.ty(), 0), op->value); } dom_analyzer_->Bind(iter->var, dom); - dom_map_.emplace(iter->var.get(), arith::IntSet::FromRange(dom)); + dom_map_.emplace(iter->var.get(), sym::IntSet::FromRange(dom)); size_t n_pending_before = pending_flat_alloc_buffers_.size(); TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(op)); CompactPendingFlatAllocBuffers(n_pending_before); @@ -383,7 +383,7 @@ class BufferAccessRegionCollector : public StmtExprVisitor { size_t n_ancestor_loops = it->second; // Step 1. Stop ancestor loop vars out of the allocation block from // being relaxed unless NeedRelaxThread() is true. - std::vector non_relaxed(n_ancestor_loops); + std::vector non_relaxed(n_ancestor_loops); for (size_t i = 0; i < n_ancestor_loops; ++i) { const IterVar& iter = ancestor_iters_[i]; const VarNode* v = iter->var.get(); @@ -469,7 +469,7 @@ class BufferAccessRegionCollector : public StmtExprVisitor { result_region.resize(nd_int_set.size()); for (size_t i = 0; i < nd_int_set.size(); ++i) { - const arith::IntSet& int_set = nd_int_set[i]; + const sym::IntSet& int_set = nd_int_set[i]; Range original = Range(/*begin=*/IntImm(original_shape[i].ty(), 0), /*end=*/original_shape[i]); Range range = int_set.CoverRange(original); @@ -505,7 +505,7 @@ class BufferAccessRegionCollector : public StmtExprVisitor { if (ffi::StructuralWalk(extent, walkfn).has_value()) { // try estimate a constant upperbound on region's extent int64_t upperbound = dom_analyzer_->const_int_bound(extent)->max_value; - if (upperbound != arith::ConstIntBound::kPosInf) { + if (upperbound != sym::ConstIntBound::kPosInf) { extent = IntImm(extent.ty(), upperbound); } else { result_region.Set(i, original); @@ -552,13 +552,13 @@ class BufferAccessRegionCollector : public StmtExprVisitor { var2buffer_; /*! \brief The map from loop vars to their iter range. */ - std::unordered_map dom_map_; + std::unordered_map dom_map_; /*! \brief Extra map from free vars to their iter range hints. */ - std::unordered_map hint_map_; + std::unordered_map hint_map_; /*! \brief Unresolved conditions within current scope. */ std::vector pending_conditions_; /*! \brief The analyzer aware of loop domains. */ - arith::Analyzer dom_analyzer_; + sym::Analyzer dom_analyzer_; /*! \brief The map from BufferVar to it's relaxed access set. */ std::unordered_map relaxed_accesses_; diff --git a/src/s_tir/transform/hoist_expression.cc b/src/s_tir/transform/hoist_expression.cc index 3d4c52df1fe4..ed5987009aed 100644 --- a/src/s_tir/transform/hoist_expression.cc +++ b/src/s_tir/transform/hoist_expression.cc @@ -20,7 +20,6 @@ /*! * \file hoist_expression.cc */ -#include #include #include #include @@ -29,6 +28,7 @@ #include #include #include +#include #include #include @@ -36,9 +36,9 @@ #include #include -#include "../../arith/interval_set.h" #include "../../runtime/thread_storage_scope.h" #include "../../s_tir/ir/ir_mutator_with_analyzer.h" +#include "../../sym/interval_set.h" #include "ir_utils.h" namespace tvm { @@ -464,7 +464,7 @@ class ExpressionHoister : public s_tir::IRMutatorWithAnalyzer { static Stmt Hoist(Stmt stmt, HoistExpressionConfig config) { auto loop_info = HoistInfoCollector::Collect(stmt, config); - arith::Analyzer analyzer; + sym::Analyzer analyzer; auto hoister = ffi::make_object(std::move(loop_info), config, analyzer); stmt = hoister->Mutate(stmt, InplaceMode::kAllow).ValueOrUnchanged(std::move(stmt)); stmt = s_tir::ConvertSSA(std::move(stmt)); @@ -476,7 +476,7 @@ class ExpressionHoister : public s_tir::IRMutatorWithAnalyzer { public: explicit ExpressionHoister(std::vector loop_info, - HoistExpressionConfig config, const arith::Analyzer& analyzer) + HoistExpressionConfig config, const sym::Analyzer& analyzer) : Parent(analyzer), config_(config) { for (auto& info : loop_info) { // Mark let bindings to use if they are enabled on their own. diff --git a/src/s_tir/transform/inject_permuted_layout.cc b/src/s_tir/transform/inject_permuted_layout.cc index 4f5fe0c50699..64b8bf645526 100644 --- a/src/s_tir/transform/inject_permuted_layout.cc +++ b/src/s_tir/transform/inject_permuted_layout.cc @@ -21,12 +21,12 @@ * \file inject_permuted_layout.cc * \brief The pass injects permuted layout for shared memory buffers to avoid bank conflicts. */ -#include #include #include #include #include #include +#include #include #include @@ -39,7 +39,7 @@ namespace tvm { namespace s_tir { using namespace tvm::tirx; -using namespace arith; +using namespace sym; using namespace runtime; namespace { diff --git a/src/s_tir/transform/inject_ptx_ldg32.cc b/src/s_tir/transform/inject_ptx_ldg32.cc index 97f34a29cd7c..3e4ba98df7a4 100644 --- a/src/s_tir/transform/inject_ptx_ldg32.cc +++ b/src/s_tir/transform/inject_ptx_ldg32.cc @@ -17,20 +17,20 @@ * under the License. */ -#include -#include #include #include #include #include #include #include +#include +#include #include #include #include -#include "../../arith/const_fold.h" -#include "../../arith/pattern_match.h" +#include "../../sym/const_fold.h" +#include "../../sym/pattern_match.h" namespace tvm { namespace s_tir { diff --git a/src/s_tir/transform/inject_software_pipeline.cc b/src/s_tir/transform/inject_software_pipeline.cc index 0dd4c8a0a63e..2117f72cec19 100644 --- a/src/s_tir/transform/inject_software_pipeline.cc +++ b/src/s_tir/transform/inject_software_pipeline.cc @@ -490,9 +490,9 @@ class PipelineRewriter : public StmtExprMutator { for (size_t i = 0; i < region1.size(); i++) { Range dim1 = region1[i]; Range dim2 = region2[i]; - auto int_set1 = arith::IntSet::FromRange(dim1); - auto int_set2 = arith::IntSet::FromRange(dim2); - if (arith::Intersect({int_set1, int_set2}).IsNothing()) { + auto int_set1 = sym::IntSet::FromRange(dim1); + auto int_set2 = sym::IntSet::FromRange(dim2); + if (sym::Intersect({int_set1, int_set2}).IsNothing()) { return false; } } @@ -652,7 +652,7 @@ class PipelineRewriter : public StmtExprMutator { // Determine where to insert async_wait and the corresponding wait count. void PopulateWaitCounts(const std::vector& new_blocks, - arith::AnalyzerObj* ana_normalized, + sym::AnalyzerObj* ana_normalized, const std::unordered_map& buffer_to_commit_group, std::map* async_states_local) { for (size_t i = 0; i < new_blocks.size(); ++i) { @@ -789,7 +789,7 @@ class PipelineRewriter : public StmtExprMutator { ffi::Array CompletePipelineLoopStatements( const std::vector& blocks, const std::map& async_states_local, - arith::AnalyzerObj* ana_normalized) const { + sym::AnalyzerObj* ana_normalized) const { std::vector new_blocks = blocks; std::vector commit_group_indices(new_blocks.size(), -1); for (const auto& [stage_id, state] : async_states_local) { @@ -892,7 +892,7 @@ class PipelineRewriter : public StmtExprMutator { // In contrast to analyzer_ which is bound to [start, end), this one is bound to // the "normalized" range, [pipeline_loop_->min, extent). - arith::Analyzer ana_normalized; + sym::Analyzer ana_normalized; if (!is_unit_loop) { ana_normalized->Bind(new_loop_var.as_or_throw(), Range(pipeline_loop_->min, extent)); } @@ -1051,7 +1051,7 @@ class PipelineRewriter : public StmtExprMutator { MakeSBlock(std::move(new_loop), buffer_data_to_buffer_)); } - arith::Analyzer analyzer_; + sym::Analyzer analyzer_; ffi::Map buffer_data_to_buffer_; const std::unordered_set& double_buffers_; ffi::Array pipeline_allocs_; diff --git a/src/s_tir/transform/inject_virtual_thread.cc b/src/s_tir/transform/inject_virtual_thread.cc index b80cee1f49c0..29bbe00de72b 100644 --- a/src/s_tir/transform/inject_virtual_thread.cc +++ b/src/s_tir/transform/inject_virtual_thread.cc @@ -236,7 +236,7 @@ class VTInjector : public s_tir::IRMutatorWithAnalyzer { using s_tir::IRMutatorWithAnalyzer::Mutate_; // constructor - VTInjector(arith::AnalyzerObj* analyzer, Var var, int num_threads, + VTInjector(sym::AnalyzerObj* analyzer, Var var, int num_threads, const std::unordered_set& touched_var, bool allow_share) : IRMutatorWithAnalyzer(analyzer), var_(var), @@ -717,7 +717,7 @@ Pass InjectVirtualThread() { auto pass_func = [](PrimFunc f, IRModule m, PassContext ctx) { auto* n = f.CopyOnWrite(); - arith::Analyzer analyzer; + sym::Analyzer analyzer; n->body = ffi::make_object(analyzer) ->Mutate(n->body, InplaceMode::kAllow) diff --git a/src/s_tir/transform/ir_utils.cc b/src/s_tir/transform/ir_utils.cc index adb3e54c5f11..7353f7de57d9 100644 --- a/src/s_tir/transform/ir_utils.cc +++ b/src/s_tir/transform/ir_utils.cc @@ -19,9 +19,9 @@ #include "ir_utils.h" -#include #include #include +#include #include namespace tvm { @@ -106,7 +106,7 @@ ffi::Array ConvertIndices(const MatchBufferRegion& match_buffer, const TensorRegion& source = match_buffer->source; TVM_FFI_ICHECK_EQ(indices.size(), target->shape.size()); - arith::Analyzer analyzer; + sym::Analyzer analyzer; ffi::Array result; result.reserve(source->region.size()); size_t offset = source->region.size() - indices.size(); @@ -128,7 +128,7 @@ Region ConvertRegion(const MatchBufferRegion& match_buffer, const Region& region const TensorRegion& source = match_buffer->source; TVM_FFI_ICHECK_EQ(region.size(), target->shape.size()); - arith::Analyzer analyzer; + sym::Analyzer analyzer; Region result; result.reserve(source->region.size()); size_t offset = source->region.size() - region.size(); diff --git a/src/s_tir/transform/loop_partition.cc b/src/s_tir/transform/loop_partition.cc index 392d821b71a8..1446dccd1bc8 100644 --- a/src/s_tir/transform/loop_partition.cc +++ b/src/s_tir/transform/loop_partition.cc @@ -20,8 +20,6 @@ /*! * \file loop_partition.cc */ -#include -#include #include #include #include @@ -34,6 +32,8 @@ #include #include #include +#include +#include #include #include @@ -41,8 +41,8 @@ #include #include -#include "../../arith/interval_set.h" #include "../../runtime/thread_storage_scope.h" +#include "../../sym/interval_set.h" #include "ir_utils.h" namespace tvm { @@ -82,9 +82,9 @@ class LoopPartitionConfig : public ffi::ObjectRef { TVM_REGISTER_PASS_CONFIG_OPTION("s_tir.LoopPartition", LoopPartitionConfig); -using arith::DeduceBound; -using arith::Intersect; -using arith::IntSet; +using sym::DeduceBound; +using sym::Intersect; +using sym::IntSet; using PartitionKey = std::pair; struct PartitionKeyHash { @@ -219,7 +219,7 @@ class CandidateSelector final : public StmtExprVisitor { bool no_split_{false}; bool partition_const_loop_{false}; std::unordered_map record_; - arith::Analyzer analyzer_; + sym::Analyzer analyzer_; }; // Finder try best to find partitions for hinted vars @@ -336,7 +336,7 @@ class PartitionFinder : public StmtExprVisitor { hint_map_, relax_map_); IntSet part2 = DeduceBound(current_var_.as_or_throw(), LE(op->a, op->b), hint_map_, relax_map_); - interval = arith::Intersect({part1, part2}); + interval = sym::Intersect({part1, part2}); if (!interval.IsNothing()) { // cond is true within interval partitions[{cond, true}] = interval; @@ -528,7 +528,7 @@ class LoopPartitioner : public StmtExprMutator { bool partition_thread_scope); std::pair GetIntervalAndCondset(const Partition& partitions, - const arith::IntervalSet& for_interval, + const sym::IntervalSet& for_interval, bool cond_value, bool has_partition_hint); inline Stmt MakeFor(const ffi::Object* op, PrimExpr extent, Stmt body); @@ -536,7 +536,7 @@ class LoopPartitioner : public StmtExprMutator { /* Candidate IRs that may be partitioned potentially */ std::unordered_map hint_map_; std::unordered_map relax_map_; - arith::Analyzer analyzer_; + sym::Analyzer analyzer_; ffi::ObjectPtr selector; bool no_unroll_loop_with_extent_one_; bool unroll_loop_with_partition_hint_no_interval_; @@ -545,15 +545,15 @@ class LoopPartitioner : public StmtExprMutator { // Returns an interval (in the first component) in which all the conditions // given in the second component provably have value given by cond_value std::pair LoopPartitioner::GetIntervalAndCondset( - const Partition& partitions, const arith::IntervalSet& for_interval, bool cond_value, + const Partition& partitions, const sym::IntervalSet& for_interval, bool cond_value, bool has_partition_hint) { ffi::Array sets; ExpressionSet cond_set; for (const auto& kv : partitions) { if (kv.first.second == cond_value) { - arith::IntervalSet interval = kv.second.as_or_throw(); - arith::IntervalSet intersection = arith::Intersect(analyzer_.get(), interval, for_interval); + sym::IntervalSet interval = kv.second.as_or_throw(); + sym::IntervalSet intersection = sym::Intersect(analyzer_.get(), interval, for_interval); if (!intersection->IsEmpty()) { sets.push_back(kv.second); @@ -566,21 +566,21 @@ std::pair LoopPartitioner::GetIntervalAndCondset( // Try to find the intersection of the cond_intervals until the intersection // is nothing when has_partition_hint is true. if (interval.IsNothing() && has_partition_hint) { - arith::IntervalSet cond_intersection = arith::IntervalSet::Everything(); + sym::IntervalSet cond_intersection = sym::IntervalSet::Everything(); cond_set.clear(); for (const auto& kv : partitions) { if (kv.first.second == cond_value) { - arith::IntervalSet cond_interval = kv.second.as_or_throw(); - arith::IntervalSet intersection = - arith::Intersect(analyzer_.get(), cond_interval, for_interval); + sym::IntervalSet cond_interval = kv.second.as_or_throw(); + sym::IntervalSet intersection = + sym::Intersect(analyzer_.get(), cond_interval, for_interval); if (!intersection->IsEmpty()) { - cond_intersection = arith::Intersect(analyzer_.get(), cond_intersection, cond_interval); + cond_intersection = sym::Intersect(analyzer_.get(), cond_intersection, cond_interval); // Return the latest interval and cond_set if the cond_intersection is nothing. if (!cond_intersection->IsEmpty()) { cond_set.insert(kv.first.first); - interval = arith::IntervalSet(analyzer_->Simplify(cond_intersection->min_value), - analyzer_->Simplify(cond_intersection->max_value)); + interval = sym::IntervalSet(analyzer_->Simplify(cond_intersection->min_value), + analyzer_->Simplify(cond_intersection->max_value)); } else { break; } @@ -640,7 +640,7 @@ std::pair LoopPartitioner::GetIntervalAndCondset( */ Stmt LoopPartitioner::TryPartition(const Stmt& stmt, Var var, PrimExpr min, PrimExpr max, Stmt body, bool partition_thread_scope) { - using namespace arith; + using namespace sym; // include hint of var. hint_map_.insert({var.get(), IntSet::Interval(min, max)}); @@ -651,7 +651,7 @@ Stmt LoopPartitioner::TryPartition(const Stmt& stmt, Var var, PrimExpr min, Prim hint_map_.erase(var.get()); if (finder->partitions.empty()) return Stmt(); - arith::IntervalSet for_interval(min, max); + sym::IntervalSet for_interval(min, max); auto [middle_interval, cond_set, opt_cond_value] = [&]() -> std::tuple> { diff --git a/src/s_tir/transform/lower_async_dma.cc b/src/s_tir/transform/lower_async_dma.cc index bc1dfc6d437b..ec1e67285da6 100644 --- a/src/s_tir/transform/lower_async_dma.cc +++ b/src/s_tir/transform/lower_async_dma.cc @@ -21,8 +21,6 @@ * \file lower_async_dma.cc */ -#include -#include #include #include #include @@ -30,6 +28,8 @@ #include #include #include +#include +#include #include #include @@ -48,7 +48,7 @@ class AsyncDMALowerer : public s_tir::IRMutatorWithAnalyzer { using s_tir::IRMutatorWithAnalyzer::Mutate; using s_tir::IRMutatorWithAnalyzer::Mutate_; - explicit AsyncDMALowerer(bool dma_bypass_cache, const arith::Analyzer& analyzer) + explicit AsyncDMALowerer(bool dma_bypass_cache, const sym::Analyzer& analyzer) : IRMutatorWithAnalyzer(analyzer), dma_bypass_cache_(dma_bypass_cache) {} // TODO(leiwang1999): split lower async DMA support for CUDA and Hexagon Backend @@ -60,7 +60,7 @@ class AsyncDMALowerer : public s_tir::IRMutatorWithAnalyzer { // if for loop is not a memcpy of a contiguous region, it might be a cuda cp.async behavior std::optional mem_copy = - s_tir::IdentifyMemCpy(ffi::GetRef(loop), ffi::GetRef(analyzer_)); + 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 s_tir::IRMutatorWithAnalyzer::Mutate_(loop, inplace_mode); @@ -194,7 +194,7 @@ namespace transform { Pass LowerAsyncDMA() { auto pass_func = [=](PrimFunc f, IRModule m, PassContext ctx) { auto fptr = f.CopyOnWrite(); - arith::Analyzer analyzer; + sym::Analyzer analyzer; bool dma_bypass_cache = ctx->GetConfig("tirx.experimental_dma_bypass_cache", false).value(); fptr->body = ffi::make_object(dma_bypass_cache, analyzer) diff --git a/src/s_tir/transform/lower_cross_thread_reduction.cc b/src/s_tir/transform/lower_cross_thread_reduction.cc index 51ff831942df..fbac8846816b 100644 --- a/src/s_tir/transform/lower_cross_thread_reduction.cc +++ b/src/s_tir/transform/lower_cross_thread_reduction.cc @@ -20,7 +20,6 @@ /*! * \file lower_cross_thread_reduction.cc */ -#include #include #include #include @@ -29,6 +28,7 @@ #include #include #include +#include #include #include @@ -113,7 +113,7 @@ bool IsDominantBlock(const SBlock& scope_block, const SBlock& block) { * check again. */ bool IsReductionBlock(const SBlockRealize& realize, const ffi::Map& loop_range_map, - const SBlock& scope_block, arith::AnalyzerObj* analyzer) { + const SBlock& scope_block, sym::AnalyzerObj* analyzer) { const auto* block = realize->block.as(); // Cond 1. The block has the `init` statement. if (!block->init.has_value()) { @@ -973,7 +973,7 @@ class CrossThreadReductionTransformer : public StmtExprMutator { std::unordered_map> block2new_buffers_; std::unordered_map loop2new_stmt_; ffi::Map loop_range_map_; - arith::Analyzer analyzer_; + sym::Analyzer analyzer_; int block_idx_depth = 0; int thread_idx_depth = 0; diff --git a/src/s_tir/transform/lower_match_buffer.cc b/src/s_tir/transform/lower_match_buffer.cc index c17dda686e5c..8574f92609b8 100644 --- a/src/s_tir/transform/lower_match_buffer.cc +++ b/src/s_tir/transform/lower_match_buffer.cc @@ -22,7 +22,6 @@ * \brief The pass for lowering match_buffer. */ -#include #include #include #include @@ -30,6 +29,7 @@ #include #include #include +#include #include #include @@ -313,7 +313,7 @@ class MatchBufferLower : public StmtExprMutator { /*! \brief BufferVar region mapping. */ ffi::Map match_buffers_; /*! \brief The analyzer */ - arith::Analyzer analyzer_; + sym::Analyzer analyzer_; }; namespace transform { diff --git a/src/s_tir/transform/lower_thread_allreduce.cc b/src/s_tir/transform/lower_thread_allreduce.cc index 6f4691e26ec7..4ee9d59067f8 100644 --- a/src/s_tir/transform/lower_thread_allreduce.cc +++ b/src/s_tir/transform/lower_thread_allreduce.cc @@ -21,7 +21,6 @@ * Lower allreduce to device implementable ir. * \file lower_thread_allreduce.cc */ -#include #include #include #include @@ -29,6 +28,7 @@ #include #include #include +#include #include #include #include @@ -881,7 +881,7 @@ class ThreadAllreduceBuilder final : public StmtExprMutator { // The load remap std::unordered_map load_remap_; // Internal analyzer - arith::Analyzer analyzer_; + sym::Analyzer analyzer_; public: const VarNode* GetAllocationKey(const VarNode* buffer) const { 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 67ad132466cb..957ff4d985e1 100644 --- a/src/s_tir/transform/manifest_shared_memory_local_stage.cc +++ b/src/s_tir/transform/manifest_shared_memory_local_stage.cc @@ -26,7 +26,6 @@ * memory. This is similar to the schedule primitive cache_read, but it bypasses the limitation * of requiring buffer access to be contiguous in each dimension. */ -#include #include #include #include @@ -34,6 +33,7 @@ #include #include #include +#include #include #include diff --git a/src/s_tir/transform/memhammer_coalesce.cc b/src/s_tir/transform/memhammer_coalesce.cc index 45e8640e7bef..d141c70a0ac6 100644 --- a/src/s_tir/transform/memhammer_coalesce.cc +++ b/src/s_tir/transform/memhammer_coalesce.cc @@ -106,7 +106,7 @@ Stmt SplitBindVectorize(const Stmt& stmt, const ConstraintSet& constraints) { int n = factors.size(); std::vector new_loop_vars; new_loop_vars.reserve(n); - arith::Analyzer analyzer; + sym::Analyzer analyzer; for (int i = 0; i < n; i++) { const PrimExpr& factor = factors[i]; PrimVar var = loop->loop_var.CopyWithSuffix("_" + std::to_string(i)); @@ -175,7 +175,7 @@ ffi::Array GetMapping(const Stmt& stmt, const ConstraintSet& constrain write_region->region.size() == write_index.size() && write_region->source.as_or_throw().same_as(buf_store->buffer)); ffi::Array result; - arith::Analyzer analyzer; + sym::Analyzer analyzer; for (int i = 0; i < static_cast(write_region->region.size()); i++) { PrimExpr pattern = analyzer->Simplify(write_index[i] - write_region->region[i]->min); if (!is_zero(pattern)) { @@ -198,12 +198,11 @@ Stmt InverseMapping::Rewrite(const Stmt& stmt, const ConstraintSet& constraints, body = loop->body; } // Step 2. Get Inverse mapping - arith::Analyzer analyzer; - auto iter_map = arith::DetectIterMap(mapping_pattern, var_range, IntImm::Bool(true), - arith::Bijective, analyzer); + sym::Analyzer analyzer; + auto iter_map = + sym::DetectIterMap(mapping_pattern, var_range, IntImm::Bool(true), sym::Bijective, analyzer); TVM_FFI_ICHECK_EQ(iter_map->indices.size(), loop_vars.size()); - ffi::Map inverse_mapping = - arith::InverseAffineIterMap(iter_map->indices, loop_vars); + ffi::Map inverse_mapping = sym::InverseAffineIterMap(iter_map->indices, loop_vars); // Step 3. Generate new body TensorRegion read_region = constraints.read_region; TensorRegion write_region = constraints.write_region; diff --git a/src/s_tir/transform/memhammer_intermediate_stage.cc b/src/s_tir/transform/memhammer_intermediate_stage.cc index 7fcdd20c66f5..baf6082cf617 100644 --- a/src/s_tir/transform/memhammer_intermediate_stage.cc +++ b/src/s_tir/transform/memhammer_intermediate_stage.cc @@ -96,7 +96,7 @@ class IndexPatternFinder : public StmtExprVisitor { static ffi::Array getRankPromotedShape(ffi::Array indices, const ffi::Map& var_range, ffi::Array* rewrite_indices) { - ffi::Map var_dom = arith::AsIntSet(var_range); + ffi::Map var_dom = sym::AsIntSet(var_range); ffi::Array new_shape; for (const PrimExpr& expr : indices) { ffi::Array indices_dim; @@ -288,7 +288,7 @@ std::pair InsertCacheStage(Stmt stmt, bool is_write_cache, ffi::S } } - arith::Analyzer analyzer; + sym::Analyzer analyzer; const TensorLoadNode* target_buffer_load = nullptr; if (is_write_cache) { auto walk_fn = [&](const TensorLoad& buffer_load) -> ffi::Expected { diff --git a/src/s_tir/transform/memhammer_lower_auto_copy.cc b/src/s_tir/transform/memhammer_lower_auto_copy.cc index f1a8a052c2c5..f5c3771e58f4 100644 --- a/src/s_tir/transform/memhammer_lower_auto_copy.cc +++ b/src/s_tir/transform/memhammer_lower_auto_copy.cc @@ -17,7 +17,6 @@ * under the License. */ -#include #include #include #include @@ -27,6 +26,7 @@ #include #include #include +#include #include #include @@ -509,7 +509,7 @@ class AutoPadder { .as_or_throw(); PrimExpr e2 = ffi::StructuralMap(e, f_substitute_one) .as_or_throw(); - arith::Analyzer analyzer; + sym::Analyzer analyzer; PrimExpr delta = ffi::StructuralMap(e2 - e1, f_substitute) .as_or_throw(); return !analyzer->CanProve(delta != 1); @@ -547,7 +547,7 @@ class AutoPadder { runtime::StorageScope scope = runtime::StorageScope::Create(op->buffer.scope()); if (scope.rank == runtime::StorageRank::kShared) { ffi::Array substitued_indices; - arith::Analyzer analyzer; + sym::Analyzer analyzer; auto f_substitute = [this](const Var& var) -> ffi::Expected> { if (auto repl = substitute_map_.Get(var)) return ffi::Any(*std::move(repl)); return ffi::Unchanged(); @@ -582,7 +582,7 @@ class AutoPadder { runtime::StorageScope scope = runtime::StorageScope::Create(buffer.scope()); if (scope.rank == runtime::StorageRank::kShared) { ffi::Array substitued_indices; - arith::Analyzer analyzer; + sym::Analyzer analyzer; auto f_substitute = [this](const Var& var) -> ffi::Expected> { if (auto repl = substitute_map_.Get(var)) return ffi::Any(*std::move(repl)); return ffi::Unchanged(); @@ -632,7 +632,7 @@ class AutoPadder { var_range_.Set(var, Range::FromMinExtent(0, region[i]->extent)); } ffi::Array substitued_indices; - arith::Analyzer analyzer; + sym::Analyzer analyzer; auto f_substitute = [this](const Var& var) -> ffi::Expected> { if (auto repl = substitute_map_.Get(var)) return ffi::Any(*std::move(repl)); diff --git a/src/s_tir/transform/memhammer_rewrite_rule.h b/src/s_tir/transform/memhammer_rewrite_rule.h index 4c0a136f37c4..224376ba6f19 100644 --- a/src/s_tir/transform/memhammer_rewrite_rule.h +++ b/src/s_tir/transform/memhammer_rewrite_rule.h @@ -19,11 +19,11 @@ #ifndef TVM_S_TIR_TRANSFORM_MEMHAMMER_REWRITE_RULE_H_ #define TVM_S_TIR_TRANSFORM_MEMHAMMER_REWRITE_RULE_H_ -#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 e21a086e34a1..8a8131abe171 100644 --- a/src/s_tir/transform/memhammer_tensorcore_rewrite.cc +++ b/src/s_tir/transform/memhammer_tensorcore_rewrite.cc @@ -46,7 +46,7 @@ std::pair> TileWmmaBlock(Stmt stmt) { PrimExpr extent_last1 = loops[n - 1]->extent; PrimExpr extent_last2 = loops[n - 2]->extent; { - arith::Analyzer analyzer; + sym::Analyzer analyzer; if (!analyzer->CanProveEqual(floormod(extent_last1, 16), 0) || !analyzer->CanProveEqual(floormod(extent_last2, 16), 0)) { return std::make_pair(stmt, std::nullopt); @@ -97,11 +97,11 @@ std::pair> TileWmmaBlock(Stmt stmt) { ffi::Array RelaxIndices(const ffi::Array& indices, const ffi::Array& shape, - const ffi::Map& var_dom) { - ffi::Array int_set; + const ffi::Map& var_dom) { + ffi::Array int_set; int_set.reserve(indices.size()); for (auto& indice : indices) { - int_set.push_back(arith::EvalSet(indice, var_dom)); + int_set.push_back(sym::EvalSet(indice, var_dom)); } int ndim = int_set.size(); ffi::Array region; @@ -118,7 +118,7 @@ ffi::Array RelaxIndices(const ffi::Array& indices, * \return The stmt after rewrite */ Stmt RewriteWmmaLoad(Stmt stmt) { - using arith::IntSet; + using sym::IntSet; const PrimType dtype_ty = PrimType::Float(16); const PrimType& dtype = dtype_ty; const PrimType int32_ty = PrimType::Int(32); @@ -217,7 +217,7 @@ Stmt RewriteWmmaLoad(Stmt stmt) { * \return The stmt after rewrite */ Stmt RewriteWmmaStore(Stmt stmt) { - using arith::IntSet; + using sym::IntSet; const PrimType int32_ty = PrimType::Int(32); Stmt body = stmt; @@ -378,7 +378,7 @@ std::pair> TileMmaToGlobalBlock(Stmt stmt) { PrimExpr extent_last1 = loops[n - 1]->extent; PrimExpr extent_last2 = loops[n - 2]->extent; { - arith::Analyzer analyzer; + sym::Analyzer analyzer; // Only tile when both extent % 8 == 0 if (!analyzer->CanProveEqual(floormod(extent_last1, 8), 0) || !analyzer->CanProveEqual(floormod(extent_last2, 8), 0)) { @@ -434,7 +434,7 @@ std::pair> TileMmaToGlobalBlock(Stmt stmt) { * \return The stmt after rewrite */ Stmt RewriteMmaStore(Stmt stmt) { - using arith::IntSet; + using sym::IntSet; const PrimType int32_ty = PrimType::Int(32); // Step 1. Get inner loop body diff --git a/src/s_tir/transform/renormalize_split_pattern.cc b/src/s_tir/transform/renormalize_split_pattern.cc index 8b5f6cff4bdb..a93d992c3c76 100644 --- a/src/s_tir/transform/renormalize_split_pattern.cc +++ b/src/s_tir/transform/renormalize_split_pattern.cc @@ -30,14 +30,14 @@ #include #include -#include "../../arith/pattern_match.h" #include "../../s_tir/ir/ir_mutator_with_analyzer.h" +#include "../../sym/pattern_match.h" namespace tvm { namespace s_tir { using namespace tvm::tirx; -using namespace arith; +using namespace sym; // macro for doing simple rewrite #define TRY_REWRITE(SrcExpr, ResExpr) \ @@ -218,7 +218,7 @@ namespace transform { Pass RenormalizeSplitPattern() { auto pass_func = [](PrimFunc f, IRModule m, PassContext ctx) { auto* n = f.CopyOnWrite(); - arith::Analyzer analyzer; + sym::Analyzer analyzer; n->body = ffi::make_object(analyzer) ->Mutate(n->body, InplaceMode::kAllow) .ValueOrUnchanged(std::move(n->body)); diff --git a/src/s_tir/transform/stmt_simplify.cc b/src/s_tir/transform/stmt_simplify.cc index 36d8cf432806..2e1246c6b229 100644 --- a/src/s_tir/transform/stmt_simplify.cc +++ b/src/s_tir/transform/stmt_simplify.cc @@ -33,7 +33,7 @@ using namespace tirx; class StmtSimplifier final : public tirx::StmtSimplifier { public: using Parent = tirx::StmtSimplifier; - StmtSimplifier(const arith::Analyzer& analyzer, tirx::StmtSimplifyConfig config) + StmtSimplifier(const sym::Analyzer& analyzer, tirx::StmtSimplifyConfig config) : Parent(GlobalVTable(), analyzer, config) {} using Parent::Mutate_; using Parent::Run; @@ -71,7 +71,7 @@ class StmtSimplifier final : public tirx::StmtSimplifier { } }; -PrimFunc StmtSimplify(PrimFunc func, const arith::Analyzer& analyzer) { +PrimFunc StmtSimplify(PrimFunc func, const sym::Analyzer& analyzer) { auto config = tvm::transform::PassConfigWithDefaults(); return ffi::make_object(analyzer, config)->Run(std::move(func)); } @@ -79,7 +79,7 @@ PrimFunc StmtSimplify(PrimFunc func, const arith::Analyzer& analyzer) { namespace transform { Pass StmtSimplify() { auto pass_func = [](PrimFunc func, IRModule, tvm::transform::PassContext ctx) { - arith::Analyzer analyzer; + sym::Analyzer analyzer; auto config = ctx->GetConfig("tirx.StmtSimplify") .value_or(tvm::transform::PassConfigWithDefaults()); return ffi::make_object(analyzer, config)->Run(std::move(func)); diff --git a/src/s_tir/transform/stmt_simplify.h b/src/s_tir/transform/stmt_simplify.h index 8dcff5bcc6d3..3149e71e1adf 100644 --- a/src/s_tir/transform/stmt_simplify.h +++ b/src/s_tir/transform/stmt_simplify.h @@ -19,11 +19,11 @@ #ifndef TVM_S_TIR_TRANSFORM_STMT_SIMPLIFY_H_ #define TVM_S_TIR_TRANSFORM_STMT_SIMPLIFY_H_ -#include +#include #include namespace tvm { namespace s_tir { -tirx::PrimFunc StmtSimplify(tirx::PrimFunc func, const arith::Analyzer& analyzer); +tirx::PrimFunc StmtSimplify(tirx::PrimFunc func, const sym::Analyzer& analyzer); } // namespace s_tir } // namespace tvm #endif // TVM_S_TIR_TRANSFORM_STMT_SIMPLIFY_H_ diff --git a/src/s_tir/transform/storage_access.cc b/src/s_tir/transform/storage_access.cc index 164bd610cfc0..2b1a62f78e09 100644 --- a/src/s_tir/transform/storage_access.cc +++ b/src/s_tir/transform/storage_access.cc @@ -60,7 +60,7 @@ ffi::Optional StorageAccessVisitor::Visit_(const TensorLoadNode* e.buffer = buf; e.dtype = op->ty.as_or_throw().WithLanes(1); for (const auto& index : op->indices) { - e.touched.push_back(arith::IntSet::Vector(index)); + e.touched.push_back(sym::IntSet::Vector(index)); } e.type = kRead; e.scope = scope; @@ -83,7 +83,7 @@ ffi::Optional StorageAccessVisitor::Visit_(const BufferStoreNode e.buffer = buf; e.dtype = op->value.ty().WithLanes(1); for (const auto& index : op->indices) { - e.touched.push_back(arith::IntSet::Vector(index)); + e.touched.push_back(sym::IntSet::Vector(index)); } e.type = kWrite; e.scope = scope; @@ -189,15 +189,15 @@ ffi::Optional StorageAccessVisitor::Visit_(const ForNode* op) { scope_.pop_back(); if (s.access.size() != 0) { // relax the touched set to contain all ranges in the loop. - std::unordered_map relax_map; + std::unordered_map relax_map; relax_map[op->loop_var.get()] = - arith::IntSet::FromRange(Range::FromMinExtent(op->min, op->extent)); + sym::IntSet::FromRange(Range::FromMinExtent(op->min, op->extent)); for (AccessEntry& e : s.access) { if (e.buffer.defined()) { TVM_FFI_ICHECK(e.touched.size()); - ffi::Array new_touched; + ffi::Array new_touched; for (const auto& touched : e.touched) { - new_touched.push_back(arith::EvalSet(touched, relax_map)); + new_touched.push_back(sym::EvalSet(touched, relax_map)); } e.touched = std::move(new_touched); } @@ -283,7 +283,7 @@ ffi::Optional StorageAccessVisitor::Visit_(const CallNode* op) { e.buffer = buf; e.dtype = value_dtype.WithLanes(1); for (size_t i = is_load ? 1 : 2; i + 1 < op->args.size(); ++i) { - e.touched.push_back(arith::IntSet::Vector(op->args[i].as_or_throw())); + e.touched.push_back(sym::IntSet::Vector(op->args[i].as_or_throw())); } e.type = is_load ? kRead : kWrite; e.scope = scope; @@ -323,7 +323,7 @@ ffi::Optional StorageAccessVisitor::Visit_(const CallNode* op) { e.threads = env_threads(); e.dtype = dtype; e.buffer = buffer; - e.touched = {arith::IntSet::FromRange(Range::FromMinExtent(offset, extent))}; + e.touched = {sym::IntSet::FromRange(Range::FromMinExtent(offset, extent))}; e.scope = scope; if (flag->value & 1) { e.type = kRead; diff --git a/src/s_tir/transform/storage_access.h b/src/s_tir/transform/storage_access.h index 83b56d3520b6..abcea29f8df6 100644 --- a/src/s_tir/transform/storage_access.h +++ b/src/s_tir/transform/storage_access.h @@ -24,10 +24,10 @@ #ifndef TVM_S_TIR_TRANSFORM_STORAGE_ACCESS_H_ #define TVM_S_TIR_TRANSFORM_STORAGE_ACCESS_H_ -#include #include #include #include +#include #include #include @@ -68,7 +68,7 @@ class StorageAccessVisitor : public StmtExprVisitor { * * Has one IntSet for each index in the buffer being accessed. */ - ffi::Array touched; + ffi::Array touched; /*! \brief The type of access */ AccessType type; /*! \brief The storage scope */ diff --git a/src/s_tir/transform/transform_mma_buffer_layout.cc b/src/s_tir/transform/transform_mma_buffer_layout.cc index 3784db34e929..2838e9089d46 100644 --- a/src/s_tir/transform/transform_mma_buffer_layout.cc +++ b/src/s_tir/transform/transform_mma_buffer_layout.cc @@ -17,13 +17,13 @@ * under the License. */ -#include #include #include #include #include #include #include +#include #include #include @@ -188,7 +188,7 @@ class MmaBufferLayoutTransformer : public StmtExprMutator { } private: - arith::Analyzer analyzer; + sym::Analyzer analyzer; }; namespace transform { diff --git a/src/s_tir/transform/unify_thread_binding.cc b/src/s_tir/transform/unify_thread_binding.cc index 658b9fe77a47..5cc70997f8f9 100644 --- a/src/s_tir/transform/unify_thread_binding.cc +++ b/src/s_tir/transform/unify_thread_binding.cc @@ -21,13 +21,13 @@ * \file unify_thread_binding.cc */ -#include #include #include #include #include #include #include +#include #include #include "../../support/utils.h" @@ -186,7 +186,7 @@ class ThreadBindingUnifier : public StmtExprMutator { /*! \brief A integer counter storing the depth of thread bindings of "blockIdx.x/y/z" */ int thread_block_depth_ = 0; /*! \brief An analyzer used for equality proof */ - arith::Analyzer ana; + sym::Analyzer ana; }; namespace transform { 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 95cc2fbfff36..0bb47ea3f7bd 100644 --- a/src/s_tir/transform/using_assume_to_reduce_branches.cc +++ b/src/s_tir/transform/using_assume_to_reduce_branches.cc @@ -48,14 +48,14 @@ #include -#include "../../arith/constraint_extract.h" #include "../../s_tir/ir/ir_mutator_with_analyzer.h" +#include "../../sym/constraint_extract.h" #include "tvm/ir/expr.h" namespace tvm { namespace s_tir { using namespace tvm::tirx; -using namespace arith; +using namespace sym; class AssumeChecker : public StmtExprVisitor { public: @@ -167,7 +167,7 @@ class ParseAssumeAndOvercompute : public IRMutatorWithAnalyzer { } ParseAssumeAndOvercompute* self{nullptr}; - With analyzer_context; + With analyzer_context; size_t old_num_constraints{0}; size_t new_num_constraints{0}; ffi::Optional assume{std::nullopt}; @@ -302,7 +302,7 @@ class ParseAssumeAndOvercompute : public IRMutatorWithAnalyzer { } void Assume(PrimExpr assumption) { - for (const auto& expr : arith::ExtractConstraints(assumption, false)) { + for (const auto& expr : sym::ExtractConstraints(assumption, false)) { AssumeConstraintComponent(expr); } } @@ -312,7 +312,7 @@ class ParseAssumeAndOvercompute : public IRMutatorWithAnalyzer { assume_struct buf_data; std::vector buffer_exprs; - for (const auto& expr : arith::ExtractComponents(assumption)) { + for (const auto& expr : sym::ExtractComponents(assumption)) { auto side_effect = SideEffect(expr); if (side_effect <= CallEffectKind::kPure) { // Pulling out portions of the assumption that do not depend @@ -389,7 +389,7 @@ namespace transform { Pass UseAssumeToReduceBranches() { auto pass_func = [](PrimFunc f, IRModule m, PassContext ctx) { auto* n = f.CopyOnWrite(); - arith::Analyzer analyzer; + sym::Analyzer analyzer; // The pass runs & eliminates pad branch with overcompute only if, // the primfunc has op_pattern defined and is an elementwise op. diff --git a/src/arith/analyzer.cc b/src/sym/analyzer.cc similarity index 88% rename from src/arith/analyzer.cc rename to src/sym/analyzer.cc index 8912c01c2917..c49bde37f29d 100644 --- a/src/arith/analyzer.cc +++ b/src/sym/analyzer.cc @@ -18,21 +18,21 @@ */ /*! - * \file tvm/arith/analyzer.cc + * \file tvm/sym/analyzer.cc */ -#include #include #include #include #include #include #include +#include #include "const_fold.h" #include "product_normal_form.h" namespace tvm { -namespace arith { +namespace sym { AnalyzerObj::AnalyzerObj() : const_int_bound(this), @@ -287,41 +287,41 @@ TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; refl::ObjectDef(); refl::GlobalDef() - .def("arith.Analyzer", []() { return Analyzer(); }) - .def("arith.AnalyzerClone", [](Analyzer analyzer) { return analyzer->Clone(); }) - .def("arith.EnterZ3ContextScope", []() { EnterZ3ContextScope(); }) - .def("arith.ExitZ3ContextScope", []() { ExitZ3ContextScope(); }) - .def("arith.AnalyzerConstIntBound", + .def("sym.Analyzer", []() { return Analyzer(); }) + .def("sym.AnalyzerClone", [](Analyzer analyzer) { return analyzer->Clone(); }) + .def("sym.EnterZ3ContextScope", []() { EnterZ3ContextScope(); }) + .def("sym.ExitZ3ContextScope", []() { ExitZ3ContextScope(); }) + .def("sym.AnalyzerConstIntBound", [](Analyzer analyzer, const PrimExpr& expr) { return analyzer->const_int_bound(expr); }) - .def("arith.AnalyzerConstIntBoundUpdate", + .def("sym.AnalyzerConstIntBoundUpdate", [](Analyzer analyzer, const Var& var, const ConstIntBound& info, bool allow_override) { analyzer->const_int_bound.Update(var, info, allow_override); }) - .def("arith.AnalyzerConstIntBoundIsBound", + .def("sym.AnalyzerConstIntBoundIsBound", [](Analyzer analyzer, const Var& var) { return analyzer->const_int_bound.IsBound(var); }) - .def("arith.AnalyzerModularSetUpdate", + .def("sym.AnalyzerModularSetUpdate", [](Analyzer analyzer, const Var& var, const ModularSet& info, bool allow_override) { analyzer->modular_set.Update(var, info, allow_override); }) - .def("arith.AnalyzerIntSetUpdate", + .def("sym.AnalyzerIntSetUpdate", [](Analyzer analyzer, const Var& var, const IntSet& info, bool allow_override) { analyzer->int_set.Update(var, info, allow_override); }) - .def("arith.AnalyzerModularSet", + .def("sym.AnalyzerModularSet", [](Analyzer analyzer, const PrimExpr& expr) { return analyzer->modular_set(expr); }) - .def("arith.AnalyzerSimplify", [](Analyzer analyzer, const PrimExpr& expr, - int steps) { return analyzer->Simplify(expr, steps); }) - .def("arith.AnalyzerRewriteSimplify", + .def("sym.AnalyzerSimplify", [](Analyzer analyzer, const PrimExpr& expr, + int steps) { return analyzer->Simplify(expr, steps); }) + .def("sym.AnalyzerRewriteSimplify", [](Analyzer analyzer, const PrimExpr& expr) { return analyzer->rewrite_simplify(expr); }) - .def("arith.AnalyzerGetRewriteSimplifyStats", + .def("sym.AnalyzerGetRewriteSimplifyStats", [](Analyzer analyzer) { return analyzer->rewrite_simplify.GetStatsCounters(); }) - .def("arith.AnalyzerResetRewriteSimplifyStats", + .def("sym.AnalyzerResetRewriteSimplifyStats", [](Analyzer analyzer) { analyzer->rewrite_simplify.ResetStatsCounters(); }) - .def("arith.AnalyzerCanonicalSimplify", + .def("sym.AnalyzerCanonicalSimplify", [](Analyzer analyzer, const PrimExpr& expr) { return analyzer->canonical_simplify(expr); }) - .def("arith.AnalyzerIntSet", + .def("sym.AnalyzerIntSet", [](Analyzer analyzer, const PrimExpr& expr, ffi::Optional> opt_dom_map) { if (opt_dom_map.has_value()) { @@ -329,7 +329,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { } return analyzer->int_set(expr); }) - .def_packed("arith.AnalyzerBind", + .def_packed("sym.AnalyzerBind", [](ffi::PackedArgs args, ffi::Any* ret) { TVM_FFI_ICHECK(args.size() == 3 || args.size() == 4) << "AnalyzerBind expects 3 or 4 arguments, but got " << args.size(); @@ -341,23 +341,23 @@ TVM_FFI_STATIC_INIT_BLOCK() { analyzer->Bind(args[1].cast(), args[2].cast(), allow_override); } }) - .def("arith.AnalyzerCanProve", + .def("sym.AnalyzerCanProve", [](Analyzer analyzer, const PrimExpr& expr, int strength) { return analyzer->CanProve(expr, static_cast(strength)); }) - .def("arith.EnterAllowUintAsIndex", []() { ++arith::uint_as_index::g_depth; }) - .def("arith.ExitAllowUintAsIndex", + .def("sym.EnterAllowUintAsIndex", []() { ++sym::uint_as_index::g_depth; }) + .def("sym.ExitAllowUintAsIndex", []() { - TVM_FFI_ICHECK(arith::uint_as_index::g_depth > 0) + TVM_FFI_ICHECK(sym::uint_as_index::g_depth > 0) << "ExitAllowUintAsIndex without a matching Enter"; - --arith::uint_as_index::g_depth; + --sym::uint_as_index::g_depth; }) - .def("arith.GetAllowUintAsIndex", []() { return arith::uint_as_index::Enabled(); }) - .def("arith.AnalyzerSetMaximumRewriteSteps", + .def("sym.GetAllowUintAsIndex", []() { return sym::uint_as_index::Enabled(); }) + .def("sym.AnalyzerSetMaximumRewriteSteps", [](Analyzer analyzer, int64_t maximum) { analyzer->rewrite_simplify.SetMaximumRewriteSteps(maximum); }) - .def("arith.AnalyzerEnterConstraintContext", + .def("sym.AnalyzerEnterConstraintContext", [](Analyzer analyzer, const PrimExpr& constraint) { // can't use make_shared due to noexcept(false) decl in destructor, // see https://stackoverflow.com/a/43907314 @@ -366,38 +366,38 @@ TVM_FFI_STATIC_INIT_BLOCK() { auto fexit = [ctx](ffi::PackedArgs, ffi::Any*) mutable { ctx.reset(); }; return ffi::Function::FromPacked(fexit); }) - .def_method("arith.AnalyzerCanProveEqual", &AnalyzerObj::CanProveEqual) - .def("arith.AnalyzerTryCompare", + .def_method("sym.AnalyzerCanProveEqual", &AnalyzerObj::CanProveEqual) + .def("sym.AnalyzerTryCompare", [](Analyzer analyzer, const PrimExpr& lhs, const PrimExpr& rhs, bool propagate_inequalities) { return static_cast( analyzer->transitive_comparisons.TryCompare(lhs, rhs, propagate_inequalities)); }) - .def("arith.AnalyzerIsZ3Enabled", + .def("sym.AnalyzerIsZ3Enabled", [](Analyzer analyzer) { return analyzer->z3_prover.IsEnabled(); }) - .def("arith.AnalyzerGetSMTLIB2", + .def("sym.AnalyzerGetSMTLIB2", [](Analyzer analyzer, ffi::Optional expr) { return analyzer->z3_prover.GetSMTLIB2(expr); }) - .def("arith.AnalyzerSetZ3TimeoutMs", + .def("sym.AnalyzerSetZ3TimeoutMs", [](Analyzer analyzer, int64_t timeout_ms) { analyzer->z3_prover.SetTimeoutMs(static_cast(timeout_ms)); }) - .def("arith.AnalyzerSetZ3RLimit", + .def("sym.AnalyzerSetZ3RLimit", [](Analyzer analyzer, int64_t rlimit) { analyzer->z3_prover.SetRLimit(static_cast(rlimit)); }) - .def("arith.AnalyzerGetZ3Stats", + .def("sym.AnalyzerGetZ3Stats", [](Analyzer analyzer) { return analyzer->z3_prover.GetStats(); }) - .def("arith.AnalyzerGetEnabledExtensions", + .def("sym.AnalyzerGetEnabledExtensions", [](Analyzer analyzer) { return static_cast(analyzer->rewrite_simplify.GetEnabledExtensions()); }) - .def("arith.AnalyzerSetEnabledExtensions", [](Analyzer analyzer, int64_t flags) { + .def("sym.AnalyzerSetEnabledExtensions", [](Analyzer analyzer, int64_t flags) { analyzer->rewrite_simplify.SetEnabledExtensions( static_cast(flags)); }); } -} // namespace arith +} // namespace sym } // namespace tvm diff --git a/src/arith/bound_deducer.cc b/src/sym/bound_deducer.cc similarity index 99% rename from src/arith/bound_deducer.cc rename to src/sym/bound_deducer.cc index 59295ca350a9..248963655037 100644 --- a/src/arith/bound_deducer.cc +++ b/src/sym/bound_deducer.cc @@ -21,12 +21,12 @@ * \file bound_deducer.cc * \brief Utility to deduce bound of expression */ -#include #include #include #include #include #include +#include #include #include @@ -35,7 +35,7 @@ #include "interval_set.h" namespace tvm { -namespace arith { +namespace sym { // Find a target path through structural expression fields, including dynamic types. // BoundDeduceInputChecker counts occurrences over the same broader domain and can @@ -406,12 +406,12 @@ IntSet DeduceBound(PrimExpr v, PrimExpr e, const ffi::Map& hint_map TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - refl::GlobalDef().def("arith.DeduceBound", + refl::GlobalDef().def("sym.DeduceBound", [](PrimExpr v, PrimExpr cond, const ffi::Map hint_map, const ffi::Map relax_map) { return DeduceBound(v, cond, hint_map, relax_map); }); } -} // namespace arith +} // namespace sym } // namespace tvm diff --git a/src/arith/canonical_simplify.cc b/src/sym/canonical_simplify.cc similarity index 99% rename from src/arith/canonical_simplify.cc rename to src/sym/canonical_simplify.cc index ca8f05ce7649..c8205226582d 100644 --- a/src/arith/canonical_simplify.cc +++ b/src/sym/canonical_simplify.cc @@ -21,13 +21,13 @@ * \file canonical_simplify.cc * \brief Canonical form based simplification. */ -#include #include #include #include #include #include #include +#include #include "const_fold.h" #include "pattern_match.h" @@ -35,7 +35,7 @@ #include "rewrite_simplify.h" namespace tvm { -namespace arith { +namespace sym { class SumExpr; class SplitExpr; @@ -56,18 +56,18 @@ class CanonicalExprNode : public ExprNode { virtual PrimExpr Normalize() const = 0; static constexpr const uint32_t _type_child_slots = 2; - TVM_FFI_DECLARE_OBJECT_INFO("arith.CanonicalExpr", CanonicalExprNode, ExprNode); + TVM_FFI_DECLARE_OBJECT_INFO("sym.CanonicalExpr", CanonicalExprNode, ExprNode); }; -} // namespace arith +} // namespace sym namespace ffi { template inline constexpr bool object_ref_contains_v = - std::is_base_of_v; + std::is_base_of_v; } // namespace ffi -namespace arith { +namespace sym { inline PrimExpr ModImpl(PrimExpr a, PrimExpr b, DivMode mode) { if (mode == kTruncDiv) { @@ -231,7 +231,7 @@ class SplitExprNode : public CanonicalExprNode { inline bool IndexEqual(const SplitExpr& other) const; inline bool DivModeCompatibleTo(DivMode mode) const; - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("arith.SplitExpr", SplitExprNode, CanonicalExprNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("sym.SplitExpr", SplitExprNode, CanonicalExprNode); }; class SplitExpr : public PrimExpr { @@ -425,7 +425,7 @@ class SumExprNode : public CanonicalExprNode { } this->ExprNode::ty = dtype; } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("arith.SumExpr", SumExprNode, CanonicalExprNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("sym.SumExpr", SumExprNode, CanonicalExprNode); private: /*! @@ -1464,5 +1464,5 @@ void CanonicalSimplifier::CopyFrom(const CanonicalSimplifier& other) { impl_->CopyFrom(*other.impl_); } -} // namespace arith +} // namespace sym } // namespace tvm diff --git a/src/arith/conjunctive_normal_form.cc b/src/sym/conjunctive_normal_form.cc similarity index 99% rename from src/arith/conjunctive_normal_form.cc rename to src/sym/conjunctive_normal_form.cc index 3194c0b33e67..72ff1ca1bc57 100644 --- a/src/arith/conjunctive_normal_form.cc +++ b/src/sym/conjunctive_normal_form.cc @@ -18,14 +18,14 @@ */ /*! - * \file tvm/arith/conjunctive_normal_form.cc + * \file tvm/sym/conjunctive_normal_form.cc */ #include "conjunctive_normal_form.h" -#include #include #include +#include #include #include @@ -37,7 +37,7 @@ #include "rewrite_simplify.h" namespace tvm { -namespace arith { +namespace sym { namespace { /* \brief A utility for simplifying expressions using conjunctive/disjuctive normal forms */ @@ -441,5 +441,5 @@ PrimExpr SimplifyAsAndOfOrs(const PrimExpr& expr, AnalyzerObj* analyzer) { return repr.ToPrimExpr(); } -} // namespace arith +} // namespace sym } // namespace tvm diff --git a/src/arith/conjunctive_normal_form.h b/src/sym/conjunctive_normal_form.h similarity index 88% rename from src/arith/conjunctive_normal_form.h rename to src/sym/conjunctive_normal_form.h index 7c477787262a..489feef7ed11 100644 --- a/src/arith/conjunctive_normal_form.h +++ b/src/sym/conjunctive_normal_form.h @@ -23,13 +23,13 @@ * \brief Centralized location for simplifying into specific forms */ -#ifndef TVM_ARITH_CONJUNCTIVE_NORMAL_FORM_H_ -#define TVM_ARITH_CONJUNCTIVE_NORMAL_FORM_H_ +#ifndef TVM_SYM_CONJUNCTIVE_NORMAL_FORM_H_ +#define TVM_SYM_CONJUNCTIVE_NORMAL_FORM_H_ #include namespace tvm { -namespace arith { +namespace sym { class AnalyzerObj; class Analyzer; @@ -44,7 +44,7 @@ class Analyzer; */ PrimExpr SimplifyAsAndOfOrs(const PrimExpr& expr, AnalyzerObj* analyzer); -} // namespace arith +} // namespace sym } // namespace tvm -#endif // TVM_ARITH_CONJUNCTIVE_NORMAL_FORM_H_ +#endif // TVM_SYM_CONJUNCTIVE_NORMAL_FORM_H_ diff --git a/src/arith/const_fold.h b/src/sym/const_fold.h similarity index 91% rename from src/arith/const_fold.h rename to src/sym/const_fold.h index 96fb29ab35be..5da41dea44d8 100644 --- a/src/arith/const_fold.h +++ b/src/sym/const_fold.h @@ -21,17 +21,17 @@ * \file const_fold.h * \brief Arithmetic aliases for shared constant folding and scoped index analysis. */ -#ifndef TVM_ARITH_CONST_FOLD_H_ -#define TVM_ARITH_CONST_FOLD_H_ +#ifndef TVM_SYM_CONST_FOLD_H_ +#define TVM_SYM_CONST_FOLD_H_ #include "../ir/prim/const_fold.h" #include "int_operator.h" -#define TVM_ARITH_CONST_PROPAGATION(BODY) TVM_PRIM_CONST_PROPAGATION(BODY) +#define TVM_SYM_CONST_PROPAGATION(BODY) TVM_PRIM_CONST_PROPAGATION(BODY) #define TVM_INDEX_CONST_PROPAGATION(BODY) TVM_PRIM_INDEX_CONST_PROPAGATION(BODY) namespace tvm { -namespace arith { +namespace sym { using prim::detail::GetFoldResultDoubleRepr; using prim::detail::is_neg_inf; @@ -70,6 +70,6 @@ class AllowUintAsIndexGuard { AllowUintAsIndexGuard& operator=(const AllowUintAsIndexGuard&) = delete; }; -} // namespace arith +} // namespace sym } // namespace tvm -#endif // TVM_ARITH_CONST_FOLD_H_ +#endif // TVM_SYM_CONST_FOLD_H_ diff --git a/src/arith/const_int_bound.cc b/src/sym/const_int_bound.cc similarity index 99% rename from src/arith/const_int_bound.cc rename to src/sym/const_int_bound.cc index 23c45454c591..2eba4ffd35f5 100644 --- a/src/arith/const_int_bound.cc +++ b/src/sym/const_int_bound.cc @@ -18,9 +18,8 @@ */ /*! - * \file tvm/arith/const_int_bound.cc + * \file tvm/sym/const_int_bound.cc */ -#include #include #include #include @@ -28,6 +27,7 @@ #include #include #include +#include #include #include @@ -37,7 +37,7 @@ #include "pattern_match.h" namespace tvm { -namespace arith { +namespace sym { TVM_FFI_STATIC_INIT_BLOCK() { ConstIntBoundNode::RegisterReflection(); } @@ -54,7 +54,7 @@ ConstIntBound MakeConstIntBound(int64_t min_value, int64_t max_value) { TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - refl::GlobalDef().def("arith.ConstIntBound", MakeConstIntBound); + refl::GlobalDef().def("sym.ConstIntBound", MakeConstIntBound); } inline void PrintBoundValue(std::ostream& os, int64_t val) { @@ -922,5 +922,5 @@ void ConstIntBoundAnalyzer::CopyFrom(const ConstIntBoundAnalyzer& other) { impl_->CopyFrom(*other.impl_); } -} // namespace arith +} // namespace sym } // namespace tvm diff --git a/src/arith/constraint_extract.cc b/src/sym/constraint_extract.cc similarity index 94% rename from src/arith/constraint_extract.cc rename to src/sym/constraint_extract.cc index 0b1450bdb426..51d025183c4d 100644 --- a/src/arith/constraint_extract.cc +++ b/src/sym/constraint_extract.cc @@ -18,18 +18,18 @@ */ /*! - * \file tvm/arith/constraint_extract.cc + * \file tvm/sym/constraint_extract.cc */ #include "constraint_extract.h" -#include #include +#include #include "pattern_match.h" namespace tvm { -namespace arith { +namespace sym { template void CollectConstraints(PrimExpr expr, F callback, bool keep_composite_constraints) { @@ -70,5 +70,5 @@ std::vector ExtractComponents(const PrimExpr& expr) { return out; } -} // namespace arith +} // namespace sym } // namespace tvm diff --git a/src/arith/constraint_extract.h b/src/sym/constraint_extract.h similarity index 94% rename from src/arith/constraint_extract.h rename to src/sym/constraint_extract.h index c14409850f7a..fa0018936f8b 100644 --- a/src/arith/constraint_extract.h +++ b/src/sym/constraint_extract.h @@ -23,15 +23,15 @@ * \brief Centralized location for extraction of constraints from a boolean expression. */ -#ifndef TVM_ARITH_CONSTRAINT_EXTRACT_H_ -#define TVM_ARITH_CONSTRAINT_EXTRACT_H_ +#ifndef TVM_SYM_CONSTRAINT_EXTRACT_H_ +#define TVM_SYM_CONSTRAINT_EXTRACT_H_ #include #include namespace tvm { -namespace arith { +namespace sym { /* \brief Returns constraints that are true if the expression is true. * @@ -81,7 +81,7 @@ std::vector ExtractConstraints(const PrimExpr& expr, */ std::vector ExtractComponents(const PrimExpr& expr); -} // namespace arith +} // namespace sym } // namespace tvm -#endif // TVM_ARITH_CONSTRAINT_EXTRACT_H_ +#endif // TVM_SYM_CONSTRAINT_EXTRACT_H_ diff --git a/src/arith/constraint_helpers.h b/src/sym/constraint_helpers.h similarity index 95% rename from src/arith/constraint_helpers.h rename to src/sym/constraint_helpers.h index 2963b2266498..f405b87e0282 100644 --- a/src/arith/constraint_helpers.h +++ b/src/sym/constraint_helpers.h @@ -18,22 +18,22 @@ */ /*! - * \file arith/constraint_helpers.h + * \file sym/constraint_helpers.h * \brief Shared constraint derivation for arithmetic and statement simplifiers. */ -#ifndef TVM_ARITH_CONSTRAINT_HELPERS_H_ -#define TVM_ARITH_CONSTRAINT_HELPERS_H_ +#ifndef TVM_SYM_CONSTRAINT_HELPERS_H_ +#define TVM_SYM_CONSTRAINT_HELPERS_H_ -#include #include #include #include +#include #include #include namespace tvm { -namespace arith { +namespace sym { namespace detail { enum class CompareKind { kEQ, kLT, kLE, kGT, kGE }; @@ -146,7 +146,7 @@ inline void EnterConstraintFacts(WithGroup* constraints, Anal } } // namespace detail -} // namespace arith +} // namespace sym } // namespace tvm -#endif // TVM_ARITH_CONSTRAINT_HELPERS_H_ +#endif // TVM_SYM_CONSTRAINT_HELPERS_H_ diff --git a/src/arith/detect_linear_equation.cc b/src/sym/detect_linear_equation.cc similarity index 97% rename from src/arith/detect_linear_equation.cc rename to src/sym/detect_linear_equation.cc index 3b890f4b1172..17891e855540 100644 --- a/src/arith/detect_linear_equation.cc +++ b/src/sym/detect_linear_equation.cc @@ -21,18 +21,18 @@ * \file detect_linear_equation.cc * \brief Utility to detect patterns in the expression. */ -#include #include #include #include #include #include #include +#include #include namespace tvm { -namespace arith { +namespace sym { using namespace tvm::prim; // Linear equation, the components can be undefined. @@ -306,10 +306,10 @@ ffi::Array DetectClipBound(const PrimExpr& e, const ffi::Array& vars) { + .def("sym.DetectLinearEquation", DetectLinearEquation) + .def("sym.DetectClipBound", [](const PrimExpr& e, const ffi::Array& vars) { return DetectClipBound(e, vars); }); } -} // namespace arith +} // namespace sym } // namespace tvm diff --git a/src/arith/int_operator.h b/src/sym/int_operator.h similarity index 97% rename from src/arith/int_operator.h rename to src/sym/int_operator.h index 70ef74f3219e..e9e6c3307370 100644 --- a/src/arith/int_operator.h +++ b/src/sym/int_operator.h @@ -21,8 +21,8 @@ * \file int_operator.h * \brief Additional useful operators for integer. */ -#ifndef TVM_ARITH_INT_OPERATOR_H_ -#define TVM_ARITH_INT_OPERATOR_H_ +#ifndef TVM_SYM_INT_OPERATOR_H_ +#define TVM_SYM_INT_OPERATOR_H_ #include #include @@ -33,7 +33,7 @@ #include "../ir/prim/int_operator.h" namespace tvm { -namespace arith { +namespace sym { /*! * \brief Check if an integer op with operand x, y will overflow. @@ -188,6 +188,6 @@ inline int64_t LeastCommonMultiple(int64_t a, int64_t b) { return (a * b) / ExtendedEuclidean(a, b, &x, &y); } -} // namespace arith +} // namespace sym } // namespace tvm -#endif // TVM_ARITH_INT_OPERATOR_H_ +#endif // TVM_SYM_INT_OPERATOR_H_ diff --git a/src/arith/int_set.cc b/src/sym/int_set.cc similarity index 97% rename from src/arith/int_set.cc rename to src/sym/int_set.cc index 78eaeb51054b..dab779437ea2 100644 --- a/src/arith/int_set.cc +++ b/src/sym/int_set.cc @@ -21,8 +21,6 @@ * \file int_set.cc * \brief The integer set functions */ -#include -#include #include #include #include @@ -31,6 +29,8 @@ #include #include #include +#include +#include #include #include @@ -43,7 +43,7 @@ #include "pattern_match.h" namespace tvm { -namespace arith { +namespace sym { using namespace tvm::prim; using prim::is_one; @@ -65,7 +65,7 @@ IntervalSet MakeIntervalSet(PrimExpr min_value, PrimExpr max_value) { TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - refl::GlobalDef().def("arith.IntervalSet", MakeIntervalSet); + refl::GlobalDef().def("sym.IntervalSet", MakeIntervalSet); } IntervalSet Intersect(AnalyzerObj* analyzer, IntervalSet a, IntervalSet b) { @@ -671,7 +671,7 @@ class IntSetAnalyzer::Impl { // bounds implied by that condition. static std::vector> DetectBoundInfo(const PrimExpr& cond); - // The parent arith::Analyzer + // The parent sym::Analyzer AnalyzerObj* analyzer_; // Map of variables to global variable bounds (e.g. loop iterator @@ -1120,12 +1120,12 @@ ExprIntSetMap EvalSetForEachSubExpr(PrimExpr e, return m.expr_map; } -ffi::Map AsIntSet(const ffi::Map& var_dom) { - ffi::Map result; +ffi::Map AsIntSet(const ffi::Map& var_dom) { + ffi::Map result; for (auto kv : var_dom) { const Var& var = kv.first; const Range& range = kv.second; - result.Set(var, arith::IntSet::FromRange(range)); + result.Set(var, sym::IntSet::FromRange(range)); } return result; } @@ -1220,7 +1220,7 @@ ffi::Array EstimateRegionUpperBound(const ffi::Array& region, const ffi::Map& var_dom, const PrimExpr& predicate, const Analyzer& analyzer) { AnalyzerObj* analyzer_ptr = analyzer.get(); - if (ffi::Optional> result = EstimateRegionStrictBound( + if (ffi::Optional> result = EstimateRegionStrictBound( /*region=*/region, /*var_dom=*/var_dom, /*predicate=*/predicate, /*analyzer=*/analyzer)) { @@ -1265,35 +1265,35 @@ ffi::Array EstimateRegionUpperBound(const ffi::Array& region, TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; refl::GlobalDef() - .def("arith.intset_single_point", IntSet::SinglePoint) - .def("arith.intset_vector", IntSet::Vector) - .def("arith.intset_interval", IntSet::Interval) - .def_method("arith.IntervalSetGetMin", &IntSet::min) - .def_method("arith.IntervalSetGetMax", &IntSet::max) - .def_method("arith.IntSetIsNothing", &IntSet::IsNothing) - .def_method("arith.IntSetIsEverything", &IntSet::IsEverything) - .def("arith.EstimateRegionLowerBound", + .def("sym.intset_single_point", IntSet::SinglePoint) + .def("sym.intset_vector", IntSet::Vector) + .def("sym.intset_interval", IntSet::Interval) + .def_method("sym.IntervalSetGetMin", &IntSet::min) + .def_method("sym.IntervalSetGetMax", &IntSet::max) + .def_method("sym.IntSetIsNothing", &IntSet::IsNothing) + .def_method("sym.IntSetIsEverything", &IntSet::IsEverything) + .def("sym.EstimateRegionLowerBound", [](ffi::Array region, ffi::Map var_dom, PrimExpr predicate, ffi::Optional opt_analyzer) -> ffi::Optional> { Analyzer analyzer = opt_analyzer.has_value() ? opt_analyzer.value() : Analyzer(); return EstimateRegionLowerBound(region, var_dom, predicate, analyzer); }) - .def("arith.EstimateRegionStrictBound", + .def("sym.EstimateRegionStrictBound", [](ffi::Array region, ffi::Map var_dom, PrimExpr predicate, ffi::Optional opt_analyzer) -> ffi::Optional> { Analyzer analyzer = opt_analyzer.has_value() ? opt_analyzer.value() : Analyzer(); return EstimateRegionStrictBound(region, var_dom, predicate, analyzer); }) - .def("arith.EstimateRegionUpperBound", + .def("sym.EstimateRegionUpperBound", [](ffi::Array region, ffi::Map var_dom, PrimExpr predicate, ffi::Optional opt_analyzer) -> ffi::Optional> { Analyzer analyzer = opt_analyzer.has_value() ? opt_analyzer.value() : Analyzer(); return EstimateRegionUpperBound(region, var_dom, predicate, analyzer); }) - .def("arith.PosInf", []() { return SymbolicLimits::pos_inf_; }) - .def("arith.NegInf", []() { return SymbolicLimits::neg_inf_; }) - .def("arith.UnionLowerBound", UnionLowerBound); + .def("sym.PosInf", []() { return SymbolicLimits::pos_inf_; }) + .def("sym.NegInf", []() { return SymbolicLimits::neg_inf_; }) + .def("sym.UnionLowerBound", UnionLowerBound); } -} // namespace arith +} // namespace sym } // namespace tvm diff --git a/src/arith/interval_set.h b/src/sym/interval_set.h similarity index 94% rename from src/arith/interval_set.h rename to src/sym/interval_set.h index b0fa9edf0203..aac3da98b911 100644 --- a/src/arith/interval_set.h +++ b/src/sym/interval_set.h @@ -21,20 +21,20 @@ * \file int_set.h * \brief Internal data structure for integer set. */ -#ifndef TVM_ARITH_INTERVAL_SET_H_ -#define TVM_ARITH_INTERVAL_SET_H_ +#ifndef TVM_SYM_INTERVAL_SET_H_ +#define TVM_SYM_INTERVAL_SET_H_ -#include #include #include #include +#include #include #include "const_fold.h" namespace tvm { -namespace arith { +namespace sym { // Acknowledgement: IntervalSet design originates from Halide. /*! @@ -76,7 +76,7 @@ class IntervalSetNode : public IntSetNode { } /*! \return whether interval represent everything */ bool IsEverything() const { return is_neg_inf(min_value) && is_pos_inf(max_value); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("arith.IntervalSet", IntervalSetNode, IntSetNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("sym.IntervalSet", IntervalSetNode, IntSetNode); }; /*! @@ -133,7 +133,7 @@ TVM_DLL IntervalSet Union(AnalyzerObj* analyzer, IntervalSet a, IntervalSet b); */ TVM_DLL IntervalSet Intersect(AnalyzerObj* analzyer, IntervalSet a, IntervalSet b); -} // namespace arith +} // namespace sym } // namespace tvm -#endif // TVM_ARITH_INTERVAL_SET_H_ +#endif // TVM_SYM_INTERVAL_SET_H_ diff --git a/src/arith/iter_affine_map.cc b/src/sym/iter_affine_map.cc similarity index 98% rename from src/arith/iter_affine_map.cc rename to src/sym/iter_affine_map.cc index f86c008e7f9b..97bf2151b274 100644 --- a/src/arith/iter_affine_map.cc +++ b/src/sym/iter_affine_map.cc @@ -18,16 +18,16 @@ */ /*! - * \file src/arith/iter_affine_map.cc + * \file src/sym/iter_affine_map.cc */ -#include -#include #include #include #include #include #include #include +#include +#include #include #include @@ -38,7 +38,7 @@ #include "product_normal_form.h" namespace tvm { -namespace arith { +namespace sym { using namespace tvm::prim; TVM_FFI_STATIC_INIT_BLOCK() { @@ -57,7 +57,7 @@ IterMark::IterMark(PrimExpr source, PrimExpr extent) { TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - refl::GlobalDef().def("arith.IterMark", + refl::GlobalDef().def("sym.IterMark", [](PrimExpr source, PrimExpr extent) { return IterMark(source, extent); }); } @@ -98,8 +98,8 @@ IterSplitExpr::IterSplitExpr(IterMark source, PrimExpr lower_factor, PrimExpr ex TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - refl::GlobalDef().def("arith.IterSplitExpr", [](IterMark source, PrimExpr lower_factor, - PrimExpr extent, PrimExpr scale) { + refl::GlobalDef().def("sym.IterSplitExpr", [](IterMark source, PrimExpr lower_factor, + PrimExpr extent, PrimExpr scale) { return IterSplitExpr(source, lower_factor, extent, scale); }); } @@ -116,7 +116,7 @@ IterSumExpr::IterSumExpr(ffi::Array args, PrimExpr base) { TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - refl::GlobalDef().def("arith.IterSumExpr", [](ffi::Array args, PrimExpr base) { + refl::GlobalDef().def("sym.IterSumExpr", [](ffi::Array args, PrimExpr base) { return IterSumExpr(args, base); }); } @@ -1382,7 +1382,7 @@ struct IterConstraint { */ bool MatchBoundConstraints(PrimExpr pred, ffi::Map* input_iters, std::vector* result) { - arith::PVar lhs, rhs, rest; + sym::PVar lhs, rhs, rest; std::unordered_set input_iter_nodes; for (const auto& [var, _] : *input_iters) { input_iter_nodes.insert(var.get()); @@ -1465,7 +1465,7 @@ bool MatchBoundConstraints(PrimExpr pred, ffi::Map* input_iters, } }; f_extract(sum_parts, true); - arith::Analyzer analyzer; + sym::Analyzer analyzer; lhs_expr = analyzer->Simplify(lhs_expr); rhs_expr = analyzer->Simplify(rhs_expr); } @@ -1534,9 +1534,9 @@ bool IterRangeSanityCheck(const ffi::Map& iter_ranges) { IterMapResult DetectIterMap(const ffi::Array& indices, const ffi::Map& input_iters, const PrimExpr& predicate, - IterMapLevel check_level, const arith::Analyzer& analyzer, + IterMapLevel check_level, const sym::Analyzer& analyzer, bool simplify_trivial_iterators) { - arith::AnalyzerObj* analyzer_ptr = analyzer.get(); + sym::AnalyzerObj* analyzer_ptr = analyzer.get(); IterMapResult result; // Overall detection algorithm is divided into two steps: @@ -1623,7 +1623,7 @@ IterMapResult DetectIterMap(const ffi::Array& indices, TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; refl::GlobalDef().def( - "arith.DetectIterMap", + "sym.DetectIterMap", [](const ffi::Array& indices, const ffi::Map& input_iters, const PrimExpr& input_pred, int check_level, bool simplify_trivial_iterators, ffi::Optional opt_analyzer) { @@ -1634,8 +1634,8 @@ TVM_FFI_STATIC_INIT_BLOCK() { } IterSumExpr NormalizeToIterSum(PrimExpr index, const ffi::Map& input_iters, - const arith::Analyzer& analyzer) { - arith::AnalyzerObj* analyzer_ptr = analyzer.get(); + const sym::Analyzer& analyzer) { + sym::AnalyzerObj* analyzer_ptr = analyzer.get(); IterMapResult result; TVM_FFI_ICHECK(IterRangeSanityCheck(input_iters)) << "Invalid iterators. Iterators may not be expressions of each other."; @@ -1653,8 +1653,8 @@ IterSumExpr NormalizeToIterSum(PrimExpr index, const ffi::Map& i TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; refl::GlobalDef().def( - "arith.NormalizeToIterSum", [](PrimExpr index, const ffi::Map& input_iters, - ffi::Optional opt_analyzer) { + "sym.NormalizeToIterSum", [](PrimExpr index, const ffi::Map& input_iters, + ffi::Optional opt_analyzer) { Analyzer ana = opt_analyzer.has_value() ? opt_analyzer.value() : Analyzer(); return NormalizeToIterSum(index, input_iters, ana); }); @@ -2297,21 +2297,21 @@ bool IterMapRewriter::CanProveDivisible(const PrimExpr& lhs, const PrimExpr& rhs } PrimExpr NormalizeIterMapToExpr(const PrimExpr& expr) { - arith::Analyzer analyzer; + sym::Analyzer analyzer; auto normalizer = ffi::make_object(analyzer.get()); return normalizer->Mutate(expr).ValueOrUnchanged(expr); } TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - refl::GlobalDef().def("arith.NormalizeIterMapToExpr", NormalizeIterMapToExpr); + refl::GlobalDef().def("sym.NormalizeIterMapToExpr", NormalizeIterMapToExpr); } ffi::Array IterMapSimplify(const ffi::Array& indices, const ffi::Map& input_iters, const PrimExpr& input_pred, IterMapLevel check_level, - const arith::Analyzer& ana, bool simplify_trivial_iterators) { - arith::AnalyzerObj* ana_ptr = ana.get(); + const sym::Analyzer& ana, bool simplify_trivial_iterators) { + sym::AnalyzerObj* ana_ptr = ana.get(); if (!IterRangeSanityCheck(input_iters)) return indices; auto res = DetectIterMap(indices, input_iters, input_pred, check_level, ana, /*simplify_trivial_iterators=*/simplify_trivial_iterators); @@ -2347,7 +2347,7 @@ ffi::Array IterMapSimplify(const ffi::Array& indices, TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; refl::GlobalDef().def( - "arith.IterMapSimplify", + "sym.IterMapSimplify", [](const ffi::Array& indices, const ffi::Map& input_iters, const PrimExpr& input_pred, int check_level, bool simplify_trivial_iterators, ffi::Optional opt_analyzer) { @@ -2658,9 +2658,9 @@ ffi::Array> SubspaceDivide(const ffi::Array& bind const ffi::Map& input_iters, const ffi::Array& sub_iters, const PrimExpr& predicate, IterMapLevel check_level, - const arith::Analyzer& analyzer, + const sym::Analyzer& analyzer, bool simplify_trivial_iterators) { - arith::AnalyzerObj* analyzer_ptr = analyzer.get(); + sym::AnalyzerObj* analyzer_ptr = analyzer.get(); if (!IterRangeSanityCheck(input_iters)) return ffi::Array>(); auto res = DetectIterMap(bindings, input_iters, predicate, check_level, analyzer, simplify_trivial_iterators); @@ -2692,7 +2692,7 @@ ffi::Array> SubspaceDivide(const ffi::Array& bind TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; refl::GlobalDef().def( - "arith.SubspaceDivide", + "sym.SubspaceDivide", [](const ffi::Array& bindings, const ffi::Map& root_iters, const ffi::Array& sub_iters, const PrimExpr& predicate, int check_level, bool simplify_trivial_iterators, ffi::Optional opt_analyzer) { @@ -2835,8 +2835,8 @@ ffi::Map InverseAffineIterMap(const ffi::Array& iter TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - refl::GlobalDef().def("arith.InverseAffineIterMap", InverseAffineIterMap); + refl::GlobalDef().def("sym.InverseAffineIterMap", InverseAffineIterMap); } -} // namespace arith +} // namespace sym } // namespace tvm diff --git a/src/arith/modular_set.cc b/src/sym/modular_set.cc similarity index 98% rename from src/arith/modular_set.cc rename to src/sym/modular_set.cc index 7b3bac363ede..6cdf85999644 100644 --- a/src/arith/modular_set.cc +++ b/src/sym/modular_set.cc @@ -21,12 +21,12 @@ * \file modular_set.cc * \brief Modular set analysis */ -#include #include #include #include #include #include +#include #include #include @@ -35,7 +35,7 @@ #include "pattern_match.h" namespace tvm { -namespace arith { +namespace sym { using namespace tvm::prim; TVM_FFI_STATIC_INIT_BLOCK() { ModularSetNode::RegisterReflection(); } @@ -49,13 +49,13 @@ ModularSet::ModularSet(int64_t coeff, int64_t base) { } // Pattern A (RM): auto-default repr from reflection produces -// "arith.ModularSet(coeff=..., base=...)" +// "sym.ModularSet(coeff=..., base=...)" ModularSet MakeModularSet(int64_t coeff, int64_t base) { return ModularSet(coeff, base); } TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - refl::GlobalDef().def("arith.ModularSet", MakeModularSet); + refl::GlobalDef().def("sym.ModularSet", MakeModularSet); } // internal entry for const int bound @@ -420,5 +420,5 @@ void ModularSetAnalyzer::CopyFrom(const ModularSetAnalyzer& other) { impl_->CopyFrom(*other.impl_); } -} // namespace arith +} // namespace sym } // namespace tvm diff --git a/src/arith/pattern_match.h b/src/sym/pattern_match.h similarity index 98% rename from src/arith/pattern_match.h rename to src/sym/pattern_match.h index 674a439758d6..9f6c9913db5b 100644 --- a/src/arith/pattern_match.h +++ b/src/sym/pattern_match.h @@ -31,7 +31,7 @@ * \code * * // max(x + z, y + z) => max(x, y) + z - * arith::PVar x, y, z; + * sym::PVar x, y, z; * * // The following code tries to match the declared pattern. * // Match will fill the result of match into PVar if successful. @@ -45,8 +45,8 @@ * } * * tvm::Var tx, ty; - * arith::PVar c; - * arith::PVar v; + * sym::PVar c; + * sym::PVar v; * // We can match integer and Var, both of which are * // special case container of Expr * TVM_FFI_ICHECK((v * c).Match(tx * 3)); @@ -62,8 +62,8 @@ * Please be aware that the filled value in a PVar * can be overriden in the next call to Match. */ -#ifndef TVM_ARITH_PATTERN_MATCH_H_ -#define TVM_ARITH_PATTERN_MATCH_H_ +#ifndef TVM_SYM_PATTERN_MATCH_H_ +#define TVM_SYM_PATTERN_MATCH_H_ #include #include @@ -76,7 +76,7 @@ #include "const_fold.h" namespace tvm { -namespace arith { +namespace sym { /*! * \brief Base class of all the patterns. * @@ -244,7 +244,7 @@ class PVar : public Pattern> { * \tparam T the type of the hole. */ template -class PVarWithCheck : public arith::Pattern> { +class PVarWithCheck : public sym::Pattern> { public: // Store by reference in the expression. using Nested = const PVarWithCheck&; @@ -269,7 +269,7 @@ class PVarWithCheck : public arith::Pattern> { T Eval() const { return pvar_.Eval(); } protected: - arith::PVar pvar_; + sym::PVar pvar_; }; /*! @@ -920,6 +920,6 @@ inline std::enable_if_t<(std::is_base_of_v, TPattern> && ... & matches_one_of(const TPattern&... patterns) { return PMatchesOneOf(patterns...); } -} // namespace arith +} // namespace sym } // namespace tvm -#endif // TVM_ARITH_PATTERN_MATCH_H_ +#endif // TVM_SYM_PATTERN_MATCH_H_ diff --git a/src/arith/presburger_set.cc b/src/sym/presburger_set.cc similarity index 98% rename from src/arith/presburger_set.cc rename to src/sym/presburger_set.cc index 6c013fc6ec2b..ce0d2a82b480 100644 --- a/src/arith/presburger_set.cc +++ b/src/sym/presburger_set.cc @@ -23,13 +23,13 @@ */ #include "presburger_set.h" -#include -#include #include #include #include #include #include +#include +#include #include #include @@ -40,7 +40,7 @@ #include "interval_set.h" namespace tvm { -namespace arith { +namespace sym { using namespace tvm::prim; #if defined(TVM_MLIR_VERSION) && TVM_MLIR_VERSION >= 150 @@ -283,10 +283,10 @@ PresburgerSet MakePresburgerSet(const PrimExpr& constraint) { return PresburgerS TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; PresburgerSetNode::RegisterReflection(); - refl::GlobalDef().def("arith.PresburgerSet", MakePresburgerSet); + refl::GlobalDef().def("sym.PresburgerSet", MakePresburgerSet); } #endif // defined(TVM_MLIR_VERSION) && TVM_MLIR_VERSION >= 150 -} // namespace arith +} // namespace sym } // namespace tvm diff --git a/src/arith/presburger_set.h b/src/sym/presburger_set.h similarity index 94% rename from src/arith/presburger_set.h rename to src/sym/presburger_set.h index a3fc96929022..88c7eba60c07 100644 --- a/src/arith/presburger_set.h +++ b/src/sym/presburger_set.h @@ -21,8 +21,8 @@ * \file presburger_set.h * \brief Integer set based on MLIR Presburger set */ -#ifndef TVM_ARITH_PRESBURGER_SET_H_ -#define TVM_ARITH_PRESBURGER_SET_H_ +#ifndef TVM_SYM_PRESBURGER_SET_H_ +#define TVM_SYM_PRESBURGER_SET_H_ #ifdef TVM_MLIR_VERSION #if TVM_MLIR_VERSION >= 150 @@ -32,10 +32,10 @@ #endif #endif -#include #include #include #include +#include #include #include @@ -43,7 +43,7 @@ #include "const_fold.h" namespace tvm { -namespace arith { +namespace sym { #ifdef TVM_MLIR_VERSION #if TVM_MLIR_VERSION >= 150 @@ -117,7 +117,7 @@ class PresburgerSetNode : public IntSetNode { return std::all_of(disjuncts.begin(), disjuncts.end(), std::mem_fn(&IntegerRelation::isIntegerEmpty)); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("arith.PresburgerSet", PresburgerSetNode, IntSetNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("sym.PresburgerSet", PresburgerSetNode, IntSetNode); private: ffi::Array vars; @@ -158,7 +158,7 @@ class PresburgerSetNode : public IntSetNode { namespace refl = tvm::ffi::reflection; refl::ObjectDef(); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("arith.PresburgerSet", PresburgerSetNode, IntSetNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("sym.PresburgerSet", PresburgerSetNode, IntSetNode); }; class PresburgerSet : public IntSet { @@ -195,7 +195,7 @@ PresburgerSet Intersect(const ffi::Array& sets); */ IntSet EvalSet(const PrimExpr& e, const PresburgerSet& set); -} // namespace arith +} // namespace sym } // namespace tvm -#endif // TVM_ARITH_PRESBURGER_SET_H_ +#endif // TVM_SYM_PRESBURGER_SET_H_ diff --git a/src/arith/product_normal_form.h b/src/sym/product_normal_form.h similarity index 94% rename from src/arith/product_normal_form.h rename to src/sym/product_normal_form.h index 2ad01de8420b..b4649244b57b 100644 --- a/src/arith/product_normal_form.h +++ b/src/sym/product_normal_form.h @@ -21,14 +21,14 @@ * \file product_normal_form.h * \brief Centralized location related to simplifying prod of results. */ -#ifndef TVM_ARITH_PRODUCT_NORMAL_FORM_H_ -#define TVM_ARITH_PRODUCT_NORMAL_FORM_H_ +#ifndef TVM_SYM_PRODUCT_NORMAL_FORM_H_ +#define TVM_SYM_PRODUCT_NORMAL_FORM_H_ #include #include namespace tvm { -namespace arith { +namespace sym { /*! * \brief Unpack reduction by calling each leaf via fleaf @@ -96,6 +96,6 @@ inline PrimExpr MulAndNormalize(const PrimExpr& lhs, const PrimExpr& rhs) { return res; } -} // namespace arith +} // namespace sym } // namespace tvm -#endif // TVM_ARITH_PRODUCT_NORMAL_FORM_H_ +#endif // TVM_SYM_PRODUCT_NORMAL_FORM_H_ diff --git a/src/arith/rewrite_simplify.cc b/src/sym/rewrite_simplify.cc similarity index 99% rename from src/arith/rewrite_simplify.cc rename to src/sym/rewrite_simplify.cc index 8f565b9df8d7..a340654c9af2 100644 --- a/src/arith/rewrite_simplify.cc +++ b/src/sym/rewrite_simplify.cc @@ -24,12 +24,12 @@ // Acknowledgement: Most rewrite-rules are from Halide. #include "rewrite_simplify.h" -#include #include #include #include #include #include +#include #include #include @@ -41,7 +41,7 @@ #include "pattern_match.h" namespace tvm { -namespace arith { +namespace sym { namespace { TVM_FFI_INLINE bool IsVectorExpr(const ExprNode* expr) { @@ -1198,27 +1198,27 @@ UnchangedOr RewriteSimplifier::Impl::Mutate_(const prim::FloorDivNode* if (uint_as_index::Enabled()) { if (floordiv(x * c1, c2).Match(ret) && c2.Eval()->value > 0 && c1.Eval()->value % c2.Eval()->value == 0) { - LOG(WARNING) << "arith: no-overflow floordiv rule on unsigned expr: " << ret; + LOG(WARNING) << "sym: no-overflow floordiv rule on unsigned expr: " << ret; return (x * floordiv(c1, c2)).Eval(); } if (floordiv(x * c1 + y, c2).Match(ret) && c2.Eval()->value > 0 && c1.Eval()->value % c2.Eval()->value == 0) { - LOG(WARNING) << "arith: no-overflow floordiv rule on unsigned expr: " << ret; + LOG(WARNING) << "sym: no-overflow floordiv rule on unsigned expr: " << ret; return (x * floordiv(c1, c2) + floordiv(y, c2)).Eval(); } if (floordiv(y + x * c1, c2).Match(ret) && c2.Eval()->value > 0 && c1.Eval()->value % c2.Eval()->value == 0) { - LOG(WARNING) << "arith: no-overflow floordiv rule on unsigned expr: " << ret; + LOG(WARNING) << "sym: no-overflow floordiv rule on unsigned expr: " << ret; return (x * floordiv(c1, c2) + floordiv(y, c2)).Eval(); } if (floordiv(x * c1 + y + z, c2).Match(ret) && c2.Eval()->value > 0 && c1.Eval()->value % c2.Eval()->value == 0) { - LOG(WARNING) << "arith: no-overflow floordiv rule on unsigned expr: " << ret; + LOG(WARNING) << "sym: no-overflow floordiv rule on unsigned expr: " << ret; return (x * floordiv(c1, c2) + floordiv(y + z, c2)).Eval(); } if (floordiv(y + x * c1 + z, c2).Match(ret) && c2.Eval()->value > 0 && c1.Eval()->value % c2.Eval()->value == 0) { - LOG(WARNING) << "arith: no-overflow floordiv rule on unsigned expr: " << ret; + LOG(WARNING) << "sym: no-overflow floordiv rule on unsigned expr: " << ret; return (x * floordiv(c1, c2) + floordiv(y + z, c2)).Eval(); } } @@ -1355,7 +1355,7 @@ UnchangedOr RewriteSimplifier::Impl::Mutate_(const prim::FloorModNode* // When c1 is a multiple of c2, x * c1 is too — but the identity is only // sound when c2 | 2^bits: otherwise the wraparound subtraction of // k*2^bits changes residues mod c2. (No non-pow2 c2 occurs in practice; - // probed across the arith/copy/kernel test suites.) + // probed across the sym/copy/kernel test suites.) TVM_TRY_REWRITE_IF(floormod(x * c1, c2), ZeroWithTypeLike(x), c2.Eval()->value != 0 && c1.Eval()->value % c2.Eval()->value == 0 && is_pow2_le_bits(c2.Eval()->value)); @@ -1379,13 +1379,13 @@ UnchangedOr RewriteSimplifier::Impl::Mutate_(const prim::FloorModNode* if (uint_as_index::Enabled()) { if (floormod(x * c1, c2).Match(ret) && c2.Eval()->value != 0 && c1.Eval()->value % c2.Eval()->value == 0) { - LOG(WARNING) << "arith: no-overflow floormod rule on unsigned expr: " << ret; + LOG(WARNING) << "sym: no-overflow floormod rule on unsigned expr: " << ret; return ZeroWithTypeLike(x).Eval(); } if (floormod(x, c2).Match(ret) && c2.Eval()->value > 0) { ModularSet mod = analyzer_->modular_set(x.Eval()); if (mod->coeff % c2.Eval()->value == 0) { - LOG(WARNING) << "arith: no-overflow floormod rule on unsigned expr: " << ret; + LOG(WARNING) << "sym: no-overflow floormod rule on unsigned expr: " << ret; return IntImm(c2.Eval().ty(), floormod(mod->base, c2.Eval()->value)); } } @@ -2607,5 +2607,5 @@ void RewriteSimplifier::CopyFrom(const RewriteSimplifier& other) { impl_->CopyFr // Pattern A (RM): auto-default repr from reflection. -} // namespace arith +} // namespace sym } // namespace tvm diff --git a/src/arith/rewrite_simplify.h b/src/sym/rewrite_simplify.h similarity index 97% rename from src/arith/rewrite_simplify.h rename to src/sym/rewrite_simplify.h index f3eb21846e96..e350840b4899 100644 --- a/src/arith/rewrite_simplify.h +++ b/src/sym/rewrite_simplify.h @@ -21,13 +21,13 @@ * \file rewrite_simplify.h * \brief Rewrite-rule based simplification. */ -#ifndef TVM_ARITH_REWRITE_SIMPLIFY_H_ -#define TVM_ARITH_REWRITE_SIMPLIFY_H_ +#ifndef TVM_SYM_REWRITE_SIMPLIFY_H_ +#define TVM_SYM_REWRITE_SIMPLIFY_H_ -#include #include #include #include +#include #include #include @@ -38,7 +38,7 @@ #include "simplify_base.h" namespace tvm { -namespace arith { +namespace sym { /* \brief Usage counters for RewriteSimplifier * @@ -63,7 +63,7 @@ struct RewriteSimplifierStatsNode : ffi::Object { .def_ro("max_recursive_depth", &RewriteSimplifierStatsNode::max_recursive_depth) .def_ro("num_recursive_rewrites", &RewriteSimplifierStatsNode::num_recursive_rewrites); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("arith.RewriteSimplifierStats", RewriteSimplifierStatsNode, + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("sym.RewriteSimplifierStats", RewriteSimplifierStatsNode, ffi::Object); }; @@ -278,6 +278,6 @@ class RewriteSimplifier::Impl : public SimplifierBase { } }; -} // namespace arith +} // namespace sym } // namespace tvm -#endif // TVM_ARITH_REWRITE_SIMPLIFY_H_ +#endif // TVM_SYM_REWRITE_SIMPLIFY_H_ diff --git a/src/arith/simplify_base.cc b/src/sym/simplify_base.cc similarity index 99% rename from src/arith/simplify_base.cc rename to src/sym/simplify_base.cc index 3427c7f6b0cd..8478bc0c73b0 100644 --- a/src/arith/simplify_base.cc +++ b/src/sym/simplify_base.cc @@ -26,7 +26,7 @@ #include "constraint_helpers.h" namespace tvm { -namespace arith { +namespace sym { using detail::EnterConstraintFacts; @@ -163,5 +163,5 @@ UnchangedOr SimplifierBase::Mutate_(const prim::SelectNode* op, return prim::Select(cond, true_value, false_value, op->span); } -} // namespace arith +} // namespace sym } // namespace tvm diff --git a/src/arith/simplify_base.h b/src/sym/simplify_base.h similarity index 90% rename from src/arith/simplify_base.h rename to src/sym/simplify_base.h index eaafa1fa0676..d7b00caea6f8 100644 --- a/src/arith/simplify_base.h +++ b/src/sym/simplify_base.h @@ -19,20 +19,20 @@ /*! * \file simplify_base.h - * \brief Expression mutator base class for arith simplifiers. + * \brief Expression mutator base class for sym simplifiers. */ -#ifndef TVM_ARITH_SIMPLIFY_BASE_H_ -#define TVM_ARITH_SIMPLIFY_BASE_H_ +#ifndef TVM_SYM_SIMPLIFY_BASE_H_ +#define TVM_SYM_SIMPLIFY_BASE_H_ -#include #include #include #include +#include #include namespace tvm { -namespace arith { +namespace sym { /*! * \brief Arithmetic semantic-operand traversal and analyzer constraint handling. @@ -68,6 +68,6 @@ class SimplifierBase : public tvm::ExprMutator { ScopeStack> constraint_scope_; }; -} // namespace arith +} // namespace sym } // namespace tvm -#endif // TVM_ARITH_SIMPLIFY_BASE_H_ +#endif // TVM_SYM_SIMPLIFY_BASE_H_ diff --git a/src/arith/transitive_comparison_analyzer.cc b/src/sym/transitive_comparison_analyzer.cc similarity index 99% rename from src/arith/transitive_comparison_analyzer.cc rename to src/sym/transitive_comparison_analyzer.cc index cd2a21e12af8..d9bfd24e7c19 100644 --- a/src/arith/transitive_comparison_analyzer.cc +++ b/src/sym/transitive_comparison_analyzer.cc @@ -17,11 +17,11 @@ * under the License. */ /*! - * \file tvm/arith/transitive_comparison_analyzer.cc + * \file tvm/sym/transitive_comparison_analyzer.cc */ -#include #include +#include #include #include @@ -31,7 +31,7 @@ #include "pattern_match.h" namespace tvm { -namespace arith { +namespace sym { using namespace tvm::prim; using prim::is_const_int; @@ -897,5 +897,5 @@ CompareResult TransitiveComparisonAnalyzer::Impl::MergeComparisons( return result; } -} // namespace arith +} // namespace sym } // namespace tvm diff --git a/src/arith/z3_prover.cc b/src/sym/z3_prover.cc similarity index 99% rename from src/arith/z3_prover.cc rename to src/sym/z3_prover.cc index 1c1f848e71ba..d20eac26bb1f 100644 --- a/src/arith/z3_prover.cc +++ b/src/sym/z3_prover.cc @@ -18,8 +18,8 @@ */ /*! - * \file src/arith/z3_prover.cc - * \brief Optional Z3 SMT solver backend for arith::Analyzer. + * \file src/sym/z3_prover.cc + * \brief Optional Z3 SMT solver backend for sym::Analyzer. * * The real implementation is compiled only when TVM_USE_Z3 is defined (set by * the USE_Z3 CMake option). Otherwise a conservative stub is compiled so the @@ -27,7 +27,6 @@ */ #ifdef TVM_USE_Z3 -#include #include #include #include @@ -36,6 +35,7 @@ #include #include #include +#include #include #include @@ -55,7 +55,7 @@ #include "tvm/ir/expr.h" #include "z3++.h" -namespace tvm::arith { +namespace tvm::sym { using namespace tvm::prim; using namespace ffi; @@ -104,7 +104,7 @@ struct Namespace { } // namespace // Deprecated no-ops, kept only for callers of -// arith.EnterZ3ContextScope / arith.ExitZ3ContextScope (Z3ContextScope in +// sym.EnterZ3ContextScope / sym.ExitZ3ContextScope (Z3ContextScope in // Python). Shared per-compile Z3 contexts are gone: every materialized solver // owns a private context (see Z3Prover::Impl::Materialize), which subsumes the // per-compilation isolation these scopes provided. Remove together with the @@ -1022,18 +1022,18 @@ TVM_DLL int64_t Z3Prover::CountSatisfyingValues(const Var& var, int64_t max_coun Z3Prover::Z3Prover(AnalyzerObj* parent) : impl_(std::make_unique(parent)) {} TVM_DLL Z3Prover::~Z3Prover() = default; -} // namespace tvm::arith +} // namespace tvm::sym #else // TVM_USE_Z3 -#include #include #include +#include #include "tvm/ffi/string.h" #include "tvm/ir/expr.h" -namespace tvm::arith { +namespace tvm::sym { using namespace ffi; @@ -1067,6 +1067,6 @@ ffi::String Z3Prover::GetStats() { return "; Z3 Prover is disabled."; } Z3Prover::Z3Prover(AnalyzerObj*) : impl_(nullptr) {} TVM_DLL Z3Prover::~Z3Prover() = default; -} // namespace tvm::arith +} // namespace tvm::sym #endif // TVM_USE_Z3 diff --git a/src/target/llvm/codegen_cpu.cc b/src/target/llvm/codegen_cpu.cc index 195d67f5e78f..b6b8f168a5c2 100644 --- a/src/target/llvm/codegen_cpu.cc +++ b/src/target/llvm/codegen_cpu.cc @@ -534,7 +534,7 @@ void CodeGenCPU::CreateComputeScope(const AttrStmtNode* op) { llvm::DISubprogram* di_subprogram_{nullptr}; std::unordered_map var_map_; std::vector> loop_frame_jump_tgts_; - arith::Analyzer analyzer_{arith::Analyzer()}; + sym::Analyzer analyzer_{sym::Analyzer()}; CodeGenCPU* parent_; }; @@ -685,7 +685,7 @@ void CodeGenCPU::CreateParallelLaunch(const Stmt& body, int num_task, std::strin builder_->CreateInBoundsGEP(t_tvm_parallel_group_env_, penv, {ConstInt32(0), ConstInt32(1)}), "num_task"); par_env.penv = penv; - auto new_analyzer = arith::Analyzer(); + auto new_analyzer = sym::Analyzer(); std::swap(function_, f); std::swap(parallel_env_, par_env); std::swap(analyzer_, new_analyzer); diff --git a/src/target/llvm/codegen_llvm.cc b/src/target/llvm/codegen_llvm.cc index 66a24f8f9c26..79ffd820e9e2 100644 --- a/src/target/llvm/codegen_llvm.cc +++ b/src/target/llvm/codegen_llvm.cc @@ -92,7 +92,7 @@ #include #include -#include "../../arith/pattern_match.h" +#include "../../sym/pattern_match.h" #include "../build_common.h" #include "codegen_params.h" #include "llvm_instance.h" @@ -247,7 +247,7 @@ void CodeGenLLVM::InitFuncState() { alias_var_set_.clear(); alloc_storage_info_.clear(); volatile_buf_.clear(); - analyzer_ = arith::Analyzer(); + analyzer_ = sym::Analyzer(); } std::tuple CodeGenLLVM::GetLinkage( @@ -679,12 +679,12 @@ void CodeGenLLVM::AddAliasInfo(llvm::Instruction* inst, const VarNode* buffer_va } int64_t base = 0, width = 0; - arith::PVar pbase, pstride; - arith::PVar planes; + sym::PVar pbase, pstride; + sym::PVar planes; // create meta-data for alias analysis // Use a group of binary tree ranges of memory banks. int64_t xwith = 0; - if (arith::ramp(pbase, pstride, planes).Match(index)) { + if (sym::ramp(pbase, pstride, planes).Match(index)) { if (auto b = pbase.Eval()->value.as(), w = (planes.Eval()->value * pstride.Eval()->value).as(); b.has_value() && w.has_value()) { @@ -751,7 +751,7 @@ void CodeGenLLVM::GetAlignment(PrimType t, const VarNode* buf_var, const PrimExp *p_native_bits = native_vector_bits_; } - arith::ModularSet me = analyzer_->modular_set(index); + sym::ModularSet me = analyzer_->modular_set(index); int64_t base = me->base; int64_t coeff = me->coeff; diff --git a/src/target/llvm/codegen_llvm.h b/src/target/llvm/codegen_llvm.h index c3e4989a5a55..34011cc2b2f4 100644 --- a/src/target/llvm/codegen_llvm.h +++ b/src/target/llvm/codegen_llvm.h @@ -40,9 +40,9 @@ #include #include #include -#include #include #include +#include #include #include #include @@ -569,7 +569,7 @@ class CodeGenLLVM : public tirx::ExprFunctor, // Whether current function is restricted bool is_restricted_{true}; // The analyzer information - arith::Analyzer analyzer_; + sym::Analyzer analyzer_; // set of var that are not restricted(can alias) std::unordered_set alias_var_set_; // set of volatile buffer. diff --git a/src/target/source/codegen_c.cc b/src/target/source/codegen_c.cc index d9a6581136c3..d93c55eddf0c 100644 --- a/src/target/source/codegen_c.cc +++ b/src/target/source/codegen_c.cc @@ -22,16 +22,16 @@ */ #include "codegen_c.h" -#include #include #include +#include #include #include #include #include -#include "../../arith/pattern_match.h" +#include "../../sym/pattern_match.h" #include "../../tirx/ir/buffer_common.h" #include "codegen_params.h" @@ -997,11 +997,11 @@ void CodeGenC::Dispatch_(const TensorLoadNode* op, std::ostream& os) { // NOLIN } } else { bool can_vector_load = false; - arith::PVar base; - if (arith::ramp(base, 1, value_ty.lanes()).Match(index)) { + sym::PVar base; + if (sym::ramp(base, 1, value_ty.lanes()).Match(index)) { const prim::RampNode* ramp = index.as(); TVM_FFI_ICHECK(ramp); - arith::ModularSet me = arith::Analyzer()->modular_set(ramp->base); + sym::ModularSet me = sym::Analyzer()->modular_set(ramp->base); // The condition: {k * coeff + base} divisible by the alignment for any k if (me->coeff % value_ty.lanes() == 0 && me->base % value_ty.lanes() == 0) { can_vector_load = true; @@ -1062,9 +1062,9 @@ void CodeGenC::Dispatch_(const BufferStoreNode* op) { this->PrintIndent(); stream << ref << " = " << value << ";\n"; } else { - arith::PVar base; + sym::PVar base; - if (arith::ramp(base, 1, value_ty.lanes()).Match(index_expr) && + if (sym::ramp(base, 1, value_ty.lanes()).Match(index_expr) && value_ty.code() != DLDataTypeCode::kDLFloat4_e2m1fn) { std::string value = this->PrintExpr(op->value); this->PrintVecStore(op->buffer.get(), value_ty, base.Eval(), value); @@ -1379,7 +1379,7 @@ void CodeGenC::Dispatch_(const AssertStmtNode* op) { void CodeGenC::Dispatch_(const ForNode* op) { std::string begin_str = PrintExpr(op->min); - PrimExpr end = is_zero(op->min) ? op->extent : arith::Analyzer()->Simplify(op->min + op->extent); + PrimExpr end = is_zero(op->min) ? op->extent : sym::Analyzer()->Simplify(op->min + op->extent); std::string end_str = PrintExpr(end); std::string step_str = op->step.has_value() ? PrintExpr(*op->step) : ""; PrintIndent(); diff --git a/src/te/operation/compute_op.cc b/src/te/operation/compute_op.cc index 874c4ce02ec0..dc45c56d60ea 100644 --- a/src/te/operation/compute_op.cc +++ b/src/te/operation/compute_op.cc @@ -22,12 +22,12 @@ * \file compute_op.cc */ -#include #include #include #include #include #include +#include #include #include #include diff --git a/src/te/operation/create_primfunc.cc b/src/te/operation/create_primfunc.cc index ca0b2045d65f..350674c28ef0 100644 --- a/src/te/operation/create_primfunc.cc +++ b/src/te/operation/create_primfunc.cc @@ -19,7 +19,6 @@ #include "create_primfunc.h" -#include #include #include #include @@ -28,6 +27,7 @@ #include #include #include +#include #include #include #include @@ -246,7 +246,7 @@ class LayoutFreePlaceholdersNormalizer : public s_tir::StmtExprMutator { using NestedIterLevels = std::vector>; NestedIterLevels GenerateNestedIterLevels(const ffi::Array& axes, - arith::AnalyzerObj* analyzer) { + sym::AnalyzerObj* analyzer) { int global_max_depth = 0; std::unordered_map depth; std::unordered_map var2iter; @@ -424,7 +424,7 @@ Stmt GenerateInitStmt(const ffi::Array& indices, const ffi::Array& indices, const ffi::Array& buffers, const ffi::Map& var_map, PrimExpr expr_body, - CreateFuncInfo* info, arith::AnalyzerObj* analyzer) { + CreateFuncInfo* info, sym::AnalyzerObj* analyzer) { auto f_substitute = [&var_map](const Var& var) -> ffi::Expected> { if (auto repl = var_map.Get(var)) return ffi::Any(*std::move(repl)); return ffi::Unchanged(); @@ -542,7 +542,7 @@ struct NestedScopeInfo { }; Stmt GenerateStmtFromCompute(const te::ComputeOp& compute_op, CreateFuncInfo* info, - arith::AnalyzerObj* analyzer) { + sym::AnalyzerObj* analyzer) { // Step 1. Collect all iter axes in original TE compute op ffi::Array axes = compute_op->axis; axes.insert(axes.end(), compute_op->reduce_axis.begin(), compute_op->reduce_axis.end()); @@ -801,7 +801,7 @@ void InitializeBufferBinds(const ffi::Array& ordered_ops, CreateF } void RewriteStageToBlock(const te::Operation& op, CreateFuncInfo* info, - ffi::Array* root_stmts, arith::AnalyzerObj* analyzer) { + ffi::Array* root_stmts, sym::AnalyzerObj* analyzer) { if (const auto* placeholder = op.as()) { // Case 1. PlaceholderOp (te.placeholder) TVM_FFI_ICHECK_EQ(op->num_outputs(), 1); @@ -858,7 +858,7 @@ PrimFunc CreatePrimFunc(const ffi::Array& arg_list, // Root body stmts. ffi::Array root_stmts; // Analyzer - arith::Analyzer analyzer; + sym::Analyzer analyzer; // Step 1. Create ordered array of operations and validate they are supported. ffi::Array order = CollectOrderedOps(arg_list); @@ -936,7 +936,7 @@ PrimFunc CreatePrimFunc(const ffi::Array& arg_list, // Root body stmts. ffi::Array root_stmts; // Analyzer - arith::Analyzer analyzer; + sym::Analyzer analyzer; // Step 1. Create ordered array of operations and validate they are supported. ffi::Array order = CollectOrderedOps(tensor_arg_list); diff --git a/src/te/operation/extern_op.cc b/src/te/operation/extern_op.cc index 9819a4865d98..0e2f9757de7d 100644 --- a/src/te/operation/extern_op.cc +++ b/src/te/operation/extern_op.cc @@ -21,10 +21,10 @@ * \brief External computation rule. * \file extern_op.cc */ -#include #include #include #include +#include #include namespace tvm { diff --git a/src/te/operation/scan_op.cc b/src/te/operation/scan_op.cc index 4cf879eb14a8..bd2496d794f9 100644 --- a/src/te/operation/scan_op.cc +++ b/src/te/operation/scan_op.cc @@ -55,7 +55,7 @@ ScanOp::ScanOp(std::string name, std::string tag, auto n = ffi::make_object(); TVM_FFI_ICHECK_EQ(init.size(), update.size()); TVM_FFI_ICHECK_EQ(init.size(), state_placeholder.size()); - arith::Analyzer analyzer; + sym::Analyzer analyzer; auto prove_equal = [&](PrimExpr lhs, PrimExpr rhs) { return is_zero(analyzer->Simplify(lhs - rhs)); }; diff --git a/src/tirx/analysis/exec_context.cc b/src/tirx/analysis/exec_context.cc index 862483b1ae30..d47114bd6ccc 100644 --- a/src/tirx/analysis/exec_context.cc +++ b/src/tirx/analysis/exec_context.cc @@ -21,9 +21,9 @@ * \brief Compile-time active-thread state backed by TileLayout. */ -#include #include #include +#include #include #include @@ -56,7 +56,7 @@ bool TryAsInt64(const PrimExpr& expr, int64_t* value) { } bool IsZero(const PrimExpr& expr) { - arith::Analyzer analyzer; + sym::Analyzer analyzer; return analyzer->CanProveEqual(expr, 0); } diff --git a/src/tirx/analysis/filter_canonical.cc b/src/tirx/analysis/filter_canonical.cc index 427275fd92bb..eeeda1878053 100644 --- a/src/tirx/analysis/filter_canonical.cc +++ b/src/tirx/analysis/filter_canonical.cc @@ -25,11 +25,11 @@ #include "filter_canonical.h" -#include #include #include #include #include +#include #include namespace tvm { @@ -102,7 +102,7 @@ CmpOp Reflect(CmpOp op) { } // Compute the half-open range [lo, hi) for `var c`. -// Uses arith::ConstIntBound sentinels for unbounded sides. +// Uses sym::ConstIntBound sentinels for unbounded sides. void OpToRange(CmpOp op, int64_t c, int64_t* lo, int64_t* hi) { switch (op) { case CmpOp::kEq: @@ -110,20 +110,20 @@ void OpToRange(CmpOp op, int64_t c, int64_t* lo, int64_t* hi) { *hi = c + 1; return; case CmpOp::kLT: - *lo = arith::ConstIntBound::kNegInf; + *lo = sym::ConstIntBound::kNegInf; *hi = c; return; case CmpOp::kLE: - *lo = arith::ConstIntBound::kNegInf; + *lo = sym::ConstIntBound::kNegInf; *hi = c + 1; return; case CmpOp::kGT: *lo = c + 1; - *hi = arith::ConstIntBound::kPosInf; + *hi = sym::ConstIntBound::kPosInf; return; case CmpOp::kGE: *lo = c; - *hi = arith::ConstIntBound::kPosInf; + *hi = sym::ConstIntBound::kPosInf; return; } } diff --git a/src/tirx/analysis/filter_canonical.h b/src/tirx/analysis/filter_canonical.h index 897f6e18358f..23fbdd7c4f73 100644 --- a/src/tirx/analysis/filter_canonical.h +++ b/src/tirx/analysis/filter_canonical.h @@ -57,7 +57,7 @@ namespace tirx { * \brief Kind of an atomic predicate in canonical form. * * All five comparison operators (==, <, <=, >, >=) are normalized into a - * single half-open range atom `[lo, hi)`. Use `arith::ConstIntBound::kNegInf` + * single half-open range atom `[lo, hi)`. Use `sym::ConstIntBound::kNegInf` * for an unbounded lower side and `kPosInf` for an unbounded upper side. */ enum class FilterAtomKind { @@ -72,7 +72,7 @@ enum class FilterAtomKind { * - `scopeid_var`: the ScopeIdDef-declared variable on the LHS of the * comparison (mirrored automatically if the input had `const var`). * - `lo`, `hi`: half-open bounds. `lo` may be - * `arith::ConstIntBound::kNegInf` for an unbounded lower side; `hi` may + * `sym::ConstIntBound::kNegInf` for an unbounded lower side; `hi` may * be `kPosInf` for an unbounded upper side. * - `elect_sync_call` is unset. * @@ -141,7 +141,7 @@ using ScopeIdPredicate = std::function; * - `c1 == c2` (two constants), `v1 == v2` (two vars), and any other * non-grammar shape causes the whole classification to fail. * - The classifier is purely syntactic: it does NOT call - * `arith::Analyzer::Simplify` on subexpressions. Callers that want + * `sym::Analyzer::Simplify` on subexpressions. Callers that want * `2 + 1` to collapse to `3` should pre-simplify their input. * - This function does NOT unwrap `tirx.filter` Calls. The caller is * responsible for extracting the inner predicate (`call->args[1]`) diff --git a/src/tirx/analysis/verify_tirx_well_formed.cc b/src/tirx/analysis/verify_tirx_well_formed.cc index 8407bdb67122..239e6e4fb1b7 100644 --- a/src/tirx/analysis/verify_tirx_well_formed.cc +++ b/src/tirx/analysis/verify_tirx_well_formed.cc @@ -22,9 +22,9 @@ * \brief Check if the TIRX program is well-formed. */ -#include #include #include +#include #include #include #include @@ -114,7 +114,7 @@ class ScopeIdVerifier : public Verifier { } Array scope_id_def_; - arith::Analyzer ana_; + sym::Analyzer ana_; }; class LayoutVerifier : public Verifier { diff --git a/src/tirx/ir/buffer.cc b/src/tirx/ir/buffer.cc index 44b91f103b58..f33cefe331ca 100644 --- a/src/tirx/ir/buffer.cc +++ b/src/tirx/ir/buffer.cc @@ -20,7 +20,6 @@ /*! * \file buffer.cc */ -#include #include #include #include @@ -28,6 +27,7 @@ #include #include #include +#include #include #include #include @@ -39,7 +39,7 @@ #include #include -#include "../../arith/pattern_match.h" +#include "../../sym/pattern_match.h" namespace tvm { namespace tirx { @@ -87,7 +87,7 @@ ffi::ObjectRef RealizeBufferSubscript( // Any slice or omitted trailing dimension denotes a region. Rejecting // steps makes the old behavior, where a stride could be silently dropped, // unrepresentable rather than giving it dimension-dependent semantics. - arith::Analyzer analyzer; + sym::Analyzer analyzer; ffi::Array region; region.reserve(buffer_ty->shape.size()); for (size_t i = 0; i < slice.size(); ++i) { @@ -282,7 +282,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { }); } -ffi::Array SimplifyArray(arith::AnalyzerObj* ana, ffi::Array array) { +ffi::Array SimplifyArray(sym::AnalyzerObj* ana, ffi::Array array) { for (size_t i = 0; i < array.size(); ++i) { array.Set(i, ana->Simplify(array[i])); } @@ -322,7 +322,7 @@ inline std::vector ExprSplitAddition(const PrimExpr& expr) { // If it can be optimized, returns (true, (a1 + a2 + ... + aj) * kt * ... * ki + c1) // Currently the we will not search the add/mult combinations exhaustively // as it will take too much computation. -inline std::pair MergeMulModInner(arith::AnalyzerObj* analyzer, +inline std::pair MergeMulModInner(sym::AnalyzerObj* analyzer, const PrimExpr& mult_expr, const PrimExpr& mod_l_expr, const PrimExpr& mod_r_expr) { @@ -419,7 +419,7 @@ inline void MergeMulModInsertElements(const std::vector& eles, // The search will be performed repeatively until no pattern is found. // Return: a pair with (false, Expr()) if cannot be optimized. // a pair with (true, optimized_expr) if can be optimized -inline PrimExpr MergeMulMod(arith::AnalyzerObj* analyzer, const PrimExpr& base) { +inline PrimExpr MergeMulMod(sym::AnalyzerObj* analyzer, const PrimExpr& base) { using namespace tirx; // 1. Prepare the lists. // We store two lists, a list that contain all the elements that match Mul and @@ -427,7 +427,7 @@ inline PrimExpr MergeMulMod(arith::AnalyzerObj* analyzer, const PrimExpr& base) // The elements in the Mod will be used to match against the elements in Mul. // The result will then be split and pushed back to these two lists. PrimExpr simplified_base = base; - arith::PVar x, y; + sym::PVar x, y; if ((floordiv(x, y) * y + floormod(x, y)).Match(simplified_base)) { simplified_base = x.Eval(); } @@ -506,7 +506,7 @@ ffi::Array BufferTypeNode::ElemOffset(ffi::Array input_indic } PrimExpr output_index = 0; - arith::Analyzer ana; + sym::Analyzer ana; for (size_t i = 0; i < input_indices.size(); i++) { if (strides.size()) { @@ -656,7 +656,7 @@ BufferVar BufferVar::MakeStrideView() const { BufferVar BufferVar::MakeSlice(ffi::Array begins, ffi::Array extents) const { const BufferTypeNode* n = operator->(); TVM_FFI_ICHECK(n != nullptr); - arith::Analyzer ana; + sym::Analyzer ana; begins = SimplifyArray(ana.get(), begins); ffi::Array elem_offset = n->ElemOffset(begins).Map([&](const PrimExpr& expr) { return ana->Simplify(expr); }); diff --git a/src/tirx/ir/exec_scope.cc b/src/tirx/ir/exec_scope.cc index 1c921ae731a8..30e9d774f75c 100644 --- a/src/tirx/ir/exec_scope.cc +++ b/src/tirx/ir/exec_scope.cc @@ -16,10 +16,10 @@ * specific language governing permissions and limitations * under the License. */ -#include #include #include #include +#include #include #include #include @@ -191,7 +191,7 @@ static ScopeIdDef FillExtents(const ScopeIdDef& existing, const ScopeIdDef& fill bool ScopeIdDefVerifier::Verify(const ffi::Array& defs, Mode mode) { id_set.clear(); - arith::Analyzer ana; + sym::Analyzer ana; std::queue queue; // Insert or upgrade a binding in id_set. @@ -316,7 +316,7 @@ static ffi::Optional Compliment(const ScopeIdDef& lhs, const ScopeId return std::nullopt; } if (is_zero(rhs.fused_extent())) return std::nullopt; - arith::Analyzer ana; + sym::Analyzer ana; auto try_compliment = [&](PrimExpr lhs_ext, PrimExpr rhs_ext, ScopeBinding scope) -> ffi::Optional { if (ana->CanProve(floormod(lhs_ext, rhs_ext) == 0)) { @@ -378,7 +378,7 @@ ffi::Array Trivial3DResolve(const LaunchParams& params, const char* pr ffi::Array ResolveCuda(ScopeBinding binding, const ffi::Optional>& extents, int out_dim, const LaunchParams& params) { - arith::Analyzer ana; + sym::Analyzer ana; switch (binding) { case ScopeBinding::kKernelCta: return Trivial3DResolve(params, "blockIdx.", out_dim); diff --git a/src/tirx/ir/index_map.cc b/src/tirx/ir/index_map.cc index 1d4c41cec5b9..6076a7aa69a1 100644 --- a/src/tirx/ir/index_map.cc +++ b/src/tirx/ir/index_map.cc @@ -21,14 +21,14 @@ * \file index_map.cc */ -#include -#include -#include #include #include #include #include #include +#include +#include +#include #include #include #include @@ -69,10 +69,10 @@ IndexMap IndexMap::FromFunc(int ndim, std::pair IndexMapInverseImpl(const IndexMap& self, const ffi::Array& initial_ranges, - arith::IterMapLevel check_level, - arith::AnalyzerObj* analyzer) { + sym::IterMapLevel check_level, + sym::AnalyzerObj* analyzer) { TVM_FFI_ICHECK(analyzer != nullptr); - arith::Analyzer analyzer_ref = ffi::GetRef(analyzer); + sym::Analyzer analyzer_ref = ffi::GetRef(analyzer); if (self->inverse_index_map.has_value()) { // return the pre-defined inverse index map if exists. In this // case, the user-defined inverse is assumed to be correct and @@ -135,7 +135,7 @@ std::pair IndexMapInverseImpl(const IndexMap& self, } PrimExpr padding_predicate = padded_iter_map->padding_predicate; - padding_predicate = arith::NormalizeIterMapToExpr(padding_predicate); + padding_predicate = sym::NormalizeIterMapToExpr(padding_predicate); auto f_substitute = [&inverse_exprs_map](const Var& var) -> ffi::Expected> { if (auto repl = inverse_exprs_map.Get(var)) return ffi::Any(*std::move(repl)); @@ -148,7 +148,7 @@ std::pair IndexMapInverseImpl(const IndexMap& self, { TVM_FFI_ICHECK_EQ(output_ranges.size(), output_vars.size()); - arith::Analyzer output_var_analyzer; + sym::Analyzer output_var_analyzer; for (size_t i = 0; i < output_vars.size(); ++i) { output_var_analyzer->Bind(output_vars[i], output_ranges[i]); } @@ -163,25 +163,24 @@ std::pair IndexMapInverseImpl(const IndexMap& self, std::pair IndexMap::NonSurjectiveInverse( ffi::Array initial_ranges) const { - arith::Analyzer analyzer; + sym::Analyzer analyzer; return NonSurjectiveInverse(initial_ranges, analyzer); } -std::pair IndexMap::NonSurjectiveInverse( - ffi::Array initial_ranges, const arith::Analyzer& analyzer) const { - return IndexMapInverseImpl(*this, initial_ranges, arith::IterMapLevel::NoCheck, analyzer.get()); +std::pair IndexMap::NonSurjectiveInverse(ffi::Array initial_ranges, + const sym::Analyzer& analyzer) const { + return IndexMapInverseImpl(*this, initial_ranges, sym::IterMapLevel::NoCheck, analyzer.get()); } IndexMap IndexMap::Inverse(ffi::Array initial_ranges) const { - arith::Analyzer analyzer; + sym::Analyzer analyzer; return Inverse(initial_ranges, analyzer); } -IndexMap IndexMap::Inverse(ffi::Array initial_ranges, - const arith::Analyzer& analyzer) const { - arith::AnalyzerObj* analyzer_ptr = analyzer.get(); +IndexMap IndexMap::Inverse(ffi::Array initial_ranges, const sym::Analyzer& analyzer) const { + sym::AnalyzerObj* analyzer_ptr = analyzer.get(); auto [inverse, padding_predicate] = - IndexMapInverseImpl(*this, initial_ranges, arith::IterMapLevel::Bijective, analyzer_ptr); + IndexMapInverseImpl(*this, initial_ranges, sym::IterMapLevel::Bijective, analyzer_ptr); TVM_FFI_ICHECK(analyzer_ptr->CanProve(!padding_predicate)) << "Bijective inverse should not contain padding, but inverse of " << *this << " over range " << initial_ranges << " resulted in a padding predicate of " << padding_predicate; @@ -189,13 +188,13 @@ IndexMap IndexMap::Inverse(ffi::Array initial_ranges, } ffi::Array IndexMapNode::MapIndices(const ffi::Array& indices) const { - arith::Analyzer analyzer; + sym::Analyzer analyzer; return MapIndices(indices, analyzer); } ffi::Array IndexMapNode::MapIndices(const ffi::Array& indices, - const arith::Analyzer& analyzer) const { - arith::AnalyzerObj* analyzer_ptr = analyzer.get(); + const sym::Analyzer& analyzer) const { + sym::AnalyzerObj* analyzer_ptr = analyzer.get(); TVM_FFI_ICHECK_EQ(indices.size(), initial_indices.size()); ffi::Map vmap; @@ -213,13 +212,13 @@ ffi::Array IndexMapNode::MapIndices(const ffi::Array& indice } ffi::Array IndexMapNode::MapRanges(const ffi::Array& ranges) const { - arith::Analyzer analyzer; + sym::Analyzer analyzer; return MapRanges(ranges, analyzer); } ffi::Array IndexMapNode::MapRanges(const ffi::Array& ranges, - const arith::Analyzer& analyzer) const { - arith::AnalyzerObj* analyzer_ptr = analyzer.get(); + const sym::Analyzer& analyzer) const { + sym::AnalyzerObj* analyzer_ptr = analyzer.get(); TVM_FFI_ICHECK_EQ(ranges.size(), initial_indices.size()); ffi::Map input_iters; @@ -227,7 +226,7 @@ ffi::Array IndexMapNode::MapRanges(const ffi::Array& ranges, input_iters.Set(initial_indices[i], ranges[i]); } auto iter_map = DetectIterMap(final_indices, input_iters, /* predicate = */ 1, - /*check_level=*/arith::IterMapLevel::NoCheck, analyzer, + /*check_level=*/sym::IterMapLevel::NoCheck, analyzer, /*simplify_trivial_iterators=*/false); ffi::Array output; if (iter_map->indices.size()) { @@ -253,13 +252,13 @@ ffi::Array IndexMapNode::MapRanges(const ffi::Array& ranges, // For example, [N] mapped through i=>[i//4,i%4] should have shape // [ceildiv(N,4), 4]. However, for N<4, this method instead // results in a shape [1, N]. - std::unordered_map dom_map; + std::unordered_map dom_map; for (size_t i = 0; i < initial_indices.size(); i++) { - dom_map[initial_indices[i].get()] = arith::IntSet::FromRange(ranges[i]); + dom_map[initial_indices[i].get()] = sym::IntSet::FromRange(ranges[i]); } for (const auto& final_index : final_indices) { - auto int_set = arith::EvalSet(final_index, dom_map); + auto int_set = sym::EvalSet(final_index, dom_map); output.push_back( Range::FromMinExtent(analyzer_ptr->Simplify(int_set.min()), analyzer_ptr->Simplify(int_set.max() - int_set.min() + 1))); @@ -284,12 +283,12 @@ ffi::Array IndexMapNode::MapRanges(const ffi::Array& ranges, } ffi::Array IndexMapNode::MapShape(const ffi::Array& shape) const { - arith::Analyzer analyzer; + sym::Analyzer analyzer; return MapShape(shape, analyzer); } ffi::Array IndexMapNode::MapShape(const ffi::Array& shape, - const arith::Analyzer& analyzer) const { + const sym::Analyzer& analyzer) const { TVM_FFI_ICHECK_EQ(shape.size(), initial_indices.size()); ffi::Array ranges; @@ -308,7 +307,7 @@ ffi::Array IndexMapNode::MapShape(const ffi::Array& shape, } runtime::Tensor IndexMapNode::MapTensor(runtime::Tensor arr_src) const { - arith::Analyzer analyzer; + sym::Analyzer analyzer; auto shape = arr_src.Shape(); TVM_FFI_ICHECK(shape.size() == initial_indices.size()) << "The rank of the input array should be " << initial_indices.size() << " but got " @@ -488,35 +487,33 @@ TVM_FFI_STATIC_INIT_BLOCK() { }) .def("tirx.IndexMapMapIndices", [](IndexMap map, ffi::Array indices, - ffi::Optional opt_analyzer) { - arith::Analyzer analyzer = - opt_analyzer.has_value() ? opt_analyzer.value() : arith::Analyzer(); + ffi::Optional opt_analyzer) { + sym::Analyzer analyzer = + opt_analyzer.has_value() ? opt_analyzer.value() : sym::Analyzer(); return map->MapIndices(indices, analyzer); }) .def("tirx.IndexMapMapShape", - [](IndexMap map, ffi::Array shape, - ffi::Optional opt_analyzer) { - arith::Analyzer analyzer = - opt_analyzer.has_value() ? opt_analyzer.value() : arith::Analyzer(); + [](IndexMap map, ffi::Array shape, ffi::Optional opt_analyzer) { + sym::Analyzer analyzer = + opt_analyzer.has_value() ? opt_analyzer.value() : sym::Analyzer(); return map->MapShape(shape, analyzer); }) .def("tirx.IndexMapInverse", [](IndexMap map, ffi::Array initial_ranges, - ffi::Optional opt_analyzer) { - arith::Analyzer analyzer = - opt_analyzer.has_value() ? opt_analyzer.value() : arith::Analyzer(); + ffi::Optional opt_analyzer) { + sym::Analyzer analyzer = + opt_analyzer.has_value() ? opt_analyzer.value() : sym::Analyzer(); return map.Inverse(initial_ranges, analyzer); }) .def("tirx.IndexMapMapTensor", [](IndexMap map, runtime::Tensor arr) { return map->MapTensor(arr); }) - .def("tirx.IndexMapNonSurjectiveInverse", - [](IndexMap forward, ffi::Array initial_ranges, - ffi::Optional opt_analyzer) { - arith::Analyzer analyzer = - opt_analyzer.has_value() ? opt_analyzer.value() : arith::Analyzer(); - auto result = forward.NonSurjectiveInverse(initial_ranges, analyzer); - return ffi::Array{result.first, result.second}; - }); + .def("tirx.IndexMapNonSurjectiveInverse", [](IndexMap forward, + ffi::Array initial_ranges, + ffi::Optional opt_analyzer) { + sym::Analyzer analyzer = opt_analyzer.has_value() ? opt_analyzer.value() : sym::Analyzer(); + auto result = forward.NonSurjectiveInverse(initial_ranges, analyzer); + return ffi::Array{result.first, result.second}; + }); } } // namespace tirx diff --git a/src/tirx/ir/ir_mutator_with_analyzer.cc b/src/tirx/ir/ir_mutator_with_analyzer.cc index 1c65f94874a3..c4194ea81277 100644 --- a/src/tirx/ir/ir_mutator_with_analyzer.cc +++ b/src/tirx/ir/ir_mutator_with_analyzer.cc @@ -22,14 +22,14 @@ */ #include "ir_mutator_with_analyzer.h" -#include #include #include +#include #include #include #include -#include "../../arith/constraint_helpers.h" +#include "../../sym/constraint_helpers.h" namespace tvm { namespace tirx { @@ -47,7 +47,7 @@ const IRMutatorWithAnalyzer::VTable* IRMutatorWithAnalyzer::GlobalVTable() { using namespace tvm::prim; -using arith::detail::EnterConstraintFacts; +using sym::detail::EnterConstraintFacts; void IRMutatorWithAnalyzer::MarkBufferParamShapes(const tirx::PrimFunc& func) { // Mark all symbolic buffer-parameter shape values as positive. @@ -69,9 +69,9 @@ ffi::Array IRMutatorWithAnalyzer::IterMapSimplifyWithContext( pred = pred && val; } int n = indices.size(); - arith::Analyzer analyzer_ref = ffi::GetRef(this->analyzer_); - ffi::Array simplified = arith::IterMapSimplify( - indices, this->iter_vars_, pred, arith::IterMapLevel::Surjective, analyzer_ref); + sym::Analyzer analyzer_ref = ffi::GetRef(this->analyzer_); + ffi::Array simplified = sym::IterMapSimplify( + indices, this->iter_vars_, pred, sym::IterMapLevel::Surjective, analyzer_ref); if (non_trivial_only) { for (int i = 0; i < n; ++i) { if (simplified[i]->IsInstance() && indices[i].as()) { diff --git a/src/tirx/ir/ir_mutator_with_analyzer.h b/src/tirx/ir/ir_mutator_with_analyzer.h index bcc976c57ac5..0e93dc140a88 100644 --- a/src/tirx/ir/ir_mutator_with_analyzer.h +++ b/src/tirx/ir/ir_mutator_with_analyzer.h @@ -24,11 +24,11 @@ #ifndef TVM_TIRX_IR_IR_MUTATOR_WITH_ANALYZER_H_ #define TVM_TIRX_IR_IR_MUTATOR_WITH_ANALYZER_H_ -#include #include #include #include #include +#include #include #include @@ -51,9 +51,9 @@ class IRMutatorWithAnalyzer : public StmtExprMutator { public: using StmtExprMutator::Mutate; using StmtExprMutator::Mutate_; - explicit IRMutatorWithAnalyzer(const arith::Analyzer& analyzer) + explicit IRMutatorWithAnalyzer(const sym::Analyzer& analyzer) : IRMutatorWithAnalyzer(analyzer.get()) {} - explicit IRMutatorWithAnalyzer(arith::AnalyzerObj* analyzer) + explicit IRMutatorWithAnalyzer(sym::AnalyzerObj* analyzer) : IRMutatorWithAnalyzer(analyzer, GlobalVTable()) {} // override functions that need to populate the context information. @@ -68,7 +68,7 @@ class IRMutatorWithAnalyzer : public StmtExprMutator { protected: static void InitVTable(VTable* vtable); - IRMutatorWithAnalyzer(arith::AnalyzerObj* analyzer, const VTable* vtable) + IRMutatorWithAnalyzer(sym::AnalyzerObj* analyzer, const VTable* vtable) : StmtExprMutator(vtable), analyzer_(analyzer) {} static const VTable* GlobalVTable(); /*! @@ -87,9 +87,9 @@ class IRMutatorWithAnalyzer : public StmtExprMutator { bool non_trivial_only); /*! \brief internal analyzer field. */ - arith::AnalyzerObj* analyzer_; + sym::AnalyzerObj* analyzer_; /*! \brief Scope stack for accumulated assert constraints. */ - ScopeStack> constraint_scope_; + ScopeStack> constraint_scope_; // the following two fields are useful in case we want // note however that iter map analysis are usually more // expensive and we only encourage doing them during diff --git a/src/tirx/ir/ir_visitor_with_analyzer.h b/src/tirx/ir/ir_visitor_with_analyzer.h index 63db9d9fc0cf..692739969611 100644 --- a/src/tirx/ir/ir_visitor_with_analyzer.h +++ b/src/tirx/ir/ir_visitor_with_analyzer.h @@ -25,10 +25,10 @@ #ifndef TVM_TIRX_IR_IR_VISITOR_WITH_ANALYZER_H_ #define TVM_TIRX_IR_IR_VISITOR_WITH_ANALYZER_H_ -#include #include #include #include +#include #include namespace tvm { @@ -58,10 +58,10 @@ class IRVisitorWithAnalyzer : public StmtExprVisitor { static void InitVTable(VTable* vtable); explicit IRVisitorWithAnalyzer(const VTable* vtable) : StmtExprVisitor(vtable) {} /*! \brief internal analyzer field. */ - arith::Analyzer analyzer_; + sym::Analyzer analyzer_; /*! \brief Scope stack for accumulated assert constraints. */ - ScopeStack> constraint_scope_; + ScopeStack> constraint_scope_; /*! \brief Extract a constraint from a conditional statement * diff --git a/src/tirx/ir/layout/axis_registry.cc b/src/tirx/ir/layout/axis_registry.cc index 6c4627a31d26..2474390789fb 100644 --- a/src/tirx/ir/layout/axis_registry.cc +++ b/src/tirx/ir/layout/axis_registry.cc @@ -163,7 +163,7 @@ void AxisRegEntry::UpdateAttr(const ffi::String& key, ffi::Any value, int plevel // register thread axis split/fuse helpers ffi::Array SplitterGen(const Iter& iter, const Axis& axis_outer, const Axis& axis_inner, const PrimExpr& e_inner) { - arith::Analyzer analyzer; + sym::Analyzer analyzer; if (analyzer->CanProve(iter->extent * iter->stride < e_inner)) { return {Iter(iter->extent, iter->stride, axis_inner)}; } else if (analyzer->CanProveEqual(floormod(e_inner, iter->stride), 0) && @@ -197,7 +197,7 @@ TVM_REGISTER_AXIS("tx") return std::nullopt; }) .set_splitter([](Target target, ffi::String scope, Iter iter) -> ffi::Array { - arith::Analyzer analyzer; + sym::Analyzer analyzer; if (target->kind->default_device_type == kDLCUDA) { if (scope == "warp") { // tx -> warpid, laneid @@ -226,7 +226,7 @@ TVM_REGISTER_AXIS("warpid") return std::nullopt; }) .set_splitter([](Target target, ffi::String scope, Iter iter) -> ffi::Array { - arith::Analyzer analyzer; + sym::Analyzer analyzer; if (target->kind->default_device_type == kDLCUDA) { if (scope == "warp") { // warpid -> wgid, wid_in_wg @@ -255,7 +255,7 @@ TVM_REGISTER_AXIS("laneid") return std::nullopt; }) .set_splitter([](Target target, ffi::String scope, Iter iter) -> ffi::Array { - arith::Analyzer analyzer; + sym::Analyzer analyzer; if (target->kind->default_device_type == kDLCUDA) { LOG(FATAL) << "laneid can not be split any more"; } @@ -279,7 +279,7 @@ TVM_REGISTER_AXIS("wgid") return std::nullopt; }) .set_splitter([](Target target, ffi::String scope, Iter iter) -> ffi::Array { - arith::Analyzer analyzer; + sym::Analyzer analyzer; if (target->kind->default_device_type == kDLCUDA) { LOG(FATAL) << "wgid can not be split any more"; } @@ -301,7 +301,7 @@ TVM_REGISTER_AXIS("tid_in_wg") return std::nullopt; }) .set_splitter([](Target target, ffi::String scope, Iter iter) -> ffi::Array { - arith::Analyzer analyzer; + sym::Analyzer analyzer; if (target->kind->default_device_type == kDLCUDA) { if (scope == "warp") { // tid_in_wg -> wid_in_wg, laneid @@ -333,7 +333,7 @@ TVM_REGISTER_AXIS("wid_in_wg") return std::nullopt; }) .set_splitter([](Target target, ffi::String scope, Iter iter) -> ffi::Array { - arith::Analyzer analyzer; + sym::Analyzer analyzer; if (target->kind->default_device_type == kDLCUDA) { LOG(FATAL) << "wid_in_wg can not be split any more"; } diff --git a/src/tirx/ir/layout/compose_layout.cc b/src/tirx/ir/layout/compose_layout.cc index b4e3fe6cbb92..1f79b13e4154 100644 --- a/src/tirx/ir/layout/compose_layout.cc +++ b/src/tirx/ir/layout/compose_layout.cc @@ -37,7 +37,7 @@ PrimExpr ApplyFullSwizzle(const ComposeLayoutNode* layout, const PrimExpr& m) { return x ^ ((x & layout->inner_mask) << layout->atom_len); }; int base = 1 << layout->per_element; - arith::Analyzer analyzer; + sym::Analyzer analyzer; PrimVar m_once("compose_m", m.ty()); PrimExpr quotient = floordiv(m_once, base); PrimVar quotient_once("compose_q", quotient.ty()); @@ -46,7 +46,7 @@ PrimExpr ApplyFullSwizzle(const ComposeLayoutNode* layout, const PrimExpr& m) { return prim::Let(m_once, m, prim::Let(quotient_once, quotient, body)); } -void AddExpr(std::optional* sum, const PrimExpr& term, const arith::Analyzer& analyzer) { +void AddExpr(std::optional* sum, const PrimExpr& term, const sym::Analyzer& analyzer) { if (is_zero(term)) return; if (sum->has_value()) { *sum = analyzer->Simplify(sum->value() + term); @@ -74,7 +74,7 @@ bool MulWithoutOverflow(int64_t lhs, int64_t rhs, int64_t* result) { } void CollectOffsetTerms(const PrimExpr& expr, int sign, std::vector* dynamic_terms, - int64_t* constant, bool* valid, const arith::Analyzer& analyzer) { + int64_t* constant, bool* valid, const sym::Analyzer& analyzer) { if (!*valid) return; PrimExpr simplified = analyzer->Simplify(expr); if (const auto* imm = simplified.as()) { @@ -104,7 +104,7 @@ void CollectOffsetTerms(const PrimExpr& expr, int sign, std::vector* d } std::optional DivideExactTerm(const PrimExpr& term, int64_t divisor, - const arith::Analyzer& analyzer) { + const sym::Analyzer& analyzer) { PrimExpr simplified = analyzer->Simplify(term); if (const auto* imm = simplified.as()) { if (imm->value % divisor != 0) return std::nullopt; @@ -146,7 +146,7 @@ ffi::Map ApplyStructured(const ComposeLayoutNode* layout, TVM_FFI_ICHECK_EQ(coord.size(), tile->shard.size()) << "Coordinate size must match the number of shard axes"; - arith::Analyzer analyzer; + sym::Analyzer analyzer; for (size_t i = 0; i < tile->shard.size(); ++i) { if (analyzer->CanProveEqual(tile->shard[i]->extent, 1)) { coord.Set(i, IntImm(coord[i].ty(), 0)); @@ -256,8 +256,8 @@ ffi::Map ApplyStructured(const ComposeLayoutNode* layout, add_high(term, quotient.value()); continue; } - arith::ConstIntBound bound = analyzer->const_int_bound(term); - if (bound->min_value < 0 || bound->max_value == arith::ConstIntBound::kPosInf || + sym::ConstIntBound bound = analyzer->const_int_bound(term); + if (bound->min_value < 0 || bound->max_value == sym::ConstIntBound::kPosInf || !add_low(term, bound->max_value)) { return fallback(); } @@ -362,7 +362,7 @@ ffi::Map ComposeLayoutNode::Apply(PrimExpr coord) const { } }; auto base = 1 << per_element; - arith::Analyzer analyzer; + sym::Analyzer analyzer; return {{"m", analyzer->Simplify((f(floordiv(m, base)) << per_element) + floormod(m, base))}}; } diff --git a/src/tirx/ir/layout/tile_canonicalize.cc b/src/tirx/ir/layout/tile_canonicalize.cc index 23a000306b97..ed6724cf21ba 100644 --- a/src/tirx/ir/layout/tile_canonicalize.cc +++ b/src/tirx/ir/layout/tile_canonicalize.cc @@ -56,7 +56,7 @@ TileLayout RemoveZeroOffsets(TileLayout layout) { TileLayout FuseContiguousShardIters(TileLayout layout) { std::vector fused_shard; - arith::Analyzer ana; + sym::Analyzer ana; const auto& shard = layout->shard; for (size_t cur = 0; cur < shard.size();) { // Find consecutive fusable axes diff --git a/src/tirx/ir/layout/tile_core.cc b/src/tirx/ir/layout/tile_core.cc index b39f1999fc82..e6681adee1fb 100644 --- a/src/tirx/ir/layout/tile_core.cc +++ b/src/tirx/ir/layout/tile_core.cc @@ -195,7 +195,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { bool TileLayoutNode::CompatibleWithShape(const Array& shape) const { return true; } bool VerifyCompactness(const std::vector& iters) { - arith::Analyzer analyzer; + sym::Analyzer analyzer; PrimExpr stride_to_find = 1; for (size_t i = 0; i < iters.size(); ++i) { auto iter = std::find_if(iters.begin(), iters.end(), [&](const Iter& iter) { @@ -251,7 +251,7 @@ PrimExpr TileLayoutNode::GetSize(ffi::Optional axis_name) const { } PrimExpr TileLayoutNode::GetSpan(ffi::Optional axis_name) const { - arith::Analyzer analyzer; + sym::Analyzer analyzer; PrimExpr result = 1; auto filter = [&](const Axis& axis) { return AxisMatchesFilter(axis, axis_name); }; @@ -281,7 +281,7 @@ ffi::Map TileLayoutNode::Apply(const ffi::Array // ``coord[d]`` against just that sub-range's *local* extents keeps the // symbolic form small (local mod/divs) and avoids the cross-dim noise of the // flatten+split-against-shard-shape round-trip. Equivalent numerical output, - // much friendlier for arith.Analyzer downstream. + // much friendlier for sym.Analyzer downstream. if (auto grouped_opt = TryGroup(ffi::GetRef(this), shape); grouped_opt.has_value()) { auto& [grouped, seps] = *grouped_opt; ffi::Array per_shard_coords; @@ -309,7 +309,7 @@ ffi::Map TileLayoutNode::Apply(const ffi::Array } ffi::Map TileLayoutNode::Apply(Array coord) const { - arith::Analyzer analyzer; + sym::Analyzer analyzer; TVM_FFI_ICHECK_EQ(coord.size(), shard.size()) << "Coordinate size must match the number of shard axes"; std::unordered_map result; diff --git a/src/tirx/ir/layout/tile_direct_sum_ops.cc b/src/tirx/ir/layout/tile_direct_sum_ops.cc index 5ee6fe9041af..a7804c03d25a 100644 --- a/src/tirx/ir/layout/tile_direct_sum_ops.cc +++ b/src/tirx/ir/layout/tile_direct_sum_ops.cc @@ -56,7 +56,7 @@ Layout TileLayoutNode::DirectSum(const TileLayout& left_in, const Arrayreplica.begin(), right->replica.end()); // Offsets add: O^A + O^B per-axis - arith::Analyzer analyzer; + sym::Analyzer analyzer; ffi::Map sum_off; for (const auto& [axis, off] : left->offset) sum_off.Set(axis, off); for (const auto& [axis, off] : right->offset) { @@ -71,7 +71,7 @@ Layout TileLayoutNode::DirectSum(const TileLayout& left_in, const ArrayCanonicalize(); } -static bool IterEqualRelaxUnit(const Iter& a, const Iter& b, arith::AnalyzerObj* analyzer) { +static bool IterEqualRelaxUnit(const Iter& a, const Iter& b, sym::AnalyzerObj* analyzer) { if (!(*analyzer).CanProveEqual(a->extent, b->extent)) return false; if (!is_one(a->extent)) { if (!(*analyzer).CanProveEqual(a->stride, b->stride)) return false; @@ -83,7 +83,7 @@ static bool IterEqualRelaxUnit(const Iter& a, const Iter& b, arith::AnalyzerObj* // Helper to subtract offsets: left = sum - right static ffi::Map SubtractOffsets(const ffi::Map& sum, const ffi::Map& rhs) { - arith::Analyzer analyzer; + sym::Analyzer analyzer; ffi::Map res; for (const auto& [axis, off] : sum) res.Set(axis, off); for (const auto& [axis, off] : rhs) { @@ -103,7 +103,7 @@ ffi::Optional TileLayoutNode::IsDirectSumRight( auto maybe_sum = sum_layout_in.as(); if (!maybe_sum) return std::nullopt; - arith::Analyzer analyzer; + sym::Analyzer analyzer; TileLayout sum_layout = maybe_sum.value()->Canonicalize().as().value(); TileLayout right = ffi::GetRef(this)->Canonicalize().as().value(); @@ -154,7 +154,7 @@ ffi::Optional TileLayoutNode::IsDirectSumLeft( auto maybe_sum = sum_layout_in.as(); if (!maybe_sum) return std::nullopt; - arith::Analyzer analyzer; + sym::Analyzer analyzer; TileLayout sum_layout = maybe_sum.value()->Canonicalize().as().value(); TileLayout left = ffi::GetRef(this)->Canonicalize().as().value(); diff --git a/src/tirx/ir/layout/tile_slice.cc b/src/tirx/ir/layout/tile_slice.cc index b2c793a93d77..ed9d61c57b4c 100644 --- a/src/tirx/ir/layout/tile_slice.cc +++ b/src/tirx/ir/layout/tile_slice.cc @@ -33,7 +33,7 @@ ffi::Optional SlicePerGroup(TileLayout layout, PrimExpr begin, PrimE return std::nullopt; } - arith::Analyzer analyzer; + sym::Analyzer analyzer; int m = static_cast(shard.size()); std::vector B(m); @@ -143,7 +143,7 @@ ffi::Optional SlicePerGroup(TileLayout layout, PrimExpr begin, PrimE ffi::Optional TileLayoutNode::Slice(const Array& shape, const Region& region) const { - arith::Analyzer analyzer; + sym::Analyzer analyzer; // Canonicalize the whole layout first so scope fusion (e.g. wid_in_wg+laneid // -> tid_in_wg) runs globally; otherwise grouping can split sibling thread // axes and SlicePerGroup's per-group fusion leaves an ill-formed mix. diff --git a/src/tirx/ir/layout/tile_tile_ops.cc b/src/tirx/ir/layout/tile_tile_ops.cc index 5f10a7a96adb..22450684f1c2 100644 --- a/src/tirx/ir/layout/tile_tile_ops.cc +++ b/src/tirx/ir/layout/tile_tile_ops.cc @@ -30,7 +30,7 @@ using namespace tvm::prim; std::pair> Group(TileLayout layout, const ffi::Array& shape) { - arith::Analyzer analyzer; + sym::Analyzer analyzer; size_t shape_idx = 0; PrimExpr prod = 1; @@ -81,7 +81,7 @@ std::pair>> GroupMany( std::vector counts; }; - arith::Analyzer analyzer; + sym::Analyzer analyzer; TVM_FFI_ICHECK(!shapes.empty()) << "group_many requires at least one shape"; std::vector> boundary_sequences; @@ -265,7 +265,7 @@ std::optional>> TryGroup( // Same algorithm as Group but returns std::nullopt instead of ICHECK-failing // on regroup impossibility. Used by Apply(coord, shape) to opportunistically // pick the group-first path with a fallback to flatten+split. - arith::Analyzer analyzer; + sym::Analyzer analyzer; size_t shape_idx = 0; PrimExpr prod = 1; @@ -336,7 +336,7 @@ Layout TileLayoutNode::Tile(const TileLayout& outer_in, const Array& o outer = grouped_outer; inner = grouped_inner; - arith::Analyzer analyzer; + sym::Analyzer analyzer; { // Scale outer axis strides by inner span on matching axes @@ -391,7 +391,7 @@ Layout TileLayoutNode::Tile(const TileLayout& outer_in, const Array& o ffi::Array TileShape(ffi::Array shape, ffi::Array factor, bool is_inner) { TVM_FFI_ICHECK_EQ(shape.size(), factor.size()) << "Shape and factor dimension must match."; - arith::Analyzer analyzer; + sym::Analyzer analyzer; ffi::Array new_shape; for (int i = 0; i < static_cast(shape.size()); ++i) { @@ -489,7 +489,7 @@ ffi::Optional TileLayoutNode::IsTileInner( } } - arith::Analyzer analyzer; + sym::Analyzer analyzer; // Get the span map of the inner layout of each axis auto inner_span_map = BuildSpanMap(layout); auto rescale_by_inner_span = [&](const Iter& iter) -> ffi::Optional { @@ -595,7 +595,7 @@ ffi::Optional TileLayoutNode::IsTileOuter(const Layout& tile_layout, } } - arith::Analyzer analyzer; + sym::Analyzer analyzer; TVM_FFI_ICHECK_EQ(tiled_shape.size(), outer_shape.size()) << "Tiled shape size must match outer shape size"; diff --git a/src/tirx/ir/layout/utils.h b/src/tirx/ir/layout/utils.h index 2051a935aaab..1c89b1732ebc 100644 --- a/src/tirx/ir/layout/utils.h +++ b/src/tirx/ir/layout/utils.h @@ -20,11 +20,11 @@ #ifndef TVM_TIRX_IR_LAYOUT_UTILS_H_ #define TVM_TIRX_IR_LAYOUT_UTILS_H_ -#include #include #include #include #include +#include #include #include diff --git a/src/tirx/ir/script/script_complete.cc b/src/tirx/ir/script/script_complete.cc index e37c2e10a497..97f8bd237fde 100644 --- a/src/tirx/ir/script/script_complete.cc +++ b/src/tirx/ir/script/script_complete.cc @@ -24,11 +24,11 @@ #include "./script_complete.h" -#include #include #include #include #include +#include #include #include diff --git a/src/tirx/ir/stmt.cc b/src/tirx/ir/stmt.cc index 179f546cd234..7d0dd90f5e5a 100644 --- a/src/tirx/ir/stmt.cc +++ b/src/tirx/ir/stmt.cc @@ -20,13 +20,13 @@ /*! * \file tvm/tirx/stmt.cc */ -#include #include #include #include #include #include #include +#include #include #include #include @@ -684,7 +684,7 @@ ffi::ObjectRef RealizeBufferRegionSubscript(Expr value, SubscriptSlice slice, Sp return BufferLoad(source->source.as_or_throw(), indices, span); } - arith::Analyzer analyzer; + sym::Analyzer analyzer; ffi::Array region; region.reserve(source->region.size()); for (size_t i = 0; i < slice.size(); ++i) { diff --git a/src/tirx/script/builder/ir.cc b/src/tirx/script/builder/ir.cc index 4fa7629cd02d..4d19237e4428 100644 --- a/src/tirx/script/builder/ir.cc +++ b/src/tirx/script/builder/ir.cc @@ -16,7 +16,6 @@ * specific language governing permissions and limitations * under the License. */ -#include #include #include #include @@ -28,6 +27,7 @@ #include #include #include +#include #include #include #include @@ -536,7 +536,7 @@ PrimExpr ConvertLoopBound(const PrimExpr& e, const PrimType& var_ty) { ffi::Optional step, ffi::Optional dtype) { \ PrimType var_ty = InferLoopVarDtype(start, stop, dtype); \ PrimExpr min = ConvertLoopBound(start, var_ty); \ - PrimExpr extent = arith::Analyzer()->Simplify(ConvertLoopBound(stop, var_ty) - min); \ + PrimExpr extent = sym::Analyzer()->Simplify(ConvertLoopBound(stop, var_ty) - min); \ if (step.has_value()) { \ step = ConvertLoopBound(step.value(), var_ty); \ } \ @@ -568,7 +568,7 @@ ForFrame ThreadBinding(PrimExpr start, PrimExpr stop, ffi::String thread, ffi::Optional> annotations) { using namespace tvm::tirx; PrimExpr min = start; - PrimExpr extent = arith::Analyzer()->Simplify(stop - start); + PrimExpr extent = sym::Analyzer()->Simplify(stop - start); ffi::ObjectPtr n = ffi::make_object(); PrimType min_ty = min.ty(); PrimType extent_ty = extent.ty(); @@ -680,7 +680,7 @@ LaunchThreadFrame LaunchThread(Var var, PrimExpr extent) { if (!iter_var->dom.defined()) { const_cast(iter_var.get())->dom = Range(tvm::IntImm(extent.ty(), 0), extent); - } else if (!arith::Analyzer()->CanProveEqual(iter_var->dom->extent, extent)) { + } else if (!sym::Analyzer()->CanProveEqual(iter_var->dom->extent, extent)) { TVM_FFI_THROW(InternalError) << "ValueError: Inconsistent extents of environment thread. " << iter_var->dom->extent << " vs " << extent; } diff --git a/src/tirx/script/printer/stmt.cc b/src/tirx/script/printer/stmt.cc index 2bc677cf1e80..9ce1c69b1fb8 100644 --- a/src/tirx/script/printer/stmt.cc +++ b/src/tirx/script/printer/stmt.cc @@ -16,7 +16,7 @@ * specific language governing permissions and limitations * under the License. */ -#include +#include #include @@ -398,7 +398,7 @@ ffi::Optional TryDeclBufferSugarWithParent(const tirx::BufferVar& child for (const PrimExpr& dim : child->shape) { child_total = child_total * dim; } - arith::Analyzer analyzer; + sym::Analyzer analyzer; bool default_physical = child_is_default && analyzer->CanProveEqual(child_total, storage_span); bool child_has_thread_axis = false; diff --git a/src/tirx/transform/flatten_buffer.cc b/src/tirx/transform/flatten_buffer.cc index 461431b5abc4..f0458210ebe5 100644 --- a/src/tirx/transform/flatten_buffer.cc +++ b/src/tirx/transform/flatten_buffer.cc @@ -21,10 +21,10 @@ * \file flatten_buffer.cc */ -#include #include #include #include +#include #include #include #include @@ -63,7 +63,7 @@ class BufferFlattener : public IRMutatorWithAnalyzer { using IRMutatorWithAnalyzer::Mutate; using IRMutatorWithAnalyzer::Mutate_; static PrimFunc Flatten(PrimFunc func) { - arith::Analyzer ana; + sym::Analyzer ana; auto pass = ffi::make_object(ana); pass->MarkBufferParamShapes(func); for (const Var& param : func->params) { @@ -98,7 +98,7 @@ class BufferFlattener : public IRMutatorWithAnalyzer { } public: - explicit BufferFlattener(const arith::Analyzer& ana) : IRMutatorWithAnalyzer(ana) {} + explicit BufferFlattener(const sym::Analyzer& ana) : IRMutatorWithAnalyzer(ana) {} private: struct FlatInfo { diff --git a/src/tirx/transform/ir_utils.cc b/src/tirx/transform/ir_utils.cc index be6d6cba09db..07eb7cc8425a 100644 --- a/src/tirx/transform/ir_utils.cc +++ b/src/tirx/transform/ir_utils.cc @@ -23,11 +23,11 @@ */ #include "ir_utils.h" -#include #include #include #include #include +#include #include #include #include @@ -654,8 +654,8 @@ ffi::Array GetBufferAllocationShape(const BufferVar& buffer) { if (buffer->strides.size()) { TVM_FFI_ICHECK_EQ(buffer->shape.size(), buffer->strides.size()); for (size_t i = buffer->strides.size() - 1; i > 0; --i) { - TVM_FFI_ICHECK(arith::Analyzer()->CanProveEqual( - floormod(buffer->strides[i - 1], buffer->strides[i]), 0)); + TVM_FFI_ICHECK( + sym::Analyzer()->CanProveEqual(floormod(buffer->strides[i - 1], buffer->strides[i]), 0)); alloc_shape.Set(i, buffer->strides[i - 1] / buffer->strides[i]); } } diff --git a/src/tirx/transform/ir_utils.h b/src/tirx/transform/ir_utils.h index 004bfa6f9d2b..4182bba09349 100644 --- a/src/tirx/transform/ir_utils.h +++ b/src/tirx/transform/ir_utils.h @@ -24,12 +24,12 @@ #ifndef TVM_TIR_TRANSFORM_IR_UTILS_H_ #define TVM_TIR_TRANSFORM_IR_UTILS_H_ -#include #include #include #include #include #include +#include #include #include #include diff --git a/src/tirx/transform/lower_intrin.cc b/src/tirx/transform/lower_intrin.cc index 53438ff2f719..c42650bb6983 100644 --- a/src/tirx/transform/lower_intrin.cc +++ b/src/tirx/transform/lower_intrin.cc @@ -37,7 +37,7 @@ #include #include -#include "../../arith/pattern_match.h" +#include "../../sym/pattern_match.h" #include "../ir/ir_mutator_with_analyzer.h" namespace tvm { @@ -132,7 +132,7 @@ class IntrinInjecter : public IRMutatorWithAnalyzer { using FLowerGeneral = ffi::TypedFunction; - IntrinInjecter(const arith::Analyzer& analyzer, const Target& tgt, bool enable_fast_math) + IntrinInjecter(const sym::Analyzer& analyzer, const Target& tgt, bool enable_fast_math) : IRMutatorWithAnalyzer(analyzer) { std::string target = tgt->kind->name; ffi::String mtriple = tgt->GetAttr("mtriple").value_or(""); @@ -348,7 +348,7 @@ class IntrinInjecter : public IRMutatorWithAnalyzer { } UnchangedOr Mutate_(const prim::MaxNode* op, InplaceMode inplace_mode) final { - using namespace arith; + using namespace sym; PVar x, y; PVar c; auto e = ffi::GetRef(op); @@ -361,7 +361,7 @@ class IntrinInjecter : public IRMutatorWithAnalyzer { } UnchangedOr Mutate_(const prim::EQNode* op, InplaceMode inplace_mode) final { - using namespace arith; + using namespace sym; PVar x, y; auto e = ffi::GetRef(op); if ((floormod(x, y) == 0).Match(e)) { @@ -372,7 +372,7 @@ class IntrinInjecter : public IRMutatorWithAnalyzer { } UnchangedOr Mutate_(const prim::NENode* op, InplaceMode inplace_mode) final { - using namespace arith; + using namespace sym; PVar x, y; auto e = ffi::GetRef(op); if ((floormod(x, y) != 0).Match(e)) { @@ -455,7 +455,7 @@ class IntrinInjecter : public IRMutatorWithAnalyzer { // NOTE: we need to be very careful in the checks below, to make sure // all the intermediate calculations in both compiler checks and runtime checks // do not overflow - arith::ConstIntBound const_int_bound_a = analyzer_->const_int_bound(a); + sym::ConstIntBound const_int_bound_a = analyzer_->const_int_bound(a); if (const_int_bound_a->min_value >= 0) { return std::nullopt; } @@ -495,7 +495,7 @@ class IntrinInjecter : public IRMutatorWithAnalyzer { }; Stmt LowerIntrinStmt(Stmt stmt, const std::string& target) { - arith::Analyzer analyzer; + sym::Analyzer analyzer; bool enable_fast_math = transform::PassContext::Current()->GetConfig("tirx.enable_fast_math", false).value(); return ffi::make_object(analyzer, Target(ffi::String(target)), enable_fast_math) @@ -510,7 +510,7 @@ Pass LowerIntrin() { auto* n = f.CopyOnWrite(); auto target = f->GetAttr(tvm::attr::kTarget); TVM_FFI_ICHECK(target.has_value()) << "LowerIntrin: Require the target attribute"; - arith::Analyzer analyzer; + sym::Analyzer analyzer; bool enable_fast_math = ctx->GetConfig("tirx.enable_fast_math", false).value(); n->body = ffi::make_object(analyzer, target.value(), enable_fast_math) ->Mutate(n->body, InplaceMode::kAllow) diff --git a/src/tirx/transform/lower_tirx_cleanup.cc b/src/tirx/transform/lower_tirx_cleanup.cc index 9a05731c93f3..a1c484988f47 100644 --- a/src/tirx/transform/lower_tirx_cleanup.cc +++ b/src/tirx/transform/lower_tirx_cleanup.cc @@ -22,8 +22,8 @@ * \brief Final cleanup stage for TIRx lowering. */ -#include #include +#include #include #include #include @@ -50,7 +50,7 @@ class LayoutApplier : public IRMutatorWithAnalyzer { using IRMutatorWithAnalyzer::Mutate_; static std::pair> Flatten(const Stmt& stmt, const ffi::Array& params, const Target& target) { - arith::Analyzer ana; + sym::Analyzer ana; auto storage_lower = ffi::make_object(ana, target); ffi::Array new_params; new_params.reserve(params.size()); @@ -81,7 +81,7 @@ class LayoutApplier : public IRMutatorWithAnalyzer { } public: - explicit LayoutApplier(const arith::Analyzer& analyzer, const Target& target) + explicit LayoutApplier(const sym::Analyzer& analyzer, const Target& target) : IRMutatorWithAnalyzer(analyzer), target_(target) {} protected: @@ -173,7 +173,7 @@ class LayoutApplier : public IRMutatorWithAnalyzer { if (auto tile_layout = buf->layout.as(); tile_layout && tile_layout->HasThreadAxis()) { // Logical alloc_buffer with thread axes: physical shape = memory-axis span - arith::Analyzer ana; + sym::Analyzer ana; PrimExpr mem_span = IntImm::Int32(1); for (const auto& iter : tile_layout->shard) { if (iter->axis->IsMemoryAxis()) { diff --git a/src/tirx/transform/lower_tirx_dedup_tensormap.cc b/src/tirx/transform/lower_tirx_dedup_tensormap.cc index 3c0613a03b86..b4229630ee37 100644 --- a/src/tirx/transform/lower_tirx_dedup_tensormap.cc +++ b/src/tirx/transform/lower_tirx_dedup_tensormap.cc @@ -22,8 +22,8 @@ * \brief Deduplicate identical cuTensorMap objects created by TIRx schedules. */ -#include #include +#include #include #include #include diff --git a/src/tirx/transform/lower_warp_memory.cc b/src/tirx/transform/lower_warp_memory.cc index 9aab893b2d32..0caa466e4a5f 100644 --- a/src/tirx/transform/lower_warp_memory.cc +++ b/src/tirx/transform/lower_warp_memory.cc @@ -25,8 +25,6 @@ */ // Thanks to Andrew Adams and Vinod Grover for // explaining the concept of warp shuffle. -#include -#include #include #include #include @@ -35,6 +33,8 @@ #include #include #include +#include +#include #include #include #include @@ -44,8 +44,8 @@ #include -#include "../../arith/pattern_match.h" #include "../../runtime/thread_storage_scope.h" +#include "../../sym/pattern_match.h" #include "ir_utils.h" #include "update_pointer_storage_scope.h" @@ -121,7 +121,7 @@ const VarNode* GetBufferVar(const Expr& expr) { class WarpStoreCoeffFinder : public StmtExprVisitor { public: - WarpStoreCoeffFinder(const VarNode* buffer, Var warp_index, arith::AnalyzerObj* analyzer) + WarpStoreCoeffFinder(const VarNode* buffer, Var warp_index, sym::AnalyzerObj* analyzer) : buffer_(buffer), warp_index_(warp_index), analyzer_(analyzer) {} // find the warp co-efficient in the statement given the warp size int Find(const Stmt& stmt) { @@ -167,8 +167,8 @@ class WarpStoreCoeffFinder : public StmtExprVisitor { PrimExpr index = op->indices[0]; PrimType value_ty = op->value.ty(); if (value_ty.lanes() != 1) { - arith::PVar base; - TVM_FFI_ICHECK(arith::ramp(base, 1, value_ty.lanes()).Match(index)) + sym::PVar base; + TVM_FFI_ICHECK(sym::ramp(base, 1, value_ty.lanes()).Match(index)) << "LowerWarpMemory failed due to store index=" << index << ", can only handle continuous store"; UpdatePattern(base.Eval()); @@ -181,8 +181,7 @@ class WarpStoreCoeffFinder : public StmtExprVisitor { } void UpdatePattern(const PrimExpr& index) { - ffi::Array m = - arith::DetectLinearEquation(index, {warp_index_.as_or_throw()}); + ffi::Array m = sym::DetectLinearEquation(index, {warp_index_.as_or_throw()}); TVM_FFI_ICHECK_EQ(m.size(), 2U) << "LowerWarpMemory failed. Could not simplify the store index `" << index << "` into the form ax + by + cz + ... Warp memory is approximated by storing values in " @@ -210,7 +209,7 @@ class WarpStoreCoeffFinder : public StmtExprVisitor { // the coefficient int64_t warp_coeff_{0}; // analyzer. - arith::AnalyzerObj* analyzer_; + sym::AnalyzerObj* analyzer_; }; // Visitor to find the warp index @@ -267,7 +266,7 @@ class WarpAccessRewriter : public StmtExprMutator { public: using StmtExprMutator::Mutate; using StmtExprMutator::Mutate_; - explicit WarpAccessRewriter(int warp_size, arith::AnalyzerObj* analyzer) + explicit WarpAccessRewriter(int warp_size, sym::AnalyzerObj* analyzer) : warp_size_(warp_size), analyzer_(analyzer) {} // Rewrite the AllocBuffer statement which transforms // warp memory to local memory. @@ -439,8 +438,8 @@ class WarpAccessRewriter : public StmtExprMutator { std::pair SplitIndexByGroup(const PrimExpr& index) { PrimType index_ty = index.ty(); if (index_ty.lanes() != 1) { - arith::PVar base; - TVM_FFI_ICHECK(arith::ramp(base, 1, index_ty.lanes()).Match(index)); + sym::PVar base; + TVM_FFI_ICHECK(sym::ramp(base, 1, index_ty.lanes()).Match(index)); auto [local_index, group] = SplitIndexByGroup(base.Eval()); local_index = prim::Ramp(local_index, IntImm(local_index.ty(), 1), index_ty.lanes()); @@ -478,7 +477,7 @@ class WarpAccessRewriter : public StmtExprMutator { // the coefficient n int warp_group_{0}; // Internal analyzer - arith::AnalyzerObj* analyzer_; + sym::AnalyzerObj* analyzer_; }; // Bind bound information of variables to make analyzer more effective @@ -490,7 +489,7 @@ class BindVarBoundInfo : public StmtExprVisitor { if (value.as()) return std::nullopt; return StmtExprVisitor::Visit(value); } - explicit BindVarBoundInfo(arith::AnalyzerObj* analyzer) : analyzer_(analyzer) {} + explicit BindVarBoundInfo(sym::AnalyzerObj* analyzer) : analyzer_(analyzer) {} ffi::Optional Visit_(const ForNode* op) final { const Var& loop_var = op->loop_var; @@ -513,7 +512,7 @@ class BindVarBoundInfo : public StmtExprVisitor { protected: // internal analyzer. - arith::AnalyzerObj* analyzer_; + sym::AnalyzerObj* analyzer_; // variable domain std::unordered_map var_dom_; }; @@ -571,7 +570,7 @@ class WarpMemoryRewriter : public StmtExprMutator { } int warp_size_{0}; - arith::Analyzer analyzer_; + sym::Analyzer analyzer_; // variable domain std::unordered_map var_dom_; }; diff --git a/src/tirx/transform/narrow_datatype.cc b/src/tirx/transform/narrow_datatype.cc index 9fa93f78d881..a59d768ec5f7 100644 --- a/src/tirx/transform/narrow_datatype.cc +++ b/src/tirx/transform/narrow_datatype.cc @@ -22,11 +22,11 @@ * \brief narrow the datatype of indexing vars */ -#include #include #include #include #include +#include #include #include #include @@ -61,8 +61,8 @@ using namespace tvm::prim; // - Use DataTypeVisitor to determine whether a Var can be narrowed or not. // - Use DataTypeRewritter to rewrite the components of an indexing expression. -using arith::Analyzer; -using arith::ConstIntBound; +using sym::Analyzer; +using sym::ConstIntBound; // Determine the result dtype for Var, IntImm and Cast, // which will be stored in `vmap` eventually. @@ -187,7 +187,7 @@ class DataTypeVisitor final : public StmtExprVisitor { protected: // internal analyzer - arith::Analyzer analyzer_; + sym::Analyzer analyzer_; private: // the maximum possible bits, which serves as an init value @@ -199,7 +199,7 @@ class DataTypeVisitor final : public StmtExprVisitor { // the extent of vars to be rewritten std::unordered_map vextent_; // the memorized bound generated by ConstIntBoundAnalyzer - arith::ConstIntBoundAnalyzer::BoundMapType bound_; + sym::ConstIntBoundAnalyzer::BoundMapType bound_; }; class NarrowDataTypeRewriter : public IndexDataTypeRewriter { diff --git a/src/tirx/transform/remove_no_op.cc b/src/tirx/transform/remove_no_op.cc index 5c6d90d0397a..2a26a32ac546 100644 --- a/src/tirx/transform/remove_no_op.cc +++ b/src/tirx/transform/remove_no_op.cc @@ -21,12 +21,12 @@ * \file remove_no_op.cc * \brief Remove no op from the stmt */ -#include #include #include #include #include #include +#include #include #include #include @@ -35,7 +35,7 @@ #include -#include "../../arith/const_fold.h" +#include "../../sym/const_fold.h" #include "../analysis/var_use_def_analysis.h" #include "../ir/ir_mutator_with_analyzer.h" #include "ir_utils.h" @@ -78,7 +78,7 @@ class NoOpRemover : public IRMutatorWithAnalyzer { public: using IRMutatorWithAnalyzer::Mutate; using IRMutatorWithAnalyzer::Mutate_; - static Stmt Apply(Stmt stmt, const arith::Analyzer& analyzer, bool ignore_profiler_call = false) { + static Stmt Apply(Stmt stmt, const sym::Analyzer& analyzer, bool ignore_profiler_call = false) { auto visitor = ffi::make_object(analyzer, ignore_profiler_call); return visitor->Mutate(stmt, InplaceMode::kAllow).ValueOrUnchanged(stmt); } @@ -87,7 +87,7 @@ class NoOpRemover : public IRMutatorWithAnalyzer { using Parent = IRMutatorWithAnalyzer; public: - NoOpRemover(const arith::Analyzer& analyzer, bool ignore_profiler_call = false) + NoOpRemover(const sym::Analyzer& analyzer, bool ignore_profiler_call = false) : Parent(analyzer), ignore_profiler_call_(ignore_profiler_call) {} private: @@ -97,7 +97,7 @@ class NoOpRemover : public IRMutatorWithAnalyzer { } 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; + sym::Analyzer ana; if (ana->CanProve(wait_cnt < 0)) { // A negative wait count can arise if it depends on a loop variable. // For example, a wait count 1 - i can be negative after loop unrolling. @@ -144,12 +144,12 @@ class NoOpRemover : public IRMutatorWithAnalyzer { } } UnchangedOr Mutate_(const ForNode* op, InplaceMode inplace_mode) final { - auto extent_range = arith::EvalSet(op->extent, var_range_map_); - if (!arith::is_neg_inf(extent_range.max()) && !arith::is_pos_inf(extent_range.max()) && + auto extent_range = sym::EvalSet(op->extent, var_range_map_); + if (!sym::is_neg_inf(extent_range.max()) && !sym::is_pos_inf(extent_range.max()) && analyzer_->CanProve(extent_range.max() <= 0)) { return Evaluate(0); } - var_range_map_[op->loop_var.get()] = arith::IntSet::FromMinExtent(op->min, op->extent); + var_range_map_[op->loop_var.get()] = sym::IntSet::FromMinExtent(op->min, op->extent); Stmt stmt = Parent::Mutate_(op, inplace_mode).ValueOrUnchanged(ffi::GetRef(op)); var_range_map_.erase(op->loop_var.get()); op = stmt.as(); @@ -267,11 +267,11 @@ class NoOpRemover : public IRMutatorWithAnalyzer { } } - std::unordered_map var_range_map_; + std::unordered_map var_range_map_; bool ignore_profiler_call_{false}; }; -Stmt RemoveNoOp(Stmt stmt, const arith::Analyzer& analyzer, bool ignore_profiler_call) { +Stmt RemoveNoOp(Stmt stmt, const sym::Analyzer& analyzer, bool ignore_profiler_call) { return NoOpRemover::Apply(std::move(stmt), analyzer, ignore_profiler_call); } @@ -283,7 +283,7 @@ Pass RemoveNoOp() { ctx->GetConfig("tirx.RemoveNoOp") .value_or(tvm::transform::PassConfigWithDefaults()); - arith::Analyzer analyzer; + sym::Analyzer analyzer; analyzer->rewrite_simplify.SetMaximumRewriteSteps(config->max_simplification_steps); bool ignore_profiler_call = config->ignore_profiler_call; diff --git a/src/tirx/transform/remove_no_op.h b/src/tirx/transform/remove_no_op.h index 03a9823aa6f6..7022d3be28ad 100644 --- a/src/tirx/transform/remove_no_op.h +++ b/src/tirx/transform/remove_no_op.h @@ -24,7 +24,7 @@ #ifndef TVM_TIR_TRANSFORM_REMOVE_NO_OP_H_ #define TVM_TIR_TRANSFORM_REMOVE_NO_OP_H_ -#include +#include #include namespace tvm { @@ -41,7 +41,7 @@ namespace tirx { * * \return The modified statement with no-ops removed */ -Stmt RemoveNoOp(Stmt stmt, const arith::Analyzer& analyzer, bool ignore_profiler_call = false); +Stmt RemoveNoOp(Stmt stmt, const sym::Analyzer& analyzer, bool ignore_profiler_call = false); } // namespace tirx } // namespace tvm diff --git a/src/tirx/transform/stmt_simplify.cc b/src/tirx/transform/stmt_simplify.cc index 6af72187811b..23908afd2c60 100644 --- a/src/tirx/transform/stmt_simplify.cc +++ b/src/tirx/transform/stmt_simplify.cc @@ -24,13 +24,13 @@ #include "stmt_simplify.h" -#include #include #include #include #include #include #include +#include #include #include #include @@ -59,19 +59,19 @@ void StmtSimplifyConfigNode::RegisterReflection() { refl::DefaultValue(false)); } -arith::RewriteSimplifier::Extension StmtSimplifyConfigNode::GetEnabledExtensions() const { - arith::RewriteSimplifier::Extension flags = arith::RewriteSimplifier::kNone; +sym::RewriteSimplifier::Extension StmtSimplifyConfigNode::GetEnabledExtensions() const { + sym::RewriteSimplifier::Extension flags = sym::RewriteSimplifier::kNone; if (transitively_prove_inequalities) { - flags = arith::RewriteSimplifier::Extension( - flags | arith::RewriteSimplifier::kTransitivelyProveInequalities); + flags = sym::RewriteSimplifier::Extension( + flags | sym::RewriteSimplifier::kTransitivelyProveInequalities); } if (convert_boolean_to_and_of_ors) { - flags = arith::RewriteSimplifier::Extension( - flags | arith::RewriteSimplifier::kConvertBooleanToAndOfOrs); + flags = sym::RewriteSimplifier::Extension(flags | + sym::RewriteSimplifier::kConvertBooleanToAndOfOrs); } if (apply_constraints_to_boolean_branches) { - flags = arith::RewriteSimplifier::Extension( - flags | arith::RewriteSimplifier::kApplyConstraintsToBooleanBranches); + flags = sym::RewriteSimplifier::Extension( + flags | sym::RewriteSimplifier::kApplyConstraintsToBooleanBranches); } return flags; } @@ -84,7 +84,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { StmtSimplifyConfigNode::RegisterReflection(); } TVM_REGISTER_PASS_CONFIG_OPTION("tirx.StmtSimplify", StmtSimplifyConfig); -PrimFunc StmtSimplifier::Apply(PrimFunc func, const arith::Analyzer& analyzer, +PrimFunc StmtSimplifier::Apply(PrimFunc func, const sym::Analyzer& analyzer, ffi::Optional config_opt) { auto config = config_opt.value_or(MakeDefaultStmtSimplifyConfig()); @@ -114,9 +114,9 @@ UnchangedOr StmtSimplifier::Mutate(ffi::AnyView input, InplaceMode inp UnchangedOr StmtSimplifier::Mutate_(const ForNode* op, InplaceMode inplace_mode) { analyzer_->Bind(op->loop_var, Range::FromMinExtent(op->min, op->extent)); - With ctx1(analyzer_, op->loop_var >= op->min); - With ctx2(analyzer_, - static_cast(op->loop_var) < op->min + op->extent); + With ctx1(analyzer_, op->loop_var >= op->min); + With ctx2(analyzer_, + static_cast(op->loop_var) < op->min + op->extent); return Parent::Mutate_(op, inplace_mode); } @@ -215,7 +215,7 @@ ffi::Optional StmtSimplifier::ProveCondition(PrimExpr condition) const { } } -PrimFunc StmtSimplify(PrimFunc func, const arith::Analyzer& analyzer) { +PrimFunc StmtSimplify(PrimFunc func, const sym::Analyzer& analyzer) { return StmtSimplifier::Apply(std::move(func), analyzer); } @@ -223,7 +223,7 @@ namespace transform { Pass StmtSimplify() { auto pass_func = [](PrimFunc f, IRModule m, PassContext ctx) { - arith::Analyzer analyzer; + sym::Analyzer analyzer; auto cfg = ctx->GetConfig("tirx.StmtSimplify"); return StmtSimplifier::Apply(f, analyzer, cfg); diff --git a/src/tirx/transform/stmt_simplify.h b/src/tirx/transform/stmt_simplify.h index 4b48c61f4a99..37694283e1c3 100644 --- a/src/tirx/transform/stmt_simplify.h +++ b/src/tirx/transform/stmt_simplify.h @@ -24,7 +24,7 @@ #ifndef TVM_TIR_TRANSFORM_STMT_SIMPLIFY_H_ #define TVM_TIR_TRANSFORM_STMT_SIMPLIFY_H_ -#include +#include #include #include "../ir/ir_mutator_with_analyzer.h" @@ -41,7 +41,7 @@ struct StmtSimplifyConfigNode : public ffi::Object { TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.transform.StmtSimplifyConfig", StmtSimplifyConfigNode, ffi::Object); - arith::RewriteSimplifier::Extension GetEnabledExtensions() const; + sym::RewriteSimplifier::Extension GetEnabledExtensions() const; }; class StmtSimplifyConfig : public ffi::ObjectRef { @@ -54,15 +54,15 @@ class StmtSimplifier : public IRMutatorWithAnalyzer { public: using IRMutatorWithAnalyzer::Mutate; using IRMutatorWithAnalyzer::Mutate_; - static PrimFunc Apply(PrimFunc func, const arith::Analyzer& analyzer, + static PrimFunc Apply(PrimFunc func, const sym::Analyzer& analyzer, ffi::Optional config_opt = std::nullopt); - explicit StmtSimplifier(const arith::Analyzer& analyzer, StmtSimplifyConfig config) + explicit StmtSimplifier(const sym::Analyzer& analyzer, StmtSimplifyConfig config) : IRMutatorWithAnalyzer(analyzer), config_(config) {} protected: using Parent = IRMutatorWithAnalyzer; - StmtSimplifier(const VTable* vtable, const arith::Analyzer& analyzer, StmtSimplifyConfig config) + StmtSimplifier(const VTable* vtable, const sym::Analyzer& analyzer, StmtSimplifyConfig config) : Parent(analyzer.get(), vtable), config_(config) {} PrimFunc Run(PrimFunc func); @@ -97,7 +97,7 @@ class StmtSimplifier : public IRMutatorWithAnalyzer { * * Applies the same behavior as the tirx.transform.StmtSimplify pass. */ -PrimFunc StmtSimplify(PrimFunc func, const arith::Analyzer& analyzer); +PrimFunc StmtSimplify(PrimFunc func, const sym::Analyzer& analyzer); } // namespace tirx } // namespace tvm diff --git a/src/tirx/transform/storage_rewrite.cc b/src/tirx/transform/storage_rewrite.cc index 44cdd6491ac5..d6a6d7bdd275 100644 --- a/src/tirx/transform/storage_rewrite.cc +++ b/src/tirx/transform/storage_rewrite.cc @@ -22,7 +22,6 @@ * \brief Memory access pattern analysis and optimization. * Re-write data access to enable memory sharing when possible. */ -#include #include #include #include @@ -32,6 +31,7 @@ #include #include #include +#include #include #include #include @@ -43,8 +43,8 @@ #include #include -#include "../../arith/int_operator.h" #include "../../runtime/thread_storage_scope.h" +#include "../../sym/int_operator.h" #include "../ir/buffer_common.h" #include "ir_utils.h" @@ -1255,7 +1255,7 @@ class StoragePlanRewriter : public StmtExprMutator { // Physical roots of buffer aliases, flattened by LinearAccessPatternFinder. ffi::Map buffer_aliases_; // analyzer - arith::Analyzer analyzer_; + sym::Analyzer analyzer_; }; /* Helper struct containing information on how a buffer is declared and used @@ -1331,8 +1331,8 @@ struct BufferVarInfo { return element_dtype; } } - arith::Analyzer analyzer_; - arith::ModularSet me = analyzer_->modular_set(extent); + sym::Analyzer analyzer_; + sym::ModularSet me = analyzer_->modular_set(extent); if ((me->coeff % lanes == 0) && (me->base % lanes == 0)) { preferred_lanes = lanes; } @@ -1571,7 +1571,7 @@ class VectorTypeAccessChecker : public StmtExprVisitor { if (ramp_index && is_one(ramp_index->stride)) { if (ramp_index->lanes->IsInstance()) { int lanes = ramp_index->lanes.as_or_throw()->value.as().value(); - arith::ModularSet me = analyzer_->modular_set(ramp_index->base); + sym::ModularSet me = analyzer_->modular_set(ramp_index->base); if ((me->coeff % lanes == 0) && (me->base % lanes == 0)) { lanes_used = lanes; } @@ -1582,7 +1582,7 @@ class VectorTypeAccessChecker : public StmtExprVisitor { if (detect_scalar_read_patterns_ && is_buffer_load && indices.size()) { const PrimExpr last_dim_index = indices[indices.size() - 1]; if (last_dim_index.ty().lanes() == 1) { - arith::ModularSet me = analyzer_->modular_set(last_dim_index); + sym::ModularSet me = analyzer_->modular_set(last_dim_index); // Fixed lane counts use 15 bits; retain the scalar access when the modular // coefficient cannot be represented by the vector type. if (int64_t lanes = me->coeff; lanes >= 0 && lanes <= 32767) { @@ -1620,7 +1620,7 @@ class VectorTypeAccessChecker : public StmtExprVisitor { bool detect_scalar_read_patterns_{true}; // internal analyzer - arith::Analyzer analyzer_; + sym::Analyzer analyzer_; }; /* \brief Rewrites buffer/pointer variables from scalar types to vectorized @@ -1761,7 +1761,7 @@ class VectorTypeRewriter : public StmtExprMutator { } indices.Set(indices.size() - 1, new_index); } else if (last_dim_index.ty().lanes() == 1 && info.factor() > 1) { - arith::ModularSet me = analyzer_->modular_set(last_dim_index); + sym::ModularSet me = analyzer_->modular_set(last_dim_index); TVM_FFI_ICHECK(me->coeff == 0 || info.factor() % me->coeff == 0); PrimExpr new_index = last_dim_index / MakeConst(last_dim_index.ty(), info.factor()); shuffle_index = me->base % info.factor(); @@ -1806,7 +1806,7 @@ class VectorTypeRewriter : public StmtExprMutator { } indices.Set(indices.size() - 1, new_index); } else if (last_dim_index.ty().lanes() == 1 && info.factor() > 1) { - arith::ModularSet me = analyzer_->modular_set(last_dim_index); + sym::ModularSet me = analyzer_->modular_set(last_dim_index); TVM_FFI_ICHECK(me->coeff == 0 || info.factor() % me->coeff == 0); PrimExpr new_index = last_dim_index / MakeConst(last_dim_index.ty(), info.factor()); shuffle_index = me->base % info.factor(); @@ -2096,7 +2096,7 @@ class VectorTypeRewriter : public StmtExprMutator { bool rewrite_indices_{true}; std::unordered_map rewrite_map_; const ffi::Map& buffer_aliases_; - arith::Analyzer analyzer_; + sym::Analyzer analyzer_; }; // Rewrite allocates, pointer parameters, and buffer parameters into vectorized versions diff --git a/src/tirx/transform/tile_primitive_dispatch.cc b/src/tirx/transform/tile_primitive_dispatch.cc index a272c2fb2469..9de41412a39b 100644 --- a/src/tirx/transform/tile_primitive_dispatch.cc +++ b/src/tirx/transform/tile_primitive_dispatch.cc @@ -23,11 +23,11 @@ * declarations and emits launch params). */ -#include -#include #include #include #include +#include +#include #include #include #include @@ -877,14 +877,14 @@ class TilePrimitiveDispatcher : public StmtExprMutator { struct ScopeIdRange { ScopeIdTarget target; - int64_t lo = arith::ConstIntBound::kNegInf; - int64_t hi = arith::ConstIntBound::kPosInf; + int64_t lo = sym::ConstIntBound::kNegInf; + int64_t hi = sym::ConstIntBound::kPosInf; }; struct PendingRangeGroup { ScopeIdTarget target; - int64_t lo = arith::ConstIntBound::kNegInf; - int64_t hi = arith::ConstIntBound::kPosInf; + int64_t lo = sym::ConstIntBound::kNegInf; + int64_t hi = sym::ConstIntBound::kPosInf; std::vector indices; }; @@ -1073,7 +1073,7 @@ class TilePrimitiveDispatcher : public StmtExprMutator { int64_t* base) { PrimExpr simplified = analyzer_->Simplify(diff); for (const auto& [var, candidate] : ScopeIdTargets()) { - ffi::Array linear = arith::DetectLinearEquation(simplified, {var}); + ffi::Array linear = sym::DetectLinearEquation(simplified, {var}); if (linear.size() != 2) continue; int64_t c = 0; int64_t b = 0; @@ -1098,8 +1098,8 @@ class TilePrimitiveDispatcher : public StmtExprMutator { if (!TryExtractLinearScopeDiff(lhs - rhs, &target, &coeff, &base)) return false; // Interpret `coeff * v + base 0` where coeff is +/- 1. - int64_t lo = arith::ConstIntBound::kNegInf; - int64_t hi = arith::ConstIntBound::kPosInf; + int64_t lo = sym::ConstIntBound::kNegInf; + int64_t hi = sym::ConstIntBound::kPosInf; if (lhs_less_rhs) { if (coeff == 1) { // v + base < 0 -> v < -base @@ -1517,7 +1517,7 @@ class TilePrimitiveDispatcher : public StmtExprMutator { } ffi::Map var_range_map_; - arith::Analyzer analyzer_; + sym::Analyzer analyzer_; const Target& target_; // List of ScopeIdDefs visible at each nesting level (one entry for the // device-entry body itself, plus one per ScopeIdDefStmt-bearing region). diff --git a/src/tirx/transform/tvm_ffi_binder.h b/src/tirx/transform/tvm_ffi_binder.h index 5cbc180dae76..4adbd0d4a73c 100644 --- a/src/tirx/transform/tvm_ffi_binder.h +++ b/src/tirx/transform/tvm_ffi_binder.h @@ -28,9 +28,9 @@ #ifndef TVM_TIR_TRANSFORM_TVM_FFI_BINDER_H_ #define TVM_TIR_TRANSFORM_TVM_FFI_BINDER_H_ -#include #include #include +#include #include #include @@ -394,7 +394,7 @@ class TVMFFIABIBuilder { /*! \brief Deferred constant-expression assertions for display-var substitution. */ std::vector pending_const_asserts_; /*! \brief internal analyzer. */ - arith::Analyzer analyzer_; + sym::Analyzer analyzer_; // Function metadata /*! \brief function name for error messages. */ diff --git a/src/tirx/transform/unroll_loop.cc b/src/tirx/transform/unroll_loop.cc index a02faa9669d0..647c8da50697 100644 --- a/src/tirx/transform/unroll_loop.cc +++ b/src/tirx/transform/unroll_loop.cc @@ -22,12 +22,12 @@ * \file unroll_loop.cc */ // Unrolls the loop as in Halide pipeline. -#include #include #include #include #include #include +#include #include #include #include @@ -293,7 +293,7 @@ class LoopUnroller : public StmtExprMutator { // set of indices touched during visit local memory std::unordered_set var_touched_local_; // analyzer - arith::Analyzer analyzer_; + sym::Analyzer analyzer_; }; Stmt UnrollLoop(Stmt stmt, UnrollLoopConfig cfg) { diff --git a/src/tirx/transform/vectorize_loop.cc b/src/tirx/transform/vectorize_loop.cc index 0a3fe701f44d..f79553c5811f 100644 --- a/src/tirx/transform/vectorize_loop.cc +++ b/src/tirx/transform/vectorize_loop.cc @@ -21,7 +21,6 @@ * \file vectorize_loop.cc */ // Loop vectorizer as in Halide pipeline. -#include #include #include #include @@ -30,6 +29,7 @@ #include #include #include +#include #include #include #include @@ -447,7 +447,7 @@ class VecAllocAccess : public StmtExprMutator { // the lanes. PrimExpr var_lanes_; // Analyzer for simplifications - arith::Analyzer analyzer_; + sym::Analyzer analyzer_; }; // Vectorization supplies its own dtype-aware expression traversal. @@ -1144,7 +1144,7 @@ class Vectorizer : public StmtExprMutator { private: // analyzer - arith::Analyzer analyzer_; + sym::Analyzer analyzer_; // deep equal prim::ExprDeepEqual deep_equal_; // variable to be replaced diff --git a/tests/cpp/pattern_match_test.cc b/tests/cpp/pattern_match_test.cc index 20854b39fc92..c1b45328a488 100644 --- a/tests/cpp/pattern_match_test.cc +++ b/tests/cpp/pattern_match_test.cc @@ -17,19 +17,19 @@ * under the License. */ -#include "../src/arith/pattern_match.h" +#include "../src/sym/pattern_match.h" #include #include TEST(Pattern, Basic) { using namespace tvm; - using namespace tvm::arith; + using namespace tvm::sym; tvm::PrimVar x("x"), y("y"), z("z"); PrimExpr scalable_lanes = prim::Mul(Call(PrimType::Int(32), prim::builtin::vscale(), {}), 4); - arith::PVar px, py, pz; - arith::PVar pt; - arith::PVar planes; + sym::PVar px, py, pz; + sym::PVar pt; + sym::PVar planes; // arithmetics auto r = 1 + (y + 1); @@ -130,8 +130,8 @@ TEST(Pattern, Basic) { TEST(Pattern, IntImm) { using namespace tvm; PrimVar tx("tx"), ty("ty"); - arith::PVar c; - arith::PVar v; + sym::PVar c; + sym::PVar v; { // We can match integer and Var, both of which are // special case container of Expr @@ -148,7 +148,7 @@ TEST(Pattern, IntImm) { TEST(Pattern, MatchWithType) { using namespace tvm; // match expr with specified dtype - arith::PVarWithDataType> pat(DLDataType{kDLFloat, 32, 1}); + sym::PVarWithDataType> pat(DLDataType{kDLFloat, 32, 1}); PrimVar x("x", PrimType::Float(32)); PrimVar y("y", PrimType::Float(32)); PrimVar x_int("x", PrimType::Int(32)); @@ -157,8 +157,8 @@ TEST(Pattern, MatchWithType) { TVM_FFI_ICHECK(!pat.Match(x_int + y_int * 2)); // match vectorized expr with specified element dtype - arith::PVecDataType vec_ty(DLDataType{kDLFloat, 32, 1}); - arith::PVarWithDataType vpat(vec_ty); + sym::PVecDataType vec_ty(DLDataType{kDLFloat, 32, 1}); + sym::PVarWithDataType vpat(vec_ty); PrimVar vx("x", PrimType::Float(32, 8)); PrimVar vy("y", PrimType::Float(32, 8)); PrimVar vx_int("x", PrimType::Int(32, 8)); diff --git a/tests/cpp/arith_simplify_test.cc b/tests/cpp/sym_simplify_test.cc similarity index 87% rename from tests/cpp/arith_simplify_test.cc rename to tests/cpp/sym_simplify_test.cc index 1e377ddd9ec4..b8d9d2a81b51 100644 --- a/tests/cpp/arith_simplify_test.cc +++ b/tests/cpp/sym_simplify_test.cc @@ -18,14 +18,14 @@ */ #include -#include #include #include +#include #include #include TEST(Simplify, MinMax) { - tvm::arith::Analyzer ana; + tvm::sym::Analyzer ana; auto x = tvm::te::var("x"); auto e1 = (tvm::max(x, 1) - tvm::max(x, 1)); auto e1s = ana->canonical_simplify(e1); @@ -37,7 +37,7 @@ TEST(Simplify, MinMax) { } TEST(Simplify, Mul) { - tvm::arith::Analyzer ana; + tvm::sym::Analyzer ana; auto x = tvm::te::var("x"); auto e = (x * x) - (x * x); auto es = ana->canonical_simplify(e); @@ -45,7 +45,7 @@ TEST(Simplify, Mul) { } TEST(Simplify, Mod) { - tvm::arith::Analyzer ana; + tvm::sym::Analyzer ana; auto x = tvm::IntImm::Int32(10); auto y = tvm::IntImm::Int32(12); // Mod::make is used instead of % to avoid constant folding during @@ -57,8 +57,8 @@ TEST(Simplify, Mod) { } TEST(AnalyzerObjectRef, CopySharesMutableState) { - tvm::arith::Analyzer analyzer; - tvm::arith::Analyzer copy = analyzer; + tvm::sym::Analyzer analyzer; + tvm::sym::Analyzer copy = analyzer; auto x = tvm::te::var("x"); copy->Bind(x, tvm::Range::FromMinExtent(0, 8)); @@ -67,8 +67,8 @@ TEST(AnalyzerObjectRef, CopySharesMutableState) { } TEST(AnalyzerObjectRef, ConstHandleRefCanMutateAnalyzerState) { - tvm::arith::Analyzer analyzer; - const tvm::arith::Analyzer& analyzer_ref = analyzer; + tvm::sym::Analyzer analyzer; + const tvm::sym::Analyzer& analyzer_ref = analyzer; auto x = tvm::te::var("x"); analyzer_ref->Bind(x, tvm::Range::FromMinExtent(0, 8)); @@ -77,19 +77,19 @@ TEST(AnalyzerObjectRef, ConstHandleRefCanMutateAnalyzerState) { } TEST(AnalyzerObjectRef, CloneIsIndependent) { - tvm::arith::Analyzer analyzer; + tvm::sym::Analyzer analyzer; auto x = tvm::te::var("x"); auto y = tvm::te::var("y"); analyzer->Bind(x, tvm::Range::FromMinExtent(0, 8)); - analyzer->modular_set.Update(x, tvm::arith::ModularSet(4, 0)); + analyzer->modular_set.Update(x, tvm::sym::ModularSet(4, 0)); - tvm::arith::Analyzer clone = analyzer->Clone(); + tvm::sym::Analyzer clone = analyzer->Clone(); TVM_FFI_ICHECK(clone->CanProve(x < 8)); TVM_FFI_ICHECK(clone->modular_set(x)->coeff == 4); clone->Bind(y, tvm::Range::FromMinExtent(0, 4)); - clone->modular_set.Update(x, tvm::arith::ModularSet(8, 0), true); + clone->modular_set.Update(x, tvm::sym::ModularSet(8, 0), true); TVM_FFI_ICHECK(clone->CanProve(y < 4)); TVM_FFI_ICHECK(!analyzer->CanProve(y < 4)); TVM_FFI_ICHECK(analyzer->CanProve(x < 8)); @@ -100,7 +100,7 @@ TEST(AnalyzerObjectRef, CloneIsIndependent) { TEST(Simplify, AssumeConstraintKeepsBufferLoadStable) { using namespace tvm; - arith::Analyzer analyzer; + sym::Analyzer analyzer; tirx::BufferVar buffer = tirx::decl_buffer({1}, PrimType::Int(32)); PrimExpr load = tirx::BufferLoad(buffer, {IntImm::Int32(0)}); PrimExpr constraint = load > 0; @@ -117,7 +117,7 @@ TEST(Simplify, AssumeConstraintKeepsBufferLoadStable) { } { - With scope(analyzer, constraint, true); + With scope(analyzer, constraint, true); if (analyzer->z3_prover.IsEnabled()) { EXPECT_TRUE(analyzer->z3_prover.CanProve(constraint)); } @@ -128,7 +128,7 @@ TEST(Simplify, AssumeConstraintKeepsBufferLoadStable) { } { - With scope(analyzer, constraint); + With scope(analyzer, constraint); if (analyzer->z3_prover.IsEnabled()) { EXPECT_FALSE(analyzer->z3_prover.CanProve(constraint)); } diff --git a/tests/python/ir/test_node_reflection.py b/tests/python/ir/test_node_reflection.py index f1aad6b98f9e..491ba2f876f7 100644 --- a/tests/python/ir/test_node_reflection.py +++ b/tests/python/ir/test_node_reflection.py @@ -39,6 +39,29 @@ def test_const_saveload_json(): tvm.ir.assert_structural_equal(zz, z, map_free_vars=True) +def test_symbolic_analysis_legacy_json_load(): + # Historical symbolic-analysis objects must load through the existing JSON upgrader. + graph = { + "root_index": 2, + "nodes": [ + {"type": "arith.ConstIntBound", "data": {"min_value": 2, "max_value": 7}}, + {"type": "arith.ModularSet", "data": {"coeff": 4, "base": 1}}, + {"type": "ffi.Array", "data": [0, 1, 0]}, + ], + "metadata": {"tvm_version": tvm.__version__}, + } + restored = tvm.ir.load_json(json.dumps(graph)) + bound, modular, shared_bound = restored + assert isinstance(bound, tvm.sym.ConstIntBound) + assert (bound.min_value, bound.max_value) == (2, 7) + assert isinstance(modular, tvm.sym.ModularSet) + assert (modular.coeff, modular.base) == (4, 1) + assert bound.same_as(shared_bound) + saved_types = {node["type"] for node in json.loads(tvm.ir.save_json(restored))["nodes"]} + assert {"sym.ConstIntBound", "sym.ModularSet"} <= saved_types + assert not any(key.startswith("arith.") for key in saved_types) + + def test_save_json_metadata_version(): obj = tvm.runtime.convert([1, 2]) json_str = tvm.ir.save_json(obj) diff --git a/tests/python/relax/test_analysis_suggest_layout_transforms.py b/tests/python/relax/test_analysis_suggest_layout_transforms.py index 21e2f0b2d050..f2a8b224e4b2 100644 --- a/tests/python/relax/test_analysis_suggest_layout_transforms.py +++ b/tests/python/relax/test_analysis_suggest_layout_transforms.py @@ -777,7 +777,7 @@ def expected( tvm.ir.assert_structural_equal(after, expected) -@pytest.mark.skip("temp disable, due to minor arith regression") +@pytest.mark.skip("temp disable, due to minor sym regression") def test_op_split_tiling_split_dim(): @T.prim_func(private=True, s_tir=True) def before( diff --git a/tests/python/relax/test_frontend_from_exported_program.py b/tests/python/relax/test_frontend_from_exported_program.py index d78629f5737b..22d5fee0d707 100644 --- a/tests/python/relax/test_frontend_from_exported_program.py +++ b/tests/python/relax/test_frontend_from_exported_program.py @@ -5580,7 +5580,7 @@ def forward(self, x, y): x_shape = mod["main"].params[0].ty.shape.values y_shape = mod["main"].params[1].ty.shape.values - assert tvm.arith.Analyzer().can_prove_equal(y_shape[1], x_shape[1] * 2) + assert tvm.sym.Analyzer().can_prove_equal(y_shape[1], x_shape[1] * 2) def test_expand_with_new_leading_dimension(): @@ -5598,9 +5598,9 @@ def forward(self, x): input_shape = mod["main"].params[0].ty.shape.values output_shape = mod["main"].ret_ty.fields[0].shape.values - assert tvm.arith.Analyzer().can_prove_equal(output_shape[0], 2) - assert tvm.arith.Analyzer().can_prove_equal(output_shape[1], input_shape[0]) - assert tvm.arith.Analyzer().can_prove_equal(output_shape[2], input_shape[1]) + assert tvm.sym.Analyzer().can_prove_equal(output_shape[0], 2) + assert tvm.sym.Analyzer().can_prove_equal(output_shape[1], input_shape[0]) + assert tvm.sym.Analyzer().can_prove_equal(output_shape[2], input_shape[1]) def test_dynamic_scalar_item_in_shape_operations(): diff --git a/tests/python/relax/test_op_manipulate.py b/tests/python/relax/test_op_manipulate.py index 42ddece4c9e6..38caaf83881f 100644 --- a/tests/python/relax/test_op_manipulate.py +++ b/tests/python/relax/test_op_manipulate.py @@ -2083,7 +2083,7 @@ def test_split_infer_ty(): # All relax shape variables are non-negative. When a scope # begins, any TIR variables that are used as shape variables are - # declared to be non-negative `tvm.arith.Analyzer`. Because + # declared to be non-negative `tvm.sym.Analyzer`. Because # `relax.op.split` clamps the indices to be within the bounds of # the axis being split, simplifying with non-negative shape # variables can result in much simpler shapes. 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 378ee7f3142a..1ba813ccaaf4 100644 --- a/tests/python/s_tir/analysis/test_sblock_access_region.py +++ b/tests/python/s_tir/analysis/test_sblock_access_region.py @@ -364,7 +364,7 @@ def test_access_of_padding_pattern(): def do_compare_buffer_region(region, expect): assert region.source == expect.source - analyzer = tvm.arith.Analyzer() + analyzer = tvm.sym.Analyzer() for observed_range, expected_range in zip(region.region, expect.region): analyzer.can_prove_equal(observed_range.min, expected_range.min) analyzer.can_prove_equal(observed_range.extent, expected_range.extent) @@ -459,7 +459,7 @@ def func( @pytest.mark.parametrize("case", ["coupled", "equal", "nonlinear", "empty", "unbounded", "rounded"]) def test_conditional_inequality_access_regions(case): - # Retain the live cases from the former arith inequality solver tests through + # Retain the live cases from the former sym inequality solver tests through # the block-access consumer, including its conservative unresolved fallback. tirx = tvm.tirx x, y, z = [tirx.Var(name, "int32") for name in ("x", "y", "z")] diff --git a/tests/python/s_tir/dlight/test_cpu_reduction.py b/tests/python/s_tir/dlight/test_cpu_reduction.py index 28e60a1d449a..042a25feaaaa 100644 --- a/tests/python/s_tir/dlight/test_cpu_reduction.py +++ b/tests/python/s_tir/dlight/test_cpu_reduction.py @@ -191,11 +191,11 @@ def test_rvv_code_size_reduction(fast): ) -# The arith analyzer no longer proves vscale-bearing inequalities via +# The sym analyzer no longer proves vscale-bearing inequalities via # substitution (CanProveVscaleExpressionFromKnownValues was deleted). This # weakens simplification of scalable-vector index expressions, which can # prevent the RVV vectorization schedule from producing scalable vector ops. -@pytest.mark.xfail(reason="arith no longer proves vscale-bearing inequalities via substitution") +@pytest.mark.xfail(reason="sym no longer proves vscale-bearing inequalities via substitution") def test_rvv_fast_softmax_vectorizes_exp(): """fast_softmax + schedule should produce RVV vector instructions for the polynomial exp approximation (no scalar exp calls).""" diff --git a/tests/python/s_tir/transform/test_s_tir_transform_inject_virtual_thread.py b/tests/python/s_tir/transform/test_s_tir_transform_inject_virtual_thread.py index 315ec04de657..acecf088c509 100644 --- a/tests/python/s_tir/transform/test_s_tir_transform_inject_virtual_thread.py +++ b/tests/python/s_tir/transform/test_s_tir_transform_inject_virtual_thread.py @@ -241,7 +241,7 @@ def visitor(node): tvm_ffi.structural_walk(after.body, visitor) assert len(masked_calls) == 4 assert all(list(call.args[0].ty.shape) == [8] for call in masked_calls) - analyzer = tvm.arith.Analyzer() + analyzer = tvm.sym.Analyzer() assert sorted(int(analyzer.simplify(call.args[-2].base)) for call in masked_calls) == [ 0, 0, diff --git a/tests/python/arith/test_arith_analyzer_object.py b/tests/python/sym/test_sym_analyzer_object.py similarity index 86% rename from tests/python/arith/test_arith_analyzer_object.py rename to tests/python/sym/test_sym_analyzer_object.py index 9edd75d7aa7b..d74fd81ecd9a 100644 --- a/tests/python/arith/test_arith_analyzer_object.py +++ b/tests/python/sym/test_sym_analyzer_object.py @@ -20,12 +20,12 @@ import tvm import tvm.testing from tvm import tirx -from tvm.arith.analyzer import CompareResult, Extension from tvm.runtime import Object +from tvm.sym.analyzer import CompareResult, Extension def test_analyzer_is_ffi_object_with_persistent_state(): - analyzer = tvm.arith.Analyzer() + analyzer = tvm.sym.Analyzer() x = tirx.Var("x", "int64") assert isinstance(analyzer, Object) @@ -41,7 +41,7 @@ def test_analyzer_is_ffi_object_with_persistent_state(): def test_analyzer_object_constraint_scope_and_override_bind(): - analyzer = tvm.arith.Analyzer() + analyzer = tvm.sym.Analyzer() x = tirx.Var("x", "int64") with analyzer.constraint_scope(x % 3 == 0): @@ -49,7 +49,7 @@ def test_analyzer_object_constraint_scope_and_override_bind(): assert analyzer.modular_set(x).coeff != 3 - analyzer = tvm.arith.Analyzer() + analyzer = tvm.sym.Analyzer() y = tirx.Var("y", "int64") analyzer.bind(y, tirx.const(4, "int64")) tvm.ir.assert_structural_equal(analyzer.simplify(y + 1), tirx.const(5, "int64")) @@ -59,10 +59,10 @@ def test_analyzer_object_constraint_scope_and_override_bind(): def test_analyzer_object_update_const_int_bound(): - analyzer = tvm.arith.Analyzer() + analyzer = tvm.sym.Analyzer() x = tirx.Var("x", "int64") - analyzer.update(x, tvm.arith.ConstIntBound(2, 5)) + analyzer.update(x, tvm.sym.ConstIntBound(2, 5)) bound = analyzer.const_int_bound(x + 1) assert bound.min_value == 3 @@ -70,11 +70,11 @@ def test_analyzer_object_update_const_int_bound(): def test_analyzer_object_update_modular_set(): - analyzer = tvm.arith.Analyzer() + analyzer = tvm.sym.Analyzer() x = tirx.Var("x", "int32") assert analyzer.modular_set(x).coeff == 1 - analyzer.update(x, tvm.arith.ModularSet(4, 0)) + analyzer.update(x, tvm.sym.ModularSet(4, 0)) result = analyzer.modular_set(x) assert result.coeff == 4 @@ -82,10 +82,10 @@ def test_analyzer_object_update_modular_set(): def test_analyzer_object_update_int_set(): - analyzer = tvm.arith.Analyzer() + analyzer = tvm.sym.Analyzer() y = tirx.Var("y", "int32") - analyzer.update(y, tvm.arith.IntervalSet(0, 8)) + analyzer.update(y, tvm.sym.IntervalSet(0, 8)) int_set = analyzer.int_set(y) assert int_set.min_value.value == 0 @@ -93,7 +93,7 @@ def test_analyzer_object_update_int_set(): def test_analyzer_object_update_rejects_unknown_info(): - analyzer = tvm.arith.Analyzer() + analyzer = tvm.sym.Analyzer() y = tirx.Var("y", "int32") with pytest.raises(TypeError): @@ -101,7 +101,7 @@ def test_analyzer_object_update_rejects_unknown_info(): def test_analyzer_object_can_prove_comparison_predicates(): - analyzer = tvm.arith.Analyzer() + analyzer = tvm.sym.Analyzer() x = tirx.Var("x", "int32") analyzer.bind(x, tvm.ir.Range(0, 8)) @@ -112,16 +112,16 @@ def test_analyzer_object_can_prove_comparison_predicates(): def test_analyzer_object_update_const_int_bound_half_space(): - analyzer = tvm.arith.Analyzer() + analyzer = tvm.sym.Analyzer() n = tirx.Var("n", "int32") assert not analyzer.can_prove(n >= 0) - analyzer.update(n, tvm.arith.ConstIntBound(0, tvm.arith.ConstIntBound.POS_INF)) + analyzer.update(n, tvm.sym.ConstIntBound(0, tvm.sym.ConstIntBound.POS_INF)) assert analyzer.can_prove(n >= 0) def test_analyzer_object_int_set_from_bound_vars(): - analyzer = tvm.arith.Analyzer() + analyzer = tvm.sym.Analyzer() x = tirx.Var("x", "int32") analyzer.bind(x, tvm.ir.Range(0, 8)) @@ -135,19 +135,19 @@ def test_analyzer_object_set_maximum_rewrite_steps(): y = tirx.Var("y", "int32") expr = (x + y) * 2 - x * 2 - y * 2 + tirx.max(x, y) - tirx.min(x, y) - capped = tvm.arith.Analyzer() + capped = tvm.sym.Analyzer() capped.set_maximum_rewrite_steps(1) with pytest.raises(RuntimeError): capped.rewrite_simplify(expr) # A generous limit must not interfere with normal simplification. - relaxed = tvm.arith.Analyzer() + relaxed = tvm.sym.Analyzer() relaxed.set_maximum_rewrite_steps(1000) relaxed.rewrite_simplify(expr) def test_analyzer_object_try_compare_transitive(): - analyzer = tvm.arith.Analyzer() + analyzer = tvm.sym.Analyzer() x = tirx.Var("x", "int32") y = tirx.Var("y", "int32") z = tirx.Var("z", "int32") @@ -164,7 +164,7 @@ def test_analyzer_object_try_compare_transitive(): def test_analyzer_object_enabled_extensions_round_trip(): - analyzer = tvm.arith.Analyzer() + analyzer = tvm.sym.Analyzer() assert analyzer.enabled_extensions == Extension.NoExtensions @@ -176,7 +176,7 @@ def test_analyzer_object_enabled_extensions_round_trip(): def test_analyzer_object_rewrite_simplify_stats(): - analyzer = tvm.arith.Analyzer() + analyzer = tvm.sym.Analyzer() x = tirx.Var("x", "int32") analyzer.reset_rewrite_simplify_stats() @@ -190,14 +190,14 @@ def test_analyzer_object_rewrite_simplify_stats(): def test_analyzer_object_state_persists_across_ffi_calls(): - analyzer = tvm.arith.Analyzer() + analyzer = tvm.sym.Analyzer() tile = tirx.Var("tile", "int32") i = tirx.Var("i", "int32") analyzer.bind(tile, tvm.tirx.const(8, "int32")) # The same analyzer object is borrowed by the C++ DetectIterMap entry point; # its binding makes the otherwise-undetectable floormod recognizable. - result = tvm.arith.detect_iter_map([i % tile], {i: tvm.ir.Range(0, 32)}, analyzer=analyzer) + result = tvm.sym.detect_iter_map([i % tile], {i: tvm.ir.Range(0, 32)}, analyzer=analyzer) assert len(result.indices) == 1 # The binding still lives in the same stateful object after the FFI call. @@ -205,7 +205,7 @@ def test_analyzer_object_state_persists_across_ffi_calls(): def test_analyzer_object_clone_is_independent(): - analyzer = tvm.arith.Analyzer() + analyzer = tvm.sym.Analyzer() x = tirx.Var("x", "int64") y = tirx.Var("y", "int64") z = tirx.Var("z", "int64") @@ -229,15 +229,15 @@ def test_analyzer_object_clone_is_independent(): def test_analyzer_object_clone_copies_every_sub_analyzer(): - analyzer = tvm.arith.Analyzer() + analyzer = tvm.sym.Analyzer() x = tirx.Var("x", "int64") w = tirx.Var("w", "int64") v = tirx.Var("v", "int64") analyzer.bind(x, tvm.ir.Range(0, 8)) - analyzer.update(x, tvm.arith.ModularSet(4, 0)) + analyzer.update(x, tvm.sym.ModularSet(4, 0)) analyzer.bind(w, tirx.const(4, "int64")) - analyzer.update(v, tvm.arith.IntervalSet(2, 9)) + analyzer.update(v, tvm.sym.IntervalSet(2, 9)) analyzer.enabled_extensions = Extension.ComparisonOfProductAndSum clone = analyzer.clone() @@ -250,8 +250,8 @@ def test_analyzer_object_clone_copies_every_sub_analyzer(): assert clone.try_compare(x, tirx.const(0, "int64")) == CompareResult.GE t = tirx.Var("t", "int64") - clone.update(x, tvm.arith.ModularSet(8, 0), override=True) - clone.update(v, tvm.arith.IntervalSet(0, 3), override=True) + clone.update(x, tvm.sym.ModularSet(8, 0), override=True) + clone.update(v, tvm.sym.IntervalSet(0, 3), override=True) clone.bind(w, tirx.const(8, "int64"), allow_override=True) clone.bind(t, tvm.ir.Range(0, 4)) clone.enabled_extensions = Extension.NoExtensions @@ -269,7 +269,7 @@ def test_analyzer_object_clone_copies_every_sub_analyzer(): def test_analyzer_object_clone_resets_rewrite_stats(): - analyzer = tvm.arith.Analyzer() + analyzer = tvm.sym.Analyzer() x = tirx.Var("x", "int64") y = tirx.Var("y", "int64") analyzer.bind(x, tvm.ir.Range(0, 8)) diff --git a/tests/python/arith/test_arith_canonical_simplify.py b/tests/python/sym/test_sym_canonical_simplify.py similarity index 93% rename from tests/python/arith/test_arith_canonical_simplify.py rename to tests/python/sym/test_sym_canonical_simplify.py index 0c8b2184c8fa..131b678e1322 100644 --- a/tests/python/arith/test_arith_canonical_simplify.py +++ b/tests/python/sym/test_sym_canonical_simplify.py @@ -25,7 +25,7 @@ class CanonicalChecker: def __init__(self): - self.analyzer = tvm.arith.Analyzer() + self.analyzer = tvm.sym.Analyzer() def _convert(self, expr): # TODO(Lunderberg): Make utility functions `tirx.convert` and @@ -88,14 +88,14 @@ def test_split_index_simplify(): ck.verify(tmod(x * 8, 2), 0) # simplify then fold - ck.analyzer.update(x, tvm.arith.ConstIntBound(0, 1000)) - ck.analyzer.update(y, tvm.arith.ConstIntBound(0, 1000)) + ck.analyzer.update(x, tvm.sym.ConstIntBound(0, 1000)) + ck.analyzer.update(y, tvm.sym.ConstIntBound(0, 1000)) ck.verify(tdiv(x * 4 + y, 2) * 2 + tmod(x * 4 + y, 2), x * 4 + y) # complex fold ck.verify(tdiv(z * 9 + y, 2) * 2 + tmod(z * 9 + y, 2), z * 9 + y) - ck.analyzer.update(x, tvm.arith.ConstIntBound(-100, 1000), True) - ck.analyzer.update(y, tvm.arith.ConstIntBound(-100, 1000), True) + ck.analyzer.update(x, tvm.sym.ConstIntBound(-100, 1000), True) + ck.analyzer.update(y, tvm.sym.ConstIntBound(-100, 1000), True) ck.verify(tdiv(x * 4 + y, 2) * 2 + tmod(x * 4 + y, 2), x * 4 + y) # floordiv @@ -114,9 +114,9 @@ def test_split_index_simplify(): d_tile = te.var("d_tile") i = te.var("i") v = te.var("v") - ck.analyzer.update(d_tile, tvm.arith.ConstIntBound(0, 7), True) - ck.analyzer.update(i, tvm.arith.ConstIntBound(0, 1), True) - ck.analyzer.update(v, tvm.arith.ConstIntBound(0, 7), True) + ck.analyzer.update(d_tile, tvm.sym.ConstIntBound(0, 7), True) + ck.analyzer.update(i, tvm.sym.ConstIntBound(0, 1), True) + ck.analyzer.update(v, tvm.sym.ConstIntBound(0, 7), True) ck.verify(fld(flm(d_tile * 16 + i * 8 + v, 64), 8), flm(d_tile * 2 + i, 8)) # cannot simplify mixed case, unless we canonicalize into one mode. @@ -130,7 +130,7 @@ def test_bigint_coefficients_and_factors(): ck = CanonicalChecker() x = tirx.Var("x", "int64") offset = 2**61 - ck.analyzer.update(x, tvm.arith.ConstIntBound(offset, offset + 1)) + ck.analyzer.update(x, tvm.sym.ConstIntBound(offset, offset + 1)) # The canonical base 16 - 48 * offset is wider than int64. ck.verify( tirx.truncdiv(16 + 48 * (x - offset), 16), @@ -159,14 +159,14 @@ def test_div_simplify(): # (17+48*x)/16 != 1+3*x ck.verify(tdiv(17 + 48 * x, 16), tdiv(x * 48 + 17, 16)) # However, when x >= 0, then 17+48*x >= 0 and (17+48*x)/16 can be simplified - ck.analyzer.update(x, tvm.arith.ConstIntBound(0, 10)) + ck.analyzer.update(x, tvm.sym.ConstIntBound(0, 10)) ck.verify(tdiv(17 + 48 * x, 16), x * 3 + 1) # Trying expressions that are not simplifiable for any values of the variables ck.verify(tdiv(17 + 47 * x, 16), tdiv(x * 47 + 17, 16)) # floordiv fld = tvm.tirx.floordiv - ck.analyzer.update(x, tvm.arith.ConstIntBound(-1000, 10000), True) + ck.analyzer.update(x, tvm.sym.ConstIntBound(-1000, 10000), True) ck.verify(fld(16 + 48 * x, 16), x * 3 + 1) ck.verify(fld(17 + 48 * x, 16), x * 3 + 1) ck.verify(fld(17 + 47 * x, 16), fld(x * 47 + 17, 16)) @@ -277,11 +277,11 @@ def test_complex_cases(): - tmod((x * 128) + y, 1296) + 1 ) - ck.analyzer.update(x, tvm.arith.ConstIntBound(0, 5)) - ck.analyzer.update(y, tvm.arith.ConstIntBound(0, 127)) + ck.analyzer.update(x, tvm.sym.ConstIntBound(0, 5)) + ck.analyzer.update(y, tvm.sym.ConstIntBound(0, 127)) ck.verify(res2, 1) - ck.analyzer.update(y, tvm.arith.ConstIntBound(0, 1024), True) + ck.analyzer.update(y, tvm.sym.ConstIntBound(0, 1024), True) res3 = ( tdiv(x * 1024 + y, 65536) + tdiv(tmod(x * 1024 + y, 65536), 256) @@ -308,21 +308,21 @@ def test_simplify_cast(): # cast(i32, i + j + 1) - cast(i32, i) i = tvm.tirx.Var("i", "int64") j = tvm.tirx.Var("j", "int64") - ck.analyzer.update(i, tvm.arith.ConstIntBound(0, 10)) - ck.analyzer.update(j, tvm.arith.ConstIntBound(0, 10)) + ck.analyzer.update(i, tvm.sym.ConstIntBound(0, 10)) + ck.analyzer.update(j, tvm.sym.ConstIntBound(0, 10)) res = tcast("int32", i + j + 1) - tcast("int32", i) ck.verify(res, tcast("int32", j) + 1) # cast(i32, i + j - 100) i = tvm.tirx.Var("i", "int64") j = tvm.tirx.Var("j", "int64") - ck.analyzer.update(i, tvm.arith.ConstIntBound(0, 2**31 - 1)) - ck.analyzer.update(j, tvm.arith.ConstIntBound(0, 10)) + ck.analyzer.update(i, tvm.sym.ConstIntBound(0, 2**31 - 1)) + ck.analyzer.update(j, tvm.sym.ConstIntBound(0, 10)) res = tcast("int32", i + j - 100) ck.verify(res, res) # cast(i32, flm(axis, 7i64) * 2i64 + 1i64) + 1i32 # - cast(i32, flm(axis, 7i64) * 2i64) axis = tvm.tirx.Var("axis", "int64") - ck.analyzer.update(axis, tvm.arith.ConstIntBound(0, 42)) + ck.analyzer.update(axis, tvm.sym.ConstIntBound(0, 42)) res = ( tcast( "int32", diff --git a/tests/python/arith/test_arith_const_int_bound.py b/tests/python/sym/test_sym_const_int_bound.py similarity index 98% rename from tests/python/arith/test_arith_const_int_bound.py rename to tests/python/sym/test_sym_const_int_bound.py index 25c5679e1f4c..fb0a47cd84e4 100644 --- a/tests/python/arith/test_arith_const_int_bound.py +++ b/tests/python/sym/test_sym_const_int_bound.py @@ -22,7 +22,7 @@ import tvm import tvm.testing -from tvm.arith import ConstIntBound +from tvm.sym import ConstIntBound NEG_INF = ConstIntBound.NEG_INF POS_INF = ConstIntBound.POS_INF @@ -46,7 +46,7 @@ def __name__(self): class BaseCompare: def test_const_bounds(self, test_case): - analyzer = tvm.arith.Analyzer() + analyzer = tvm.sym.Analyzer() for var, bounds in test_case.known_bounds.items(): analyzer.update(var, ConstIntBound(*bounds)) @@ -77,7 +77,7 @@ class TestDataType(BaseCompare): def test_plain_var_non_negative_bound_requires_context(): var = tvm.tirx.Var("x", "int64") - analyzer = tvm.arith.Analyzer() + analyzer = tvm.sym.Analyzer() assert analyzer.const_int_bound(var).min_value == NEG_INF with analyzer.constraint_scope(var >= 0): @@ -355,7 +355,7 @@ class TestRampBound(BaseCompare): class TestModularSetBound(BaseCompare): - analyzer = tvm.arith.Analyzer() + analyzer = tvm.sym.Analyzer() tx = tvm.tirx.Var("tx", "int32") bx = tvm.tirx.Var("bx", "int32") diff --git a/tests/python/arith/test_arith_deduce_bound.py b/tests/python/sym/test_sym_deduce_bound.py similarity index 76% rename from tests/python/arith/test_arith_deduce_bound.py rename to tests/python/sym/test_sym_deduce_bound.py index 8f465b2b219d..e76c8551209f 100644 --- a/tests/python/arith/test_arith_deduce_bound.py +++ b/tests/python/sym/test_sym_deduce_bound.py @@ -28,94 +28,94 @@ def test_deduce(): c = tvm.tirx.Var("c", "int32") d = tvm.tirx.Var("d", "int32") - b_s = tvm.arith.IntervalSet(2, 3) - c_s = tvm.arith.IntervalSet(10, 15) - d_s = tvm.arith.IntervalSet(-3, -1) + b_s = tvm.sym.IntervalSet(2, 3) + c_s = tvm.sym.IntervalSet(10, 15) + d_s = tvm.sym.IntervalSet(-3, -1) zero = tvm.tirx.const(0, "int32") fdiv = tvm.tirx.floordiv e0 = (-b) * a + c - d - res0 = tvm.arith.deduce_bound(a, e0 >= 0, {b: b_s, c: c_s, d: d_s}, {}) + res0 = tvm.sym.deduce_bound(a, e0 >= 0, {b: b_s, c: c_s, d: d_s}, {}) ans0 = fdiv(d - c, b * -1) tvm.testing.assert_prim_expr_equal(res0.max_value, ans0) # expression containing variable a is on rhs - res0 = tvm.arith.deduce_bound(a, zero <= e0, {b: b_s, c: c_s, d: d_s}, {}) + res0 = tvm.sym.deduce_bound(a, zero <= e0, {b: b_s, c: c_s, d: d_s}, {}) tvm.testing.assert_prim_expr_equal(res0.max_value, ans0) e0 = d * a + c - d - res0 = tvm.arith.deduce_bound(a, e0 >= 0, {b: b_s, c: c_s, d: d_s}, {}) + res0 = tvm.sym.deduce_bound(a, e0 >= 0, {b: b_s, c: c_s, d: d_s}, {}) ans0 = fdiv(d - c, d) tvm.testing.assert_prim_expr_equal(res0.max_value, ans0) # expression containing variable a is on rhs - res0 = tvm.arith.deduce_bound(a, zero <= e0, {b: b_s, c: c_s, d: d_s}, {}) + res0 = tvm.sym.deduce_bound(a, zero <= e0, {b: b_s, c: c_s, d: d_s}, {}) tvm.testing.assert_prim_expr_equal(res0.max_value, ans0) e1 = a * 4 + b < c - res1 = tvm.arith.deduce_bound(a, e1, {b: b_s, c: c_s, d: d_s}, {}) + res1 = tvm.sym.deduce_bound(a, e1, {b: b_s, c: c_s, d: d_s}, {}) ans1 = fdiv(c - 1 - b, 4) tvm.testing.assert_prim_expr_equal(res1.max_value, ans1) # expression containing variable a is on rhs e1 = c > a * 4 + b - res1 = tvm.arith.deduce_bound(a, e1, {b: b_s, c: c_s, d: d_s}, {}) + res1 = tvm.sym.deduce_bound(a, e1, {b: b_s, c: c_s, d: d_s}, {}) tvm.testing.assert_prim_expr_equal(res1.max_value, ans1) e2 = tvm.tirx.max(5, a * 4) < 0 - res2 = tvm.arith.deduce_bound(a, e2, {b: b_s, c: c_s, d: d_s}, {}) + res2 = tvm.sym.deduce_bound(a, e2, {b: b_s, c: c_s, d: d_s}, {}) assert res2.max_value.name == "neg_inf" assert res2.min_value.name == "pos_inf" # expression containing variable a is on rhs e2 = zero < tvm.tirx.max(5, a * 4) - res2 = tvm.arith.deduce_bound(a, e2, {b: b_s, c: c_s, d: d_s}, {}) + res2 = tvm.sym.deduce_bound(a, e2, {b: b_s, c: c_s, d: d_s}, {}) assert res2.max_value.name == "neg_inf" assert res2.min_value.name == "pos_inf" e3 = (-b) + a * c - d - res3 = tvm.arith.deduce_bound(a, e3 >= 0, {b: b_s, c: c_s, d: d_s}, {b: b_s, d: d_s}) + res3 = tvm.sym.deduce_bound(a, e3 >= 0, {b: b_s, c: c_s, d: d_s}, {b: b_s, d: d_s}) ans3 = fdiv(2, c) + 1 tvm.testing.assert_prim_expr_equal(res3.min_value, ans3) - res3 = tvm.arith.deduce_bound(a, zero <= e3, {b: b_s, c: c_s, d: d_s}, {b: b_s, d: d_s}) + res3 = tvm.sym.deduce_bound(a, zero <= e3, {b: b_s, c: c_s, d: d_s}, {b: b_s, d: d_s}) tvm.testing.assert_prim_expr_equal(res3.min_value, ans3) # tests for `EQ` op - res4 = tvm.arith.deduce_bound(a, a == b, {}, {}) + res4 = tvm.sym.deduce_bound(a, a == b, {}, {}) tvm.testing.assert_prim_expr_equal(res4.max_value, b) tvm.testing.assert_prim_expr_equal(res4.min_value, b) # Unsatisfiable `EQ`, variable as one of the Operand - res5 = tvm.arith.deduce_bound(a, (a == b), {b: b_s}, {b: b_s}) + res5 = tvm.sym.deduce_bound(a, (a == b), {b: b_s}, {b: b_s}) assert res5.max_value.name == "neg_inf" assert res5.min_value.name == "pos_inf" # variable `a` on the RHS side - res6 = tvm.arith.deduce_bound(a, 10 == a, {}, {}) + res6 = tvm.sym.deduce_bound(a, 10 == a, {}, {}) tvm.testing.assert_prim_expr_equal(res6.max_value, 10) tvm.testing.assert_prim_expr_equal(res6.min_value, 10) # Add, Sub in `EQ` e4 = (a - c) == (b + d) ans4 = b + d + c - res7 = tvm.arith.deduce_bound(a, e4, {b: b_s, c: c_s, d: d_s}, {}) + res7 = tvm.sym.deduce_bound(a, e4, {b: b_s, c: c_s, d: d_s}, {}) tvm.testing.assert_prim_expr_equal(res7.max_value, ans4) tvm.testing.assert_prim_expr_equal(res7.min_value, ans4) # Satisfiable Mul in `EQ` with negative sign - res8 = tvm.arith.deduce_bound(a, (5 * a == -10), {}, {}) + res8 = tvm.sym.deduce_bound(a, (5 * a == -10), {}, {}) tvm.testing.assert_prim_expr_equal(res8.max_value, -2) tvm.testing.assert_prim_expr_equal(res8.min_value, -2) # Unsatisfiable Mul in `EQ` e5 = 4 * a == b - res9 = tvm.arith.deduce_bound(a, e5, {b: b_s}, {}) + res9 = tvm.sym.deduce_bound(a, e5, {b: b_s}, {}) assert res9.max_value.name == "neg_inf" assert res9.min_value.name == "pos_inf" - res10 = tvm.arith.deduce_bound(a, (b * a == b), {b: b_s}, {}) + res10 = tvm.sym.deduce_bound(a, (b * a == b), {b: b_s}, {}) # simplifier is now able to prove symbolic relation (b * a % b == 0) tvm.testing.assert_prim_expr_equal(res10.max_value, 1) tvm.testing.assert_prim_expr_equal(res10.min_value, 1) @@ -127,20 +127,20 @@ def test_check(): c = tvm.tirx.Var("c", "int32") d = tvm.tirx.Var("d", "int32") - b_s = tvm.arith.IntervalSet(2, 3) - c_s = tvm.arith.IntervalSet(5, 7) - d_s = tvm.arith.IntervalSet(-3, -1) + b_s = tvm.sym.IntervalSet(2, 3) + c_s = tvm.sym.IntervalSet(5, 7) + d_s = tvm.sym.IntervalSet(-3, -1) # no compare operator - res1 = tvm.arith.deduce_bound(a, a + b, {b: b_s}, {}) + res1 = tvm.sym.deduce_bound(a, a + b, {b: b_s}, {}) assert res1.is_nothing() # multiple compare operators - res2 = tvm.arith.deduce_bound(a, (a + b > 3).astype(c.ty) > c, {b: b_s, c: c_s}, {}) + res2 = tvm.sym.deduce_bound(a, (a + b > 3).astype(c.ty) > c, {b: b_s, c: c_s}, {}) assert res2.is_nothing() # multiple target variable - res2 = tvm.arith.deduce_bound(a, a * 2 - a > b, {b: b_s}, {}) + res2 = tvm.sym.deduce_bound(a, a * 2 - a > b, {b: b_s}, {}) assert res2.is_nothing() @@ -148,25 +148,25 @@ def test_deduce_basic(): def test_basic(a1, a2, coff): a = tvm.tirx.Var("a", "int32") b = tvm.tirx.Var("b", "int32") - b_s = tvm.arith.IntervalSet(a1, a2) + b_s = tvm.sym.IntervalSet(a1, a2) e0 = b + a * coff + 3 - res1 = tvm.arith.deduce_bound(a, e0 < 17, {b: b_s}, {b: b_s}) + res1 = tvm.sym.deduce_bound(a, e0 < 17, {b: b_s}, {b: b_s}) [x, y] = [res1.max_value, b_s.max_value] if coff > 0 else [res1.min_value, b_s.min_value] tvm.testing.assert_prim_expr_equal((x * coff + 3 + y) < 17, True) # expression containing variable a is on rhs - res1 = tvm.arith.deduce_bound(a, tvm.tirx.const(17, "int32") < e0, {b: b_s}, {b: b_s}) + res1 = tvm.sym.deduce_bound(a, tvm.tirx.const(17, "int32") < e0, {b: b_s}, {b: b_s}) [x, y] = [res1.max_value, b_s.max_value] if coff < 0 else [res1.min_value, b_s.min_value] tvm.testing.assert_prim_expr_equal((x * coff + 3 + y) > 17, True) # expression containing variable a is on rhs - res1 = tvm.arith.deduce_bound(a, tvm.tirx.const(17, "int32") >= e0, {b: b_s}, {b: b_s}) + res1 = tvm.sym.deduce_bound(a, tvm.tirx.const(17, "int32") >= e0, {b: b_s}, {b: b_s}) [x, y] = [res1.max_value, b_s.max_value] if coff > 0 else [res1.min_value, b_s.min_value] tvm.testing.assert_prim_expr_equal((x * coff + 3 + y) <= 17, True) - res1 = tvm.arith.deduce_bound(a, e0 >= 17, {b: b_s}, {b: b_s}) + res1 = tvm.sym.deduce_bound(a, e0 >= 17, {b: b_s}, {b: b_s}) [x, y] = [res1.max_value, b_s.max_value] if coff < 0 else [res1.min_value, b_s.min_value] tvm.testing.assert_prim_expr_equal((x * coff + 3 + y) >= 17, True) @@ -182,24 +182,24 @@ def test_deduce_complex(): def test_complex(a1, a2, coff): a = tvm.tirx.Var("a", "int32") b = tvm.tirx.Var("b", "int32") - b_s = tvm.arith.IntervalSet(a1, a2) + b_s = tvm.sym.IntervalSet(a1, a2) e0 = (b * 3 + a * coff) * 4 - res1 = tvm.arith.deduce_bound(a, e0 < 63, {b: b_s}, {b: b_s}) + res1 = tvm.sym.deduce_bound(a, e0 < 63, {b: b_s}, {b: b_s}) [t, x] = [res1.max_value, b_s.max_value] if coff > 0 else [res1.min_value, b_s.min_value] tvm.testing.assert_prim_expr_equal(((x * 3 + t * coff) * 4) < 63, True) # expression containing variable a is on rhs - res1 = tvm.arith.deduce_bound(a, tvm.tirx.const(63, "int32") >= e0, {b: b_s}, {b: b_s}) + res1 = tvm.sym.deduce_bound(a, tvm.tirx.const(63, "int32") >= e0, {b: b_s}, {b: b_s}) [t, x] = [res1.max_value, b_s.max_value] if coff > 0 else [res1.min_value, b_s.min_value] tvm.testing.assert_prim_expr_equal(((x * 3 + t * coff) * 4) <= 63, True) - res1 = tvm.arith.deduce_bound(a, e0 > 63, {b: b_s}, {b: b_s}) + res1 = tvm.sym.deduce_bound(a, e0 > 63, {b: b_s}, {b: b_s}) [t, x] = [res1.max_value, b_s.max_value] if coff < 0 else [res1.min_value, b_s.min_value] tvm.testing.assert_prim_expr_equal(((x * 3 + t * coff) * 4) > 63, True) # expression containing variable a is on rhs - res1 = tvm.arith.deduce_bound(a, tvm.tirx.const(63, "int32") <= e0, {b: b_s}, {b: b_s}) + res1 = tvm.sym.deduce_bound(a, tvm.tirx.const(63, "int32") <= e0, {b: b_s}, {b: b_s}) [t, x] = [res1.max_value, b_s.max_value] if coff < 0 else [res1.min_value, b_s.min_value] tvm.testing.assert_prim_expr_equal(((x * 3 + t * coff) * 4) >= 63, True) @@ -215,7 +215,7 @@ def test_deduce_non_support(): a = tvm.tirx.Var("a", "int32") def test_non_support(lhs): - res = tvm.arith.deduce_bound(a, lhs < 10, {}, {}) + res = tvm.sym.deduce_bound(a, lhs < 10, {}, {}) assert res.is_nothing() test_non_support(tvm.tirx.floormod(a, 16)) @@ -235,7 +235,7 @@ def test_deduce_floordiv(): def do_test(gen_expr, dom_map, expect_min, expect_max): a = tvm.tirx.Var("a", "int32") expr = gen_expr(a) - res = tvm.arith.deduce_bound(a, expr, dom_map, dom_map) + res = tvm.sym.deduce_bound(a, expr, dom_map, dom_map) if isinstance(expect_min, str): assert res.min_value.name == expect_min else: @@ -261,7 +261,7 @@ def do_test(gen_expr, dom_map, expect_min, expect_max): # test nested cases b = tvm.tirx.Var("b", "int32") - bs = {b: tvm.arith.IntervalSet(2, 6)} + bs = {b: tvm.sym.IntervalSet(2, 6)} do_test(lambda a: b * 3 + a // 8 < 63, bs, "neg_inf", 359) do_test(lambda a: b * 3 + a // 8 <= 63, bs, "neg_inf", 367) do_test(lambda a: b * 3 + a // 8 > 63, bs, 464, "pos_inf") diff --git a/tests/python/arith/test_arith_detect_clip_bound.py b/tests/python/sym/test_sym_detect_clip_bound.py similarity index 75% rename from tests/python/arith/test_arith_detect_clip_bound.py rename to tests/python/sym/test_sym_detect_clip_bound.py index c93130491cbe..eca33c54f273 100644 --- a/tests/python/arith/test_arith_detect_clip_bound.py +++ b/tests/python/sym/test_sym_detect_clip_bound.py @@ -22,17 +22,17 @@ def test_basic(): a = tvm.tirx.Var("a", "int32") b = tvm.tirx.Var("b", "int32") c = tvm.tirx.Var("c", "int32") - m = tvm.arith.detect_clip_bound(tvm.tirx.all(a * 1 < b * 6, a - 1 > 0), [a]) + m = tvm.sym.detect_clip_bound(tvm.tirx.all(a * 1 < b * 6, a - 1 > 0), [a]) tvm.testing.assert_prim_expr_equal(m[1], b * 6 - 1) assert m[0].value == 2 - m = tvm.arith.detect_clip_bound(tvm.tirx.all(a * 1 < b * 6, a - 1 > 0), [a, b]) + m = tvm.sym.detect_clip_bound(tvm.tirx.all(a * 1 < b * 6, a - 1 > 0), [a, b]) assert len(m) == 0 - m = tvm.arith.detect_clip_bound(tvm.tirx.all(a + 10 * c <= 20, b - 1 > 0), [a, b]) + m = tvm.sym.detect_clip_bound(tvm.tirx.all(a + 10 * c <= 20, b - 1 > 0), [a, b]) tvm.testing.assert_prim_expr_equal(m[1], 20 - 10 * c) tvm.testing.assert_prim_expr_equal(m[2], 2) - m = tvm.arith.detect_clip_bound(tvm.tirx.all(tvm.tirx.Not(a * 1 > b * 6), a - 1 > 0), [a]) + m = tvm.sym.detect_clip_bound(tvm.tirx.all(tvm.tirx.Not(a * 1 > b * 6), a - 1 > 0), [a]) tvm.testing.assert_prim_expr_equal(m[1], b * 6) - m = tvm.arith.detect_clip_bound(tvm.tirx.all(tvm.tirx.Min(a, b) > 3, a - 10 < 0), [a, b]) + m = tvm.sym.detect_clip_bound(tvm.tirx.all(tvm.tirx.Min(a, b) > 3, a - 10 < 0), [a, b]) tvm.testing.assert_prim_expr_equal(m[0], 4) tvm.testing.assert_prim_expr_equal(m[1], 9) tvm.testing.assert_prim_expr_equal(m[2], 4) @@ -41,10 +41,10 @@ def test_basic(): def test_trivial_eq(): a = tvm.tirx.Var("a", "int32") b = tvm.tirx.Var("b", "int32") - m = tvm.arith.detect_clip_bound(b == 3, [a, b]) + m = tvm.sym.detect_clip_bound(b == 3, [a, b]) tvm.testing.assert_prim_expr_equal(m[2], 3) tvm.testing.assert_prim_expr_equal(m[3], 3) - m = tvm.arith.detect_clip_bound(tvm.tirx.all(a == 4, b == 3), [a, b]) + m = tvm.sym.detect_clip_bound(tvm.tirx.all(a == 4, b == 3), [a, b]) tvm.testing.assert_prim_expr_equal(m[0], 4) tvm.testing.assert_prim_expr_equal(m[1], 4) tvm.testing.assert_prim_expr_equal(m[2], 3) diff --git a/tests/python/arith/test_arith_detect_linear_equation.py b/tests/python/sym/test_sym_detect_linear_equation.py similarity index 67% rename from tests/python/arith/test_arith_detect_linear_equation.py rename to tests/python/sym/test_sym_detect_linear_equation.py index 08332cec9760..9e3154b6e9d5 100644 --- a/tests/python/arith/test_arith_detect_linear_equation.py +++ b/tests/python/sym/test_sym_detect_linear_equation.py @@ -21,58 +21,58 @@ def test_basic(): a = tvm.tirx.Var("a", "int32") b = tvm.tirx.Var("b", "int32") - m = tvm.arith.detect_linear_equation(a * 4 + b * 6 + 7, [a]) + m = tvm.sym.detect_linear_equation(a * 4 + b * 6 + 7, [a]) assert m[0].value == 4 tvm.testing.assert_prim_expr_equal(m[1], b * 6 + 7) - m = tvm.arith.detect_linear_equation(a * 4 * (a + 1) + b * 6 + 7, [a]) + m = tvm.sym.detect_linear_equation(a * 4 * (a + 1) + b * 6 + 7, [a]) assert len(m) == 0 - m = tvm.arith.detect_linear_equation(a * 4 + (a + 1) + b * 6 + 7, [a]) + m = tvm.sym.detect_linear_equation(a * 4 + (a + 1) + b * 6 + 7, [a]) assert m[0].value == 5 tvm.testing.assert_prim_expr_equal(m[1], b * 6 + 7 + 1) - m = tvm.arith.detect_linear_equation(a * b + 7, [a]) + m = tvm.sym.detect_linear_equation(a * b + 7, [a]) assert m[0] == b - m = tvm.arith.detect_linear_equation(b * 7, [a]) + m = tvm.sym.detect_linear_equation(b * 7, [a]) assert m[0].value == 0 - m = tvm.arith.detect_linear_equation(b * 7, []) + m = tvm.sym.detect_linear_equation(b * 7, []) assert len(m) == 1 tvm.testing.assert_prim_expr_equal(m[0], b * 7) c = tvm.tirx.Var("c", "uint32") - m = tvm.arith.detect_linear_equation(128 - c, [c]) + m = tvm.sym.detect_linear_equation(128 - c, [c]) assert m[0].value == -1 def test_multivariate(): v = [tvm.tirx.Var(f"v{i}", "int32") for i in range(4)] b = tvm.tirx.Var("b", "int32") - m = tvm.arith.detect_linear_equation(v[0] * (b + 4) + v[0] + v[1] * 8, v) + m = tvm.sym.detect_linear_equation(v[0] * (b + 4) + v[0] + v[1] * 8, v) tvm.testing.assert_prim_expr_equal(m[0], b + 5) assert m[1].value == 8 - m = tvm.arith.detect_linear_equation(v[0] * (b + 4) + v[0] + v[1] * 8 * v[2], v) + m = tvm.sym.detect_linear_equation(v[0] * (b + 4) + v[0] + v[1] * 8 * v[2], v) assert len(m) == 0 - m = tvm.arith.detect_linear_equation(v[0] * (b + 4) + v[0] + v[1] * 8 * v[1] + v[3], v) + m = tvm.sym.detect_linear_equation(v[0] * (b + 4) + v[0] + v[1] * 8 * v[1] + v[3], v) assert len(m) == 0 - m = tvm.arith.detect_linear_equation(((v[0] * b + v[1]) * 8 + v[2] + 1) * 2, v) + m = tvm.sym.detect_linear_equation(((v[0] * b + v[1]) * 8 + v[2] + 1) * 2, v) assert m[1].value == 16 assert m[2].value == 2 assert m[len(m) - 1].value == 2 - m = tvm.arith.detect_linear_equation((v[0] - v[1]), [v[2]]) + m = tvm.sym.detect_linear_equation((v[0] - v[1]), [v[2]]) assert m[0].value == 0 tvm.testing.assert_prim_expr_equal(m[1], v[0] - v[1]) - m = tvm.arith.detect_linear_equation((v[0] - v[1]), []) + m = tvm.sym.detect_linear_equation((v[0] - v[1]), []) assert len(m) == 1 tvm.testing.assert_prim_expr_equal(m[0], v[0] - v[1]) diff --git a/tests/python/arith/test_arith_intset.py b/tests/python/sym/test_sym_intset.py similarity index 77% rename from tests/python/arith/test_arith_intset.py rename to tests/python/sym/test_sym_intset.py index 223b73f49dba..99cc81a472ac 100644 --- a/tests/python/arith/test_arith_intset.py +++ b/tests/python/sym/test_sym_intset.py @@ -18,12 +18,12 @@ import tvm import tvm.testing from tvm import tirx -from tvm.arith.analyzer import Analyzer +from tvm.sym.analyzer import Analyzer class IntSetChecker: def __init__(self): - self.analyzer = tvm.arith.Analyzer() + self.analyzer = tvm.sym.Analyzer() def verify(self, data, dmap, expected): res = self.analyzer.int_set(data, dmap) @@ -36,11 +36,11 @@ def err_msg(): def test_basic(): - s = tvm.arith.IntervalSet(2, 3) + s = tvm.sym.IntervalSet(2, 3) assert s.min_value.value == 2 assert s.max_value.value == 3 - s = tvm.arith.IntSet.single_point(2) + s = tvm.sym.IntSet.single_point(2) assert s.min_value.value == 2 assert s.max_value.value == 2 @@ -49,25 +49,25 @@ def test_vector(): base = 10 stride = 3 lanes = 2 - s = tvm.arith.IntSet.vector(tvm.tirx.Ramp(base, stride, lanes)) + s = tvm.sym.IntSet.vector(tvm.tirx.Ramp(base, stride, lanes)) assert s.min_value.value == base assert s.max_value.value == base + stride * (lanes - 1) def test_scalable_vector(): base = 5 - s = tvm.arith.IntSet.vector(tvm.tirx.Ramp(base, 2, tvm.tirx.vscale() * 4)) + s = tvm.sym.IntSet.vector(tvm.tirx.Ramp(base, 2, tvm.tirx.vscale() * 4)) - assert s.min_value.same_as(tvm.arith.int_set.neg_inf()) - assert s.max_value.same_as(tvm.arith.int_set.pos_inf()) + assert s.min_value.same_as(tvm.sym.int_set.neg_inf()) + assert s.max_value.same_as(tvm.sym.int_set.pos_inf()) def test_add_sub(): ck = IntSetChecker() x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32") - ck.verify(x + y, {x: tvm.arith.IntervalSet(0, 10)}, (y, 10 + y)) - ck.verify(x + y, {x: tvm.arith.IntervalSet(0, 10), y: tvm.arith.IntervalSet(1, 11)}, (1, 21)) - ck.verify(x - y, {x: tvm.arith.IntervalSet(0, 10), y: tvm.arith.IntervalSet(1, 11)}, (-11, 9)) + ck.verify(x + y, {x: tvm.sym.IntervalSet(0, 10)}, (y, 10 + y)) + ck.verify(x + y, {x: tvm.sym.IntervalSet(0, 10), y: tvm.sym.IntervalSet(1, 11)}, (1, 21)) + ck.verify(x - y, {x: tvm.sym.IntervalSet(0, 10), y: tvm.sym.IntervalSet(1, 11)}, (-11, 9)) def test_mul_div(): @@ -75,41 +75,41 @@ def test_mul_div(): x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32") tdiv = tvm.tirx.truncdiv - ck.analyzer.update(y, tvm.arith.ConstIntBound(1, 100), override=True) - ck.verify(x * y, {x: tvm.arith.IntervalSet(0, 10)}, (0, 10 * y)) - ck.verify(x * 2, {x: tvm.arith.IntervalSet(1, 10)}, (2, 20)) - ck.verify(x * -2, {x: tvm.arith.IntervalSet(1, 10)}, (-20, -2)) + ck.analyzer.update(y, tvm.sym.ConstIntBound(1, 100), override=True) + ck.verify(x * y, {x: tvm.sym.IntervalSet(0, 10)}, (0, 10 * y)) + ck.verify(x * 2, {x: tvm.sym.IntervalSet(1, 10)}, (2, 20)) + ck.verify(x * -2, {x: tvm.sym.IntervalSet(1, 10)}, (-20, -2)) - ck.verify(tdiv(x, y), {x: tvm.arith.IntervalSet(0, 10)}, (0, tdiv(10, y))) - ck.verify(tdiv(x, 2), {x: tvm.arith.IntervalSet(1, 10)}, (0, 5)) + ck.verify(tdiv(x, y), {x: tvm.sym.IntervalSet(0, 10)}, (0, tdiv(10, y))) + ck.verify(tdiv(x, 2), {x: tvm.sym.IntervalSet(1, 10)}, (0, 5)) fld = tvm.tirx.floordiv - ck.verify(fld(x, y), {x: tvm.arith.IntervalSet(0, 10)}, (0, fld(10, y))) - ck.verify(fld(x, 2), {x: tvm.arith.IntervalSet(-1, 10)}, (-1, 5)) + ck.verify(fld(x, y), {x: tvm.sym.IntervalSet(0, 10)}, (0, fld(10, y))) + ck.verify(fld(x, 2), {x: tvm.sym.IntervalSet(-1, 10)}, (-1, 5)) def test_mod(): ck = IntSetChecker() x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32") tmod = tvm.tirx.truncmod - ck.analyzer.update(y, tvm.arith.ConstIntBound(1, 100), override=True) - ck.verify(tmod(x, y), {x: tvm.arith.IntervalSet(0, 10)}, (0, y - 1)) - ck.verify(tmod(x, 10), {x: tvm.arith.IntervalSet(1, 10)}, (0, 9)) + ck.analyzer.update(y, tvm.sym.ConstIntBound(1, 100), override=True) + ck.verify(tmod(x, y), {x: tvm.sym.IntervalSet(0, 10)}, (0, y - 1)) + ck.verify(tmod(x, 10), {x: tvm.sym.IntervalSet(1, 10)}, (0, 9)) flm = tvm.tirx.floormod - ck.verify(flm(x, 10), {x: tvm.arith.IntervalSet(-10, 10)}, (0, 9)) - ck.verify(flm(x, 10), {x: tvm.arith.IntervalSet(3, 5)}, (3, 5)) - ck.verify(flm(x, 10), {x: tvm.arith.IntervalSet(13, 15)}, (3, 5)) - ck.verify(flm(x, 10), {x: tvm.arith.IntervalSet(3, 15)}, (0, 9)) - ck.verify(flm(x, 10), {x: tvm.arith.IntervalSet(3, 11)}, (0, 9)) - ck.verify(flm(x, 10), {x: tvm.arith.IntervalSet(1, 21)}, (0, 9)) + ck.verify(flm(x, 10), {x: tvm.sym.IntervalSet(-10, 10)}, (0, 9)) + ck.verify(flm(x, 10), {x: tvm.sym.IntervalSet(3, 5)}, (3, 5)) + ck.verify(flm(x, 10), {x: tvm.sym.IntervalSet(13, 15)}, (3, 5)) + ck.verify(flm(x, 10), {x: tvm.sym.IntervalSet(3, 15)}, (0, 9)) + ck.verify(flm(x, 10), {x: tvm.sym.IntervalSet(3, 11)}, (0, 9)) + ck.verify(flm(x, 10), {x: tvm.sym.IntervalSet(1, 21)}, (0, 9)) fld = tvm.tirx.floordiv z = tvm.tirx.Var("z", "int32") ck.analyzer.bind(x, tvm.ir.Range.from_min_extent(0, 3)) ck.verify( flm(y, 8), - {y: tvm.arith.IntervalSet(z * 8 + x * 4, z * 8 + x * 4 + 3)}, + {y: tvm.sym.IntervalSet(z * 8 + x * 4, z * 8 + x * 4 + 3)}, ( z * 8 + x * 4 - 8 * fld(z * 8 + x * 4, 8), z * 8 + x * 4 + 3 - 8 * fld(z * 8 + x * 4, 8), @@ -118,15 +118,15 @@ def test_mod(): ck1 = IntSetChecker() ck1.analyzer.bind(x, tvm.ir.Range.from_min_extent(0, 2)) ck1.verify( - flm(y, 8), {y: tvm.arith.IntervalSet(z * 8 + x * 4, z * 8 + x * 4 + 3)}, (x * 4, x * 4 + 3) + flm(y, 8), {y: tvm.sym.IntervalSet(z * 8 + x * 4, z * 8 + x * 4 + 3)}, (x * 4, x * 4 + 3) ) def test_max_min(): ck = IntSetChecker() x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32") - ck.verify(tvm.tirx.max(x, x + 1), {x: tvm.arith.IntervalSet(0, 10)}, (1, 11)) - ck.verify(tvm.tirx.min(x - 1, x + 1), {x: tvm.arith.IntervalSet(0, 10)}, (-1, 9)) + ck.verify(tvm.tirx.max(x, x + 1), {x: tvm.sym.IntervalSet(0, 10)}, (1, 11)) + ck.verify(tvm.tirx.min(x - 1, x + 1), {x: tvm.sym.IntervalSet(0, 10)}, (-1, 9)) ck.verify(tvm.tirx.min(x, y), {}, (tvm.tirx.min(x, y), tvm.tirx.min(x, y))) ck.verify(tvm.tirx.max(x, y), {}, (tvm.tirx.max(x, y), tvm.tirx.max(x, y))) @@ -134,7 +134,7 @@ def test_max_min(): def test_select(): ck = IntSetChecker() x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32") - ck.verify(tvm.tirx.Select(x > 0, x - 1, x + 1), {x: tvm.arith.IntervalSet(0, 10)}, (-1, 11)) + ck.verify(tvm.tirx.Select(x > 0, x - 1, x + 1), {x: tvm.sym.IntervalSet(0, 10)}, (-1, 11)) def check_region_bound(expect_region, var_dom, mode, predicate=None): @@ -165,15 +165,15 @@ def check_region_bound(expect_region, var_dom, mode, predicate=None): region.append(tvm.ir.Range.from_min_extent(k[0], Analyzer().simplify(k[1] - k[0]))) expect.append(v) if mode == "lowerbound": - result = tvm.arith.estimate_region_lower_bound( + result = tvm.sym.estimate_region_lower_bound( region=region, var_dom=var_dom, predicate=predicate ) elif mode == "upperbound": - result = tvm.arith.estimate_region_upper_bound( + result = tvm.sym.estimate_region_upper_bound( region=region, var_dom=var_dom, predicate=predicate ) else: - result = tvm.arith.estimate_region_strict_bound( + result = tvm.sym.estimate_region_strict_bound( region=region, var_dom=var_dom, predicate=predicate ) if result is None: @@ -371,15 +371,15 @@ def test_region_lower_bound_unfusable(): def test_union_lower_bound(): - neg_inf = tvm.arith.int_set.neg_inf() - pos_inf = tvm.arith.int_set.pos_inf() - set_0 = tvm.arith.IntervalSet(min_value=neg_inf, max_value=0) - set_1 = tvm.arith.IntervalSet(min_value=1, max_value=pos_inf) - result = tvm.arith.int_set.union_lower_bound([set_0, set_1]) + neg_inf = tvm.sym.int_set.neg_inf() + pos_inf = tvm.sym.int_set.pos_inf() + set_0 = tvm.sym.IntervalSet(min_value=neg_inf, max_value=0) + set_1 = tvm.sym.IntervalSet(min_value=1, max_value=pos_inf) + result = tvm.sym.int_set.union_lower_bound([set_0, set_1]) assert result.min_value.same_as(neg_inf) assert result.max_value.same_as(pos_inf) - set_2 = tvm.arith.IntervalSet(min_value=pos_inf, max_value=neg_inf) - result = tvm.arith.int_set.union_lower_bound([set_0, set_1, set_2]) + set_2 = tvm.sym.IntervalSet(min_value=pos_inf, max_value=neg_inf) + result = tvm.sym.int_set.union_lower_bound([set_0, set_1, set_2]) assert result.min_value.same_as(neg_inf) assert result.max_value.same_as(pos_inf) @@ -389,9 +389,7 @@ def test_modular_set(): x = tvm.tirx.Var("x", "int32") y = tvm.tirx.Var("y", "int32") expr = (x * 2048 + y * 16) % 7168 - ck.verify( - expr, {x: tvm.arith.IntervalSet(0, 128), y: tvm.arith.IntervalSet(0, 3584)}, (0, 7152) - ) + ck.verify(expr, {x: tvm.sym.IntervalSet(0, 128), y: tvm.sym.IntervalSet(0, 3584)}, (0, 7152)) def test_relax_deep_variable_dependency_chain(): @@ -407,19 +405,19 @@ def test_relax_deep_variable_dependency_chain(): ck = IntSetChecker() n = 64 # 2^64 expansions without memoization; trivially fast with it. xs = [tvm.tirx.Var(f"x{i}", "int32") for i in range(n + 1)] - dmap = {xs[i]: tvm.arith.IntervalSet(xs[i + 1] - 1, xs[i + 1] + 1) for i in range(n)} - dmap[xs[n]] = tvm.arith.IntervalSet(0, 100) + dmap = {xs[i]: tvm.sym.IntervalSet(xs[i + 1] - 1, xs[i + 1] + 1) for i in range(n)} + dmap[xs[n]] = tvm.sym.IntervalSet(0, 100) # x0 relaxes through the whole chain: [0 - n, 100 + n]. ck.verify(xs[0], dmap, (-n, 100 + n)) def test_relax_cyclic_variable_dependency(): """A cyclic variable dependency must terminate (and stay symbolic).""" - ana = tvm.arith.Analyzer() + ana = tvm.sym.Analyzer() x = tvm.tirx.Var("x", "int32") y = tvm.tirx.Var("y", "int32") # x depends on y and y depends on x: relaxation must not loop forever. - dmap = {x: tvm.arith.IntervalSet(y, y), y: tvm.arith.IntervalSet(x, x)} + dmap = {x: tvm.sym.IntervalSet(y, y), y: tvm.sym.IntervalSet(x, x)} res = ana.int_set(x, dmap) assert res is not None @@ -431,17 +429,17 @@ def test_estimate_region_accepts_external_analyzer(): dom = {i: tvm.ir.Range(0, 16)} # Without knowing `tile`, the affine detection fails for exact bounds. - assert tvm.arith.estimate_region_lower_bound(region, dom, True) is None - assert tvm.arith.estimate_region_strict_bound(region, dom, True) is None - upper_without_analyzer = tvm.arith.estimate_region_upper_bound(region, dom, True) + assert tvm.sym.estimate_region_lower_bound(region, dom, True) is None + assert tvm.sym.estimate_region_strict_bound(region, dom, True) is None + upper_without_analyzer = tvm.sym.estimate_region_upper_bound(region, dom, True) - analyzer = tvm.arith.Analyzer() + analyzer = tvm.sym.Analyzer() analyzer.bind(tile, tvm.tirx.const(4, "int32")) # The external binding lets the affine detection succeed. for estimate_region in [ - tvm.arith.estimate_region_lower_bound, - tvm.arith.estimate_region_strict_bound, - tvm.arith.estimate_region_upper_bound, + tvm.sym.estimate_region_lower_bound, + tvm.sym.estimate_region_strict_bound, + tvm.sym.estimate_region_upper_bound, ]: result = estimate_region(region, dom, True, analyzer=analyzer) assert result is not None diff --git a/tests/python/arith/test_arith_iter_affine_map.py b/tests/python/sym/test_sym_iter_affine_map.py similarity index 85% rename from tests/python/arith/test_arith_iter_affine_map.py rename to tests/python/sym/test_sym_iter_affine_map.py index 6a5b809ac430..61a22fde093e 100644 --- a/tests/python/arith/test_arith_iter_affine_map.py +++ b/tests/python/sym/test_sym_iter_affine_map.py @@ -46,14 +46,14 @@ def var_dom(iters): def convert_iter_expr(expr): - return tvm.arith.normalize_iter_map_to_expr(expr) + return tvm.sym.normalize_iter_map_to_expr(expr) def assert_iter_sum_pattern( expect_dict, dom_map, predicate=True, check_level="surjective", simplify_trivial_iterators=True ): keys = list(expect_dict.keys()) - res = tvm.arith.detect_iter_map( + res = tvm.sym.detect_iter_map( keys, dom_map, predicate=predicate, @@ -71,7 +71,7 @@ def assert_iter_sum_pattern( scale = spec[2] if len(spec) > 2 else 1 expect_iter = spec[3] if len(spec) > 3 else None sum_expr = indices[i] - assert isinstance(sum_expr, tvm.arith.IterSumExpr) + assert isinstance(sum_expr, tvm.sym.IterSumExpr) if extent == 1: assert len(sum_expr.args) == 0 else: @@ -80,7 +80,7 @@ def assert_iter_sum_pattern( tvm.testing.assert_prim_expr_equal(sum_expr.args[0].scale, scale) tvm.testing.assert_prim_expr_equal(sum_expr.base, base) if expect_iter is not None: - if not isinstance(expect_iter, tvm.arith.IterMapExpr): + if not isinstance(expect_iter, tvm.sym.IterMapExpr): sum_expr = convert_iter_expr(sum_expr) tvm.ir.assert_structural_equal(sum_expr, expect_iter) @@ -89,14 +89,14 @@ def assert_iter_map_simplify( expect_dict, dom_map, predicate=True, check_level="surjective", simplify_trivial_iterators=True ): keys = list(expect_dict.keys()) - imap = tvm.arith.detect_iter_map( + imap = tvm.sym.detect_iter_map( keys, dom_map, predicate=predicate, check_level=check_level, simplify_trivial_iterators=simplify_trivial_iterators, ) - res = tvm.arith.iter_map_simplify( + res = tvm.sym.iter_map_simplify( keys, dom_map, predicate=predicate, @@ -109,7 +109,7 @@ def assert_iter_map_simplify( def assert_iter_sum_failure(iters, dom_map, predicate=True, check_level="surjective"): - res = tvm.arith.detect_iter_map( + res = tvm.sym.detect_iter_map( list(iters), dom_map, predicate=predicate, check_level=check_level ).indices assert len(res) == 0 @@ -206,12 +206,12 @@ def test_split_simplified_modulo(): i = tvm.tirx.Var("i", "int32") j = tvm.tirx.Var("j", "int32") dom = var_dom([(i, 64), (j, 192)]) - analyzer = tvm.arith.Analyzer() + analyzer = tvm.sym.Analyzer() for flat in [i * 192 + j, j + i * 192]: lane = analyzer.simplify(floormod(flat, 128)) quotient = floordiv(flat, 128) - res = tvm.arith.detect_iter_map([lane, quotient], dom, check_level="bijective") + res = tvm.sym.detect_iter_map([lane, quotient], dom, check_level="bijective") assert len(res.errors) == 0, res.errors assert len(res.indices) == 2 @@ -225,16 +225,16 @@ def test_compound(): z = ifuse([yo, xo, yi]) # reconstruct the pattern manually - mx = tvm.arith.IterMark(x, 10) - my = tvm.arith.IterMark(y, 9) + mx = tvm.sym.IterMark(x, 10) + my = tvm.sym.IterMark(y, 9) xoscale = 3 yoscale = 6 yiscale = 1 - mxo = tvm.arith.IterSplitExpr(mx, 5, 2, xoscale) - myo = tvm.arith.IterSplitExpr(my, 3, 3, yoscale) - myi = tvm.arith.IterSplitExpr(my, 1, 3, yiscale) - mz = tvm.arith.IterMark(tvm.arith.IterSumExpr([myo, mxo, myi], 0), 18) - sz = tvm.arith.IterSumExpr([tvm.arith.IterSplitExpr(mz, 1, 18, 1)], 0) + mxo = tvm.sym.IterSplitExpr(mx, 5, 2, xoscale) + myo = tvm.sym.IterSplitExpr(my, 3, 3, yoscale) + myi = tvm.sym.IterSplitExpr(my, 1, 3, yiscale) + mz = tvm.sym.IterMark(tvm.sym.IterSumExpr([myo, mxo, myi], 0), 18) + sz = tvm.sym.IterSumExpr([tvm.sym.IterSplitExpr(mz, 1, 18, 1)], 0) assert_iter_sum_pattern({z[0]: (18, 0, 1, sz), xi[0]: (5, 0)}, var_dom([(x, 10), (y, 9)])) @@ -563,8 +563,8 @@ def convert_division(divisions): for division in divisions[:-1]: res.append( [ - tvm.arith.normalize_iter_map_to_expr(division[0].source), - tvm.arith.normalize_iter_map_to_expr(division[1].source), + tvm.sym.normalize_iter_map_to_expr(division[0].source), + tvm.sym.normalize_iter_map_to_expr(division[1].source), ] ) res.append([divisions[-1][0].extent, divisions[-1][1].extent]) @@ -582,16 +582,14 @@ def test_subspace_division(): c = tvm.tirx.Var("c", "int32") # simple 1.1 - res = tvm.arith.subspace_divide( - [z * 12 + y * 3 + x + c], var_dom([(x, 3), (y, 4), (z, 5)]), [x] - ) + res = tvm.sym.subspace_divide([z * 12 + y * 3 + x + c], var_dom([(x, 3), (y, 4), (z, 5)]), [x]) res = convert_division(res) assert len(res) == 2 tvm.ir.assert_structural_equal(res[0][0], z * 4 + y) tvm.ir.assert_structural_equal(res[0][1], x + c) # simple 1.2 - res = tvm.arith.subspace_divide( + res = tvm.sym.subspace_divide( [z * 12 + y * 3 + x + c], var_dom([(x, 3), (y, 4), (z, 5)]), [x], z * 4 + y < 18 ) res = convert_division(res) @@ -611,7 +609,7 @@ def test_subspace_division(): k1 = ifuse([i2, i3]) # compound 1.1 - res = tvm.arith.subspace_divide([k0[0], k1[0]], var_dom([i0, j0, i3]), [i3[0]]) + res = tvm.sym.subspace_divide([k0[0], k1[0]], var_dom([i0, j0, i3]), [i3[0]]) res = convert_division(res) assert len(res) == 3 tvm.ir.assert_structural_equal(res[0][0], (i0[0] * 2) + floordiv(j0[0], 4)) @@ -620,13 +618,13 @@ def test_subspace_division(): tvm.ir.assert_structural_equal(res[1][1], i3[0]) assert_iter_sum_pattern - res1 = tvm.arith.detect_iter_map([res[0][1], res[1][1]], var_dom([i3])).indices + res1 = tvm.sym.detect_iter_map([res[0][1], res[1][1]], var_dom([i3])).indices assert len(res1) == 2 - res2 = tvm.arith.detect_iter_map([res[0][0], res[1][0]], var_dom([i0, j0])).indices + res2 = tvm.sym.detect_iter_map([res[0][0], res[1][0]], var_dom([i0, j0])).indices assert len(res2) == 2 # compound 1.2 - res = tvm.arith.subspace_divide([k0[0], k1[0]], var_dom([i0, j0, i3]), [j0[0], i3[0]]) + res = tvm.sym.subspace_divide([k0[0], k1[0]], var_dom([i0, j0, i3]), [j0[0], i3[0]]) res = convert_division(res) assert len(res) == 3 tvm.ir.assert_structural_equal(res[0][0], i0[0]) @@ -634,18 +632,18 @@ def test_subspace_division(): tvm.ir.assert_structural_equal(res[1][0], T.int32(0)) tvm.ir.assert_structural_equal(res[1][1], (floormod(j0[0], 4) * 2) + i3[0]) - res1 = tvm.arith.detect_iter_map([res[0][1], res[1][1]], var_dom([j0, i3])).indices + res1 = tvm.sym.detect_iter_map([res[0][1], res[1][1]], var_dom([j0, i3])).indices assert len(res1) == 2 - res2 = tvm.arith.detect_iter_map([res[0][0], res[1][0]], var_dom([i0])).indices + res2 = tvm.sym.detect_iter_map([res[0][0], res[1][0]], var_dom([i0])).indices assert len(res2) == 2 # compound 1.3 - res = tvm.arith.subspace_divide([k0[0], k1[0]], var_dom([i0, j0, i3]), [i0[0], i3[0]]) + res = tvm.sym.subspace_divide([k0[0], k1[0]], var_dom([i0, j0, i3]), [i0[0], i3[0]]) res = convert_division(res) assert len(res) == 0 # compound 1.4 - res = tvm.arith.subspace_divide([k0[0], k1[0]], var_dom([i0, j0, i3]), [i3[0]], k0[0] < 7) + res = tvm.sym.subspace_divide([k0[0], k1[0]], var_dom([i0, j0, i3]), [i3[0]], k0[0] < 7) res = convert_division(res) assert len(res) == 3 tvm.ir.assert_structural_equal(res[0][0], (i0[0] * 2) + floordiv(j0[0], 4)) @@ -655,15 +653,13 @@ def test_subspace_division(): tvm.ir.assert_structural_equal(res[2][0], (i0[0] * 2) + floordiv(j0[0], 4) < 7) tvm.ir.assert_structural_equal(res[2][1], T.bool(True)) - res1 = tvm.arith.detect_iter_map([res[0][1], res[1][1]], var_dom([i3])).indices + res1 = tvm.sym.detect_iter_map([res[0][1], res[1][1]], var_dom([i3])).indices assert len(res1) == 2 - res2 = tvm.arith.detect_iter_map([res[0][0], res[1][0]], var_dom([i0, j0])).indices + res2 = tvm.sym.detect_iter_map([res[0][0], res[1][0]], var_dom([i0, j0])).indices assert len(res2) == 2 # compound 1.5 - res = tvm.arith.subspace_divide( - [k0[0], k1[0]], var_dom([i0, j0, i3]), [j0[0], i3[0]], k1[0] < 7 - ) + res = tvm.sym.subspace_divide([k0[0], k1[0]], var_dom([i0, j0, i3]), [j0[0], i3[0]], k1[0] < 7) res = convert_division(res) assert len(res) == 3 tvm.ir.assert_structural_equal(res[0][0], i0[0]) @@ -673,13 +669,13 @@ def test_subspace_division(): tvm.ir.assert_structural_equal(res[2][0], T.bool(True)) tvm.ir.assert_structural_equal(res[2][1], (floormod(j0[0], 4) * 2) + i3[0] < 7) - res1 = tvm.arith.detect_iter_map([res[0][1], res[1][1]], var_dom([j0, i3])).indices + res1 = tvm.sym.detect_iter_map([res[0][1], res[1][1]], var_dom([j0, i3])).indices assert len(res1) == 2 - res2 = tvm.arith.detect_iter_map([res[0][0], res[1][0]], var_dom([i0])).indices + res2 = tvm.sym.detect_iter_map([res[0][0], res[1][0]], var_dom([i0])).indices assert len(res2) == 2 # compound 1.6 - res = tvm.arith.subspace_divide( + res = tvm.sym.subspace_divide( [k0[0], k1[0]], var_dom([i0, j0, i3]), [i3[0]], tvm.tirx.all(k0[0] < 7, k1[0] < 7) ) res = convert_division(res) @@ -698,9 +694,7 @@ def test_subspace_division(): i2 = ifuse([j2, j3]) # compound 2.1 - res = tvm.arith.subspace_divide( - [i0[0], i1[0], i2[0]], var_dom([j0, l0, l1, j3]), [l1[0], j3[0]] - ) + res = tvm.sym.subspace_divide([i0[0], i1[0], i2[0]], var_dom([j0, l0, l1, j3]), [l1[0], j3[0]]) res = convert_division(res) assert len(res) == 4 tvm.ir.assert_structural_equal(res[0][0], (j0[0] * 2) + l0[0]) @@ -710,13 +704,13 @@ def test_subspace_division(): tvm.ir.assert_structural_equal(res[2][0], T.int32(0)) tvm.ir.assert_structural_equal(res[2][1], (floormod(l1[0], 3) * 3) + j3[0]) - res1 = tvm.arith.detect_iter_map([res[0][1], res[1][1], res[2][1]], var_dom([l1, j3])).indices + res1 = tvm.sym.detect_iter_map([res[0][1], res[1][1], res[2][1]], var_dom([l1, j3])).indices assert len(res1) == 3 - res2 = tvm.arith.detect_iter_map([res[0][0], res[1][0], res[2][0]], var_dom([j0, l0])).indices + res2 = tvm.sym.detect_iter_map([res[0][0], res[1][0], res[2][0]], var_dom([j0, l0])).indices assert len(res2) == 3 # compound 2.2 - res = tvm.arith.subspace_divide( + res = tvm.sym.subspace_divide( [i0[0], i1[0], i2[0]], var_dom([j0, l0, l1, j3]), [l0[0], l1[0], j3[0]] ) res = convert_division(res) @@ -728,22 +722,18 @@ def test_subspace_division(): tvm.ir.assert_structural_equal(res[2][0], T.int32(0)) tvm.ir.assert_structural_equal(res[2][1], (floormod(l0[0] * 6 + l1[0], 3) * 3) + j3[0]) - res1 = tvm.arith.detect_iter_map( - [res[0][1], res[1][1], res[2][1]], var_dom([l0, l1, j3]) - ).indices + res1 = tvm.sym.detect_iter_map([res[0][1], res[1][1], res[2][1]], var_dom([l0, l1, j3])).indices assert len(res1) == 3 - res2 = tvm.arith.detect_iter_map([res[0][0], res[1][0], res[2][0]], var_dom([j0])).indices + res2 = tvm.sym.detect_iter_map([res[0][0], res[1][0], res[2][0]], var_dom([j0])).indices assert len(res2) == 3 # compound 2.3 - res = tvm.arith.subspace_divide( - [i0[0], i1[0], i2[0]], var_dom([j0, l0, l1, j3]), [l0[0], j3[0]] - ) + res = tvm.sym.subspace_divide([i0[0], i1[0], i2[0]], var_dom([j0, l0, l1, j3]), [l0[0], j3[0]]) res = convert_division(res) assert len(res) == 0 # compound 2.4 - res = tvm.arith.subspace_divide( + res = tvm.sym.subspace_divide( [i0[0], i1[0], i2[0]], var_dom([j0, l0, l1, j3]), [l1[0], j3[0]], @@ -760,13 +750,13 @@ def test_subspace_division(): tvm.ir.assert_structural_equal(res[3][0], (j0[0] * 2) + l0[0] < 7) tvm.ir.assert_structural_equal(res[3][1], (floormod(l1[0], 3) * 3) + j3[0] < 8) - res1 = tvm.arith.detect_iter_map([res[0][1], res[1][1], res[2][1]], var_dom([l1, j3])).indices + res1 = tvm.sym.detect_iter_map([res[0][1], res[1][1], res[2][1]], var_dom([l1, j3])).indices assert len(res1) == 3 - res2 = tvm.arith.detect_iter_map([res[0][0], res[1][0], res[2][0]], var_dom([j0, l0])).indices + res2 = tvm.sym.detect_iter_map([res[0][0], res[1][0], res[2][0]], var_dom([j0, l0])).indices assert len(res2) == 3 # compound 2.5 - res = tvm.arith.subspace_divide( + res = tvm.sym.subspace_divide( [i0[0], i1[0], i2[0]], var_dom([j0, l0, l1, j3]), [j3[0]], i2[0] < 8 ) res = convert_division(res) @@ -780,11 +770,11 @@ def test_subspace_divide_accepts_external_analyzer(): root_iters = {i: tvm.ir.Range(0, 4), j: tvm.ir.Range(0, tile)} bindings = [j * tile + i] - assert len(tvm.arith.subspace_divide(bindings, root_iters, [i])) == 0 + assert len(tvm.sym.subspace_divide(bindings, root_iters, [i])) == 0 - analyzer = tvm.arith.Analyzer() + analyzer = tvm.sym.Analyzer() analyzer.bind(tile, T.int32(4)) - res = tvm.arith.subspace_divide(bindings, root_iters, [i], analyzer=analyzer) + res = tvm.sym.subspace_divide(bindings, root_iters, [i], analyzer=analyzer) res = convert_division(res) assert len(res) == 2 @@ -798,7 +788,7 @@ def test_subspace_divide_trivial_iters(): z = tvm.tirx.Var("z", "int32") # trivial 1.1 - res = tvm.arith.subspace_divide( + res = tvm.sym.subspace_divide( [x * 16 + y], var_dom([(x, 1), (y, 16)]), [y], simplify_trivial_iterators=False ) res = convert_division(res) @@ -807,7 +797,7 @@ def test_subspace_divide_trivial_iters(): tvm.ir.assert_structural_equal(res[0][1], y) # trivial 1.2 - res = tvm.arith.subspace_divide( + res = tvm.sym.subspace_divide( [x, y], var_dom([(x, 1), (y, 1)]), [y], @@ -846,52 +836,48 @@ def test_complex(): i0 = ifuse([j0, j1], 200) i1 = ifuse([j2, j3], 50) - n0_mark = tvm.arith.IterMark(n0[0], n0[1]) - n1_mark = tvm.arith.IterMark(n1[0], n1[1]) - l0_mark = tvm.arith.IterMark(l0[0], l0[1]) - l1_mark = tvm.arith.IterMark(l1[0], l1[1]) - m1_mark = tvm.arith.IterMark(m1[0], m1[1]) - l3_mark = tvm.arith.IterMark(l3[0], l3[1]) + n0_mark = tvm.sym.IterMark(n0[0], n0[1]) + n1_mark = tvm.sym.IterMark(n1[0], n1[1]) + l0_mark = tvm.sym.IterMark(l0[0], l0[1]) + l1_mark = tvm.sym.IterMark(l1[0], l1[1]) + m1_mark = tvm.sym.IterMark(m1[0], m1[1]) + l3_mark = tvm.sym.IterMark(l3[0], l3[1]) - m0_expr = tvm.arith.IterSumExpr( + m0_expr = tvm.sym.IterSumExpr( [ - tvm.arith.IterSplitExpr(n0_mark, 1, n0[1], 4), - tvm.arith.IterSplitExpr(n1_mark, 1, n1[1], 1), + tvm.sym.IterSplitExpr(n0_mark, 1, n0[1], 4), + tvm.sym.IterSplitExpr(n1_mark, 1, n1[1], 1), ], 0, ) - m0_mark = tvm.arith.IterMark(m0_expr, 6) - l2_expr = tvm.arith.IterSumExpr( - [tvm.arith.IterSplitExpr(m0_mark, 1, 6, 3), tvm.arith.IterSplitExpr(m1_mark, 1, m1[1], 1)], + m0_mark = tvm.sym.IterMark(m0_expr, 6) + l2_expr = tvm.sym.IterSumExpr( + [tvm.sym.IterSplitExpr(m0_mark, 1, 6, 3), tvm.sym.IterSplitExpr(m1_mark, 1, m1[1], 1)], 0, ) - l2_mark = tvm.arith.IterMark(l2_expr, 16) - k0_expr = tvm.arith.IterSplitExpr(l0_mark, 2, 2, 4) - k1_expr = tvm.arith.IterSplitExpr(l1_mark, 2, 4, 1) - k2_expr = tvm.arith.IterSplitExpr(l2_mark, 4, 4, 8) - k3_expr = tvm.arith.IterSplitExpr(l3_mark, 4, 8, 1) - k4_expr = tvm.arith.IterSplitExpr(l0_mark, 1, 2, 30) - k5_expr = tvm.arith.IterSplitExpr(l1_mark, 1, 2, 15) - k6_expr = tvm.arith.IterSplitExpr(l2_mark, 1, 4, 4) - k7_expr = tvm.arith.IterSplitExpr(l3_mark, 1, 4, 1) + l2_mark = tvm.sym.IterMark(l2_expr, 16) + k0_expr = tvm.sym.IterSplitExpr(l0_mark, 2, 2, 4) + k1_expr = tvm.sym.IterSplitExpr(l1_mark, 2, 4, 1) + k2_expr = tvm.sym.IterSplitExpr(l2_mark, 4, 4, 8) + k3_expr = tvm.sym.IterSplitExpr(l3_mark, 4, 8, 1) + k4_expr = tvm.sym.IterSplitExpr(l0_mark, 1, 2, 30) + k5_expr = tvm.sym.IterSplitExpr(l1_mark, 1, 2, 15) + k6_expr = tvm.sym.IterSplitExpr(l2_mark, 1, 4, 4) + k7_expr = tvm.sym.IterSplitExpr(l3_mark, 1, 4, 1) - j0_expr = tvm.arith.IterSumExpr([k0_expr, k1_expr], 0) - j0_mark = tvm.arith.IterMark(j0_expr, 7) - i0_expr = tvm.arith.IterSumExpr( - [tvm.arith.IterSplitExpr(j0_mark, 1, 7, 32), k2_expr, k3_expr], 0 - ) + j0_expr = tvm.sym.IterSumExpr([k0_expr, k1_expr], 0) + j0_mark = tvm.sym.IterMark(j0_expr, 7) + i0_expr = tvm.sym.IterSumExpr([tvm.sym.IterSplitExpr(j0_mark, 1, 7, 32), k2_expr, k3_expr], 0) - j3_expr = tvm.arith.IterSumExpr([k6_expr, k7_expr], 0) - j3_mark = tvm.arith.IterMark(j3_expr, 15) - i1_expr = tvm.arith.IterSumExpr( - [k4_expr, k5_expr, tvm.arith.IterSplitExpr(j3_mark, 1, 15, 1)], 0 - ) + j3_expr = tvm.sym.IterSumExpr([k6_expr, k7_expr], 0) + j3_mark = tvm.sym.IterMark(j3_expr, 15) + i1_expr = tvm.sym.IterSumExpr([k4_expr, k5_expr, tvm.sym.IterSplitExpr(j3_mark, 1, 15, 1)], 0) - i0_mark = tvm.arith.IterMark(i0_expr, i0[1]) - i1_mark = tvm.arith.IterMark(i1_expr, i1[1]) + i0_mark = tvm.sym.IterMark(i0_expr, i0[1]) + i1_mark = tvm.sym.IterMark(i1_expr, i1[1]) - i0_final = tvm.arith.IterSumExpr([tvm.arith.IterSplitExpr(i0_mark, 1, i0[1], 1)], 0) - i1_final = tvm.arith.IterSumExpr([tvm.arith.IterSplitExpr(i1_mark, 1, i1[1], 1)], 0) + i0_final = tvm.sym.IterSumExpr([tvm.sym.IterSplitExpr(i0_mark, 1, i0[1], 1)], 0) + i1_final = tvm.sym.IterSumExpr([tvm.sym.IterSplitExpr(i1_mark, 1, i1[1], 1)], 0) assert_iter_sum_pattern( {i0[0]: (200, 0, 1, i0_final), i1[0]: (50, 0, 1, i1_final)}, @@ -909,7 +895,7 @@ def test_complex(): ) # subspace_division - res = tvm.arith.subspace_divide( + res = tvm.sym.subspace_divide( [i0[0], i1[0]], var_dom([l0, l1, n0, n1, m1, l3]), [n0[0], n1[0], m1[0], l3[0]], @@ -951,21 +937,21 @@ def test_normalize_iter_map_to_expr(): xo, xi = isplit((x, 10), 5) yo, yi = isplit((y, 9), 3) z = ifuse([yo, xo, yi]) - res = tvm.arith.detect_iter_map([z[0], xi[0]], var_dom([(x, 10), (y, 9)])) + res = tvm.sym.detect_iter_map([z[0], xi[0]], var_dom([(x, 10), (y, 9)])) tvm.ir.assert_structural_equal( - tvm.arith.normalize_iter_map_to_expr(res.indices[0]), + tvm.sym.normalize_iter_map_to_expr(res.indices[0]), fld(y, 3) * 6 + fld(x, 5) * 3 + flm(y, 3), ) - tvm.ir.assert_structural_equal(tvm.arith.normalize_iter_map_to_expr(res.indices[1]), flm(x, 5)) + tvm.ir.assert_structural_equal(tvm.sym.normalize_iter_map_to_expr(res.indices[1]), flm(x, 5)) # iter mark wrap a complex expr - split = tvm.arith.IterSplitExpr(tvm.arith.IterMark(x * y + 1, 1024), 1, 1024, 1) - tvm.ir.assert_structural_equal(tvm.arith.normalize_iter_map_to_expr(split), x * y + 1) + split = tvm.sym.IterSplitExpr(tvm.sym.IterMark(x * y + 1, 1024), 1, 1024, 1) + tvm.ir.assert_structural_equal(tvm.sym.normalize_iter_map_to_expr(split), x * y + 1) def test_inverse_affine_iter_map(): - analyzer = tvm.arith.Analyzer() + analyzer = tvm.sym.Analyzer() l0 = create_iter("l0", 64) l1 = create_iter("l1", 64) l2 = create_iter("l2", 64) @@ -975,11 +961,11 @@ def test_inverse_affine_iter_map(): l1_0, l1_1 = isplit(l1, 4) l0_1_l1_1_fused = ifuse([l0_1, l1_1]) - iter_map = tvm.arith.detect_iter_map( + iter_map = tvm.sym.detect_iter_map( [l0_1_l1_1_fused[0], l0_0[0], l1_0[0]], var_dom([l0, l1]) ).indices outputs = [tvm.tirx.Var(f"output_{i}", "int32") for i in range(len(iter_map))] - res = tvm.arith.inverse_affine_iter_map(iter_map, outputs) + res = tvm.sym.inverse_affine_iter_map(iter_map, outputs) assert len(res) == 2 l0_inverse = floordiv(outputs[0], 4) + outputs[1] * 16 l1_inverse = floormod(outputs[0], 4) + outputs[2] * 4 @@ -994,11 +980,11 @@ def test_inverse_affine_iter_map(): l0_1_l2_1_l1_1_l2_0_fused = ifuse([l0_1, l2_1, l1_1, l2_0]) - iter_map = tvm.arith.detect_iter_map( + iter_map = tvm.sym.detect_iter_map( [l0_1_l2_1_l1_1_l2_0_fused[0], l0_0[0], l2_2[0], l1_0[0]], var_dom([l0, l1, l2]) ).indices outputs = [tvm.tirx.Var(f"output_{i}", "int32") for i in range(len(iter_map))] - res = tvm.arith.inverse_affine_iter_map(iter_map, outputs) + res = tvm.sym.inverse_affine_iter_map(iter_map, outputs) assert len(res) == 3 l0_inverse = floordiv(outputs[0], 64) + outputs[1] * 16 l1_inverse = floormod(floordiv(outputs[0], 4), 4) + outputs[3] * 4 @@ -1016,9 +1002,9 @@ def test_inverse_affine_iter_map(): l1_0, l1_1 = isplit(l1, 8) l2 = ifuse([l1_1, l1_0]) - iter_map = tvm.arith.detect_iter_map([l2[0]], var_dom([l0])).indices + iter_map = tvm.sym.detect_iter_map([l2[0]], var_dom([l0])).indices outputs = [tvm.tirx.Var(f"output_{i}", "int32") for i in range(len(iter_map))] - res = tvm.arith.inverse_affine_iter_map(iter_map, outputs) + res = tvm.sym.inverse_affine_iter_map(iter_map, outputs) assert len(res) == 1 l1_inverse = floormod(outputs[0], 8) * 8 + floordiv(outputs[0], 8) l0_inverse = floormod(l1_inverse, 4) * 16 + floordiv(l1_inverse, 4) @@ -1027,12 +1013,12 @@ def test_inverse_affine_iter_map(): def test_inverse_affine_map_trivial_iter(): - analyzer = tvm.arith.Analyzer() + analyzer = tvm.sym.Analyzer() l0 = create_iter("l0", 64) l1 = create_iter("l1", 64) - iter_map = tvm.arith.detect_iter_map([0, l0[0], l1[0]], var_dom([l0, l1])).indices + iter_map = tvm.sym.detect_iter_map([0, l0[0], l1[0]], var_dom([l0, l1])).indices outputs = [tvm.tirx.Var(f"output_{i}", "int32") for i in range(len(iter_map))] - res = tvm.arith.inverse_affine_iter_map(iter_map, outputs) + res = tvm.sym.inverse_affine_iter_map(iter_map, outputs) # output_0 is expected to be constant and it is not included in the inverse map assert len(res) == 2 assert analyzer.can_prove_equal(res[l0[0]], outputs[1]) @@ -1264,7 +1250,7 @@ def test_iter_map_simplify_predicate_fallback_requires_no_padding(): fused = tvm.tirx.Var("fused", "int64") predicate = fused % 2 == 0 unpadded_index = fused // 4 * 4 + fused % 4 - simplified = tvm.arith.iter_map_simplify( + simplified = tvm.sym.iter_map_simplify( [unpadded_index], var_dom([(fused, 1024)]), predicate=predicate, @@ -1279,7 +1265,7 @@ def test_iter_map_simplify_predicate_fallback_requires_no_padding(): # The parity predicate is not a bound constraint, so IterMapSimplify falls back to # detecting the map without it. That fallback requires left-padding the iterator; # discarding the corresponding padding predicate would change the index expression. - simplified = tvm.arith.iter_map_simplify( + simplified = tvm.sym.iter_map_simplify( [index], var_dom([(fused, 1024), (kernel, 3)]), predicate=predicate, @@ -1335,7 +1321,7 @@ def test_iter_map_simplify_unit_loop_order(): def assert_normalize_to_iter_sum(index, input_iters, args, base): - """Assert the result of arith.normalize_to_iter_sum is correct + """Assert the result of sym.normalize_to_iter_sum is correct Parameters ---------- @@ -1343,24 +1329,24 @@ def assert_normalize_to_iter_sum(index, input_iters, args, base): The index to be normalized input_iters : Mapping[Var, Range] The input iterators - args : List[Union[tvm.arith.IterSplitExpr, Tuple[Expr, Expr]]] + args : List[Union[tvm.sym.IterSplitExpr, Tuple[Expr, Expr]]] The expected result. Ordered list of args of the expected IterSumExpr. Each arg can be either IterSplitExpr or a tuple of (Expr, Expr) where the first element is the iterator normalized to Expr and the second element is the scale. base : tvm.tirx.Expr The expected base """ - res = tvm.arith.normalize_to_iter_sum(index, input_iters) + res = tvm.sym.normalize_to_iter_sum(index, input_iters) - assert isinstance(res, tvm.arith.IterSumExpr) + assert isinstance(res, tvm.sym.IterSumExpr) assert len(res.args) == len(args) for split, item in zip(res.args, args): - if isinstance(item, tvm.arith.IterSplitExpr): + if isinstance(item, tvm.sym.IterSplitExpr): tvm.ir.assert_structural_equal(split, item) continue tvm.testing.assert_prim_expr_equal(split.scale, item[1]) tvm.testing.assert_prim_expr_equal( - tvm.arith.normalize_iter_map_to_expr(split), item[0] * item[1] + tvm.sym.normalize_iter_map_to_expr(split), item[0] * item[1] ) tvm.testing.assert_prim_expr_equal(res.base, base) @@ -1417,8 +1403,8 @@ def test_normalize_to_iter_sum(): x // 5, var_dom([(x, 4096)]), [ - tvm.arith.IterSplitExpr( - tvm.arith.IterMark(x, 4096), + tvm.sym.IterSplitExpr( + tvm.sym.IterMark(x, 4096), lower_factor=tvm.tirx.const(5, "int64"), extent=tvm.tirx.const(820, "int64"), scale=tvm.tirx.const(1, "int64"), @@ -1441,9 +1427,9 @@ def test_normalize_to_iter_sum_accepts_external_analyzer(): tile = tvm.tirx.Var("tile", "int32") input_iters = {i: tvm.ir.Range(0, 16)} - analyzer = tvm.arith.Analyzer() + analyzer = tvm.sym.Analyzer() analyzer.bind(tile, T.int32(4)) - res = tvm.arith.normalize_to_iter_sum(i // tile, input_iters, analyzer=analyzer) + res = tvm.sym.normalize_to_iter_sum(i // tile, input_iters, analyzer=analyzer) assert len(res.args) == 1 tvm.testing.assert_prim_expr_equal(res.args[0].lower_factor, tile) @@ -1466,7 +1452,7 @@ def test_detect_iter_map_with_bufferload_recursion(): j: tvm.ir.Range(tvm.tirx.const(0, "int32"), m), } - result = tvm.arith.detect_iter_map(indices, iter_vars) + result = tvm.sym.detect_iter_map(indices, iter_vars) assert len(result.indices) == 0 @@ -1476,12 +1462,12 @@ def test_detect_iter_map_accepts_external_analyzer(): iter_vars = {i: tvm.ir.Range(0, 16)} # Without knowing `tile`, the floormod cannot be recognized as an iterator. - assert len(tvm.arith.detect_iter_map([i % tile], iter_vars).indices) == 0 + assert len(tvm.sym.detect_iter_map([i % tile], iter_vars).indices) == 0 - analyzer = tvm.arith.Analyzer() + analyzer = tvm.sym.Analyzer() analyzer.bind(tile, T.int32(4)) # The external analyzer supplies `tile == 4`, allowing detection to succeed. - assert len(tvm.arith.detect_iter_map([i % tile], iter_vars, analyzer=analyzer).indices) == 1 + assert len(tvm.sym.detect_iter_map([i % tile], iter_vars, analyzer=analyzer).indices) == 1 def test_iter_map_simplify_accepts_external_analyzer(): @@ -1489,9 +1475,9 @@ def test_iter_map_simplify_accepts_external_analyzer(): tile = tvm.tirx.Var("tile", "int32") iter_vars = {i: tvm.ir.Range(0, 32)} - analyzer = tvm.arith.Analyzer() + analyzer = tvm.sym.Analyzer() analyzer.bind(tile, T.int32(8)) - simplified = tvm.arith.iter_map_simplify([i % tile], iter_vars, analyzer=analyzer) + simplified = tvm.sym.iter_map_simplify([i % tile], iter_vars, analyzer=analyzer) tvm.ir.assert_structural_equal(simplified, [i % 8]) diff --git a/tests/python/arith/test_arith_modular_set.py b/tests/python/sym/test_sym_modular_set.py similarity index 91% rename from tests/python/arith/test_arith_modular_set.py rename to tests/python/sym/test_sym_modular_set.py index 9a9d35b48397..a8b52512cefd 100644 --- a/tests/python/arith/test_arith_modular_set.py +++ b/tests/python/sym/test_sym_modular_set.py @@ -21,7 +21,7 @@ def test_cast(): - analyzer = tvm.arith.Analyzer() + analyzer = tvm.sym.Analyzer() x = tvm.tirx.Var("x", "int8") m = analyzer.modular_set((x * 3).astype("uint32")) assert m.coeff == 3 @@ -32,7 +32,7 @@ def test_cast(): def test_add_sub(): - analyzer = tvm.arith.Analyzer() + analyzer = tvm.sym.Analyzer() x, y = tvm.tirx.Var("x", "int64"), tvm.tirx.Var("y", "int64") m = analyzer.modular_set(x * 6 + y * 4) assert m.coeff == 2 @@ -45,7 +45,7 @@ def test_add_sub(): def test_mul(): - analyzer = tvm.arith.Analyzer() + analyzer = tvm.sym.Analyzer() x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32") m = analyzer.modular_set((x * 4 + 2) * (y * 6 + 1)) assert m.coeff == 4 @@ -53,7 +53,7 @@ def test_mul(): def test_shift_left(): - analyzer = tvm.arith.Analyzer() + analyzer = tvm.sym.Analyzer() x, y = te.var("x"), te.var("y") m = analyzer.modular_set((x * 4 + 2) << 2) assert m.coeff == 16 @@ -61,7 +61,7 @@ def test_shift_left(): def test_floormod(): - analyzer = tvm.arith.Analyzer() + analyzer = tvm.sym.Analyzer() x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32") m = analyzer.modular_set(tvm.tirx.floormod(x * 128 + y * 4, 256)) assert m.coeff == 4 @@ -69,7 +69,7 @@ def test_floormod(): def test_div_shift(): - analyzer = tvm.arith.Analyzer() + analyzer = tvm.sym.Analyzer() x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32") # not sure if x is non-negative tdiv = tvm.tirx.truncdiv @@ -85,14 +85,14 @@ def test_div_shift(): assert m.coeff == 2 assert m.base == 1 # x is non-negative - analyzer.update(x, tvm.arith.ConstIntBound(0, 100)) + analyzer.update(x, tvm.sym.ConstIntBound(0, 100)) m = analyzer.modular_set(tdiv(x * 4 + 2, 2)) assert m.coeff == 2 assert m.base == 1 def test_mod(): - analyzer = tvm.arith.Analyzer() + analyzer = tvm.sym.Analyzer() x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32") tmod = tvm.tirx.truncmod fmod = tvm.tirx.floormod @@ -112,14 +112,14 @@ def test_mod(): assert m.coeff == 4 assert m.base == 3 # x is non-negative - analyzer.update(x, tvm.arith.ConstIntBound(0, 100)) + analyzer.update(x, tvm.sym.ConstIntBound(0, 100)) m = analyzer.modular_set(tmod(x * 4 + 3, 2)) assert m.coeff == 2 assert m.base == 1 def test_min_max_select(): - analyzer = tvm.arith.Analyzer() + analyzer = tvm.sym.Analyzer() x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32") m = analyzer.modular_set(tvm.tirx.min(x * 3, y * 9)) assert m.coeff == 3 @@ -137,7 +137,7 @@ def test_min_max_select(): def test_mix_index(): a = tvm.tirx.Var("a", "int32") b = tvm.tirx.Var("b", "int32") - analyzer = tvm.arith.Analyzer() + analyzer = tvm.sym.Analyzer() tdiv = tvm.tirx.truncdiv m = analyzer.modular_set(a * 4 + b * 6 + 7) assert m.coeff == 2 @@ -167,7 +167,7 @@ def test_mix_index(): def test_constraint_scope(): a = tvm.tirx.Var("a", "int32") b = tvm.tirx.Var("b", "int32") - analyzer = tvm.arith.Analyzer() + analyzer = tvm.sym.Analyzer() tmod = tvm.tirx.truncmod with analyzer.constraint_scope(tmod(b, 4) == 2): @@ -189,7 +189,7 @@ def test_constraint_scope(): def test_intersect(): a = tvm.tirx.Var("a", "int32") - analyzer = tvm.arith.Analyzer() + analyzer = tvm.sym.Analyzer() tmod = tvm.tirx.truncmod with analyzer.constraint_scope(tmod(a, 4) == 1): with analyzer.constraint_scope(tmod(a, 3) == 1): @@ -206,7 +206,7 @@ def test_intersect(): def test_let(): - analyzer = tvm.arith.Analyzer() + analyzer = tvm.sym.Analyzer() x = tvm.tirx.Var("x", "int32") y = tvm.tirx.Var("y", "int32") m = analyzer.modular_set(tvm.tirx.Let(x, y * 10, x + 1)) @@ -215,7 +215,7 @@ def test_let(): def test_bitwise_and(): - analyzer = tvm.arith.Analyzer() + analyzer = tvm.sym.Analyzer() x = tvm.tirx.Var("x", "int32") y = tvm.tirx.Var("y", "int32") diff --git a/tests/python/arith/test_arith_rewrite_simplify.py b/tests/python/sym/test_sym_rewrite_simplify.py similarity index 99% rename from tests/python/arith/test_arith_rewrite_simplify.py rename to tests/python/sym/test_sym_rewrite_simplify.py index 6dbe2a65a1e2..5236fd0938dd 100644 --- a/tests/python/arith/test_arith_rewrite_simplify.py +++ b/tests/python/sym/test_sym_rewrite_simplify.py @@ -68,10 +68,10 @@ def __name__(self): class BaseCompare: - extensions = tvm.arith.Extension.NoExtensions + extensions = tvm.sym.Extension.NoExtensions def test_simplify(self, test_case): - analyzer = tvm.arith.Analyzer() + analyzer = tvm.sym.Analyzer() analyzer.enabled_extensions = self.extensions if inspect.isclass(test_case.expected) and issubclass(test_case.expected, Exception): @@ -752,7 +752,7 @@ class TestFloorModPadded(BaseCompare): def test_uint_floormod_const_fold(): - analyzer = tvm.arith.Analyzer() + analyzer = tvm.sym.Analyzer() expr = flm(8192, T.uint32(128)) tvm.ir.assert_structural_equal(expr, T.uint32(0)) assert analyzer.can_prove_equal(expr, 0) @@ -1173,7 +1173,7 @@ class TestComparisons(BaseCompare): class TestComparisonOfProductAndSum(BaseCompare): - extensions = tvm.arith.Extension.ComparisonOfProductAndSum + extensions = tvm.sym.Extension.ComparisonOfProductAndSum x, y, z = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32"), tvm.tirx.Var("z", "int32") @@ -1352,7 +1352,7 @@ def test_allow_uint_as_index(): w = tirx.Var("w", "int32") lane = tirx.Var("lane", "int32") x = tirx.Var("x", "uint32") - analyzer = tvm.arith.Analyzer() + analyzer = tvm.sym.Analyzer() analyzer.bind(lane, tvm.ir.Range.from_min_extent(0, 32)) analyzer.bind(w, tvm.ir.Range.from_min_extent(0, 4)) s_off = ( @@ -1366,12 +1366,12 @@ def test_allow_uint_as_index(): # the flag is on: under the no-overflow assertion any c2 != 0 works. np2 = flm(x * T.uint32(12), T.uint32(12)) assert not analyzer.can_prove_equal(np2, 0) - with tvm.arith.allow_uint_as_index(): + with tvm.sym.allow_uint_as_index(): assert analyzer.can_prove_equal(check, 0) assert analyzer.can_prove_equal(flm(fld(s_off, T.uint32(32)), T.uint32(2)), 0) assert analyzer.can_prove_equal(np2, 0) # nested scopes compose - with tvm.arith.allow_uint_as_index(): + with tvm.sym.allow_uint_as_index(): assert analyzer.can_prove_equal(check, 0) assert analyzer.can_prove_equal(check, 0) # the mode is scoped: disabled again afterwards diff --git a/tests/python/arith/test_arith_simplify.py b/tests/python/sym/test_sym_simplify.py similarity index 92% rename from tests/python/arith/test_arith_simplify.py rename to tests/python/sym/test_sym_simplify.py index 09b6a7e1aa0a..29b5b5342ccc 100644 --- a/tests/python/arith/test_arith_simplify.py +++ b/tests/python/sym/test_sym_simplify.py @@ -24,7 +24,7 @@ def test_simplify_reshape_flattened_index(): - ana = tvm.arith.Analyzer() + ana = tvm.sym.Analyzer() i0 = tirx.Var("i0", "int64") i1 = tirx.Var("i1", "int64") @@ -54,29 +54,29 @@ def test_simplify_reshape_flattened_index(): def test_can_prove_self_identity(dtype): - ana = tvm.arith.Analyzer() + ana = tvm.sym.Analyzer() n = tirx.Var("n", dtype) assert ana.can_prove(n == n) def test_can_prove_self_equal_to_self(dtype): - ana = tvm.arith.Analyzer() + ana = tvm.sym.Analyzer() n = tirx.Var("n", dtype) assert ana.can_prove_equal(n, n) def test_simplify_symbolic_comparison(): - ana = tvm.arith.Analyzer() + ana = tvm.sym.Analyzer() i0 = tirx.Var("i0", "int64") i1 = tirx.Var("i1", "int64") n, m = tvm.tirx.Var("n", "int64"), tvm.tirx.Var("m", "int64") outer = (n + 31) // 32 - PS = tvm.arith.ProofStrength + PS = tvm.sym.ProofStrength - non_negative = tvm.arith.ConstIntBound(0, tvm.arith.ConstIntBound.POS_INF) + non_negative = tvm.sym.ConstIntBound(0, tvm.sym.ConstIntBound.POS_INF) ana.update(n, non_negative) ana.update(m, non_negative) ana.bind(i0, tvm.ir.Range(0, outer)) @@ -90,7 +90,7 @@ def test_simplify_symbolic_comparison(): def test_regression_simplify_inf_recursion(): - ana = tvm.arith.Analyzer() + ana = tvm.sym.Analyzer() cond = tirx.Var("cond", "int32") res = (tvm.tirx.NE(cond, 0).astype("int8") - tvm.tirx.NE(cond, 0).astype("int8")).astype( @@ -102,7 +102,7 @@ def test_regression_simplify_inf_recursion(): def test_bind_allow_override(): - ana = tvm.arith.Analyzer() + ana = tvm.sym.Analyzer() x = tirx.Var("x", "int64") ana.bind(x, tvm.ir.Range(0, 10)) @@ -120,7 +120,7 @@ def test_simplify_floor_mod_with_linear_offset(): """ Test that the floor_mod is simplified correctly when the offset is linear. """ - ana = tvm.arith.Analyzer() + ana = tvm.sym.Analyzer() past_decoder_sequence_length = tirx.Var("past_decoder_sequence_length", "int64") expr1 = (past_decoder_sequence_length + 1) * 64 divisor1 = (past_decoder_sequence_length + 1) * 32 @@ -131,7 +131,7 @@ def test_simplify_floor_mod_with_linear_offset(): def test_simplify_uint_floormod_const_scale_divisible(): """uint32 floormod(x * c1, c2) -> 0 when c1 % c2 == 0 (overflow-free).""" - ana = tvm.arith.Analyzer() + ana = tvm.sym.Analyzer() q = tirx.Var("q_stage_idx", "uint32") expr = q * tirx.Cast("uint32", 128) mod = expr % tirx.const(4, "uint32") @@ -142,7 +142,7 @@ def test_simplify_uint_floormod_const_scale_divisible(): def test_simplify_float_division(): # Test for the discussion: # https://discuss.tvm.apache.org/t/discuss-is-constant-division-to-multiplication-rewrite-in-tvm-necessary/18615 - ana = tvm.arith.Analyzer() + ana = tvm.sym.Analyzer() x = tirx.Var("x", "float32") ry = x / 27 # in old version, the division will be rewritten into x * T.float32(1 / 27) diff --git a/tests/python/arith/test_arith_z3.py b/tests/python/sym/test_sym_z3.py similarity index 99% rename from tests/python/arith/test_arith_z3.py rename to tests/python/sym/test_sym_z3.py index 85bbc2c3c9a0..34bc23a2b140 100644 --- a/tests/python/arith/test_arith_z3.py +++ b/tests/python/sym/test_sym_z3.py @@ -24,7 +24,7 @@ import tvm import tvm.testing from tvm import tirx -from tvm.arith import Analyzer, ProofStrength +from tvm.sym import Analyzer, ProofStrength # The Z3 prover is only consulted at the kSymbolicBound strength so the common # default path never pays the prover cost. @@ -101,7 +101,7 @@ def test_z3_context_scope_clone_lifetime(): c = tirx.Var("c", "int32") expr = ((b - a) // c) * c + a <= b - with tvm.arith.Z3ContextScope(): + with tvm.sym.Z3ContextScope(): analyzer = Analyzer() analyzer.bind(a, tvm.ir.Range(1, 100000)) analyzer.bind(b, tvm.ir.Range(1, 100000)) @@ -110,7 +110,7 @@ def test_z3_context_scope_clone_lifetime(): # Clone while a different scope is active. The clone must adopt the # source Analyzer's context before copying any Z3 handles. - with tvm.arith.Z3ContextScope(): + with tvm.sym.Z3ContextScope(): cloned = analyzer.clone() del analyzer diff --git a/tests/python/tirx-base/test_tir_index_map.py b/tests/python/tirx-base/test_tir_index_map.py index 75f5332f7393..13f07c184971 100644 --- a/tests/python/tirx-base/test_tir_index_map.py +++ b/tests/python/tirx-base/test_tir_index_map.py @@ -32,7 +32,7 @@ def assert_equal_index_map(map1: IndexMap, map2: IndexMap) -> None: iters_2 = map2.final_indices assert len(iters_1) == len(iters_2) - analyzer = tvm.arith.Analyzer() + analyzer = tvm.sym.Analyzer() for iter1, iter2 in zip(iters_1, iters_2): assert analyzer.can_prove_equal(iter1, iter2) @@ -50,7 +50,7 @@ def test_index_mapping(): def test_map_indices_accepts_external_analyzer(): tile = tvm.tirx.Var("tile", "int32") index_map = IndexMap.from_func(lambda i: [i // tile], index_dtype="int32") - analyzer = tvm.arith.Analyzer() + analyzer = tvm.sym.Analyzer() unsimplified = index_map.map_indices([T.int32(32)])[0] analyzer.bind(tile, T.int32(16)) @@ -63,7 +63,7 @@ def test_map_indices_accepts_external_analyzer(): def test_map_shape_accepts_external_analyzer(): tile = tvm.tirx.Var("tile", "int32") index_map = IndexMap.from_func(lambda i: [i // tile, i % tile], index_dtype="int32") - analyzer = tvm.arith.Analyzer() + analyzer = tvm.sym.Analyzer() analyzer.bind(tile, T.int32(16)) mapped_shape = index_map.map_shape([T.int32(32)], analyzer=analyzer) @@ -79,7 +79,7 @@ def test_is_equivalent_to_accepts_external_analyzer(): # Without binding `tile`, the symbolic map cannot be proven equivalent. assert not concrete.is_equivalent_to(symbolic) - analyzer = tvm.arith.Analyzer() + analyzer = tvm.sym.Analyzer() analyzer.bind(tile, T.int32(4)) assert concrete.is_equivalent_to(symbolic, analyzer=analyzer) @@ -111,7 +111,7 @@ def test_inverse_preserves_passthrough_var_names(): def test_inverse_accepts_external_analyzer(): tile = tvm.tirx.Var("tile", "int32") index_map = IndexMap.from_func(lambda i: [i // tile, i % tile], index_dtype="int32") - analyzer = tvm.arith.Analyzer() + analyzer = tvm.sym.Analyzer() analyzer.bind(tile, T.int32(16)) inverse = index_map.inverse([T.int32(32)], analyzer=analyzer) @@ -248,7 +248,7 @@ def test_nonsurjective_inverse(padding_test_case): # Can't use analyzer.can_prove_equal, because it can't simplify # expressions like `(4*i+j >= 14) - (4*i+j >= 14)`. - analyzer = tvm.arith.Analyzer() + analyzer = tvm.sym.Analyzer() expected_predicate = analyzer.simplify(expected_predicate) padding_predicate = analyzer.simplify(padding_predicate) tvm.ir.assert_structural_equal(padding_predicate, expected_predicate) @@ -257,7 +257,7 @@ def test_nonsurjective_inverse(padding_test_case): def test_non_surjective_inverse_accepts_external_analyzer(): tile = tvm.tirx.Var("tile", "int32") index_map = IndexMap.from_func(lambda i: [i // tile, i % tile], index_dtype="int32") - analyzer = tvm.arith.Analyzer() + analyzer = tvm.sym.Analyzer() analyzer.bind(tile, T.int32(16)) inverse, padding_predicate = index_map.non_surjective_inverse([T.int32(31)], analyzer=analyzer) @@ -290,7 +290,7 @@ def test_non_surjective_inverse_accepts_external_analyzer(): def test_non_surjective_inverse_does_not_bind_output_vars_to_external_analyzer(): tile = tvm.tirx.Var("tile", "int32") index_map = IndexMap.from_func(lambda i: [i // tile, i % tile], index_dtype="int32") - analyzer = tvm.arith.Analyzer() + analyzer = tvm.sym.Analyzer() analyzer.bind(tile, T.int32(16)) inverse, _ = index_map.non_surjective_inverse([T.int32(31)], analyzer=analyzer) diff --git a/tests/python/tirx-base/test_tir_op_types.py b/tests/python/tirx-base/test_tir_op_types.py index c20f051871e1..eaee54309b3c 100644 --- a/tests/python/tirx-base/test_tir_op_types.py +++ b/tests/python/tirx-base/test_tir_op_types.py @@ -232,7 +232,7 @@ def test_op_ptx_cp_async(): assert access_ptr.op.name == "tirx.tvm_access_ptr" assert access_ptr.args[1].op.name == "tirx.buffer_data" assert isinstance(access_ptr.args[1].args[0], tirx.Var) - simplified_offset = tvm.arith.Analyzer().simplify(access_ptr.args[2]) + simplified_offset = tvm.sym.Analyzer().simplify(access_ptr.args[2]) assert int(simplified_offset) == expected_offset diff --git a/tests/python/tirx-transform/test_tir_transform_lower_intrin.py b/tests/python/tirx-transform/test_tir_transform_lower_intrin.py index 86d986250e26..98cb476b19ae 100644 --- a/tests/python/tirx-transform/test_tir_transform_lower_intrin.py +++ b/tests/python/tirx-transform/test_tir_transform_lower_intrin.py @@ -112,7 +112,7 @@ def collect(node): assert len(address_calls) == 1 load = address_calls[0].args[0] assert isinstance(load, tvm.ir.TensorLoad) - assert int(tvm.arith.Analyzer().simplify(load.indices[0])) == 5 + assert int(tvm.sym.Analyzer().simplify(load.indices[0])) == 5 targets = ["c"] if env.has_llvm(): diff --git a/tests/python/tirx-transform/test_tir_transform_simplify.py b/tests/python/tirx-transform/test_tir_transform_simplify.py index f62dcf77bac7..fc74467177cf 100644 --- a/tests/python/tirx-transform/test_tir_transform_simplify.py +++ b/tests/python/tirx-transform/test_tir_transform_simplify.py @@ -710,7 +710,7 @@ def test_remove_transitively_provable_condition(): (tvm.tirx.all(i < j + 5, j < k + 7), i < k + 10, False), ] - analyzer = tvm.arith.Analyzer() + analyzer = tvm.sym.Analyzer() for priors, postulate, provable in test_cases: # well formed checker complains of undefined variables in condition diff --git a/tests/python/tirx/operator/tile_primitive/cuda/copy/test_gmem_smem.py b/tests/python/tirx/operator/tile_primitive/cuda/copy/test_gmem_smem.py index 33ae4252b77a..cd459e7b5870 100644 --- a/tests/python/tirx/operator/tile_primitive/cuda/copy/test_gmem_smem.py +++ b/tests/python/tirx/operator/tile_primitive/cuda/copy/test_gmem_smem.py @@ -480,7 +480,7 @@ def test_layout_permute_copy_preserves_smem_strides(): # S is K-tiled : s_off(tid) = (tid // 8) * 8 + (tid % 8) * 1024. # For tid=1 the two MUST differ — they're identical iff S was # collapsed to row-major (the regression). - analyzer = tvm.arith.Analyzer() + analyzer = tvm.sym.Analyzer() value_map = {tid_var: _IntImm("int32", 1)} s_off_at_1 = analyzer.simplify( tvm_ffi.structural_map( diff --git a/tests/python/tirx/operator/tile_primitive/cuda/copy/test_reg.py b/tests/python/tirx/operator/tile_primitive/cuda/copy/test_reg.py index 72702ec57ad0..06e75399517b 100644 --- a/tests/python/tirx/operator/tile_primitive/cuda/copy/test_reg.py +++ b/tests/python/tirx/operator/tile_primitive/cuda/copy/test_reg.py @@ -852,7 +852,6 @@ def _eval_const_layout_expr(expr, values): @pytest.mark.parametrize("case", ["wg", "wg_slice", "tcgen05"]) def test_reg_synthetic_tile_matches_thread_base_plus_outer_delta(case): - from tvm.arith import Analyzer from tvm.backend.cuda.tile_primitive.copy.vec_auto_reg import ( _build_atoms, _build_s_apply_layout, @@ -864,6 +863,7 @@ def test_reg_synthetic_tile_matches_thread_base_plus_outer_delta(case): _split_thread_loop, align_layouts_raw, ) + from tvm.sym import Analyzer from tvm.tirx.exec_scope import ExecScope from tvm.tirx.layout import ComposeLayout, wg_local_layout from tvm.tirx.operator.tile_primitive import DispatchContext diff --git a/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_tma.py b/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_tma.py index 62f19ad0682e..19d2a3c070e9 100644 --- a/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_tma.py +++ b/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_tma.py @@ -26,10 +26,10 @@ import tvm import tvm.testing -from tvm.arith import Analyzer from tvm.ir import PointerType, PrimType, Range from tvm.script import tirx as T from tvm.script.tirx import tile as Tx +from tvm.sym import Analyzer from tvm.testing import env from tvm.tirx import IntImm, StringImm, Var from tvm.tirx.cuda.tile_primitive.copy_async.tma import ( diff --git a/tests/python/tirx/test_layout.py b/tests/python/tirx/test_layout.py index aa4d24124bef..de616be5e4be 100644 --- a/tests/python/tirx/test_layout.py +++ b/tests/python/tirx/test_layout.py @@ -22,12 +22,12 @@ import pytest import tvm -from tvm.arith import Analyzer from tvm.ir import assert_structural_equal from tvm.ir.type import PointerType, PrimType from tvm.script import tirx as T from tvm.script.ir_builder import IRBuilder from tvm.script.ir_builder import tirx as Tx_builder +from tvm.sym import Analyzer from tvm.tirx import Var from tvm.tirx.cuda.tile_primitive.tma_utils import ( SwizzleMode, diff --git a/tests/python/tirx/test_parser_printer.py b/tests/python/tirx/test_parser_printer.py index 4fc72f04144b..93fa3898885a 100644 --- a/tests/python/tirx/test_parser_printer.py +++ b/tests/python/tirx/test_parser_printer.py @@ -1903,7 +1903,7 @@ def func() -> None: bufs = _collect_buffers(func) a_buf, b_buf = bufs["A"], bufs["B"] # 5 -> (5 // 4, 5 % 4) = (1, 1) -> 1 * 1024 + 1 * 64 - assert int(tvm.arith.Analyzer().simplify(b_buf.elem_offset - a_buf.elem_offset)) == 1088 + assert int(tvm.sym.Analyzer().simplify(b_buf.elem_offset - a_buf.elem_offset)) == 1088 assert [int(s) for s in b_buf.shape] == [16] assert_structural_equal(b_buf.layout, tvm.tirx.layout.TileLayout(T.S[(16,) : (1,)])) @@ -1941,11 +1941,11 @@ def func() -> None: a_buf, b_buf, c_buf = bufs["A"], bufs["B"], bufs["C"] # sub[1, 2:6]: drop dim 0 at 1 (1 * 256) then narrow dim 1 to [2, 6) (2 * 16) assert [int(s) for s in b_buf.shape] == [4, 16] - assert int(tvm.arith.Analyzer().simplify(b_buf.elem_offset - a_buf.elem_offset)) == 288 + assert int(tvm.sym.Analyzer().simplify(b_buf.elem_offset - a_buf.elem_offset)) == 288 assert_structural_equal(b_buf.layout, tvm.tirx.layout.TileLayout(T.S[(4, 16) : (16, 1)])) # sub[:, 1::2]: keep dim 0, split dim 1 into (4, 2) and fix the remainder at 1 assert [int(s) for s in c_buf.shape] == [4, 4, 16] - assert int(tvm.arith.Analyzer().simplify(c_buf.elem_offset - a_buf.elem_offset)) == 16 + assert int(tvm.sym.Analyzer().simplify(c_buf.elem_offset - a_buf.elem_offset)) == 16 assert_structural_equal( c_buf.layout, tvm.tirx.layout.TileLayout(T.S[(4, 4, 16) : (256, 32, 1)]) ) @@ -1993,14 +1993,14 @@ def test_buffer_sub_swizzle_commutation(): placements must be address-equivalent to the parent layout.""" def addr(buf, base, *coords): - analyzer = tvm.arith.Analyzer() + analyzer = tvm.sym.Analyzer() if len(coords) == 1: rel = buf.layout.apply(coords[0])["m"] else: rel = buf.layout.apply(*coords, shape=[int(s) for s in buf.shape])["m"] return int(analyzer.simplify((buf.elem_offset - base) + rel)) - analyzer = tvm.arith.Analyzer() + analyzer = tvm.sym.Analyzer() compose = T.ComposeLayout( 3, 3, 3, T.TileLayout(T.S[(4, 1024) : (1024, 1)]) ) # period = 2^(3+3+3) = 512 elements diff --git a/tests/scripts/release/make_notes.py b/tests/scripts/release/make_notes.py index abee4e85db67..d5fca9543afb 100644 --- a/tests/scripts/release/make_notes.py +++ b/tests/scripts/release/make_notes.py @@ -49,7 +49,8 @@ "wasm": "web", "runtime": "Runtime", "aot": "AOT", - "arith": "Arith", + "arith": "Symbolic Analysis", + "sym": "Symbolic Analysis", "byoc": "BYOC", "community": "Community", "tensorir": "TIR",