diff --git a/3rdparty/tvm-tirx b/3rdparty/tvm-tirx index 198a04573ea4..336af7322399 160000 --- a/3rdparty/tvm-tirx +++ b/3rdparty/tvm-tirx @@ -1 +1 @@ -Subproject commit 198a04573ea44e060c1087c8223cf082c64a6c03 +Subproject commit 336af73223991bd5586efe528bd59f888c0efc46 diff --git a/CMakeLists.txt b/CMakeLists.txt index fe2defca49ce..80d2d7e8263a 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -149,8 +149,8 @@ if(TVM_USE_INSTALLED_TIRX) "TVM_OPENCL_DLL=__attribute__((visibility(\"default\")))") endif() endif() - if(NOT TVM_TIRX_VERSION STREQUAL "0.1.0a2") - message(FATAL_ERROR "TVM requires apache-tvm-tirx 0.1.0a2") + if(NOT TVM_TIRX_VERSION STREQUAL "0.1.0a4") + message(FATAL_ERROR "TVM requires apache-tvm-tirx 0.1.0a4") endif() else() # Keep source-built core libraries beside the parent libraries. diff --git a/apps/android_rpc/app/src/main/jni/tvm_runtime.h b/apps/android_rpc/app/src/main/jni/tvm_runtime.h index 03c9ea98e2f3..8f36d167e9e8 100644 --- a/apps/android_rpc/app/src/main/jni/tvm_runtime.h +++ b/apps/android_rpc/app/src/main/jni/tvm_runtime.h @@ -64,7 +64,6 @@ #include "../3rdparty/tvm-tirx/src/runtime/tensor.cc" #include "../3rdparty/tvm-tirx/src/runtime/thread_pool.cc" #include "../3rdparty/tvm-tirx/src/runtime/threading_backend.cc" -#include "../3rdparty/tvm-tirx/src/runtime/workspace_pool.cc" // clang-format on #ifdef TVM_OPENCL_RUNTIME diff --git a/docs/arch/runtime.rst b/docs/arch/runtime.rst index 0ee894362c7a..fa535a4d3eb2 100644 --- a/docs/arch/runtime.rst +++ b/docs/arch/runtime.rst @@ -216,35 +216,35 @@ We can also define subclass ``ObjectRef`` to hold each subtypes of ``Object``. E RegisterReflection function. -Each ``Object`` subclass will override this to register its members. Here is an example implementation of IntImmNode. +Each ``Object`` subclass will override this to register its members. Here is an example implementation of IntImmObj. .. code:: c - class IntImmNode : public PrimExprNode { + class IntImmObj : public PrimExprNode { public: /*! \brief the Internal value. */ int64_t value; static void RegisterReflection() { namespace refl = tvm::ffi::reflection; - refl::ObjectDef().def_ro("value", &IntImmNode::value); + refl::ObjectDef().def_ro("value", &IntImmObj::value); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.IntImm", IntImmNode, PrimExprNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.IntImm", IntImmObj, PrimExprNode); }; // in cc file - TVM_FFI_STATIC_INIT_BLOCK() { IntImmNode::RegisterReflection(); } + TVM_FFI_STATIC_INIT_BLOCK() { IntImmObj::RegisterReflection(); } The RegisterReflection gives us a reflection API to register each member of the object. We can use this function to visit the node and serialize any language object recursively. It also allows us to get members of an object easily in front-end language. -For example, we can access the value field of the IntImmNode. +For example, we can access the value field of the IntImmObj. .. code:: python import tvm x = tvm_tirx.ir.IntImm("int32", 1) - # access the value field of IntImmNode + # access the value field of IntImmObj print(x.value) New ``Object`` can be added to C++ without changing the front-end runtime, making it easy to make extensions to the compiler stack. diff --git a/include/tvm/relax/analysis.h b/include/tvm/relax/analysis.h index 6273e1b967d2..26ed6c1ef4b9 100644 --- a/include/tvm/relax/analysis.h +++ b/include/tvm/relax/analysis.h @@ -544,7 +544,7 @@ VarUsageInfo CollectVarUsage(const Expr& expr); * * \return A set of variable nodes that are used in the expression */ -TVM_DLL std::set GetUsedVars(const Expr& expr); +TVM_DLL std::set GetUsedVars(const Expr& expr); /*! * \brief Remove unused statements inside DataflowBlocks. diff --git a/include/tvm/relax/binding_rewrite.h b/include/tvm/relax/binding_rewrite.h index d116478570c6..a038f307ca8a 100644 --- a/include/tvm/relax/binding_rewrite.h +++ b/include/tvm/relax/binding_rewrite.h @@ -82,7 +82,7 @@ class DataflowBlockRewriteObj : public ffi::Object { DataflowBlock dfb_; //!< The rewritten dataflow block. ffi::Optional root_fn_; //!< The rewritten function. - const FunctionNode* original_fn_ptr_; //!< Pointer to the original function. + const FunctionObj* original_fn_ptr_; //!< Pointer to the original function. ffi::Map> to_users_; //!< Map from variable to its users. ffi::Array fn_outputs_; //!< Variables required by function outputs. diff --git a/include/tvm/relax/dataflow_pattern.h b/include/tvm/relax/dataflow_pattern.h index 1c7c688403ef..b090e32ccd31 100644 --- a/include/tvm/relax/dataflow_pattern.h +++ b/include/tvm/relax/dataflow_pattern.h @@ -90,15 +90,15 @@ TVM_DLL PatternSeq operator>>(const PatternSeq& lhs, const PatternSeq& rhs); * \brief Base type of all dataflow patterns. * \sa DFPattern */ -class DFPatternNode : public ffi::Object { +class DFPatternObj : public ffi::Object { public: static constexpr const uint32_t _type_child_slots = 21; - TVM_FFI_DECLARE_OBJECT_INFO("relax.dpl.DFPattern", DFPatternNode, ffi::Object); + TVM_FFI_DECLARE_OBJECT_INFO("relax.dpl.DFPattern", DFPatternObj, ffi::Object); }; /*! * \brief Managed reference to dataflow patterns. - * \sa DFPatternNode + * \sa DFPatternObj */ class DFPattern : public ffi::ObjectRef { public: @@ -131,7 +131,7 @@ class DFPattern : public ffi::ObjectRef { /*! \brief Implicit conversion from DFPattern to PatternSeq */ TVM_DLL operator PatternSeq() const; - TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(DFPattern, ffi::ObjectRef, DFPatternNode); + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(DFPattern, ffi::ObjectRef, DFPatternObj); }; /*! \brief Constraint of a DFPattern edge (producer -> consumer) in graph-level matching */ @@ -163,7 +163,7 @@ struct PairCons { * example, constraining the two branches of an elementwise operation * to have the same shape. */ -class DFConstraintNode : public ffi::Object { +class DFConstraintObj : public ffi::Object { public: /*! \brief Return the patterns on which the constraint depends */ virtual ffi::Array GetDependentPatterns() const = 0; @@ -196,36 +196,36 @@ class DFConstraintNode : public ffi::Object { * sufficient for the constraint to be satisfied. */ virtual std::tuple AsCondition( - std::function(const DFPatternNode*)> match_state) const = 0; + std::function(const DFPatternObj*)> match_state) const = 0; static constexpr const uint32_t _type_child_slots = 1; - TVM_FFI_DECLARE_OBJECT_INFO("relax.dpl.DFConstraint", DFConstraintNode, ffi::Object); + TVM_FFI_DECLARE_OBJECT_INFO("relax.dpl.DFConstraint", DFConstraintObj, ffi::Object); }; class DFConstraint : public ffi::ObjectRef { public: - TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(DFConstraint, ffi::ObjectRef, DFConstraintNode); + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(DFConstraint, ffi::ObjectRef, DFConstraintObj); }; /*! * \brief A sequence of DFPatterns that the previous DFPattern is connected to the next one. * \sa PatternSeq */ -class PatternSeqNode final : public ffi::Object { +class PatternSeqObj final : public ffi::Object { public: tvm::ffi::Array patterns; /*!< The sequence of DFPatterns */ std::vector pair_constraints; /*!< Constraints between the previous and next patterns */ static void RegisterReflection() { namespace refl = tvm::ffi::reflection; - refl::ObjectDef().def_ro("patterns", &PatternSeqNode::patterns); + refl::ObjectDef().def_ro("patterns", &PatternSeqObj::patterns); } - TVM_FFI_DECLARE_OBJECT_INFO("relax.dpl.PatternSeq", PatternSeqNode, ffi::Object); + TVM_FFI_DECLARE_OBJECT_INFO("relax.dpl.PatternSeq", PatternSeqObj, ffi::Object); }; /*! * \brief Managed reference to pattern sequences. - * \sa PatternSeqNode + * \sa PatternSeqObj */ class PatternSeq final : public ffi::ObjectRef { public: @@ -242,7 +242,7 @@ class PatternSeq final : public ffi::ObjectRef { friend PatternSeq UsedBy(const PatternSeq& lhs, const PatternSeq& rhs, int index); friend PatternSeq OnlyUsedBy(const PatternSeq& lhs, const PatternSeq& rhs, int index); - TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(PatternSeq, ffi::ObjectRef, PatternSeqNode); + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(PatternSeq, ffi::ObjectRef, PatternSeqObj); }; /*! @@ -341,28 +341,28 @@ class PatternContext : public ffi::ObjectRef { * \brief Pattern for Relax Expression. * \sa ExprPattern */ -class ExprPatternNode : public DFPatternNode { +class ExprPatternObj : public DFPatternObj { public: - explicit ExprPatternNode(Expr expr) : expr(std::move(expr)) {} - explicit ExprPatternNode(ffi::UnsafeInit) : expr(ffi::UnsafeInit{}) {} + explicit ExprPatternObj(Expr expr) : expr(std::move(expr)) {} + explicit ExprPatternObj(ffi::UnsafeInit) : expr(ffi::UnsafeInit{}) {} Expr expr; /*!< The expression to match */ static void RegisterReflection() { namespace refl = tvm::ffi::reflection; - refl::ObjectDef().def_ro("expr", &ExprPatternNode::expr); + refl::ObjectDef().def_ro("expr", &ExprPatternObj::expr); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.ExprPattern", ExprPatternNode, DFPatternNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.ExprPattern", ExprPatternObj, DFPatternObj); }; /*! * \brief Managed reference to an ExprPattern. - * \sa ExprPatternNode + * \sa ExprPatternObj */ class ExprPattern : public DFPattern { public: TVM_DLL explicit ExprPattern(Expr expr); - TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ExprPattern, DFPattern, ExprPatternNode); + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ExprPattern, DFPattern, ExprPatternObj); }; /*! @@ -370,23 +370,23 @@ class ExprPattern : public DFPattern { * \note The name field matches any string if it is empty. * \sa VarPattern */ -class VarPatternNode : public DFPatternNode { +class VarPatternObj : public DFPatternObj { public: ffi::String name; const ffi::String& name_hint() const { return name; } static void RegisterReflection() { namespace refl = tvm::ffi::reflection; - refl::ObjectDef().def_ro("name", &VarPatternNode::name); + refl::ObjectDef().def_ro("name", &VarPatternObj::name); } static constexpr const uint32_t _type_child_slots = 1; - TVM_FFI_DECLARE_OBJECT_INFO("relax.dpl.VarPattern", VarPatternNode, DFPatternNode); + TVM_FFI_DECLARE_OBJECT_INFO("relax.dpl.VarPattern", VarPatternObj, DFPatternObj); }; /*! * \brief Managed reference to a VarPattern. - * \sa VarPatternNode + * \sa VarPatternObj */ class VarPattern : public DFPattern { public: @@ -396,90 +396,89 @@ class VarPattern : public DFPattern { * \param name_hint Variable name to match. Any if empty (""). */ TVM_DLL VarPattern(ffi::String name_hint); - TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(VarPattern, DFPattern, VarPatternNode); + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(VarPattern, DFPattern, VarPatternObj); }; /*! * \brief A Pattern to Match a Relax Dataflow Variable * \sa DataflowVarPattern */ -class DataflowVarPatternNode : public VarPatternNode { +class DataflowVarPatternObj : public VarPatternObj { public: static void RegisterReflection() { namespace refl = tvm::ffi::reflection; - refl::ObjectDef(); + refl::ObjectDef(); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.DataflowVarPattern", DataflowVarPatternNode, - VarPatternNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.DataflowVarPattern", DataflowVarPatternObj, + VarPatternObj); }; /*! * \brief Managed reference to a DataflowVarPattern. - * \sa DataflowVarPatternNode + * \sa DataflowVarPatternObj */ class DataflowVarPattern : public DFPattern { public: /*! \sa VarPattern::VarPattern */ TVM_DLL DataflowVarPattern(ffi::String name_hint); TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(DataflowVarPattern, DFPattern, - DataflowVarPatternNode); + DataflowVarPatternObj); }; /*! * \brief A Pattern to Match a Relax Global Variable * \sa GlobalVarPattern */ -class GlobalVarPatternNode : public VarPatternNode { +class GlobalVarPatternObj : public VarPatternObj { public: - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.GlobalVarPattern", GlobalVarPatternNode, - DFPatternNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.GlobalVarPattern", GlobalVarPatternObj, + DFPatternObj); }; /*! * \brief Managed reference to a GlobalVarPattern. - * \sa GlobalVarPatternNode + * \sa GlobalVarPatternObj */ class GlobalVarPattern : public DFPattern { public: TVM_DLL GlobalVarPattern(ffi::String name_hint); - TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(GlobalVarPattern, DFPattern, GlobalVarPatternNode); + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(GlobalVarPattern, DFPattern, GlobalVarPatternObj); }; /*! * \brief A Pattern to Match a Relax GenericConst. * \sa ConstantPattern */ -class ConstantPatternNode : public DFPatternNode { +class ConstantPatternObj : public DFPatternObj { public: static void RegisterReflection() { namespace refl = tvm::ffi::reflection; - refl::ObjectDef(); + refl::ObjectDef(); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.ConstantPattern", ConstantPatternNode, - DFPatternNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.ConstantPattern", ConstantPatternObj, DFPatternObj); }; /*! * \brief Managed reference to a ConstantPattern. - * \sa ConstantPatternNode + * \sa ConstantPatternObj */ class ConstantPattern : public DFPattern { public: - explicit ConstantPattern(ffi::UnsafeInit tag, ffi::ObjectPtr n) + explicit ConstantPattern(ffi::UnsafeInit tag, ffi::ObjectPtr n) : DFPattern(tag) { data_ = std::move(n); } - TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ConstantPattern, DFPattern, ConstantPatternNode); + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ConstantPattern, DFPattern, ConstantPatternObj); }; /*! * \brief A pattern to match a callable node in Relax. * \sa CallPattern */ -class CallPatternNode : public DFPatternNode { +class CallPatternObj : public DFPatternObj { public: - explicit CallPatternNode(DFPattern op) : op(std::move(op)) {} - explicit CallPatternNode(ffi::UnsafeInit) : op(ffi::UnsafeInit{}) {} + explicit CallPatternObj(DFPattern op) : op(std::move(op)) {} + explicit CallPatternObj(ffi::UnsafeInit) : op(ffi::UnsafeInit{}) {} /*! * \note The op field can be: @@ -500,17 +499,17 @@ class CallPatternNode : public DFPatternNode { static void RegisterReflection() { namespace refl = tvm::ffi::reflection; - refl::ObjectDef() - .def_ro("op", &CallPatternNode::op) - .def_ro("args", &CallPatternNode::args); + refl::ObjectDef() + .def_ro("op", &CallPatternObj::op) + .def_ro("args", &CallPatternObj::args); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.CallPattern", CallPatternNode, DFPatternNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.CallPattern", CallPatternObj, DFPatternObj); }; class CallPattern : public DFPattern { public: TVM_DLL CallPattern(DFPattern op, ffi::Array args, bool varg_default_wildcard = false); - TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(CallPattern, DFPattern, CallPatternNode); + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(CallPattern, DFPattern, CallPatternObj); }; /*! @@ -518,25 +517,25 @@ class CallPattern : public DFPattern { * \sa PrimArrPattern * \note This is often used to match shapes specified as arguments to a function. */ -class PrimArrPatternNode : public DFPatternNode { +class PrimArrPatternObj : public DFPatternObj { public: ffi::Array fields; /*!< The array to match */ static void RegisterReflection() { namespace refl = tvm::ffi::reflection; - refl::ObjectDef().def_ro("fields", &PrimArrPatternNode::fields); + refl::ObjectDef().def_ro("fields", &PrimArrPatternObj::fields); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.PrimArrPattern", PrimArrPatternNode, DFPatternNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.PrimArrPattern", PrimArrPatternObj, DFPatternObj); }; /*! * \brief Managed reference to a PrimArrPattern. - * \sa PrimArrPatternNode + * \sa PrimArrPatternObj */ class PrimArrPattern : public DFPattern { public: TVM_DLL PrimArrPattern(ffi::Array arr); - TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(PrimArrPattern, DFPattern, PrimArrPatternNode); + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(PrimArrPattern, DFPattern, PrimArrPatternObj); }; /*! @@ -544,33 +543,32 @@ class PrimArrPattern : public DFPattern { * \sa Function * \sa FunctionPattern */ -class FunctionPatternNode : public DFPatternNode { +class FunctionPatternObj : public DFPatternObj { public: - explicit FunctionPatternNode(DFPattern body) : body(std::move(body)) {} - explicit FunctionPatternNode(ffi::UnsafeInit) : body(ffi::UnsafeInit{}) {} + explicit FunctionPatternObj(DFPattern body) : body(std::move(body)) {} + explicit FunctionPatternObj(ffi::UnsafeInit) : body(ffi::UnsafeInit{}) {} tvm::ffi::Array params; /*!< The parameters of the function */ /*! * \note Note that in Relax, the function body is a SeqExpr which contains - * 1) SeqExprNode::blocks, which is a list of blocks of statements; and 2) - * SeqExprNode::body, which is an Expr that can be anything. FunctionPattern + * 1) SeqExprObj::blocks, which is a list of blocks of statements; and 2) + * SeqExprObj::body, which is an Expr that can be anything. FunctionPattern * only matches the body of the function (writing patterns to statements is tricky). */ DFPattern body; /*!< The body of the function */ static void RegisterReflection() { namespace refl = tvm::ffi::reflection; - refl::ObjectDef() - .def_ro("params", &FunctionPatternNode::params) - .def_ro("body", &FunctionPatternNode::body); + refl::ObjectDef() + .def_ro("params", &FunctionPatternObj::params) + .def_ro("body", &FunctionPatternObj::body); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.FunctionPattern", FunctionPatternNode, - DFPatternNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.FunctionPattern", FunctionPatternObj, DFPatternObj); }; /*! - * \brief Managed reference to FunctionPatternNode. - * \sa FunctionPatternNode + * \brief Managed reference to FunctionPatternObj. + * \sa FunctionPatternObj */ class FunctionPattern : public DFPattern { public: @@ -581,60 +579,59 @@ class FunctionPattern : public DFPattern { */ TVM_DLL FunctionPattern(tvm::ffi::Array params, DFPattern body); - TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(FunctionPattern, DFPattern, FunctionPatternNode); + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(FunctionPattern, DFPattern, FunctionPatternObj); }; /*! * \brief Pattern to match a tuple of ordered expressions. * \sa TuplePattern */ -class TuplePatternNode : public DFPatternNode { +class TuplePatternObj : public DFPatternObj { public: tvm::ffi::Array fields; /*!< The fields of the tuple */ static void RegisterReflection() { namespace refl = tvm::ffi::reflection; - refl::ObjectDef().def_ro("fields", &TuplePatternNode::fields); + refl::ObjectDef().def_ro("fields", &TuplePatternObj::fields); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.TuplePattern", TuplePatternNode, DFPatternNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.TuplePattern", TuplePatternObj, DFPatternObj); }; /*! - * \brief Managed reference to TuplePatternNode. - * \sa TuplePatternNode + * \brief Managed reference to TuplePatternObj. + * \sa TuplePatternObj */ class TuplePattern : public DFPattern { public: TVM_DLL explicit TuplePattern(tvm::ffi::Array fields); - TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(TuplePattern, DFPattern, TuplePatternNode); + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(TuplePattern, DFPattern, TuplePatternObj); }; /*! * \brief A pattern to match multiple expressions unorderedly. * \sa UnorderedTuplePattern */ -class UnorderedTuplePatternNode : public DFPatternNode { +class UnorderedTuplePatternObj : public DFPatternObj { public: tvm::ffi::Array fields; /*!< The fields of the tuple */ static void RegisterReflection() { namespace refl = tvm::ffi::reflection; - refl::ObjectDef().def_ro("fields", - &UnorderedTuplePatternNode::fields); + refl::ObjectDef().def_ro("fields", &UnorderedTuplePatternObj::fields); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.UnorderedTuplePattern", UnorderedTuplePatternNode, - DFPatternNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.UnorderedTuplePattern", UnorderedTuplePatternObj, + DFPatternObj); }; /*! - * \brief Managed reference to UnorderedTuplePatternNode. - * \sa UnorderedTuplePatternNode + * \brief Managed reference to UnorderedTuplePatternObj. + * \sa UnorderedTuplePatternObj */ class UnorderedTuplePattern : public DFPattern { public: TVM_DLL explicit UnorderedTuplePattern(tvm::ffi::Array fields); TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(UnorderedTuplePattern, DFPattern, - UnorderedTuplePatternNode); + UnorderedTuplePatternObj); }; /*! @@ -642,150 +639,148 @@ class UnorderedTuplePattern : public DFPattern { * \sa TupleGetItem * \sa TupleGetItemPattern */ -class TupleGetItemPatternNode : public DFPatternNode { +class TupleGetItemPatternObj : public DFPatternObj { public: - explicit TupleGetItemPatternNode(DFPattern tuple) : tuple(std::move(tuple)) {} - explicit TupleGetItemPatternNode(ffi::UnsafeInit) : tuple(ffi::UnsafeInit{}) {} + explicit TupleGetItemPatternObj(DFPattern tuple) : tuple(std::move(tuple)) {} + explicit TupleGetItemPatternObj(ffi::UnsafeInit) : tuple(ffi::UnsafeInit{}) {} DFPattern tuple; /*!< The tuple Expression */ int index; /*!< The index of the tuple with -1 meaning arbitrary */ static void RegisterReflection() { namespace refl = tvm::ffi::reflection; - refl::ObjectDef() - .def_ro("tuple", &TupleGetItemPatternNode::tuple) - .def_ro("index", &TupleGetItemPatternNode::index); + refl::ObjectDef() + .def_ro("tuple", &TupleGetItemPatternObj::tuple) + .def_ro("index", &TupleGetItemPatternObj::index); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.TupleGetItemPattern", TupleGetItemPatternNode, - DFPatternNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.TupleGetItemPattern", TupleGetItemPatternObj, + DFPatternObj); }; /*! - * \brief Managed reference to TupleGetItemPatternNode. - * \sa TupleGetItemPatternNode + * \brief Managed reference to TupleGetItemPatternObj. + * \sa TupleGetItemPatternObj */ class TupleGetItemPattern : public DFPattern { public: TVM_DLL TupleGetItemPattern(DFPattern tuple, int index); TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(TupleGetItemPattern, DFPattern, - TupleGetItemPatternNode); + TupleGetItemPatternObj); }; /*! * \brief Match a conjunction of other patterns. * \sa AndPattern */ -class AndPatternNode : public DFPatternNode { +class AndPatternObj : public DFPatternObj { public: - explicit AndPatternNode(DFPattern left, DFPattern right) + explicit AndPatternObj(DFPattern left, DFPattern right) : left(std::move(left)), right(std::move(right)) {} - explicit AndPatternNode(ffi::UnsafeInit) : left(ffi::UnsafeInit{}), right(ffi::UnsafeInit{}) {} + explicit AndPatternObj(ffi::UnsafeInit) : left(ffi::UnsafeInit{}), right(ffi::UnsafeInit{}) {} DFPattern left; /*!< The left hand side of the conjunction */ DFPattern right; /*!< The right hand side of the conjunction */ static void RegisterReflection() { namespace refl = tvm::ffi::reflection; - refl::ObjectDef() - .def_ro("left", &AndPatternNode::left) - .def_ro("right", &AndPatternNode::right); + refl::ObjectDef() + .def_ro("left", &AndPatternObj::left) + .def_ro("right", &AndPatternObj::right); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.AndPattern", AndPatternNode, DFPatternNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.AndPattern", AndPatternObj, DFPatternObj); }; /*! - * \brief Managed reference to AndPatternNode. - * \sa AndPatternNode + * \brief Managed reference to AndPatternObj. + * \sa AndPatternObj */ class AndPattern : public DFPattern { public: TVM_DLL AndPattern(DFPattern lhs, DFPattern rhs); - TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(AndPattern, DFPattern, AndPatternNode); + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(AndPattern, DFPattern, AndPatternObj); }; /*! * \brief Match a disjunction of other patterns. * \sa OrPattern */ -class OrPatternNode : public DFPatternNode { +class OrPatternObj : public DFPatternObj { public: - explicit OrPatternNode(DFPattern left, DFPattern right) + explicit OrPatternObj(DFPattern left, DFPattern right) : left(std::move(left)), right(std::move(right)) {} - explicit OrPatternNode(ffi::UnsafeInit) : left(ffi::UnsafeInit{}), right(ffi::UnsafeInit{}) {} + explicit OrPatternObj(ffi::UnsafeInit) : left(ffi::UnsafeInit{}), right(ffi::UnsafeInit{}) {} DFPattern left; /*!< The left hand side of the disjunction */ DFPattern right; /*!< The right hand side of the disjunction */ static void RegisterReflection() { namespace refl = tvm::ffi::reflection; - refl::ObjectDef() - .def_ro("left", &OrPatternNode::left) - .def_ro("right", &OrPatternNode::right); + refl::ObjectDef() + .def_ro("left", &OrPatternObj::left) + .def_ro("right", &OrPatternObj::right); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.OrPattern", OrPatternNode, DFPatternNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.OrPattern", OrPatternObj, DFPatternObj); }; /*! - * \brief Managed reference to OrPatternNode. - * \sa OrPatternNode + * \brief Managed reference to OrPatternObj. + * \sa OrPatternObj */ class OrPattern : public DFPattern { public: TVM_DLL OrPattern(DFPattern left, DFPattern right); - TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(OrPattern, DFPattern, OrPatternNode); + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(OrPattern, DFPattern, OrPatternObj); }; /*! * \brief Pattern for rejecting a certain pattern. * \sa NotPattern */ -class NotPatternNode : public DFPatternNode { +class NotPatternObj : public DFPatternObj { public: - explicit NotPatternNode(DFPattern reject) : reject(std::move(reject)) {} - explicit NotPatternNode(ffi::UnsafeInit) : reject(ffi::UnsafeInit{}) {} + explicit NotPatternObj(DFPattern reject) : reject(std::move(reject)) {} + explicit NotPatternObj(ffi::UnsafeInit) : reject(ffi::UnsafeInit{}) {} DFPattern reject; /*!< The pattern to reject */ static void RegisterReflection() { namespace refl = tvm::ffi::reflection; - refl::ObjectDef().def_ro("reject", &NotPatternNode::reject); + refl::ObjectDef().def_ro("reject", &NotPatternObj::reject); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.NotPattern", NotPatternNode, DFPatternNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.NotPattern", NotPatternObj, DFPatternObj); }; /*! - * \brief Managed reference to NotPatternNode. - * \sa NotPatternNode + * \brief Managed reference to NotPatternObj. + * \sa NotPatternObj */ class NotPattern : public DFPattern { public: TVM_DLL NotPattern(DFPattern reject); - TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(NotPattern, DFPattern, NotPatternNode); + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(NotPattern, DFPattern, NotPatternObj); }; /*! * \brief Wildcard Pattern is a pattern that can match anything. * \sa WildcardPattern */ -class WildcardPatternNode : public DFPatternNode { +class WildcardPatternObj : public DFPatternObj { public: static void RegisterReflection() { namespace refl = tvm::ffi::reflection; - refl::ObjectDef(); + refl::ObjectDef(); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.WildcardPattern", WildcardPatternNode, - DFPatternNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.WildcardPattern", WildcardPatternObj, DFPatternObj); }; /*! - * \brief Managed reference to WildcardPatternNode. - * \sa WildcardPatternNode + * \brief Managed reference to WildcardPatternObj. + * \sa WildcardPatternObj */ class WildcardPattern : public DFPattern { public: WildcardPattern(); - explicit WildcardPattern(ffi::ObjectPtr data) - : DFPattern(ffi::UnsafeInit{}) { + explicit WildcardPattern(ffi::ObjectPtr data) : DFPattern(ffi::UnsafeInit{}) { TVM_FFI_ICHECK(data != nullptr); data_ = std::move(data); } @@ -795,86 +790,86 @@ class WildcardPattern : public DFPattern { // nullptr`. This allows a zero-parameter constructor to be // declared here, to create a valid wildcard instance. - TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(WildcardPattern, DFPattern, WildcardPatternNode); + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(WildcardPattern, DFPattern, WildcardPatternObj); }; /*! * \brief Pattern for matching a certain type. * \sa TypePattern */ -class TypePatternNode : public DFPatternNode { +class TypePatternObj : public DFPatternObj { public: - explicit TypePatternNode(DFPattern pattern) : pattern(std::move(pattern)) {} - explicit TypePatternNode(ffi::UnsafeInit) : pattern(ffi::UnsafeInit{}) {} + explicit TypePatternObj(DFPattern pattern) : pattern(std::move(pattern)) {} + explicit TypePatternObj(ffi::UnsafeInit) : pattern(ffi::UnsafeInit{}) {} DFPattern pattern; /*!< The pattern to match */ Type ty = Type::Missing(); /*!< The type to match */ static void RegisterReflection() { namespace refl = tvm::ffi::reflection; - refl::ObjectDef() - .def_ro("pattern", &TypePatternNode::pattern) - .def_ro("ty", &TypePatternNode::ty); + refl::ObjectDef() + .def_ro("pattern", &TypePatternObj::pattern) + .def_ro("ty", &TypePatternObj::ty); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.TypePattern", TypePatternNode, DFPatternNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.TypePattern", TypePatternObj, DFPatternObj); }; class TypePattern : public DFPattern { public: TVM_DLL TypePattern(DFPattern pattern, Type ty); - TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(TypePattern, DFPattern, TypePatternNode); + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(TypePattern, DFPattern, TypePatternObj); }; /*! * \brief A pattern that asserting a root pattern has a certain shape. * \sa ShapePattern */ -class ShapePatternNode : public DFPatternNode { +class ShapePatternObj : public DFPatternObj { public: - explicit ShapePatternNode(DFPattern pattern) : pattern(std::move(pattern)) {} - explicit ShapePatternNode(ffi::UnsafeInit) : pattern(ffi::UnsafeInit{}) {} + explicit ShapePatternObj(DFPattern pattern) : pattern(std::move(pattern)) {} + explicit ShapePatternObj(ffi::UnsafeInit) : pattern(ffi::UnsafeInit{}) {} DFPattern pattern; /*!< The root pattern to match */ ffi::Array shape; /*!< The shape to match */ static void RegisterReflection() { namespace refl = tvm::ffi::reflection; - refl::ObjectDef() - .def_ro("pattern", &ShapePatternNode::pattern) - .def_ro("shape", &ShapePatternNode::shape); + refl::ObjectDef() + .def_ro("pattern", &ShapePatternObj::pattern) + .def_ro("shape", &ShapePatternObj::shape); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.ShapePattern", ShapePatternNode, DFPatternNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.ShapePattern", ShapePatternObj, DFPatternObj); }; /*! - * \brief Managed reference to ShapePatternNode. - * \sa ShapePatternNode + * \brief Managed reference to ShapePatternObj. + * \sa ShapePatternObj */ class ShapePattern : public DFPattern { public: TVM_DLL ShapePattern(DFPattern pattern, ffi::Array type); - TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ShapePattern, DFPattern, ShapePatternNode); + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ShapePattern, DFPattern, ShapePatternObj); }; /*! * \brief A pattern that asserting multiple root patterns have the same shape * \sa SameShapePattern */ -class SameShapeConstraintNode : public DFConstraintNode { +class SameShapeConstraintObj : public DFConstraintObj { public: ffi::Array args; /*!< The patterns with matching shapes */ ffi::Array GetDependentPatterns() const override { return args; } std::tuple AsCondition( - std::function(const DFPatternNode*)> match_state) const override; + std::function(const DFPatternObj*)> match_state) const override; static void RegisterReflection() { namespace refl = tvm::ffi::reflection; - refl::ObjectDef().def_ro("args", &SameShapeConstraintNode::args); + refl::ObjectDef().def_ro("args", &SameShapeConstraintObj::args); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.SameShapeConstraint", SameShapeConstraintNode, - DFConstraintNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.SameShapeConstraint", SameShapeConstraintObj, + DFConstraintObj); }; /*! @@ -885,70 +880,69 @@ class SameShapeConstraint : public DFConstraint { public: TVM_DLL SameShapeConstraint(ffi::Array args); TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(SameShapeConstraint, DFConstraint, - SameShapeConstraintNode); + SameShapeConstraintObj); }; /*! * \brief A pattern that asserting a root pattern has a certain data type. * \sa DataTypePattern */ -class DataTypePatternNode : public DFPatternNode { +class DataTypePatternObj : public DFPatternObj { public: - explicit DataTypePatternNode(DFPattern pattern) : pattern(std::move(pattern)) {} - explicit DataTypePatternNode(ffi::UnsafeInit) : pattern(ffi::UnsafeInit{}) {} + explicit DataTypePatternObj(DFPattern pattern) : pattern(std::move(pattern)) {} + explicit DataTypePatternObj(ffi::UnsafeInit) : pattern(ffi::UnsafeInit{}) {} DFPattern pattern; /*!< The root pattern to match */ DLDataType dtype; /*!< The data type to match */ static void RegisterReflection() { namespace refl = tvm::ffi::reflection; - refl::ObjectDef() - .def_ro("pattern", &DataTypePatternNode::pattern) - .def_ro("dtype", &DataTypePatternNode::dtype); + refl::ObjectDef() + .def_ro("pattern", &DataTypePatternObj::pattern) + .def_ro("dtype", &DataTypePatternObj::dtype); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.DataTypePattern", DataTypePatternNode, - DFPatternNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.DataTypePattern", DataTypePatternObj, DFPatternObj); }; /*! - * \brief Managed reference to DataTypePatternNode. - * \sa DataTypePatternNode + * \brief Managed reference to DataTypePatternObj. + * \sa DataTypePatternObj */ class DataTypePattern : public DFPattern { public: TVM_DLL DataTypePattern(DFPattern pattern, DLDataType dtype); - TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(DataTypePattern, DFPattern, DataTypePatternNode); + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(DataTypePattern, DFPattern, DataTypePatternObj); }; /*! * \brief A pattern that asserting a root pattern has certain attributes. * \sa AttrPattern */ -class AttrPatternNode : public DFPatternNode { +class AttrPatternObj : public DFPatternObj { public: - explicit AttrPatternNode(DFPattern pattern) : pattern(std::move(pattern)) {} - explicit AttrPatternNode(ffi::UnsafeInit) : pattern(ffi::UnsafeInit{}) {} + explicit AttrPatternObj(DFPattern pattern) : pattern(std::move(pattern)) {} + explicit AttrPatternObj(ffi::UnsafeInit) : pattern(ffi::UnsafeInit{}) {} DFPattern pattern; /*!< The root pattern to match */ DictAttrs attrs; /*!< The attributes (a map/dictionary) to match */ static void RegisterReflection() { namespace refl = tvm::ffi::reflection; - refl::ObjectDef() - .def_ro("pattern", &AttrPatternNode::pattern) - .def_ro("attrs", &AttrPatternNode::attrs); + refl::ObjectDef() + .def_ro("pattern", &AttrPatternObj::pattern) + .def_ro("attrs", &AttrPatternObj::attrs); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.AttrPattern", AttrPatternNode, DFPatternNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.AttrPattern", AttrPatternObj, DFPatternObj); }; /*! - * \brief Managed reference to AttrPatternNode. - * \sa AttrPatternNode + * \brief Managed reference to AttrPatternObj. + * \sa AttrPatternObj */ class AttrPattern : public DFPattern { public: TVM_DLL AttrPattern(DFPattern pattern, DictAttrs attrs); - TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(AttrPattern, DFPattern, AttrPatternNode); + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(AttrPattern, DFPattern, AttrPatternObj); }; /*! @@ -956,7 +950,7 @@ class AttrPattern : public DFPattern { * \sa ExternFunc * \sa ExternFuncPattern */ -class ExternFuncPatternNode : public DFPatternNode { +class ExternFuncPatternObj : public DFPatternObj { public: ffi::String global_symbol_; /*!< The global symbol name of the external function */ @@ -965,22 +959,21 @@ class ExternFuncPatternNode : public DFPatternNode { static void RegisterReflection() { namespace refl = tvm::ffi::reflection; - refl::ObjectDef().def_ro("global_symbol", - &ExternFuncPatternNode::global_symbol_); + refl::ObjectDef().def_ro("global_symbol", + &ExternFuncPatternObj::global_symbol_); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.ExternFuncPattern", ExternFuncPatternNode, - DFPatternNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.ExternFuncPattern", ExternFuncPatternObj, + DFPatternObj); }; /*! - * \brief Managed reference to ExternFuncPatternNode. - * \sa ExternFuncPatternNode + * \brief Managed reference to ExternFuncPatternObj. + * \sa ExternFuncPatternObj */ class ExternFuncPattern : public DFPattern { public: TVM_DLL ExternFuncPattern(ffi::String global_symbol); - TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ExternFuncPattern, DFPattern, - ExternFuncPatternNode); + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ExternFuncPattern, DFPattern, ExternFuncPatternObj); }; /*! \brief Syntatic Sugar for creating a VarPattern with a name */ diff --git a/include/tvm/relax/dataflow_pattern_functor.h b/include/tvm/relax/dataflow_pattern_functor.h index a97984c4d5b9..ddd02c4e2b3c 100644 --- a/include/tvm/relax/dataflow_pattern_functor.h +++ b/include/tvm/relax/dataflow_pattern_functor.h @@ -83,30 +83,29 @@ class DFPatternFunctor { return vtable(n, this, std::forward(args)...); } // Functions that can be overriden by subclass - virtual R VisitDFPattern_(const OrPatternNode* op, Args... args) DFPATTERN_FUNCTOR_DEFAULT; - virtual R VisitDFPattern_(const AndPatternNode* op, Args... args) DFPATTERN_FUNCTOR_DEFAULT; - virtual R VisitDFPattern_(const NotPatternNode* op, Args... args) DFPATTERN_FUNCTOR_DEFAULT; - virtual R VisitDFPattern_(const AttrPatternNode* op, Args... args) DFPATTERN_FUNCTOR_DEFAULT; - virtual R VisitDFPattern_(const CallPatternNode* op, Args... args) DFPATTERN_FUNCTOR_DEFAULT; - virtual R VisitDFPattern_(const ConstantPatternNode* op, Args... args) DFPATTERN_FUNCTOR_DEFAULT; - virtual R VisitDFPattern_(const DataTypePatternNode* op, Args... args) DFPATTERN_FUNCTOR_DEFAULT; - virtual R VisitDFPattern_(const ExprPatternNode* op, Args... args) DFPATTERN_FUNCTOR_DEFAULT; - virtual R VisitDFPattern_(const FunctionPatternNode* op, Args... args) DFPATTERN_FUNCTOR_DEFAULT; - virtual R VisitDFPattern_(const ShapePatternNode* op, Args... args) DFPATTERN_FUNCTOR_DEFAULT; - virtual R VisitDFPattern_(const TupleGetItemPatternNode* op, + virtual R VisitDFPattern_(const OrPatternObj* op, Args... args) DFPATTERN_FUNCTOR_DEFAULT; + virtual R VisitDFPattern_(const AndPatternObj* op, Args... args) DFPATTERN_FUNCTOR_DEFAULT; + virtual R VisitDFPattern_(const NotPatternObj* op, Args... args) DFPATTERN_FUNCTOR_DEFAULT; + virtual R VisitDFPattern_(const AttrPatternObj* op, Args... args) DFPATTERN_FUNCTOR_DEFAULT; + virtual R VisitDFPattern_(const CallPatternObj* op, Args... args) DFPATTERN_FUNCTOR_DEFAULT; + virtual R VisitDFPattern_(const ConstantPatternObj* op, Args... args) DFPATTERN_FUNCTOR_DEFAULT; + virtual R VisitDFPattern_(const DataTypePatternObj* op, Args... args) DFPATTERN_FUNCTOR_DEFAULT; + virtual R VisitDFPattern_(const ExprPatternObj* op, Args... args) DFPATTERN_FUNCTOR_DEFAULT; + virtual R VisitDFPattern_(const FunctionPatternObj* op, Args... args) DFPATTERN_FUNCTOR_DEFAULT; + virtual R VisitDFPattern_(const ShapePatternObj* op, Args... args) DFPATTERN_FUNCTOR_DEFAULT; + virtual R VisitDFPattern_(const TupleGetItemPatternObj* op, Args... args) DFPATTERN_FUNCTOR_DEFAULT; - virtual R VisitDFPattern_(const TuplePatternNode* op, Args... args) DFPATTERN_FUNCTOR_DEFAULT; - virtual R VisitDFPattern_(const TypePatternNode* op, Args... args) DFPATTERN_FUNCTOR_DEFAULT; - virtual R VisitDFPattern_(const WildcardPatternNode* op, Args... args) DFPATTERN_FUNCTOR_DEFAULT; - virtual R VisitDFPattern_(const VarPatternNode* op, Args... args) DFPATTERN_FUNCTOR_DEFAULT; + virtual R VisitDFPattern_(const TuplePatternObj* op, Args... args) DFPATTERN_FUNCTOR_DEFAULT; + virtual R VisitDFPattern_(const TypePatternObj* op, Args... args) DFPATTERN_FUNCTOR_DEFAULT; + virtual R VisitDFPattern_(const WildcardPatternObj* op, Args... args) DFPATTERN_FUNCTOR_DEFAULT; + virtual R VisitDFPattern_(const VarPatternObj* op, Args... args) DFPATTERN_FUNCTOR_DEFAULT; - virtual R VisitDFPattern_(const DataflowVarPatternNode* op, + virtual R VisitDFPattern_(const DataflowVarPatternObj* op, Args... args) DFPATTERN_FUNCTOR_DEFAULT; - virtual R VisitDFPattern_(const GlobalVarPatternNode* op, Args... args) DFPATTERN_FUNCTOR_DEFAULT; - virtual R VisitDFPattern_(const ExternFuncPatternNode* op, - Args... args) DFPATTERN_FUNCTOR_DEFAULT; - virtual R VisitDFPattern_(const PrimArrPatternNode* op, Args... args) DFPATTERN_FUNCTOR_DEFAULT; - virtual R VisitDFPattern_(const UnorderedTuplePatternNode* op, + virtual R VisitDFPattern_(const GlobalVarPatternObj* op, Args... args) DFPATTERN_FUNCTOR_DEFAULT; + virtual R VisitDFPattern_(const ExternFuncPatternObj* op, Args... args) DFPATTERN_FUNCTOR_DEFAULT; + virtual R VisitDFPattern_(const PrimArrPatternObj* op, Args... args) DFPATTERN_FUNCTOR_DEFAULT; + virtual R VisitDFPattern_(const UnorderedTuplePatternObj* op, Args... args) DFPATTERN_FUNCTOR_DEFAULT; virtual R VisitDFPatternDefault_(const ffi::Object* op, Args...) { @@ -119,26 +118,26 @@ class DFPatternFunctor { static FType InitVTable() { FType vtable; // Set dispatch - RELAX_DFPATTERN_FUNCTOR_DISPATCH(OrPatternNode); - RELAX_DFPATTERN_FUNCTOR_DISPATCH(AndPatternNode); - RELAX_DFPATTERN_FUNCTOR_DISPATCH(NotPatternNode); - RELAX_DFPATTERN_FUNCTOR_DISPATCH(AttrPatternNode); - RELAX_DFPATTERN_FUNCTOR_DISPATCH(CallPatternNode); - RELAX_DFPATTERN_FUNCTOR_DISPATCH(ConstantPatternNode); - RELAX_DFPATTERN_FUNCTOR_DISPATCH(DataTypePatternNode); - RELAX_DFPATTERN_FUNCTOR_DISPATCH(ExprPatternNode); - RELAX_DFPATTERN_FUNCTOR_DISPATCH(FunctionPatternNode); - RELAX_DFPATTERN_FUNCTOR_DISPATCH(ShapePatternNode); - RELAX_DFPATTERN_FUNCTOR_DISPATCH(TupleGetItemPatternNode); - RELAX_DFPATTERN_FUNCTOR_DISPATCH(TuplePatternNode); - RELAX_DFPATTERN_FUNCTOR_DISPATCH(TypePatternNode); - RELAX_DFPATTERN_FUNCTOR_DISPATCH(WildcardPatternNode); - RELAX_DFPATTERN_FUNCTOR_DISPATCH(VarPatternNode); - RELAX_DFPATTERN_FUNCTOR_DISPATCH(DataflowVarPatternNode); - RELAX_DFPATTERN_FUNCTOR_DISPATCH(GlobalVarPatternNode); - RELAX_DFPATTERN_FUNCTOR_DISPATCH(ExternFuncPatternNode); - RELAX_DFPATTERN_FUNCTOR_DISPATCH(PrimArrPatternNode); - RELAX_DFPATTERN_FUNCTOR_DISPATCH(UnorderedTuplePatternNode); + RELAX_DFPATTERN_FUNCTOR_DISPATCH(OrPatternObj); + RELAX_DFPATTERN_FUNCTOR_DISPATCH(AndPatternObj); + RELAX_DFPATTERN_FUNCTOR_DISPATCH(NotPatternObj); + RELAX_DFPATTERN_FUNCTOR_DISPATCH(AttrPatternObj); + RELAX_DFPATTERN_FUNCTOR_DISPATCH(CallPatternObj); + RELAX_DFPATTERN_FUNCTOR_DISPATCH(ConstantPatternObj); + RELAX_DFPATTERN_FUNCTOR_DISPATCH(DataTypePatternObj); + RELAX_DFPATTERN_FUNCTOR_DISPATCH(ExprPatternObj); + RELAX_DFPATTERN_FUNCTOR_DISPATCH(FunctionPatternObj); + RELAX_DFPATTERN_FUNCTOR_DISPATCH(ShapePatternObj); + RELAX_DFPATTERN_FUNCTOR_DISPATCH(TupleGetItemPatternObj); + RELAX_DFPATTERN_FUNCTOR_DISPATCH(TuplePatternObj); + RELAX_DFPATTERN_FUNCTOR_DISPATCH(TypePatternObj); + RELAX_DFPATTERN_FUNCTOR_DISPATCH(WildcardPatternObj); + RELAX_DFPATTERN_FUNCTOR_DISPATCH(VarPatternObj); + RELAX_DFPATTERN_FUNCTOR_DISPATCH(DataflowVarPatternObj); + RELAX_DFPATTERN_FUNCTOR_DISPATCH(GlobalVarPatternObj); + RELAX_DFPATTERN_FUNCTOR_DISPATCH(ExternFuncPatternObj); + RELAX_DFPATTERN_FUNCTOR_DISPATCH(PrimArrPatternObj); + RELAX_DFPATTERN_FUNCTOR_DISPATCH(UnorderedTuplePatternObj); vtable.Finalize(); return vtable; } @@ -153,27 +152,27 @@ class DFPatternFunctor { class DFPatternVisitor : public DFPatternFunctor { public: void VisitDFPattern(const DFPattern& pattern) override; - void VisitDFPattern_(const OrPatternNode* op) override; - void VisitDFPattern_(const AndPatternNode* op) override; - void VisitDFPattern_(const NotPatternNode* op) override; - void VisitDFPattern_(const AttrPatternNode* op) override; - void VisitDFPattern_(const CallPatternNode* op) override; - void VisitDFPattern_(const ConstantPatternNode* op) override; - void VisitDFPattern_(const DataTypePatternNode* op) override; - void VisitDFPattern_(const ExprPatternNode* op) override; - void VisitDFPattern_(const FunctionPatternNode* op) override; - void VisitDFPattern_(const ShapePatternNode* op) override; - void VisitDFPattern_(const TupleGetItemPatternNode* op) override; - void VisitDFPattern_(const TuplePatternNode* op) override; - void VisitDFPattern_(const TypePatternNode* op) override; - void VisitDFPattern_(const WildcardPatternNode* op) override; - void VisitDFPattern_(const VarPatternNode* op) override; - - void VisitDFPattern_(const DataflowVarPatternNode* op) override; - void VisitDFPattern_(const GlobalVarPatternNode* op) override; - void VisitDFPattern_(const ExternFuncPatternNode* op) override; - void VisitDFPattern_(const PrimArrPatternNode* op) override; - void VisitDFPattern_(const UnorderedTuplePatternNode* op) override; + void VisitDFPattern_(const OrPatternObj* op) override; + void VisitDFPattern_(const AndPatternObj* op) override; + void VisitDFPattern_(const NotPatternObj* op) override; + void VisitDFPattern_(const AttrPatternObj* op) override; + void VisitDFPattern_(const CallPatternObj* op) override; + void VisitDFPattern_(const ConstantPatternObj* op) override; + void VisitDFPattern_(const DataTypePatternObj* op) override; + void VisitDFPattern_(const ExprPatternObj* op) override; + void VisitDFPattern_(const FunctionPatternObj* op) override; + void VisitDFPattern_(const ShapePatternObj* op) override; + void VisitDFPattern_(const TupleGetItemPatternObj* op) override; + void VisitDFPattern_(const TuplePatternObj* op) override; + void VisitDFPattern_(const TypePatternObj* op) override; + void VisitDFPattern_(const WildcardPatternObj* op) override; + void VisitDFPattern_(const VarPatternObj* op) override; + + void VisitDFPattern_(const DataflowVarPatternObj* op) override; + void VisitDFPattern_(const GlobalVarPatternObj* op) override; + void VisitDFPattern_(const ExternFuncPatternObj* op) override; + void VisitDFPattern_(const PrimArrPatternObj* op) override; + void VisitDFPattern_(const UnorderedTuplePatternObj* op) override; protected: // set of already-visited nodes diff --git a/include/tvm/relax/distributed/axis_group_graph.h b/include/tvm/relax/distributed/axis_group_graph.h index 307a0cf5aea3..df86576a7bfe 100644 --- a/include/tvm/relax/distributed/axis_group_graph.h +++ b/include/tvm/relax/distributed/axis_group_graph.h @@ -74,14 +74,14 @@ class TensorAxisGraphExtractor : public s_tir::StmtExprVisitor { extractor->Visit(function->body); ffi::Map inverse_tensor_map; for (const Var& param : function->params) { - if (param->ty.as()) { + if (param->ty.as()) { inverse_tensor_map.Set(param.as_or_throw(), param); } } std::vector> tir_var_axis_group_list; std::unordered_set visited; for (const Var& param : function->params) { - if (!param->ty.as()) { + if (!param->ty.as()) { continue; } TensorVar tensor = param.as_or_throw(); @@ -120,14 +120,14 @@ class TensorAxisGraphExtractor : public s_tir::StmtExprVisitor { } private: - ffi::Optional Visit_(const TensorStoreNode* op) final { + ffi::Optional Visit_(const TensorStoreObj* op) final { TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(s_tir::StmtExprVisitor::Visit_(op)); tensor_access_indices_.push_back({op->dest.as_or_throw(), op->indices}); return std::nullopt; } - ffi::Optional Visit_(const TensorLoadNode* op) final { + ffi::Optional Visit_(const TensorLoadObj* op) final { TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(s_tir::StmtExprVisitor::Visit_(op)); tensor_access_indices_.push_back({op->source.as_or_throw(), op->indices}); @@ -145,7 +145,6 @@ class TensorAxisGraphExtractor : public s_tir::StmtExprVisitor { return false; } Var var = *prim_var; - analyzer->Bind(iter_var_range_); b = analyzer->Simplify(b); // index var `a` must access whole range of a specific tensor dimension sym::IntSet intset_b = sym::EvalSet(b, sym::AsIntSet(iter_var_range_)); @@ -160,7 +159,7 @@ class TensorAxisGraphExtractor : public s_tir::StmtExprVisitor { return true; } - ffi::Optional Visit_(const s_tir::SBlockNode* op) final { + ffi::Optional Visit_(const s_tir::SBlockObj* op) final { if (op->name_hint == "root") { return s_tir::StmtExprVisitor::Visit_(op); } @@ -171,6 +170,11 @@ class TensorAxisGraphExtractor : public s_tir::StmtExprVisitor { iter_var_range_.Set(iter_var->var, iter_var->dom.value()); } sym::Analyzer analyzer; + sym::ConstraintContext context; + for (const auto& [var, range] : iter_var_range_) { + context->Append(sym::RangeConstraint(var, range)); + } + analyzer->SyncConstraintContext(context); for (const auto& access_pr : tensor_access_indices_) { TensorVar tensor = access_pr.first; ffi::Array indices = access_pr.second; @@ -219,13 +223,13 @@ namespace distributed { /*! \brief tensor axis*/ struct Axis { - const ExprNode* tensor; + const ExprObj* tensor; int dim = 0; int tuple_index = 0; - Axis(const ExprNode* tensor, int dim, int tuple_index = 0) + Axis(const ExprObj* tensor, int dim, int tuple_index = 0) : tensor(tensor), dim(dim), tuple_index(tuple_index) { - TVM_FFI_ICHECK(tensor->IsInstance() || tensor->IsInstance()); + TVM_FFI_ICHECK(tensor->IsInstance() || tensor->IsInstance()); } bool operator==(const Axis& other) const { @@ -236,7 +240,7 @@ struct Axis { class AxisHash { public: size_t operator()(const Axis& axis) const { - size_t const h1(std::hash()(axis.tensor)); + size_t const h1(std::hash()(axis.tensor)); size_t const h2(std::hash()(axis.dim)); size_t const h3(std::hash()(axis.tuple_index)); return h1 ^ (h2 << 1) ^ (h3 << 2); diff --git a/include/tvm/relax/distributed/global_info.h b/include/tvm/relax/distributed/global_info.h index 36d7d3b3037a..16a1501aaa62 100644 --- a/include/tvm/relax/distributed/global_info.h +++ b/include/tvm/relax/distributed/global_info.h @@ -36,7 +36,7 @@ namespace distributed { * \brief Device mesh express a view of topology of devices, represented by an n-d matrix of * device ids */ -class DeviceMeshNode : public GlobalInfoNode { +class DeviceMeshObj : public GlobalInfoObj { public: /*! \brief logical shape of the mesh*/ ffi::Shape shape; @@ -49,23 +49,23 @@ class DeviceMeshNode : public GlobalInfoNode { static void RegisterReflection() { namespace refl = tvm::ffi::reflection; - refl::ObjectDef() - .def_ro("shape", &DeviceMeshNode::shape) - .def_ro("device_ids", &DeviceMeshNode::device_ids) - .def_ro("device_range", &DeviceMeshNode::device_range); + refl::ObjectDef() + .def_ro("shape", &DeviceMeshObj::shape) + .def_ro("device_ids", &DeviceMeshObj::device_ids) + .def_ro("device_range", &DeviceMeshObj::device_range); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.distributed.DeviceMesh", DeviceMeshNode, GlobalInfoNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.distributed.DeviceMesh", DeviceMeshObj, GlobalInfoObj); }; /*! * \brief Managed reference to a DeviceMesh. - * \sa DeviceMeshNode + * \sa DeviceMeshObj */ class DeviceMesh : public GlobalInfo { public: TVM_DLL DeviceMesh(ffi::Shape shape, ffi::Array device_ids); TVM_DLL DeviceMesh(ffi::Shape shape, Range device_range); - TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(DeviceMesh, GlobalInfo, DeviceMeshNode); + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(DeviceMesh, GlobalInfo, DeviceMeshObj); }; } // namespace distributed diff --git a/include/tvm/relax/distributed/type.h b/include/tvm/relax/distributed/type.h index a6a1316dd5e6..ca3c7943fe37 100644 --- a/include/tvm/relax/distributed/type.h +++ b/include/tvm/relax/distributed/type.h @@ -38,7 +38,7 @@ namespace distributed { enum class PlacementSpecKind : int { kSharding = 0, kReplica = 1 }; /*! \brief Describes how data is distributed in one dimension of the device mesh*/ -class PlacementSpecNode : public ffi::Object { +class PlacementSpecObj : public ffi::Object { public: /*! \brief If the kind is sharding, this value represents the tensor dimension to shard. * otherwise, axis is -1. @@ -50,18 +50,18 @@ class PlacementSpecNode : public ffi::Object { static void RegisterReflection() { namespace refl = tvm::ffi::reflection; - refl::ObjectDef() - .def_ro("axis", &PlacementSpecNode::axis) - .def_ro("kind", &PlacementSpecNode::kind); + refl::ObjectDef() + .def_ro("axis", &PlacementSpecObj::axis) + .def_ro("kind", &PlacementSpecObj::kind); } static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindConstTreeNode; - TVM_FFI_DECLARE_OBJECT_INFO("relax.distributed.PlacementSpec", PlacementSpecNode, ffi::Object); + TVM_FFI_DECLARE_OBJECT_INFO("relax.distributed.PlacementSpec", PlacementSpecObj, ffi::Object); }; /*! - * \brief Managed reference to PlacementSpecNode. - * \sa PlacementSpecNode + * \brief Managed reference to PlacementSpecObj. + * \sa PlacementSpecObj */ class PlacementSpec : public ffi::ObjectRef { public: @@ -69,28 +69,28 @@ class PlacementSpec : public ffi::ObjectRef { TVM_DLL static PlacementSpec Replica(); - explicit PlacementSpec(ffi::UnsafeInit tag, ffi::ObjectPtr n) + explicit PlacementSpec(ffi::UnsafeInit tag, ffi::ObjectPtr n) : ffi::ObjectRef(tag) { data_ = std::move(n); } - TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(PlacementSpec, ffi::ObjectRef, PlacementSpecNode); + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(PlacementSpec, ffi::ObjectRef, PlacementSpecObj); }; -class ShardingNode : public PlacementSpecNode { +class ShardingObj : public PlacementSpecObj { public: /*! \brief The dimension of tensor we shard*/ int64_t sharding_dim; static void RegisterReflection() { namespace refl = tvm::ffi::reflection; - refl::ObjectDef().def_ro("sharding_dim", &ShardingNode::sharding_dim); + refl::ObjectDef().def_ro("sharding_dim", &ShardingObj::sharding_dim); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.distributed.Sharding", ShardingNode, PlacementSpecNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.distributed.Sharding", ShardingObj, PlacementSpecObj); }; /*! \brief Describes how data is distributed in each dimension of the device mesh*/ -class PlacementNode : public ffi::Object { +class PlacementObj : public ffi::Object { public: /*! \brief specs for each dim of device mesh.*/ ffi::Array dim_specs; @@ -99,36 +99,36 @@ class PlacementNode : public ffi::Object { static void RegisterReflection() { namespace refl = tvm::ffi::reflection; - refl::ObjectDef().def_ro("dim_specs", &PlacementNode::dim_specs); + refl::ObjectDef().def_ro("dim_specs", &PlacementObj::dim_specs); } static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindConstTreeNode; - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.distributed.Placement", PlacementNode, ffi::Object); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.distributed.Placement", PlacementObj, ffi::Object); }; /*! * \brief Managed reference to a Placement. - * \sa PlacementNode + * \sa PlacementObj */ class Placement : public ffi::ObjectRef { public: TVM_DLL explicit Placement(ffi::Array dim_specs); /*! \brief replica dim is printed as "R" and sharding dim is printed as "S[i]".]*/ static Placement FromText(ffi::String text_repr); - TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Placement, ffi::ObjectRef, PlacementNode); + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Placement, ffi::ObjectRef, PlacementObj); }; /*! * \brief Type of DTensor (Distributed Tensor). */ -class DTensorTypeNode : public TypeNode { +class DTensorTypeObj : public TypeObj { public: - explicit DTensorTypeNode(ffi::UnsafeInit) + explicit DTensorTypeObj(ffi::UnsafeInit) : tensor_ty(ffi::UnsafeInit{}), device_mesh(ffi::UnsafeInit{}), placement(ffi::UnsafeInit{}) {} - DTensorTypeNode(TensorType tensor_ty, DeviceMesh device_mesh, Placement placement) + DTensorTypeObj(TensorType tensor_ty, DeviceMesh device_mesh, Placement placement) : tensor_ty(std::move(tensor_ty)), device_mesh(std::move(device_mesh)), placement(std::move(placement)) {} @@ -148,17 +148,17 @@ class DTensorTypeNode : public TypeNode { static void RegisterReflection() { namespace refl = tvm::ffi::reflection; - refl::ObjectDef() - .def_ro("device_mesh", &DTensorTypeNode::device_mesh) - .def_ro("placement", &DTensorTypeNode::placement) - .def_ro("tensor_ty", &DTensorTypeNode::tensor_ty); + refl::ObjectDef() + .def_ro("device_mesh", &DTensorTypeObj::device_mesh) + .def_ro("placement", &DTensorTypeObj::placement) + .def_ro("tensor_ty", &DTensorTypeObj::tensor_ty); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.DTensorType", DTensorTypeNode, TypeNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.DTensorType", DTensorTypeObj, TypeObj); }; /*! - * \brief Managed reference to DTensorTypeNode. - * \sa DTensorTypeNode + * \brief Managed reference to DTensorTypeObj. + * \sa DTensorTypeObj */ class DTensorType : public Type { public: @@ -172,7 +172,7 @@ class DTensorType : public Type { TVM_DLL DTensorType(TensorType tensor_ty, DeviceMesh device_mesh, Placement placement, Location loc = UnknownLoc()); - TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(DTensorType, Type, DTensorTypeNode); + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(DTensorType, Type, DTensorTypeObj); }; } // namespace distributed diff --git a/include/tvm/relax/expr.h b/include/tvm/relax/expr.h index 7374ab124d92..13823c7cdce8 100644 --- a/include/tvm/relax/expr.h +++ b/include/tvm/relax/expr.h @@ -42,44 +42,44 @@ namespace relax { // Compatibility aliases. Tuple expressions are defined in the common IR. using ::tvm::Tuple; using ::tvm::TupleGetItem; -using ::tvm::TupleGetItemNode; -using ::tvm::TupleNode; +using ::tvm::TupleGetItemObj; +using ::tvm::TupleObj; /*! \brief A shape expression which allows users to construct a shape containing PrimExpr. */ -class ShapeExprNode : public ExprNode { +class ShapeExprObj : public ExprObj { public: /*! The values of the shape expression. */ ffi::Array values; static void RegisterReflection() { namespace refl = tvm::ffi::reflection; - refl::ObjectDef().def_ro("values", &ShapeExprNode::values); + refl::ObjectDef().def_ro("values", &ShapeExprObj::values); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.expr.ShapeExpr", ShapeExprNode, ExprNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.expr.ShapeExpr", ShapeExprObj, ExprObj); }; class ShapeExpr : public Expr { public: TVM_DLL explicit ShapeExpr(ffi::Array values, Location loc = UnknownLoc()); - explicit ShapeExpr(ffi::ObjectPtr node) : Expr(std::move(node)) {} + explicit ShapeExpr(ffi::ObjectPtr node) : Expr(std::move(node)) {} - TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ShapeExpr, Expr, ShapeExprNode); - TVM_DEFINE_OBJECT_REF_COW_METHOD(ShapeExprNode); + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ShapeExpr, Expr, ShapeExprObj); + TVM_DEFINE_OBJECT_REF_COW_METHOD(ShapeExprObj); }; /*! \brief A sub-type of the variable node used to mark dataflow variables from * normal visible "function local" bindings. */ -class DataflowVarNode : public VarNode { +class DataflowVarObj : public VarObj { public: static void RegisterReflection() { namespace refl = tvm::ffi::reflection; - refl::ObjectDef(); + refl::ObjectDef(); } static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindFreeVar; - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.expr.DataflowVar", DataflowVarNode, VarNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.expr.DataflowVar", DataflowVarObj, VarObj); }; class DataflowVar : public Var { @@ -87,9 +87,9 @@ class DataflowVar : public Var { TVM_DLL explicit DataflowVar(ffi::String name, ffi::Optional ty_annotation, Location loc = UnknownLoc()); - explicit DataflowVar(ffi::ObjectPtr node) : Var(std::move(node)) {} + explicit DataflowVar(ffi::ObjectPtr node) : Var(std::move(node)) {} - TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(DataflowVar, Var, DataflowVarNode); + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(DataflowVar, Var, DataflowVarObj); }; /*! \brief Construct a tensor constant, inferring shape and dtype when type is omitted. */ @@ -98,10 +98,10 @@ TVM_DLL GenericConst MakeTensorConst(ffi::Tensor data, Location loc = UnknownLoc()); /*! \brief The base class of a variable binding in Relax. */ -class BindingNode : public ffi::Object { +class BindingObj : public ffi::Object { public: - explicit BindingNode(Var var) : var(std::move(var)) {} - explicit BindingNode(ffi::UnsafeInit) : var(ffi::UnsafeInit{}) {} + explicit BindingObj(Var var) : var(std::move(var)) {} + explicit BindingObj(ffi::UnsafeInit) : var(ffi::UnsafeInit{}) {} mutable Location loc = UnknownLoc(); /*! \brief The return variable to bound to. */ @@ -109,15 +109,15 @@ class BindingNode : public ffi::Object { static void RegisterReflection() { namespace refl = tvm::ffi::reflection; - refl::ObjectDef() - .def_ro("loc", &BindingNode::loc, refl::AttachFieldFlag::SEqHashIgnore()) + refl::ObjectDef() + .def_ro("loc", &BindingObj::loc, refl::AttachFieldFlag::SEqHashIgnore()) // TODO(tqchen): use SEqHashDefSimple after the next pypi tvm-ffi release - .def_ro("var", &BindingNode::var, refl::AttachFieldFlag::SEqHashDefPattern()); + .def_ro("var", &BindingObj::var, refl::AttachFieldFlag::SEqHashDefPattern()); } static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode; - TVM_FFI_DECLARE_OBJECT_INFO("relax.expr.Binding", BindingNode, ffi::Object); + TVM_FFI_DECLARE_OBJECT_INFO("relax.expr.Binding", BindingObj, ffi::Object); }; class Binding : public ffi::ObjectRef { @@ -130,10 +130,10 @@ class Binding : public ffi::ObjectRef { Binding(Binding&&) = default; Binding& operator=(const Binding&) = default; Binding& operator=(Binding&&) = default; - const BindingNode* operator->() const { return static_cast(data_.get()); } - const BindingNode* get() const { return operator->(); } + const BindingObj* operator->() const { return static_cast(data_.get()); } + const BindingObj* get() const { return operator->(); } static constexpr bool _type_is_nullable = false; - using ContainerType = BindingNode; + using ContainerType = BindingObj; }; /*! @@ -143,12 +143,12 @@ class Binding : public ffi::ObjectRef { * and vars in ty in first occurance, and insert equality assertions in * other cases. */ -class MatchCastNode : public BindingNode { +class MatchCastObj : public BindingObj { public: - explicit MatchCastNode(Var var, Expr value) - : BindingNode(std::move(var)), value(std::move(value)) {} - explicit MatchCastNode(ffi::UnsafeInit) - : BindingNode(ffi::UnsafeInit{}), value(ffi::UnsafeInit{}) {} + explicit MatchCastObj(Var var, Expr value) + : BindingObj(std::move(var)), value(std::move(value)) {} + explicit MatchCastObj(ffi::UnsafeInit) + : BindingObj(ffi::UnsafeInit{}), value(ffi::UnsafeInit{}) {} /*! \brief The input value to match cast. */ Expr value; @@ -157,131 +157,130 @@ class MatchCastNode : public BindingNode { static void RegisterReflection() { namespace refl = tvm::ffi::reflection; - refl::ObjectDef() - .def_ro("value", &MatchCastNode::value) + refl::ObjectDef() + .def_ro("value", &MatchCastObj::value) // TODO(tqchen): use SEqHashDefSimple after the next pypi tvm-ffi release - .def_ro("ty", &MatchCastNode::ty, refl::AttachFieldFlag::SEqHashDefPattern()); + .def_ro("ty", &MatchCastObj::ty, refl::AttachFieldFlag::SEqHashDefPattern()); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.expr.MatchCast", MatchCastNode, BindingNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.expr.MatchCast", MatchCastObj, BindingObj); }; /*! - * \brief Managed reference to MatchCastNode. - * \sa MatchCastNode + * \brief Managed reference to MatchCastObj. + * \sa MatchCastObj */ class MatchCast : public Binding { public: TVM_DLL explicit MatchCast(Var var, Expr value, Type ty, Location loc = UnknownLoc()); - TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(MatchCast, Binding, MatchCastNode); - TVM_DEFINE_OBJECT_REF_COW_METHOD(MatchCastNode); + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(MatchCast, Binding, MatchCastObj); + TVM_DEFINE_OBJECT_REF_COW_METHOD(MatchCastObj); }; -class VarBindingNode : public BindingNode { +class VarBindingObj : public BindingObj { public: - explicit VarBindingNode(Var var, Expr value) - : BindingNode(std::move(var)), value(std::move(value)) {} - explicit VarBindingNode(ffi::UnsafeInit) - : BindingNode(ffi::UnsafeInit{}), value(ffi::UnsafeInit{}) {} + explicit VarBindingObj(Var var, Expr value) + : BindingObj(std::move(var)), value(std::move(value)) {} + explicit VarBindingObj(ffi::UnsafeInit) + : BindingObj(ffi::UnsafeInit{}), value(ffi::UnsafeInit{}) {} /*! \brief The binding value. */ Expr value; static void RegisterReflection() { namespace refl = tvm::ffi::reflection; - refl::ObjectDef().def_ro("value", &VarBindingNode::value); + refl::ObjectDef().def_ro("value", &VarBindingObj::value); // customize the SEqual and SHash methods for better error messages - refl::TypeAttrDef() - .def(tvm::ffi::type_attr::kSEqual.name, &VarBindingNode::SEqual) - .def(tvm::ffi::type_attr::kSHash.name, &VarBindingNode::SHash); + refl::TypeAttrDef() + .def(tvm::ffi::type_attr::kSEqual.name, &VarBindingObj::SEqual) + .def(tvm::ffi::type_attr::kSHash.name, &VarBindingObj::SHash); } - bool SEqual(const VarBindingNode* other, + bool SEqual(const VarBindingObj* other, ffi::TypedFunction equal) const; int64_t SHash(int64_t init_hash, ffi::TypedFunction hash) const; - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.expr.VarBinding", VarBindingNode, BindingNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.expr.VarBinding", VarBindingObj, BindingObj); }; class VarBinding : public Binding { public: TVM_DLL explicit VarBinding(Var var, Expr value, Location loc = UnknownLoc()); - TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(VarBinding, Binding, VarBindingNode); - TVM_DEFINE_OBJECT_REF_COW_METHOD(VarBindingNode); + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(VarBinding, Binding, VarBindingObj); + TVM_DEFINE_OBJECT_REF_COW_METHOD(VarBindingObj); }; -class BindingBlockNode : public ffi::Object { +class BindingBlockObj : public ffi::Object { public: ffi::Array bindings; mutable Location loc = UnknownLoc(); static void RegisterReflection() { namespace refl = tvm::ffi::reflection; - refl::ObjectDef() - .def_ro("bindings", &BindingBlockNode::bindings) - .def_ro("loc", &BindingBlockNode::loc, refl::AttachFieldFlag::SEqHashIgnore(), + refl::ObjectDef() + .def_ro("bindings", &BindingBlockObj::bindings) + .def_ro("loc", &BindingBlockObj::loc, refl::AttachFieldFlag::SEqHashIgnore(), refl::DefaultValue(UnknownLoc())); } static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode; - TVM_FFI_DECLARE_OBJECT_INFO("relax.expr.BindingBlock", BindingBlockNode, ffi::Object); + TVM_FFI_DECLARE_OBJECT_INFO("relax.expr.BindingBlock", BindingBlockObj, ffi::Object); }; class BindingBlock : public ffi::ObjectRef { public: TVM_DLL explicit BindingBlock(ffi::Array bindings, Location loc = UnknownLoc()); - explicit BindingBlock(ffi::UnsafeInit tag, ffi::ObjectPtr n) + explicit BindingBlock(ffi::UnsafeInit tag, ffi::ObjectPtr n) : ffi::ObjectRef(tag) { data_ = std::move(n); } - TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(BindingBlock, ffi::ObjectRef, BindingBlockNode); + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(BindingBlock, ffi::ObjectRef, BindingBlockObj); - BindingBlockNode* CopyOnWrite(); + BindingBlockObj* CopyOnWrite(); }; -class DataflowBlockNode : public BindingBlockNode { +class DataflowBlockObj : public BindingBlockObj { public: static void RegisterReflection() { namespace refl = tvm::ffi::reflection; - refl::ObjectDef(); + refl::ObjectDef(); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.expr.DataflowBlock", DataflowBlockNode, - BindingBlockNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.expr.DataflowBlock", DataflowBlockObj, BindingBlockObj); }; class DataflowBlock : public BindingBlock { public: TVM_DLL explicit DataflowBlock(ffi::Array bindings, Location loc = UnknownLoc()); - explicit DataflowBlock(ffi::UnsafeInit tag, ffi::ObjectPtr n) + explicit DataflowBlock(ffi::UnsafeInit tag, ffi::ObjectPtr n) : BindingBlock(tag) { data_ = std::move(n); } - TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(DataflowBlock, BindingBlock, DataflowBlockNode); - TVM_DEFINE_OBJECT_REF_COW_METHOD(DataflowBlockNode); + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(DataflowBlock, BindingBlock, DataflowBlockObj); + TVM_DEFINE_OBJECT_REF_COW_METHOD(DataflowBlockObj); }; /*! \brief A sequence of blocks followed by an expression. * * The order of blocks enforces scoping and ordering. */ -class SeqExprNode : public ExprNode { +class SeqExprObj : public ExprObj { public: - explicit SeqExprNode(Expr body) : body(std::move(body)) {} - explicit SeqExprNode(ffi::UnsafeInit) : body(ffi::UnsafeInit{}) {} + explicit SeqExprObj(Expr body) : body(std::move(body)) {} + explicit SeqExprObj(ffi::UnsafeInit) : body(ffi::UnsafeInit{}) {} ffi::Array blocks; Expr body; static void RegisterReflection() { namespace refl = tvm::ffi::reflection; - refl::ObjectDef() - .def_ro("blocks", &SeqExprNode::blocks) - .def_ro("body", &SeqExprNode::body); - refl::TypeAttrDef() - .def(tvm::ffi::type_attr::kSEqual.name, &SeqExprNode::SEqual) - .def(tvm::ffi::type_attr::kSHash.name, &SeqExprNode::SHash); + refl::ObjectDef() + .def_ro("blocks", &SeqExprObj::blocks) + .def_ro("body", &SeqExprObj::body); + refl::TypeAttrDef() + .def(tvm::ffi::type_attr::kSEqual.name, &SeqExprObj::SEqual) + .def(tvm::ffi::type_attr::kSHash.name, &SeqExprObj::SHash); } - bool SEqual(const SeqExprNode* other, + bool SEqual(const SeqExprObj* other, ffi::TypedFunction equal) const { // Establish mappings for symbolic variables defined by bindings before // comparing their uses in the SeqExpr result type and body. @@ -296,7 +295,7 @@ class SeqExprNode : public ExprNode { hash_value = hash(body, hash_value, false); return hash_value; } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.expr.SeqExpr", SeqExprNode, ExprNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.expr.SeqExpr", SeqExprObj, ExprObj); }; class SeqExpr : public Expr { @@ -310,15 +309,15 @@ class SeqExpr : public Expr { * a transform may use `func.CopyOnWrite()->body = expr;`. * * If the expression is already a `relax::SeqExpr`, the same - * underlying `relax::SeqExprNode` is used, and no copies are made. + * underlying `relax::SeqExprObj` is used, and no copies are made. */ TVM_DLL SeqExpr(Expr body); // NOLINT(*) TVM_DLL explicit SeqExpr(ffi::Array blocks, Expr body, Location loc = UnknownLoc()); - explicit SeqExpr(ffi::ObjectPtr node) : Expr(std::move(node)) {} + explicit SeqExpr(ffi::ObjectPtr node) : Expr(std::move(node)) {} - TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(SeqExpr, Expr, SeqExprNode); - TVM_DEFINE_OBJECT_REF_COW_METHOD(SeqExprNode); + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(SeqExpr, Expr, SeqExprObj); + TVM_DEFINE_OBJECT_REF_COW_METHOD(SeqExprObj); }; /*! @@ -332,13 +331,13 @@ class SeqExpr : public Expr { * * \note This is similar to C's ternary operator. */ -class IfExprNode : public ExprNode { +class IfExprObj : public ExprObj { public: - explicit IfExprNode(Expr cond, SeqExpr true_branch, SeqExpr false_branch) + explicit IfExprObj(Expr cond, SeqExpr true_branch, SeqExpr false_branch) : cond(std::move(cond)), true_branch(std::move(true_branch)), false_branch(std::move(false_branch)) {} - explicit IfExprNode(ffi::UnsafeInit) + explicit IfExprObj(ffi::UnsafeInit) : cond(ffi::UnsafeInit{}), true_branch(ffi::UnsafeInit{}), false_branch(ffi::UnsafeInit{}) {} /*! \brief The condition. */ @@ -350,14 +349,14 @@ class IfExprNode : public ExprNode { static void RegisterReflection() { namespace refl = tvm::ffi::reflection; - refl::ObjectDef() - .def_ro("cond", &IfExprNode::cond) - .def_ro("true_branch", &IfExprNode::true_branch) - .def_ro("false_branch", &IfExprNode::false_branch); + refl::ObjectDef() + .def_ro("cond", &IfExprObj::cond) + .def_ro("true_branch", &IfExprObj::true_branch) + .def_ro("false_branch", &IfExprObj::false_branch); } static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindDAGNode; - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.expr.IfExpr", IfExprNode, ExprNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.expr.IfExpr", IfExprObj, ExprObj); }; class IfExpr : public Expr { @@ -381,17 +380,17 @@ class IfExpr : public Expr { */ TVM_DLL IfExpr(Expr cond, Expr true_branch, Expr false_branch, Location loc = UnknownLoc()); - explicit IfExpr(ffi::ObjectPtr node) : Expr(std::move(node)) {} + explicit IfExpr(ffi::ObjectPtr node) : Expr(std::move(node)) {} - TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(IfExpr, Expr, IfExprNode); - TVM_DEFINE_OBJECT_REF_COW_METHOD(IfExprNode); + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(IfExpr, Expr, IfExprObj); + TVM_DEFINE_OBJECT_REF_COW_METHOD(IfExprObj); }; /*! \brief A Relax function. */ -class FunctionNode : public BaseFuncNode { +class FunctionObj : public BaseFuncObj { public: - explicit FunctionNode(SeqExpr body) : body(std::move(body)) {} - explicit FunctionNode(ffi::UnsafeInit) : body(ffi::UnsafeInit{}) {} + explicit FunctionObj(SeqExpr body) : body(std::move(body)) {} + explicit FunctionObj(ffi::UnsafeInit) : body(ffi::UnsafeInit{}) {} /*! \brief The parameters to the function. */ ffi::Array params; @@ -404,15 +403,15 @@ class FunctionNode : public BaseFuncNode { static void RegisterReflection() { namespace refl = tvm::ffi::reflection; - refl::ObjectDef() - .def_ro("params", &FunctionNode::params, refl::AttachFieldFlag::SEqHashDefPattern()) - .def_ro("body", &FunctionNode::body) - .def_ro("ret_ty", &FunctionNode::ret_ty) - .def_ro("is_pure", &FunctionNode::is_pure); + refl::ObjectDef() + .def_ro("params", &FunctionObj::params, refl::AttachFieldFlag::SEqHashDefPattern()) + .def_ro("body", &FunctionObj::body) + .def_ro("ret_ty", &FunctionObj::ret_ty) + .def_ro("is_pure", &FunctionObj::is_pure); } static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindDAGNode; - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.expr.Function", FunctionNode, BaseFuncNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.expr.Function", FunctionObj, BaseFuncObj); }; class Function : public BaseFunc { @@ -449,23 +448,23 @@ class Function : public BaseFunc { TVM_DLL static Function CreateEmpty(ffi::Array params, Type ret_ty, bool is_pure = true, DictAttrs attrs = DictAttrs(), Location loc = UnknownLoc()); - explicit Function(ffi::ObjectPtr node) : BaseFunc(std::move(node)) {} + explicit Function(ffi::ObjectPtr node) : BaseFunc(std::move(node)) {} - TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Function, BaseFunc, FunctionNode); - TVM_DEFINE_OBJECT_REF_COW_METHOD(FunctionNode); + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Function, BaseFunc, FunctionObj); + TVM_DEFINE_OBJECT_REF_COW_METHOD(FunctionObj); }; /*! \brief The extern function, which can represent packed function. */ -class ExternFuncNode : public BaseFuncNode { +class ExternFuncObj : public BaseFuncObj { public: /*! \brief The name of global symbol. */ ffi::String global_symbol; static void RegisterReflection() { namespace refl = tvm::ffi::reflection; - refl::ObjectDef().def_ro("global_symbol", &ExternFuncNode::global_symbol); + refl::ObjectDef().def_ro("global_symbol", &ExternFuncObj::global_symbol); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.expr.ExternFunc", ExternFuncNode, BaseFuncNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.expr.ExternFunc", ExternFuncObj, BaseFuncObj); }; class ExternFunc : public BaseFunc { @@ -473,10 +472,10 @@ class ExternFunc : public BaseFunc { TVM_DLL ExternFunc(ffi::String global_symbol, Location loc = UnknownLoc()); TVM_DLL ExternFunc(ffi::String global_symbol, Type ty, Location loc = UnknownLoc()); - explicit ExternFunc(ffi::ObjectPtr node) : BaseFunc(std::move(node)) {} + explicit ExternFunc(ffi::ObjectPtr node) : BaseFunc(std::move(node)) {} - TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ExternFunc, BaseFunc, ExternFuncNode); - TVM_DEFINE_OBJECT_REF_COW_METHOD(ExternFuncNode); + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ExternFunc, BaseFunc, ExternFuncObj); + TVM_DEFINE_OBJECT_REF_COW_METHOD(ExternFuncObj); }; /*! diff --git a/include/tvm/relax/expr_functor.h b/include/tvm/relax/expr_functor.h index e623f0ca8381..f5d627e689ae 100644 --- a/include/tvm/relax/expr_functor.h +++ b/include/tvm/relax/expr_functor.h @@ -146,55 +146,55 @@ class ExprFunctor { // Functions that can be overriden by subclass // NOTE: cross dialect calls are invoked through global var // We do not expect inline tirx::Function to appear in relax IR. - virtual R VisitExpr_(const GenericConstNode* op, Args... args) EXPR_FUNCTOR_DEFAULT; - virtual R VisitExpr_(const TupleNode* op, Args... args) EXPR_FUNCTOR_DEFAULT; - virtual R VisitExpr_(const VarNode* op, Args... args) EXPR_FUNCTOR_DEFAULT; - virtual R VisitExpr_(const DataflowVarNode* op, Args... args) EXPR_FUNCTOR_DEFAULT; - virtual R VisitExpr_(const ShapeExprNode* op, Args... args) EXPR_FUNCTOR_DEFAULT; - virtual R VisitExpr_(const ExternFuncNode* op, Args... args) EXPR_FUNCTOR_DEFAULT; - virtual R VisitExpr_(const GlobalVarNode* op, Args... args) EXPR_FUNCTOR_DEFAULT; - virtual R VisitExpr_(const FunctionNode* op, Args... args) EXPR_FUNCTOR_DEFAULT; - virtual R VisitExpr_(const CallNode* op, Args... args) EXPR_FUNCTOR_DEFAULT; - virtual R VisitExpr_(const TensorLoadNode* op, Args... args) EXPR_FUNCTOR_DEFAULT; - virtual R VisitExpr_(const tvm::IntImmNode* op, Args... args) EXPR_FUNCTOR_DEFAULT; - virtual R VisitExpr_(const tvm::FloatImmNode* op, Args... args) EXPR_FUNCTOR_DEFAULT; - virtual R VisitExpr_(const SeqExprNode* op, Args... args) EXPR_FUNCTOR_DEFAULT; - virtual R VisitExpr_(const IfExprNode* op, Args... args) EXPR_FUNCTOR_DEFAULT; - virtual R VisitExpr_(const OpNode* op, Args... args) EXPR_FUNCTOR_DEFAULT; - virtual R VisitExpr_(const TupleGetItemNode* op, Args... args) EXPR_FUNCTOR_DEFAULT; - virtual R VisitExpr_(const StringImmNode* op, Args... args) EXPR_FUNCTOR_DEFAULT; - virtual R VisitExpr_(const DataTypeImmNode* op, Args... args) EXPR_FUNCTOR_DEFAULT; - virtual R VisitExpr_(const prim::LetNode* op, Args...) EXPR_FUNCTOR_DISABLED; - virtual R VisitExpr_(const prim::AddNode* op, Args... args) EXPR_FUNCTOR_DEFAULT; - virtual R VisitExpr_(const prim::LShiftNode* op, Args... args) EXPR_FUNCTOR_DEFAULT; - virtual R VisitExpr_(const prim::RShiftNode* op, Args... args) EXPR_FUNCTOR_DEFAULT; - virtual R VisitExpr_(const prim::BitwiseAndNode* op, Args... args) EXPR_FUNCTOR_DEFAULT; - virtual R VisitExpr_(const prim::BitwiseOrNode* op, Args... args) EXPR_FUNCTOR_DEFAULT; - virtual R VisitExpr_(const prim::BitwiseXorNode* op, Args... args) EXPR_FUNCTOR_DEFAULT; - virtual R VisitExpr_(const prim::BitwiseNotNode* op, Args... args) EXPR_FUNCTOR_DEFAULT; - virtual R VisitExpr_(const prim::SubNode* op, Args... args) EXPR_FUNCTOR_DEFAULT; - virtual R VisitExpr_(const prim::MulNode* op, Args... args) EXPR_FUNCTOR_DEFAULT; - virtual R VisitExpr_(const prim::DivNode* op, Args... args) EXPR_FUNCTOR_DEFAULT; - virtual R VisitExpr_(const prim::ModNode* op, Args... args) EXPR_FUNCTOR_DEFAULT; - virtual R VisitExpr_(const prim::FloorDivNode* op, Args... args) EXPR_FUNCTOR_DEFAULT; - virtual R VisitExpr_(const prim::FloorModNode* op, Args... args) EXPR_FUNCTOR_DEFAULT; - virtual R VisitExpr_(const prim::MinNode* op, Args... args) EXPR_FUNCTOR_DEFAULT; - virtual R VisitExpr_(const prim::MaxNode* op, Args... args) EXPR_FUNCTOR_DEFAULT; - virtual R VisitExpr_(const prim::EQNode* op, Args... args) EXPR_FUNCTOR_DEFAULT; - virtual R VisitExpr_(const prim::NENode* op, Args... args) EXPR_FUNCTOR_DEFAULT; - virtual R VisitExpr_(const prim::LTNode* op, Args... args) EXPR_FUNCTOR_DEFAULT; - virtual R VisitExpr_(const prim::LENode* op, Args... args) EXPR_FUNCTOR_DEFAULT; - virtual R VisitExpr_(const prim::GTNode* op, Args... args) EXPR_FUNCTOR_DEFAULT; - virtual R VisitExpr_(const prim::GENode* op, Args... args) EXPR_FUNCTOR_DEFAULT; - virtual R VisitExpr_(const prim::AndNode* op, Args... args) EXPR_FUNCTOR_DEFAULT; - virtual R VisitExpr_(const prim::OrNode* op, Args... args) EXPR_FUNCTOR_DEFAULT; - virtual R VisitExpr_(const prim::CastNode* op, Args... args) EXPR_FUNCTOR_DEFAULT; - virtual R VisitExpr_(const prim::NotNode* op, Args... args) EXPR_FUNCTOR_DEFAULT; - virtual R VisitExpr_(const prim::SelectNode* op, Args... args) EXPR_FUNCTOR_DEFAULT; - virtual R VisitExpr_(const prim::RampNode* op, Args... args) EXPR_FUNCTOR_DEFAULT; - virtual R VisitExpr_(const prim::BroadcastNode* op, Args... args) EXPR_FUNCTOR_DEFAULT; - virtual R VisitExpr_(const prim::ShuffleNode* op, Args... args) EXPR_FUNCTOR_DEFAULT; - virtual R VisitExprFallback_(const ExprNode* op, Args... args) EXPR_FUNCTOR_DEFAULT; + virtual R VisitExpr_(const GenericConstObj* op, Args... args) EXPR_FUNCTOR_DEFAULT; + virtual R VisitExpr_(const TupleObj* op, Args... args) EXPR_FUNCTOR_DEFAULT; + virtual R VisitExpr_(const VarObj* op, Args... args) EXPR_FUNCTOR_DEFAULT; + virtual R VisitExpr_(const DataflowVarObj* op, Args... args) EXPR_FUNCTOR_DEFAULT; + virtual R VisitExpr_(const ShapeExprObj* op, Args... args) EXPR_FUNCTOR_DEFAULT; + virtual R VisitExpr_(const ExternFuncObj* op, Args... args) EXPR_FUNCTOR_DEFAULT; + virtual R VisitExpr_(const GlobalVarObj* op, Args... args) EXPR_FUNCTOR_DEFAULT; + virtual R VisitExpr_(const FunctionObj* op, Args... args) EXPR_FUNCTOR_DEFAULT; + virtual R VisitExpr_(const CallObj* op, Args... args) EXPR_FUNCTOR_DEFAULT; + virtual R VisitExpr_(const TensorLoadObj* op, Args... args) EXPR_FUNCTOR_DEFAULT; + virtual R VisitExpr_(const tvm::IntImmObj* op, Args... args) EXPR_FUNCTOR_DEFAULT; + virtual R VisitExpr_(const tvm::FloatImmObj* op, Args... args) EXPR_FUNCTOR_DEFAULT; + virtual R VisitExpr_(const SeqExprObj* op, Args... args) EXPR_FUNCTOR_DEFAULT; + virtual R VisitExpr_(const IfExprObj* op, Args... args) EXPR_FUNCTOR_DEFAULT; + virtual R VisitExpr_(const OpObj* op, Args... args) EXPR_FUNCTOR_DEFAULT; + virtual R VisitExpr_(const TupleGetItemObj* op, Args... args) EXPR_FUNCTOR_DEFAULT; + virtual R VisitExpr_(const StringImmObj* op, Args... args) EXPR_FUNCTOR_DEFAULT; + virtual R VisitExpr_(const DataTypeImmObj* op, Args... args) EXPR_FUNCTOR_DEFAULT; + virtual R VisitExpr_(const prim::LetObj* op, Args...) EXPR_FUNCTOR_DISABLED; + virtual R VisitExpr_(const prim::AddObj* op, Args... args) EXPR_FUNCTOR_DEFAULT; + virtual R VisitExpr_(const prim::LShiftObj* op, Args... args) EXPR_FUNCTOR_DEFAULT; + virtual R VisitExpr_(const prim::RShiftObj* op, Args... args) EXPR_FUNCTOR_DEFAULT; + virtual R VisitExpr_(const prim::BitwiseAndObj* op, Args... args) EXPR_FUNCTOR_DEFAULT; + virtual R VisitExpr_(const prim::BitwiseOrObj* op, Args... args) EXPR_FUNCTOR_DEFAULT; + virtual R VisitExpr_(const prim::BitwiseXorObj* op, Args... args) EXPR_FUNCTOR_DEFAULT; + virtual R VisitExpr_(const prim::BitwiseNotObj* op, Args... args) EXPR_FUNCTOR_DEFAULT; + virtual R VisitExpr_(const prim::SubObj* op, Args... args) EXPR_FUNCTOR_DEFAULT; + virtual R VisitExpr_(const prim::MulObj* op, Args... args) EXPR_FUNCTOR_DEFAULT; + virtual R VisitExpr_(const prim::DivObj* op, Args... args) EXPR_FUNCTOR_DEFAULT; + virtual R VisitExpr_(const prim::ModObj* op, Args... args) EXPR_FUNCTOR_DEFAULT; + virtual R VisitExpr_(const prim::FloorDivObj* op, Args... args) EXPR_FUNCTOR_DEFAULT; + virtual R VisitExpr_(const prim::FloorModObj* op, Args... args) EXPR_FUNCTOR_DEFAULT; + virtual R VisitExpr_(const prim::MinObj* op, Args... args) EXPR_FUNCTOR_DEFAULT; + virtual R VisitExpr_(const prim::MaxObj* op, Args... args) EXPR_FUNCTOR_DEFAULT; + virtual R VisitExpr_(const prim::EQObj* op, Args... args) EXPR_FUNCTOR_DEFAULT; + virtual R VisitExpr_(const prim::NEObj* op, Args... args) EXPR_FUNCTOR_DEFAULT; + virtual R VisitExpr_(const prim::LTObj* op, Args... args) EXPR_FUNCTOR_DEFAULT; + virtual R VisitExpr_(const prim::LEObj* op, Args... args) EXPR_FUNCTOR_DEFAULT; + virtual R VisitExpr_(const prim::GTObj* op, Args... args) EXPR_FUNCTOR_DEFAULT; + virtual R VisitExpr_(const prim::GEObj* op, Args... args) EXPR_FUNCTOR_DEFAULT; + virtual R VisitExpr_(const prim::AndObj* op, Args... args) EXPR_FUNCTOR_DEFAULT; + virtual R VisitExpr_(const prim::OrObj* op, Args... args) EXPR_FUNCTOR_DEFAULT; + virtual R VisitExpr_(const prim::CastObj* op, Args... args) EXPR_FUNCTOR_DEFAULT; + virtual R VisitExpr_(const prim::NotObj* op, Args... args) EXPR_FUNCTOR_DEFAULT; + virtual R VisitExpr_(const prim::SelectObj* op, Args... args) EXPR_FUNCTOR_DEFAULT; + virtual R VisitExpr_(const prim::RampObj* op, Args... args) EXPR_FUNCTOR_DEFAULT; + virtual R VisitExpr_(const prim::BroadcastObj* op, Args... args) EXPR_FUNCTOR_DEFAULT; + virtual R VisitExpr_(const prim::ShuffleObj* op, Args... args) EXPR_FUNCTOR_DEFAULT; + virtual R VisitExprFallback_(const ExprObj* op, Args... args) EXPR_FUNCTOR_DEFAULT; virtual R VisitExprDefault_(const ffi::Object* op, Args...) { TVM_FFI_THROW(InternalError) << "Do not have a default for " << op->GetTypeKey(); throw; @@ -205,54 +205,54 @@ class ExprFunctor { static FType InitVTable() { FType vtable; // Set dispatch - RELAX_EXPR_FUNCTOR_DISPATCH(GenericConstNode); - RELAX_EXPR_FUNCTOR_DISPATCH(TupleNode); - RELAX_EXPR_FUNCTOR_DISPATCH(VarNode); - RELAX_EXPR_FUNCTOR_DISPATCH(DataflowVarNode); - RELAX_EXPR_FUNCTOR_DISPATCH(ShapeExprNode); - RELAX_EXPR_FUNCTOR_DISPATCH(ExternFuncNode); - RELAX_EXPR_FUNCTOR_DISPATCH(GlobalVarNode); - RELAX_EXPR_FUNCTOR_DISPATCH(FunctionNode); - RELAX_EXPR_FUNCTOR_DISPATCH(CallNode); - RELAX_EXPR_FUNCTOR_DISPATCH(prim::LetNode); - RELAX_EXPR_FUNCTOR_DISPATCH(TensorLoadNode); - RELAX_EXPR_FUNCTOR_DISPATCH(prim::AddNode); - RELAX_EXPR_FUNCTOR_DISPATCH(prim::LShiftNode); - RELAX_EXPR_FUNCTOR_DISPATCH(prim::RShiftNode); - RELAX_EXPR_FUNCTOR_DISPATCH(prim::BitwiseAndNode); - RELAX_EXPR_FUNCTOR_DISPATCH(prim::BitwiseOrNode); - RELAX_EXPR_FUNCTOR_DISPATCH(prim::BitwiseXorNode); - RELAX_EXPR_FUNCTOR_DISPATCH(prim::BitwiseNotNode); - RELAX_EXPR_FUNCTOR_DISPATCH(prim::SubNode); - RELAX_EXPR_FUNCTOR_DISPATCH(prim::MulNode); - RELAX_EXPR_FUNCTOR_DISPATCH(prim::DivNode); - RELAX_EXPR_FUNCTOR_DISPATCH(prim::ModNode); - RELAX_EXPR_FUNCTOR_DISPATCH(prim::FloorDivNode); - RELAX_EXPR_FUNCTOR_DISPATCH(prim::FloorModNode); - RELAX_EXPR_FUNCTOR_DISPATCH(prim::MinNode); - RELAX_EXPR_FUNCTOR_DISPATCH(prim::MaxNode); - RELAX_EXPR_FUNCTOR_DISPATCH(prim::EQNode); - RELAX_EXPR_FUNCTOR_DISPATCH(prim::NENode); - RELAX_EXPR_FUNCTOR_DISPATCH(prim::LTNode); - RELAX_EXPR_FUNCTOR_DISPATCH(prim::LENode); - RELAX_EXPR_FUNCTOR_DISPATCH(prim::GTNode); - RELAX_EXPR_FUNCTOR_DISPATCH(prim::GENode); - RELAX_EXPR_FUNCTOR_DISPATCH(prim::AndNode); - RELAX_EXPR_FUNCTOR_DISPATCH(prim::OrNode); - RELAX_EXPR_FUNCTOR_DISPATCH(prim::CastNode); - RELAX_EXPR_FUNCTOR_DISPATCH(prim::NotNode); - RELAX_EXPR_FUNCTOR_DISPATCH(prim::SelectNode); - RELAX_EXPR_FUNCTOR_DISPATCH(prim::RampNode); - RELAX_EXPR_FUNCTOR_DISPATCH(prim::BroadcastNode); - RELAX_EXPR_FUNCTOR_DISPATCH(prim::ShuffleNode); - RELAX_EXPR_FUNCTOR_DISPATCH(tvm::IntImmNode); - RELAX_EXPR_FUNCTOR_DISPATCH(tvm::FloatImmNode); - RELAX_EXPR_FUNCTOR_DISPATCH(SeqExprNode); - RELAX_EXPR_FUNCTOR_DISPATCH(IfExprNode); - RELAX_EXPR_FUNCTOR_DISPATCH(OpNode); - RELAX_EXPR_FUNCTOR_DISPATCH(TupleGetItemNode); - RELAX_EXPR_FUNCTOR_DISPATCH(StringImmNode); - RELAX_EXPR_FUNCTOR_DISPATCH(DataTypeImmNode); + RELAX_EXPR_FUNCTOR_DISPATCH(GenericConstObj); + RELAX_EXPR_FUNCTOR_DISPATCH(TupleObj); + RELAX_EXPR_FUNCTOR_DISPATCH(VarObj); + RELAX_EXPR_FUNCTOR_DISPATCH(DataflowVarObj); + RELAX_EXPR_FUNCTOR_DISPATCH(ShapeExprObj); + RELAX_EXPR_FUNCTOR_DISPATCH(ExternFuncObj); + RELAX_EXPR_FUNCTOR_DISPATCH(GlobalVarObj); + RELAX_EXPR_FUNCTOR_DISPATCH(FunctionObj); + RELAX_EXPR_FUNCTOR_DISPATCH(CallObj); + RELAX_EXPR_FUNCTOR_DISPATCH(prim::LetObj); + RELAX_EXPR_FUNCTOR_DISPATCH(TensorLoadObj); + RELAX_EXPR_FUNCTOR_DISPATCH(prim::AddObj); + RELAX_EXPR_FUNCTOR_DISPATCH(prim::LShiftObj); + RELAX_EXPR_FUNCTOR_DISPATCH(prim::RShiftObj); + RELAX_EXPR_FUNCTOR_DISPATCH(prim::BitwiseAndObj); + RELAX_EXPR_FUNCTOR_DISPATCH(prim::BitwiseOrObj); + RELAX_EXPR_FUNCTOR_DISPATCH(prim::BitwiseXorObj); + RELAX_EXPR_FUNCTOR_DISPATCH(prim::BitwiseNotObj); + RELAX_EXPR_FUNCTOR_DISPATCH(prim::SubObj); + RELAX_EXPR_FUNCTOR_DISPATCH(prim::MulObj); + RELAX_EXPR_FUNCTOR_DISPATCH(prim::DivObj); + RELAX_EXPR_FUNCTOR_DISPATCH(prim::ModObj); + RELAX_EXPR_FUNCTOR_DISPATCH(prim::FloorDivObj); + RELAX_EXPR_FUNCTOR_DISPATCH(prim::FloorModObj); + RELAX_EXPR_FUNCTOR_DISPATCH(prim::MinObj); + RELAX_EXPR_FUNCTOR_DISPATCH(prim::MaxObj); + RELAX_EXPR_FUNCTOR_DISPATCH(prim::EQObj); + RELAX_EXPR_FUNCTOR_DISPATCH(prim::NEObj); + RELAX_EXPR_FUNCTOR_DISPATCH(prim::LTObj); + RELAX_EXPR_FUNCTOR_DISPATCH(prim::LEObj); + RELAX_EXPR_FUNCTOR_DISPATCH(prim::GTObj); + RELAX_EXPR_FUNCTOR_DISPATCH(prim::GEObj); + RELAX_EXPR_FUNCTOR_DISPATCH(prim::AndObj); + RELAX_EXPR_FUNCTOR_DISPATCH(prim::OrObj); + RELAX_EXPR_FUNCTOR_DISPATCH(prim::CastObj); + RELAX_EXPR_FUNCTOR_DISPATCH(prim::NotObj); + RELAX_EXPR_FUNCTOR_DISPATCH(prim::SelectObj); + RELAX_EXPR_FUNCTOR_DISPATCH(prim::RampObj); + RELAX_EXPR_FUNCTOR_DISPATCH(prim::BroadcastObj); + RELAX_EXPR_FUNCTOR_DISPATCH(prim::ShuffleObj); + RELAX_EXPR_FUNCTOR_DISPATCH(tvm::IntImmObj); + RELAX_EXPR_FUNCTOR_DISPATCH(tvm::FloatImmObj); + RELAX_EXPR_FUNCTOR_DISPATCH(SeqExprObj); + RELAX_EXPR_FUNCTOR_DISPATCH(IfExprObj); + RELAX_EXPR_FUNCTOR_DISPATCH(OpObj); + RELAX_EXPR_FUNCTOR_DISPATCH(TupleGetItemObj); + RELAX_EXPR_FUNCTOR_DISPATCH(StringImmObj); + RELAX_EXPR_FUNCTOR_DISPATCH(DataTypeImmObj); vtable.Finalize(); return vtable; } @@ -270,54 +270,54 @@ class ExprVisitor : public ExprFunctor { */ void VisitExpr(const Expr& expr) override; // specific leaf level visitor functions - void VisitExpr_(const GenericConstNode* op) override; - void VisitExpr_(const TupleNode* op) override; - void VisitExpr_(const VarNode* op) override; - void VisitExpr_(const DataflowVarNode* op) override; - void VisitExpr_(const ShapeExprNode* op) override; - void VisitExpr_(const ExternFuncNode* op) override; - void VisitExpr_(const GlobalVarNode* op) override; - void VisitExpr_(const FunctionNode* op) override; - void VisitExpr_(const CallNode* op) override; - void VisitExpr_(const TensorLoadNode* op) override; - void VisitExpr_(const tvm::IntImmNode* op) override; - void VisitExpr_(const tvm::FloatImmNode* op) override; - void VisitExpr_(const SeqExprNode* op) override; - void VisitExpr_(const IfExprNode* op) override; - void VisitExpr_(const OpNode* op) override; - void VisitExpr_(const TupleGetItemNode* op) override; - void VisitExpr_(const StringImmNode* op) override; - void VisitExpr_(const DataTypeImmNode* op) override; - void VisitExpr_(const prim::AddNode* op) override; - void VisitExpr_(const prim::LShiftNode* op) override; - void VisitExpr_(const prim::RShiftNode* op) override; - void VisitExpr_(const prim::BitwiseAndNode* op) override; - void VisitExpr_(const prim::BitwiseOrNode* op) override; - void VisitExpr_(const prim::BitwiseXorNode* op) override; - void VisitExpr_(const prim::BitwiseNotNode* op) override; - void VisitExpr_(const prim::SubNode* op) override; - void VisitExpr_(const prim::MulNode* op) override; - void VisitExpr_(const prim::DivNode* op) override; - void VisitExpr_(const prim::ModNode* op) override; - void VisitExpr_(const prim::FloorDivNode* op) override; - void VisitExpr_(const prim::FloorModNode* op) override; - void VisitExpr_(const prim::MinNode* op) override; - void VisitExpr_(const prim::MaxNode* op) override; - void VisitExpr_(const prim::EQNode* op) override; - void VisitExpr_(const prim::NENode* op) override; - void VisitExpr_(const prim::LTNode* op) override; - void VisitExpr_(const prim::LENode* op) override; - void VisitExpr_(const prim::GTNode* op) override; - void VisitExpr_(const prim::GENode* op) override; - void VisitExpr_(const prim::AndNode* op) override; - void VisitExpr_(const prim::OrNode* op) override; - void VisitExpr_(const prim::CastNode* op) override; - void VisitExpr_(const prim::NotNode* op) override; - void VisitExpr_(const prim::SelectNode* op) override; - void VisitExpr_(const prim::RampNode* op) override; - void VisitExpr_(const prim::BroadcastNode* op) override; - void VisitExpr_(const prim::ShuffleNode* op) override; - void VisitExprFallback_(const ExprNode* op) override; + void VisitExpr_(const GenericConstObj* op) override; + void VisitExpr_(const TupleObj* op) override; + void VisitExpr_(const VarObj* op) override; + void VisitExpr_(const DataflowVarObj* op) override; + void VisitExpr_(const ShapeExprObj* op) override; + void VisitExpr_(const ExternFuncObj* op) override; + void VisitExpr_(const GlobalVarObj* op) override; + void VisitExpr_(const FunctionObj* op) override; + void VisitExpr_(const CallObj* op) override; + void VisitExpr_(const TensorLoadObj* op) override; + void VisitExpr_(const tvm::IntImmObj* op) override; + void VisitExpr_(const tvm::FloatImmObj* op) override; + void VisitExpr_(const SeqExprObj* op) override; + void VisitExpr_(const IfExprObj* op) override; + void VisitExpr_(const OpObj* op) override; + void VisitExpr_(const TupleGetItemObj* op) override; + void VisitExpr_(const StringImmObj* op) override; + void VisitExpr_(const DataTypeImmObj* op) override; + void VisitExpr_(const prim::AddObj* op) override; + void VisitExpr_(const prim::LShiftObj* op) override; + void VisitExpr_(const prim::RShiftObj* op) override; + void VisitExpr_(const prim::BitwiseAndObj* op) override; + void VisitExpr_(const prim::BitwiseOrObj* op) override; + void VisitExpr_(const prim::BitwiseXorObj* op) override; + void VisitExpr_(const prim::BitwiseNotObj* op) override; + void VisitExpr_(const prim::SubObj* op) override; + void VisitExpr_(const prim::MulObj* op) override; + void VisitExpr_(const prim::DivObj* op) override; + void VisitExpr_(const prim::ModObj* op) override; + void VisitExpr_(const prim::FloorDivObj* op) override; + void VisitExpr_(const prim::FloorModObj* op) override; + void VisitExpr_(const prim::MinObj* op) override; + void VisitExpr_(const prim::MaxObj* op) override; + void VisitExpr_(const prim::EQObj* op) override; + void VisitExpr_(const prim::NEObj* op) override; + void VisitExpr_(const prim::LTObj* op) override; + void VisitExpr_(const prim::LEObj* op) override; + void VisitExpr_(const prim::GTObj* op) override; + void VisitExpr_(const prim::GEObj* op) override; + void VisitExpr_(const prim::AndObj* op) override; + void VisitExpr_(const prim::OrObj* op) override; + void VisitExpr_(const prim::CastObj* op) override; + void VisitExpr_(const prim::NotObj* op) override; + void VisitExpr_(const prim::SelectObj* op) override; + void VisitExpr_(const prim::RampObj* op) override; + void VisitExpr_(const prim::BroadcastObj* op) override; + void VisitExpr_(const prim::ShuffleObj* op) override; + void VisitExprFallback_(const ExprObj* op) override; /*! * \brief Generic dispatcher for bindings. @@ -325,39 +325,39 @@ class ExprVisitor : public ExprFunctor { */ virtual void VisitBinding(const Binding& binding); // specific leaf level visitor functions - virtual void VisitBinding_(const VarBindingNode* binding); - virtual void VisitBinding_(const MatchCastNode* binding); + virtual void VisitBinding_(const VarBindingObj* binding); + virtual void VisitBinding_(const MatchCastObj* binding); // second level dispatching based on binding value type. // these dispatching functions get called from first-level dispatch on VarBinding - virtual void VisitBinding_(const VarBindingNode* binding, const GenericConstNode* val); - virtual void VisitBinding_(const VarBindingNode* binding, const TupleNode* val); - virtual void VisitBinding_(const VarBindingNode* binding, const VarNode* val); - virtual void VisitBinding_(const VarBindingNode* binding, const DataflowVarNode* val); - virtual void VisitBinding_(const VarBindingNode* binding, const ShapeExprNode* val); - virtual void VisitBinding_(const VarBindingNode* binding, const ExternFuncNode* val); - virtual void VisitBinding_(const VarBindingNode* binding, const GlobalVarNode* val); - virtual void VisitBinding_(const VarBindingNode* binding, const FunctionNode* val); - virtual void VisitBinding_(const VarBindingNode* binding, const CallNode* val); - virtual void VisitBinding_(const VarBindingNode* binding, const SeqExprNode* val); - virtual void VisitBinding_(const VarBindingNode* binding, const IfExprNode* val); - virtual void VisitBinding_(const VarBindingNode* binding, const OpNode* val); - virtual void VisitBinding_(const VarBindingNode* binding, const TupleGetItemNode* val); - virtual void VisitBinding_(const VarBindingNode* binding, const ExprNode* val); - virtual void VisitBinding_(const VarBindingNode* binding, const StringImmNode* val); - virtual void VisitBinding_(const VarBindingNode* binding, const DataTypeImmNode* val); + virtual void VisitBinding_(const VarBindingObj* binding, const GenericConstObj* val); + virtual void VisitBinding_(const VarBindingObj* binding, const TupleObj* val); + virtual void VisitBinding_(const VarBindingObj* binding, const VarObj* val); + virtual void VisitBinding_(const VarBindingObj* binding, const DataflowVarObj* val); + virtual void VisitBinding_(const VarBindingObj* binding, const ShapeExprObj* val); + virtual void VisitBinding_(const VarBindingObj* binding, const ExternFuncObj* val); + virtual void VisitBinding_(const VarBindingObj* binding, const GlobalVarObj* val); + virtual void VisitBinding_(const VarBindingObj* binding, const FunctionObj* val); + virtual void VisitBinding_(const VarBindingObj* binding, const CallObj* val); + virtual void VisitBinding_(const VarBindingObj* binding, const SeqExprObj* val); + virtual void VisitBinding_(const VarBindingObj* binding, const IfExprObj* val); + virtual void VisitBinding_(const VarBindingObj* binding, const OpObj* val); + virtual void VisitBinding_(const VarBindingObj* binding, const TupleGetItemObj* val); + virtual void VisitBinding_(const VarBindingObj* binding, const ExprObj* val); + virtual void VisitBinding_(const VarBindingObj* binding, const StringImmObj* val); + virtual void VisitBinding_(const VarBindingObj* binding, const DataTypeImmObj* val); /*! * \brief Generic dispatcher for binding blocks. * \param block The binding block to be visited. */ virtual void VisitBindingBlock(const BindingBlock& block); // specific leaf level visitor functions - virtual void VisitBindingBlock_(const BindingBlockNode* block); - virtual void VisitBindingBlock_(const DataflowBlockNode* block); + virtual void VisitBindingBlock_(const BindingBlockObj* block); + virtual void VisitBindingBlock_(const DataflowBlockObj* block); /*! * \brief Generic dispatcher for visiting the var definition site. * \param var The var to be visited. - * \note VisitExpr_(const VarNode*) will only visit the usage site of an Var + * \note VisitExpr_(const VarObj*) will only visit the usage site of an Var */ virtual void VisitVarDef(const Var& var); @@ -379,8 +379,8 @@ class ExprVisitor : public ExprFunctor { virtual void VisitExprDepTypeField(const Type& ty); // specific leaf level visitor functions - virtual void VisitVarDef_(const VarNode* var); - virtual void VisitVarDef_(const DataflowVarNode* var); + virtual void VisitVarDef_(const VarObj* var); + virtual void VisitVarDef_(const DataflowVarObj* var); virtual void VisitLoc(const Location& loc); virtual void VisitTypePrimExprField(const PrimExpr& expr); @@ -388,7 +388,7 @@ class ExprVisitor : public ExprFunctor { private: using TSelf = ExprVisitor; using VisitBindingVTable = tvm::ObjectFunctor; + const VarBindingObj* binding)>; // initialize the vtable. static VisitBindingVTable InitVisitBindingVTable(); /*! @@ -407,7 +407,7 @@ class ExprVisitor : public ExprFunctor { // Override defaults in type visitor. void VisitTypeExprField(const Expr& expr) final; void VisitTypeExprField(const PrimExpr& expr) final; - void VisitType_(const FuncTypeNode* op) final; + void VisitType_(const FuncTypeObj* op) final; private: ExprVisitor* parent_; @@ -430,54 +430,54 @@ void PostOrderVisit(const Expr& node, std::function fvisit); class ExprMutatorBase : public ExprFunctor { public: Expr VisitExpr(const Expr& expr) override; - Expr VisitExpr_(const GenericConstNode* op) override; - Expr VisitExpr_(const TupleNode* op) override; - Expr VisitExpr_(const VarNode* op) override; - Expr VisitExpr_(const DataflowVarNode* op) override; - Expr VisitExpr_(const ShapeExprNode* op) override; - Expr VisitExpr_(const ExternFuncNode* op) override; - Expr VisitExpr_(const GlobalVarNode* op) override; - Expr VisitExpr_(const FunctionNode* op) override; - Expr VisitExpr_(const CallNode* op) override; - Expr VisitExpr_(const TensorLoadNode* op) override; - Expr VisitExpr_(const tvm::IntImmNode* op) override; - Expr VisitExpr_(const tvm::FloatImmNode* op) override; - Expr VisitExpr_(const SeqExprNode* op) override; - Expr VisitExpr_(const IfExprNode* op) override; - Expr VisitExpr_(const OpNode* op) override; - Expr VisitExpr_(const TupleGetItemNode* op) override; - Expr VisitExpr_(const StringImmNode* op) override; - Expr VisitExpr_(const DataTypeImmNode* op) override; - Expr VisitExpr_(const prim::AddNode* op) override; - Expr VisitExpr_(const prim::LShiftNode* op) override; - Expr VisitExpr_(const prim::RShiftNode* op) override; - Expr VisitExpr_(const prim::BitwiseAndNode* op) override; - Expr VisitExpr_(const prim::BitwiseOrNode* op) override; - Expr VisitExpr_(const prim::BitwiseXorNode* op) override; - Expr VisitExpr_(const prim::BitwiseNotNode* op) override; - Expr VisitExpr_(const prim::SubNode* op) override; - Expr VisitExpr_(const prim::MulNode* op) override; - Expr VisitExpr_(const prim::DivNode* op) override; - Expr VisitExpr_(const prim::ModNode* op) override; - Expr VisitExpr_(const prim::FloorDivNode* op) override; - Expr VisitExpr_(const prim::FloorModNode* op) override; - Expr VisitExpr_(const prim::MinNode* op) override; - Expr VisitExpr_(const prim::MaxNode* op) override; - Expr VisitExpr_(const prim::EQNode* op) override; - Expr VisitExpr_(const prim::NENode* op) override; - Expr VisitExpr_(const prim::LTNode* op) override; - Expr VisitExpr_(const prim::LENode* op) override; - Expr VisitExpr_(const prim::GTNode* op) override; - Expr VisitExpr_(const prim::GENode* op) override; - Expr VisitExpr_(const prim::AndNode* op) override; - Expr VisitExpr_(const prim::OrNode* op) override; - Expr VisitExpr_(const prim::CastNode* op) override; - Expr VisitExpr_(const prim::NotNode* op) override; - Expr VisitExpr_(const prim::SelectNode* op) override; - Expr VisitExpr_(const prim::RampNode* op) override; - Expr VisitExpr_(const prim::BroadcastNode* op) override; - Expr VisitExpr_(const prim::ShuffleNode* op) override; - Expr VisitExprFallback_(const ExprNode* op) override; + Expr VisitExpr_(const GenericConstObj* op) override; + Expr VisitExpr_(const TupleObj* op) override; + Expr VisitExpr_(const VarObj* op) override; + Expr VisitExpr_(const DataflowVarObj* op) override; + Expr VisitExpr_(const ShapeExprObj* op) override; + Expr VisitExpr_(const ExternFuncObj* op) override; + Expr VisitExpr_(const GlobalVarObj* op) override; + Expr VisitExpr_(const FunctionObj* op) override; + Expr VisitExpr_(const CallObj* op) override; + Expr VisitExpr_(const TensorLoadObj* op) override; + Expr VisitExpr_(const tvm::IntImmObj* op) override; + Expr VisitExpr_(const tvm::FloatImmObj* op) override; + Expr VisitExpr_(const SeqExprObj* op) override; + Expr VisitExpr_(const IfExprObj* op) override; + Expr VisitExpr_(const OpObj* op) override; + Expr VisitExpr_(const TupleGetItemObj* op) override; + Expr VisitExpr_(const StringImmObj* op) override; + Expr VisitExpr_(const DataTypeImmObj* op) override; + Expr VisitExpr_(const prim::AddObj* op) override; + Expr VisitExpr_(const prim::LShiftObj* op) override; + Expr VisitExpr_(const prim::RShiftObj* op) override; + Expr VisitExpr_(const prim::BitwiseAndObj* op) override; + Expr VisitExpr_(const prim::BitwiseOrObj* op) override; + Expr VisitExpr_(const prim::BitwiseXorObj* op) override; + Expr VisitExpr_(const prim::BitwiseNotObj* op) override; + Expr VisitExpr_(const prim::SubObj* op) override; + Expr VisitExpr_(const prim::MulObj* op) override; + Expr VisitExpr_(const prim::DivObj* op) override; + Expr VisitExpr_(const prim::ModObj* op) override; + Expr VisitExpr_(const prim::FloorDivObj* op) override; + Expr VisitExpr_(const prim::FloorModObj* op) override; + Expr VisitExpr_(const prim::MinObj* op) override; + Expr VisitExpr_(const prim::MaxObj* op) override; + Expr VisitExpr_(const prim::EQObj* op) override; + Expr VisitExpr_(const prim::NEObj* op) override; + Expr VisitExpr_(const prim::LTObj* op) override; + Expr VisitExpr_(const prim::LEObj* op) override; + Expr VisitExpr_(const prim::GTObj* op) override; + Expr VisitExpr_(const prim::GEObj* op) override; + Expr VisitExpr_(const prim::AndObj* op) override; + Expr VisitExpr_(const prim::OrObj* op) override; + Expr VisitExpr_(const prim::CastObj* op) override; + Expr VisitExpr_(const prim::NotObj* op) override; + Expr VisitExpr_(const prim::SelectObj* op) override; + Expr VisitExpr_(const prim::RampObj* op) override; + Expr VisitExpr_(const prim::BroadcastObj* op) override; + Expr VisitExpr_(const prim::ShuffleObj* op) override; + Expr VisitExprFallback_(const ExprObj* op) override; /*! * \brief Mutate BindingBlock. @@ -522,7 +522,7 @@ class ExprMutatorBase : public ExprFunctor { * node without ty, and trigger normalizer to re-derive. */ bool VisitAndCheckTypeFieldUnchanged(const ffi::ObjectRef& ty) { - if (const TypeNode* ty_node = ty.as()) { + if (const TypeObj* ty_node = ty.as()) { Type type = ffi::GetRef(ty_node); return type.as().has_value() || this->VisitExprDepTypeField(type).same_as(ty); } else { @@ -545,7 +545,7 @@ class ExprMutatorBase : public ExprFunctor { // Override defaults in type visitor. Expr VisitTypeExprField(const Expr& expr) final; PrimExpr VisitTypeExprField(const PrimExpr& expr) final; - Type VisitType_(const FuncTypeNode* op) final; + Type VisitType_(const FuncTypeObj* op) final; private: ExprMutatorBase* parent_; @@ -568,11 +568,11 @@ class ExprMutator : public ExprMutatorBase { ExprMutator(ffi::Optional mod = std::nullopt) : builder_(BlockBuilder::Create(mod)) {} Expr VisitExpr(const Expr& expr) override; - Expr VisitExpr_(const VarNode* op) override; - Expr VisitExpr_(const DataflowVarNode* op) override; - Expr VisitExpr_(const FunctionNode* op) override; - Expr VisitExpr_(const SeqExprNode* op) override; - Expr VisitExpr_(const IfExprNode* op) override; + Expr VisitExpr_(const VarObj* op) override; + Expr VisitExpr_(const DataflowVarObj* op) override; + Expr VisitExpr_(const FunctionObj* op) override; + Expr VisitExpr_(const SeqExprObj* op) override; + Expr VisitExpr_(const IfExprObj* op) override; /*! * \brief Generic dispatcher for bindings. @@ -580,26 +580,26 @@ class ExprMutator : public ExprMutatorBase { */ virtual void VisitBinding(const Binding& binding); // specific leaf level visitor functions - virtual void VisitBinding_(const VarBindingNode* binding); - virtual void VisitBinding_(const MatchCastNode* binding); + virtual void VisitBinding_(const VarBindingObj* binding); + virtual void VisitBinding_(const MatchCastObj* binding); // second level dispatching based on binding value type. // these dispatching functions get called from first-level dispatch on VarBinding - virtual void VisitBinding_(const VarBindingNode* binding, const GenericConstNode* val); - virtual void VisitBinding_(const VarBindingNode* binding, const TupleNode* val); - virtual void VisitBinding_(const VarBindingNode* binding, const VarNode* val); - virtual void VisitBinding_(const VarBindingNode* binding, const DataflowVarNode* val); - virtual void VisitBinding_(const VarBindingNode* binding, const ShapeExprNode* val); - virtual void VisitBinding_(const VarBindingNode* binding, const ExternFuncNode* val); - virtual void VisitBinding_(const VarBindingNode* binding, const GlobalVarNode* val); - virtual void VisitBinding_(const VarBindingNode* binding, const FunctionNode* val); - virtual void VisitBinding_(const VarBindingNode* binding, const CallNode* val); - virtual void VisitBinding_(const VarBindingNode* binding, const SeqExprNode* val); - virtual void VisitBinding_(const VarBindingNode* binding, const IfExprNode* val); - virtual void VisitBinding_(const VarBindingNode* binding, const OpNode* val); - virtual void VisitBinding_(const VarBindingNode* binding, const TupleGetItemNode* val); - virtual void VisitBinding_(const VarBindingNode* binding, const ExprNode* val); - virtual void VisitBinding_(const VarBindingNode* binding, const StringImmNode* val); - virtual void VisitBinding_(const VarBindingNode* binding, const DataTypeImmNode* val); + virtual void VisitBinding_(const VarBindingObj* binding, const GenericConstObj* val); + virtual void VisitBinding_(const VarBindingObj* binding, const TupleObj* val); + virtual void VisitBinding_(const VarBindingObj* binding, const VarObj* val); + virtual void VisitBinding_(const VarBindingObj* binding, const DataflowVarObj* val); + virtual void VisitBinding_(const VarBindingObj* binding, const ShapeExprObj* val); + virtual void VisitBinding_(const VarBindingObj* binding, const ExternFuncObj* val); + virtual void VisitBinding_(const VarBindingObj* binding, const GlobalVarObj* val); + virtual void VisitBinding_(const VarBindingObj* binding, const FunctionObj* val); + virtual void VisitBinding_(const VarBindingObj* binding, const CallObj* val); + virtual void VisitBinding_(const VarBindingObj* binding, const SeqExprObj* val); + virtual void VisitBinding_(const VarBindingObj* binding, const IfExprObj* val); + virtual void VisitBinding_(const VarBindingObj* binding, const OpObj* val); + virtual void VisitBinding_(const VarBindingObj* binding, const TupleGetItemObj* val); + virtual void VisitBinding_(const VarBindingObj* binding, const ExprObj* val); + virtual void VisitBinding_(const VarBindingObj* binding, const StringImmObj* val); + virtual void VisitBinding_(const VarBindingObj* binding, const DataTypeImmObj* val); /*! * \brief Generic dispatcher for binding blocks. * \param block The binding block to be visited. @@ -607,19 +607,19 @@ class ExprMutator : public ExprMutatorBase { */ virtual BindingBlock VisitBindingBlock(const BindingBlock& block) override; // NOLINT(*) // specific leaf level visitor functions - virtual BindingBlock VisitBindingBlock_(const BindingBlockNode* block); - virtual BindingBlock VisitBindingBlock_(const DataflowBlockNode* block); + virtual BindingBlock VisitBindingBlock_(const BindingBlockObj* block); + virtual BindingBlock VisitBindingBlock_(const DataflowBlockObj* block); /*! * \brief Generic dispatcher for rewriting the var definition site. * \param var The var to be visited. * \return The var after post-order rewritten. - * \note VisitExpr_(const VarNode*) will only visit the usage site of an Var + * \note VisitExpr_(const VarObj*) will only visit the usage site of an Var */ virtual Var VisitVarDef(const Var& var); // specific leaf level visitor functions - virtual Var VisitVarDef_(const VarNode* var); - virtual Var VisitVarDef_(const DataflowVarNode* var); + virtual Var VisitVarDef_(const VarObj* var); + virtual Var VisitVarDef_(const DataflowVarObj* var); protected: /*! @@ -633,7 +633,7 @@ class ExprMutator : public ExprMutatorBase { * Otherwise, a new binding will be emitted to replace the var specified in * the current binding. */ - void ReEmitBinding(const VarBindingNode* binding, Expr new_value); + void ReEmitBinding(const VarBindingObj* binding, Expr new_value); /*! * \brief Rewrite the expr with a new scope, used in a Function's body. @@ -703,7 +703,7 @@ class ExprMutator : public ExprMutatorBase { private: using TSelf = ExprMutator; using VisitBindingVTable = tvm::ObjectFunctor; + const VarBindingObj* binding)>; // initialize the vtable. static VisitBindingVTable InitVisitBindingVTable(); }; diff --git a/include/tvm/relax/global_info.h b/include/tvm/relax/global_info.h index 6ffc691822a8..416a46fff3b8 100644 --- a/include/tvm/relax/global_info.h +++ b/include/tvm/relax/global_info.h @@ -40,7 +40,7 @@ using runtime::MemoryScope; /*! * \brief A global info subclass for virtual devices. */ -class VDeviceNode : public GlobalInfoNode { +class VDeviceObj : public GlobalInfoObj { public: /*! \brief The \p Target describing how to compile for the virtual device. */ ffi::Optional target; @@ -52,49 +52,49 @@ class VDeviceNode : public GlobalInfoNode { static void RegisterReflection() { namespace refl = tvm::ffi::reflection; - refl::ObjectDef() - .def_ro("target", &VDeviceNode::target) - .def_ro("vdevice_id", &VDeviceNode::vdevice_id) - .def_ro("memory_scope", &VDeviceNode::memory_scope); + refl::ObjectDef() + .def_ro("target", &VDeviceObj::target) + .def_ro("vdevice_id", &VDeviceObj::vdevice_id) + .def_ro("memory_scope", &VDeviceObj::memory_scope); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.VDevice", VDeviceNode, GlobalInfoNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.VDevice", VDeviceObj, GlobalInfoObj); }; /*! - * \brief Managed reference to VDeviceNode. - * \sa VDeviceNode + * \brief Managed reference to VDeviceObj. + * \sa VDeviceObj */ class VDevice : public GlobalInfo { public: TVM_DLL explicit VDevice(ffi::Optional tgt, int dev_id, MemoryScope mem_scope); - TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(VDevice, GlobalInfo, VDeviceNode); + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(VDevice, GlobalInfo, VDeviceObj); }; /*! * \brief A dummy global info sub-class for testing purpose. */ -class DummyGlobalInfoNode : public GlobalInfoNode { +class DummyGlobalInfoObj : public GlobalInfoObj { public: static void RegisterReflection() { namespace refl = tvm::ffi::reflection; - refl::ObjectDef(); + refl::ObjectDef(); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.DummyGlobalInfo", DummyGlobalInfoNode, GlobalInfoNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.DummyGlobalInfo", DummyGlobalInfoObj, GlobalInfoObj); }; /*! - * \brief Managed reference to DummyGlobalInfoNode. - * \sa DummyGlobalInfoNode + * \brief Managed reference to DummyGlobalInfoObj. + * \sa DummyGlobalInfoObj */ class DummyGlobalInfo : public GlobalInfo { public: - explicit DummyGlobalInfo(ffi::UnsafeInit tag, ffi::ObjectPtr n) + explicit DummyGlobalInfo(ffi::UnsafeInit tag, ffi::ObjectPtr n) : GlobalInfo(tag) { data_ = std::move(n); } - TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(DummyGlobalInfo, GlobalInfo, DummyGlobalInfoNode); + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(DummyGlobalInfo, GlobalInfo, DummyGlobalInfoObj); }; } // namespace relax diff --git a/include/tvm/relax/nested_msg.h b/include/tvm/relax/nested_msg.h index ea2f5e4a7aec..3c39873f4394 100644 --- a/include/tvm/relax/nested_msg.h +++ b/include/tvm/relax/nested_msg.h @@ -183,7 +183,7 @@ class NestedMsg { */ T LeafValue() const { TVM_FFI_ICHECK(IsLeaf()); - return ffi::details::AnyUnsafe::CopyFromAnyViewAfterCheck(data_); + return ffi::detail::AnyUnsafe::CopyFromAnyViewAfterCheck(data_); } /*! @@ -191,8 +191,7 @@ class NestedMsg { * \note This checks if the underlying data type is array. */ ffi::Array, void> NestedArray() const { - return ffi::details::AnyUnsafe::CopyFromAnyViewAfterCheck, void>>( - data_); + return ffi::detail::AnyUnsafe::CopyFromAnyViewAfterCheck, void>>(data_); } private: @@ -264,7 +263,7 @@ bool Equal(const NestedMsg& lhs, const NestedMsg& rhs, FType fequal) { */ template NestedMsg MapToNestedMsg(Expr expr, FType fmapleaf) { - if (auto* tuple = expr.as()) { + if (auto* tuple = expr.as()) { ffi::Array> res; res.reserve(tuple->fields.size()); for (Expr x : tuple->fields) { @@ -291,7 +290,7 @@ NestedMsg MapToNestedMsg(Expr expr, FType fmapleaf) { */ template NestedMsg MapToNestedMsg(Type ty, FType fmapleaf) { - if (auto* tuple = ty.as()) { + if (auto* tuple = ty.as()) { ffi::Array> res; res.reserve(tuple->fields.size()); for (Type x : tuple->fields) { @@ -320,12 +319,12 @@ NestedMsg MapToNestedMsg(Type ty, FType fmapleaf) { template NestedMsg MapToNestedMsgByType(Expr expr, FType fmapleaf) { auto ty = GetType(expr); - if (auto* tuple = ty.as()) { + if (auto* tuple = ty.as()) { ffi::Array> res; res.reserve(tuple->fields.size()); for (size_t i = 0; i < tuple->fields.size(); ++i) { Expr field = [&]() -> Expr { - if (const auto* expr_tuple = expr.as()) { + if (const auto* expr_tuple = expr.as()) { return expr_tuple->fields[i]; } return TupleGetItem(expr, i); @@ -393,7 +392,7 @@ Expr NestedMsgToExpr(NestedMsg msg, FType fmapleaf) { if (arr.size() >= 1) { simplified_flag = true; for (size_t i = 0; i < arr.size() && simplified_flag; ++i) { - auto* node = arr[i].as(); + auto* node = arr[i].as(); if (node == nullptr || node->index != static_cast(i)) { simplified_flag = false; } else { @@ -491,7 +490,7 @@ NestedMsg MapNestedMsg(NestedMsg msg, FType fmapleaf) { */ template void DecomposeNestedMsg(Expr expr, NestedMsg msg, FType fvisitleaf) { - if (auto* tuple = expr.as()) { + if (auto* tuple = expr.as()) { TVM_FFI_ICHECK(msg.IsNested()) << "Expected nested to match tuple"; ffi::Array> arr = msg.NestedArray(); TVM_FFI_ICHECK_EQ(arr.size(), tuple->fields.size()) @@ -521,7 +520,7 @@ void DecomposeNestedMsg(Expr expr, NestedMsg msg, FType fvisitleaf) { template Expr TransformTupleLeaf(Expr expr, std::array, N> msgs, FType ftransleaf) { Type ty = GetType(expr); - if (const auto* tuple = ty.as()) { + if (const auto* tuple = ty.as()) { std::array>, N> msg_arrays; for (size_t i = 0; i < N; ++i) { TVM_FFI_ICHECK(msgs[i].IsNested()) << "Expected nested to match tuple"; @@ -532,7 +531,7 @@ Expr TransformTupleLeaf(Expr expr, std::array, N> msgs, FType ftran fields.reserve(tuple->fields.size()); for (size_t i = 0; i < tuple->fields.size(); ++i) { Expr field = [&]() -> Expr { - if (const auto* expr_tuple = expr.as()) { + if (const auto* expr_tuple = expr.as()) { return expr_tuple->fields[i]; } return TupleGetItem(expr, i); @@ -569,7 +568,7 @@ Expr TransformTupleLeaf(Expr expr, std::array, N> msgs, FType ftran */ template Type TransformTupleLeaf(Type ty, std::array, N> msgs, FType ftransleaf) { - if (const auto* tuple = ty.as()) { + if (const auto* tuple = ty.as()) { std::array>, N> msg_arrays; for (size_t i = 0; i < N; ++i) { TVM_FFI_ICHECK(msgs[i].IsNested()) << "Expected nested to match tuple"; diff --git a/include/tvm/relax/op/ccl.h b/include/tvm/relax/op/ccl.h index a0c986e23721..55be6043c7e6 100644 --- a/include/tvm/relax/op/ccl.h +++ b/include/tvm/relax/op/ccl.h @@ -31,31 +31,31 @@ namespace tvm { namespace relax { /*! \brief Attributes used in allreduce operators */ -struct AllReduceAttrs : public tvm::AttrsNode { +struct AllReduceAttrs : public tvm::AttrsObj { ffi::String op_type; bool in_group; static void RegisterReflection(); - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.AllReduceAttrs", AllReduceAttrs, AttrsNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.AllReduceAttrs", AllReduceAttrs, AttrsObj); }; // struct AllReduceAttrs /*! \brief Attributes used in allgather operators */ -struct AllGatherAttrs : public tvm::AttrsNode { +struct AllGatherAttrs : public tvm::AttrsObj { int num_workers; bool in_group; static void RegisterReflection(); - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.AllGatherAttrs", AllGatherAttrs, AttrsNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.AllGatherAttrs", AllGatherAttrs, AttrsObj); }; // struct AllGatherAttrs /*! \brief Attributes used in scatter operators */ -struct ScatterCollectiveAttrs : public tvm::AttrsNode { +struct ScatterCollectiveAttrs : public tvm::AttrsObj { int num_workers; int axis; static void RegisterReflection(); TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.ScatterCollectiveAttrs", ScatterCollectiveAttrs, - AttrsNode); + AttrsObj); }; // struct ScatterCollectiveAttrs } // namespace relax diff --git a/include/tvm/relax/op/create.h b/include/tvm/relax/op/create.h index 9b4d0a50620e..aa6fbb4a1483 100644 --- a/include/tvm/relax/op/create.h +++ b/include/tvm/relax/op/create.h @@ -30,19 +30,19 @@ namespace tvm { namespace relax { /*! \brief Attributes used in full/full_like, ones/ones_like, and zeros/zeros_like operators */ -struct InitAttrs : public AttrsNode { +struct InitAttrs : public AttrsObj { ffi::Optional dtype; static void RegisterReflection(); - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.InitAttrs", InitAttrs, AttrsNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.InitAttrs", InitAttrs, AttrsObj); }; // struct InitAttrs /*! \brief Attributes used in tril and triu operator */ -struct TriluAttrs : public AttrsNode { +struct TriluAttrs : public AttrsObj { int k; static void RegisterReflection(); - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.TriluAttrs", TriluAttrs, AttrsNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.TriluAttrs", TriluAttrs, AttrsObj); }; // struct TriluAttrs } // namespace relax diff --git a/include/tvm/relax/op/datatype.h b/include/tvm/relax/op/datatype.h index 9c04904acd65..7817e8247ec5 100644 --- a/include/tvm/relax/op/datatype.h +++ b/include/tvm/relax/op/datatype.h @@ -30,19 +30,19 @@ namespace tvm { namespace relax { /*! \brief Attributes used in astype operator */ -struct AstypeAttrs : public AttrsNode { +struct AstypeAttrs : public AttrsObj { DLDataType dtype; static void RegisterReflection(); - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.AstypeAttrs", AstypeAttrs, AttrsNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.AstypeAttrs", AstypeAttrs, AttrsObj); }; // struct AstypeAttrs. /*! \brief Attributes used in wrap_param operator */ -struct WrapParamAttrs : public AttrsNode { +struct WrapParamAttrs : public AttrsObj { DLDataType dtype; static void RegisterReflection(); - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.WrapParamAttrs", WrapParamAttrs, AttrsNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.WrapParamAttrs", WrapParamAttrs, AttrsObj); }; // struct WrapParamAttrs. } // namespace relax diff --git a/include/tvm/relax/op/distributed.h b/include/tvm/relax/op/distributed.h index 7ad1f171099c..92b368c43175 100644 --- a/include/tvm/relax/op/distributed.h +++ b/include/tvm/relax/op/distributed.h @@ -32,7 +32,7 @@ namespace tvm { namespace relax { /*! \brief Attributes for redistribute and annotate_sharding operator */ -struct DistributionAttrs : public AttrsNode { +struct DistributionAttrs : public AttrsObj { DistributionAttrs(distributed::DeviceMesh mesh, distributed::Placement place) : device_mesh(std::move(mesh)), placement(std::move(place)) {} explicit DistributionAttrs(ffi::UnsafeInit) @@ -41,7 +41,7 @@ struct DistributionAttrs : public AttrsNode { distributed::Placement placement; static void RegisterReflection(); - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.DistributionAttrs", DistributionAttrs, AttrsNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.DistributionAttrs", DistributionAttrs, AttrsObj); }; // struct DistributionAttrs } // namespace relax diff --git a/include/tvm/relax/op/image.h b/include/tvm/relax/op/image.h index 9aa1beb98bf2..b3b2283f3740 100644 --- a/include/tvm/relax/op/image.h +++ b/include/tvm/relax/op/image.h @@ -30,7 +30,7 @@ namespace tvm { namespace relax { /*! \brief Attributes used in image resize2d operator */ -struct Resize2DAttrs : public AttrsNode { +struct Resize2DAttrs : public AttrsObj { ffi::Array roi; ffi::String layout; ffi::String method; @@ -42,11 +42,11 @@ struct Resize2DAttrs : public AttrsNode { ffi::Optional out_dtype; static void RegisterReflection(); - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.Resize2DAttrs", Resize2DAttrs, AttrsNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.Resize2DAttrs", Resize2DAttrs, AttrsObj); }; // struct Resize2dAttrs /*! \brief Attributes used in image resize3d operator */ -struct Resize3DAttrs : public AttrsNode { +struct Resize3DAttrs : public AttrsObj { ffi::Array roi; ffi::String layout; ffi::String method; @@ -58,26 +58,26 @@ struct Resize3DAttrs : public AttrsNode { ffi::Optional out_dtype; static void RegisterReflection(); - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.Resize3DAttrs", Resize3DAttrs, AttrsNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.Resize3DAttrs", Resize3DAttrs, AttrsObj); }; // struct Resize3DAttrs /*! \brief Attributes used in image grid_sample operator */ -struct GridSampleAttrs : public AttrsNode { +struct GridSampleAttrs : public AttrsObj { ffi::String method; ffi::String layout; ffi::String padding_mode; bool align_corners; static void RegisterReflection(); - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.GridSampleAttrs", GridSampleAttrs, AttrsNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.GridSampleAttrs", GridSampleAttrs, AttrsObj); }; // struct GridSampleAttrs /*! \brief Attributes used in image affine_grid operator */ -struct AffineGridAttrs : public AttrsNode { +struct AffineGridAttrs : public AttrsObj { bool align_corners; static void RegisterReflection(); - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.AffineGridAttrs", AffineGridAttrs, AttrsNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.AffineGridAttrs", AffineGridAttrs, AttrsObj); }; // struct AffineGridAttrs } // namespace relax diff --git a/include/tvm/relax/op/index.h b/include/tvm/relax/op/index.h index adff401a8bd9..253534989ee1 100644 --- a/include/tvm/relax/op/index.h +++ b/include/tvm/relax/op/index.h @@ -30,28 +30,28 @@ namespace tvm { namespace relax { /*! \brief Attributes used in take operator */ -struct TakeAttrs : public AttrsNode { +struct TakeAttrs : public AttrsObj { ffi::Optional axis; ffi::String mode; static void RegisterReflection(); - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.TakeAttrs", TakeAttrs, AttrsNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.TakeAttrs", TakeAttrs, AttrsObj); }; // struct TakeAttrs /*! \brief Attributes used in take_backward operator */ -struct TakeBackwardAttrs : public AttrsNode { +struct TakeBackwardAttrs : public AttrsObj { ffi::Optional axis; static void RegisterReflection(); - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.TakeBackwardAttrs", TakeBackwardAttrs, AttrsNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.TakeBackwardAttrs", TakeBackwardAttrs, AttrsObj); }; /*! \brief Attributes used in strided_slice operator */ -struct StridedSliceAttrs : public AttrsNode { +struct StridedSliceAttrs : public AttrsObj { bool assume_inbound; static void RegisterReflection(); - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.StridedSliceAttrs", StridedSliceAttrs, AttrsNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.StridedSliceAttrs", StridedSliceAttrs, AttrsObj); }; // struct StridedSliceAttrs } // namespace relax diff --git a/include/tvm/relax/op/linear_algebra.h b/include/tvm/relax/op/linear_algebra.h index 75415b2f1a50..2c472f8aeaa2 100644 --- a/include/tvm/relax/op/linear_algebra.h +++ b/include/tvm/relax/op/linear_algebra.h @@ -30,19 +30,19 @@ namespace tvm { namespace relax { /*! \brief Attributes for matmul operator */ -struct MatmulAttrs : public AttrsNode { +struct MatmulAttrs : public AttrsObj { ffi::Optional out_dtype; static void RegisterReflection(); - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.MatmulAttrs", MatmulAttrs, AttrsNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.MatmulAttrs", MatmulAttrs, AttrsObj); }; // struct MatmulAttrs /*! \brief Attributes used in einsum operator */ -struct EinsumAttrs : public AttrsNode { +struct EinsumAttrs : public AttrsObj { ffi::String subscripts; static void RegisterReflection(); - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.EinsumAttrs", EinsumAttrs, AttrsNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.EinsumAttrs", EinsumAttrs, AttrsObj); }; // struct EinsumAttrs } // namespace relax diff --git a/include/tvm/relax/op/manipulate.h b/include/tvm/relax/op/manipulate.h index b8a40bcee924..69f581d99bea 100644 --- a/include/tvm/relax/op/manipulate.h +++ b/include/tvm/relax/op/manipulate.h @@ -31,23 +31,23 @@ namespace tvm { namespace relax { /*! \brief Attributes used in concat operators */ -struct ConcatAttrs : public AttrsNode { +struct ConcatAttrs : public AttrsObj { ffi::Optional axis; static void RegisterReflection(); - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.ConcatAttrs", ConcatAttrs, AttrsNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.ConcatAttrs", ConcatAttrs, AttrsObj); }; // struct ConcatAttrs /*! \brief Attributes used in expand_dims operators */ -struct ExpandDimsAttrs : public AttrsNode { +struct ExpandDimsAttrs : public AttrsObj { ffi::Array axis; static void RegisterReflection(); - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.ExpandDimsAttrs", ExpandDimsAttrs, AttrsNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.ExpandDimsAttrs", ExpandDimsAttrs, AttrsObj); }; // struct ExpandDimsAttrs /*! \brief Attributes used in layout_transform operator */ -struct LayoutTransformAttrs : public AttrsNode { +struct LayoutTransformAttrs : public AttrsObj { explicit LayoutTransformAttrs(tirx::IndexMap value) : index_map(std::move(value)) {} explicit LayoutTransformAttrs(ffi::UnsafeInit) : index_map(ffi::UnsafeInit{}) {} tirx::IndexMap index_map; @@ -57,143 +57,143 @@ struct LayoutTransformAttrs : public AttrsNode { static void RegisterReflection(); TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.LayoutTransformAttrs", LayoutTransformAttrs, - AttrsNode); + AttrsObj); }; // struct LayoutTransformAttrs /*! \brief Attributes used in permute_dims operator */ -struct PermuteDimsAttrs : public AttrsNode { +struct PermuteDimsAttrs : public AttrsObj { ffi::Optional> axes; static void RegisterReflection(); - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.PermuteDimsAttrs", PermuteDimsAttrs, AttrsNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.PermuteDimsAttrs", PermuteDimsAttrs, AttrsObj); }; // struct PermuteDimsAttrs /*! \brief Attributes used in split operator */ -struct SplitAttrs : public AttrsNode { +struct SplitAttrs : public AttrsObj { ffi::ObjectRef indices_or_sections; int axis; static void RegisterReflection(); - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.SplitAttrs", SplitAttrs, AttrsNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.SplitAttrs", SplitAttrs, AttrsObj); }; // struct SplitAttrs /*! \brief Attributes used in squeeze operators */ -struct SqueezeAttrs : public AttrsNode { +struct SqueezeAttrs : public AttrsObj { ffi::Optional> axis; static void RegisterReflection(); - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.SqueezeAttrs", SqueezeAttrs, AttrsNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.SqueezeAttrs", SqueezeAttrs, AttrsObj); }; // struct SqueezeAttrs /*! \brief Attributes used in stack operators */ -struct StackAttrs : public AttrsNode { +struct StackAttrs : public AttrsObj { ffi::Optional axis; static void RegisterReflection(); - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.StackAttrs", StackAttrs, AttrsNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.StackAttrs", StackAttrs, AttrsObj); }; // struct StackAttrs /*! \brief Attributes used in repeat operators */ -struct RepeatAttrs : public AttrsNode { +struct RepeatAttrs : public AttrsObj { int repeats; ffi::Optional axis; static void RegisterReflection(); - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.RepeatAttrs", RepeatAttrs, AttrsNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.RepeatAttrs", RepeatAttrs, AttrsObj); }; // struct RepeatAttrs /*! \brief Attributes used in tile operators */ -struct TileAttrs : public AttrsNode { +struct TileAttrs : public AttrsObj { ffi::Array repeats; static void RegisterReflection(); - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.TileAttrs", TileAttrs, AttrsNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.TileAttrs", TileAttrs, AttrsObj); }; // struct TileAttrs /*! \brief Attributes used in flip operators */ -struct FlipAttrs : public AttrsNode { +struct FlipAttrs : public AttrsObj { int64_t axis; static void RegisterReflection(); - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.FlipAttrs", FlipAttrs, AttrsNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.FlipAttrs", FlipAttrs, AttrsObj); }; // struct FlipAttrs /*! \brief Attributes used in reverse_sequence operators */ -struct ReverseSequenceAttrs : public AttrsNode { +struct ReverseSequenceAttrs : public AttrsObj { int64_t seq_axis; int64_t batch_axis; static void RegisterReflection(); TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.ReverseSequenceAttrs", ReverseSequenceAttrs, - AttrsNode); + AttrsObj); }; // struct ReverseSequenceAttrs /*! \brief Attributes used in gather_elements operators */ -struct GatherElementsAttrs : public AttrsNode { +struct GatherElementsAttrs : public AttrsObj { int64_t axis; static void RegisterReflection(); TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.GatherElementsAttrs", GatherElementsAttrs, - AttrsNode); + AttrsObj); }; // struct GatherElementsAttrs /*! \brief Attributes used in gather_nd operators */ -struct GatherNDAttrs : public AttrsNode { +struct GatherNDAttrs : public AttrsObj { int64_t batch_dims; static void RegisterReflection(); - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.GatherNDAttrs", GatherNDAttrs, AttrsNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.GatherNDAttrs", GatherNDAttrs, AttrsObj); }; // struct GatherNDAttrs /*! \brief Attributes used in index_put operator */ -struct IndexPutAttrs : public AttrsNode { +struct IndexPutAttrs : public AttrsObj { bool accumulate; static void RegisterReflection(); - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.IndexPutAttrs", IndexPutAttrs, AttrsNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.IndexPutAttrs", IndexPutAttrs, AttrsObj); }; // struct IndexPutAttrs /*! \brief Attribute used in meshgrid operator */ -struct MeshgridAttrs : public AttrsNode { +struct MeshgridAttrs : public AttrsObj { ffi::Optional indexing; static void RegisterReflection(); - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.MeshgridAttrs", MeshgridAttrs, AttrsNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.MeshgridAttrs", MeshgridAttrs, AttrsObj); }; /*! \brief Attributes used in scatter_elements operators */ -struct ScatterElementsAttrs : public AttrsNode { +struct ScatterElementsAttrs : public AttrsObj { int64_t axis; ffi::String reduction; static void RegisterReflection(); TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.ScatterElementsAttrs", ScatterElementsAttrs, - AttrsNode); + AttrsObj); }; // struct ScatterElementsAttrs /*! \brief Attributes used in scatter_nd operators */ -struct ScatterNDAttrs : public AttrsNode { +struct ScatterNDAttrs : public AttrsObj { ffi::String reduction; static void RegisterReflection(); - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.ScatterNDAttrs", ScatterNDAttrs, AttrsNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.ScatterNDAttrs", ScatterNDAttrs, AttrsObj); }; // struct ScatterNDAttrs /*! \brief Attributes used in slice_scatter operator */ -struct SliceScatterAttrs : public AttrsNode { +struct SliceScatterAttrs : public AttrsObj { int axis; static void RegisterReflection(); - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.SliceScatterAttrs", SliceScatterAttrs, AttrsNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.SliceScatterAttrs", SliceScatterAttrs, AttrsObj); }; // struct SliceScatterAttrs /*! \brief Attributes used in one_hot operator */ -struct OneHotAttrs : public AttrsNode { +struct OneHotAttrs : public AttrsObj { int depth; int axis; static void RegisterReflection(); - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.OneHotAttrs", OneHotAttrs, AttrsNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.OneHotAttrs", OneHotAttrs, AttrsObj); }; // struct OneHotAttrs } // namespace relax diff --git a/include/tvm/relax/op/nn.h b/include/tvm/relax/op/nn.h index 3a147ac54a6a..d8dd0b99f087 100644 --- a/include/tvm/relax/op/nn.h +++ b/include/tvm/relax/op/nn.h @@ -30,7 +30,7 @@ namespace tvm { namespace relax { /*! \brief Attributes used in Conv1d operator */ -struct Conv1DAttrs : public AttrsNode { +struct Conv1DAttrs : public AttrsObj { ffi::Array strides; ffi::Array padding; ffi::Array dilation; @@ -41,11 +41,11 @@ struct Conv1DAttrs : public AttrsNode { ffi::Optional out_dtype; static void RegisterReflection(); - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.Conv1DAttrs", Conv1DAttrs, AttrsNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.Conv1DAttrs", Conv1DAttrs, AttrsObj); }; // struct Conv1dAttrs /*! \brief Attributes used in Conv2d operator */ -struct Conv2DAttrs : public AttrsNode { +struct Conv2DAttrs : public AttrsObj { ffi::Array strides; ffi::Array padding; ffi::Array dilation; @@ -56,11 +56,11 @@ struct Conv2DAttrs : public AttrsNode { ffi::Optional out_dtype; static void RegisterReflection(); - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.Conv2DAttrs", Conv2DAttrs, AttrsNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.Conv2DAttrs", Conv2DAttrs, AttrsObj); }; // struct Conv2dAttrs /*! \brief Attributes used in Conv3d operator */ -struct Conv3DAttrs : public AttrsNode { +struct Conv3DAttrs : public AttrsObj { ffi::Array strides; ffi::Array padding; ffi::Array dilation; @@ -71,11 +71,11 @@ struct Conv3DAttrs : public AttrsNode { ffi::Optional out_dtype; static void RegisterReflection(); - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.Conv3DAttrs", Conv3DAttrs, AttrsNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.Conv3DAttrs", Conv3DAttrs, AttrsObj); }; // struct Conv3dAttrs /*! \brief Attributes used in Conv1DTranspose operator */ -struct Conv1DTransposeAttrs : public AttrsNode { +struct Conv1DTransposeAttrs : public AttrsObj { ffi::Array strides; ffi::Array padding; ffi::Array output_padding; @@ -88,11 +88,11 @@ struct Conv1DTransposeAttrs : public AttrsNode { static void RegisterReflection(); TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.Conv1DTransposeAttrs", Conv1DTransposeAttrs, - AttrsNode); + AttrsObj); }; // struct Conv1DTransposeAttrs /*! \brief Attributes used in Conv2d operator */ -struct Conv2DTransposeAttrs : public AttrsNode { +struct Conv2DTransposeAttrs : public AttrsObj { ffi::Array strides; ffi::Array padding; ffi::Array output_padding; @@ -105,11 +105,11 @@ struct Conv2DTransposeAttrs : public AttrsNode { static void RegisterReflection(); TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.Conv2DTransposeAttrs", Conv2DTransposeAttrs, - AttrsNode); + AttrsObj); }; // struct Conv2DTransposeAttrs /*! \brief Attributes used in Conv3dTranspose operator */ -struct Conv3DTransposeAttrs : public AttrsNode { +struct Conv3DTransposeAttrs : public AttrsObj { ffi::Array strides; ffi::Array padding; ffi::Array output_padding; @@ -122,11 +122,11 @@ struct Conv3DTransposeAttrs : public AttrsNode { static void RegisterReflection(); TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.Conv3DTransposeAttrs", Conv3DTransposeAttrs, - AttrsNode); + AttrsObj); }; // struct Conv3DTransposeAttrs /*! \brief Attributes used in max_pool1d and avg_pool1d operator */ -struct Pool1DAttrs : public AttrsNode { +struct Pool1DAttrs : public AttrsObj { ffi::Array pool_size; ffi::Array strides; ffi::Array padding; @@ -137,11 +137,11 @@ struct Pool1DAttrs : public AttrsNode { ffi::String out_layout; static void RegisterReflection(); - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.Pool1DAttrs", Pool1DAttrs, AttrsNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.Pool1DAttrs", Pool1DAttrs, AttrsObj); }; // struct Pool1dAttrs /*! \brief Attributes used in max_pool2d and avg_pool2d operator */ -struct Pool2DAttrs : public AttrsNode { +struct Pool2DAttrs : public AttrsObj { ffi::Array pool_size; ffi::Array strides; ffi::Array padding; @@ -152,11 +152,11 @@ struct Pool2DAttrs : public AttrsNode { ffi::String out_layout; static void RegisterReflection(); - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.Pool2DAttrs", Pool2DAttrs, AttrsNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.Pool2DAttrs", Pool2DAttrs, AttrsObj); }; // struct Pool2dAttrs /*! \brief Attributes used in max_pool3d and avg_pool3d operator */ -struct Pool3DAttrs : public AttrsNode { +struct Pool3DAttrs : public AttrsObj { ffi::Array pool_size; ffi::Array strides; ffi::Array padding; @@ -167,77 +167,77 @@ struct Pool3DAttrs : public AttrsNode { ffi::String out_layout; static void RegisterReflection(); - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.Pool3DAttrs", Pool3DAttrs, AttrsNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.Pool3DAttrs", Pool3DAttrs, AttrsObj); }; // struct Pool3dAttrs /*! \brief Attributes for 1d adaptive pool operator */ -struct AdaptivePool1DAttrs : public AttrsNode { +struct AdaptivePool1DAttrs : public AttrsObj { ffi::Optional> output_size; ffi::String layout; ffi::String out_layout; static void RegisterReflection(); TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.AdaptivePool1DAttrs", AdaptivePool1DAttrs, - AttrsNode); + AttrsObj); }; // struct AdaptivePool1DAttrs /*! \brief Attributes for 2d adaptive pool operator */ -struct AdaptivePool2DAttrs : public AttrsNode { +struct AdaptivePool2DAttrs : public AttrsObj { ffi::Optional> output_size; ffi::String layout; ffi::String out_layout; static void RegisterReflection(); TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.AdaptivePool2DAttrs", AdaptivePool2DAttrs, - AttrsNode); + AttrsObj); }; // struct AdaptivePool2DAttrs /*! \brief Attributes for 3d adaptive pool operator */ -struct AdaptivePool3DAttrs : public AttrsNode { +struct AdaptivePool3DAttrs : public AttrsObj { ffi::Optional> output_size; ffi::String layout; ffi::String out_layout; static void RegisterReflection(); TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.AdaptivePool3DAttrs", AdaptivePool3DAttrs, - AttrsNode); + AttrsObj); }; // struct AdaptivePool3DAttrs /*! \brief Attributes used in softmax operators */ -struct SoftmaxAttrs : public AttrsNode { +struct SoftmaxAttrs : public AttrsObj { int axis; static void RegisterReflection(); - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.SoftmaxAttrs", SoftmaxAttrs, AttrsNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.SoftmaxAttrs", SoftmaxAttrs, AttrsObj); }; /*! \brief Attributes used in softmax operators */ -struct LeakyReluAttrs : public AttrsNode { +struct LeakyReluAttrs : public AttrsObj { double alpha; static void RegisterReflection(); - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.LeakyReluAttrs", LeakyReluAttrs, AttrsNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.LeakyReluAttrs", LeakyReluAttrs, AttrsObj); }; /*! \brief Attributes used in softplus operators */ -struct SoftplusAttrs : public AttrsNode { +struct SoftplusAttrs : public AttrsObj { double beta; double threshold; static void RegisterReflection(); - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.SoftplusAttrs", SoftplusAttrs, AttrsNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.SoftplusAttrs", SoftplusAttrs, AttrsObj); }; /*! \brief Attributes used in PReLU operator */ -struct PReluAttrs : public AttrsNode { +struct PReluAttrs : public AttrsObj { int axis; static void RegisterReflection(); - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.PReluAttrs", PReluAttrs, AttrsNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.PReluAttrs", PReluAttrs, AttrsObj); }; /*! \brief Attributes used in batch_norm operator */ -struct BatchNormAttrs : public AttrsNode { +struct BatchNormAttrs : public AttrsObj { int axis; double epsilon; bool center; @@ -246,22 +246,22 @@ struct BatchNormAttrs : public AttrsNode { bool training; static void RegisterReflection(); - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.BatchNormAttrs", BatchNormAttrs, AttrsNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.BatchNormAttrs", BatchNormAttrs, AttrsObj); }; // struct BatchNormAttrs /*! \brief Attributes used in layer_norm operator */ -struct LayerNormAttrs : public AttrsNode { +struct LayerNormAttrs : public AttrsObj { ffi::Array axes; double epsilon; bool center; bool scale; static void RegisterReflection(); - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.LayerNormAttrs", LayerNormAttrs, AttrsNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.LayerNormAttrs", LayerNormAttrs, AttrsObj); }; // struct LayerNormAttrs /*! \brief Attributes used in group_norm operator */ -struct GroupNormAttrs : public AttrsNode { +struct GroupNormAttrs : public AttrsObj { int num_groups; int channel_axis; ffi::Array axes; @@ -270,11 +270,11 @@ struct GroupNormAttrs : public AttrsNode { bool scale; static void RegisterReflection(); - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.GroupNormAttrs", GroupNormAttrs, AttrsNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.GroupNormAttrs", GroupNormAttrs, AttrsObj); }; // struct GroupNormAttrs /*! \brief Attributes used in instance_norm operator */ -struct InstanceNormAttrs : public AttrsNode { +struct InstanceNormAttrs : public AttrsObj { int channel_axis; ffi::Array axes; double epsilon; @@ -282,61 +282,61 @@ struct InstanceNormAttrs : public AttrsNode { bool scale; static void RegisterReflection(); - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.InstanceNormAttrs", InstanceNormAttrs, AttrsNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.InstanceNormAttrs", InstanceNormAttrs, AttrsObj); }; // struct InstanceNormAttrs /*! \brief Attributes used in rms_norm operator */ -struct RMSNormAttrs : public AttrsNode { +struct RMSNormAttrs : public AttrsObj { ffi::Array axes; double epsilon; static void RegisterReflection(); - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.RMSNormAttrs", RMSNormAttrs, AttrsNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.RMSNormAttrs", RMSNormAttrs, AttrsObj); }; // struct RMSNormAttrs /*! \brief Attributes used in nll_loss operator */ -struct NLLLossAttrs : public AttrsNode { +struct NLLLossAttrs : public AttrsObj { ffi::String reduction; int ignore_index; static void RegisterReflection(); - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.NLLLossAttrs", NLLLossAttrs, AttrsNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.NLLLossAttrs", NLLLossAttrs, AttrsObj); }; // struct NLLLossAttrs /*! \brief Attributes used in dropout operator */ -struct DropoutAttrs : public AttrsNode { +struct DropoutAttrs : public AttrsObj { double rate; static void RegisterReflection(); - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.DropoutAttrs", DropoutAttrs, AttrsNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.DropoutAttrs", DropoutAttrs, AttrsObj); }; // struct DropoutAttrs /*! \brief Attributes used in Attention operator */ -struct AttentionAttrs : public AttrsNode { +struct AttentionAttrs : public AttrsObj { ffi::Optional scale; ffi::Optional causal_mask; ffi::Optional window_size; static void RegisterReflection(); - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.AttentionAttrs", AttentionAttrs, AttrsNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.AttentionAttrs", AttentionAttrs, AttrsObj); }; // struct AttentionAttrs /*! \brief Attributes used for the padding operator */ -struct PadAttrs : public AttrsNode { +struct PadAttrs : public AttrsObj { ffi::Array pad_width; double pad_value = 0.0; tvm::ffi::String pad_mode; static void RegisterReflection(); - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.PadAttrs", PadAttrs, AttrsNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.PadAttrs", PadAttrs, AttrsObj); }; /*! \brief Attributes used for the pixel shuffle operator */ -struct PixelShuffleAttrs : public AttrsNode { +struct PixelShuffleAttrs : public AttrsObj { int upscale_factor; static void RegisterReflection(); - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.PixelShuffleAttrs", PixelShuffleAttrs, AttrsNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.PixelShuffleAttrs", PixelShuffleAttrs, AttrsObj); }; } // namespace relax diff --git a/include/tvm/relax/op/op.h b/include/tvm/relax/op/op.h index 768373f85eaa..3657fc673109 100644 --- a/include/tvm/relax/op/op.h +++ b/include/tvm/relax/op/op.h @@ -31,17 +31,17 @@ namespace tvm { namespace relax { /*! \brief Attributes used in call_tir_with_grad */ -struct CallTIRWithGradAttrs : public AttrsNode { +struct CallTIRWithGradAttrs : public AttrsObj { ffi::String te_grad_name; ffi::Map te_grad_kwargs; static void RegisterReflection(); TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.CallTIRWithGradAttrs", CallTIRWithGradAttrs, - AttrsNode); + AttrsObj); }; // struct CallTIRAttrs /*! \brief Attributes used in call_tir_inplace */ -struct CallTIRInplaceAttrs : public AttrsNode { +struct CallTIRInplaceAttrs : public AttrsObj { /*! * \brief Indices that describe which input corresponds to which output. * @@ -53,11 +53,11 @@ struct CallTIRInplaceAttrs : public AttrsNode { static void RegisterReflection(); TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.CallTIRInplaceAttrs", CallTIRInplaceAttrs, - AttrsNode); + AttrsObj); }; // struct CallTIRInplaceAttrs /*! \brief Attributes used in call_inplace_packed */ -struct CallInplacePackedAttrs : public AttrsNode { +struct CallInplacePackedAttrs : public AttrsObj { /*! * \brief Indices that describe which input corresponds to which output. * @@ -69,27 +69,27 @@ struct CallInplacePackedAttrs : public AttrsNode { static void RegisterReflection(); TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.CallInplacePackedAttrs", CallInplacePackedAttrs, - AttrsNode); + AttrsObj); }; // struct CallInplacePackedAttrs /*! \brief Attributes used in to_vdevice */ -struct ToVDeviceAttrs : public AttrsNode { +struct ToVDeviceAttrs : public AttrsObj { explicit ToVDeviceAttrs(VDevice value) : dst_vdevice(std::move(value)) {} explicit ToVDeviceAttrs(ffi::UnsafeInit) : dst_vdevice(ffi::UnsafeInit{}) {} VDevice dst_vdevice; static void RegisterReflection(); - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.ToVDeviceAttrs", ToVDeviceAttrs, AttrsNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.ToVDeviceAttrs", ToVDeviceAttrs, AttrsObj); }; // struct ToVDeviceAttrs /*! \brief Attributes used in hint_on_device */ -struct HintOnDeviceAttrs : public AttrsNode { +struct HintOnDeviceAttrs : public AttrsObj { int32_t device_type; int32_t index; MemoryScope memory_scope; static void RegisterReflection(); - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.HintOnDeviceAttrs", HintOnDeviceAttrs, AttrsNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.HintOnDeviceAttrs", HintOnDeviceAttrs, AttrsObj); }; // struct HintOnDeviceAttrs } // namespace relax diff --git a/include/tvm/relax/op/qdq.h b/include/tvm/relax/op/qdq.h index 978028c916cf..428100f6ee1b 100644 --- a/include/tvm/relax/op/qdq.h +++ b/include/tvm/relax/op/qdq.h @@ -30,12 +30,12 @@ namespace tvm { namespace relax { /*! \brief Attributes for relax.quantize/relax.dequantize operator */ -struct QuantizeAttrs : public AttrsNode { +struct QuantizeAttrs : public AttrsObj { DLDataType out_dtype; int axis; static void RegisterReflection(); - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.QuantizeAttrs", QuantizeAttrs, AttrsNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.QuantizeAttrs", QuantizeAttrs, AttrsObj); }; // QuantizeAttrs } // namespace relax diff --git a/include/tvm/relax/op/sampling.h b/include/tvm/relax/op/sampling.h index 771e75cce25c..5ed50153d7cb 100644 --- a/include/tvm/relax/op/sampling.h +++ b/include/tvm/relax/op/sampling.h @@ -30,12 +30,12 @@ namespace tvm { namespace relax { /*! \brief Attributes used in multinomial_from_uniform operator */ -struct MultinomialFromUniformAttrs : public AttrsNode { +struct MultinomialFromUniformAttrs : public AttrsObj { DLDataType dtype; static void RegisterReflection(); TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.MultinomialFromUniformAttrs", - MultinomialFromUniformAttrs, AttrsNode); + MultinomialFromUniformAttrs, AttrsObj); }; // struct MultinomialFromUniformAttrs } // namespace relax diff --git a/include/tvm/relax/op/search.h b/include/tvm/relax/op/search.h index 6bfe61ad5ffa..b10413821ab3 100644 --- a/include/tvm/relax/op/search.h +++ b/include/tvm/relax/op/search.h @@ -30,21 +30,21 @@ namespace tvm { namespace relax { /*! \brief Attributes for search operators */ -struct ArgmaxArgminAttrs : public AttrsNode { +struct ArgmaxArgminAttrs : public AttrsObj { ffi::Optional axis; bool keepdims; static void RegisterReflection(); - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.ArgmaxArgminAttrs", ArgmaxArgminAttrs, AttrsNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.ArgmaxArgminAttrs", ArgmaxArgminAttrs, AttrsObj); }; // struct ArgmaxArgminAttrs /*! \brief Attributes for bucketize operator */ -struct BucketizeAttrs : public tvm::AttrsNode { +struct BucketizeAttrs : public tvm::AttrsObj { bool out_int32; bool right; static void RegisterReflection(); - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.BucketizeAttrs", BucketizeAttrs, AttrsNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.BucketizeAttrs", BucketizeAttrs, AttrsObj); }; // struct BucketizeAttrs } // namespace relax diff --git a/include/tvm/relax/op/sorting.h b/include/tvm/relax/op/sorting.h index 4be091cd36cd..1a9aa424e77f 100644 --- a/include/tvm/relax/op/sorting.h +++ b/include/tvm/relax/op/sorting.h @@ -31,26 +31,26 @@ namespace tvm { namespace relax { /*! \brief Attributes used in sort operator */ -struct SortAttrs : public AttrsNode { +struct SortAttrs : public AttrsObj { int axis; bool descending; static void RegisterReflection(); - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.SortAttrs", SortAttrs, AttrsNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.SortAttrs", SortAttrs, AttrsObj); }; // struct SortAttrs /*! \brief Attributes used in argsort operator */ -struct ArgsortAttrs : public AttrsNode { +struct ArgsortAttrs : public AttrsObj { int axis; bool descending; ffi::Optional dtype; static void RegisterReflection(); - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.ArgsortAttrs", ArgsortAttrs, AttrsNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.ArgsortAttrs", ArgsortAttrs, AttrsObj); }; // struct ArgsortAttrs /*! \brief Attributes used in topk operator */ -struct TopKAttrs : public AttrsNode { +struct TopKAttrs : public AttrsObj { int k; int axis; bool largest; @@ -58,7 +58,7 @@ struct TopKAttrs : public AttrsNode { ffi::Optional dtype; static void RegisterReflection(); - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.TopKAttrs", TopKAttrs, AttrsNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.TopKAttrs", TopKAttrs, AttrsObj); }; // struct TopKAttrs } // namespace relax diff --git a/include/tvm/relax/op/statistical.h b/include/tvm/relax/op/statistical.h index 2d7ec0c733af..ff6e43c821b5 100644 --- a/include/tvm/relax/op/statistical.h +++ b/include/tvm/relax/op/statistical.h @@ -30,22 +30,22 @@ namespace tvm { namespace relax { /*! \brief Attributes for statistical operators */ -struct StatisticalAttrs : public AttrsNode { +struct StatisticalAttrs : public AttrsObj { ffi::Optional> axis; bool keepdims; static void RegisterReflection(); - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.StatisticalAttrs", StatisticalAttrs, AttrsNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.StatisticalAttrs", StatisticalAttrs, AttrsObj); }; // struct StatisticalAttrs /*! \brief Attributes used in scan operators like cumsum, cumprod */ -struct ScanopAttrs : public AttrsNode { +struct ScanopAttrs : public AttrsObj { ffi::Optional axis; ffi::Optional dtype; bool exclusive = false; static void RegisterReflection(); - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.ScanopAttrs", ScanopAttrs, AttrsNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.ScanopAttrs", ScanopAttrs, AttrsObj); }; // struct ScanopAttrs } // namespace relax diff --git a/include/tvm/relax/op/vision.h b/include/tvm/relax/op/vision.h index ee9f35ed164f..5fea80fa4cce 100644 --- a/include/tvm/relax/op/vision.h +++ b/include/tvm/relax/op/vision.h @@ -32,16 +32,16 @@ namespace tvm { namespace relax { /*! \brief Attributes used in AllClassNonMaximumSuppression operator */ -struct AllClassNonMaximumSuppressionAttrs : public AttrsNode { +struct AllClassNonMaximumSuppressionAttrs : public AttrsObj { ffi::String output_format; static void RegisterReflection(); TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.AllClassNonMaximumSuppressionAttrs", - AllClassNonMaximumSuppressionAttrs, AttrsNode); + AllClassNonMaximumSuppressionAttrs, AttrsObj); }; // struct AllClassNonMaximumSuppressionAttrs /*! \brief Attributes used in ROIAlign operator */ -struct ROIAlignAttrs : public AttrsNode { +struct ROIAlignAttrs : public AttrsObj { ffi::Array pooled_size; double spatial_scale; int sample_ratio; @@ -50,32 +50,32 @@ struct ROIAlignAttrs : public AttrsNode { ffi::String mode; static void RegisterReflection(); - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.ROIAlignAttrs", ROIAlignAttrs, AttrsNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.ROIAlignAttrs", ROIAlignAttrs, AttrsObj); }; // struct ROIAlignAttrs /*! \brief Attributes used in ROIPool operator */ -struct ROIPoolAttrs : public AttrsNode { +struct ROIPoolAttrs : public AttrsObj { ffi::Array pooled_size; double spatial_scale; ffi::String layout; static void RegisterReflection(); - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.ROIPoolAttrs", ROIPoolAttrs, AttrsNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.ROIPoolAttrs", ROIPoolAttrs, AttrsObj); }; // struct ROIPoolAttrs /*! \brief Attributes used in GetValidCounts operator */ -struct GetValidCountsAttrs : public AttrsNode { +struct GetValidCountsAttrs : public AttrsObj { double score_threshold; int id_index; int score_index; static void RegisterReflection(); TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.GetValidCountsAttrs", GetValidCountsAttrs, - AttrsNode); + AttrsObj); }; // struct GetValidCountsAttrs /*! \brief Attributes used in NonMaximumSuppression operator */ -struct NonMaximumSuppressionAttrs : public AttrsNode { +struct NonMaximumSuppressionAttrs : public AttrsObj { int max_output_size; double iou_threshold; bool force_suppress; @@ -90,11 +90,11 @@ struct NonMaximumSuppressionAttrs : public AttrsNode { static void RegisterReflection(); TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.NonMaximumSuppressionAttrs", - NonMaximumSuppressionAttrs, AttrsNode); + NonMaximumSuppressionAttrs, AttrsObj); }; // struct NonMaximumSuppressionAttrs /*! \brief Attributes for multibox_transform_loc (SSD / TFLite-style box decode). */ -struct MultiboxTransformLocAttrs : public AttrsNode { +struct MultiboxTransformLocAttrs : public AttrsObj { bool clip; double threshold; ffi::Array variances; @@ -103,7 +103,7 @@ struct MultiboxTransformLocAttrs : public AttrsNode { static void RegisterReflection(); TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.MultiboxTransformLocAttrs", - MultiboxTransformLocAttrs, AttrsNode); + MultiboxTransformLocAttrs, AttrsObj); }; // struct MultiboxTransformLocAttrs } // namespace relax diff --git a/include/tvm/relax/type.h b/include/tvm/relax/type.h index c2c4095a1151..f668754648d6 100644 --- a/include/tvm/relax/type.h +++ b/include/tvm/relax/type.h @@ -40,7 +40,7 @@ namespace tvm { namespace relax { using Expr = tvm::Expr; -using ExprNode = tvm::ExprNode; +using ExprObj = tvm::ExprObj; class BlockBuilder; @@ -48,28 +48,28 @@ class BlockBuilder; static constexpr int kUnknownNDim = -1; using tvm::AnyType; -using tvm::AnyTypeNode; +using tvm::AnyTypeObj; using tvm::TupleType; -using tvm::TupleTypeNode; +using tvm::TupleTypeObj; // Compatibility aliases for existing C++ callers. New code should use AnyType. -using ObjectTypeNode = AnyTypeNode; +using ObjectTypeObj = AnyTypeObj; using ObjectType = AnyType; -class PackedFuncTypeNode : public TypeNode { +class PackedFuncTypeObj : public TypeObj { public: static void RegisterReflection() { namespace refl = tvm::ffi::reflection; - refl::ObjectDef(); + refl::ObjectDef(); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.PackedFuncType", PackedFuncTypeNode, TypeNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.PackedFuncType", PackedFuncTypeObj, TypeObj); }; class PackedFuncType : public Type { public: TVM_DLL PackedFuncType(Location loc = UnknownLoc()); - TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(PackedFuncType, Type, PackedFuncTypeNode); + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(PackedFuncType, Type, PackedFuncTypeObj); }; /*! @@ -100,7 +100,7 @@ class PackedFuncType : public Type { /*! * \brief Type of shape value. */ -class ShapeTypeNode : public TypeNode { +class ShapeTypeObj : public TypeObj { public: /*! \brief optionally stores the symbolic value patterns of the shape */ ffi::Optional> values; @@ -115,16 +115,16 @@ class ShapeTypeNode : public TypeNode { static void RegisterReflection() { namespace refl = tvm::ffi::reflection; - refl::ObjectDef() - .def_ro("values", &ShapeTypeNode::values) - .def_ro("ndim", &ShapeTypeNode::ndim); + refl::ObjectDef() + .def_ro("values", &ShapeTypeObj::values) + .def_ro("ndim", &ShapeTypeObj::ndim); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.ShapeType", ShapeTypeNode, TypeNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.ShapeType", ShapeTypeObj, TypeObj); }; /*! - * \brief Managed reference to ShapeTypeNode. - * \sa ShapeTypeNode + * \brief Managed reference to ShapeTypeObj. + * \sa ShapeTypeObj */ class ShapeType : public Type { public: @@ -141,13 +141,13 @@ class ShapeType : public Type { */ TVM_DLL ShapeType(int ndim, Location loc = UnknownLoc()); - TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ShapeType, Type, ShapeTypeNode); + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ShapeType, Type, ShapeTypeObj); }; /*! * \brief Type of Tensor. */ -class TensorTypeNode : public TypeNode { +class TensorTypeObj : public TypeObj { public: /*! * \brief optionally store the shape expression of the tensor. @@ -177,7 +177,7 @@ class TensorTypeNode : public TypeNode { if (!shape.has_value()) return {}; const Expr& shape_expr = this->shape.value(); if (shape_expr->ty.as().has_value()) return {}; - if (const auto* shape_ty = shape_expr->ty.as()) { + if (const auto* shape_ty = shape_expr->ty.as()) { return shape_ty->values; } return {}; @@ -185,22 +185,22 @@ class TensorTypeNode : public TypeNode { static void RegisterReflection() { namespace refl = tvm::ffi::reflection; - refl::ObjectDef() - .def_ro("shape", &TensorTypeNode::shape) - .def_ro("dtype", &TensorTypeNode::dtype) - .def_ro("vdevice", &TensorTypeNode::vdevice) - .def_ro("ndim", &TensorTypeNode::ndim); + refl::ObjectDef() + .def_ro("shape", &TensorTypeObj::shape) + .def_ro("dtype", &TensorTypeObj::dtype) + .def_ro("vdevice", &TensorTypeObj::vdevice) + .def_ro("ndim", &TensorTypeObj::ndim); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.TensorType", TensorTypeNode, TypeNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.TensorType", TensorTypeObj, TypeObj); }; /*! - * \brief Managed reference to TensorTypeNode. - * \sa TensorTypeNode + * \brief Managed reference to TensorTypeObj. + * \sa TensorTypeObj */ class TensorType : public Type { public: - explicit TensorType(ffi::ObjectPtr data) : Type(ffi::UnsafeInit{}) { + explicit TensorType(ffi::ObjectPtr data) : Type(ffi::UnsafeInit{}) { TVM_FFI_ICHECK(data != nullptr); data_ = std::move(data); } @@ -227,7 +227,7 @@ class TensorType : public Type { TVM_DLL TensorType(ffi::Optional dtype, int ndim, ffi::Optional vdevice = std::nullopt, Location loc = UnknownLoc()); - TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(TensorType, Type, TensorTypeNode); + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(TensorType, Type, TensorTypeObj); }; /*! @@ -244,7 +244,7 @@ using TypeDeriveFunc = TypedEnvFunc() - .def_ro("params", &FuncTypeNode::params, refl::AttachFieldFlag::SEqHashDefPattern()) - .def_ro("ret", &FuncTypeNode::ret) - .def_ro("derive_func", &FuncTypeNode::derive_func) - .def_ro("purity", &FuncTypeNode::purity); + refl::ObjectDef() + .def_ro("params", &FuncTypeObj::params, refl::AttachFieldFlag::SEqHashDefPattern()) + .def_ro("ret", &FuncTypeObj::ret) + .def_ro("derive_func", &FuncTypeObj::derive_func) + .def_ro("purity", &FuncTypeObj::purity); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.FuncType", FuncTypeNode, TypeNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.FuncType", FuncTypeObj, TypeObj); }; /*! - * \brief Managed reference to FuncTypeNode. - * \sa FuncTypeNode + * \brief Managed reference to FuncTypeObj. + * \sa FuncTypeObj */ class FuncType : public Type { public: - explicit FuncType(ffi::ObjectPtr data) : Type(ffi::UnsafeInit{}) { + explicit FuncType(ffi::ObjectPtr data) : Type(ffi::UnsafeInit{}) { TVM_FFI_ICHECK(data != nullptr); data_ = std::move(data); } @@ -337,7 +337,7 @@ class FuncType : public Type { TVM_DLL static FuncType OpaqueFunc(Type ret = AnyType(), bool purity = false, Location loc = UnknownLoc()); - TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(FuncType, Type, FuncTypeNode); + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(FuncType, Type, FuncTypeObj); }; /*! @@ -393,7 +393,7 @@ inline Type GetType(const Expr& expr) { * \return Whether the expr has void type. */ inline bool HasVoidType(const Expr& expr) { - auto* ptr = expr->ty.as(); + auto* ptr = expr->ty.as(); return ptr != nullptr && ptr->fields.size() == 0; } diff --git a/include/tvm/relax/type_functor.h b/include/tvm/relax/type_functor.h index 9374c8c849e8..9d367e01c261 100644 --- a/include/tvm/relax/type_functor.h +++ b/include/tvm/relax/type_functor.h @@ -79,16 +79,16 @@ class TypeFunctor { return vtable(n, this, std::forward(args)...); } // Functions that can be overriden by subclass - virtual R VisitType_(const AnyTypeNode* op, Args... args) RELAX_TYPE_FUNCTOR_DEFAULT; - virtual R VisitType_(const PrimTypeNode* op, Args... args) RELAX_TYPE_FUNCTOR_DEFAULT; - virtual R VisitType_(const StringTypeNode* op, Args... args) RELAX_TYPE_FUNCTOR_DEFAULT; - virtual R VisitType_(const ShapeTypeNode* op, Args... args) RELAX_TYPE_FUNCTOR_DEFAULT; - virtual R VisitType_(const TensorTypeNode* op, Args... args) RELAX_TYPE_FUNCTOR_DEFAULT; - virtual R VisitType_(const distributed::DTensorTypeNode* op, + virtual R VisitType_(const AnyTypeObj* op, Args... args) RELAX_TYPE_FUNCTOR_DEFAULT; + virtual R VisitType_(const PrimTypeObj* op, Args... args) RELAX_TYPE_FUNCTOR_DEFAULT; + virtual R VisitType_(const StringTypeObj* op, Args... args) RELAX_TYPE_FUNCTOR_DEFAULT; + virtual R VisitType_(const ShapeTypeObj* op, Args... args) RELAX_TYPE_FUNCTOR_DEFAULT; + virtual R VisitType_(const TensorTypeObj* op, Args... args) RELAX_TYPE_FUNCTOR_DEFAULT; + virtual R VisitType_(const distributed::DTensorTypeObj* op, Args... args) RELAX_TYPE_FUNCTOR_DEFAULT; - virtual R VisitType_(const TupleTypeNode* op, Args... args) RELAX_TYPE_FUNCTOR_DEFAULT; - virtual R VisitType_(const tvm::FuncTypeNode* op, Args... args) RELAX_TYPE_FUNCTOR_DEFAULT; - virtual R VisitType_(const FuncTypeNode* op, Args... args) RELAX_TYPE_FUNCTOR_DEFAULT; + virtual R VisitType_(const TupleTypeObj* op, Args... args) RELAX_TYPE_FUNCTOR_DEFAULT; + virtual R VisitType_(const tvm::FuncTypeObj* op, Args... args) RELAX_TYPE_FUNCTOR_DEFAULT; + virtual R VisitType_(const FuncTypeObj* op, Args... args) RELAX_TYPE_FUNCTOR_DEFAULT; virtual R VisitTypeDefault_(const ffi::Object* op, Args...) { TVM_FFI_THROW(InternalError) << "Do not have a default for " << op->GetTypeKey(); throw; // unreachable, written to stop compiler warning @@ -99,15 +99,15 @@ class TypeFunctor { static FType InitVTable() { FType vtable; // Set dispatch - TVM_RELAX_TYPE_FUNCTOR_DISPATCH(AnyTypeNode); - TVM_RELAX_TYPE_FUNCTOR_DISPATCH(PrimTypeNode); - TVM_RELAX_TYPE_FUNCTOR_DISPATCH(StringTypeNode); - TVM_RELAX_TYPE_FUNCTOR_DISPATCH(ShapeTypeNode); - TVM_RELAX_TYPE_FUNCTOR_DISPATCH(TensorTypeNode); - TVM_RELAX_TYPE_FUNCTOR_DISPATCH(distributed::DTensorTypeNode); - TVM_RELAX_TYPE_FUNCTOR_DISPATCH(TupleTypeNode); - TVM_RELAX_TYPE_FUNCTOR_DISPATCH(tvm::FuncTypeNode); - TVM_RELAX_TYPE_FUNCTOR_DISPATCH(FuncTypeNode); + TVM_RELAX_TYPE_FUNCTOR_DISPATCH(AnyTypeObj); + TVM_RELAX_TYPE_FUNCTOR_DISPATCH(PrimTypeObj); + TVM_RELAX_TYPE_FUNCTOR_DISPATCH(StringTypeObj); + TVM_RELAX_TYPE_FUNCTOR_DISPATCH(ShapeTypeObj); + TVM_RELAX_TYPE_FUNCTOR_DISPATCH(TensorTypeObj); + TVM_RELAX_TYPE_FUNCTOR_DISPATCH(distributed::DTensorTypeObj); + TVM_RELAX_TYPE_FUNCTOR_DISPATCH(TupleTypeObj); + TVM_RELAX_TYPE_FUNCTOR_DISPATCH(tvm::FuncTypeObj); + TVM_RELAX_TYPE_FUNCTOR_DISPATCH(FuncTypeObj); vtable.Finalize(); return vtable; } @@ -120,15 +120,15 @@ class TypeFunctor { */ class TVM_DLL TypeVisitor : public TypeFunctor { public: - void VisitType_(const AnyTypeNode* op) override; - void VisitType_(const PrimTypeNode* op) override; - void VisitType_(const StringTypeNode* op) override; - void VisitType_(const ShapeTypeNode* op) override; - void VisitType_(const TensorTypeNode* op) override; - void VisitType_(const distributed::DTensorTypeNode* op) override; - void VisitType_(const TupleTypeNode* op) override; - void VisitType_(const tvm::FuncTypeNode* op) override; - void VisitType_(const FuncTypeNode* op) override; + void VisitType_(const AnyTypeObj* op) override; + void VisitType_(const PrimTypeObj* op) override; + void VisitType_(const StringTypeObj* op) override; + void VisitType_(const ShapeTypeObj* op) override; + void VisitType_(const TensorTypeObj* op) override; + void VisitType_(const distributed::DTensorTypeObj* op) override; + void VisitType_(const TupleTypeObj* op) override; + void VisitType_(const tvm::FuncTypeObj* op) override; + void VisitType_(const FuncTypeObj* op) override; protected: // two functions to override when visit expr fields in type nodes. @@ -141,15 +141,15 @@ class TVM_DLL TypeVisitor : public TypeFunctor { */ class TVM_DLL TypeMutator : public TypeFunctor { public: - Type VisitType_(const AnyTypeNode* op) override; - Type VisitType_(const PrimTypeNode* op) override; - Type VisitType_(const StringTypeNode* op) override; - Type VisitType_(const ShapeTypeNode* op) override; - Type VisitType_(const TensorTypeNode* op) override; - Type VisitType_(const distributed::DTensorTypeNode* op) override; - Type VisitType_(const TupleTypeNode* op) override; - Type VisitType_(const tvm::FuncTypeNode* op) override; - Type VisitType_(const FuncTypeNode* op) override; + Type VisitType_(const AnyTypeObj* op) override; + Type VisitType_(const PrimTypeObj* op) override; + Type VisitType_(const StringTypeObj* op) override; + Type VisitType_(const ShapeTypeObj* op) override; + Type VisitType_(const TensorTypeObj* op) override; + Type VisitType_(const distributed::DTensorTypeObj* op) override; + Type VisitType_(const TupleTypeObj* op) override; + Type VisitType_(const tvm::FuncTypeObj* op) override; + Type VisitType_(const FuncTypeObj* op) override; protected: // two functions to override when visit expr fields in type nodes. diff --git a/include/tvm/s_tir/analysis.h b/include/tvm/s_tir/analysis.h index bdc5bc14c871..29e7e6cdcf91 100644 --- a/include/tvm/s_tir/analysis.h +++ b/include/tvm/s_tir/analysis.h @@ -87,7 +87,7 @@ TVM_DLL ffi::Map> DetectTensorAccessLCA(const Fun * \param mod The input TIR module. * \return The anchor block if found, nullptr otherwise. */ -const s_tir::SBlockNode* FindAnchorBlock(const IRModule& mod); +const s_tir::SBlockObj* FindAnchorBlock(const IRModule& mod); } // namespace tirx diff --git a/include/tvm/s_tir/data_layout.h b/include/tvm/s_tir/data_layout.h index f71e21b3e93f..7b8a3fa7a092 100644 --- a/include/tvm/s_tir/data_layout.h +++ b/include/tvm/s_tir/data_layout.h @@ -101,7 +101,7 @@ class SLayoutAxis { * Here subordinate axis channel_block=16 is the factor size of the primal axis C (channel). * SLayout for scalar is defined, while both its name and axes have size 0. */ -class SLayoutNode : public ffi::Object { +class SLayoutObj : public ffi::Object { public: /*! \brief string representation of layout, "" for scalar. */ ffi::String name; @@ -115,16 +115,16 @@ class SLayoutNode : public ffi::Object { static void RegisterReflection() { namespace refl = tvm::ffi::reflection; - refl::ObjectDef() - .def_ro("name", &SLayoutNode::name) - .def_ro("axes", &SLayoutNode::axes); + refl::ObjectDef() + .def_ro("name", &SLayoutObj::name) + .def_ro("axes", &SLayoutObj::axes); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("s_tir.SLayout", SLayoutNode, ffi::Object); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("s_tir.SLayout", SLayoutObj, ffi::Object); }; /*! - * \brief Managed reference to SLayoutNode - * \sa SLayoutNode + * \brief Managed reference to SLayoutObj + * \sa SLayoutObj */ class SLayout : public ffi::ObjectRef { public: @@ -154,7 +154,7 @@ class SLayout : public ffi::ObjectRef { * \brief access the internal node container * \return the pointer to the internal node container */ - SLayoutNode* operator->() { return static_cast(get_mutable()); } + SLayoutObj* operator->() { return static_cast(get_mutable()); } /*! * \brief Return an undefined layout. @@ -342,15 +342,15 @@ class SLayout : public ffi::ObjectRef { return os; } - TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(SLayout, ffi::ObjectRef, SLayoutNode); + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(SLayout, ffi::ObjectRef, SLayoutObj); }; // Internal node container SBijectiveLayout -class SBijectiveLayoutNode : public ffi::Object { +class SBijectiveLayoutObj : public ffi::Object { public: - explicit SBijectiveLayoutNode(ffi::UnsafeInit tag) : src_layout(tag), dst_layout(tag) {} + explicit SBijectiveLayoutObj(ffi::UnsafeInit tag) : src_layout(tag), dst_layout(tag) {} - SBijectiveLayoutNode(SLayout src, SLayout dst) + SBijectiveLayoutObj(SLayout src, SLayout dst) : src_layout(std::move(src)), dst_layout(std::move(dst)) {} /*! \brief Describes how source axes can be mapped to the destination axes, * e.g., [i0 / 16, i1, i0 % 16] can describe NC -> NC16n @@ -370,15 +370,15 @@ class SBijectiveLayoutNode : public ffi::Object { static void RegisterReflection() { namespace refl = tvm::ffi::reflection; - refl::ObjectDef() - .def_ro("src_layout", &SBijectiveLayoutNode::src_layout) - .def_ro("dst_layout", &SBijectiveLayoutNode::dst_layout) - .def_ro("index_forward_rule", &SBijectiveLayoutNode::index_forward_rule) - .def_ro("index_backward_rule", &SBijectiveLayoutNode::index_backward_rule) - .def_ro("shape_forward_rule", &SBijectiveLayoutNode::shape_forward_rule) - .def_ro("shape_backward_rule", &SBijectiveLayoutNode::shape_backward_rule); + refl::ObjectDef() + .def_ro("src_layout", &SBijectiveLayoutObj::src_layout) + .def_ro("dst_layout", &SBijectiveLayoutObj::dst_layout) + .def_ro("index_forward_rule", &SBijectiveLayoutObj::index_forward_rule) + .def_ro("index_backward_rule", &SBijectiveLayoutObj::index_backward_rule) + .def_ro("shape_forward_rule", &SBijectiveLayoutObj::shape_forward_rule) + .def_ro("shape_backward_rule", &SBijectiveLayoutObj::shape_backward_rule); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("s_tir.SBijectiveLayout", SBijectiveLayoutNode, ffi::Object); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("s_tir.SBijectiveLayout", SBijectiveLayoutObj, ffi::Object); }; /*! @@ -389,7 +389,7 @@ class SBijectiveLayoutNode : public ffi::Object { */ class SBijectiveLayout : public ffi::ObjectRef { public: - explicit SBijectiveLayout(ffi::UnsafeInit tag, ffi::ObjectPtr data) + explicit SBijectiveLayout(ffi::UnsafeInit tag, ffi::ObjectPtr data) : ffi::ObjectRef(tag) { data_ = std::move(data); } @@ -412,7 +412,7 @@ class SBijectiveLayout : public ffi::ObjectRef { TVM_DLL ffi::Array BackwardIndex(const ffi::Array& dst_index) const; TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(SBijectiveLayout, ffi::ObjectRef, - SBijectiveLayoutNode); + SBijectiveLayoutObj); }; } // namespace tirx diff --git a/include/tvm/s_tir/iter_var.h b/include/tvm/s_tir/iter_var.h index 96ddb9ff6183..27a7aa9f2652 100644 --- a/include/tvm/s_tir/iter_var.h +++ b/include/tvm/s_tir/iter_var.h @@ -106,11 +106,11 @@ enum IterVarType : int { * * The dtype of the extent of the `dom` of the IterVar must match the dtype of the internal Var. */ -class IterVarNode : public OpaqueExprNode { +class IterVarObj : public OpaqueExprObj { public: - explicit IterVarNode(ffi::UnsafeInit tag) : var(tag) {} + explicit IterVarObj(ffi::UnsafeInit tag) : var(tag) {} - explicit IterVarNode(PrimVar var) : var(std::move(var)) { ty = this->var.ty(); } + explicit IterVarObj(PrimVar var) : var(std::move(var)) { ty = this->var.ty(); } /*! * \brief the domain of iteration, if known, can be None @@ -128,15 +128,15 @@ class IterVarNode : public OpaqueExprNode { ffi::String thread_tag; static void RegisterReflection() { namespace refl = tvm::ffi::reflection; - refl::ObjectDef() - .def_ro("dom", &IterVarNode::dom) - .def_ro("var", &IterVarNode::var, refl::AttachFieldFlag::SEqHashDefSimple()) - .def_ro("iter_type", &IterVarNode::iter_type) - .def_ro("thread_tag", &IterVarNode::thread_tag); + refl::ObjectDef() + .def_ro("dom", &IterVarObj::dom) + .def_ro("var", &IterVarObj::var, refl::AttachFieldFlag::SEqHashDefSimple()) + .def_ro("iter_type", &IterVarObj::iter_type) + .def_ro("thread_tag", &IterVarObj::thread_tag); } static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode; - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("s_tir.IterVar", IterVarNode, OpaqueExprNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("s_tir.IterVar", IterVarObj, OpaqueExprObj); }; /*! @@ -150,11 +150,11 @@ class IterVar : public PrimExpr { TVM_DLL IterVar(ffi::Optional dom, PrimVar var, IterVarType iter_type, ffi::String thread_tag = "", Location loc = UnknownLoc()); - explicit IterVar(ffi::ObjectPtr node) : PrimExpr(std::move(node)) {} + explicit IterVar(ffi::ObjectPtr node) : PrimExpr(std::move(node)) {} - TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(IterVar, PrimExpr, IterVarNode); + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(IterVar, PrimExpr, IterVarObj); static constexpr bool _type_container_is_exact = true; - TVM_DEFINE_OBJECT_REF_COW_METHOD(IterVarNode); + TVM_DEFINE_OBJECT_REF_COW_METHOD(IterVarObj); }; inline const char* IterVarType2String(IterVarType t) { diff --git a/include/tvm/s_tir/sblock_dependence_info.h b/include/tvm/s_tir/sblock_dependence_info.h index 815d85c6b988..b10204aad496 100644 --- a/include/tvm/s_tir/sblock_dependence_info.h +++ b/include/tvm/s_tir/sblock_dependence_info.h @@ -60,7 +60,7 @@ class SBlockDependenceInfoObj : public ffi::Object { */ std::unordered_map sref2scope; /*! \brief The reverse mapping from block/for-loop to their corresponding srefs */ - std::unordered_map stmt2ref; + std::unordered_map stmt2ref; static void RegisterReflection() { namespace refl = tvm::ffi::reflection; diff --git a/include/tvm/s_tir/sblock_scope.h b/include/tvm/s_tir/sblock_scope.h index 67281c3f1f65..76aa072e9001 100644 --- a/include/tvm/s_tir/sblock_scope.h +++ b/include/tvm/s_tir/sblock_scope.h @@ -60,7 +60,7 @@ class StmtSRefObj : public ffi::Object { * \note Non-owned reference (raw pointer) is used here, so that we can perform copy-on-write * optimization on statements when possible. The strong reference is held in the ScheduleState. */ - const StmtNode* stmt; + const StmtObj* stmt; /*! \brief The parent sref. */ StmtSRefObj* parent; /*! @@ -89,7 +89,7 @@ class StmtSRefObj : public ffi::Object { * It serves the same purpose as `ffi::ObjectRef::as`, but does not acquire strong reference to * `stmt` * \tparam StmtType The type that `this->stmt` to be downcasted to. Presumably - * tvm::s_tir::SBlockNode or tvm::ForNode + * tvm::s_tir::SBlockObj or tvm::ForObj * \return nullptr if type check fails, otherwise the casted result for `this->stmt` */ template @@ -115,7 +115,7 @@ class StmtSRef : public ffi::ObjectRef { * \param seq_index The location in an array if the parent of the stmt contains multiple children. * -1 if the parent does not contain multiple children. */ - TVM_DLL explicit StmtSRef(const StmtNode* stmt, StmtSRefObj* parent, int64_t seq_index); + TVM_DLL explicit StmtSRef(const StmtObj* stmt, StmtSRefObj* parent, int64_t seq_index); TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(StmtSRef, ffi::ObjectRef, StmtSRefObj); @@ -151,7 +151,7 @@ class SRefTreeCreator : public s_tir::StmtExprVisitor { using s_tir::StmtExprVisitor::Visit_; ffi::Optional Visit(ffi::AnyView value) override { - if (value.as()) return std::nullopt; + if (value.as()) return std::nullopt; return s_tir::StmtExprVisitor::Visit(value); } @@ -160,8 +160,8 @@ class SRefTreeCreator : public s_tir::StmtExprVisitor { * \param mod The module being scheduled. * \param include_loops Ignore ForNodes if this value is false */ - static std::unordered_map Create(IRModule mod, - bool include_loops = true) { + static std::unordered_map Create(IRModule mod, + bool include_loops = true) { auto creator = ffi::make_object(include_loops); for (const auto& kv : mod->functions) { const BaseFunc& base_func = kv.second; @@ -180,20 +180,20 @@ class SRefTreeCreator : public s_tir::StmtExprVisitor { * \brief Add a new statement to the stack, which becomes the current scope * \param stmt A for-loop statement or a block statement */ - void PushSRef(const StmtNode* stmt); + void PushSRef(const StmtObj* stmt); /*! \brief Pop the top of the scope and record it in stmt2ref map */ void PopAndRecordSRef(); - ffi::Optional Visit_(const ForNode* loop) final; + ffi::Optional Visit_(const ForObj* loop) final; - ffi::Optional Visit_(const s_tir::SBlockRealizeNode* realize) final; + ffi::Optional Visit_(const s_tir::SBlockRealizeObj* realize) final; - ffi::Optional Visit_(const SeqStmtNode* seq_stmt) final; + ffi::Optional Visit_(const SeqStmtObj* seq_stmt) final; bool include_loops_; /*! \brief The result ScheduleStateObj */ - std::unordered_map stmt2ref_; + std::unordered_map stmt2ref_; /*! \brief The stack frame used to indicate the current scope */ std::vector srefs_; }; diff --git a/include/tvm/s_tir/schedule/instruction.h b/include/tvm/s_tir/schedule/instruction.h index 1bf368c90cf2..f9c6eef966bd 100644 --- a/include/tvm/s_tir/schedule/instruction.h +++ b/include/tvm/s_tir/schedule/instruction.h @@ -80,7 +80,7 @@ using FInstructionAttrsFromJSON = ffi::TypedFunction(ffi::Object * 4) A functor that serialize its attributes to JSON * 5) A functor that deserialize its attributes from JSON * - * Unlike `tvm::OpNode`, `InstructionKindObj` doesn't support unstructured properties, + * Unlike `tvm::OpObj`, `InstructionKindObj` doesn't support unstructured properties, * mainly because there is no such usecase yet to add any other property. */ class InstructionKindObj : public ffi::Object { @@ -143,11 +143,11 @@ class InstructionKind : public ffi::ObjectRef { }; /*! \brief Schedule instructions each corresponds to a schedule primitive */ -class InstructionNode : public ffi::Object { +class InstructionObj : public ffi::Object { public: - explicit InstructionNode(ffi::UnsafeInit tag) : kind(tag) {} + explicit InstructionObj(ffi::UnsafeInit tag) : kind(tag) {} - explicit InstructionNode(InstructionKind kind) : kind(std::move(kind)) {} + explicit InstructionObj(InstructionKind kind) : kind(std::move(kind)) {} /*! \brief The kind of the instruction */ InstructionKind kind; /*! @@ -178,18 +178,18 @@ class InstructionNode : public ffi::Object { static void RegisterReflection() { namespace refl = tvm::ffi::reflection; - refl::ObjectDef() - .def_ro("kind", &InstructionNode::kind) - .def_ro("inputs", &InstructionNode::inputs) - .def_ro("attrs", &InstructionNode::attrs) - .def_ro("outputs", &InstructionNode::outputs); + refl::ObjectDef() + .def_ro("kind", &InstructionObj::kind) + .def_ro("inputs", &InstructionObj::inputs) + .def_ro("attrs", &InstructionObj::attrs) + .def_ro("outputs", &InstructionObj::outputs); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("s_tir.Instruction", InstructionNode, ffi::Object); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("s_tir.Instruction", InstructionObj, ffi::Object); }; /*! - * \brief Managed reference to InstructionNode - * \sa InstructionNode + * \brief Managed reference to InstructionObj + * \sa InstructionObj */ class Instruction : public ffi::ObjectRef { public: @@ -203,7 +203,7 @@ class Instruction : public ffi::ObjectRef { explicit Instruction(InstructionKind kind, ffi::Array inputs, ffi::Array attrs, ffi::Array outputs); - TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Instruction, ffi::ObjectRef, InstructionNode); + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Instruction, ffi::ObjectRef, InstructionObj); }; /*! \brief Temporary builder for a canonical named instruction kind. */ diff --git a/include/tvm/s_tir/schedule/schedule.h b/include/tvm/s_tir/schedule/schedule.h index 46fcea90f810..01bc71eb2ac5 100644 --- a/include/tvm/s_tir/schedule/schedule.h +++ b/include/tvm/s_tir/schedule/schedule.h @@ -52,47 +52,47 @@ enum class TensorIndexType : int32_t { /**************** Random variable: SBlockRV ****************/ /*! \brief A random variable that evaluates to a TensorIR block */ -class SBlockRVNode : public ffi::Object { +class SBlockRVObj : public ffi::Object { public: static void RegisterReflection() { namespace refl = tvm::ffi::reflection; - refl::ObjectDef(); + refl::ObjectDef(); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("s_tir.SBlockRV", SBlockRVNode, ffi::Object); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("s_tir.SBlockRV", SBlockRVObj, ffi::Object); }; /*! - * \brief Managed reference to SBlockRVNode - * \sa SBlockRVNode + * \brief Managed reference to SBlockRVObj + * \sa SBlockRVObj */ class SBlockRV : public ffi::ObjectRef { public: /*! \brief Constructor */ TVM_DLL SBlockRV(); - TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(SBlockRV, ffi::ObjectRef, SBlockRVNode); + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(SBlockRV, ffi::ObjectRef, SBlockRVObj); }; /**************** Random variable: LoopRV ****************/ /*! \brief A random variable that evaluates to a TensorIR for loop */ -class LoopRVNode : public ffi::Object { +class LoopRVObj : public ffi::Object { public: static void RegisterReflection() { namespace refl = tvm::ffi::reflection; - refl::ObjectDef(); + refl::ObjectDef(); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("s_tir.LoopRV", LoopRVNode, ffi::Object); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("s_tir.LoopRV", LoopRVObj, ffi::Object); }; /*! - * \brief Managed reference to LoopRVNode - * \sa LoopRVNode + * \brief Managed reference to LoopRVObj + * \sa LoopRVObj */ class LoopRV : public ffi::ObjectRef { public: /*! \brief Constructor */ TVM_DLL LoopRV(); - TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(LoopRV, ffi::ObjectRef, LoopRVNode); + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(LoopRV, ffi::ObjectRef, LoopRVObj); }; /**************** Random variable: ExprRV ****************/ @@ -100,7 +100,7 @@ class LoopRV : public ffi::ObjectRef { /*! \brief An expr random variable */ using ExprRV = PrimExpr; -using ExprRVNode = ExprNode; +using ExprRVObj = ExprObj; /**************** The Schedule class ****************/ @@ -200,7 +200,7 @@ class ScheduleObj : public ffi::Object { * \param stmt The statement to be looked up * \return The corresponding block/loop sref */ - virtual StmtSRef GetSRef(const StmtNode* stmt) const; + virtual StmtSRef GetSRef(const StmtObj* stmt) const; /*! * \brief Get the block/loop sref corresponding to the specific statement * \param stmt The statement to be looked up diff --git a/include/tvm/s_tir/schedule/state.h b/include/tvm/s_tir/schedule/state.h index 27d3d0d2a35e..6d2dbf19751b 100644 --- a/include/tvm/s_tir/schedule/state.h +++ b/include/tvm/s_tir/schedule/state.h @@ -111,7 +111,7 @@ class ScheduleStateObj : public ffi::Object { */ std::unordered_map block_info; /*! \brief The reverse mapping from block/for-loop to their corresponding srefs */ - std::unordered_map stmt2ref; + std::unordered_map stmt2ref; /*! * \brief Do extra correctness checking after the class creation * and each time after calling the Replace method. diff --git a/include/tvm/s_tir/schedule/trace.h b/include/tvm/s_tir/schedule/trace.h index dba4f6c80c0c..aa7f3efce32e 100644 --- a/include/tvm/s_tir/schedule/trace.h +++ b/include/tvm/s_tir/schedule/trace.h @@ -57,7 +57,7 @@ using FTraceDecisionProvider = * their decisions accordingly. Re-sampling is invoked if a sampling instruction doesn't have its * corresponding decision; Otherwise the existing decision will be reused accordingly. */ -class TraceNode : public ffi::Object { +class TraceObj : public ffi::Object { public: /*! \brief The instructions invoked so far in the program execution */ ffi::Array insts; @@ -66,13 +66,13 @@ class TraceNode : public ffi::Object { static void RegisterReflection() { namespace refl = tvm::ffi::reflection; - refl::ObjectDef() - .def_ro("insts", &TraceNode::insts) - .def_ro("decisions", &TraceNode::decisions); + refl::ObjectDef() + .def_ro("insts", &TraceObj::insts) + .def_ro("decisions", &TraceObj::decisions); } static constexpr const bool _type_mutable = true; - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("s_tir.Trace", TraceNode, ffi::Object); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("s_tir.Trace", TraceObj, ffi::Object); public: /*! @@ -139,8 +139,8 @@ class TraceNode : public ffi::Object { }; /*! - * \brief Managed reference to TraceNode - * \sa TraceNode + * \brief Managed reference to TraceObj + * \sa TraceObj */ class Trace : public ffi::ObjectRef { public: @@ -159,7 +159,7 @@ class Trace : public ffi::ObjectRef { */ static void ApplyJSONToSchedule(ffi::ObjectRef json, Schedule sch); - TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Trace, ffi::ObjectRef, TraceNode); + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Trace, ffi::ObjectRef, TraceObj); }; } // namespace s_tir diff --git a/include/tvm/s_tir/stmt.h b/include/tvm/s_tir/stmt.h index 2315dfe6cd11..c44b0e9844ae 100644 --- a/include/tvm/s_tir/stmt.h +++ b/include/tvm/s_tir/stmt.h @@ -47,11 +47,11 @@ using TensorAllocatedAddresses = ffi::Array * low-level hardware primitives in the IR and defer the check after the sequence of * transformations. */ -class MatchTensorRegionNode : public ffi::Object { +class MatchTensorRegionObj : public ffi::Object { public: - explicit MatchTensorRegionNode(ffi::UnsafeInit tag) : tensor(tag), source(tag) {} + explicit MatchTensorRegionObj(ffi::UnsafeInit tag) : tensor(tag), source(tag) {} - MatchTensorRegionNode(tirx::TensorVar tensor, TensorRegion source) + MatchTensorRegionObj(tirx::TensorVar tensor, TensorRegion source) : tensor(std::move(tensor)), source(std::move(source)) {} /*! \brief The target tensor. */ @@ -61,23 +61,22 @@ class MatchTensorRegionNode : public ffi::Object { static void RegisterReflection() { namespace refl = tvm::ffi::reflection; - refl::ObjectDef() - .def_ro("tensor", &MatchTensorRegionNode::tensor, - refl::AttachFieldFlag::SEqHashDefPattern()) - .def_ro("source", &MatchTensorRegionNode::source); + refl::ObjectDef() + .def_ro("tensor", &MatchTensorRegionObj::tensor, refl::AttachFieldFlag::SEqHashDefPattern()) + .def_ro("source", &MatchTensorRegionObj::source); } static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode; - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("s_tir.MatchTensorRegion", MatchTensorRegionNode, ffi::Object); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("s_tir.MatchTensorRegion", MatchTensorRegionObj, ffi::Object); }; /*! - * \brief Managed reference to MatchTensorRegionNode. - * \sa MatchTensorRegionNode + * \brief Managed reference to MatchTensorRegionObj. + * \sa MatchTensorRegionObj */ class MatchTensorRegion : public ffi::ObjectRef { public: - explicit MatchTensorRegion(ffi::UnsafeInit tag, ffi::ObjectPtr data) + explicit MatchTensorRegion(ffi::UnsafeInit tag, ffi::ObjectPtr data) : ffi::ObjectRef(tag) { data_ = std::move(data); } @@ -85,8 +84,8 @@ class MatchTensorRegion : public ffi::ObjectRef { TVM_DLL explicit MatchTensorRegion(tirx::TensorVar tensor, TensorRegion source); TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(MatchTensorRegion, ffi::ObjectRef, - MatchTensorRegionNode); - TVM_DEFINE_OBJECT_REF_COW_METHOD(MatchTensorRegionNode); + MatchTensorRegionObj); + TVM_DEFINE_OBJECT_REF_COW_METHOD(MatchTensorRegionObj); }; /*! @@ -110,11 +109,11 @@ class MatchTensorRegion : public ffi::ObjectRef { * * \endcode */ -class SBlockNode : public StmtNode { +class SBlockObj : public StmtObj { public: - explicit SBlockNode(ffi::UnsafeInit tag) : body(tag) {} + explicit SBlockObj(ffi::UnsafeInit tag) : body(tag) {} - explicit SBlockNode(SeqStmt body) : body(std::move(body)) {} + explicit SBlockObj(SeqStmt body) : body(std::move(body)) {} /*! \brief The variables of the block. */ ffi::Array iter_vars; @@ -143,24 +142,24 @@ class SBlockNode : public StmtNode { static void RegisterReflection() { namespace refl = tvm::ffi::reflection; - refl::ObjectDef() - .def_ro("iter_vars", &SBlockNode::iter_vars) - .def_ro("reads", &SBlockNode::reads) - .def_ro("writes", &SBlockNode::writes) - .def_ro("name_hint", &SBlockNode::name_hint, refl::AttachFieldFlag::SEqHashIgnore()) - .def_ro("alloc_tensors", &SBlockNode::alloc_tensors, + refl::ObjectDef() + .def_ro("iter_vars", &SBlockObj::iter_vars) + .def_ro("reads", &SBlockObj::reads) + .def_ro("writes", &SBlockObj::writes) + .def_ro("name_hint", &SBlockObj::name_hint, refl::AttachFieldFlag::SEqHashIgnore()) + .def_ro("alloc_tensors", &SBlockObj::alloc_tensors, refl::AttachFieldFlag::SEqHashDefSimple()) - .def_ro("match_tensors", &SBlockNode::match_tensors) - .def_ro("annotations", &SBlockNode::annotations) - .def_ro("init", &SBlockNode::init) - .def_ro("body", &SBlockNode::body); + .def_ro("match_tensors", &SBlockObj::match_tensors) + .def_ro("annotations", &SBlockObj::annotations) + .def_ro("init", &SBlockObj::init) + .def_ro("body", &SBlockObj::body); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("s_tir.SBlock", SBlockNode, StmtNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("s_tir.SBlock", SBlockObj, StmtObj); }; /*! - * \brief Managed reference to SBlockNode. - * \sa SBlockNode + * \brief Managed reference to SBlockObj. + * \sa SBlockObj */ class SBlock : public Stmt { public: @@ -177,20 +176,20 @@ class SBlock : public Stmt { ffi::Array alloc_tensors = ffi::Array(), Location loc = UnknownLoc()); - explicit SBlock(ffi::ObjectPtr node) : Stmt(std::move(node)) {} + explicit SBlock(ffi::ObjectPtr node) : Stmt(std::move(node)) {} - TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(SBlock, Stmt, SBlockNode); - TVM_DEFINE_OBJECT_REF_COW_METHOD(SBlockNode); + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(SBlock, Stmt, SBlockObj); + TVM_DEFINE_OBJECT_REF_COW_METHOD(SBlockObj); }; /*! * \brief A block realization node represents execution of the block at the binding values. */ -class SBlockRealizeNode : public StmtNode { +class SBlockRealizeObj : public StmtObj { public: - explicit SBlockRealizeNode(ffi::UnsafeInit tag) : predicate(tag), block(tag) {} + explicit SBlockRealizeObj(ffi::UnsafeInit tag) : predicate(tag), block(tag) {} - SBlockRealizeNode(PrimExpr predicate, SBlock block) + SBlockRealizeObj(PrimExpr predicate, SBlock block) : predicate(std::move(predicate)), block(std::move(block)) {} /*! \brief The corresponding values of the iter vars. */ @@ -205,12 +204,12 @@ class SBlockRealizeNode : public StmtNode { static void RegisterReflection() { namespace refl = tvm::ffi::reflection; - refl::ObjectDef() - .def_ro("iter_values", &SBlockRealizeNode::iter_values) - .def_ro("predicate", &SBlockRealizeNode::predicate) - .def_ro("block", &SBlockRealizeNode::block); + refl::ObjectDef() + .def_ro("iter_values", &SBlockRealizeObj::iter_values) + .def_ro("predicate", &SBlockRealizeObj::predicate) + .def_ro("block", &SBlockRealizeObj::block); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("s_tir.SBlockRealize", SBlockRealizeNode, StmtNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("s_tir.SBlockRealize", SBlockRealizeObj, StmtObj); }; /*! @@ -222,10 +221,10 @@ class SBlockRealize : public Stmt { TVM_DLL explicit SBlockRealize(ffi::Array iter_values, PrimExpr predicate, SBlock block, Location loc = UnknownLoc()); - explicit SBlockRealize(ffi::ObjectPtr node) : Stmt(std::move(node)) {} + explicit SBlockRealize(ffi::ObjectPtr node) : Stmt(std::move(node)) {} - TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(SBlockRealize, Stmt, SBlockRealizeNode); - TVM_DEFINE_OBJECT_REF_COW_METHOD(SBlockRealizeNode); + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(SBlockRealize, Stmt, SBlockRealizeObj); + TVM_DEFINE_OBJECT_REF_COW_METHOD(SBlockRealizeObj); }; /*! \brief Region marking copies eligible for asynchronous lowering. */ @@ -386,7 +385,7 @@ inline constexpr ffi::reflection::AttrKey> kExplicitReadRegi inline constexpr ffi::reflection::AttrKey> kExplicitWriteRegion{ "explicit_write_region"}; -/*! \brief ,ark a ForNode represent an irregular loop of non-structural control flow edges. */ +/*! \brief ,ark a ForObj represent an irregular loop of non-structural control flow edges. */ inline constexpr ffi::reflection::AttrKey kIrregularLoopMark{"irregular_loop_mark"}; /*! \brief Mark auto copy for memhammer */ @@ -414,7 +413,7 @@ inline constexpr ffi::reflection::AttrKey> kMetaScheduleWrit } // namespace attr /*! \brief Interpret and validate S-TIR thread placement on a shared loop. */ -inline ffi::Optional GetThreadBinding(const ForNode* loop) { +inline ffi::Optional GetThreadBinding(const ForObj* loop) { if (auto tag = loop->annotations.Get(tvm::s_tir::attr::kThreadBinding.name)) { TVM_FFI_CHECK(loop->kind == ForKind::kParallel, ValueError) << "thread_binding requires a parallel loop"; diff --git a/include/tvm/s_tir/stmt_functor.h b/include/tvm/s_tir/stmt_functor.h index cf089d6b53ce..2af7059e710f 100644 --- a/include/tvm/s_tir/stmt_functor.h +++ b/include/tvm/s_tir/stmt_functor.h @@ -45,10 +45,10 @@ class StmtFunctor : public tirx::StmtFunctorDispatchDefault_(op, std::forward(args)...); } - virtual R Dispatch_(const SBlockRealizeNode* op, Args... args) { + virtual R Dispatch_(const SBlockRealizeObj* op, Args... args) { return this->DispatchDefault_(op, std::forward(args)...); } @@ -57,8 +57,8 @@ class StmtFunctor : public tirx::StmtFunctor(vtable); - Parent::template SetDispatch(vtable); + Parent::template SetDispatch(vtable); + Parent::template SetDispatch(vtable); } }; @@ -80,14 +80,14 @@ class TVM_DLL StmtExprVisitor : public tirx::StmtExprVisitor { using tirx::StmtExprVisitor::Visit; using tirx::StmtExprVisitor::Visit_; - virtual ffi::Optional Visit_(const SBlockNode* op); - virtual ffi::Optional Visit_(const SBlockRealizeNode* op); + virtual ffi::Optional Visit_(const SBlockObj* op); + virtual ffi::Optional Visit_(const SBlockRealizeObj* op); // Shared native traversal for specialized TIRX helpers extended by S-TIR. - static ffi::Optional VisitBlock(tirx::StmtExprVisitor* visitor, - const SBlockNode* op); - static ffi::Optional VisitBlockRealize(tirx::StmtExprVisitor* visitor, - const SBlockRealizeNode* op); + static ffi::Optional VisitBlock(ffi::StructuralVisitorObj* visitor, + const SBlockObj* op); + static ffi::Optional VisitBlockRealize(ffi::StructuralVisitorObj* visitor, + const SBlockRealizeObj* op); protected: explicit StmtExprVisitor(const VTable* vtable) : tirx::StmtExprVisitor(vtable) {} @@ -110,15 +110,14 @@ class TVM_DLL StmtExprMutator : public tirx::StmtExprMutator { using tirx::StmtExprMutator::Mutate; using tirx::StmtExprMutator::Mutate_; - virtual UnchangedOr Mutate_(const SBlockNode* op, InplaceMode inplace_mode); - virtual UnchangedOr Mutate_(const SBlockRealizeNode* op, InplaceMode inplace_mode); + virtual UnchangedOr Mutate_(const SBlockObj* op, InplaceMode inplace_mode); + virtual UnchangedOr Mutate_(const SBlockRealizeObj* op, InplaceMode inplace_mode); // Share block ownership and binder rules with specialized TIRX helpers. - static UnchangedOr MutateBlock(tirx::StmtExprMutator* mutator, const SBlockNode* op, + static UnchangedOr MutateBlock(ffi::StructuralMutatorObj* mutator, const SBlockObj* op, InplaceMode inplace_mode); - static UnchangedOr MutateBlockRealize(tirx::StmtExprMutator* mutator, - const SBlockRealizeNode* op, - InplaceMode inplace_mode); + static UnchangedOr MutateBlockRealize(ffi::StructuralMutatorObj* mutator, + const SBlockRealizeObj* op, InplaceMode inplace_mode); protected: explicit StmtExprMutator(const VTable* vtable) : tirx::StmtExprMutator(vtable) {} diff --git a/include/tvm/s_tir/stmt_functor_with_analyzer.h b/include/tvm/s_tir/stmt_functor_with_constraint_context.h similarity index 50% rename from include/tvm/s_tir/stmt_functor_with_analyzer.h rename to include/tvm/s_tir/stmt_functor_with_constraint_context.h index d1452c561c5b..9f0cb4fda45c 100644 --- a/include/tvm/s_tir/stmt_functor_with_analyzer.h +++ b/include/tvm/s_tir/stmt_functor_with_constraint_context.h @@ -18,58 +18,67 @@ */ /*! - * \file tvm/s_tir/stmt_functor_with_analyzer.h - * \brief Statement and expression functors with an analyzer context. + * \file tvm/s_tir/stmt_functor_with_constraint_context.h + * \brief Statement and expression functors with a constraint context. */ -#ifndef TVM_S_TIR_STMT_FUNCTOR_WITH_ANALYZER_H_ -#define TVM_S_TIR_STMT_FUNCTOR_WITH_ANALYZER_H_ +#ifndef TVM_S_TIR_STMT_FUNCTOR_WITH_CONSTRAINT_CONTEXT_H_ +#define TVM_S_TIR_STMT_FUNCTOR_WITH_CONSTRAINT_CONTEXT_H_ #include #include -#include +#include namespace tvm { namespace s_tir { -class TVM_DLL StmtExprVisitorWithAnalyzer : public tirx::StmtExprVisitorWithAnalyzer { +class TVM_DLL StmtExprVisitorWithConstraintContext + : public tirx::StmtExprVisitorWithConstraintContext { public: - using Parent = tirx::StmtExprVisitorWithAnalyzer; + using Parent = tirx::StmtExprVisitorWithConstraintContext; using Parent::Visit; using Parent::Visit_; - TVM_DEFINE_OBJECT_FUNCTOR_DEFAULT_CONSTRUCTOR(StmtExprVisitorWithAnalyzer, Parent) - virtual ffi::Optional Visit_(const SBlockNode* op); - virtual ffi::Optional Visit_(const SBlockRealizeNode* op) { + PrimExpr Simplify(const PrimExpr& expr) { return GetSyncedAnalyzer()->Simplify(expr); } + TVM_DEFINE_OBJECT_FUNCTOR_DEFAULT_CONSTRUCTOR(StmtExprVisitorWithConstraintContext, Parent) + virtual ffi::Optional Visit_(const SBlockObj* op); + virtual ffi::Optional Visit_(const SBlockRealizeObj* op) { return s_tir::StmtExprVisitor::VisitBlockRealize(this, op); } protected: + sym::AnalyzerObj* GetSyncedAnalyzer() { + analyzer_->SyncConstraintContext(constraints); + return analyzer_.get(); + } + sym::Analyzer analyzer_; + static void InitVTable(VTable* vtable) { Parent::InitVTable(vtable); - SetDispatch(vtable); - SetDispatch(vtable); + SetDispatch(vtable); + SetDispatch(vtable); } }; -class TVM_DLL StmtExprMutatorWithAnalyzer : public tirx::StmtExprMutatorWithAnalyzer { +class TVM_DLL StmtExprMutatorWithConstraintContext + : public tirx::StmtExprMutatorWithConstraintContext { public: - using Parent = tirx::StmtExprMutatorWithAnalyzer; + using Parent = tirx::StmtExprMutatorWithConstraintContext; using Parent::Mutate; using Parent::Mutate_; - explicit StmtExprMutatorWithAnalyzer(const sym::Analyzer& analyzer) - : StmtExprMutatorWithAnalyzer(analyzer.get()) {} - explicit StmtExprMutatorWithAnalyzer(sym::AnalyzerObj* analyzer) - : Parent(analyzer, GlobalVTable()) {} - virtual UnchangedOr Mutate_(const SBlockNode* op, InplaceMode inplace_mode); - virtual UnchangedOr Mutate_(const SBlockRealizeNode* op, InplaceMode inplace_mode) { + explicit StmtExprMutatorWithConstraintContext(const sym::ConstraintContext& context) + : Parent(GlobalVTable()) { + constraints = context->Snapshot(); + } + virtual UnchangedOr Mutate_(const SBlockObj* op, InplaceMode inplace_mode); + virtual UnchangedOr Mutate_(const SBlockRealizeObj* op, InplaceMode inplace_mode) { return s_tir::StmtExprMutator::MutateBlockRealize(this, op, inplace_mode); } protected: static void InitVTable(VTable* vtable) { Parent::InitVTable(vtable); - SetDispatch(vtable); - SetDispatch(vtable); + SetDispatch(vtable); + SetDispatch(vtable); } static const VTable* GlobalVTable() { static const VTable table = [] { @@ -84,4 +93,4 @@ class TVM_DLL StmtExprMutatorWithAnalyzer : public tirx::StmtExprMutatorWithAnal } // namespace s_tir } // namespace tvm -#endif // TVM_S_TIR_STMT_FUNCTOR_WITH_ANALYZER_H_ +#endif // TVM_S_TIR_STMT_FUNCTOR_WITH_CONSTRAINT_CONTEXT_H_ diff --git a/include/tvm/s_tir/tensor_intrin.h b/include/tvm/s_tir/tensor_intrin.h index c583ec87fba6..af0ee6e7d4d4 100644 --- a/include/tvm/s_tir/tensor_intrin.h +++ b/include/tvm/s_tir/tensor_intrin.h @@ -33,11 +33,11 @@ namespace s_tir { /*! * \brief Tensor intrinsics for tensorization */ -class TensorIntrinNode : public ffi::Object { +class TensorIntrinObj : public ffi::Object { public: - explicit TensorIntrinNode(ffi::UnsafeInit tag) : desc(tag), impl(tag) {} + explicit TensorIntrinObj(ffi::UnsafeInit tag) : desc(tag), impl(tag) {} - TensorIntrinNode(tirx::Function desc, tirx::Function impl) + TensorIntrinObj(tirx::Function desc, tirx::Function impl) : desc(std::move(desc)), impl(std::move(impl)) {} /*! \brief The function to describe the computation. */ @@ -47,15 +47,15 @@ class TensorIntrinNode : public ffi::Object { static void RegisterReflection() { namespace refl = tvm::ffi::reflection; - refl::ObjectDef() - .def_ro("desc", &TensorIntrinNode::desc) - .def_ro("impl", &TensorIntrinNode::impl); + refl::ObjectDef() + .def_ro("desc", &TensorIntrinObj::desc) + .def_ro("impl", &TensorIntrinObj::impl); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("s_tir.TensorIntrin", TensorIntrinNode, ffi::Object); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("s_tir.TensorIntrin", TensorIntrinObj, ffi::Object); }; /*! - * \brief Managed reference to TensorIntrinNode. + * \brief Managed reference to TensorIntrinObj. */ class TensorIntrin : public ffi::ObjectRef { public: @@ -88,7 +88,7 @@ class TensorIntrin : public ffi::ObjectRef { */ TVM_DLL static ffi::Optional Get(ffi::String name, bool allow_missing = false); - TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(TensorIntrin, ffi::ObjectRef, TensorIntrinNode); + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(TensorIntrin, ffi::ObjectRef, TensorIntrinObj); }; } // namespace s_tir diff --git a/include/tvm/s_tir/utils.h b/include/tvm/s_tir/utils.h index f70b39892ffc..e4048263271f 100644 --- a/include/tvm/s_tir/utils.h +++ b/include/tvm/s_tir/utils.h @@ -49,7 +49,7 @@ namespace tirx { */ #define TVM_SREF_TO_SBLOCK(SRef) \ [&]() { \ - auto result = TVM_SREF_AS_OR_ERR(result, (SRef), ::tvm::s_tir::SBlockNode) \ + auto result = TVM_SREF_AS_OR_ERR(result, (SRef), ::tvm::s_tir::SBlockObj) \ << "Expects StmtSRef `" << #SRef << "` points to `Block`, but gets: " \ << ((SRef)->stmt ? (SRef)->stmt->GetTypeKey() : "None"); \ return result; \ @@ -65,7 +65,7 @@ namespace tirx { */ #define TVM_SREF_TO_FOR(SRef) \ [&]() { \ - auto result = TVM_SREF_AS_OR_ERR(result, (SRef), ::tvm::ForNode) \ + auto result = TVM_SREF_AS_OR_ERR(result, (SRef), ::tvm::ForObj) \ << "Expects StmtSRef `" << #SRef << "` points to `Loop`, but gets: " \ << ((SRef)->stmt ? (SRef)->stmt->GetTypeKey() : "None"); \ return result; \ @@ -104,16 +104,16 @@ namespace tirx { * \param include_loops Ignore ForNodes if this value is false * \note The method is NOP for statements that are not schedulable, i.e. not For or s_tir::SBlock */ -inline void SetSeqIndex(std::unordered_map& stmt2ref, // NOLINT(*) +inline void SetSeqIndex(std::unordered_map& stmt2ref, // NOLINT(*) const Stmt& stmt, int seq_index, bool include_loops = true) { - if (const auto* realize = stmt.as()) { - const s_tir::SBlockNode* block = realize->block.get(); + if (const auto* realize = stmt.as()) { + const s_tir::SBlockObj* block = realize->block.get(); TVM_FFI_ICHECK(stmt2ref.count(block)); stmt2ref.at(block)->seq_index = seq_index; - } else if (const auto* block = stmt.as()) { + } else if (const auto* block = stmt.as()) { TVM_FFI_ICHECK(stmt2ref.count(block)); stmt2ref.at(block)->seq_index = seq_index; - } else if (const auto* loop = stmt.as()) { + } else if (const auto* loop = stmt.as()) { if (!include_loops) return; TVM_FFI_ICHECK(stmt2ref.count(loop)); stmt2ref.at(loop)->seq_index = seq_index; @@ -127,8 +127,8 @@ inline void SetSeqIndex(std::unordered_map& stmt2ref, * \param include_loops Ignore ForNodes if this value is false */ inline void SetSeqIndexInChildren( - std::unordered_map& stmt2ref, // NOLINT(*) - const SeqStmtNode* seq_stmt, bool include_loops = true) { + std::unordered_map& stmt2ref, // NOLINT(*) + const SeqStmtObj* seq_stmt, bool include_loops = true) { int i = 0; for (const Stmt& stmt : seq_stmt->seq) { SetSeqIndex(stmt2ref, stmt, i, include_loops); diff --git a/include/tvm/te/operation.h b/include/tvm/te/operation.h index f6cb430af84b..02368ae33787 100644 --- a/include/tvm/te/operation.h +++ b/include/tvm/te/operation.h @@ -49,7 +49,7 @@ namespace te { * \brief A commutative reducer node to represent a commutative * binary operator with identity element */ -class CommReducerNode : public ffi::Object { +class CommReducerObj : public ffi::Object { public: /*! \brief The left argument of reducer */ ffi::Array lhs; @@ -73,36 +73,35 @@ class CommReducerNode : public ffi::Object { static void RegisterReflection() { namespace refl = tvm::ffi::reflection; - refl::ObjectDef() - .def_ro("lhs", &CommReducerNode::lhs, refl::AttachFieldFlag::SEqHashDefPattern()) - .def_ro("rhs", &CommReducerNode::rhs, refl::AttachFieldFlag::SEqHashDefPattern()) - .def_ro("result", &CommReducerNode::result) - .def_ro("identity_element", &CommReducerNode::identity_element) - .def_ro("loc", &CommReducerNode::loc, refl::AttachFieldFlag::SEqHashIgnore()); + refl::ObjectDef() + .def_ro("lhs", &CommReducerObj::lhs, refl::AttachFieldFlag::SEqHashDefPattern()) + .def_ro("rhs", &CommReducerObj::rhs, refl::AttachFieldFlag::SEqHashDefPattern()) + .def_ro("result", &CommReducerObj::result) + .def_ro("identity_element", &CommReducerObj::identity_element) + .def_ro("loc", &CommReducerObj::loc, refl::AttachFieldFlag::SEqHashIgnore()); } static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode; - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("te.CommReducer", CommReducerNode, ffi::Object); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("te.CommReducer", CommReducerObj, ffi::Object); }; /*! - * \brief Managed reference to CommReducerNode - * \sa CommReducerNode + * \brief Managed reference to CommReducerObj + * \sa CommReducerObj */ class CommReducer : public ffi::ObjectRef { public: TVM_DLL CommReducer(ffi::Array lhs, ffi::Array rhs, ffi::Array result, ffi::Array identity_element, Location loc = UnknownLoc()); - TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(CommReducer, ffi::ObjectRef, CommReducerNode); + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(CommReducer, ffi::ObjectRef, CommReducerObj); }; /*! \brief Reduction operator */ -class ReduceNode : public OpaqueExprNode { +class ReduceObj : public OpaqueExprObj { public: - explicit ReduceNode(CommReducer combiner, PrimExpr condition) + explicit ReduceObj(CommReducer combiner, PrimExpr condition) : combiner(std::move(combiner)), condition(std::move(condition)) {} - explicit ReduceNode(ffi::UnsafeInit) - : combiner(ffi::UnsafeInit{}), condition(ffi::UnsafeInit{}) {} + explicit ReduceObj(ffi::UnsafeInit) : combiner(ffi::UnsafeInit{}), condition(ffi::UnsafeInit{}) {} /*! \brief The commutative combiner */ CommReducer combiner; @@ -122,31 +121,31 @@ class ReduceNode : public OpaqueExprNode { static void RegisterReflection() { namespace refl = tvm::ffi::reflection; - refl::ObjectDef() - .def_ro("combiner", &ReduceNode::combiner) - .def_ro("source", &ReduceNode::source) - .def_ro("init", &ReduceNode::init) - .def_ro("axis", &ReduceNode::axis) - .def_ro("condition", &ReduceNode::condition) - .def_ro("value_index", &ReduceNode::value_index); + refl::ObjectDef() + .def_ro("combiner", &ReduceObj::combiner) + .def_ro("source", &ReduceObj::source) + .def_ro("init", &ReduceObj::init) + .def_ro("axis", &ReduceObj::axis) + .def_ro("condition", &ReduceObj::condition) + .def_ro("value_index", &ReduceObj::value_index); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("te.Reduce", ReduceNode, OpaqueExprNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("te.Reduce", ReduceObj, OpaqueExprObj); }; /*! - * \brief Managed reference to ReduceNode - * \sa ReduceNode + * \brief Managed reference to ReduceObj + * \sa ReduceObj */ class Reduce : public PrimExpr { public: TVM_DLL Reduce(CommReducer combiner, ffi::Array src, ffi::Array rdom, ffi::Optional condition, int value_index, ffi::Array init, Location loc = UnknownLoc()); - explicit Reduce(ffi::ObjectPtr node) : PrimExpr(std::move(node)) {} + explicit Reduce(ffi::ObjectPtr node) : PrimExpr(std::move(node)) {} - TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Reduce, PrimExpr, ReduceNode); + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Reduce, PrimExpr, ReduceObj); static constexpr bool _type_container_is_exact = true; - TVM_DEFINE_OBJECT_REF_COW_METHOD(ReduceNode); + TVM_DEFINE_OBJECT_REF_COW_METHOD(ReduceObj); }; /*! @@ -163,7 +162,7 @@ struct TensorDom { /*! * \brief Base class of all operation nodes */ -class TVM_DLL OperationNode : public ffi::Object { +class TVM_DLL OperationObj : public ffi::Object { public: /*! \brief optional name of the operation */ std::string name; @@ -172,7 +171,7 @@ class TVM_DLL OperationNode : public ffi::Object { /*! \brief additional attributes of the operation*/ ffi::Map attrs; // virtual destructor. - virtual ~OperationNode() {} + virtual ~OperationObj() {} /*! \return number of outputs */ virtual int num_outputs() const = 0; /*! @@ -195,18 +194,18 @@ class TVM_DLL OperationNode : public ffi::Object { static void RegisterReflection() { namespace refl = tvm::ffi::reflection; - refl::ObjectDef() - .def_ro("name", &OperationNode::name) - .def_ro("tag", &OperationNode::tag) - .def_ro("attrs", &OperationNode::attrs); + refl::ObjectDef() + .def_ro("name", &OperationObj::name) + .def_ro("tag", &OperationObj::tag) + .def_ro("attrs", &OperationObj::attrs); } - TVM_FFI_DECLARE_OBJECT_INFO("te.Operation", OperationNode, ffi::Object); + TVM_FFI_DECLARE_OBJECT_INFO("te.Operation", OperationObj, ffi::Object); }; /*! * \brief A placeholder op represents an input placeholder. */ -class PlaceholderOpNode : public OperationNode { +class PlaceholderOpObj : public OperationObj { public: /*! \brief The shape of the input */ ffi::Array shape; @@ -220,34 +219,34 @@ class PlaceholderOpNode : public OperationNode { static void RegisterReflection() { namespace refl = tvm::ffi::reflection; - refl::ObjectDef() - .def_ro("shape", &PlaceholderOpNode::shape) - .def_ro("dtype", &PlaceholderOpNode::dtype); + refl::ObjectDef() + .def_ro("shape", &PlaceholderOpObj::shape) + .def_ro("dtype", &PlaceholderOpObj::dtype); } - TVM_FFI_DECLARE_OBJECT_INFO("te.PlaceholderOp", PlaceholderOpNode, OperationNode); + TVM_FFI_DECLARE_OBJECT_INFO("te.PlaceholderOp", PlaceholderOpObj, OperationObj); }; /*! - * \brief Managed reference to PlaceholderOpNode - * \sa PlaceholderOpNode + * \brief Managed reference to PlaceholderOpObj + * \sa PlaceholderOpObj */ class PlaceholderOp : public Operation { public: TVM_DLL PlaceholderOp(std::string name, ffi::Array shape, PrimType dtype); - explicit PlaceholderOp(ffi::UnsafeInit tag, ffi::ObjectPtr node) + explicit PlaceholderOp(ffi::UnsafeInit tag, ffi::ObjectPtr node) : Operation(tag) { data_ = std::move(node); } - TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(PlaceholderOp, Operation, PlaceholderOpNode); + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(PlaceholderOp, Operation, PlaceholderOpObj); }; /*! * \brief A Compute op that compute a tensor on certain domain. * This is the base class for ComputeOp (operating on a scalar at a time) */ -class TVM_DLL BaseComputeOpNode : public OperationNode { +class TVM_DLL BaseComputeOpObj : public OperationObj { public: /*! \brief IterVar on each axis */ ffi::Array axis; @@ -258,22 +257,22 @@ class TVM_DLL BaseComputeOpNode : public OperationNode { static void RegisterReflection() { namespace refl = tvm::ffi::reflection; - refl::ObjectDef() - .def_ro("axis", &BaseComputeOpNode::axis) - .def_ro("reduce_axis", &BaseComputeOpNode::reduce_axis); + refl::ObjectDef() + .def_ro("axis", &BaseComputeOpObj::axis) + .def_ro("reduce_axis", &BaseComputeOpObj::reduce_axis); } - TVM_FFI_DECLARE_OBJECT_INFO("te.BaseComputeOp", BaseComputeOpNode, OperationNode); + TVM_FFI_DECLARE_OBJECT_INFO("te.BaseComputeOp", BaseComputeOpObj, OperationObj); }; /*! * \brief A Compute op that compute a tensor on certain domain. */ -class TVM_DLL ComputeOpNode : public BaseComputeOpNode { +class TVM_DLL ComputeOpObj : public BaseComputeOpObj { public: /*! \brief the compute expression */ ffi::Array body; /*! \brief constructor */ - ComputeOpNode() {} + ComputeOpObj() {} // override functions int num_outputs() const final; PrimType output_dtype(size_t i) const final; @@ -281,28 +280,28 @@ class TVM_DLL ComputeOpNode : public BaseComputeOpNode { static void RegisterReflection() { namespace refl = tvm::ffi::reflection; - refl::ObjectDef().def_ro("body", &ComputeOpNode::body); + refl::ObjectDef().def_ro("body", &ComputeOpObj::body); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("te.ComputeOp", ComputeOpNode, BaseComputeOpNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("te.ComputeOp", ComputeOpObj, BaseComputeOpObj); }; /*! - * \brief Managed reference to ComputeOpNode - * \sa ComputeOpNode + * \brief Managed reference to ComputeOpObj + * \sa ComputeOpObj */ class ComputeOp : public Operation { public: TVM_DLL ComputeOp(std::string name, std::string tag, ffi::Map attrs, ffi::Array axis, ffi::Array body); - TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ComputeOp, Operation, ComputeOpNode); - TVM_DEFINE_OBJECT_REF_COW_METHOD(ComputeOpNode); + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ComputeOp, Operation, ComputeOpObj); + TVM_DEFINE_OBJECT_REF_COW_METHOD(ComputeOpObj); }; /*! * \brief Symbolic scan. */ -class ScanOpNode : public OperationNode { +class ScanOpObj : public OperationObj { public: /*! \brief IterVar to scan over */ IterVar scan_axis{ffi::UnsafeInit{}}; @@ -328,7 +327,7 @@ class ScanOpNode : public OperationNode { */ ffi::Array spatial_axis_; /*! \brief constructor */ - ScanOpNode() {} + ScanOpObj() {} // override behavior. int num_outputs() const final; PrimType output_dtype(size_t i) const final; @@ -337,20 +336,20 @@ class ScanOpNode : public OperationNode { static void RegisterReflection() { namespace refl = tvm::ffi::reflection; - refl::ObjectDef() - .def_ro("scan_axis", &ScanOpNode::scan_axis) - .def_ro("init", &ScanOpNode::init) - .def_ro("update", &ScanOpNode::update) - .def_ro("state_placeholder", &ScanOpNode::state_placeholder) - .def_ro("inputs", &ScanOpNode::inputs) - .def_ro("spatial_axis_", &ScanOpNode::spatial_axis_); + refl::ObjectDef() + .def_ro("scan_axis", &ScanOpObj::scan_axis) + .def_ro("init", &ScanOpObj::init) + .def_ro("update", &ScanOpObj::update) + .def_ro("state_placeholder", &ScanOpObj::state_placeholder) + .def_ro("inputs", &ScanOpObj::inputs) + .def_ro("spatial_axis_", &ScanOpObj::spatial_axis_); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("te.ScanOp", ScanOpNode, OperationNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("te.ScanOp", ScanOpObj, OperationObj); }; /*! - * \brief Managed reference to ScanOpNode - * \sa ScanOpNode + * \brief Managed reference to ScanOpObj + * \sa ScanOpObj */ class ScanOp : public Operation { public: @@ -359,16 +358,16 @@ class ScanOp : public Operation { ffi::Array init, ffi::Array update, ffi::Array state_placeholder, ffi::Array input); - TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ScanOp, Operation, ScanOpNode); + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ScanOp, Operation, ScanOpObj); }; /*! * \brief External computation that cannot be splitted. */ -class ExternOpNode : public OperationNode { +class ExternOpObj : public OperationObj { public: - explicit ExternOpNode(SeqStmt body) : body(std::move(body)) {} - explicit ExternOpNode(ffi::UnsafeInit) : body(ffi::UnsafeInit{}) {} + explicit ExternOpObj(SeqStmt body) : body(std::move(body)) {} + explicit ExternOpObj(ffi::UnsafeInit) : body(ffi::UnsafeInit{}) {} /*! \brief The input tensors */ ffi::Array inputs; @@ -387,18 +386,18 @@ class ExternOpNode : public OperationNode { static void RegisterReflection() { namespace refl = tvm::ffi::reflection; - refl::ObjectDef() - .def_ro("inputs", &ExternOpNode::inputs) - .def_ro("input_placeholders", &ExternOpNode::input_placeholders) - .def_ro("output_placeholders", &ExternOpNode::output_placeholders) - .def_ro("body", &ExternOpNode::body); + refl::ObjectDef() + .def_ro("inputs", &ExternOpObj::inputs) + .def_ro("input_placeholders", &ExternOpObj::input_placeholders) + .def_ro("output_placeholders", &ExternOpObj::output_placeholders) + .def_ro("body", &ExternOpObj::body); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("te.ExternOp", ExternOpNode, OperationNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("te.ExternOp", ExternOpObj, OperationObj); }; /*! - * \brief Managed reference to ExternOpNode - * \sa ExternOpNode + * \brief Managed reference to ExternOpObj + * \sa ExternOpObj */ class ExternOp : public Operation { public: @@ -406,7 +405,7 @@ class ExternOp : public Operation { ffi::Array inputs, ffi::Array input_placeholders, ffi::Array output_placeholders, SeqStmt body); - TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ExternOp, Operation, ExternOpNode); + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ExternOp, Operation, ExternOpObj); }; /*! @@ -519,15 +518,15 @@ inline Tensor compute(ffi::Array shape, } // inline function. -inline const OperationNode* Operation::operator->() const { - return static_cast(get()); +inline const OperationObj* Operation::operator->() const { + return static_cast(get()); } } // namespace te namespace ffi { template <> -inline constexpr bool object_ref_contains_v = true; +inline constexpr bool object_ref_contains_v = true; } // namespace ffi } // namespace tvm diff --git a/include/tvm/te/tensor.h b/include/tvm/te/tensor.h index 1fafd7e4a8c7..baa3b3572068 100644 --- a/include/tvm/te/tensor.h +++ b/include/tvm/te/tensor.h @@ -40,7 +40,7 @@ namespace tvm { namespace te { using s_tir::IterVar; -using s_tir::IterVarNode; +using s_tir::IterVarObj; using s_tir::IterVarType; using s_tir::IterVarType2String; using s_tir::kCommReduce; @@ -57,7 +57,7 @@ using sym::IntSet; using namespace tvm::tirx; // internal node container for Operation -class OperationNode; +class OperationObj; class Tensor; /*! \brief Operation that produces tensors */ @@ -69,7 +69,7 @@ class Operation : public ffi::ObjectRef { * \brief access the internal node container * \return the pointer to the internal node container */ - inline const OperationNode* operator->() const; + inline const OperationObj* operator->() const; /*! * \brief get the i-th output of the operation. * \param i the output index. @@ -77,11 +77,11 @@ class Operation : public ffi::ObjectRef { */ TVM_DLL Tensor output(size_t i) const; /*! \brief specify container node */ - using ContainerType = OperationNode; + using ContainerType = OperationObj; }; /*! \brief Opaque construction-time node that represents a tensor. */ -class TensorNode : public OpaqueExprNode { +class TensorObj : public OpaqueExprObj { public: /*! \brief The shape of the tensor */ ffi::Array shape; @@ -102,7 +102,7 @@ class TensorNode : public OpaqueExprNode { static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindConstTreeNode; - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("te.Tensor", TensorNode, OpaqueExprNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("te.Tensor", TensorObj, OpaqueExprObj); }; /*! @@ -219,9 +219,9 @@ class Tensor : public OpaqueExpr { */ inline Slice operator[](PrimExpr i) const { return Slice(*this, {i}); } - explicit Tensor(ffi::ObjectPtr node) : OpaqueExpr(std::move(node)) {} + explicit Tensor(ffi::ObjectPtr node) : OpaqueExpr(std::move(node)) {} - TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Tensor, OpaqueExpr, TensorNode); + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Tensor, OpaqueExpr, TensorObj); }; /*! \brief Return whether an expression calls the registered te.tensor_load operator. */ diff --git a/include/tvm/topi/broadcast.h b/include/tvm/topi/broadcast.h index 646b5cfc317d..3940ed09d15e 100644 --- a/include/tvm/topi/broadcast.h +++ b/include/tvm/topi/broadcast.h @@ -56,7 +56,7 @@ inline tvm::te::Tensor broadcast_to(const tvm::te::Tensor& t, TVM_FFI_ICHECK_EQ(output_shape.size(), bh.common_shape.size()); ffi::Array oshape; for (size_t i = 0; i < output_shape.size(); ++i) { - if (output_shape[i].as() == nullptr) { + if (output_shape[i].as() == nullptr) { oshape.push_back(output_shape[i]); } else { TVM_FFI_ICHECK(topi::detail::EqualCheck(output_shape[i], bh.common_shape[i])); diff --git a/include/tvm/topi/detail/broadcast.h b/include/tvm/topi/detail/broadcast.h index 670651f7bee7..99bbe6c35ef9 100644 --- a/include/tvm/topi/detail/broadcast.h +++ b/include/tvm/topi/detail/broadcast.h @@ -62,8 +62,8 @@ inline BroadcastHelper BroadcastShape(const tvm::ffi::Array& shap for (i = 1; i <= std::min(s1_size, s2_size); ++i) { // TODO(@icemelon9): Need to revisit this part - const IntImmNode* static_size1 = shape1[s1_size - i].as(); - const IntImmNode* static_size2 = shape2[s2_size - i].as(); + const IntImmObj* static_size1 = shape1[s1_size - i].as(); + const IntImmObj* static_size2 = shape2[s2_size - i].as(); PrimType common_type = CommonType(shape1[s1_size - i].ty(), shape2[s2_size - i].ty()); bh.all_vars.push_front(tvm::PrimVar("dim", common_type)); diff --git a/include/tvm/topi/detail/constant_utils.h b/include/tvm/topi/detail/constant_utils.h index 3f850f347925..8b1151836c8a 100644 --- a/include/tvm/topi/detail/constant_utils.h +++ b/include/tvm/topi/detail/constant_utils.h @@ -45,7 +45,7 @@ using namespace tvm::te; * * \return true if the given expr is a constant int or uint, false otherwise. */ -inline bool IsConstInt(PrimExpr expr) { return expr->IsInstance(); } +inline bool IsConstInt(PrimExpr expr) { return expr->IsInstance(); } /*! * \brief Test whether the given Array has every element as constant integer. @@ -58,7 +58,7 @@ inline bool IsConstInt(PrimExpr expr) { return expr->IsInstance inline bool IsConstIntArray(ffi::Array array) { bool IsConstInt = true; for (auto const& elem : array) { - IsConstInt &= !elem.defined() || elem->IsInstance(); + IsConstInt &= !elem.defined() || elem->IsInstance(); } return IsConstInt; } @@ -72,8 +72,8 @@ inline bool IsConstIntArray(ffi::Array array) { * \return The integer value. */ inline int64_t GetConstInt(PrimExpr expr) { - if (expr->IsInstance()) { - return static_cast(expr.as()->value); + if (expr->IsInstance()) { + return static_cast(expr.as()->value); } LOG(ERROR) << "expr must be a constant integer"; return -1; @@ -133,7 +133,7 @@ inline bool EqualCheck(PrimExpr lhs, PrimExpr rhs) { bool result = expr_equal(lhs, rhs); if (!result) { PrimExpr t = tvm::sym::Analyzer()->Simplify(lhs - rhs); - if (const IntImmNode* i = t.as()) { + if (const IntImmObj* i = t.as()) { result = i->value == 0; } } diff --git a/include/tvm/topi/detail/strided_slice.h b/include/tvm/topi/detail/strided_slice.h index a30940f09df1..a7099b78751d 100644 --- a/include/tvm/topi/detail/strided_slice.h +++ b/include/tvm/topi/detail/strided_slice.h @@ -97,7 +97,7 @@ inline ffi::Array StridedSliceCanonicalizeBegin(const ffi::Array begin_expr; for (size_t i = 0; i < axes.size(); ++i) { int64_t ax = axes[i]; - if (ishape[ax]->IsInstance()) { + if (ishape[ax]->IsInstance()) { int64_t dim_i = GetConstInt(ishape[ax]); int64_t begin_i = CanonicalizeIndex(begin[i], dim_i, strides[i]); begin_expr.push_back(MakeConst(dtype, begin_i)); @@ -131,9 +131,9 @@ inline ffi::Array StridedSliceOutputShape( for (size_t i = 0; i < axes.size(); ++i) { int64_t ax = axes[i]; - if (ishape[ax]->IsInstance()) { + if (ishape[ax]->IsInstance()) { const int64_t dim_i = GetConstInt(ishape[ax]); - TVM_FFI_ICHECK(begin_canonicalized[i]->IsInstance()); + TVM_FFI_ICHECK(begin_canonicalized[i]->IsInstance()); int64_t begin_i = GetConstInt(begin_canonicalized[i]); int64_t end_i = CanonicalizeIndex(end[i], dim_i, strides[i]); int interval = std::abs(end_i - begin_i); diff --git a/include/tvm/topi/detail/tensor_utils.h b/include/tvm/topi/detail/tensor_utils.h index 32ad3c27e7c9..11867287b58b 100644 --- a/include/tvm/topi/detail/tensor_utils.h +++ b/include/tvm/topi/detail/tensor_utils.h @@ -43,7 +43,7 @@ using namespace tvm::te; inline bool is_empty_shape(const ffi::Array& x) { bool is_empty = false; for (const auto& dim : x) { - if (auto int_dim = dim.as()) { + if (auto int_dim = dim.as()) { if (int_dim->value == 0) { is_empty = true; break; diff --git a/include/tvm/topi/nn.h b/include/tvm/topi/nn.h index 0fb2bbcc6164..8ceb1331d63e 100644 --- a/include/tvm/topi/nn.h +++ b/include/tvm/topi/nn.h @@ -636,11 +636,11 @@ inline tvm::te::Tensor batch_to_space_nd(const tvm::te::Tensor& data, strides.push_back(IntImm(index_ty, 1)); if (i > 0 && i <= num_block_dims) { // prepare begin and end index for spatial dimensions - const auto* begin_i_imm = crop_begin_list[i - 1].as(); + const auto* begin_i_imm = crop_begin_list[i - 1].as(); ffi::BigInt begin_i = begin_i_imm ? begin_i_imm->value : GetConstInt(crop_begin_list[i - 1]); - const auto* end_i_imm = crop_end_list[i - 1].as(); + const auto* end_i_imm = crop_end_list[i - 1].as(); ffi::BigInt end_i = end_i_imm ? end_i_imm->value : GetConstInt(crop_end_list[i - 1]); - const auto* out_i_imm = r_p_shape[i].as(); + const auto* out_i_imm = r_p_shape[i].as(); ffi::BigInt out_i = out_i_imm ? out_i_imm->value : GetConstInt(r_p_shape[i]); TVM_FFI_ICHECK_GT(out_i, (begin_i + end_i)) << "Incorrect crop sizes for (" << i << ")th dim, can not crop more than" @@ -650,7 +650,7 @@ inline tvm::te::Tensor batch_to_space_nd(const tvm::te::Tensor& data, } else { // ignore the batch and remaining dimension begin_idx.push_back(IntImm(index_ty, 0)); - const auto* extent = r_p_shape[i].as(); + const auto* extent = r_p_shape[i].as(); end_idx.push_back(IntImm(index_ty, extent ? extent->value : GetConstInt(r_p_shape[i]))); } } diff --git a/include/tvm/topi/nn/group_norm.h b/include/tvm/topi/nn/group_norm.h index 4cd116de0c5d..2eb3186a6e61 100644 --- a/include/tvm/topi/nn/group_norm.h +++ b/include/tvm/topi/nn/group_norm.h @@ -141,8 +141,8 @@ inline Tensor group_norm(const Tensor& data, const ffi::Optional& gamma, gamma_indices = {indices[channel_axis], indices[channel_axis + 1]}; auto mean = temp_x(non_reduce_indices) / reduce_extent; auto var = temp_x2(non_reduce_indices) / reduce_extent - mean * mean; - PrimExpr group_norm = (data_reshaped(indices) - mean) * - tvm::prim::rsqrt(var + MakeConst(PrimType(data->dtype), epsilon)); + PrimExpr group_norm = + (data_reshaped(indices) - mean) * tvm::prim::rsqrt(var + MakeConst(var.ty(), epsilon)); if (is_float16) { group_norm = prim::Cast(PrimType::Float(16), group_norm); } diff --git a/include/tvm/topi/nn/pooling.h b/include/tvm/topi/nn/pooling.h index 3e50ffd0d82c..c89e7afe6376 100644 --- a/include/tvm/topi/nn/pooling.h +++ b/include/tvm/topi/nn/pooling.h @@ -101,10 +101,10 @@ inline Tensor pool_grad_impl(const Tensor& out_grad, const Tensor& x, out_shape.Set(height_axis, out_height); out_shape.Set(width_axis, out_width); - const auto* padding_h0 = pad_top.as(); - const auto* padding_w0 = pad_left.as(); - const auto* padding_h1 = pad_bottom.as(); - const auto* padding_w1 = pad_right.as(); + const auto* padding_h0 = pad_top.as(); + const auto* padding_w0 = pad_left.as(); + const auto* padding_h1 = pad_bottom.as(); + const auto* padding_w1 = pad_right.as(); const bool do_pad = ((padding_h0 && padding_h0->value != 0) || (padding_w0 && padding_w0->value != 0)) || ((padding_h1 && padding_h1->value != 0) || (padding_w1 && padding_w1->value != 0)); @@ -559,8 +559,8 @@ inline Tensor pool_impl_nd(const Tensor& x, const ffi::Array& kernel_s pad_tail[i] += offset[i]; } - const auto* padding0 = pad_head[i].as(); - const auto* padding1 = pad_tail[i].as(); + const auto* padding0 = pad_head[i].as(); + const auto* padding1 = pad_tail[i].as(); do_pad = do_pad || (padding0 && padding0->value != 0) || (padding1 && padding1->value != 0); daxis.push_back(tvm::te::reduce_axis(Range(0, kernel[i]), "rv" + std::to_string(i))); diff --git a/include/tvm/topi/nn/rms_norm.h b/include/tvm/topi/nn/rms_norm.h index a4c6152a7ad4..8f04d1fe9b06 100644 --- a/include/tvm/topi/nn/rms_norm.h +++ b/include/tvm/topi/nn/rms_norm.h @@ -76,7 +76,7 @@ inline Tensor rms_norm(const Tensor& data, const Tensor& weight, const ffi::Arra } } auto output = tvm::prim::rsqrt(square_sum(non_reduce_indices) / reduce_extent + - MakeConst(PrimType(data_type), epsilon)); + MakeConst(PrimType(data_fp32->dtype), epsilon)); return output; }; auto rsqrt_shape = ffi::Array(); diff --git a/include/tvm/topi/transform.h b/include/tvm/topi/transform.h index d1ae30b2cd47..b2d53198f643 100644 --- a/include/tvm/topi/transform.h +++ b/include/tvm/topi/transform.h @@ -597,8 +597,8 @@ inline ffi::Array split_indices_array(const Tensor& x, ffi::Array(); - auto back_node = begin_ids.back().as(); + auto idx_node = idx.as(); + auto back_node = begin_ids.back().as(); if (idx_node && back_node) { TVM_FFI_ICHECK_GT(idx_node->value, back_node->value) << "split_indices must be sorted"; } @@ -653,7 +653,7 @@ inline PrimExpr DynamicCanonicalizeIndex(PrimExpr index, PrimExpr extent, PrimEx PrimExpr begin_range = tvm::if_then_else(stride < 0, -1, 0); PrimExpr end_range = tvm::if_then_else(stride < 0, extent - 1, extent); - if (!(index->IsInstance() && index.as_or_throw()->value >= 0)) { + if (!(index->IsInstance() && index.as_or_throw()->value >= 0)) { index = tvm::if_then_else(index < 0, index + extent, index); } @@ -670,8 +670,8 @@ inline int64_t StaticCanonicalizeIndex(int64_t index, int64_t extent, int64_t st } inline PrimExpr CanonicalizeIndex(PrimExpr index, PrimExpr extent, PrimExpr stride) { - if (index->IsInstance() && extent->IsInstance() && - stride->IsInstance()) { + if (index->IsInstance() && extent->IsInstance() && + stride->IsInstance()) { return tvm::IntImm( tvm::PrimType::Int(64), StaticCanonicalizeIndex(GetConstInt(index), GetConstInt(extent), GetConstInt(stride))); @@ -837,9 +837,9 @@ inline te::Tensor dynamic_strided_slice(const te::Tensor& x, const te::Tensor& b std::string name = "T_strided_slice_dynamic", std::string tag = topi::kInjective) { PrimType index_ty = begin->shape[0].ty(); - const int64_t num_dynamic_axes = static_cast(begin->shape[0].as()->value); - TVM_FFI_ICHECK_EQ(end->shape[0].as()->value, num_dynamic_axes); - TVM_FFI_ICHECK_EQ(strides->shape[0].as()->value, num_dynamic_axes); + const int64_t num_dynamic_axes = static_cast(begin->shape[0].as()->value); + TVM_FFI_ICHECK_EQ(end->shape[0].as()->value, num_dynamic_axes); + TVM_FFI_ICHECK_EQ(strides->shape[0].as()->value, num_dynamic_axes); ffi::Array begin_expr, end_expr, strides_expr; for (int64_t i = 0; i < num_dynamic_axes; ++i) { @@ -1152,7 +1152,7 @@ inline Tensor take(const Tensor& a, ffi::Variant indices, int TVM_FFI_ICHECK_LT(axis, a->shape.size()) << "axis out of bounds"; auto axis_dim = a->shape[axis]; auto indices_shape = [&]() -> ffi::Array { - if (auto tensor = indices.as()) { + if (auto tensor = indices.as()) { return tensor->shape; } else { return {}; @@ -1526,7 +1526,7 @@ inline Tensor gather(const Tensor& data, int axis, const Tensor& indices, } TVM_FFI_ICHECK_GE(axis, 0); TVM_FFI_ICHECK_LT(axis, ndim_d); - if (indices->shape[axis].as()) { + if (indices->shape[axis].as()) { size_t indices_dim_i = static_cast(GetConstInt(indices->shape[axis])); TVM_FFI_ICHECK_GE(indices_dim_i, 1); } @@ -2308,9 +2308,9 @@ inline te::Tensor dynamic_strided_slice(const te::Tensor& x, const te::Tensor& b TVM_FFI_ICHECK_EQ(begin.ndim(), 1); TVM_FFI_ICHECK_EQ(end.ndim(), 1); TVM_FFI_ICHECK_EQ(strides.ndim(), 1); - const auto* len_begin = begin->shape[0].as(); - const auto* len_end = end->shape[0].as(); - const auto* len_strides = strides->shape[0].as(); + const auto* len_begin = begin->shape[0].as(); + const auto* len_end = end->shape[0].as(); + const auto* len_strides = strides->shape[0].as(); TVM_FFI_ICHECK(len_begin); TVM_FFI_ICHECK(len_end); TVM_FFI_ICHECK(len_strides); diff --git a/jvm/native/src/main/native/org_apache_tvm_native_c_api.cc b/jvm/native/src/main/native/org_apache_tvm_native_c_api.cc index 152185fd1259..dafb6fd774a6 100644 --- a/jvm/native/src/main/native/org_apache_tvm_native_c_api.cc +++ b/jvm/native/src/main/native/org_apache_tvm_native_c_api.cc @@ -75,7 +75,7 @@ JNIEXPORT jint JNICALL Java_org_apache_tvm_LibInfo_shutdown(JNIEnv* env, jobject } JNIEXPORT jstring JNICALL Java_org_apache_tvm_LibInfo_tvmFFIGetLastError(JNIEnv* env, jobject obj) { - std::string err_msg = ::tvm::ffi::details::MoveFromSafeCallRaised().what(); + std::string err_msg = ::tvm::ffi::detail::MoveFromSafeCallRaised().what(); return env->NewStringUTF(err_msg.c_str()); } @@ -416,7 +416,7 @@ JNIEXPORT jint JNICALL Java_org_apache_tvm_LibInfo_tvmTensorEmpty( env->ReleaseLongArrayElements(jshape, shapeArray, 0); static auto fempty = tvm::ffi::Function::GetGlobalRequired("runtime.TensorAllocWithScope"); tvm::ffi::Tensor out = fempty(shape, dtype, device, nullptr).cast(); - void* handle = tvm::ffi::details::ObjectUnsafe::MoveObjectRefToTVMFFIObjectPtr(std::move(out)); + void* handle = tvm::ffi::detail::ObjectUnsafe::MoveObjectRefToTVMFFIObjectPtr(std::move(out)); setLongField(env, jret, reinterpret_cast(handle)); TVM_FFI_SAFE_CALL_END(); } diff --git a/pyproject.toml b/pyproject.toml index 66bf3b4e73fe..c92dd7fa7ad2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -20,7 +20,7 @@ requires = [ "scikit-build-core>=0.11", "setuptools-scm>=8", - "apache-tvm-tirx==0.1.0a2", + "apache-tvm-tirx==0.1.0a4", ] build-backend = "scikit_build_core.build" @@ -49,7 +49,7 @@ classifiers = [ "Topic :: Scientific/Engineering :: Artificial Intelligence", "Topic :: Software Development :: Libraries :: Python Modules", ] -dependencies = ["apache-tvm-tirx==0.1.0a2"] +dependencies = ["apache-tvm-tirx==0.1.0a4"] [project.scripts] tvm-config = "tvm.config:__main__" diff --git a/python/tvm/contrib/emcc.py b/python/tvm/contrib/emcc.py index 296f9ee9a21e..eb1b7c47632d 100644 --- a/python/tvm/contrib/emcc.py +++ b/python/tvm/contrib/emcc.py @@ -142,4 +142,5 @@ def create_wasm(output, objects, options=None, cc="emcc", libs=None): raise RuntimeError(msg) -create_wasm.object_format = "bc" +# Pass target objects across LLVM versions instead of version-specific bitcode. +create_wasm.object_format = "o" diff --git a/python/tvm/relax/dpl/pattern.py b/python/tvm/relax/dpl/pattern.py index f9e0c78018fd..75e7c7ab9fa8 100644 --- a/python/tvm/relax/dpl/pattern.py +++ b/python/tvm/relax/dpl/pattern.py @@ -317,7 +317,7 @@ class VarPattern(DFPattern): ---------- name_hint: str The name of the variable. Optional, if not provided, - the pattern will match any VarNode. + the pattern will match any VarObj. """ def __init__(self, name_hint: str = ""): @@ -332,7 +332,7 @@ class DataflowVarPattern(VarPattern): ---------- name_hint: str The name of the variable. Optional, if not provided, - the pattern will match any VarNode. + the pattern will match any VarObj. """ def __init__(self, name_hint: str = ""): @@ -347,7 +347,7 @@ class GlobalVarPattern(DFPattern): ---------- name_hint: str The name of the variable. Optional, if not provided, - the pattern will match any GlobalVarNode. + the pattern will match any GlobalVarObj. """ def __init__(self, name_hint: str = ""): @@ -362,7 +362,7 @@ class ExternFuncPattern(DFPattern): ---------- global_symbol: str The name of the function. Optional, if not provided, - the pattern will match any ExternFuncNode. + the pattern will match any ExternFuncObj. """ def __init__(self, global_symbol: str = ""): diff --git a/python/tvm/relax/expr_functor.py b/python/tvm/relax/expr_functor.py index 9341fdaea5d4..f6ced8ad80a0 100644 --- a/python/tvm/relax/expr_functor.py +++ b/python/tvm/relax/expr_functor.py @@ -647,7 +647,7 @@ def visit_var_def(self, var: Var) -> None: def visit_generic_const_(self, op: GenericConst) -> None: """Visit GenericConst. - Users can customized this function to overwrite VisitExpr_(const GenericConstNode* op) + Users can customized this function to overwrite VisitExpr_(const GenericConstObj* op) on the C++ side. Parameters @@ -660,7 +660,7 @@ def visit_generic_const_(self, op: GenericConst) -> None: def visit_tuple_(self, op: Tuple) -> None: """Visit Tuple. - Users can customized this function to overwrite VisitExpr_(const TupleNode* op) + Users can customized this function to overwrite VisitExpr_(const TupleObj* op) on the C++ side. Parameters @@ -673,7 +673,7 @@ def visit_tuple_(self, op: Tuple) -> None: def visit_var_(self, op: Var) -> None: """Visit Var. - Users can customized this function to overwrite VisitExpr_(const VarNode* op) + Users can customized this function to overwrite VisitExpr_(const VarObj* op) on the C++ side. Parameters @@ -686,7 +686,7 @@ def visit_var_(self, op: Var) -> None: def visit_dataflow_var_(self, op: DataflowVar) -> None: """Visit DataflowVar. - Users can customized this function to overwrite VisitExpr_(const DataflowVarNode* op) + Users can customized this function to overwrite VisitExpr_(const DataflowVarObj* op) on the C++ side. Parameters @@ -699,7 +699,7 @@ def visit_dataflow_var_(self, op: DataflowVar) -> None: def visit_shape_expr_(self, op: ShapeExpr) -> None: """Visit ShapeExpr. - Users can customized this function to overwrite VisitExpr_(const ShapeExprNode* op) + Users can customized this function to overwrite VisitExpr_(const ShapeExprObj* op) on the C++ side. Parameters @@ -712,7 +712,7 @@ def visit_shape_expr_(self, op: ShapeExpr) -> None: def visit_extern_func_(self, op: ExternFunc) -> None: """Visit ExternFunc. - Users can customized this function to overwrite VisitExpr_(const ExternFuncNode* op) + Users can customized this function to overwrite VisitExpr_(const ExternFuncObj* op) on the C++ side. Parameters @@ -725,7 +725,7 @@ def visit_extern_func_(self, op: ExternFunc) -> None: def visit_global_var_(self, op: GlobalVar) -> None: """Visit GlobalVar. - Users can customized this function to overwrite VisitExpr_(const GlobalVarNode* op) + Users can customized this function to overwrite VisitExpr_(const GlobalVarObj* op) on the C++ side. Parameters @@ -738,7 +738,7 @@ def visit_global_var_(self, op: GlobalVar) -> None: def visit_function_(self, op: Function) -> None: """Visit Function. - Users can customized this function to overwrite VisitExpr_(const FunctionNode* op) + Users can customized this function to overwrite VisitExpr_(const FunctionObj* op) on the C++ side. Parameters @@ -751,7 +751,7 @@ def visit_function_(self, op: Function) -> None: def visit_call_(self, op: Call) -> None: """Visit Call. - Users can customized this function to overwrite VisitExpr_(const CallNode* op) + Users can customized this function to overwrite VisitExpr_(const CallObj* op) on the C++ side. Parameters @@ -764,7 +764,7 @@ def visit_call_(self, op: Call) -> None: def visit_seq_expr_(self, op: SeqExpr) -> None: """Visit SeqExpr. - Users can customized this function to overwrite VisitExpr_(const SeqExprNode* op) + Users can customized this function to overwrite VisitExpr_(const SeqExprObj* op) on the C++ side. Parameters @@ -777,7 +777,7 @@ def visit_seq_expr_(self, op: SeqExpr) -> None: def visit_if_expr_(self, op: IfExpr) -> None: """Visit IfExpr. - Users can customized this function to overwrite VisitExpr_(const IfExprNode* op) + Users can customized this function to overwrite VisitExpr_(const IfExprObj* op) on the C++ side. Parameters @@ -790,7 +790,7 @@ def visit_if_expr_(self, op: IfExpr) -> None: def visit_op_(self, op: Op) -> None: """Visit Op. - Users can customized this function to overwrite VisitExpr_(const OpNode* op) + Users can customized this function to overwrite VisitExpr_(const OpObj* op) on the C++ side. Parameters @@ -803,7 +803,7 @@ def visit_op_(self, op: Op) -> None: def visit_tuple_getitem_(self, op: TupleGetItem) -> None: """Visit TupleGetItem. - Users can customized this function to overwrite VisitExpr_(const TupleGetItemNode* op) + Users can customized this function to overwrite VisitExpr_(const TupleGetItemObj* op) on the C++ side. Parameters @@ -827,7 +827,7 @@ def visit_expr_fallback_(self, op: Expr) -> None: def visit_string_imm_(self, op: StringImm) -> None: """Visit StringImm. - Users can customized this function to overwrite VisitExpr_(const StringImmNode* op) + Users can customized this function to overwrite VisitExpr_(const StringImmObj* op) on the C++ side. Parameters @@ -840,7 +840,7 @@ def visit_string_imm_(self, op: StringImm) -> None: def visit_data_type_imm_(self, op: DataTypeImm) -> None: """Visit DataTypeImm. - Users can customized this function to overwrite VisitExpr_(const DataTypeImmNode* op) + Users can customized this function to overwrite VisitExpr_(const DataTypeImmObj* op) on the C++ side. Parameters @@ -853,7 +853,7 @@ def visit_data_type_imm_(self, op: DataTypeImm) -> None: def visit_var_binding_(self, binding: VarBinding) -> None: """Visit VarBinding. - Users can customized this function to overwrite VisitBinding_(const VarBindingNode* binding) + Users can customized this function to overwrite VisitBinding_(const VarBindingObj* binding) on the C++ side. Parameters @@ -866,7 +866,7 @@ def visit_var_binding_(self, binding: VarBinding) -> None: def visit_match_cast_(self, binding: MatchCast) -> None: """Visit MatchCast. - Users can customized this function to overwrite VisitBinding_(const MatchCastNode* binding) + Users can customized this function to overwrite VisitBinding_(const MatchCastObj* binding) on the C++ side. Parameters @@ -879,7 +879,7 @@ def visit_match_cast_(self, binding: MatchCast) -> None: def visit_binding_block_(self, block: BindingBlock) -> None: """Visit BindingBlock. - Users can customized this function to overwrite VisitBindingBlock_(const BindingBlockNode* + Users can customized this function to overwrite VisitBindingBlock_(const BindingBlockObj* block) on the C++ side. Parameters @@ -892,7 +892,7 @@ def visit_binding_block_(self, block: BindingBlock) -> None: def visit_dataflow_block_(self, block: DataflowBlock) -> None: """Visit DataflowBlock. - Users can customized this function to overwrite VisitBindingBlock_(const DataflowBlockNode* + Users can customized this function to overwrite VisitBindingBlock_(const DataflowBlockObj* block) on the C++ side. Parameters @@ -905,7 +905,7 @@ def visit_dataflow_block_(self, block: DataflowBlock) -> None: def visit_var_def_(self, var: Var) -> None: """Visit the Var definition site. - Users can customized this function to overwrite VisitVarDef_(const VarNode* var) + Users can customized this function to overwrite VisitVarDef_(const VarObj* var) on the C++ side. Parameters @@ -918,7 +918,7 @@ def visit_var_def_(self, var: Var) -> None: def visit_dataflow_var_def_(self, var: DataflowVar) -> None: """Visit the DataflowVar definition site. - Users can customized this function to overwrite VisitVarDef_(const DataflowVarNode* var) + Users can customized this function to overwrite VisitVarDef_(const DataflowVarObj* var) on the C++ side. Parameters @@ -1202,7 +1202,7 @@ def visit_var_def(self, var: Var) -> Var: def visit_generic_const_(self, op: GenericConst) -> Expr: """Visit GenericConst. - Users can customized this function to overwrite VisitExpr_(const GenericConstNode* op) + Users can customized this function to overwrite VisitExpr_(const GenericConstObj* op) on the C++ side. Parameters @@ -1220,7 +1220,7 @@ def visit_generic_const_(self, op: GenericConst) -> Expr: def visit_tuple_(self, op: Tuple) -> Expr: """Visit Tuple. - Users can customized this function to overwrite VisitExpr_(const TupleNode* op) + Users can customized this function to overwrite VisitExpr_(const TupleObj* op) on the C++ side. Parameters @@ -1238,7 +1238,7 @@ def visit_tuple_(self, op: Tuple) -> Expr: def visit_var_(self, op: Var) -> Expr: """Visit Var. - Users can customized this function to overwrite VisitExpr_(const VarNode* op) + Users can customized this function to overwrite VisitExpr_(const VarObj* op) on the C++ side. Parameters @@ -1256,7 +1256,7 @@ def visit_var_(self, op: Var) -> Expr: def visit_dataflow_var_(self, op: DataflowVar) -> Expr: """Visit DataflowVar. - Users can customized this function to overwrite VisitExpr_(const DataflowVarNode* op) + Users can customized this function to overwrite VisitExpr_(const DataflowVarObj* op) on the C++ side. Parameters @@ -1274,7 +1274,7 @@ def visit_dataflow_var_(self, op: DataflowVar) -> Expr: def visit_shape_expr_(self, op: ShapeExpr) -> Expr: """Visit ShapeExpr. - Users can customized this function to overwrite VisitExpr_(const ShapeExprNode* op) + Users can customized this function to overwrite VisitExpr_(const ShapeExprObj* op) on the C++ side. Parameters @@ -1292,7 +1292,7 @@ def visit_shape_expr_(self, op: ShapeExpr) -> Expr: def visit_extern_func_(self, op: ExternFunc) -> Expr: """Visit ExternFunc. - Users can customized this function to overwrite VisitExpr_(const ExternFuncNode* op) + Users can customized this function to overwrite VisitExpr_(const ExternFuncObj* op) on the C++ side. Parameters @@ -1310,7 +1310,7 @@ def visit_extern_func_(self, op: ExternFunc) -> Expr: def visit_global_var_(self, op: GlobalVar) -> Expr: """Visit GlobalVar. - Users can customized this function to overwrite VisitExpr_(const GlobalVarNode* op) + Users can customized this function to overwrite VisitExpr_(const GlobalVarObj* op) on the C++ side. Parameters @@ -1328,7 +1328,7 @@ def visit_global_var_(self, op: GlobalVar) -> Expr: def visit_function_(self, op: Function) -> Expr: """Visit Function. - Users can customized this function to overwrite VisitExpr_(const FunctionNode* op) + Users can customized this function to overwrite VisitExpr_(const FunctionObj* op) on the C++ side. Parameters @@ -1346,7 +1346,7 @@ def visit_function_(self, op: Function) -> Expr: def visit_call_(self, op: Call) -> Expr: """Visit Call. - Users can customized this function to overwrite VisitExpr_(const CallNode* op) + Users can customized this function to overwrite VisitExpr_(const CallObj* op) on the C++ side. Parameters @@ -1364,7 +1364,7 @@ def visit_call_(self, op: Call) -> Expr: def visit_seq_expr_(self, op: SeqExpr) -> Expr: """Visit SeqExpr. - Users can customized this function to overwrite VisitExpr_(const SeqExprNode* op) + Users can customized this function to overwrite VisitExpr_(const SeqExprObj* op) on the C++ side. Parameters @@ -1382,7 +1382,7 @@ def visit_seq_expr_(self, op: SeqExpr) -> Expr: def visit_if_expr_(self, op: IfExpr) -> Expr: """Visit IfExpr. - Users can customized this function to overwrite VisitExpr_(const IfExprNode* op) + Users can customized this function to overwrite VisitExpr_(const IfExprObj* op) on the C++ side. Parameters @@ -1400,7 +1400,7 @@ def visit_if_expr_(self, op: IfExpr) -> Expr: def visit_op_(self, op: Op) -> Expr: """Visit Op. - Users can customized this function to overwrite VisitExpr_(const OpNode* op) + Users can customized this function to overwrite VisitExpr_(const OpObj* op) on the C++ side. Parameters @@ -1418,7 +1418,7 @@ def visit_op_(self, op: Op) -> Expr: def visit_tuple_getitem_(self, op: TupleGetItem) -> Expr: """Visit TupleGetItem. - Users can customized this function to overwrite VisitExpr_(const TupleGetItemNode* op) + Users can customized this function to overwrite VisitExpr_(const TupleGetItemObj* op) on the C++ side. Parameters @@ -1452,7 +1452,7 @@ def visit_expr_fallback_(self, op: Expr) -> Expr: def visit_string_imm_(self, op: StringImm) -> Expr: """Visit StringImm. - Users can customized this function to overwrite VisitExpr_(const StringImmNode* op) + Users can customized this function to overwrite VisitExpr_(const StringImmObj* op) on the C++ side. Parameters @@ -1470,7 +1470,7 @@ def visit_string_imm_(self, op: StringImm) -> Expr: def visit_data_type_imm_(self, op: DataTypeImm) -> Expr: """Visit DataTypeImm. - Users can customized this function to overwrite VisitExpr_(const DataTypeImmNode* op) + Users can customized this function to overwrite VisitExpr_(const DataTypeImmObj* op) on the C++ side. Parameters @@ -1488,7 +1488,7 @@ def visit_data_type_imm_(self, op: DataTypeImm) -> Expr: def visit_var_binding_(self, binding: VarBinding) -> None: """Visit VarBinding. - Users can customized this function to overwrite VisitBinding_(const VarBindingNode* binding) + Users can customized this function to overwrite VisitBinding_(const VarBindingObj* binding) on the C++ side. Parameters @@ -1501,7 +1501,7 @@ def visit_var_binding_(self, binding: VarBinding) -> None: def visit_match_cast_(self, binding: MatchCast) -> None: """Visit MatchCast. - Users can customized this function to overwrite VisitBinding_(const MatchCastNode* binding) + Users can customized this function to overwrite VisitBinding_(const MatchCastObj* binding) on the C++ side. Parameters @@ -1514,7 +1514,7 @@ def visit_match_cast_(self, binding: MatchCast) -> None: def visit_binding_block_(self, block: BindingBlock) -> BindingBlock: """Visit BindingBlock. - Users can customized this function to overwrite VisitBindingBlock_(const BindingBlockNode* + Users can customized this function to overwrite VisitBindingBlock_(const BindingBlockObj* block) on the C++ side. Parameters @@ -1532,7 +1532,7 @@ def visit_binding_block_(self, block: BindingBlock) -> BindingBlock: def visit_dataflow_block_(self, block: DataflowBlock) -> BindingBlock: """Visit DataflowBlock. - Users can customized this function to overwrite VisitBindingBlock_(const DataflowBlockNode* + Users can customized this function to overwrite VisitBindingBlock_(const DataflowBlockObj* block) on the C++ side. Parameters @@ -1550,7 +1550,7 @@ def visit_dataflow_block_(self, block: DataflowBlock) -> BindingBlock: def visit_var_def_(self, var: Var) -> Var: """Visit the Var definition site. - Users can customized this function to overwrite VisitVarDef_(const VarNode* var) + Users can customized this function to overwrite VisitVarDef_(const VarObj* var) on the C++ side. Parameters @@ -1568,7 +1568,7 @@ def visit_var_def_(self, var: Var) -> Var: def visit_dataflow_var_def_(self, var: DataflowVar) -> Var: """Visit the DataflowVar definition site. - Users can customized this function to overwrite VisitVarDef_(const DataflowVarNode* var) + Users can customized this function to overwrite VisitVarDef_(const DataflowVarObj* var) on the C++ side. Parameters diff --git a/python/tvm/relax/op/base.py b/python/tvm/relax/op/base.py index 22915461d5cc..955381c2b44a 100644 --- a/python/tvm/relax/op/base.py +++ b/python/tvm/relax/op/base.py @@ -406,7 +406,7 @@ def invoke_closure( The input arguments. type_args: Union[List[Type], Type] - The type information arguments of the CallNode + The type information arguments of the CallObj Returns ------- @@ -808,7 +808,7 @@ def invoke_pure_closure( The input arguments. type_args: Union[List[Type], Type] - The type information arguments of the CallNode + The type information arguments of the CallObj Returns ------- diff --git a/python/tvm/relax/op/vm/vm.py b/python/tvm/relax/op/vm/vm.py index 4e68272a6fca..9d3be890209b 100644 --- a/python/tvm/relax/op/vm/vm.py +++ b/python/tvm/relax/op/vm/vm.py @@ -135,7 +135,7 @@ def kill_object(obj: Expr, *, ty=None, loc: Location = UNKNOWN_LOC) -> Call: Returns ------- result : Call - CallNode that kills the input object. + CallObj that kills the input object. """ return _Call("relax.vm.kill_object", [obj], loc=loc, ty=ty) # type: ignore diff --git a/python/tvm/relax/pipeline.py b/python/tvm/relax/pipeline.py index 53ef83f5ed22..ab04aa9e191e 100644 --- a/python/tvm/relax/pipeline.py +++ b/python/tvm/relax/pipeline.py @@ -37,7 +37,7 @@ def zero_pipeline(*, enable_warning: bool = False): ---------- enable_warning : bool A boolean value indicating if to print warnings - * in LegalizeOps pass, for CallNode whose op's legalization function is + * in LegalizeOps pass, for CallObj whose op's legalization function is not registered, * in MetaScheduleApplyDatabase pass, for TIR functions now showing up in the database. By default we don't print warning. diff --git a/python/tvm/relax/script/ir_builder/__init__.py b/python/tvm/relax/script/ir_builder/__init__.py index 25892e45abe7..b2977cb98580 100644 --- a/python/tvm/relax/script/ir_builder/__init__.py +++ b/python/tvm/relax/script/ir_builder/__init__.py @@ -17,8 +17,11 @@ # pylint: disable=redefined-builtin, wrong-import-order, no-member, invalid-name """Public exports for the Relax builder.""" -from tvm_tirx.ir.type import _init_type_api +import sys as _sys + from tvm_tirx.script.ir_builder.ir import constexpr +from tvm_tirx.script.parser import register_namespace as _register_namespace +from tvm_tirx.tirx.script.ir_builder import expose_script_op as _expose_script_op from . import distributed as dist from . import ir as _native @@ -34,5 +37,16 @@ __all__ = [*_native.__all__, *_op.__all__, *_protocol.__all__, "constexpr", "dist"] - -_init_type_api("relax", __name__) +# Install explicit builders before requesting generated inspection operators. +_register_namespace("relax", _sys.modules[__name__]) +for _name in ( + "tensor_dtype_code", + "tensor_dtype_bits", + "tensor_dtype_lanes", + "tensor_ndim", + "tensor_byte_offset", + "tensor_elem_offset", + "tensor_shape_i", + "tensor_stride_i", +): + _expose_script_op(f"relax.inspect.{_name}") diff --git a/python/tvm/relax/script/ir_builder/ir.py b/python/tvm/relax/script/ir_builder/ir.py index 72739a0888d8..26355c4b3682 100644 --- a/python/tvm/relax/script/ir_builder/ir.py +++ b/python/tvm/relax/script/ir_builder/ir.py @@ -33,10 +33,18 @@ from tvm.relax.distributed import device_mesh from tvm.relax.dpl import PatternMatchingRewriter from tvm.relax.global_info import DummyGlobalInfo, VDevice +from tvm.relax.type import FuncType as Func from tvm_tirx import ir as _ir from tvm_tirx.ir import IRModule from tvm_tirx.ir.location import UNKNOWN_LOC, Location -from tvm_tirx.ir.type import _init_type_api, _normalize_type +from tvm_tirx.ir.type import MissingType as Missing +from tvm_tirx.ir.type import NoneType as NoneType +from tvm_tirx.ir.type import OpaqueType as Opaque +from tvm_tirx.ir.type import PrimType as Prim +from tvm_tirx.ir.type import PtrType as Ptr +from tvm_tirx.ir.type import StringType as String +from tvm_tirx.ir.type import TensorRegionType as TensorRegion +from tvm_tirx.ir.type import _normalize_type from tvm_tirx.runtime._tensor import ( cpu, cuda, @@ -555,12 +563,20 @@ def match_cast(value, ty, *, loc: _LocationEntry | Location = UNKNOWN_LOC): "Callable", "DTensor", "ExternFunc", + "Func", + "Missing", + "NoneType", "Object", + "Opaque", "PackedFunc", + "Prim", + "Ptr", "Range", "Shape", "ShapeExpr", + "String", "Tensor", + "TensorRegion", "Tuple", "TupleGetItem", "const", @@ -587,6 +603,3 @@ def match_cast(value, ty, *, loc: _LocationEntry | Location = UNKNOWN_LOC): "vulkan", "webgpu", ] - - -_init_type_api("relax", __name__) diff --git a/python/tvm/relax/script/ir_builder/op.py b/python/tvm/relax/script/ir_builder/op.py index 43e6eb5ec66d..74dde62d0b01 100644 --- a/python/tvm/relax/script/ir_builder/op.py +++ b/python/tvm/relax/script/ir_builder/op.py @@ -232,10 +232,8 @@ from tvm_tirx import ir as _ir from tvm_tirx import tirx as _tir from tvm_tirx.ir.location import UNKNOWN_LOC, Location -from tvm_tirx.ir.op import _init_op_api from tvm_tirx.ir.prim import _ffi_api as _prim_ffi from tvm_tirx.script.ir_builder import base as _base -from tvm_tirx.script.ir_builder.base import auto_emit_void_call as _auto_emit_void_call from . import distributed as dist # noqa: F401 from .ir import _value, lookup_vdevice @@ -574,7 +572,6 @@ def ne_(lhs: Any, rhs: Any, *, loc: _Loc = None) -> _ir.Expr: inspect = _Namespace() -_init_op_api("relax", __name__, wrapper=_auto_emit_void_call) __all__ = [ "ArgmaxArgminAttrs", diff --git a/python/tvm/relax/transform/legalize_ops/create.py b/python/tvm/relax/transform/legalize_ops/create.py index f193a4c1e5b0..f42a6eef5c39 100644 --- a/python/tvm/relax/transform/legalize_ops/create.py +++ b/python/tvm/relax/transform/legalize_ops/create.py @@ -34,9 +34,7 @@ def _full(is_like: bool, fill_value: float | None, function_name: str) -> LegalizeFunc: def full_call_te(bb: BlockBuilder, call: Call) -> Expr: _fill_value = ( - _try_convert_to_scalar_const(call.args[1], python_native=True) - if fill_value is None - else fill_value + _try_convert_to_scalar_const(call.args[1]) if fill_value is None else fill_value ) shape = call.args[0].ty.shape if is_like else call.args[0] diff --git a/python/tvm/relax/transform/legalize_ops/index.py b/python/tvm/relax/transform/legalize_ops/index.py index 87e83e99be37..d954d86ed8d1 100644 --- a/python/tvm/relax/transform/legalize_ops/index.py +++ b/python/tvm/relax/transform/legalize_ops/index.py @@ -88,19 +88,21 @@ def _dynamic_strided_slice(bb: BlockBuilder, call: Call) -> Expr: # 1. Insert shape function def shape_func(data, begin, end, strides): def _compute(i): - def canonicalize_index(index, extent, strides): - begin_range = tirx.Select( - strides < 0, tirx.const(-1, "int64"), tirx.const(0, "int64") - ) - end_range = tirx.Select(strides < 0, extent - 1, extent) + def canonicalize_index(index, extent, negative_stride): + begin_range = tirx.const(-1 if negative_stride else 0, "int64") + end_range = extent - 1 if negative_stride else extent index = tirx.Select(index < 0, index + extent, index) return tirx.Min(tirx.Max(index, begin_range), end_range) def get_length(begin, end, strides, length): - begin = canonicalize_index(begin, length, strides) - end = canonicalize_index(end, length, strides) - len1 = tirx.ceildiv(begin - end, -strides) - len2 = tirx.ceildiv(end - begin, strides) + # Construct each branch with its known stride sign. Tensor reads are + # mutable, so the analyzer cannot assume repeated reads agree. + begin_neg = canonicalize_index(begin, length, True) + end_neg = canonicalize_index(end, length, True) + begin_pos = canonicalize_index(begin, length, False) + end_pos = canonicalize_index(end, length, False) + len1 = tirx.ceildiv(begin_neg - end_neg, -strides) + len2 = tirx.ceildiv(end_pos - begin_pos, strides) return tirx.Select(strides < 0, len1, len2) length = tirx.const(-1, "int64") diff --git a/python/tvm/relax/transform/transform.py b/python/tvm/relax/transform/transform.py index 445c3bd94a65..2b34ce429f6b 100644 --- a/python/tvm/relax/transform/transform.py +++ b/python/tvm/relax/transform/transform.py @@ -1088,7 +1088,7 @@ def LegalizeOps( List of ops that need to be skipped from legalization enable_warning : bool - A boolean value indicating if to print warnings for CallNode whose op's + A boolean value indicating if to print warnings for CallObj whose op's legalization function is not registered. By default we don't print warnings. diff --git a/python/tvm/s_tir/meta_schedule/feature_extractor/per_store_feature.py b/python/tvm/s_tir/meta_schedule/feature_extractor/per_store_feature.py index 7d0d94ee1aa6..320f4103226c 100644 --- a/python/tvm/s_tir/meta_schedule/feature_extractor/per_store_feature.py +++ b/python/tvm/s_tir/meta_schedule/feature_extractor/per_store_feature.py @@ -15,7 +15,7 @@ # specific language governing permissions and limitations # under the License. # pylint: disable=invalid-name -"""We extract one feature vector per TensorStoreNode statement in a TIR Stmt, +"""We extract one feature vector per TensorStoreObj statement in a TIR Stmt, so we call this feature as "per-store" feature. """ @@ -27,7 +27,7 @@ @register_object("s_tir.meta_schedule.PerStoreFeature") class PerStoreFeature(FeatureExtractor): - """PerStoreFeature extracts one feature vector per TensorStoreNode + """PerStoreFeature extracts one feature vector per TensorStoreObj Parameters ---------- diff --git a/python/tvm/s_tir/meta_schedule/profiler.py b/python/tvm/s_tir/meta_schedule/profiler.py index 7ba3cbec10b2..8e226121627a 100644 --- a/python/tvm/s_tir/meta_schedule/profiler.py +++ b/python/tvm/s_tir/meta_schedule/profiler.py @@ -37,7 +37,8 @@ def __init__(self) -> None: def get(self) -> dict[str, float]: """Get the profiling results in seconds""" - return _ffi_api.ProfilerGet(self) # type: ignore # pylint: disable=no-member + results = _ffi_api.ProfilerGet(self) # type: ignore # pylint: disable=no-member + return {str(name): float(value.value) for name, value in results.items()} def table(self) -> str: """Get the profiling results in a table format""" diff --git a/src/relax/analysis/analysis.cc b/src/relax/analysis/analysis.cc index 07578ab82cc1..dcc01bfb0396 100644 --- a/src/relax/analysis/analysis.cc +++ b/src/relax/analysis/analysis.cc @@ -102,9 +102,9 @@ class VarVisitor : protected ExprVisitor { void VisitExprDepTypeField(const Type&) final {} - void VisitExpr_(const VarNode* var) final { vars_.Insert(ffi::GetRef(var)); } + void VisitExpr_(const VarObj* var) final { vars_.Insert(ffi::GetRef(var)); } - void VisitExpr_(const FunctionNode* op) final { + void VisitExpr_(const FunctionObj* op) final { for (const auto& param : op->params) { MarkBounded(param); MarkTypeDefinitionsBound(GetType(param)); @@ -112,11 +112,9 @@ class VarVisitor : protected ExprVisitor { VisitExpr(op->body); } - void VisitExpr_(const GlobalVarNode* op) final { - global_vars_.Insert(ffi::GetRef(op)); - } + void VisitExpr_(const GlobalVarObj* op) final { global_vars_.Insert(ffi::GetRef(op)); } - void VisitExpr_(const CallNode* call_node) final { + void VisitExpr_(const CallObj* call_node) final { VisitLoc(call_node->loc); VisitExpr(call_node->op); @@ -129,13 +127,13 @@ class VarVisitor : protected ExprVisitor { } } - void VisitBinding_(const VarBindingNode* binding) final { + void VisitBinding_(const VarBindingObj* binding) final { MarkBounded(binding->var); VisitExpr(binding->value); VisitVarDef(binding->var); } - void VisitBinding_(const MatchCastNode* binding) final { + void VisitBinding_(const MatchCastObj* binding) final { MarkBounded(binding->var); MarkTypeDefinitionsBound(binding->ty); ExprVisitor::VisitBinding_(binding); @@ -176,7 +174,7 @@ ffi::Optional FindImpureCall(const Expr& expr, const ffi::Optional& } } - void VisitExpr_(const FunctionNode* func) override { + void VisitExpr_(const FunctionObj* func) override { // we don't visit inner functions because an impure call in an inner function // does *not* mean the outer function contains an impure call } @@ -187,7 +185,7 @@ ffi::Optional FindImpureCall(const Expr& expr, const ffi::Optional& // as part of the Relax expression being checked. } - void VisitExpr_(const CallNode* call) override { + void VisitExpr_(const CallObj* call) override { // ignore recursive calls if we find one bool is_recursive = (own_name_ && own_name_.value().same_as(call->op)); auto expr = ffi::GetRef(call); @@ -204,12 +202,12 @@ ffi::Optional FindImpureCall(const Expr& expr, const ffi::Optional& }; if (own_name) { - TVM_FFI_ICHECK(own_name.value().as() || own_name.value().as()) + TVM_FFI_ICHECK(own_name.value().as() || own_name.value().as()) << "Must pass a Var or GlobalVar for own_name"; } Expr to_check = expr; - if (auto func = to_check.as()) { + if (auto func = to_check.as()) { to_check = func->body; } return ImpureCallChecker::Check(to_check, own_name); diff --git a/src/relax/analysis/computable_at_compile_time.cc b/src/relax/analysis/computable_at_compile_time.cc index e23b78c9ffae..bb3cb770a88f 100644 --- a/src/relax/analysis/computable_at_compile_time.cc +++ b/src/relax/analysis/computable_at_compile_time.cc @@ -43,7 +43,7 @@ class CompileTimeCollector : ExprVisitor { } private: - void VisitExpr_(const FunctionNode* func) override { + void VisitExpr_(const FunctionObj* func) override { if (auto opt_num_input = func->attrs.GetAttr(tvm::relax::attr::kNumInput)) { size_t num_input = opt_num_input.value(); for (size_t i = num_input; i < func->params.size(); i++) { diff --git a/src/relax/analysis/detect_recursion.cc b/src/relax/analysis/detect_recursion.cc index 8ced3ad37a13..5916ef6bd010 100644 --- a/src/relax/analysis/detect_recursion.cc +++ b/src/relax/analysis/detect_recursion.cc @@ -86,9 +86,9 @@ class DependencyGatherer : public ExprVisitor { return deps_; } - void VisitExpr_(const GlobalVarNode* gv) override { + void VisitExpr_(const GlobalVarObj* gv) override { // disregard Functions - if (!m_->Lookup(ffi::GetRef(gv)).as()) { + if (!m_->Lookup(ffi::GetRef(gv)).as()) { return; } deps_.insert(gv->name_hint); @@ -106,7 +106,7 @@ using adjacency_index = std::vector; adjacency_map GatherDependencyGraph(const IRModule& m) { adjacency_map ret; for (auto gv_func : m->functions) { - const relax::FunctionNode* func = gv_func.second.as(); + const relax::FunctionObj* func = gv_func.second.as(); // disregard Functions and the like if (!func) { continue; diff --git a/src/relax/analysis/graph_partitioner.cc b/src/relax/analysis/graph_partitioner.cc index 3add82b0ef5e..01c4f0d51a0a 100644 --- a/src/relax/analysis/graph_partitioner.cc +++ b/src/relax/analysis/graph_partitioner.cc @@ -236,7 +236,7 @@ size_t GraphPartitioner::CountArgs_(IndexedForwardGraph::Node* src, visited_groups.insert(gnode->FindRoot()); auto calc_args_number = [this, src, &graph, &visited_groups, update_postpone](const Expr& arg) -> size_t { - if (arg.as()) return 0; + if (arg.as()) return 0; auto* node = graph.node_map.at(arg.get()); Group* prev_group = groups_[node->index]->FindRoot(); if (visited_groups.count(prev_group) == 0) { @@ -254,11 +254,11 @@ size_t GraphPartitioner::CountArgs_(IndexedForwardGraph::Node* src, } return 0; }; - if (auto call_node = ffi::GetRef(src->ref).as()) { + if (auto call_node = ffi::GetRef(src->ref).as()) { for (auto& it : call_node->args) { sum += calc_args_number(it); } - } else if (auto tuple_node = ffi::GetRef(src->ref).as()) { + } else if (auto tuple_node = ffi::GetRef(src->ref).as()) { for (auto& it : tuple_node->fields) { sum += calc_args_number(it); } @@ -290,19 +290,19 @@ size_t GraphPartitioner::CountFusedArgs(const IndexedForwardGraph& graph, void GraphPartitioner::InitGroups(const IndexedForwardGraph& graph) { auto args_counter = [](const tvm::ffi::Object* obj) { size_t args_num = 0; - if (auto call_node = ffi::GetRef(obj).as()) { + if (auto call_node = ffi::GetRef(obj).as()) { for (auto& it : call_node->args) { - if (it.as() || it.as()) { + if (it.as() || it.as()) { args_num++; } } - } else if (auto tuple_node = ffi::GetRef(obj).as()) { + } else if (auto tuple_node = ffi::GetRef(obj).as()) { for (auto& it : tuple_node->fields) { - if (it.as() || it.as()) { + if (it.as() || it.as()) { args_num++; } } - } else if (ffi::GetRef(obj).as()) { + } else if (ffi::GetRef(obj).as()) { args_num++; } return args_num; diff --git a/src/relax/analysis/layout_transformation.cc b/src/relax/analysis/layout_transformation.cc index 7ee6f2981b21..5d42a979f90f 100644 --- a/src/relax/analysis/layout_transformation.cc +++ b/src/relax/analysis/layout_transformation.cc @@ -71,11 +71,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)); @@ -485,14 +485,14 @@ class BlockAnalyzer : public s_tir::StmtExprVisitor { } } - ffi::Optional Visit_(const s_tir::SBlockNode* op) final { + ffi::Optional Visit_(const s_tir::SBlockObj* op) final { // Blocks with nested blocks cannot be handled yet. LOG(WARNING) << "[LayoutInference] Nested blocks are not supported for layout inference yet"; can_transform_block_ = false; return std::nullopt; } - ffi::Optional Visit_(const TensorStoreNode* op) final { + ffi::Optional Visit_(const TensorStoreObj* op) final { TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(s_tir::StmtExprVisitor::Visit_(op)); TensorAccessInfo& access_info = @@ -526,7 +526,7 @@ class BlockAnalyzer : public s_tir::StmtExprVisitor { return std::nullopt; } - ffi::Optional Visit_(const TensorLoadNode* op) final { + ffi::Optional Visit_(const TensorLoadObj* op) final { TensorVar read_tensor = op->source.as_or_throw(); TensorAccessInfo& access_info = tensor_access_info_[op->source.as_or_throw()]; @@ -604,7 +604,7 @@ class FunctionAnalyzer : public s_tir::StmtExprVisitor { } private: - ffi::Optional Visit_(const s_tir::SBlockNode* op) final { + ffi::Optional Visit_(const s_tir::SBlockObj* op) final { if (op->name_hint == "root") { // Skip the root block return s_tir::StmtExprVisitor::Visit_(op); diff --git a/src/relax/analysis/shape_analysis.cc b/src/relax/analysis/shape_analysis.cc index 126654cf1d37..bd7a17d6dd94 100644 --- a/src/relax/analysis/shape_analysis.cc +++ b/src/relax/analysis/shape_analysis.cc @@ -41,8 +41,8 @@ bool CanProveShapeEqual(const ffi::Array& lhs, const ffi::Array(); - auto* rhs_shape = rhs.as(); + auto* lhs_shape = lhs.as(); + auto* rhs_shape = rhs.as(); if (lhs_shape && rhs_shape) { return CanProveShapeEqual(lhs_shape->values, rhs_shape->values, ana); diff --git a/src/relax/analysis/tir_op_pattern_kind.cc b/src/relax/analysis/tir_op_pattern_kind.cc index d5a7220d4bfd..70fd1af6c10b 100644 --- a/src/relax/analysis/tir_op_pattern_kind.cc +++ b/src/relax/analysis/tir_op_pattern_kind.cc @@ -51,7 +51,7 @@ class PatternKindAnalyzer : public s_tir::StmtExprVisitor { } private: - bool IsOutputBlock(const s_tir::SBlockNode* block) { + bool IsOutputBlock(const s_tir::SBlockObj* block) { for (const TensorRegion& write_region : block->writes) { if (param_tensors_.count(write_region->source.as_or_throw())) { return true; @@ -60,7 +60,7 @@ class PatternKindAnalyzer : public s_tir::StmtExprVisitor { return false; } - ffi::Optional Visit_(const TensorStoreNode* op) final { + ffi::Optional Visit_(const TensorStoreObj* op) final { // We only support one tensor store in a block (usually generated by TE compute) // If we have already seen tensor store in the current block, classify as Opaque. if (store_.has_value() && !IsSameArray(op->indices, store_.value()->indices)) { @@ -71,12 +71,12 @@ class PatternKindAnalyzer : public s_tir::StmtExprVisitor { return s_tir::StmtExprVisitor::Visit_(op); } - ffi::Optional Visit_(const TensorLoadNode* op) final { + ffi::Optional Visit_(const TensorLoadObj* op) final { loads_.push_back(ffi::GetRef(op)); return s_tir::StmtExprVisitor::Visit_(op); } - ffi::Optional Visit_(const s_tir::SBlockNode* op) final { + ffi::Optional Visit_(const s_tir::SBlockObj* op) final { if (op->name_hint == "root") { // Skip the root block return s_tir::StmtExprVisitor::Visit(op->body); @@ -232,7 +232,7 @@ class PatternKindAnalyzer : public s_tir::StmtExprVisitor { * A[i, j] = B[i - j] is injective since the load index vars are only i, j */ static bool IsInjectivePattern(const TensorStore& store, const TensorLoad& load) { - std::unordered_set vars; + std::unordered_set vars; for (const PrimExpr& store_index : store->indices) { if (auto var = store_index.as()) { vars.insert(var.value().get()); @@ -260,7 +260,7 @@ class PatternKindAnalyzer : public s_tir::StmtExprVisitor { * Store = A[i, j] and Load = B[i, j + k] allow data reuse. */ static bool IsAllowReusePattern(const TensorStore& store, const TensorLoad& load) { - std::unordered_set vars; + std::unordered_set vars; for (const PrimExpr& index : store->indices) { if (auto var = index.as()) { vars.insert(var.value().get()); @@ -282,7 +282,7 @@ class PatternKindAnalyzer : public s_tir::StmtExprVisitor { static PrimExpr RemoveCast(PrimExpr e) { for (;;) { - if (const auto* cast = e.as()) { + if (const auto* cast = e.as()) { e = cast->value; } else { break; @@ -294,17 +294,17 @@ class PatternKindAnalyzer : public s_tir::StmtExprVisitor { /*! \brief Checking if the stmt is multiply add. E.g. C[i, j] += A[i, k] * B[j, k] */ static bool IsFMA(const SeqStmt& body) { if (body->size() != 1) return false; - if (const auto* store = body->seq[0].as()) { - if (const auto* add = RemoveCast(store->value).as()) { - if (const auto* mul = RemoveCast(add->b).as()) { - const auto* store_lhs = RemoveCast(add->a).as(); + if (const auto* store = body->seq[0].as()) { + if (const auto* add = RemoveCast(store->value).as()) { + if (const auto* mul = RemoveCast(add->b).as()) { + const auto* store_lhs = RemoveCast(add->a).as(); if (!store_lhs || !store->dest.same_as(store_lhs->source.as_or_throw()) || !IsSameArray(store->indices, store_lhs->indices)) { return false; } - const auto* lhs = RemoveCast(mul->a).as(); - const auto* rhs = RemoveCast(mul->b).as(); + const auto* lhs = RemoveCast(mul->a).as(); + const auto* rhs = RemoveCast(mul->b).as(); if (!lhs || !rhs) { return false; } @@ -375,7 +375,7 @@ bool HasReshapePattern(const tirx::Function& func) { class ReshapeDetector : public s_tir::StmtExprVisitor { public: ffi::Optional Visit(ffi::AnyView value) override { - if (value.as()) return std::nullopt; + if (value.as()) return std::nullopt; return s_tir::StmtExprVisitor::Visit(value); } @@ -389,19 +389,20 @@ bool HasReshapePattern(const tirx::Function& func) { : is_reshape_(false), src_tensor_(src_tensor), dst_tensor_(dst_tensor) {} private: - ffi::Optional Visit_(const ForNode* loop) final { - ana_->Bind(loop->var, Range(loop->begin, loop->end)); + ffi::Optional Visit_(const ForObj* loop) final { + auto context = ana_->GetConstraintContext(); + context->Append(sym::RangeConstraint(loop->var, Range(loop->begin, loop->end))); + ana_->SyncConstraintContext(context); // To detect the reshape pattern, we require each For to have // either another For or a BlockRealize as body. - if (loop->body->size() != 1 || - !(loop->body->seq[0]->IsInstance() || - loop->body->seq[0]->IsInstance())) { + if (loop->body->size() != 1 || !(loop->body->seq[0]->IsInstance() || + loop->body->seq[0]->IsInstance())) { return std::nullopt; } return this->Visit(loop->body); } - ffi::Optional Visit_(const s_tir::SBlockRealizeNode* block_realize) final { + ffi::Optional Visit_(const s_tir::SBlockRealizeObj* block_realize) final { // Constructing the mapping from block iterators to iterator // binding values. The mapping will be used in the substitution of // the flattened tensor access index. @@ -421,27 +422,29 @@ bool HasReshapePattern(const tirx::Function& func) { return this->Visit(block); } - ffi::Optional Visit_(const s_tir::SBlockNode* block) final { + ffi::Optional Visit_(const s_tir::SBlockObj* block) final { if (block->body->size() != 1) return std::nullopt; - // Step 0. If the block body is a ForNode, recurse into it. - if (block->body->seq[0]->IsInstance()) { + // Step 0. If the block body is a ForObj, recurse into it. + if (block->body->seq[0]->IsInstance()) { return this->Visit(block->body); } ffi::Map var_range; for (const s_tir::IterVar& v : block->iter_vars) { - ana_->Bind(v->var, v->dom.value()); + auto context = ana_->GetConstraintContext(); + context->Append(sym::RangeConstraint(v->var, v->dom.value())); + ana_->SyncConstraintContext(context); var_range.Set(v->var, v->dom.value()); } // Step 1. Get the load/store pattern of the block body. // To detect the reshape pattern, we require the block body to be a // TensorStore, which has a TensorLoad as value. - const auto* tensor_store = block->body->seq[0].as(); + const auto* tensor_store = block->body->seq[0].as(); if (tensor_store == nullptr) { return std::nullopt; } - const auto* tensor_load = tensor_store->value.as(); + const auto* tensor_load = tensor_store->value.as(); if (tensor_load == nullptr) { return std::nullopt; } @@ -581,7 +584,7 @@ bool HasReshapePattern(const tirx::Function& func) { // To detect the reshape pattern, we require each For to have // either another For or a BlockRealize as body. TVM_FFI_ICHECK(func->body.value()->size() == 1 && - func->body.value()->seq[0].as()); + func->body.value()->seq[0].as()); return ReshapeDetector::Detect(src_tensor, dst_tensor, func->body.value()); } diff --git a/src/relax/analysis/type_analysis.cc b/src/relax/analysis/type_analysis.cc index 3e9cb9f58d83..78e348f4ccf7 100644 --- a/src/relax/analysis/type_analysis.cc +++ b/src/relax/analysis/type_analysis.cc @@ -44,29 +44,29 @@ using namespace tvm::prim; //-------------------------- class StaticTypeDeriver : public TypeFunctor { public: - Type VisitType_(const AnyTypeNode* op) final { return AnyType(op->loc); } + Type VisitType_(const AnyTypeObj* op) final { return AnyType(op->loc); } - Type VisitType_(const PrimTypeNode* op) final { return tvm::PrimType(op->dtype); } + Type VisitType_(const PrimTypeObj* op) final { return tvm::PrimType(op->dtype); } - Type VisitType_(const StringTypeNode* op) final { return StringType(); } + Type VisitType_(const StringTypeObj* op) final { return StringType(); } - Type VisitType_(const ShapeTypeNode* op) final { return ShapeType(op->ndim, op->loc); } + Type VisitType_(const ShapeTypeObj* op) final { return ShapeType(op->ndim, op->loc); } - Type VisitType_(const TensorTypeNode* op) final { + Type VisitType_(const TensorTypeObj* op) final { return TensorType(op->dtype, op->ndim, op->vdevice, op->loc); } // module: distributed - Type VisitType_(const distributed::DTensorTypeNode* op) final { return AnyType(); } + Type VisitType_(const distributed::DTensorTypeObj* op) final { return AnyType(); } // end-module: distributed - Type VisitType_(const TupleTypeNode* op) final { + Type VisitType_(const TupleTypeObj* op) final { ffi::Array fields = op->fields.Map([this](const Type& ty) { return this->VisitType(ty); }); return TupleType(fields, op->loc); } - Type VisitType_(const FuncTypeNode* op) final { + Type VisitType_(const FuncTypeObj* op) final { if (op->IsOpaque()) return PackedFuncType(op->loc); ffi::Array params = op->params.value().Map([this](const Type& ty) { return this->VisitType(ty); }); @@ -88,31 +88,31 @@ TVM_FFI_STATIC_INIT_BLOCK() { //-------------------------- Type TypeFromStaticType(const Type& type) { - if (type.as()) { + if (type.as()) { return AnyType(type->loc); - } else if (type.as()) { + } else if (type.as()) { return StringType(); - } else if (const PrimTypeNode* prim_type = type.as()) { + } else if (const PrimTypeObj* prim_type = type.as()) { return tvm::PrimType(prim_type->dtype); - } else if (const tvm::PrimTypeNode* prim_type = type.as()) { + } else if (const tvm::PrimTypeObj* prim_type = type.as()) { return tvm::PrimType(prim_type->dtype); - } else if (const ShapeTypeNode* shape_type = type.as()) { + } else if (const ShapeTypeObj* shape_type = type.as()) { return ShapeType(shape_type->ndim, type->loc); - } else if (const TensorTypeNode* tensor_type = type.as()) { + } else if (const TensorTypeObj* tensor_type = type.as()) { return TensorType(tensor_type->dtype, tensor_type->ndim); - } else if (const TupleTypeNode* tuple_type = type.as()) { + } else if (const TupleTypeObj* tuple_type = type.as()) { ffi::Array fields; for (const Type& field : tuple_type->fields) { fields.push_back(TypeFromStaticType(field)); } return TupleType(fields, type->loc); - } else if (const FuncTypeNode* func_type = type.as()) { + } else if (const FuncTypeObj* func_type = type.as()) { if (func_type->IsOpaque()) return FuncType::OpaqueFunc(func_type->ret, func_type->purity); ffi::Array params = func_type->params.value().Map([](const Type& param) { return TypeFromStaticType(param); }); Type ret = TypeFromStaticType(func_type->ret); return FuncType(params, ret, func_type->purity, func_type->loc); - } else if (const tvm::FuncTypeNode* func_type = type.as()) { + } else if (const tvm::FuncTypeObj* func_type = type.as()) { ffi::Array params = func_type->arg_types.Map([](const Type& param) { return TypeFromStaticType(param); }); Type ret = TypeFromStaticType(func_type->ret_type); @@ -133,9 +133,9 @@ class WellDefinedEraser : public TypeMutator, public ExprMutatorBase { sym::AnalyzerObj* ana) : f_var_map_(f_var_map), ana_(ana) {} - Type VisitType_(const PrimTypeNode* op) final { return ffi::GetRef(op); } + Type VisitType_(const PrimTypeObj* op) final { return ffi::GetRef(op); } - Type VisitType_(const ShapeTypeNode* op) final { + Type VisitType_(const ShapeTypeObj* op) final { bool has_undefined = false; ffi::Optional> values; @@ -156,7 +156,7 @@ class WellDefinedEraser : public TypeMutator, public ExprMutatorBase { } } - Type VisitType_(const TensorTypeNode* op) final { + Type VisitType_(const TensorTypeObj* op) final { bool has_undefined = false; ffi::Optional shape; @@ -184,7 +184,7 @@ class WellDefinedEraser : public TypeMutator, public ExprMutatorBase { } } - Type VisitType_(const FuncTypeNode* op) final { + Type VisitType_(const FuncTypeObj* op) final { // NOTE: we always require func type to be well-defined. // // All the occuring symbolic variables are defined in parameters' @@ -196,7 +196,7 @@ class WellDefinedEraser : public TypeMutator, public ExprMutatorBase { PrimExpr VisitPrimitiveExpr(const PrimExpr& expr) { auto f_substitute = [this](const Var& var) -> ffi::Expected> { - if (var.as()) { + if (var.as()) { has_undefined_ = true; return ffi::Unchanged(); } @@ -205,7 +205,7 @@ class WellDefinedEraser : public TypeMutator, public ExprMutatorBase { if (!ret.has_value()) return ffi::Unchanged(); PrimExpr value = ret.value().as_or_throw(); - if (value->IsInstance()) { + if (value->IsInstance()) { return ffi::Any(tvm::prim::cast(value, PrimType::Int(64))); } TVM_FFI_ICHECK(value.ty().MatchesElementType(DLDataTypeCode::kDLInt, 64)) @@ -223,13 +223,13 @@ class WellDefinedEraser : public TypeMutator, public ExprMutatorBase { } } - Expr VisitExpr_(const ShapeExprNode* op) final { + Expr VisitExpr_(const ShapeExprObj* op) final { ffi::Array values = op->values.Map([this](const PrimExpr& expr) { return VisitPrimitiveExpr(expr); }); return values.same_as(op->values) ? ffi::GetRef(op) : ShapeExpr(values, op->loc); } - Expr VisitExpr_(const VarNode* var) final { + Expr VisitExpr_(const VarObj* var) final { Var id = ffi::GetRef(var); ffi::Optional ret; if (f_var_map_ != nullptr) { @@ -238,13 +238,13 @@ class WellDefinedEraser : public TypeMutator, public ExprMutatorBase { has_undefined_ = has_undefined_ || !ret.has_value(); if (ret.has_value()) { - TVM_FFI_ICHECK((ret.as() && !ret.as()) || ret.as()) + TVM_FFI_ICHECK((ret.as() && !ret.as()) || ret.as()) << "Only allow Expr in Type to be ShapeExpr or Var"; } return ret.value_or(ffi::GetRef(var)); } - Expr VisitExpr_(const DataflowVarNode* var) final { + Expr VisitExpr_(const DataflowVarObj* var) final { has_undefined_ = true; return ffi::GetRef(var); } @@ -317,19 +317,19 @@ class TypeBaseChecker : public TypeFunctor()) return BaseCheckResult::kPass; - return other.as() ? BaseCheckResult::kFailL1 : BaseCheckResult::kFailL0; + BaseCheckResult VisitType_(const StringTypeObj* lhs, const Type& other) final { + if (other.as()) return BaseCheckResult::kPass; + return other.as() ? BaseCheckResult::kFailL1 : BaseCheckResult::kFailL0; } - BaseCheckResult VisitType_(const PrimTypeNode* lhs, const Type& other) final { - auto* rhs = other.as(); + BaseCheckResult VisitType_(const PrimTypeObj* lhs, const Type& other) final { + auto* rhs = other.as(); if (rhs == nullptr) { - if (other.as()) return BaseCheckResult::kFailL1; + if (other.as()) return BaseCheckResult::kFailL1; return BaseCheckResult::kFailL0; } @@ -340,10 +340,10 @@ class TypeBaseChecker : public TypeFunctor(); + BaseCheckResult VisitType_(const ShapeTypeObj* lhs, const Type& other) final { + auto* rhs = other.as(); if (rhs == nullptr) { - if (other.as()) return BaseCheckResult::kFailL1; + if (other.as()) return BaseCheckResult::kFailL1; return BaseCheckResult::kFailL0; } // lhs have unknown ndim @@ -364,10 +364,10 @@ class TypeBaseChecker : public TypeFunctorvalues.value(), rhs->values.value()); } - BaseCheckResult VisitType_(const TensorTypeNode* lhs, const Type& other) final { - auto* rhs = other.as(); + BaseCheckResult VisitType_(const TensorTypeObj* lhs, const Type& other) final { + auto* rhs = other.as(); if (rhs == nullptr) { - if (other.as()) return BaseCheckResult::kFailL1; + if (other.as()) return BaseCheckResult::kFailL1; return BaseCheckResult::kFailL0; } // dtype mismatch @@ -409,10 +409,10 @@ class TypeBaseChecker : public TypeFunctor(); + BaseCheckResult VisitType_(const distributed::DTensorTypeObj* lhs, const Type& other) final { + auto* rhs = other.as(); if (rhs == nullptr) { - if (other.as()) return BaseCheckResult::kFailL1; + if (other.as()) return BaseCheckResult::kFailL1; return BaseCheckResult::kFailL0; } BaseCheckResult tensor_ty_check_result = this->VisitType(lhs->tensor_ty, rhs->tensor_ty); @@ -427,19 +427,19 @@ class TypeBaseChecker : public TypeFunctor(); + BaseCheckResult VisitType_(const TupleTypeObj* lhs, const Type& other) final { + auto* rhs = other.as(); if (rhs == nullptr) { - if (other.as()) return BaseCheckResult::kFailL1; + if (other.as()) return BaseCheckResult::kFailL1; return BaseCheckResult::kFailL0; } return ArrayCheck(lhs->fields, rhs->fields); } - BaseCheckResult VisitType_(const FuncTypeNode* lhs, const Type& other) override { - auto* rhs = other.as(); + BaseCheckResult VisitType_(const FuncTypeObj* lhs, const Type& other) override { + auto* rhs = other.as(); if (rhs == nullptr) { - if (other.as()) return BaseCheckResult::kFailL1; + if (other.as()) return BaseCheckResult::kFailL1; return BaseCheckResult::kFailL0; } @@ -499,8 +499,8 @@ class TypeBaseChecker : public TypeFunctor(); - auto* int_rhs = rhs.as(); + auto* int_lhs = lhs.as(); + auto* int_rhs = rhs.as(); if (int_lhs && int_rhs) { if (int_lhs->value == int_rhs->value) { return BaseCheckResult::kPass; @@ -537,8 +537,8 @@ class TypeBaseChecker : public TypeFunctor(); - auto* rhs_shape = rhs.as(); + auto* lhs_shape = lhs.as(); + auto* rhs_shape = rhs.as(); if (lhs_shape && rhs_shape) { return ShapeMatchCheck(lhs_shape->values, rhs_shape->values); } else { @@ -645,16 +645,14 @@ class TypeBasePreconditionCollector : public TypeFunctor() != nullptr); + PrimExpr VisitType_(const StringTypeObj* lhs, const Type& other) final { + return IntImm::Bool(other.as() != nullptr); } - PrimExpr VisitType_(const PrimTypeNode* lhs, const Type& other) final { - auto* rhs = other.as(); + PrimExpr VisitType_(const PrimTypeObj* lhs, const Type& other) final { + auto* rhs = other.as(); if (rhs == nullptr) { return IntImm::Bool(false); } @@ -666,8 +664,8 @@ class TypeBasePreconditionCollector : public TypeFunctor(); + PrimExpr VisitType_(const ShapeTypeObj* lhs, const Type& other) final { + auto* rhs = other.as(); if (rhs == nullptr) { return IntImm::Bool(false); } @@ -690,8 +688,8 @@ class TypeBasePreconditionCollector : public TypeFunctor(); + PrimExpr VisitType_(const TensorTypeObj* lhs, const Type& other) final { + auto* rhs = other.as(); if (rhs == nullptr) { return IntImm::Bool(false); } @@ -734,8 +732,8 @@ class TypeBasePreconditionCollector : public TypeFunctorshape.as(); - auto* rhs_shape = rhs->shape.as(); + auto* lhs_shape = lhs->shape.as(); + auto* rhs_shape = rhs->shape.as(); if (lhs_shape && rhs_shape) { return ArrayCheck(lhs_shape->values, rhs_shape->values); } else if (lhs_shape && !rhs_shape) { @@ -745,8 +743,8 @@ class TypeBasePreconditionCollector : public TypeFunctor(); + PrimExpr VisitType_(const distributed::DTensorTypeObj* lhs, const Type& other) final { + auto* rhs = other.as(); if (rhs == nullptr) { return IntImm::Bool(false); } @@ -760,16 +758,16 @@ class TypeBasePreconditionCollector : public TypeFunctorVisitType(lhs->tensor_ty, rhs->tensor_ty); } - PrimExpr VisitType_(const TupleTypeNode* lhs, const Type& other) final { - auto* rhs = other.as(); + PrimExpr VisitType_(const TupleTypeObj* lhs, const Type& other) final { + auto* rhs = other.as(); if (rhs == nullptr) { return IntImm::Bool(false); } return ArrayCheck(lhs->fields, rhs->fields); } - PrimExpr VisitType_(const FuncTypeNode* lhs, const Type& other) override { - auto* rhs = other.as(); + PrimExpr VisitType_(const FuncTypeObj* lhs, const Type& other) override { + auto* rhs = other.as(); if (rhs == nullptr) { return IntImm::Bool(false); } @@ -924,7 +922,7 @@ class CallRetTypeDeriver : public TypeBaseChecker { return TypeBaseChecker::ShapeMatchCheck(lhs, rhs); } - if (auto* ptr = lhs.as(); ptr && !lhs.as() && !lhs.as()) { + if (auto* ptr = lhs.as(); ptr && !lhs.as() && !lhs.as()) { auto var = ffi::GetRef(ptr); auto it = var_map_.find(var); // not populated @@ -940,8 +938,8 @@ class CallRetTypeDeriver : public TypeBaseChecker { return BaseCheckResult::kFailL2; } } - auto lhs_shape = lhs.as(); - auto rhs_shape = rhs.as(); + auto lhs_shape = lhs.as(); + auto rhs_shape = rhs.as(); TVM_FFI_ICHECK(lhs_shape) << "lhs must have a shape"; if (!rhs_shape) return BaseCheckResult::kFailL2; return ShapeMatchCheck(lhs_shape->values, rhs_shape->values); @@ -1003,16 +1001,14 @@ class TypeLCAFinder : public TypeFunctor { } // AnyType is base of every Relax type, unify to Any. - Type VisitType_(const AnyTypeNode* lhs, const Type& other) final { - return ffi::GetRef(lhs); - } + Type VisitType_(const AnyTypeObj* lhs, const Type& other) final { return ffi::GetRef(lhs); } - Type VisitType_(const StringTypeNode* lhs, const Type& other) final { - return other.as() ? ffi::GetRef(lhs) : AnyType(lhs->loc); + Type VisitType_(const StringTypeObj* lhs, const Type& other) final { + return other.as() ? ffi::GetRef(lhs) : AnyType(lhs->loc); } - Type VisitType_(const PrimTypeNode* lhs, const Type& other) final { - auto* rhs = other.as(); + Type VisitType_(const PrimTypeObj* lhs, const Type& other) final { + auto* rhs = other.as(); if (rhs == nullptr) return AnyType(lhs->loc); if (lhs->dtype != rhs->dtype) { // PrimType will be treated as their boxed Any values @@ -1022,8 +1018,8 @@ class TypeLCAFinder : public TypeFunctor { return ffi::GetRef(lhs); } - Type VisitType_(const ShapeTypeNode* lhs, const Type& other) final { - auto* rhs = other.as(); + Type VisitType_(const ShapeTypeObj* lhs, const Type& other) final { + auto* rhs = other.as(); if (rhs == nullptr) return AnyType(lhs->loc); int ndim = lhs->ndim == rhs->ndim ? lhs->ndim : kUnknownNDim; @@ -1041,8 +1037,8 @@ class TypeLCAFinder : public TypeFunctor { return ffi::GetRef(lhs); } - Type VisitType_(const TensorTypeNode* lhs, const Type& other) final { - auto* rhs = other.as(); + Type VisitType_(const TensorTypeObj* lhs, const Type& other) final { + auto* rhs = other.as(); if (rhs == nullptr) return AnyType(lhs->loc); // find the target dtype, ndim, and vdevice. @@ -1077,8 +1073,8 @@ class TypeLCAFinder : public TypeFunctor { } } - Type VisitType_(const TupleTypeNode* lhs, const Type& other) final { - auto* rhs = other.as(); + Type VisitType_(const TupleTypeObj* lhs, const Type& other) final { + auto* rhs = other.as(); if (rhs == nullptr) return AnyType(lhs->loc); ffi::Optional> fields = UnifyArray(lhs->fields, rhs->fields); // tuple length not the same. @@ -1092,8 +1088,8 @@ class TypeLCAFinder : public TypeFunctor { } } - Type VisitType_(const FuncTypeNode* lhs, const Type& other) final { - auto* rhs = other.as(); + Type VisitType_(const FuncTypeObj* lhs, const Type& other) final { + auto* rhs = other.as(); if (rhs == nullptr) return AnyType(lhs->loc); // the unified function is pure only if both are pure @@ -1226,15 +1222,15 @@ class TIRVarsDetector : public TypeVisitor { } } - void VisitType_(const PrimTypeNode* prim_ty) final {} + void VisitType_(const PrimTypeObj* prim_ty) final {} - void VisitType_(const ShapeTypeNode* shape_ty) final { + void VisitType_(const ShapeTypeObj* shape_ty) final { if (shape_ty->values.has_value()) { VisitShape(shape_ty->values.value()); } } - void VisitType_(const TensorTypeNode* tensor_ty) final { + void VisitType_(const TensorTypeObj* tensor_ty) final { if (tensor_ty->shape.has_value()) { VisitType(GetType(tensor_ty->shape.value())); } @@ -1248,7 +1244,7 @@ class TIRVarsDetector : public TypeVisitor { } ffi::Array tir_vars_; - std::unordered_set used_tir_vars_dedup_; + std::unordered_set used_tir_vars_dedup_; VarType collection_type; }; @@ -1281,13 +1277,13 @@ class NonNegativeExpressionCollector : relax::TypeVisitor { } private: - void VisitType_(const TensorTypeNode* op) override { + void VisitType_(const TensorTypeObj* op) override { if (op->shape.has_value()) { VisitType(GetType(op->shape.value())); } } - void VisitType_(const PrimTypeNode* op) override { + void VisitType_(const PrimTypeObj* op) override { // Unlike the expressions in TensorType or ShapeType, // PrimType may contain negative values. This override // prevents calling VisitTypeExprField from the default @@ -1295,7 +1291,7 @@ class NonNegativeExpressionCollector : relax::TypeVisitor { } void VisitTypeExprField(const PrimExpr& size_expr) override { - if (auto size_int = size_expr.as(); size_int && size_int->value >= 0) { + if (auto size_int = size_expr.as(); size_int && size_int->value >= 0) { // Avoid cluttering the result with non-negative integers return; } @@ -1362,7 +1358,7 @@ class SymbolicVarCollector : public relax::ExprVisitor, public relax::TypeVisito kRequireDefinition = 2, }; - void VisitExpr_(const FunctionNode* op) final { + void VisitExpr_(const FunctionObj* op) final { WithMode(VisitMode::kProvideDefinition, [&]() { for (Var param : op->params) { relax::TypeVisitor::VisitType(GetType(param)); @@ -1378,7 +1374,7 @@ class SymbolicVarCollector : public relax::ExprVisitor, public relax::TypeVisito relax::ExprVisitor::VisitExpr_(op); } - void VisitBinding_(const MatchCastNode* binding) final { + void VisitBinding_(const MatchCastObj* binding) final { WithMode(VisitMode(VisitMode::kProvideDefinition | VisitMode::kRequireDefinition), [&]() { this->VisitType(binding->ty); }); @@ -1387,7 +1383,7 @@ class SymbolicVarCollector : public relax::ExprVisitor, public relax::TypeVisito void VisitExprDepTypeField(const Type& ty) { return this->VisitType(ty); } - void VisitType_(const FuncTypeNode* op) final { + void VisitType_(const FuncTypeObj* op) final { if (op->params.has_value()) { // Visit the parameters once to collect bindings, and another // time to collect usages. Otherwise, a symbolic variable @@ -1409,7 +1405,7 @@ class SymbolicVarCollector : public relax::ExprVisitor, public relax::TypeVisito } void VisitTypeExprField(const Expr& expr) final { - if (auto* shape = expr.as()) { + if (auto* shape = expr.as()) { for (const auto& val : shape->values) { this->VisitTypeExprField(val); } @@ -1432,8 +1428,8 @@ class SymbolicVarCollector : public relax::ExprVisitor, public relax::TypeVisito } } - void VisitExpr_(const VarNode* op) final { - if (!op->ty.as()) { + void VisitExpr_(const VarObj* op) final { + if (!op->ty.as()) { return; } PrimVar var = ffi::GetRef(op).as_or_throw(); @@ -1443,17 +1439,17 @@ class SymbolicVarCollector : public relax::ExprVisitor, public relax::TypeVisito } } - void VisitExpr_(const DataflowVarNode*) final {} + void VisitExpr_(const DataflowVarObj*) final {} - void VisitVarDef_(const VarNode* op) final { - if (op->ty.as()) { + void VisitVarDef_(const VarObj* op) final { + if (op->ty.as()) { defined_symbolic_var_.insert(ffi::GetRef(op).as_or_throw()); } relax::ExprVisitor::VisitVarDef_(op); } - void VisitVarDef_(const DataflowVarNode* op) final { - relax::ExprVisitor::VisitVarDef_(static_cast(op)); + void VisitVarDef_(const DataflowVarObj* op) final { + relax::ExprVisitor::VisitVarDef_(static_cast(op)); } // Run callback with mode. diff --git a/src/relax/analysis/udchain.cc b/src/relax/analysis/udchain.cc index f993a27c8944..62dba55dd032 100644 --- a/src/relax/analysis/udchain.cc +++ b/src/relax/analysis/udchain.cc @@ -64,7 +64,7 @@ class UDChain : relax::ExprVisitor { ffi::Optional cur_user_; - void VisitBinding_(const VarBindingNode* binding) override { + void VisitBinding_(const VarBindingObj* binding) override { TVM_FFI_ICHECK(!bound_values.count(binding->var)) << "Variable " << binding->var << " was defined multiple times"; bound_values.Set(binding->var, binding->value); @@ -75,7 +75,7 @@ class UDChain : relax::ExprVisitor { cur_user_ = cache; } - void VisitBinding_(const VarBindingNode* binding, const FunctionNode* func) override { + void VisitBinding_(const VarBindingObj* binding, const FunctionObj* func) override { // A local Relax function may be recursively defined. References to // `binding->var` that appear within `func` are valid. DefineVar(binding->var); @@ -90,7 +90,7 @@ class UDChain : relax::ExprVisitor { DefineVar(var); } } - void VisitExpr_(const VarNode* op) override { + void VisitExpr_(const VarObj* op) override { auto var = ffi::GetRef(op); if (cur_user_) { @@ -100,7 +100,7 @@ class UDChain : relax::ExprVisitor { } } - void VisitExpr_(const FunctionNode* op) override { + void VisitExpr_(const FunctionObj* op) override { cur_user_ = std::nullopt; ExprVisitor::VisitExpr_(op); } @@ -129,7 +129,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { .def("relax.analysis.used_vars", [](const Expr& expr) { auto used_vars = GetUsedVars(expr); ffi::Array result; - for (const VarNode* var_node : used_vars) { + for (const VarObj* var_node : used_vars) { result.push_back(ffi::GetRef(var_node)); } return result; @@ -138,12 +138,12 @@ TVM_FFI_STATIC_INIT_BLOCK() { VarUsageInfo CollectVarUsage(const Expr& expr) { return UDChain::Collect(expr); } -std::set GetUsedVars(const Expr& expr) { +std::set GetUsedVars(const Expr& expr) { class UsedVars : public ExprVisitor { public: - std::set used_vars; + std::set used_vars; void VisitExprDepTypeField(const Type&) final {} - void VisitExpr_(const VarNode* op) override { used_vars.insert(op); } + void VisitExpr_(const VarObj* op) override { used_vars.insert(op); } } visitor; visitor.VisitExpr(expr); return std::move(visitor.used_vars); diff --git a/src/relax/analysis/var2value.cc b/src/relax/analysis/var2value.cc index e18d0e7fcabb..00c39ecb65ff 100644 --- a/src/relax/analysis/var2value.cc +++ b/src/relax/analysis/var2value.cc @@ -30,7 +30,7 @@ namespace relax { class Var2ValAnalysis : public relax::ExprVisitor { public: ffi::Map var2value_; - void VisitBinding_(const VarBindingNode* binding) override { + void VisitBinding_(const VarBindingObj* binding) override { var2value_.Set(binding->var, binding->value); // Recursively visit the value to handle local functions. VisitExpr(binding->value); @@ -54,7 +54,7 @@ ffi::Map AnalyzeVar2Value(const IRModule& m) { for (const auto& it : m->functions) { // visit relax.Function - if (auto* n = it.second.as()) { + if (auto* n = it.second.as()) { var2val_analysis.VisitExpr(ffi::GetRef(n)); } } @@ -73,12 +73,12 @@ class Name2BindingAnalysis : public relax::ExprVisitor { // Map is not suitable for doing in-place update. // so we use standard container for internal usage. std::map> name2bindings_; - void VisitBinding_(const VarBindingNode* binding) override { + void VisitBinding_(const VarBindingObj* binding) override { const auto& vname = binding->var->name; name2bindings_[vname].push_back(ffi::GetRef(binding)); } - void VisitBinding_(const MatchCastNode* binding) override { + void VisitBinding_(const MatchCastObj* binding) override { const auto& vname = binding->var->name; name2bindings_[vname].push_back(ffi::GetRef(binding)); } diff --git a/src/relax/analysis/well_formed.cc b/src/relax/analysis/well_formed.cc index 37cb612e8571..147f1b781a3a 100644 --- a/src/relax/analysis/well_formed.cc +++ b/src/relax/analysis/well_formed.cc @@ -35,10 +35,10 @@ * 5. Vars, including those introduced in definitional type positions, are defined before use. * 6. Vars are defined exactly once. * 7. DataflowVars cannot be defined inside BindingBlock. - * 8. Vars defined in IfExprNode, except the return Var, are invisible + * 8. Vars defined in IfExprObj, except the return Var, are invisible * out of the If body.(May change for new AST designs) * 9. SeqExpr only serves as function body, or in the true and - * false branches in IfExprNode. + * false branches in IfExprObj. * 10. The IR is in ANF: * (a) Expressions cannot contain nested complex expressions. * Here are the expressions that may be nested inside other expressions: @@ -99,17 +99,17 @@ class WellFormedChecker : public relax::ExprVisitor, public relax::TypeVisitor { static void Check(ffi::Variant obj, bool check_ty) { WellFormedChecker well_formed_checker = WellFormedChecker(obj.as(), check_ty); - if (const auto* mod = obj.as()) { + if (const auto* mod = obj.as()) { for (const auto& it : mod->functions) { // visit relax.Function - if (auto* n = it.second.as()) { + if (auto* n = it.second.as()) { Function func = ffi::GetRef(n); well_formed_checker.func_name_map_[n] = it.first->name_hint; well_formed_checker.CheckGlobalVarAndGsymbolConsistency(it.first, func); well_formed_checker.VisitExpr(func); } } - } else if (const auto* func = obj.as()) { + } else if (const auto* func = obj.as()) { well_formed_checker.VisitExpr(ffi::GetRef(func)); } else { TVM_FFI_THROW(InternalError) << "Unreachable, " @@ -151,21 +151,21 @@ class WellFormedChecker : public relax::ExprVisitor, public relax::TypeVisitor { parent_(parent) {} private: - ffi::Optional Visit_(const tvm::VarNode* op) final { + ffi::Optional Visit_(const tvm::VarObj* op) final { parent_->VisitExpr(ffi::GetRef(op)); return std::nullopt; } static void InitVTable(VTable* table) { tirx::StmtExprVisitor::InitVTable(table); - SetDispatch(table); + SetDispatch(table); } WellFormedChecker* parent_{nullptr}; }; /*! \brief Get the name of a function for use in error messages. */ - std::string FuncName(const FunctionNode* func) const { + std::string FuncName(const FunctionObj* func) const { auto it = func_name_map_.find(func); if (it != func_name_map_.end()) { return "\"" + it->second + "\""; @@ -187,13 +187,13 @@ class WellFormedChecker : public relax::ExprVisitor, public relax::TypeVisitor { } void VisitExpr(const Expr& expr) final { - if (!expr.as() && expr->ty.as().has_value()) { + if (!expr.as() && expr->ty.as().has_value()) { TVM_FFI_VISIT_THROW(TypeError, expr) << "The ty of Expr " << expr << " is missing."; } relax::ExprVisitor::VisitExpr(expr); } - void VisitExpr_(const GlobalVarNode* op) final { + void VisitExpr_(const GlobalVarObj* op) final { GlobalVar var = ffi::GetRef(op); if (mod_.has_value()) { if (!(mod_.value()->ContainGlobalVar(var->name_hint) && @@ -204,12 +204,12 @@ class WellFormedChecker : public relax::ExprVisitor, public relax::TypeVisitor { } if (!op->ty.as().has_value()) { - if (op->ty.as()) { + if (op->ty.as()) { // A native callee owns its parameter symbols and tensor contracts. // Explicit call bridges validate their Relax-facing use separately. return; } - if (!op->ty->IsInstance()) { + if (!op->ty->IsInstance()) { TVM_FFI_VISIT_THROW(TypeError, var) << "The ty of GlobalVar " << ffi::GetRef(op) << " must be either FuncType."; } @@ -218,7 +218,7 @@ class WellFormedChecker : public relax::ExprVisitor, public relax::TypeVisitor { CheckType(op); } - void VisitExpr_(const TupleNode* op) final { + void VisitExpr_(const TupleObj* op) final { TVM_FFI_VISIT_BEGIN(); for (size_t i = 0; i < op->fields.size(); i++) { Expr expr = op->fields[i]; @@ -234,7 +234,7 @@ class WellFormedChecker : public relax::ExprVisitor, public relax::TypeVisitor { TVM_FFI_VISIT_END(ffi::GetRef(op)); } - void VisitExpr_(const TupleGetItemNode* op) final { + void VisitExpr_(const TupleGetItemObj* op) final { if (IsLeafOrTuple(op->tuple)) { this->VisitExpr(op->tuple); } else { @@ -244,7 +244,7 @@ class WellFormedChecker : public relax::ExprVisitor, public relax::TypeVisitor { CheckType(op); } - void VisitExpr_(const VarNode* op) final { + void VisitExpr_(const VarObj* op) final { Var var = ffi::GetRef(op); if (var_set_.count(var) == 0 && recur_vars_.count(var) == 0) { TVM_FFI_VISIT_THROW(ValueError, var) << "Var " << ffi::GetRef(op) << " is not defined."; @@ -252,7 +252,7 @@ class WellFormedChecker : public relax::ExprVisitor, public relax::TypeVisitor { CheckType(op); } - void VisitExpr_(const DataflowVarNode* op) final { + void VisitExpr_(const DataflowVarObj* op) final { DataflowVar var = ffi::GetRef(op); if (!is_dataflow_) { TVM_FFI_VISIT_THROW(ValueError, var) @@ -265,7 +265,7 @@ class WellFormedChecker : public relax::ExprVisitor, public relax::TypeVisitor { CheckType(op); } - void VisitExpr_(const FunctionNode* op) final { + void VisitExpr_(const FunctionObj* op) final { TVM_FFI_VISIT_BEGIN(); // set current visited function. // for nested functions, we only set the outermost function. @@ -304,18 +304,18 @@ class WellFormedChecker : public relax::ExprVisitor, public relax::TypeVisitor { // Primitive parameters may supply dimensions anywhere in the signature. // Register each once so duplicate parameters remain an error. for (Var param : op->params) { - if (GetType(param).as()) { + if (GetType(param).as()) { RegisterVarDefinition(param); } } // Other runtime parameters retain source-order definition scope. for (Var param : op->params) { - if (auto* dataflow_var = param.as()) { + if (auto* dataflow_var = param.as()) { TVM_FFI_VISIT_THROW(ValueError, ffi::GetRef(dataflow_var)) << "DataflowVar " << param << " is defined outside DataflowBlock."; } - if (!GetType(param).as()) { + if (!GetType(param).as()) { RegisterVarDefinition(param); } @@ -359,13 +359,13 @@ class WellFormedChecker : public relax::ExprVisitor, public relax::TypeVisitor { TVM_FFI_VISIT_END(ffi::GetRef(op)); } - void VisitExpr_(const CallNode* call) final { + void VisitExpr_(const CallObj* call) final { TVM_FFI_VISIT_BEGIN(); - TVM_FFI_CHECK(call->op.as() || !call->op->ty.as(), TypeError) + TVM_FFI_CHECK(call->op.as() || !call->op->ty.as(), TypeError) << "Ordinary Relax calls cannot invoke a native TIRx function; " << "use R.call_tir or R.call_tir_packed"; if (IsLeafOrTuple(call->op)) { - const FunctionNode* prev_visited_func = cur_visited_func_; + const FunctionObj* prev_visited_func = cur_visited_func_; cur_visited_func_ = nullptr; // do not attribute the callee to its caller this->VisitExpr(call->op); cur_visited_func_ = prev_visited_func; @@ -438,7 +438,7 @@ class WellFormedChecker : public relax::ExprVisitor, public relax::TypeVisitor { op_map_infer_type_.count(op.value()) || op_map_infer_type_with_builder_.count(op.value()); } if (check_ty && !call->ty.as().has_value() && - (!call->ty.as() || has_infer_type)) { + (!call->ty.as() || has_infer_type)) { // The `InferType` method isn't currently exposed by the // Normalizer, and can only be called indirectly by normalizing // an expression that does not yet have `Type`. @@ -482,7 +482,7 @@ class WellFormedChecker : public relax::ExprVisitor, public relax::TypeVisitor { TVM_FFI_VISIT_END(ffi::GetRef(call)); } - void VisitExpr_(const IfExprNode* op) final { + void VisitExpr_(const IfExprObj* op) final { TVM_FFI_VISIT_BEGIN(); if (is_dataflow_) { TVM_FFI_VISIT_THROW(ValueError, ffi::GetRef(op)) @@ -505,7 +505,7 @@ class WellFormedChecker : public relax::ExprVisitor, public relax::TypeVisitor { TVM_FFI_VISIT_END(ffi::GetRef(op)); } - void VisitExpr_(const ShapeExprNode* op) final { + void VisitExpr_(const ShapeExprObj* op) final { for (PrimExpr expr : op->values) { // check if the symbolic vars in the expr are defined, e.g, 2 * m VisitPrimitiveExpr(expr); @@ -517,16 +517,16 @@ class WellFormedChecker : public relax::ExprVisitor, public relax::TypeVisitor { CheckType(op); } - void VisitExpr_(const SeqExprNode* op) final { + void VisitExpr_(const SeqExprObj* op) final { TVM_FFI_VISIT_THROW(ValueError, ffi::GetRef(op)) - << "SeqExpr only serves as the function body in FunctionNode, " - "or the true/false branch body in IfExprNode."; + << "SeqExpr only serves as the function body in FunctionObj, " + "or the true/false branch body in IfExprObj."; } - void VisitSeqExpr(const SeqExprNode* op) { + void VisitSeqExpr(const SeqExprObj* op) { TVM_FFI_VISIT_BEGIN(); // a special call only if SeqExpr is the function body - // in FunctionNode or the true/false branch body in IfExprNode + // in FunctionObj or the true/false branch body in IfExprObj for (BindingBlock block : op->blocks) { this->VisitBindingBlock(block); } @@ -539,13 +539,13 @@ class WellFormedChecker : public relax::ExprVisitor, public relax::TypeVisitor { TVM_FFI_VISIT_END(ffi::GetRef(op)); } - void VisitBinding_(const VarBindingNode* binding) final { + void VisitBinding_(const VarBindingObj* binding) final { bool is_lambda = false; - if (binding->value->IsInstance()) { + if (binding->value->IsInstance()) { is_lambda = true; recur_vars_.insert(binding->var); } - if (binding->value->IsInstance()) { + if (binding->value->IsInstance()) { TVM_FFI_VISIT_THROW(ValueError, binding->value) << "Inline tirx::Function is disallowed in Relax IR."; } else { @@ -570,7 +570,7 @@ class WellFormedChecker : public relax::ExprVisitor, public relax::TypeVisitor { } } - void VisitBinding_(const MatchCastNode* binding) final { + void VisitBinding_(const MatchCastObj* binding) final { this->VisitExpr(binding->value); // define the vars WithMode(VisitMode::kMatchVarDef, [&]() { this->VisitType(binding->ty); }); @@ -579,7 +579,7 @@ class WellFormedChecker : public relax::ExprVisitor, public relax::TypeVisitor { this->VisitVarDef(binding->var); } - void VisitBindingBlock_(const DataflowBlockNode* block) final { + void VisitBindingBlock_(const DataflowBlockObj* block) final { bool old_is_dataflow_ = is_dataflow_; is_dataflow_ = true; for (Binding binding : block->bindings) { @@ -589,7 +589,7 @@ class WellFormedChecker : public relax::ExprVisitor, public relax::TypeVisitor { dataflow_var_set_.clear(); } - void VisitVarDef_(const DataflowVarNode* var) final { + void VisitVarDef_(const DataflowVarObj* var) final { if (!is_dataflow_) { TVM_FFI_VISIT_THROW(ValueError, ffi::GetRef(var)) << "DataflowVar " << var << " is defined outside DataflowBlock."; @@ -603,7 +603,7 @@ class WellFormedChecker : public relax::ExprVisitor, public relax::TypeVisitor { CheckType(var); } - void VisitVarDef_(const VarNode* var) final { + void VisitVarDef_(const VarObj* var) final { Var gv = ffi::GetRef(var); RegisterVarDefinition(gv); CheckType(var); @@ -616,9 +616,9 @@ class WellFormedChecker : public relax::ExprVisitor, public relax::TypeVisitor { var_set_.insert(var); } - void VisitType_(const FuncTypeNode* op) final { + void VisitType_(const FuncTypeObj* op) final { auto prev_var_set = var_set_; - const FunctionNode* prev_visited_func = cur_visited_func_; + const FunctionObj* prev_visited_func = cur_visited_func_; cur_visited_func_ = nullptr; if (op->params.has_value()) { WithMode(VisitMode::kMatchVarDef, [&]() { @@ -634,16 +634,16 @@ class WellFormedChecker : public relax::ExprVisitor, public relax::TypeVisitor { } void VisitTypeExprField(const Expr& expr) final { - if (auto* op = expr.as()) { + if (auto* op = expr.as()) { TVM_FFI_VISIT_THROW(ValueError, ffi::GetRef(op)) << "DataflowVar cannot be used as a dependent type expression."; } if (mode_ == VisitMode::kMatchVarDef) { // Populate a canonical Var definition on first occurrence. - if (auto* op = expr.as()) { + if (auto* op = expr.as()) { MarkTypeVarDefinition(ffi::GetRef(op)); } - if (auto* shape = expr.as()) { + if (auto* shape = expr.as()) { for (auto val : shape->values) { this->VisitTypeExprField(val); } @@ -654,7 +654,7 @@ class WellFormedChecker : public relax::ExprVisitor, public relax::TypeVisitor { } void VisitTypeExprField(const PrimExpr& expr) final { - if (auto* op = expr.as()) { + if (auto* op = expr.as()) { TVM_FFI_VISIT_THROW(ValueError, ffi::GetRef(op)) << "DataflowVar cannot be used as a dependent type expression."; } @@ -662,7 +662,7 @@ class WellFormedChecker : public relax::ExprVisitor, public relax::TypeVisitor { // A primitive Var is the same canonical Var as a Relax binding. A // first occurrence in a definitional type position enters the one // lexical definition set used by all Var occurrences. - if (auto* op = expr.as()) { + if (auto* op = expr.as()) { MarkTypeVarDefinition(ffi::GetRef(op)); } } else { @@ -688,12 +688,12 @@ class WellFormedChecker : public relax::ExprVisitor, public relax::TypeVisitor { param_var_func_map_.insert({var, cur_visited_func_}); } - void CheckType(const ExprNode* op) { + void CheckType(const ExprObj* op) { if (!check_ty) { return; } - if (auto* ty = op->ty.as()) { + if (auto* ty = op->ty.as()) { this->VisitType(ffi::GetRef(ty)); } else { TVM_FFI_VISIT_THROW(TypeError, ffi::GetRef(op)) @@ -714,16 +714,16 @@ class WellFormedChecker : public relax::ExprVisitor, public relax::TypeVisitor { const bool check_ty; bool is_dataflow_; // Current visited function. - const FunctionNode* cur_visited_func_; + const FunctionObj* cur_visited_func_; // Map from function pointer to its global name (for error messages). - std::unordered_map func_name_map_; + std::unordered_map func_name_map_; // Current visit mode. VisitMode mode_ = VisitMode::kDefault; // set of context variables. std::unordered_set var_set_; std::unordered_set recur_vars_; std::unordered_set dataflow_var_set_; - std::unordered_map param_var_func_map_; + std::unordered_map param_var_func_map_; decltype(Op::GetAttrMap(tvm::relax::op_attr::kNormalize)) op_map_normalize_ = Op::GetAttrMap(tvm::relax::op_attr::kNormalize); diff --git a/src/relax/backend/adreno/annotate_custom_storage.cc b/src/relax/backend/adreno/annotate_custom_storage.cc index a07e403bcbd6..33916c16ac41 100644 --- a/src/relax/backend/adreno/annotate_custom_storage.cc +++ b/src/relax/backend/adreno/annotate_custom_storage.cc @@ -121,7 +121,7 @@ * on consumer scope information. It does evaluating mutiple consumer cases and conflicts. * 3: DefineVDevice: Pass does injects hint_on_device for each argument. It also tries to update * out Type containing VDevice information. This update for tirx calls is straight forward - * as ty_args in CallNode is meant for this purpose. This ty_args for other calls by + * as ty_args in CallObj is meant for this purpose. This ty_args for other calls by * design is invalid as we do this by per-op type inference. * Another issue with type inference per op is that it cannot decide this * memory scope information which is done by this pass based on consumer demand. @@ -309,8 +309,8 @@ class CollectConsumerScopeInfo : public ExprVisitor { return std::make_pair(call_scope_info, scope_info); } - void VisitBinding_(const VarBindingNode* binding, - const TupleGetItemNode* tuple_get_item_node) final { + void VisitBinding_(const VarBindingObj* binding, + const TupleGetItemObj* tuple_get_item_node) final { /* * lv9 = R.call_tir(add, (lv7, m * lv: R.Tuple( @@ -335,11 +335,11 @@ class CollectConsumerScopeInfo : public ExprVisitor { tuple_item_to_binding.Set(ffi::GetRef(binding->var.get()), field_map); } - void VisitBinding_(const VarBindingNode* binding, const TupleNode* tuple) final { + void VisitBinding_(const VarBindingObj* binding, const TupleObj* tuple) final { tuples_to_binding.Set(ffi::GetRef(binding->var.get()), tuple->fields); } - void VisitExpr_(const CallNode* call) final { + void VisitExpr_(const CallObj* call) final { static const Op call_tir_op = Op::Get("relax.call_tir"); ffi::Optional gv; ffi::Array op_attrs; @@ -367,7 +367,7 @@ class CollectConsumerScopeInfo : public ExprVisitor { i < is_texture_supported.size() ? is_texture_supported[i] : is_texture_supported[0]; auto scope = is_texture ? Scope(GetShapeFromTensorType(tensor_ty.value())) : "global"; ffi::Map> ent_call; - const VarNode* arg_var = func_args->fields[i].as(); + const VarObj* arg_var = func_args->fields[i].as(); if (scope_info.find(ffi::GetRef(arg_var)) != scope_info.end()) { ent_call = scope_info[ffi::GetRef(arg_var)]; } @@ -429,9 +429,9 @@ class CollectConsumerScopeInfo : public ExprVisitor { // currently we support only textures been made from 5d tensors // 5d requirement is not limitation of textures in general, it is limitation how // we are representing memory scopes/layout and flattening of textures in tirx - if (shape.size() == 5 && shape[4].as()->value == 4) { + if (shape.size() == 5 && shape[4].as()->value == 4) { for (auto ind : shape) { - if (!ind.as()) { + if (!ind.as()) { // Dynamic tensors return "global.texture-nchw"; } @@ -441,10 +441,10 @@ class CollectConsumerScopeInfo : public ExprVisitor { static_cast(target_->GetAttr("texture_spatial_limit").value_or(16384)); int depth_limit = static_cast(target_->GetAttr("texture_depth_limit").value_or(2048)); - int a0 = shape[0].as()->value.as().value(); - int a1 = shape[1].as()->value.as().value(); - int a2 = shape[2].as()->value.as().value(); - int a3 = shape[3].as()->value.as().value(); + int a0 = shape[0].as()->value.as().value(); + int a1 = shape[1].as()->value.as().value(); + int a2 = shape[2].as()->value.as().value(); + int a3 = shape[3].as()->value.as().value(); int d1r = a0 * a1; int d2r = a2 * a3; @@ -496,7 +496,7 @@ class CollectProducerScopeInfo : public ExprVisitor { return producer_ty; } - void VisitBinding_(const VarBindingNode* binding, const CallNode* call) final { + void VisitBinding_(const VarBindingObj* binding, const CallObj* call) final { ExprVisitor::VisitBinding_(binding, call); static const Op call_tir_op = Op::Get("relax.call_tir"); @@ -505,7 +505,7 @@ class CollectProducerScopeInfo : public ExprVisitor { if (call->op.same_as(call_tir_op)) { out_ty = call->ty_args[0]; } else { - auto* op_ptr = call->op.as(); + auto* op_ptr = call->op.as(); Op op = ffi::GetRef(op_ptr); static auto op_map_context_free = Op::GetAttrMap(tvm::op_attr::kInferType); if (op_map_context_free.count(op)) { @@ -521,7 +521,7 @@ class CollectProducerScopeInfo : public ExprVisitor { std::unordered_map scope_count; // Decide the final scope based on the max consumer demand. Rest will use to_device. - auto arg_var = binding->var.as(); + auto arg_var = binding->var.as(); if (scope_info_.find(ffi::GetRef(arg_var)) != scope_info_.end()) { for (const auto& val : scope_info_[ffi::GetRef(arg_var)]) { auto call_node = val.first.as_or_throw(); @@ -548,14 +548,14 @@ class CollectProducerScopeInfo : public ExprVisitor { private: Type UpdateOutputType(const Type& out_ty, ffi::Array scope) { - if (out_ty->IsInstance()) { + if (out_ty->IsInstance()) { auto tensor_ty = out_ty.as_or_throw(); auto shape_arr = GetShapeFromTensorType(tensor_ty); return TensorType(ShapeExpr(shape_arr), tensor_ty->dtype, VDevice(target_, 0, runtime::MemoryScope::Get(scope[0]))); } - TVM_FFI_ICHECK(out_ty->IsInstance()) + TVM_FFI_ICHECK(out_ty->IsInstance()) << "Expect output type of call_tir to be either TupleType or " "TensorType, but got " << out_ty; @@ -563,7 +563,7 @@ class CollectProducerScopeInfo : public ExprVisitor { const auto& tuple_ty = out_ty.as_or_throw(); ffi::Array ty_fields; for (const auto& si : tuple_ty->fields) { - TVM_FFI_ICHECK(si->IsInstance()) + TVM_FFI_ICHECK(si->IsInstance()) << "Fields of TupleType must be TensorType for call_tir " "output structinfo, but got " << si; @@ -596,7 +596,7 @@ class DefineVDevice : ExprMutator { IRModule Run(IRModule& mod) { mod_ = mod; for (const auto& [gv, func] : mod_->functions) { - if (func->IsInstance()) { + if (func->IsInstance()) { const auto& base_func = mod_->Lookup(gv); // Only non primitive relax functions if (base_func->HasNonzeroAttr(tvm::relax::attr::kPrimitive.name)) { @@ -630,7 +630,7 @@ class DefineVDevice : ExprMutator { using ExprMutator::VisitExpr_; - Expr VisitExpr_(const CallNode* call_node) override { + Expr VisitExpr_(const CallObj* call_node) override { auto call = ExprMutator::VisitExpr_(call_node).as_or_throw(); static const Op call_tir_op = Op::Get("relax.call_tir"); @@ -647,7 +647,7 @@ class DefineVDevice : ExprMutator { ffi::Array new_args; Type updated_ret_ty = producer_ty_[ffi::GetRef(call_node)]; - if (updated_ret_ty->IsInstance()) { + if (updated_ret_ty->IsInstance()) { auto tensor_ty = updated_ret_ty.as_or_throw(); auto shape = tensor_ty->shape.value(); auto dtype = tensor_ty->dtype; @@ -657,7 +657,7 @@ class DefineVDevice : ExprMutator { updated_ret_ty = TensorType(shape, dtype, vdev_global); } } else { - TVM_FFI_ICHECK(updated_ret_ty->IsInstance()) + TVM_FFI_ICHECK(updated_ret_ty->IsInstance()) << "Expect output type of call_tir to be either TupleType or " "TensorType, but got " << updated_ret_ty; @@ -665,7 +665,7 @@ class DefineVDevice : ExprMutator { const auto& tuple_ty = updated_ret_ty.as_or_throw(); ffi::Array ty_fields; for (const auto& si : tuple_ty->fields) { - TVM_FFI_ICHECK(si->IsInstance()) + TVM_FFI_ICHECK(si->IsInstance()) << "Fields of TupleType must be TensorType for call_tir " "output structinfo, but got " << si; @@ -725,8 +725,8 @@ class DefineVDevice : ExprMutator { } Expr HintArg(const Expr& arg, ffi::String scope) { - if (arg->IsInstance()) { - if (auto tensor_ty = arg->ty.as()) { + if (arg->IsInstance()) { + if (auto tensor_ty = arg->ty.as()) { if (!tensor_ty->vdevice.has_value()) { const VDevice& vdev = MakeGlobalVDevice(VDevice(target_, 0, runtime::MemoryScope::Get(scope))); @@ -749,7 +749,7 @@ class DefineVDevice : ExprMutator { } ffi::Optional GetTarget(const Type& ty) { - auto tinfo = ty.as(); + auto tinfo = ty.as(); if (tinfo->vdevice.has_value()) { auto vdevice = tinfo->vdevice.value(); if (vdevice->target.has_value()) { diff --git a/src/relax/backend/adreno/fold_vdevice_scope_change.cc b/src/relax/backend/adreno/fold_vdevice_scope_change.cc index 1312ea608ccf..c6b1c5227291 100644 --- a/src/relax/backend/adreno/fold_vdevice_scope_change.cc +++ b/src/relax/backend/adreno/fold_vdevice_scope_change.cc @@ -54,13 +54,13 @@ std::tuple)>> auto pattern_out = IsOp("relax.to_vdevice")(pat_call_tir); auto rewriter = [=](Expr expr, ffi::Map matches) -> Expr { - const auto* call_tir = matches[pat_call_tir].as(); + const auto* call_tir = matches[pat_call_tir].as(); TVM_FFI_CHECK(call_tir, InternalError) << "Match of relax.call_tir operator should produce Call, " << "but instead produces " << matches[pat_call_tir] << " with type " << matches[pat_call_tir]->GetTypeKey(); - const auto* out = matches[pattern_out].as(); + const auto* out = matches[pattern_out].as(); TVM_FFI_CHECK(out, InternalError) << "Match of relax.to_vdevice operator should produce Call, " << "but instead produces " << matches[pattern_out] << " with type " << matches[pattern_out]->GetTypeKey(); @@ -70,12 +70,12 @@ std::tuple)>> << "Attributes for relax.to_vdevice operator should be ToVDeviceAttrs, " << "but were instead " << out->attrs << " with type " << out->GetTypeKey(); - const auto* tir_out_ty = call_tir->ty_args[0].as(); + const auto* tir_out_ty = call_tir->ty_args[0].as(); if (!tir_out_ty) return expr; if (!tir_out_ty->vdevice.has_value()) return expr; - const VarNode* arg_var = out->args[0].as(); + const VarObj* arg_var = out->args[0].as(); if (consumers.find(ffi::GetRef(arg_var)) != consumers.end()) { if (consumers[ffi::GetRef(arg_var)].size() > 1) { /* Don't do to_device optimization as we are not the only consumer */ @@ -125,15 +125,15 @@ class CollectConsumerDetails : public ExprVisitor { return consumers; } - void VisitBinding_(const VarBindingNode* binding, - const TupleGetItemNode* tuple_get_item_node) final { + void VisitBinding_(const VarBindingObj* binding, + const TupleGetItemObj* tuple_get_item_node) final { if (arg_to_binding.find(ffi::GetRef(binding->var.get())) == arg_to_binding.end()) { arg_to_binding.Set(ffi::GetRef(binding->var.get()), ffi::GetRef(tuple_get_item_node->tuple.get())); } } - void VisitExpr_(const CallNode* call) final { + void VisitExpr_(const CallObj* call) final { static const Op call_tir_op = Op::Get("relax.call_tir"); Tuple func_args = call->op.same_as(call_tir_op) ? call->args[1].as_or_throw() : Tuple(call->args); @@ -143,7 +143,7 @@ class CollectConsumerDetails : public ExprVisitor { if (auto tensor_ty = ty.as()) { ffi::Array call_list; - const VarNode* arg_var = arg.as(); + const VarObj* arg_var = arg.as(); if (consumers.find(ffi::GetRef(arg_var)) != consumers.end()) { call_list = consumers[ffi::GetRef(arg_var)]; diff --git a/src/relax/backend/contrib/clml/codegen.cc b/src/relax/backend/contrib/clml/codegen.cc index cb8a678366e5..8b0fe2919ac1 100644 --- a/src/relax/backend/contrib/clml/codegen.cc +++ b/src/relax/backend/contrib/clml/codegen.cc @@ -85,10 +85,10 @@ class CollectCLMLFromCompositeFunctionBody : public ExprVisitor { explicit CollectCLMLFromCompositeFunctionBody(OpenCLMLJSONSerializer* serializer) : serializer_(serializer), node_(std::make_shared()) {} - void VisitExpr_(const GenericConstNode* constant_node) final; - void VisitExpr_(const CallNode* call_node) final; + void VisitExpr_(const GenericConstObj* constant_node) final; + void VisitExpr_(const CallObj* call_node) final; - void SetGenericAttributes(const CallNode* call_node) { + void SetGenericAttributes(const CallObj* call_node) { if (backend::IsOp(call_node, "relax.nn.relu")) { node_->SetAttr("activation_type", ffi::String("relu")); } @@ -124,19 +124,19 @@ class OpenCLMLJSONSerializer : public JSONSerializer { * convolution. Supports nn.conv2d */ struct CompositeConvNode { - const CallNode* pad = nullptr; - const CallNode* conv = nullptr; - const CallNode* bn = nullptr; - const CallNode* bias = nullptr; - const CallNode* activation = nullptr; + const CallObj* pad = nullptr; + const CallObj* conv = nullptr; + const CallObj* bn = nullptr; + const CallObj* bias = nullptr; + const CallObj* activation = nullptr; std::string act_type; }; using JSONSerializer::VisitExpr_; - std::vector VisitExpr_(const CallNode* call_node) final { + std::vector VisitExpr_(const CallObj* call_node) final { // The call must be to an inline "Composite" function - const auto* fn_var = call_node->op.as(); + const auto* fn_var = call_node->op.as(); TVM_FFI_ICHECK(fn_var); const auto fn = bindings_[ffi::GetRef(fn_var)].as_or_throw(); @@ -177,7 +177,7 @@ class OpenCLMLJSONSerializer : public JSONSerializer { VLOG(1) << name << " has " << node->GetInputs().size() << " inputs"; } - return AddNode(node, ffi::GetRef(call_node)); + return AddObj(node, ffi::GetRef(call_node)); } /*! @@ -186,10 +186,10 @@ class OpenCLMLJSONSerializer : public JSONSerializer { * \param cn The call node of the composite function. * \return Extracted composite convolution nodes. */ - CompositeConvNode UnpackCompositeConvolution(const CallNode* cn) { + CompositeConvNode UnpackCompositeConvolution(const CallObj* cn) { CompositeConvNode nodes{}; - const auto* fn_var = cn->op.as(); + const auto* fn_var = cn->op.as(); TVM_FFI_ICHECK(fn_var); const auto fn = bindings_[ffi::GetRef(fn_var)].as_or_throw(); auto opt_composite = fn->GetAttr(tvm::relax::attr::kComposite); @@ -215,10 +215,10 @@ class OpenCLMLJSONSerializer : public JSONSerializer { * \param cn The call to be represented. * \return A JSON representation of a specific operator. */ - std::shared_ptr CreateCompositeConvJSONNode(const CallNode* cn) { + std::shared_ptr CreateCompositeConvJSONNode(const CallObj* cn) { CompositeConvNode nodes = UnpackCompositeConvolution(cn); - const auto* fn_var = cn->op.as(); + const auto* fn_var = cn->op.as(); TVM_FFI_ICHECK(fn_var); const auto fn = bindings_[ffi::GetRef(fn_var)].as_or_throw(); auto opt_composite = fn->GetAttr(tvm::relax::attr::kComposite); @@ -278,13 +278,13 @@ class OpenCLMLJSONSerializer : public JSONSerializer { ffi::Map bindings_; }; -void CollectCLMLFromCompositeFunctionBody::VisitExpr_(const GenericConstNode* constant_node) { +void CollectCLMLFromCompositeFunctionBody::VisitExpr_(const GenericConstObj* constant_node) { for (const auto& entry : serializer_->VisitExpr(ffi::GetRef(constant_node))) { args_.emplace_back(entry); } } -void CollectCLMLFromCompositeFunctionBody::VisitExpr_(const CallNode* call_node) { +void CollectCLMLFromCompositeFunctionBody::VisitExpr_(const CallObj* call_node) { SetGenericAttributes(call_node); ExprVisitor::VisitExpr_(call_node); } diff --git a/src/relax/backend/contrib/codegen_c/codegen_c.h b/src/relax/backend/contrib/codegen_c/codegen_c.h index 3b907d24e9c9..5796b3b12847 100644 --- a/src/relax/backend/contrib/codegen_c/codegen_c.h +++ b/src/relax/backend/contrib/codegen_c/codegen_c.h @@ -332,19 +332,19 @@ class CodegenCBase { * \return The dtype string. */ std::string GetDtypeString(const Var& var) { - auto tensor_ty = var->ty.as(); - TVM_FFI_ICHECK(tensor_ty) << "Expect TensorTypeNode"; + auto tensor_ty = var->ty.as(); + TVM_FFI_ICHECK(tensor_ty) << "Expect TensorTypeObj"; return GetDtypeString(tensor_ty); } /*! * \brief Returns dtype string * - * \param ttype TensorTypeNode* to get the dtype of + * \param ttype TensorTypeObj* to get the dtype of * * \return The dtype string. */ - std::string GetDtypeString(const TensorTypeNode* tensor_ty) { + std::string GetDtypeString(const TensorTypeObj* tensor_ty) { std::string dtype; DLDataType raw_dtype = tensor_ty->dtype.value()->dtype; if (raw_dtype == DLDataType{kDLFloat, 32, 1}) { diff --git a/src/relax/backend/contrib/codegen_json/codegen_json.h b/src/relax/backend/contrib/codegen_json/codegen_json.h index e6fe066d9a9a..ea36b4f7402a 100644 --- a/src/relax/backend/contrib/codegen_json/codegen_json.h +++ b/src/relax/backend/contrib/codegen_json/codegen_json.h @@ -112,8 +112,8 @@ class OpAttrExtractor { bool can_be_double = true; for (size_t i = 0; i < an->size(); ++i) { const auto& elem = (*an)[i]; - bool is_int = elem.try_cast() || elem.as(); - bool is_float = elem.try_cast() || elem.as(); + bool is_int = elem.try_cast() || elem.as(); + bool is_float = elem.try_cast() || elem.as(); if (!is_int) { can_be_int = false; } @@ -127,7 +127,7 @@ class OpAttrExtractor { for (size_t i = 0; i < an->size(); ++i) { if (auto opt_int = (*an)[i].try_cast()) { attr.push_back(opt_int.value()); - } else if (const auto* im = (*an)[i].as()) { + } else if (const auto* im = (*an)[i].as()) { attr.push_back(static_cast(im->value)); } } @@ -139,11 +139,11 @@ class OpAttrExtractor { const auto& elem = (*an)[i]; if (auto opt_float = elem.try_cast()) { attr.push_back(opt_float.value()); - } else if (const auto* fm = elem.as()) { + } else if (const auto* fm = elem.as()) { attr.push_back(fm->value); } else if (auto opt_int = elem.try_cast()) { attr.push_back(static_cast(opt_int.value())); - } else if (const auto* im = elem.as()) { + } else if (const auto* im = elem.as()) { attr.push_back(static_cast(im->value)); } } @@ -162,9 +162,9 @@ class OpAttrExtractor { } } else if (*value == nullptr) { // Skip NullValue SetNodeAttr(key, ffi::String("")); - } else if (const auto* im = (*value).as()) { + } else if (const auto* im = (*value).as()) { SetNodeAttr(key, static_cast(im->value)); - } else if (const auto* fm = (*value).as()) { + } else if (const auto* fm = (*value).as()) { SetNodeAttr(key, static_cast(fm->value)); } else if (const auto opt_str = (*value).as()) { SetNodeAttr(key, *opt_str); @@ -242,7 +242,7 @@ class JSONSerializer : public relax::MemoizedExprTranslator { // First we convert all the parameters into input nodes. for (const auto& param : func->params) { auto node_ptr = std::make_shared(param->name, "input" /* op_type_ */); - memo_[param] = AddNode(node_ptr, param); + memo_[param] = AddObj(node_ptr, param); } heads_ = VisitExpr(func->body); } @@ -267,7 +267,7 @@ class JSONSerializer : public relax::MemoizedExprTranslator { * \return A list of graph entry nodes. It the relax expr is a tuple type, we * will flatten it. */ - NodeEntries AddNode(JSONGraphObjectPtr node, const Expr& expr) { + NodeEntries AddObj(JSONGraphObjectPtr node, const Expr& expr) { auto ty = GetType(expr); auto node_id = nodes_.size(); nodes_.push_back(node); @@ -276,9 +276,9 @@ class JSONSerializer : public relax::MemoizedExprTranslator { TypeVector dtype; // Flatten tuple node. - if (const auto* tuple_ty = ty.as()) { + if (const auto* tuple_ty = ty.as()) { for (size_t i = 0; i < tuple_ty->fields.size(); ++i) { - const auto* tensor_ty = tuple_ty->fields[i].as(); + const auto* tensor_ty = tuple_ty->fields[i].as(); TVM_FFI_ICHECK(tensor_ty) << "Expect TensorType, but received: ." << tuple_ty->fields[i]->GetTypeKey(); TVM_FFI_ICHECK(tensor_ty->shape.has_value()) << "Expect shape to be defined."; @@ -289,7 +289,7 @@ class JSONSerializer : public relax::MemoizedExprTranslator { } node->SetNumOutput(tuple_ty->fields.size()); } else { - const auto* tensor_ty = ty.as(); + const auto* tensor_ty = ty.as(); TVM_FFI_ICHECK(tensor_ty) << "Expect TensorType, but received: " << ty->GetTypeKey(); TVM_FFI_ICHECK(tensor_ty->shape.has_value()) << "Expect shape to be defined."; ShapeExpr output_shape = tensor_ty->shape.value().as_or_throw(); @@ -309,12 +309,12 @@ class JSONSerializer : public relax::MemoizedExprTranslator { return ret; } - void SetCallNodeAttribute(JSONGraphObjectPtr node, const CallNode* cn) { - if (cn->op.as()) { + void SetCallNodeAttribute(JSONGraphObjectPtr node, const CallObj* cn) { + if (cn->op.as()) { OpAttrExtractor extractor(node); const ffi::Object* call_attr = cn->attrs.get(); extractor.Extract(const_cast(call_attr)); - } else if (const auto* fn = cn->op.as()) { + } else if (const auto* fn = cn->op.as()) { TVM_FFI_ICHECK(false); auto pattern = fn->GetAttr(tvm::relax::attr::kPartitionedFromPattern); TVM_FFI_ICHECK(pattern.has_value()); @@ -322,17 +322,17 @@ class JSONSerializer : public relax::MemoizedExprTranslator { } } - NodeEntries VisitBinding_(const MatchCastNode* binding) { + NodeEntries VisitBinding_(const MatchCastObj* binding) { TVM_FFI_THROW(InternalError) << "JSON runtime currently doesn't match cast\n"; return {}; } NodeEntries VisitBinding(const Binding& binding) { NodeEntries nodes; - if (const auto* node = binding.as()) { + if (const auto* node = binding.as()) { auto from_b = VisitBinding_(node); nodes.insert(nodes.end(), from_b.begin(), from_b.end()); - } else if (const auto* node = binding.as()) { + } else if (const auto* node = binding.as()) { auto from_b = VisitBinding_(node); nodes.insert(nodes.end(), from_b.begin(), from_b.end()); } else { @@ -343,10 +343,10 @@ class JSONSerializer : public relax::MemoizedExprTranslator { NodeEntries VisitBindingBlock(const BindingBlock& block) { NodeEntries nodes; - if (const auto* node = block.as()) { + if (const auto* node = block.as()) { auto from_bb = VisitBindingBlock_(node); nodes.insert(nodes.end(), from_bb.begin(), from_bb.end()); - } else if (const auto* node = block.as()) { + } else if (const auto* node = block.as()) { auto from_bb = VisitBindingBlock_(node); nodes.insert(nodes.end(), from_bb.begin(), from_bb.end()); } else { @@ -355,7 +355,7 @@ class JSONSerializer : public relax::MemoizedExprTranslator { return nodes; } - NodeEntries VisitBindingBlock_(const BindingBlockNode* block) { + NodeEntries VisitBindingBlock_(const BindingBlockObj* block) { NodeEntries nodes; for (Binding binding : block->bindings) { auto from_b = VisitBinding(binding); @@ -364,7 +364,7 @@ class JSONSerializer : public relax::MemoizedExprTranslator { return nodes; } - NodeEntries VisitBindingBlock_(const DataflowBlockNode* block) { + NodeEntries VisitBindingBlock_(const DataflowBlockObj* block) { NodeEntries nodes; for (Binding binding : block->bindings) { auto from_b = VisitBinding(binding); @@ -373,7 +373,7 @@ class JSONSerializer : public relax::MemoizedExprTranslator { return nodes; } - NodeEntries VisitExpr_(const SeqExprNode* op) { + NodeEntries VisitExpr_(const SeqExprObj* op) { NodeEntries nodes; for (BindingBlock block : op->blocks) { VisitBindingBlock(block); @@ -388,16 +388,16 @@ class JSONSerializer : public relax::MemoizedExprTranslator { return {}; } - NodeEntries VisitExpr_(const GenericConstNode* cn) { + NodeEntries VisitExpr_(const GenericConstObj* cn) { auto name = constant_names_.find(ffi::GetRef(cn)); TVM_FFI_ICHECK(name != constant_names_.end()) << "Cannot find the name of the constant: " << ffi::GetRef(cn); constants_used_.push_back((*name).second); auto node = std::make_shared((*name).second, "const" /* op_type_ */); - return AddNode(node, ffi::GetRef(cn)); + return AddObj(node, ffi::GetRef(cn)); } - NodeEntries VisitExpr_(const TupleNode* tn) { + NodeEntries VisitExpr_(const TupleObj* tn) { NodeEntries fields; for (const auto& field : tn->fields) { auto ref = VisitExpr(field); @@ -406,12 +406,12 @@ class JSONSerializer : public relax::MemoizedExprTranslator { return fields; } - NodeEntries VisitExpr_(const CallNode* cn) { + NodeEntries VisitExpr_(const CallObj* cn) { Expr expr = ffi::GetRef(cn); std::string name; - if (const auto* op_node = cn->op.as()) { + if (const auto* op_node = cn->op.as()) { name = op_node->name; - } else if (const auto* fn = cn->op.as()) { + } else if (const auto* fn = cn->op.as()) { auto comp = fn->GetAttr(tvm::relax::attr::kComposite); TVM_FFI_ICHECK(comp.has_value()) << "JSON runtime only supports composite functions."; name = comp.value(); @@ -433,19 +433,19 @@ class JSONSerializer : public relax::MemoizedExprTranslator { "kernel", /* op_type_ */ inputs, 1 /* num_outputs_ */); SetCallNodeAttribute(node, cn); - return AddNode(node, ffi::GetRef(cn)); + return AddObj(node, ffi::GetRef(cn)); } - NodeEntries VisitExpr_(const TupleGetItemNode* gtn) { + NodeEntries VisitExpr_(const TupleGetItemObj* gtn) { auto vtuple = VisitExpr(gtn->tuple); return {vtuple[gtn->index]}; } - NodeEntries VisitExpr_(const FunctionNode* fn) { + NodeEntries VisitExpr_(const FunctionObj* fn) { TVM_FFI_ICHECK(fn->GetAttr(tvm::relax::attr::kComposite).has_value()) << "JSON runtime only supports composite functions"; - // FunctionNode should be handled by the caller. + // FunctionObj should be handled by the caller. return {}; } diff --git a/src/relax/backend/contrib/cublas/codegen.cc b/src/relax/backend/contrib/cublas/codegen.cc index 969ec1dc0ba4..2e1f145b881e 100644 --- a/src/relax/backend/contrib/cublas/codegen.cc +++ b/src/relax/backend/contrib/cublas/codegen.cc @@ -50,8 +50,8 @@ class CublasJSONSerializer : public JSONSerializer { using JSONSerializer::VisitExpr_; - NodeEntries VisitExpr_(const CallNode* call_node) final { - const auto* fn_var = call_node->op.as(); + NodeEntries VisitExpr_(const CallObj* call_node) final { + const auto* fn_var = call_node->op.as(); TVM_FFI_ICHECK(fn_var); const auto fn = bindings_[ffi::GetRef(fn_var)].as_or_throw(); TVM_FFI_ICHECK(fn.defined()) << "Expects the callee to be a function."; @@ -84,9 +84,9 @@ class CublasJSONSerializer : public JSONSerializer { "kernel", /* op_type_ */ inputs, 1 /* num_outputs_ */); if (composite_name.find("dequantize") != std::string::npos) { - const CallNode* dequantize_call = backend::GetOpInFunction(fn, "relax.dequantize"); - if (dequantize_call->args[1]->IsInstance()) { - const auto* const_expr = dequantize_call->args[1].as(); + const CallObj* dequantize_call = backend::GetOpInFunction(fn, "relax.dequantize"); + if (dequantize_call->args[1]->IsInstance()) { + const auto* const_expr = dequantize_call->args[1].as(); auto ty = const_expr->ty.as_or_throw(); ffi::Tensor scale = const_expr->value.cast(); TVM_FFI_ICHECK(scale.defined()); @@ -109,9 +109,9 @@ class CublasJSONSerializer : public JSONSerializer { } } - const CallNode* root_call = backend::GetOpInFunction(fn, "relax.matmul"); + const CallObj* root_call = backend::GetOpInFunction(fn, "relax.matmul"); SetCallNodeAttribute(node, root_call); - return AddNode(node, ffi::GetRef(call_node)); + return AddObj(node, ffi::GetRef(call_node)); } private: diff --git a/src/relax/backend/contrib/cudnn/codegen.cc b/src/relax/backend/contrib/cudnn/codegen.cc index 939d5e029ee1..730f162f30bd 100644 --- a/src/relax/backend/contrib/cudnn/codegen.cc +++ b/src/relax/backend/contrib/cudnn/codegen.cc @@ -48,8 +48,8 @@ class cuDNNJSONSerializer : public JSONSerializer { using JSONSerializer::VisitExpr_; - NodeEntries VisitExpr_(const CallNode* call_node) final { - const auto* fn_var = call_node->op.as(); + NodeEntries VisitExpr_(const CallObj* call_node) final { + const auto* fn_var = call_node->op.as(); TVM_FFI_ICHECK(fn_var); const auto fn = bindings_[ffi::GetRef(fn_var)].as_or_throw(); TVM_FFI_ICHECK(fn.defined()) << "Expects the callee to be a function."; @@ -68,7 +68,7 @@ class cuDNNJSONSerializer : public JSONSerializer { } } - NodeEntries HandleConv2D(const CallNode* call_node, const Function& fn, + NodeEntries HandleConv2D(const CallObj* call_node, const Function& fn, const std::string& composite_name) { NodeEntries inputs_tmp; for (const auto& arg : call_node->args) { @@ -90,12 +90,12 @@ class cuDNNJSONSerializer : public JSONSerializer { "kernel", /* op_type_ */ inputs, 1 /* num_outputs_ */); - const CallNode* root_call = backend::GetOpInFunction(fn, "relax.nn.conv2d"); + const CallObj* root_call = backend::GetOpInFunction(fn, "relax.nn.conv2d"); SetCallNodeAttribute(node, root_call); - return AddNode(node, ffi::GetRef(call_node)); + return AddObj(node, ffi::GetRef(call_node)); } - NodeEntries HandleAttention(const CallNode* call_node, const Function& fn, + NodeEntries HandleAttention(const CallObj* call_node, const Function& fn, const std::string& composite_name) { std::string layout = composite_name.substr(composite_name.find_last_of(".") + 1); NodeEntries inputs; @@ -107,17 +107,17 @@ class cuDNNJSONSerializer : public JSONSerializer { auto node = std::make_shared(composite_name, /* name_ */ "kernel", /* op_type_ */ inputs, 1 /* num_outputs_ */); - const CallNode* root_call = backend::GetOpInFunction(fn, "relax.nn.attention"); + const CallObj* root_call = backend::GetOpInFunction(fn, "relax.nn.attention"); auto q_shape = root_call->args[0]->ty.as_or_throw()->shape.value().as_or_throw(); auto k_shape = root_call->args[1]->ty.as_or_throw()->shape.value().as_or_throw(); auto v_shape = root_call->args[2]->ty.as_or_throw()->shape.value().as_or_throw(); - int num_heads = q_shape->values[2].as()->value.as().value(); - int num_kv_heads = k_shape->values[2].as()->value.as().value(); - int head_size = q_shape->values[3].as()->value.as().value(); - int head_size_v = v_shape->values[3].as()->value.as().value(); + int num_heads = q_shape->values[2].as()->value.as().value(); + int num_kv_heads = k_shape->values[2].as()->value.as().value(); + int head_size = q_shape->values[3].as()->value.as().value(); + int head_size_v = v_shape->values[3].as()->value.as().value(); SetCallNodeAttribute(node, root_call); node->SetAttr("num_heads", static_cast(num_heads)); @@ -125,7 +125,7 @@ class cuDNNJSONSerializer : public JSONSerializer { node->SetAttr("head_size", static_cast(head_size)); node->SetAttr("head_size_v", static_cast(head_size_v)); node->SetAttr("layout", ffi::String(layout)); - return AddNode(node, ffi::GetRef(call_node)); + return AddObj(node, ffi::GetRef(call_node)); } private: diff --git a/src/relax/backend/contrib/cutlass/codegen.cc b/src/relax/backend/contrib/cutlass/codegen.cc index 5539eaacdfe2..4dd29b411e5d 100644 --- a/src/relax/backend/contrib/cutlass/codegen.cc +++ b/src/relax/backend/contrib/cutlass/codegen.cc @@ -168,9 +168,9 @@ class CodegenCutlass : public relax::MemoizedExprTranslator, for (const auto& arg : ext_func_args_) { auto ty = GetType(arg); - if (const auto* tensor_ty = ty.as()) { + if (const auto* tensor_ty = ty.as()) { arg_types.emplace_back(backend::DType2String(tensor_ty->dtype.value()->dtype)); - } else if (const auto* shape_ty = ty.as()) { + } else if (const auto* shape_ty = ty.as()) { arg_types.emplace_back(backend::DType2String(shape_ty->values.value()[0].ty()->dtype)); } else { TVM_FFI_THROW(InternalError) << "Unimplemented"; @@ -202,7 +202,7 @@ class CodegenCutlass : public relax::MemoizedExprTranslator, ffi::Array GetHeaders() { return headers_; } protected: - OutputType VisitExpr_(const VarNode* node) final { + OutputType VisitExpr_(const VarObj* node) final { Output output; auto it = var_name_map_.find(node); TVM_FFI_ICHECK(it != var_name_map_.end()); @@ -210,8 +210,8 @@ class CodegenCutlass : public relax::MemoizedExprTranslator, return {output}; } - OutputType VisitExpr_(const CallNode* call) final { - const auto* fn_var = call->op.as(); + OutputType VisitExpr_(const CallObj* call) final { + const auto* fn_var = call->op.as(); TVM_FFI_ICHECK(fn_var); const auto func = bindings_[ffi::GetRef(fn_var)].as_or_throw(); const auto pattern_name_opt = func->GetAttr(tvm::relax::attr::kComposite); @@ -222,16 +222,16 @@ class CodegenCutlass : public relax::MemoizedExprTranslator, return ret.outputs; } - OutputType VisitExpr_(const FunctionNode* fn) final { + OutputType VisitExpr_(const FunctionObj* fn) final { TVM_FFI_ICHECK(fn->GetAttr(tvm::relax::attr::kComposite).has_value()) << "JSON runtime only supports composite functions"; - // FunctionNode should be handled by the caller. + // FunctionObj should be handled by the caller. return {}; } OutputType VisitBinding(const Binding& binding) { OutputType outputs; - if (const auto* node = binding.as()) { + if (const auto* node = binding.as()) { auto from_b = VisitBinding_(node); outputs.insert(outputs.end(), from_b.begin(), from_b.end()); } else { @@ -242,10 +242,10 @@ class CodegenCutlass : public relax::MemoizedExprTranslator, OutputType VisitBindingBlock(const BindingBlock& block) { OutputType outputs; - if (const auto* node = block.as()) { + if (const auto* node = block.as()) { auto from_bb = VisitBindingBlock_(node); outputs.insert(outputs.end(), from_bb.begin(), from_bb.end()); - } else if (const auto* node = block.as()) { + } else if (const auto* node = block.as()) { auto from_bb = VisitBindingBlock_(node); outputs.insert(outputs.end(), from_bb.begin(), from_bb.end()); } else { @@ -254,7 +254,7 @@ class CodegenCutlass : public relax::MemoizedExprTranslator, return outputs; } - OutputType VisitBindingBlock_(const BindingBlockNode* block) { + OutputType VisitBindingBlock_(const BindingBlockObj* block) { OutputType outputs; for (Binding binding : block->bindings) { auto from_b = VisitBinding(binding); @@ -263,7 +263,7 @@ class CodegenCutlass : public relax::MemoizedExprTranslator, return outputs; } - OutputType VisitBindingBlock_(const DataflowBlockNode* block) { + OutputType VisitBindingBlock_(const DataflowBlockObj* block) { OutputType outputs; for (Binding binding : block->bindings) { auto from_b = VisitBinding(binding); @@ -272,7 +272,7 @@ class CodegenCutlass : public relax::MemoizedExprTranslator, return outputs; } - OutputType VisitExpr_(const SeqExprNode* op) final { + OutputType VisitExpr_(const SeqExprObj* op) final { OutputType outputs; for (BindingBlock block : op->blocks) { @@ -286,7 +286,7 @@ class CodegenCutlass : public relax::MemoizedExprTranslator, } private: - ffi::Array GetArgumentNames(const CallNode* call) { + ffi::Array GetArgumentNames(const CallObj* call) { ffi::Array arg_names; for (size_t i = 0; i < call->args.size(); ++i) { auto res = VisitExpr(call->args[i]); @@ -297,13 +297,13 @@ class CodegenCutlass : public relax::MemoizedExprTranslator, return arg_names; } - GenerateBodyOutput GenerateBody(const CallNode* call, const std::string& func_name, + GenerateBodyOutput GenerateBody(const CallObj* call, const std::string& func_name, const ffi::Map& attrs) { auto func_args = GetArgumentNames(call); auto ty = GetType(ffi::GetRef(call)); std::vector out_types; - if (const auto* tensor_ty = ty.as()) { + if (const auto* tensor_ty = ty.as()) { out_types.emplace_back(backend::DType2String(tensor_ty->dtype.value()->dtype)); } else { TVM_FFI_THROW(InternalError) << "Unimplemented ty type: " << ty; @@ -334,7 +334,7 @@ class CodegenCutlass : public relax::MemoizedExprTranslator, * We use this since sometimes different parameters to the same function end up having the same * name. */ - std::unordered_map var_name_map_; + std::unordered_map var_name_map_; /*! \brief A unique name supply to generate a unique name for each parameter. */ UniqueNameSupply name_sup_; }; diff --git a/src/relax/backend/contrib/dnnl/codegen.cc b/src/relax/backend/contrib/dnnl/codegen.cc index a3fb320227ff..a6c5e199163d 100644 --- a/src/relax/backend/contrib/dnnl/codegen.cc +++ b/src/relax/backend/contrib/dnnl/codegen.cc @@ -48,8 +48,8 @@ class DNNLJSONSerializer : public JSONSerializer { using JSONSerializer::VisitExpr_; - NodeEntries VisitExpr_(const CallNode* call_node) final { - const auto* fn_var = call_node->op.as(); + NodeEntries VisitExpr_(const CallObj* call_node) final { + const auto* fn_var = call_node->op.as(); TVM_FFI_ICHECK(fn_var); const auto fn = bindings_[ffi::GetRef(fn_var)].as_or_throw(); TVM_FFI_ICHECK(fn.defined()) << "Expects the callee to be a function."; @@ -68,7 +68,7 @@ class DNNLJSONSerializer : public JSONSerializer { "kernel", /* op_type_ */ inputs, 1 /* num_outputs_ */); - const CallNode* root_call = nullptr; + const CallObj* root_call = nullptr; if (composite_name.find("conv2d") != std::string::npos) { root_call = backend::GetOpInFunction(fn, "relax.nn.conv2d"); } else { @@ -76,7 +76,7 @@ class DNNLJSONSerializer : public JSONSerializer { } SetCallNodeAttribute(node, root_call); - return AddNode(node, ffi::GetRef(call_node)); + return AddObj(node, ffi::GetRef(call_node)); } private: diff --git a/src/relax/backend/contrib/example_npu/codegen.cc b/src/relax/backend/contrib/example_npu/codegen.cc index e731e3640ac4..99eda0f6080d 100644 --- a/src/relax/backend/contrib/example_npu/codegen.cc +++ b/src/relax/backend/contrib/example_npu/codegen.cc @@ -48,8 +48,8 @@ class ExampleNPUJSONSerializer : public JSONSerializer { using JSONSerializer::VisitExpr_; - NodeEntries VisitExpr_(const CallNode* call_node) final { - const auto* fn_var = call_node->op.as(); + NodeEntries VisitExpr_(const CallObj* call_node) final { + const auto* fn_var = call_node->op.as(); TVM_FFI_ICHECK(fn_var); const auto fn = bindings_[ffi::GetRef(fn_var)].as_or_throw(); TVM_FFI_ICHECK(fn.defined()) << "Expects the callee to be a function."; @@ -65,7 +65,7 @@ class ExampleNPUJSONSerializer : public JSONSerializer { inputs.insert(inputs.end(), res.begin(), res.end()); } auto node = std::make_shared(composite_name, "kernel", inputs, 1); - return AddNode(node, ffi::GetRef(call_node)); + return AddObj(node, ffi::GetRef(call_node)); } private: diff --git a/src/relax/backend/contrib/hipblas/codegen.cc b/src/relax/backend/contrib/hipblas/codegen.cc index 905a9fb37f05..7ff6b7302413 100644 --- a/src/relax/backend/contrib/hipblas/codegen.cc +++ b/src/relax/backend/contrib/hipblas/codegen.cc @@ -48,8 +48,8 @@ class HipblasJSONSerializer : public JSONSerializer { using JSONSerializer::VisitExpr_; - NodeEntries VisitExpr_(const CallNode* call_node) final { - const auto* fn_var = call_node->op.as(); + NodeEntries VisitExpr_(const CallObj* call_node) final { + const auto* fn_var = call_node->op.as(); TVM_FFI_ICHECK(fn_var); const auto fn = bindings_[ffi::GetRef(fn_var)].as_or_throw(); TVM_FFI_ICHECK(fn.defined()) << "Expects the callee to be a function."; @@ -79,9 +79,9 @@ class HipblasJSONSerializer : public JSONSerializer { "kernel", /* op_type_ */ inputs, 1 /* num_outputs_ */); - const CallNode* root_call = backend::GetOpInFunction(fn, "relax.matmul"); + const CallObj* root_call = backend::GetOpInFunction(fn, "relax.matmul"); SetCallNodeAttribute(node, root_call); - return AddNode(node, ffi::GetRef(call_node)); + return AddObj(node, ffi::GetRef(call_node)); } private: diff --git a/src/relax/backend/contrib/nnapi/codegen.cc b/src/relax/backend/contrib/nnapi/codegen.cc index 7cc485926c10..f285d9d336d2 100644 --- a/src/relax/backend/contrib/nnapi/codegen.cc +++ b/src/relax/backend/contrib/nnapi/codegen.cc @@ -51,9 +51,9 @@ class CollectFromCompositeFunctionBody : public ExprVisitor { explicit CollectFromCompositeFunctionBody(NNAPIJSONSerializer* serializer) : serializer_(serializer), node_(std::make_shared()) {} - void VisitExpr_(const CallNode* call_node) override; + void VisitExpr_(const CallObj* call_node) override; - void SetPermuteDimsAttribute(const CallNode* call_node) { + void SetPermuteDimsAttribute(const CallObj* call_node) { const auto* permute_dims_attr = call_node->attrs.as(); TVM_FFI_ICHECK(permute_dims_attr); if (permute_dims_attr->axes) { @@ -65,13 +65,13 @@ class CollectFromCompositeFunctionBody : public ExprVisitor { } } - void SetAstypeAttribute(const CallNode* call_node) { + void SetAstypeAttribute(const CallObj* call_node) { const auto* astype_attrs = call_node->attrs.as(); TVM_FFI_ICHECK(astype_attrs); node_->SetAttr("astype_dtype", ffi::String(ffi::DLDataTypeToString(astype_attrs->dtype))); } - void SetMeanAttribute(const CallNode* call_node) { + void SetMeanAttribute(const CallObj* call_node) { const auto* mean_attrs = call_node->attrs.as(); TVM_FFI_ICHECK(mean_attrs); TVM_FFI_ICHECK(mean_attrs->axis.has_value()); @@ -86,7 +86,7 @@ class CollectFromCompositeFunctionBody : public ExprVisitor { node_->SetAttr("keepdims", static_cast(mean_attrs->keepdims ? 1 : 0)); } - void SetConv2dAttribute(const CallNode* call_node) { + void SetConv2dAttribute(const CallObj* call_node) { const auto* conv2d_attr = call_node->attrs.as(); TVM_FFI_ICHECK(conv2d_attr) << "didn't catch attributes"; @@ -110,7 +110,7 @@ class CollectFromCompositeFunctionBody : public ExprVisitor { node_->SetAttr("group", static_cast(conv2d_attr->groups)); } - void SetMaxPool2dAttribute(const CallNode* call_node) { + void SetMaxPool2dAttribute(const CallObj* call_node) { const auto* max_pool_2d_attr = call_node->attrs.as(); TVM_FFI_ICHECK(max_pool_2d_attr) << "didn't catch attributes"; @@ -150,8 +150,8 @@ class NNAPIJSONSerializer : public JSONSerializer { : JSONSerializer(constant_names), bindings_(bindings) {} using JSONSerializer::VisitExpr_; - std::vector VisitExpr_(const CallNode* call_node) final { - const auto* fn_var = call_node->op.as(); + std::vector VisitExpr_(const CallObj* call_node) final { + const auto* fn_var = call_node->op.as(); TVM_FFI_ICHECK(fn_var); const auto fn = bindings_[ffi::GetRef(fn_var)].as_or_throw(); TVM_FFI_ICHECK(fn.defined()) << "Expects the callee to be a function."; @@ -177,15 +177,15 @@ class NNAPIJSONSerializer : public JSONSerializer { VLOG(1) << "Adding node " << composite_name << " with " << node->GetInputs().size() << " inputs"; - return AddNode(node, ffi::GetRef(call_node)); + return AddObj(node, ffi::GetRef(call_node)); } private: ffi::Map bindings_; }; -void CollectFromCompositeFunctionBody::VisitExpr_(const CallNode* call_node) { - const auto* op_node = call_node->op.as(); +void CollectFromCompositeFunctionBody::VisitExpr_(const CallObj* call_node) { + const auto* op_node = call_node->op.as(); TVM_FFI_ICHECK(op_node != nullptr); std::string name = op_node->name; if (name == "relax.permute_dims") { diff --git a/src/relax/backend/contrib/tensorrt/codegen.cc b/src/relax/backend/contrib/tensorrt/codegen.cc index 2221678b6429..92b5207f201a 100644 --- a/src/relax/backend/contrib/tensorrt/codegen.cc +++ b/src/relax/backend/contrib/tensorrt/codegen.cc @@ -115,9 +115,9 @@ class CollectFromCompositeFunctionBody : public ExprVisitor { public: CollectFromCompositeFunctionBody() : node_(std::make_shared()) {} - void VisitExpr_(const CallNode* call_node) final; + void VisitExpr_(const CallObj* call_node) final; - void SetGenericAttributes(const CallNode* call_node) { + void SetGenericAttributes(const CallObj* call_node) { OpAttrExtractor extractor(node_); const ffi::Object* attr_obj = call_node->attrs.get(); extractor.Extract(const_cast(attr_obj)); @@ -125,8 +125,8 @@ class CollectFromCompositeFunctionBody : public ExprVisitor { // Serialize an op's non-tensor arguments (scalars/shapes) as "arg_" attributes; the "arg_" // prefix avoids JSONGraphNode's reserved "shape"/"dtype". - void SetArgumentAttributes(const CallNode* call_node) { - const auto* op_node = call_node->op.as(); + void SetArgumentAttributes(const CallObj* call_node) { + const auto* op_node = call_node->op.as(); if (op_node == nullptr) return; const ffi::Array& arg_infos = op_node->args_info; for (size_t i = 0; i < call_node->args.size() && i < arg_infos.size(); ++i) { @@ -134,12 +134,12 @@ class CollectFromCompositeFunctionBody : public ExprVisitor { const std::string key = "arg_" + std::string(arg_infos[i]->name); if (auto prim_value = arg.as()) { PrimExpr value = prim_value.value(); - if (const auto* imm = value.as()) { + if (const auto* imm = value.as()) { node_->SetAttr(key, static_cast(imm->value)); - } else if (const auto* fimm = value.as()) { + } else if (const auto* fimm = value.as()) { node_->SetAttr(key, static_cast(fimm->value)); } - } else if (const auto* shape_expr = arg.as()) { + } else if (const auto* shape_expr = arg.as()) { SetIntArrayAttr(key, shape_expr->values); } } @@ -147,12 +147,12 @@ class CollectFromCompositeFunctionBody : public ExprVisitor { // Relax reduce axis is optional; materialize the all-axes default (it otherwise serializes as // ""). - void MaybeFillReduceAxes(const CallNode* call_node) { + void MaybeFillReduceAxes(const CallObj* call_node) { const auto* attrs = call_node->attrs.as(); if (attrs == nullptr || attrs->axis.has_value()) return; - const auto* tensor_ty = GetType(call_node->args[0]).as(); + const auto* tensor_ty = GetType(call_node->args[0]).as(); if (tensor_ty == nullptr || !tensor_ty->shape.has_value()) return; - const auto* shape = tensor_ty->shape.value().as(); + const auto* shape = tensor_ty->shape.value().as(); if (shape == nullptr) return; ffi::Array all_axes; for (size_t i = 0; i < shape->values.size(); ++i) all_axes.push_back(static_cast(i)); @@ -161,12 +161,12 @@ class CollectFromCompositeFunctionBody : public ExprVisitor { // strided_slice's axes/begin/end/strides are tuple args the op does not name; serialize by // position. - void SetStridedSliceArguments(const CallNode* call_node) { - const auto* op_node = call_node->op.as(); + void SetStridedSliceArguments(const CallObj* call_node) { + const auto* op_node = call_node->op.as(); if (op_node == nullptr || op_node->name != "relax.strided_slice") return; static constexpr const char* kNames[] = {"arg_axes", "arg_begin", "arg_end", "arg_strides"}; for (size_t i = 1; i < call_node->args.size() && i <= 4; ++i) { - const auto* tuple = call_node->args[i].as(); + const auto* tuple = call_node->args[i].as(); if (tuple == nullptr) continue; ffi::Array values; for (const Expr& field : tuple->fields) { @@ -182,7 +182,7 @@ class CollectFromCompositeFunctionBody : public ExprVisitor { void SetIntArrayAttr(const std::string& key, const ffi::Array& exprs) { ffi::Array values; for (const PrimExpr& expr : exprs) { - const auto* imm = expr.as(); + const auto* imm = expr.as(); if (imm == nullptr) return; values.push_back(static_cast(imm->value)); } @@ -191,8 +191,8 @@ class CollectFromCompositeFunctionBody : public ExprVisitor { // layout_transform's IndexMap is not generically serializable; emit a pure permutation as // "arg_axes". Returns true for layout_transform (so generic extraction is skipped for it). - bool TrySetLayoutTransformAttributes(const CallNode* call_node) { - const auto* op_node = call_node->op.as(); + bool TrySetLayoutTransformAttributes(const CallObj* call_node) { + const auto* op_node = call_node->op.as(); if (op_node == nullptr || op_node->name != "relax.layout_transform") return false; const auto* attrs = call_node->attrs.as(); if (attrs == nullptr) return true; @@ -219,7 +219,7 @@ class CollectFromCompositeFunctionBody : public ExprVisitor { } /*! \brief The primitive operator call in the composite function. */ - const CallNode* operator_call_{nullptr}; + const CallObj* operator_call_{nullptr}; /*! * \brief Temporary node into which we'll accumulate attributes. Ideally this would be the * final JSONGraphNode however we don't yet know how many inputs that will have. @@ -240,9 +240,9 @@ class TensorRTJSONSerializer : public JSONSerializer { using JSONSerializer::VisitExpr_; - std::vector VisitExpr_(const CallNode* call_node) final { + std::vector VisitExpr_(const CallObj* call_node) final { // The call must be to an inline "Composite" function - const auto* fn_var = call_node->op.as(); + const auto* fn_var = call_node->op.as(); TVM_FFI_ICHECK(fn_var); const auto fn = bindings_[ffi::GetRef(fn_var)].as_or_throw(); @@ -269,12 +269,12 @@ class TensorRTJSONSerializer : public JSONSerializer { std::vector inputs; auto append_tensor_inputs = [&](const auto& self, const Expr& expr) -> void { - if (const auto* tuple = expr.as()) { + if (const auto* tuple = expr.as()) { for (const Expr& field : tuple->fields) self(self, field); return; } Type type = GetType(expr); - if (type->IsInstance() || type->IsInstance()) { + if (type->IsInstance() || type->IsInstance()) { auto entries = VisitExpr(expr); inputs.insert(inputs.end(), entries.begin(), entries.end()); } @@ -296,7 +296,7 @@ class TensorRTJSONSerializer : public JSONSerializer { VLOG(1) << name << " has " << node->GetInputs().size() << " inputs"; - return AddNode(node, ffi::GetRef(call_node)); + return AddObj(node, ffi::GetRef(call_node)); } static void SaveGlobalAttributes(std::shared_ptr node) { @@ -317,8 +317,8 @@ class TensorRTJSONSerializer : public JSONSerializer { ffi::Map bindings_; }; -void CollectFromCompositeFunctionBody::VisitExpr_(const CallNode* call_node) { - TVM_FFI_ICHECK(call_node->op->IsInstance()) +void CollectFromCompositeFunctionBody::VisitExpr_(const CallObj* call_node) { + TVM_FFI_ICHECK(call_node->op->IsInstance()) << "TensorRT Composite functions must contain exactly one primitive Relax operator call"; TVM_FFI_ICHECK(operator_call_ == nullptr) << "TensorRT Composite functions must contain exactly one primitive Relax operator call"; @@ -356,18 +356,18 @@ ffi::Array TensorRTCompiler(ffi::Array functions, if (build_at_compile_time) { // Reject dynamic inputs before the JSON serializer requires integer shapes. auto check_input_type = [&](const auto& self, const Type& type) -> void { - if (const auto* tuple = type.as()) { + if (const auto* tuple = type.as()) { for (const auto& field : tuple->fields) self(self, field); return; } - const auto* tensor = type.as(); + const auto* tensor = type.as(); TVM_FFI_CHECK(tensor != nullptr && tensor->shape.has_value(), ValueError) << "TensorRT compile-time engine building requires static positive input dimensions"; - const auto* shape = tensor->shape.value().as(); + const auto* shape = tensor->shape.value().as(); TVM_FFI_CHECK(shape != nullptr, ValueError) << "TensorRT compile-time engine building requires static positive input dimensions"; for (const auto& dim : shape->values) { - const auto* value = dim.as(); + const auto* value = dim.as(); TVM_FFI_CHECK(value != nullptr && value->value > 0, ValueError) << "TensorRT compile-time engine building requires static positive input dimensions"; } diff --git a/src/relax/backend/contrib/utils.cc b/src/relax/backend/contrib/utils.cc index e49fac44511e..1366abc5d6a3 100644 --- a/src/relax/backend/contrib/utils.cc +++ b/src/relax/backend/contrib/utils.cc @@ -55,7 +55,7 @@ ffi::Map ExtractArgIdx(ffi::String pattern_name, Function f for (const auto& [name, pat] : pattern.value()->annotation_patterns) { auto exp = matched_expr.value()[pat]; - if (auto arg_var = exp.as()) { + if (auto arg_var = exp.as()) { if (auto idx = find_index(f->params, ffi::GetRef(arg_var))) { arg_idx.Set(name, IntImm::Int64(*idx)); } diff --git a/src/relax/backend/contrib/utils.h b/src/relax/backend/contrib/utils.h index e0b5ff706e98..604d9abebef5 100644 --- a/src/relax/backend/contrib/utils.h +++ b/src/relax/backend/contrib/utils.h @@ -48,7 +48,7 @@ namespace backend { inline std::vector GetIntShape(const ffi::Array& shape) { std::vector ret; for (const auto& dim : shape) { - const auto* imm = dim.as(); + const auto* imm = dim.as(); auto pval = imm ? imm->value.as() : std::nullopt; ret.push_back(pval.has_value() ? *pval : -1); } @@ -69,8 +69,8 @@ inline std::string DType2String(DLDataType dtype) { return tvm::ffi::DLDataTypeT * \param op_name The name of the op * \return true if the callee op matches with the op name */ -inline bool IsOp(const CallNode* call, const std::string& op_name) { - const auto* op_node = call->op.as(); +inline bool IsOp(const CallObj* call, const std::string& op_name) { + const auto* op_node = call->op.as(); if (!op_node) return false; Op op = ffi::GetRef(op_node); return op.same_as(Op::Get(op_name)); @@ -83,10 +83,10 @@ inline bool IsOp(const CallNode* call, const std::string& op_name) { * \param op_name The name of the op * \return A call node which calls an op with the given name or nullptr if not */ -inline const CallNode* TryGetOpInFunction(Function f, const std::string& op_name) { +inline const CallObj* TryGetOpInFunction(Function f, const std::string& op_name) { auto local_bindings = AnalyzeVar2Value(f); for (const auto& entry : local_bindings) { - if (auto call = entry.second.as(); call && backend::IsOp(call, op_name)) { + if (auto call = entry.second.as(); call && backend::IsOp(call, op_name)) { return call; } } @@ -100,8 +100,8 @@ inline const CallNode* TryGetOpInFunction(Function f, const std::string& op_name * \param op_name The name of the op * \return A call node which calls an op with the given name */ -inline const CallNode* GetOpInFunction(Function f, const std::string& op_name) { - const CallNode* op = TryGetOpInFunction(f, op_name); +inline const CallObj* GetOpInFunction(Function f, const std::string& op_name) { + const CallObj* op = TryGetOpInFunction(f, op_name); TVM_FFI_ICHECK(op) << op_name << " not found in the function:\n" << f; return op; } diff --git a/src/relax/backend/task_extraction.cc b/src/relax/backend/task_extraction.cc index 4a08ddccb77e..72c14985a197 100644 --- a/src/relax/backend/task_extraction.cc +++ b/src/relax/backend/task_extraction.cc @@ -54,7 +54,7 @@ using s_tir::meta_schedule::ModuleHash; class BlockCounter : public s_tir::StmtExprVisitor { public: ffi::Optional Visit(ffi::AnyView value) override { - if (value.as()) return std::nullopt; + if (value.as()) return std::nullopt; return s_tir::StmtExprVisitor::Visit(value); } @@ -65,7 +65,7 @@ class BlockCounter : public s_tir::StmtExprVisitor { } private: - ffi::Optional Visit_(const s_tir::SBlockNode* op) final { + ffi::Optional Visit_(const s_tir::SBlockObj* op) final { ++count; return s_tir::StmtExprVisitor::Visit_(op); } @@ -79,7 +79,7 @@ class TaskExtractor : public ExprVisitor { TaskExtractor extractor(mod, target, mod_eq_name); // We go through each Relax function in the module. for (const auto& kv : mod->functions) { - if (const auto* func = kv.second.as()) { + if (const auto* func = kv.second.as()) { extractor(ffi::GetRef(func)); } } @@ -100,7 +100,7 @@ class TaskExtractor : public ExprVisitor { TVM_FFI_ICHECK(normalize_mod_func_.has_value()) << "Normalization function is not found."; } - void VisitExpr_(const CallNode* call) final { + void VisitExpr_(const CallObj* call) final { static const Op call_tir_op = Op::Get("relax.call_tir"); // TODO(@tvm-team): When we differentiate the call for tirx function and packed function, diff --git a/src/relax/backend/vm/codegen_vm.cc b/src/relax/backend/vm/codegen_vm.cc index 1367b6aed37e..59e1b04a7a3e 100644 --- a/src/relax/backend/vm/codegen_vm.cc +++ b/src/relax/backend/vm/codegen_vm.cc @@ -62,7 +62,7 @@ class CodeGenVM : public ExprFunctor { CodeGenVM codegen(builder, mod); // Remove relax function and turn into TIR func. for (const auto& [gvar, f] : mod->functions) { - if (auto* func = f.as()) { + if (auto* func = f.as()) { codegen.Codegen(ffi::GetRef(func)); res_mod->Remove(gvar); } @@ -110,13 +110,13 @@ class CodeGenVM : public ExprFunctor { var_arg_map_.clear(); } - Instruction::Arg VisitExpr_(const SeqExprNode* op) final { + Instruction::Arg VisitExpr_(const SeqExprObj* op) final { for (auto block : op->blocks) { for (Binding binding : block->bindings) { Expr expr = GetBoundValue(binding); Instruction::Arg value = VisitExpr(expr); - if (expr.as()) { + if (expr.as()) { // For a normalized relax module, there should be one // register for each relax::Binding. This makes the Relax // semantics of R.vm.kill_* operate the same as the Python @@ -135,7 +135,7 @@ class CodeGenVM : public ExprFunctor { return ret_reg; } - Instruction::Arg VisitExpr_(const CallNode* call_node) final { + Instruction::Arg VisitExpr_(const CallObj* call_node) final { Call call = ffi::GetRef(call_node); // ComputePrimValue may introduce pure native calls after purity removal. @@ -154,7 +154,7 @@ class CodeGenVM : public ExprFunctor { // allocate dst register. RegName dst_reg = HasVoidType(call) ? Instruction::kVoidRegister : NewRegister(); - if (call->op.as()) { + if (call->op.as()) { if (call_node->op.same_as(call_builtin_with_ctx_op_)) { // TODO(relax-team) migrate most handling of op to // directly map to call_builtin_with_ctx before codegen and simplify vm codegen. @@ -183,7 +183,7 @@ class CodeGenVM : public ExprFunctor { return Instruction::Arg::Register(dst_reg); } - Instruction::Arg VisitExpr_(const IfExprNode* op) final { + Instruction::Arg VisitExpr_(const IfExprObj* op) final { const IfExpr& ife = ffi::GetRef(op); Instruction::Arg cond_value = this->VisitExpr(ife->cond); @@ -227,21 +227,21 @@ class CodeGenVM : public ExprFunctor { return Instruction::Arg::Register(merge_register); } - Instruction::Arg VisitExpr_(const VarNode* op) final { + Instruction::Arg VisitExpr_(const VarObj* op) final { Var var = ffi::GetRef(op); auto it = this->var_arg_map_.find(var); TVM_FFI_ICHECK(it != this->var_arg_map_.end()) << "Var " << var << " is not defined"; return it->second; } - Instruction::Arg VisitExpr_(const DataflowVarNode* op) final { - return VisitExpr_(static_cast(op)); + Instruction::Arg VisitExpr_(const DataflowVarObj* op) final { + return VisitExpr_(static_cast(op)); } - Instruction::Arg VisitExpr_(const GenericConstNode* op) final { + Instruction::Arg VisitExpr_(const GenericConstObj* op) final { auto arg = builder_->ConvertConstant(op->value); - if (auto tensor_ty = op->ty.as()) { + if (auto tensor_ty = op->ty.as()) { if (tensor_ty->vdevice.has_value()) { VDevice vdev = tensor_ty->vdevice.value(); builder_->SaveMemoryScope(arg, vdev->memory_scope); @@ -250,10 +250,10 @@ class CodeGenVM : public ExprFunctor { return arg; } - Instruction::Arg VisitExpr_(const ShapeExprNode* op) final { + Instruction::Arg VisitExpr_(const ShapeExprObj* op) final { std::vector shape; for (PrimExpr e : op->values) { - if (auto* int_value = e.as()) { + if (auto* int_value = e.as()) { shape.push_back(static_cast(int_value->value)); } else { TVM_FFI_THROW(InternalError) @@ -263,22 +263,22 @@ class CodeGenVM : public ExprFunctor { return builder_->ConvertConstant(ffi::Shape(shape)); } - Instruction::Arg VisitExpr_(const IntImmNode* op) final { + Instruction::Arg VisitExpr_(const IntImmObj* op) final { return builder_->ConvertConstant(op->value); } - Instruction::Arg VisitExpr_(const FloatImmNode* op) final { + Instruction::Arg VisitExpr_(const FloatImmObj* op) final { return builder_->ConvertConstant(op->value); } - Instruction::Arg VisitExpr_(const StringImmNode* op) final { + Instruction::Arg VisitExpr_(const StringImmObj* op) final { return builder_->ConvertConstant(op->value); } - Instruction::Arg VisitExpr_(const DataTypeImmNode* op) final { + Instruction::Arg VisitExpr_(const DataTypeImmObj* op) final { return builder_->ConvertConstant(op->value); } - Instruction::Arg VisitExpr_(const TupleNode* op) final { + Instruction::Arg VisitExpr_(const TupleObj* op) final { Tuple tuple = ffi::GetRef(op); std::vector args; for (Expr arg : tuple->fields) { @@ -290,7 +290,7 @@ class CodeGenVM : public ExprFunctor { return Instruction::Arg::Register(dst_register); } - Instruction::Arg VisitExpr_(const TupleGetItemNode* op) final { + Instruction::Arg VisitExpr_(const TupleGetItemObj* op) final { TupleGetItem expr = ffi::GetRef(op); std::vector args = {this->VisitExpr(expr->tuple)}; @@ -302,7 +302,7 @@ class CodeGenVM : public ExprFunctor { return Instruction::Arg::Register(dst_register); } - Instruction::Arg VisitExpr_(const GlobalVarNode* op) final { + Instruction::Arg VisitExpr_(const GlobalVarObj* op) final { GlobalVar gvar = ffi::GetRef(op); ffi::Optional symbol; VMFuncInfo::FuncKind kind = VMFuncInfo::FuncKind::kPackedFunc; @@ -311,10 +311,10 @@ class CodeGenVM : public ExprFunctor { auto it = ctx_mod_->functions.find(gvar); if (it != ctx_mod_->functions.end()) { BaseFunc func = (*it).second; - if (auto* efunc = func.as()) { + if (auto* efunc = func.as()) { symbol = efunc->global_symbol; kind = VMFuncInfo::FuncKind::kPackedFunc; - } else if (func.as()) { + } else if (func.as()) { symbol = gvar->name_hint; kind = VMFuncInfo::FuncKind::kVMFunc; } @@ -333,7 +333,7 @@ class CodeGenVM : public ExprFunctor { return builder_->GetFunction(symbol.value()); } - Instruction::Arg VisitExpr_(const ExternFuncNode* op) final { + Instruction::Arg VisitExpr_(const ExternFuncObj* op) final { if (ffi::Optional opt_code = op->attrs.GetAttr(tvm::relax::attr::kCSource)) { ffi::String sym = op->global_symbol; ffi::String fmt = op->attrs.GetAttr(tvm::relax::attr::kCSourceFmt).value_or("c"); @@ -370,7 +370,7 @@ class CodeGenVM : public ExprFunctor { args.push_back(this->VisitExpr(call_node->args[i])); } int64_t vdevice_index = -1; - if (const auto* int_imm = call_node->args[4].as()) { + if (const auto* int_imm = call_node->args[4].as()) { vdevice_index = int_imm->value.as().value(); } auto vdevice = GetGlobalVDevice(ctx_mod_, vdevice_index); diff --git a/src/relax/backend/vm/lower_runtime_builtin.cc b/src/relax/backend/vm/lower_runtime_builtin.cc index bd40b086ed19..75d42117b681 100644 --- a/src/relax/backend/vm/lower_runtime_builtin.cc +++ b/src/relax/backend/vm/lower_runtime_builtin.cc @@ -42,7 +42,7 @@ class LowerRuntimeBuiltinMutator : public ExprMutator { public: using ExprMutator::VisitExpr_; - Expr VisitExpr_(const CallNode* call_node) final { + Expr VisitExpr_(const CallObj* call_node) final { static const auto& lower_builtin_fmap = Op::GetAttrMap(tvm::relax::op_attr::kLowerBuiltin); // post-order mutation Call call = VisitExprPostOrder_(call_node).as_or_throw(); @@ -74,7 +74,7 @@ class LowerRuntimeBuiltinMutator : public ExprMutator { return MakeMemAllocTensor(call); } else if (call->op.same_as(mem_kill_storage_op_) || call->op.same_as(mem_kill_tensor_op_)) { return MakeMemKillObject(call); - } else if (const auto* op_node = call->op.as()) { + } else if (const auto* op_node = call->op.as()) { Op op = ffi::GetRef(op_node); if (lower_builtin_fmap.count(op)) { return lower_builtin_fmap[op](builder_, call); @@ -110,8 +110,8 @@ class LowerRuntimeBuiltinMutator : public ExprMutator { Expr CallTIRDyn(const Call& call_node) { TVM_FFI_ICHECK(call_node->args.size() == 2); - TVM_FFI_ICHECK(call_node->args[0]->IsInstance()); - TVM_FFI_ICHECK(call_node->args[1]->IsInstance()); + TVM_FFI_ICHECK(call_node->args[0]->IsInstance()); + TVM_FFI_ICHECK(call_node->args[1]->IsInstance()); ffi::Array args; auto tir_args = call_node->args[1].as_or_throw(); @@ -127,9 +127,9 @@ class LowerRuntimeBuiltinMutator : public ExprMutator { TVM_FFI_ICHECK(!call_node->ty.as().has_value()); auto arg = call_node->args[1]; - TVM_FFI_CHECK(arg->ty->IsInstance(), TypeError) + TVM_FFI_CHECK(arg->ty->IsInstance(), TypeError) << "LowerRuntimeBuiltin expects the shape arg of R.reshape " - << "to be a ShapeExpr or VarNode bound to a ShapeExpr. " + << "to be a ShapeExpr or VarObj bound to a ShapeExpr. " << "However, in expression " << call_node << ", the shape argument " << arg << " has type " << arg->ty; @@ -187,8 +187,8 @@ class LowerRuntimeBuiltinMutator : public ExprMutator { Expr MakeClosure(const Call& call_node) { TVM_FFI_ICHECK(call_node->args.size() == 2); - TVM_FFI_ICHECK(call_node->args[0]->IsInstance()); - TVM_FFI_ICHECK(call_node->args[1]->IsInstance()); + TVM_FFI_ICHECK(call_node->args[0]->IsInstance()); + TVM_FFI_ICHECK(call_node->args[1]->IsInstance()); ffi::Array args; auto func = call_node->args[0]; @@ -204,8 +204,8 @@ class LowerRuntimeBuiltinMutator : public ExprMutator { Expr InvokeClosure(const Call& call_node) { TVM_FFI_ICHECK(call_node->args.size() == 2); - TVM_FFI_ICHECK(call_node->args[0]->IsInstance()); - TVM_FFI_ICHECK(call_node->args[1]->IsInstance()); + TVM_FFI_ICHECK(call_node->args[0]->IsInstance()); + TVM_FFI_ICHECK(call_node->args[1]->IsInstance()); ffi::Array args; diff --git a/src/relax/backend/vm/vm_shape_lower.cc b/src/relax/backend/vm/vm_shape_lower.cc index 977c3739d185..86ef0b4050d6 100644 --- a/src/relax/backend/vm/vm_shape_lower.cc +++ b/src/relax/backend/vm/vm_shape_lower.cc @@ -85,7 +85,7 @@ using PrimExprSlotMap = using LiveVarSet = std::unordered_set; -static bool IsRelaxOwnedCall(const CallNode* call) { +static bool IsRelaxOwnedCall(const CallObj* call) { auto op = call->op.as(); if (!op) return true; static auto infer_type_map = Op::GetAttrMap(tvm::op_attr::kInferType); @@ -120,7 +120,7 @@ class PrimExprSlotCollector : public ExprVisitor, public TypeVisitor { } private: - void VisitBinding_(const VarBindingNode* op) final { + void VisitBinding_(const VarBindingObj* op) final { if (!collect_slots_) { VisitVarDef(op->var); return; @@ -139,13 +139,13 @@ class PrimExprSlotCollector : public ExprVisitor, public TypeVisitor { void MarkLive(const Var& var) { if (!live_vars_.insert(var).second) return; if (auto it = bound_values_->find(var); it != bound_values_->end()) { - for (const VarNode* used : GetUsedVars((*it).second)) MarkLive(ffi::GetRef(used)); + for (const VarObj* used : GetUsedVars((*it).second)) MarkLive(ffi::GetRef(used)); } } - void VisitExpr_(const VarNode* op) final { + void VisitExpr_(const VarObj* op) final { Var var = ffi::GetRef(op); - if (collect_scalar_ && !var.as()) { + if (collect_scalar_ && !var.as()) { if (auto prim_var = var.as(); prim_var && prim_var.value().ty()->dtype == DLDataType{kDLInt, 64, 1}) { HandlePrimExpr(prim_var.value()); @@ -154,7 +154,7 @@ class PrimExprSlotCollector : public ExprVisitor, public TypeVisitor { ExprVisitor::VisitExpr_(op); } - void VisitExpr_(const ShapeExprNode* op) final { + void VisitExpr_(const ShapeExprObj* op) final { for (const PrimExpr& value : op->values) HandlePrimExpr(value); } @@ -163,8 +163,8 @@ class PrimExprSlotCollector : public ExprVisitor, public TypeVisitor { } void CollectPrimExprSlot(const PrimExpr& expr) { - if (expr->IsInstance()) return; - if (const auto* call = expr.as()) { + if (expr->IsInstance()) return; + if (const auto* call = expr.as()) { TVM_FFI_CHECK(!IsRelaxOwnedCall(call), ValueError) << "VM shape expressions cannot compile a Relax-owned Call: " << ffi::GetRef(call); } @@ -180,17 +180,17 @@ class PrimExprSlotCollector : public ExprVisitor, public TypeVisitor { } } - void VisitBinding_(const MatchCastNode* op) final { + void VisitBinding_(const MatchCastObj* op) final { // Visit the match cast type so we can define // the symbolic variables here. this->VisitType(op->ty); } - void VisitExpr_(const FunctionNode* op) final { + void VisitExpr_(const FunctionObj* op) final { // Do not recurse into function node as it is self-contained } - void VisitType_(const FuncTypeNode* op) final { + void VisitType_(const FuncTypeObj* op) final { // Do not recurse into function type as it is self-contained } @@ -259,7 +259,7 @@ class PrimExprSlotCollector : public ExprVisitor, public TypeVisitor { * - CheckMatchCast: recursively structinfo unpacking, generate checks and match items. * - RunMatch: generate symbolic shape matches * - EmitOutstandingPrimExprCompute: tracks the variables to be computed and emit shape computation - * - VisitExpr_(ShapeExprNode*): makes symbolic shape tuple. + * - VisitExpr_(ShapeExprObj*): makes symbolic shape tuple. * * The checks and symbolic shape all maps to runtime builtin functions. Please checkout * runtime/vm/builtin.cc for their definitions. @@ -292,7 +292,7 @@ class VMShapeLowerMutator VMShapeLowerMutator mutator(mod, emit_err_ctx); for (auto& kv : mod->functions) { - if (auto* func = kv.second.as()) { + if (auto* func = kv.second.as()) { Function updated_func = mutator.Rewrite(kv.first, ffi::GetRef(func)); mutator.builder_->UpdateFunction(kv.first, updated_func); } @@ -306,9 +306,9 @@ class VMShapeLowerMutator using ExprMutator::VisitExpr_; - Expr VisitExpr_(const VarNode* op) final { + Expr VisitExpr_(const VarObj* op) final { Var var = ffi::GetRef(op); - if (!var.as()) { + if (!var.as()) { if (auto prim_var = var.as(); prim_var && slot_map_.count(*prim_var)) { return RewritePrimValue(*prim_var); } @@ -316,7 +316,7 @@ class VMShapeLowerMutator return ExprMutator::VisitExpr_(op); } - Expr VisitExpr_(const CallNode* op) final { + Expr VisitExpr_(const CallObj* op) final { if (auto prim_expr = ffi::GetRef(op).as(); prim_expr && slot_map_.count(prim_expr.value()) && !IsRelaxOwnedCall(op)) { return RewritePrimValue(prim_expr.value()); @@ -324,7 +324,7 @@ class VMShapeLowerMutator return ExprMutator::VisitExpr_(op); } - Expr VisitExprFallback_(const ExprNode* op) final { + Expr VisitExprFallback_(const ExprObj* op) final { if (auto prim_expr = ffi::GetRef(op).as(); prim_expr && slot_map_.count(prim_expr.value())) { return RewritePrimValue(prim_expr.value()); @@ -332,8 +332,8 @@ class VMShapeLowerMutator return ExprMutator::VisitExprFallback_(op); } - void VisitBinding_(const VarBindingNode* binding) final { - if (binding->var.as()) { + void VisitBinding_(const VarBindingObj* binding) final { + if (binding->var.as()) { ExprMutator::VisitBinding_(binding); return; } @@ -426,7 +426,7 @@ class VMShapeLowerMutator static DLDataType ShapeDType() { return DLDataType{kDLInt, 64, 1}; } PrimExprSlot* GetPrimValueSlot(const Var& var) const { - if (var.as()) return nullptr; + if (var.as()) return nullptr; auto prim_var = var.as(); if (!prim_var) return nullptr; auto it = slot_map_.find(PrimExpr(*prim_var)); @@ -452,7 +452,7 @@ class VMShapeLowerMutator void PopulateSlotInfo() { for (auto& kv : slot_map_) { auto* slot = kv.second; - if (!slot->expr.as()) { + if (!slot->expr.as()) { ffi::Array dep_vars = tvm::UndefinedVars(slot->expr); for (auto var : dep_vars) { auto it = slot_map_.find(var.as_or_throw()); @@ -495,7 +495,7 @@ class VMShapeLowerMutator //------------------------------------------------------- // Expr mutation overloading. //------------------------------------------------------- - Expr VisitExpr_(const FunctionNode* op) final { + Expr VisitExpr_(const FunctionObj* op) final { TVM_FFI_THROW(InternalError) << "VMShapeLower do not work for local functions, make sure " << " to run it after LambdaLift"; return ffi::GetRef(op); @@ -504,7 +504,7 @@ class VMShapeLowerMutator std::pair MakeSymbolicShapeArg(const PrimExpr& expr) { using runtime::vm::MakeShapeCode; - if (auto* int_expr = expr.as()) { + if (auto* int_expr = expr.as()) { return {IntImm::Int64(static_cast(MakeShapeCode::kUseImm)), IntImm::Int64(int_expr->value)}; } else { @@ -521,7 +521,7 @@ class VMShapeLowerMutator Expr RewritePrimValue(const PrimExpr& value) { using runtime::vm::MakeShapeCode; // Constant shape can be preserved. - bool is_const_value = value->IsInstance() || value->IsInstance(); + bool is_const_value = value->IsInstance() || value->IsInstance(); if (is_const_value) { return value; } @@ -536,12 +536,11 @@ class VMShapeLowerMutator return call; } - Expr VisitExpr_(const ShapeExprNode* op) final { + Expr VisitExpr_(const ShapeExprObj* op) final { using runtime::vm::MakeShapeCode; // Constant shape can be preserved. - bool is_const_shape = std::all_of(op->values.begin(), op->values.end(), [](const PrimExpr& e) { - return e->IsInstance(); - }); + bool is_const_shape = std::all_of(op->values.begin(), op->values.end(), + [](const PrimExpr& e) { return e->IsInstance(); }); if (is_const_shape) { return ffi::GetRef(op); } @@ -559,7 +558,7 @@ class VMShapeLowerMutator return call; } - void VisitBinding_(const MatchCastNode* binding) final { + void VisitBinding_(const MatchCastObj* binding) final { std::ostringstream err_ctx; err_ctx << "ErrorContext(match_cast, ty=" << binding->ty << ") "; Expr value = this->VisitExpr(binding->value); @@ -599,7 +598,7 @@ class VMShapeLowerMutator bool require_value_computed) { using runtime::vm::MatchShapeCode; - if (auto* int_expr = expr.as()) { + if (auto* int_expr = expr.as()) { return {MatchShapeCode::kAssertEqualToImm, IntImm::Int64(int_expr->value)}; } @@ -612,7 +611,7 @@ class VMShapeLowerMutator // the value is not yet computed TVM_FFI_ICHECK(!require_value_computed) << "PrimExpr " << expr << " is not computed"; - if (expr.as()) { + if (expr.as()) { // It is a var we will populate it in this round. slot->value_computed = true; @@ -654,9 +653,9 @@ class VMShapeLowerMutator ffi::Array args = {item.input, shape_heap_}; - Expr match_op = item.input->ty.as() ? Expr(builtin_match_prim_value_) - : Expr(builtin_match_shape_); - if (item.input->ty.as()) { + Expr match_op = item.input->ty.as() ? Expr(builtin_match_prim_value_) + : Expr(builtin_match_shape_); + if (item.input->ty.as()) { TVM_FFI_ICHECK_EQ(item.pattern.size(), 1); } else { args.push_back(IntImm::Int64(item.pattern.size())); @@ -795,10 +794,10 @@ class VMShapeLowerMutator return TypeFunctor::VisitType(ty, value, always_check, dynamic_only, err_ctx, match_todos); } - void VisitType_(const AnyTypeNode* op, Expr value, bool always_check, bool dynamic_only, + void VisitType_(const AnyTypeObj* op, Expr value, bool always_check, bool dynamic_only, const ffi::String& err_ctx, std::vector* match_todos) final {} - void VisitType_(const StringTypeNode* op, Expr value, bool always_check, bool dynamic_only, + void VisitType_(const StringTypeObj* op, Expr value, bool always_check, bool dynamic_only, const ffi::String& err_ctx, std::vector* match_todos) final { if (always_check || !IsBaseOf(StringType(), GetType(value))) { builder_->Emit(Call(Type::Missing(), ExternFunc("vm.builtin.check_string_info"), @@ -807,7 +806,7 @@ class VMShapeLowerMutator } } - void VisitType_(const PrimTypeNode* op, Expr value, bool always_check, bool dynamic_only, + void VisitType_(const PrimTypeObj* op, Expr value, bool always_check, bool dynamic_only, const ffi::String& err_ctx, std::vector* match_todos) final { // emit runtime check of shape if (always_check || !IsBaseOf(PrimType(op->dtype), GetType(value))) { @@ -818,7 +817,7 @@ class VMShapeLowerMutator } } - void VisitType_(const ShapeTypeNode* op, Expr value, bool always_check, bool dynamic_only, + void VisitType_(const ShapeTypeObj* op, Expr value, bool always_check, bool dynamic_only, const ffi::String& err_ctx, std::vector* match_todos) final { // emit runtime check of shape if (always_check || !IsBaseOf(ShapeType(op->ndim), GetType(value))) { @@ -832,13 +831,12 @@ class VMShapeLowerMutator } } - void VisitType_(const TensorTypeNode* op, Expr value, bool always_check, bool dynamic_only, + void VisitType_(const TensorTypeObj* op, Expr value, bool always_check, bool dynamic_only, const ffi::String& err_ctx, std::vector* match_todos) final { // emit runtime check of shape - auto* shape_expr = op->shape.as(); - if (dynamic_only && - std::all_of(shape_expr->values.begin(), shape_expr->values.end(), - [](const PrimExpr& e) { return e->IsInstance(); })) { + auto* shape_expr = op->shape.as(); + if (dynamic_only && std::all_of(shape_expr->values.begin(), shape_expr->values.end(), + [](const PrimExpr& e) { return e->IsInstance(); })) { // if we only check dynamic shapes, and the shape is static, we can skip. return; } @@ -854,7 +852,7 @@ class VMShapeLowerMutator if (shape_expr != nullptr) { match_todos->push_back(MatchShapeTodoItem{value, shape_expr->values, err_ctx}); - } else if (op->shape.as()) { + } else if (op->shape.as()) { // NOTE: This part of the logic is left empty for future support as it is less common. // Future implementors: we can emit a binding here and assert here. TVM_FFI_THROW(InternalError) @@ -868,9 +866,9 @@ class VMShapeLowerMutator // This function will try to simplify constant tuples // the return value **always** have type. Expr MakeTupleGetItem(Expr value, int64_t index) { - if (auto* tuple_expr = value.as()) { + if (auto* tuple_expr = value.as()) { return tuple_expr->fields[index]; - } else if (GetTypeAs(value)) { + } else if (GetTypeAs(value)) { // value is tuple type, it is OK to run tuple get item. return TupleGetItem(value, index); } else { @@ -882,9 +880,9 @@ class VMShapeLowerMutator } } - void VisitType_(const TupleTypeNode* op, Expr value, bool always_check, bool dynamic_only, + void VisitType_(const TupleTypeObj* op, Expr value, bool always_check, bool dynamic_only, const ffi::String& err_ctx, std::vector* match_todos) final { - auto* value_tinfo = GetTypeAs(value); + auto* value_tinfo = GetTypeAs(value); if (value_tinfo) { TVM_FFI_CHECK_EQ(value_tinfo->fields.size(), op->fields.size(), TypeError) << err_ctx << " during match-cast we find tuple size mismatch"; @@ -904,7 +902,7 @@ class VMShapeLowerMutator } } - void VisitType_(const FuncTypeNode* op, Expr value, bool always_check, bool dynamic_only, + void VisitType_(const FuncTypeObj* op, Expr value, bool always_check, bool dynamic_only, const ffi::String& err_ctx, std::vector* match_todos) final { // we only check function is callable. if (!always_check && MatchType(value)) return; diff --git a/src/relax/distributed/axis_group_graph.cc b/src/relax/distributed/axis_group_graph.cc index 94526605eac3..8aa23c65d6e1 100644 --- a/src/relax/distributed/axis_group_graph.cc +++ b/src/relax/distributed/axis_group_graph.cc @@ -70,12 +70,12 @@ namespace tvm { namespace relax { namespace distributed { -const TensorTypeNode* GetTensorType(Expr tensor) { - const auto* tensor_ty = GetTypeAs(tensor); +const TensorTypeObj* GetTensorType(Expr tensor) { + const auto* tensor_ty = GetTypeAs(tensor); if (tensor_ty) { return tensor_ty; } - const auto* dtensor_ty = GetTypeAs(tensor); + const auto* dtensor_ty = GetTypeAs(tensor); if (dtensor_ty) { return dtensor_ty->tensor_ty.get(); } @@ -100,7 +100,7 @@ void UnaryOpHelper(ffi::Array tensor_list, distributed::AxisGroupGraph* ax void BuildAxisGraphUnary(const Var& output_var, const Call& call, distributed::AxisGroupGraph* axis_group_graph) { ffi::Array tensor_list; // vars in param and output - if (call->args[0]->IsInstance()) { + if (call->args[0]->IsInstance()) { tensor_list.push_back(call->args[0]); } tensor_list.push_back(output_var); @@ -110,10 +110,10 @@ void BuildAxisGraphUnary(const Var& output_var, const Call& call, void BuildAxisGraphBinary(const Var& output_var, const Call& call, distributed::AxisGroupGraph* axis_group_graph) { ffi::Array tensor_list; // vars in param and output - if (call->args[0]->ty.as() || call->args[0]->ty.as()) { + if (call->args[0]->ty.as() || call->args[0]->ty.as()) { tensor_list.push_back(call->args[0]); } - if (call->args[1]->ty.as() || call->args[1]->ty.as()) { + if (call->args[1]->ty.as() || call->args[1]->ty.as()) { tensor_list.push_back(call->args[1]); } tensor_list.push_back(output_var); @@ -125,8 +125,8 @@ void BuildAxisGraphBinary(const Var& output_var, const Call& call, const auto* x2_ty = GetTensorType(tensor_list[1]); int x1_ndim = x1_ty->ndim; int x2_ndim = x2_ty->ndim; - const auto* x1_shape = x1_ty->shape.as(); - const auto* x2_shape = x2_ty->shape.as(); + const auto* x1_shape = x1_ty->shape.as(); + const auto* x2_shape = x2_ty->shape.as(); TVM_FFI_ICHECK(x1_shape && x2_shape); sym::Analyzer analyzer; for (int i = 1; i <= std::min(x1_ndim, x2_ndim); ++i) { @@ -231,8 +231,8 @@ void BuildAxisGraphMatmul(const Var& output_var, const Call& call, x2_ndim = 2; x2_appended = 1; } - const auto* x1_shape = x1_ty->shape.as(); - const auto* x2_shape = x2_ty->shape.as(); + const auto* x1_shape = x1_ty->shape.as(); + const auto* x2_shape = x2_ty->shape.as(); TVM_FFI_ICHECK(x1_shape && x2_shape); ffi::Array x1_shape_prefix{x1_shape->values.begin(), x1_shape->values.end() - 2 + x1_prepended}; @@ -314,8 +314,8 @@ void BuildAxisGraphReshape(const Var& output_var, const Call& call, distributed::AxisGroupGraph* axis_group_graph) { Expr input_tensor = call->args[0]; const auto* tensor_ty = GetTensorType(input_tensor); - const auto* new_shape_ty = GetTypeAs(call->args[1]); - const auto* old_shape_ty = GetTypeAs(tensor_ty->shape.value()); + const auto* new_shape_ty = GetTypeAs(call->args[1]); + const auto* old_shape_ty = GetTypeAs(tensor_ty->shape.value()); TVM_FFI_ICHECK_NOTNULL(old_shape_ty); ffi::Array old_shape_values = old_shape_ty->values.value(); ffi::Array new_shape_values = new_shape_ty->values.value(); @@ -345,7 +345,7 @@ void BuildAxisGraphReshape(const Var& output_var, const Call& call, inline int GetNumOutput(Call call) { Type output_ty = call->ty_args[0]; - if (const auto* tuple_ty = output_ty.as()) { + if (const auto* tuple_ty = output_ty.as()) { return tuple_ty->fields.size(); } else { return 1; @@ -359,7 +359,7 @@ void BuildAxisGraphCallTIR(const Var& output_var, const Call& call, const tirx:: ffi::Array input_list = call->args[1].as_or_throw()->fields; input_list.push_back(output_var); for (int i = 0; i < static_cast(input_list.size()); i++) { - if (func->params[i]->ty.as()) { + if (func->params[i]->ty.as()) { input_var_to_relax_expr.Set(func->params[i], input_list[i]); } } diff --git a/src/relax/distributed/global_info.cc b/src/relax/distributed/global_info.cc index edd43cc86a79..fa6738a5f948 100644 --- a/src/relax/distributed/global_info.cc +++ b/src/relax/distributed/global_info.cc @@ -24,7 +24,7 @@ namespace tvm { namespace relax { namespace distributed { -TVM_FFI_STATIC_INIT_BLOCK() { DeviceMeshNode::RegisterReflection(); } +TVM_FFI_STATIC_INIT_BLOCK() { DeviceMeshObj::RegisterReflection(); } DeviceMesh::DeviceMesh(ffi::Shape shape, ffi::Array device_ids) : GlobalInfo(ffi::UnsafeInit{}) { @@ -32,7 +32,7 @@ DeviceMesh::DeviceMesh(ffi::Shape shape, ffi::Array device_ids) for (int i = 0; i < static_cast(shape.size()); i++) { prod *= shape[i]; } - ffi::ObjectPtr n = ffi::make_object(); + ffi::ObjectPtr n = ffi::make_object(); TVM_FFI_ICHECK_EQ(prod, static_cast(device_ids.size())) << "The number of device ids must match the product of the shape"; n->shape = std::move(shape); @@ -41,10 +41,10 @@ DeviceMesh::DeviceMesh(ffi::Shape shape, ffi::Array device_ids) } DeviceMesh::DeviceMesh(ffi::Shape shape, Range device_range) : GlobalInfo(ffi::UnsafeInit{}) { - ffi::ObjectPtr n = ffi::make_object(); + ffi::ObjectPtr n = ffi::make_object(); ffi::Array device_ids; - int range_start = device_range->begin.as()->value.as().value(); - int range_extent = device_range->extent.as()->value.as().value(); + int range_start = device_range->begin.as()->value.as().value(); + int range_extent = device_range->extent.as()->value.as().value(); for (int i = range_start; i < range_start + range_extent; i++) { device_ids.push_back(i); } diff --git a/src/relax/distributed/transform/legalize_redistribute.cc b/src/relax/distributed/transform/legalize_redistribute.cc index ec090a340ce0..ee8ac03dbb27 100644 --- a/src/relax/distributed/transform/legalize_redistribute.cc +++ b/src/relax/distributed/transform/legalize_redistribute.cc @@ -50,25 +50,25 @@ class RedistributeLegalizer : public ExprMutator { IRModule Legalize() { auto mod = builder_->GetContextIRModule(); for (const auto& [gv, base_func] : mod->functions) { - const auto* func_ = base_func.as(); + const auto* func_ = base_func.as(); if (func_ == nullptr) { continue; } Expr new_func_body = VisitExpr(func_->body); - auto new_func = ffi::make_object(*func_); + auto new_func = ffi::make_object(*func_); new_func->body = new_func_body; builder_->UpdateFunction(gv, Function(new_func)); } return builder_->GetContextIRModule(); } using ExprMutator::VisitExpr_; - Expr VisitExpr_(const CallNode* op) final { + Expr VisitExpr_(const CallObj* op) final { Call call = ExprMutator::VisitExpr_(op).as_or_throw(); static Op redistribute_op = Op::Get("relax.dist.redistribute"); if (call->op.same_as(redistribute_op)) { const auto* attrs = call->attrs.as(); TVM_FFI_ICHECK(attrs); - const auto* input_ty = call->args[0]->ty.as(); + const auto* input_ty = call->args[0]->ty.as(); TVM_FFI_ICHECK(input_ty); // As the first step, we only support redistribute in the same device mesh, // and the device mesh must be 1d diff --git a/src/relax/distributed/transform/lower_distir.cc b/src/relax/distributed/transform/lower_distir.cc index e692b8c1dfea..9c561e46cbf1 100644 --- a/src/relax/distributed/transform/lower_distir.cc +++ b/src/relax/distributed/transform/lower_distir.cc @@ -54,7 +54,7 @@ class DistIRSharder : public ExprMutator { IRModule Lower() { auto mod = builder_->GetContextIRModule(); for (const auto& [gv, base_func] : mod->functions) { - const auto* func_ = base_func.as(); + const auto* func_ = base_func.as(); if (func_ == nullptr || !IsDistIRFunc(ffi::GetRef(func_))) { continue; } @@ -80,24 +80,24 @@ class DistIRSharder : public ExprMutator { TensorType ShardDTensorType(DTensorType orig_ty) { TensorType tensor_ty = orig_ty->tensor_ty; TVM_FFI_ICHECK(tensor_ty->shape); - const auto* orig_shape = tensor_ty->shape.as(); - auto new_tensor_ty = ffi::make_object(*tensor_ty.get()); + const auto* orig_shape = tensor_ty->shape.as(); + auto new_tensor_ty = ffi::make_object(*tensor_ty.get()); new_tensor_ty->shape = ShardShape(ffi::GetRef(orig_shape), orig_ty->device_mesh, orig_ty->placement); return TensorType(new_tensor_ty); } Type ConvertType(Type orig_ty, bool shard_shape) { - if (const auto* dtensor_ty = orig_ty.as()) { + if (const auto* dtensor_ty = orig_ty.as()) { if (shard_shape) { return ShardDTensorType(ffi::GetRef(dtensor_ty)); } else { return dtensor_ty->tensor_ty; } - } else if (const auto* tuple_ty = orig_ty.as()) { + } else if (const auto* tuple_ty = orig_ty.as()) { ffi::Array new_fields; for (const auto& field_ty : tuple_ty->fields) { - if (const auto* dtensor_ty = field_ty.as()) { + if (const auto* dtensor_ty = field_ty.as()) { if (shard_shape) { new_fields.push_back(ShardDTensorType(ffi::GetRef(dtensor_ty))); } else { @@ -117,10 +117,10 @@ class DistIRSharder : public ExprMutator { TVM_FFI_ICHECK(!input->ty.as().has_value()); Type old_ty = GetType(input); Type new_ty = ConvertType(old_ty, false); - if (const auto* var = input.as()) { + if (const auto* var = input.as()) { Var new_param(var->name, new_ty); return new_param; - } else if (const auto* constant = input.as()) { + } else if (const auto* constant = input.as()) { for (const auto& spec : old_ty.as_or_throw()->placement->dim_specs) { TVM_FFI_ICHECK(spec->kind == PlacementSpecKind::kReplica); } @@ -139,7 +139,7 @@ class DistIRSharder : public ExprMutator { PlacementSpec sharding_spec = dtensor_ty->placement->dim_specs[0]; if (sharding_spec->kind == PlacementSpecKind::kReplica) { Var new_var = builder_->Emit(broadcast_from_worker0(new_expr)); - if (const auto* var = old_expr.as()) { + if (const auto* var = old_expr.as()) { var_remap_.insert_or_assign(ffi::GetRef(var), new_var); } else { tuple_getitem_remap_.insert_or_assign(old_expr.as_or_throw(), new_var); @@ -147,7 +147,7 @@ class DistIRSharder : public ExprMutator { } else if (sharding_spec->kind == PlacementSpecKind::kSharding) { Var scatter_var = builder_->Emit( scatter_from_worker0(new_expr, dtensor_ty->device_mesh->shape[0], sharding_spec->axis)); - if (const auto* var = old_expr.as()) { + if (const auto* var = old_expr.as()) { var_remap_.insert_or_assign(ffi::GetRef(var), scatter_var); } else { tuple_getitem_remap_.insert_or_assign(old_expr.as_or_throw(), scatter_var); @@ -160,11 +160,11 @@ class DistIRSharder : public ExprMutator { void InputPreprocessing() { for (int i = 0; i < static_cast(func_->params.size()); i++) { Var param = func_->params[i]; - if (const auto* dtensor_ty = GetTypeAs(param)) { + if (const auto* dtensor_ty = GetTypeAs(param)) { EmitBroadcastOrScatter(param, new_params_[i], ffi::GetRef(dtensor_ty)); - } else if (const auto* tuple_ty = GetTypeAs(param)) { + } else if (const auto* tuple_ty = GetTypeAs(param)) { for (int j = 0; j < static_cast(tuple_ty->fields.size()); j++) { - if (const auto* dtensor_ty = tuple_ty->fields[j].as()) { + if (const auto* dtensor_ty = tuple_ty->fields[j].as()) { EmitBroadcastOrScatter(TupleGetItem(param, j), TupleGetItem(new_params_[i], j), ffi::GetRef(dtensor_ty)); } @@ -187,7 +187,7 @@ class DistIRSharder : public ExprMutator { return new_func; } - void VisitBinding_(const VarBindingNode* binding, const TupleGetItemNode* val) { + void VisitBinding_(const VarBindingObj* binding, const TupleGetItemObj* val) { if (tuple_getitem_remap_.count(ffi::GetRef(val))) { var_remap_.insert_or_assign(binding->var, tuple_getitem_remap_.at(ffi::GetRef(val))); @@ -196,7 +196,7 @@ class DistIRSharder : public ExprMutator { } } - BindingBlock VisitBindingBlock_(const BindingBlockNode* block) { + BindingBlock VisitBindingBlock_(const BindingBlockObj* block) { builder_->BeginBindingBlock(); InputPreprocessing(); for (Binding binding : block->bindings) { @@ -205,7 +205,7 @@ class DistIRSharder : public ExprMutator { return builder_->EndBlock(); } - BindingBlock VisitBindingBlock_(const DataflowBlockNode* block) { + BindingBlock VisitBindingBlock_(const DataflowBlockObj* block) { builder_->BeginDataflowBlock(); InputPreprocessing(); for (auto binding : block->bindings) { @@ -214,34 +214,34 @@ class DistIRSharder : public ExprMutator { return builder_->EndBlock(); } - Call HandleSpecialCaseinDTensorLowering(const CallNode* call, Var binding_var) { + Call HandleSpecialCaseinDTensorLowering(const CallObj* call, Var binding_var) { static Op reshape_op = Op::Get("relax.reshape"); static Op call_tir_op = Op::Get("relax.call_tir"); static Op call_tir_local_view_op = Op::Get("relax.dist.call_tir_local_view"); if (call->op.same_as(reshape_op)) { - TVM_FFI_ICHECK(call->args[1].as()); - const auto* out_ty = GetTypeAs(binding_var); + TVM_FFI_ICHECK(call->args[1].as()); + const auto* out_ty = GetTypeAs(binding_var); TVM_FFI_ICHECK(out_ty); - auto new_call_node = ffi::make_object(*call); + auto new_call_node = ffi::make_object(*call); new_call_node->args.Set(1, ShardShape(call->args[1].as_or_throw(), out_ty->device_mesh, out_ty->placement)); return Call(new_call_node); } else if (call->op.same_as(call_tir_local_view_op)) { - auto new_call_node = ffi::make_object(*call); + auto new_call_node = ffi::make_object(*call); new_call_node->op = call_tir_op; new_call_node->ty_args = {ConvertType(GetType(binding_var), true)}; return Call(new_call_node); } else if (call->op.same_as(call_tir_op)) { TVM_FFI_THROW(InternalError) << "call_tir should be lowered to call_tir_local_view before lowering to relax"; - } else if (const auto* extern_func = call->op.as()) { - auto new_call_node = ffi::make_object(*call); + } else if (const auto* extern_func = call->op.as()) { + auto new_call_node = ffi::make_object(*call); if (extern_func->global_symbol == "vm.builtin.distributed.attention_kv_cache_append") { new_call_node->op = ExternFunc("vm.builtin.attention_kv_cache_append"); } else if (extern_func->global_symbol == "vm.builtin.distributed.attention_kv_cache_view") { new_call_node->op = ExternFunc("vm.builtin.attention_kv_cache_view"); auto orig_shape = call->args[1].as_or_throw(); - const auto* out_ty = GetTypeAs(binding_var); + const auto* out_ty = GetTypeAs(binding_var); TVM_FFI_ICHECK(out_ty); ShapeExpr new_shape = ShardShape(orig_shape, out_ty->device_mesh, out_ty->placement); new_call_node->args.Set(1, new_shape); @@ -252,7 +252,7 @@ class DistIRSharder : public ExprMutator { return ffi::GetRef(call); } - void VisitBinding_(const VarBindingNode* binding, const CallNode* val) { + void VisitBinding_(const VarBindingObj* binding, const CallObj* val) { Call new_call = this->VisitExpr(HandleSpecialCaseinDTensorLowering(val, binding->var)).as_or_throw(); ReEmitBinding(binding, builder_->Normalize(new_call)); 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 41e5ed3da717..230e71f0a553 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 @@ -55,17 +55,17 @@ class DistTensorReplacer : public s_tir::StmtExprMutator { class DistSBlockInfoCollector : public s_tir::StmtExprVisitor { private: - ffi::Optional Visit_(const TensorStoreNode* op) final { + ffi::Optional Visit_(const TensorStoreObj* op) final { tensor_access_indices[op->dest.as_or_throw()].push_back(op->indices); return s_tir::StmtExprVisitor::Visit_(op); } - ffi::Optional Visit_(const TensorLoadNode* op) final { + ffi::Optional Visit_(const TensorLoadObj* op) final { tensor_access_indices[op->source.as_or_throw()].push_back(op->indices); return s_tir::StmtExprVisitor::Visit_(op); } - ffi::Optional Visit_(const s_tir::SBlockNode* op) final { + ffi::Optional Visit_(const s_tir::SBlockObj* op) final { for (const auto& iter_var : op->iter_vars) { if (iter_var->iter_type == s_tir::kCommReduce) { TVM_FFI_ICHECK(op->writes.size() == 1); @@ -76,7 +76,7 @@ class DistSBlockInfoCollector : public s_tir::StmtExprVisitor { } bool IsReduceTensorAccess(const PrimExpr& expr) { - if (const auto* tensor_load = expr.as()) { + if (const auto* tensor_load = expr.as()) { return reduce_tensor_.has_value() && tensor_load->source.as_or_throw().same_as( reduce_tensor_.value()); @@ -84,28 +84,28 @@ class DistSBlockInfoCollector : public s_tir::StmtExprVisitor { return false; } - ffi::Optional Visit_(const prim::AddNode* op) final { + ffi::Optional Visit_(const prim::AddObj* op) final { if (IsReduceTensorAccess(op->a) || IsReduceTensorAccess(op->b)) { reduce_kind = "sum"; } return s_tir::StmtExprVisitor::Visit_(op); } - ffi::Optional Visit_(const prim::MulNode* op) final { + ffi::Optional Visit_(const prim::MulObj* op) final { if (IsReduceTensorAccess(op->a) || IsReduceTensorAccess(op->b)) { reduce_kind = "prod"; } return s_tir::StmtExprVisitor::Visit_(op); } - ffi::Optional Visit_(const prim::MinNode* op) final { + ffi::Optional Visit_(const prim::MinObj* op) final { if (IsReduceTensorAccess(op->a) || IsReduceTensorAccess(op->b)) { reduce_kind = "min"; } return s_tir::StmtExprVisitor::Visit_(op); } - ffi::Optional Visit_(const prim::MaxNode* op) final { + ffi::Optional Visit_(const prim::MaxObj* op) final { if (IsReduceTensorAccess(op->a) || IsReduceTensorAccess(op->b)) { reduce_kind = "max"; } @@ -134,7 +134,7 @@ class DistributedTensorCompactor : public s_tir::StmtExprMutator { ffi::Array new_params; ffi::Map replace_tensor_map; for (const Var& param : function->params) { - if (!param->ty.as()) { + if (!param->ty.as()) { new_params.push_back(param); continue; } @@ -169,7 +169,7 @@ class DistributedTensorCompactor : public s_tir::StmtExprMutator { std::unordered_set visited; for (int i = 0, j = 0; i < static_cast(function->params.size()); i++) { Var param_var = function->params[i]; - if (!param_var->ty.as()) { + if (!param_var->ty.as()) { continue; } TensorVar param_tensor = param_var.as_or_throw(); @@ -260,13 +260,13 @@ class DistributedTensorCompactor : public s_tir::StmtExprMutator { shape.push_back(tensor.shape()[i]); } } - auto new_type_node = ffi::make_object(*tensor.ty().get()); + auto new_type_node = ffi::make_object(*tensor.ty().get()); new_type_node->shape = std::move(shape); TensorType new_type(std::move(new_type_node)); return TensorVar(tensor->name, std::move(new_type), tensor->loc); } - UnchangedOr Mutate_(const s_tir::SBlockNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const s_tir::SBlockObj* op, InplaceMode inplace_mode) final { s_tir::SBlock block = s_tir::StmtExprMutator::Mutate_(op, inplace_mode) .ValueOrUnchanged(ffi::GetRef(op)) .as_or_throw(); @@ -292,8 +292,8 @@ class DistributedTensorCompactor : public s_tir::StmtExprMutator { break; } } - ffi::ObjectPtr new_block = - ffi::make_object(*block.operator->()); + ffi::ObjectPtr new_block = + ffi::make_object(*block.operator->()); new_block->iter_vars = new_iter_vars; new_block->alloc_tensors = new_alloc_tensors; if (new_block->name_hint == "root") { @@ -307,7 +307,7 @@ class DistributedTensorCompactor : public s_tir::StmtExprMutator { void AddAllReduceBlock(std::string reduce_kind) { add_allreduce_kind_ = reduce_kind; } - UnchangedOr Mutate_(const s_tir::SBlockRealizeNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const s_tir::SBlockRealizeObj* op, InplaceMode inplace_mode) final { s_tir::SBlockRealize realize = s_tir::StmtExprMutator::Mutate_(op, inplace_mode) .ValueOrUnchanged(ffi::GetRef(op)) .as_or_throw(); @@ -325,7 +325,7 @@ class DistributedTensorCompactor : public s_tir::StmtExprMutator { return realize; } - UnchangedOr Mutate_(const ForNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const ForObj* op, InplaceMode inplace_mode) final { For new_loop = s_tir::StmtExprMutator::Mutate_(op, inplace_mode) .ValueOrUnchanged(ffi::GetRef(op)) .as_or_throw(); @@ -368,12 +368,12 @@ class LowerTIRToLocalView : public ExprMutator { IRModule Lower() { auto mod = builder_->GetContextIRModule(); for (const auto& [gv, base_func] : mod->functions) { - const auto* func_ = base_func.as(); + const auto* func_ = base_func.as(); if (func_ == nullptr || !IsDistIRFunc(ffi::GetRef(func_))) { continue; } Expr new_func_body = this->VisitExpr(func_->body); - ffi::ObjectPtr new_func = ffi::make_object(*func_); + ffi::ObjectPtr new_func = ffi::make_object(*func_); new_func->body = new_func_body; builder_->UpdateFunction(gv, Function(new_func)); } @@ -382,9 +382,9 @@ class LowerTIRToLocalView : public ExprMutator { private: inline ffi::Array ExtractDTensorType(Var var) { - if (const auto* dtensor_ty = GetTypeAs(var)) { + if (const auto* dtensor_ty = GetTypeAs(var)) { return {ffi::GetRef(dtensor_ty)}; - } else if (const auto* tuple_ty = GetTypeAs(var)) { + } else if (const auto* tuple_ty = GetTypeAs(var)) { ffi::Array ret; for (const auto& field : tuple_ty->fields) { ret.push_back(field.as_or_throw()); @@ -396,7 +396,7 @@ class LowerTIRToLocalView : public ExprMutator { } } - void VisitBinding_(const VarBindingNode* binding, const CallNode* val) final { + void VisitBinding_(const VarBindingObj* binding, const CallObj* val) final { static const Op call_tir_op = Op::Get("relax.call_tir"); if (!val->op.same_as(call_tir_op)) { ExprMutator::VisitBinding_(binding, val); @@ -410,8 +410,8 @@ class LowerTIRToLocalView : public ExprMutator { for (size_t i = 0; i < args.size(); ++i) { const Expr& arg = args[i]; const tvm::Var& param = function->params[i]; - if (param->ty.as()) { - const auto* ty = GetTypeAs(arg); + if (param->ty.as()) { + const auto* ty = GetTypeAs(arg); TVM_FFI_CHECK(ty, TypeError) << "Expected tensor parameter " << param << " to receive a distributed tensor, but " << arg << " has type " << GetType(arg); @@ -432,7 +432,7 @@ class LowerTIRToLocalView : public ExprMutator { tirx::DistributedTensorCompactor::DistTensorCompact(sharding_specs, function); auto new_gvar = builder_->AddFunction(new_function, gvar->name_hint); Call call = this->VisitExpr(binding->value).as_or_throw(); - ffi::ObjectPtr new_call_node = ffi::make_object(*call.get()); + ffi::ObjectPtr new_call_node = ffi::make_object(*call.get()); new_call_node->op = Op::Get("relax.dist.call_tir_local_view"); new_call_node->args.Set(0, new_gvar); Call new_call(new_call_node); diff --git a/src/relax/distributed/transform/propagate_sharding.cc b/src/relax/distributed/transform/propagate_sharding.cc index 73d612640053..9144a581d99b 100644 --- a/src/relax/distributed/transform/propagate_sharding.cc +++ b/src/relax/distributed/transform/propagate_sharding.cc @@ -41,7 +41,7 @@ namespace tvm { namespace relax { namespace distributed { -void CollectAxisGraphBinary(const VarBindingNode* binding, const CallNode* call, +void CollectAxisGraphBinary(const VarBindingObj* binding, const CallObj* call, AxisGroupGraph* axis_group_graph) { const std::vector binary_op_names = { "add", "subtract", "multiply", "divide", "power", "floor_divide", "equal", @@ -55,7 +55,7 @@ void CollectAxisGraphBinary(const VarBindingNode* binding, const CallNode* call, } } -void CollectAxisGraphUnary(const VarBindingNode* binding, const CallNode* call, +void CollectAxisGraphUnary(const VarBindingObj* binding, const CallObj* call, AxisGroupGraph* axis_group_graph) { const std::vector unary_op_names = { "abs", "acos", "acosh", @@ -77,7 +77,7 @@ void CollectAxisGraphUnary(const VarBindingNode* binding, const CallNode* call, } } -void CollectAxisGraphReduce(const VarBindingNode* binding, const CallNode* call, +void CollectAxisGraphReduce(const VarBindingObj* binding, const CallObj* call, AxisGroupGraph* axis_group_graph) { const std::vector reduction_op_names = {"sum", "max", "min", "prod", "mean", "std", "variance", "nn.softmax"}; @@ -90,7 +90,7 @@ void CollectAxisGraphReduce(const VarBindingNode* binding, const CallNode* call, } } -void CollectAxisGraphMatmul(const VarBindingNode* binding, const CallNode* call, +void CollectAxisGraphMatmul(const VarBindingObj* binding, const CallObj* call, AxisGroupGraph* axis_group_graph) { static const Op matmul_op = Op::Get("relax.matmul"); if (call->op.same_as(matmul_op)) { @@ -98,7 +98,7 @@ void CollectAxisGraphMatmul(const VarBindingNode* binding, const CallNode* call, } } -void CollectAxisGraphPermuteDims(const VarBindingNode* binding, const CallNode* call, +void CollectAxisGraphPermuteDims(const VarBindingObj* binding, const CallObj* call, AxisGroupGraph* axis_group_graph) { static const Op permute_dims_op = Op::Get("relax.permute_dims"); if (call->op.same_as(permute_dims_op)) { @@ -106,7 +106,7 @@ void CollectAxisGraphPermuteDims(const VarBindingNode* binding, const CallNode* } } -void CollectAxisGraphReshape(const VarBindingNode* binding, const CallNode* call, +void CollectAxisGraphReshape(const VarBindingObj* binding, const CallObj* call, AxisGroupGraph* axis_group_graph) { static const Op reshape_op = Op::Get("relax.reshape"); if (call->op.same_as(reshape_op)) { @@ -114,7 +114,7 @@ void CollectAxisGraphReshape(const VarBindingNode* binding, const CallNode* call } } -void CollectAxisGraphForDeviceMesh(const VarBindingNode* binding, const CallNode* call, +void CollectAxisGraphForDeviceMesh(const VarBindingObj* binding, const CallObj* call, AxisGroupGraph* axis_group_graph) { ffi::Array tensor_list; static const Op call_tir_op = Op::Get("relax.call_tir"); @@ -125,7 +125,7 @@ void CollectAxisGraphForDeviceMesh(const VarBindingNode* binding, const CallNode args = call->args; } for (const auto& arg : args) { - if (arg->ty.as()) { + if (arg->ty.as()) { tensor_list.push_back(arg); } } @@ -150,7 +150,7 @@ class AxisGroupGraphBuilder : public ExprVisitor { explicit AxisGroupGraphBuilder(AxisGroupGraph* axis_group_graph, IRModule mod) : axis_group_graph_(axis_group_graph), mod_(mod) {} - void VisitBinding_(const VarBindingNode* binding, const CallNode* val) { + void VisitBinding_(const VarBindingObj* binding, const CallObj* val) { CollectAxisGraphBinary(binding, val, axis_group_graph_); CollectAxisGraphUnary(binding, val, axis_group_graph_); CollectAxisGraphReduce(binding, val, axis_group_graph_); @@ -168,10 +168,10 @@ class AxisGroupGraphBuilder : public ExprVisitor { ExprVisitor::VisitBinding_(binding, val); } - void VisitBinding_(const VarBindingNode* binding, const TupleGetItemNode* val) { + void VisitBinding_(const VarBindingObj* binding, const TupleGetItemObj* val) { axis_group_graph_->JoinAxis(Axis(val->tuple.get(), -1, val->index), {binding->var.get(), -1}, distributed::AxisGroupGraph::EdgeType::kDescend); - const auto* tensor_ty = GetTypeAs(binding->var); + const auto* tensor_ty = GetTypeAs(binding->var); if (!tensor_ty) { ExprVisitor::VisitBinding_(binding, val); return; @@ -184,11 +184,11 @@ class AxisGroupGraphBuilder : public ExprVisitor { ExprVisitor::VisitBinding_(binding, val); } - void VisitBinding_(const VarBindingNode* binding, const VarNode* val) { + void VisitBinding_(const VarBindingObj* binding, const VarObj* val) { ffi::Array tensor_tys; - if (const auto* tensor_ty = binding->var->ty.as()) { + if (const auto* tensor_ty = binding->var->ty.as()) { tensor_tys.push_back(ffi::GetRef(tensor_ty)); - } else if (const auto* tuple_ty = binding->var->ty.as()) { + } else if (const auto* tuple_ty = binding->var->ty.as()) { TVM_FFI_ICHECK(tuple_ty); for (const auto& field_ty : tuple_ty->fields) { tensor_tys.push_back(field_ty.as_or_throw()); @@ -224,7 +224,7 @@ class ShardingAnnotationCollector : public ExprVisitor { private: explicit ShardingAnnotationCollector(AxisGroupGraph* axis_group_graph) : axis_group_graph_(axis_group_graph) {} - void VisitBinding_(const VarBindingNode* binding, const CallNode* val) { + void VisitBinding_(const VarBindingObj* binding, const CallObj* val) { static const Op annotate_sharding_op = Op::Get("relax.dist.annotate_sharding"); if (val->op.same_as(annotate_sharding_op)) { const auto* attrs = val->attrs.as(); @@ -255,7 +255,7 @@ class ShardingConflictHandler : public ExprVisitor { ShardingConflictHandler handler(axis_group_graph); handler.VisitExpr(function); for (const Var& var : function->params) { - if (GetTypeAs(var)) { + if (GetTypeAs(var)) { handler.CheckTensorShardingCompatible(var); } } @@ -267,9 +267,9 @@ class ShardingConflictHandler : public ExprVisitor { : axis_group_graph_(axis_group_graph) {} void CheckTensorShardingCompatible(Var var) { - const auto* tensor_ty = GetTypeAs(var); + const auto* tensor_ty = GetTypeAs(var); TVM_FFI_ICHECK(tensor_ty); - const auto* shape = tensor_ty->shape.as(); + const auto* shape = tensor_ty->shape.as(); TVM_FFI_ICHECK(shape); int ndim = tensor_ty->ndim; std::unordered_set sharded_mesh_dim; @@ -297,7 +297,7 @@ class ShardingConflictHandler : public ExprVisitor { << ": Replicate sharding device mesh axis " << sharding_dim << ". Conflict Handling logic will be added in the future."; sharded_mesh_dim.insert(sharding_dim); - if (const auto* val = shape->values[i].as()) { + if (const auto* val = shape->values[i].as()) { if (val->value < device_mesh.value()->shape[sharding_spec->second]) { axis_group_graph_->AddPropagationCutPoint({var.get(), i}, sharding_spec.value()); } @@ -307,7 +307,7 @@ class ShardingConflictHandler : public ExprVisitor { } void CheckConstantNoSharding(GenericConst constant) { - const auto* tensor_ty = GetTypeAs(constant); + const auto* tensor_ty = GetTypeAs(constant); for (int i = 0; i < tensor_ty->ndim; i++) { std::optional sharding_spec; int has_sharding_spec; @@ -318,18 +318,18 @@ class ShardingConflictHandler : public ExprVisitor { } } - void VisitExpr_(const CallNode* op) final { + void VisitExpr_(const CallObj* op) final { ffi::Array args = GetCallArgs(ffi::GetRef(op)); for (const auto& arg : args) { - if (arg.as() && arg->ty.as()) { + if (arg.as() && arg->ty.as()) { CheckConstantNoSharding(arg.as_or_throw()); } } ExprVisitor::VisitExpr_(op); } - void VisitBinding_(const VarBindingNode* binding) final { - if (GetTypeAs(binding->var)) { + void VisitBinding_(const VarBindingObj* binding) final { + if (GetTypeAs(binding->var)) { CheckTensorShardingCompatible(binding->var); } ExprVisitor::VisitBinding_(binding); @@ -348,7 +348,7 @@ class DistributedIRBuilder : public ExprMutator { IRModule BuildDistributedIR() { auto mod = builder_->GetContextIRModule(); for (const auto& [gv, base_func] : mod->functions) { - const auto* func_ = base_func.as(); + const auto* func_ = base_func.as(); if (func_ == nullptr || !IsShardingAnnotatedFunc(ffi::GetRef(func_))) { continue; } @@ -385,12 +385,12 @@ class DistributedIRBuilder : public ExprMutator { Expr RewriteInputTensorAndConstant(Expr tensor) { Type new_ty = Type::Missing(); - if (tensor->ty.as()) { + if (tensor->ty.as()) { new_ty = ConvertToDTensorType(tensor->ty.as_or_throw(), tensor); - } else if (const auto* tuple = tensor->ty.as()) { + } else if (const auto* tuple = tensor->ty.as()) { ffi::Array tuple_ty_fields; for (int i = 0; i < static_cast(tuple->fields.size()); i++) { - if (tuple->fields[i].as()) { + if (tuple->fields[i].as()) { tuple_ty_fields.push_back( ConvertToDTensorType(tuple->fields[i].as_or_throw(), tensor, i)); } else { @@ -400,10 +400,10 @@ class DistributedIRBuilder : public ExprMutator { new_ty = TupleType(tuple_ty_fields); } - if (const auto* var = tensor.as()) { + if (const auto* var = tensor.as()) { Var new_param(var->name, new_ty); return new_param; - } else if (const auto* constant = tensor.as()) { + } else if (const auto* constant = tensor.as()) { GenericConst new_constant(new_ty, constant->value.cast()); return new_constant; } else { @@ -422,7 +422,7 @@ class DistributedIRBuilder : public ExprMutator { // Step 4. Rewrite Function ffi::Array new_params; for (const Var& var : func->params) { - if (GetTypeAs(var) || GetTypeAs(var)) { + if (GetTypeAs(var) || GetTypeAs(var)) { Var new_param = RewriteInputTensorAndConstant(var).as_or_throw(); input_tensor_remap_.Set(var, new_param); new_params.push_back(new_param); @@ -435,7 +435,7 @@ class DistributedIRBuilder : public ExprMutator { return new_func; } - Expr VisitExpr_(const CallNode* call) final { + Expr VisitExpr_(const CallObj* call) final { static const Op call_tir_op = Op::Get("relax.call_tir"); FBuildAxisGraph f = [&](const Var& var, const Call& call, AxisGroupGraph* axis_group_graph) { ffi::Optional function = @@ -446,12 +446,12 @@ class DistributedIRBuilder : public ExprMutator { Call new_call = ExprMutator::VisitExpr_(call).as_or_throw(); ffi::Array args = GetCallArgs(new_call); for (int i = 0; i < static_cast(args.size()); i++) { - if (args[i].as() && args[i]->ty.as()) { + if (args[i].as() && args[i]->ty.as()) { args.Set(i, RewriteInputTensorAndConstant(args[i])); } } - ffi::ObjectPtr n = ffi::make_object(*new_call.get()); + ffi::ObjectPtr n = ffi::make_object(*new_call.get()); if (new_call->op.same_as(call_tir_op)) { // do not infer output type when arg size is 0 if (!args.empty()) { @@ -462,7 +462,7 @@ class DistributedIRBuilder : public ExprMutator { n->args = args; } - if (const auto* extern_func = new_call->op.as()) { + if (const auto* extern_func = new_call->op.as()) { if (extern_func->global_symbol == "vm.builtin.attention_kv_cache_append") { n->op = ExternFunc("vm.builtin.distributed.attention_kv_cache_append"); } else if (extern_func->global_symbol == "vm.builtin.attention_kv_cache_view") { @@ -489,9 +489,9 @@ class DistributedIRBuilder : public ExprMutator { // In cases when inference fails (for example, arg size is 0), use the propagated output type. Call new_call = call; static Op call_tir_op = Op::Get("relax.call_tir"); - if (const auto* extern_func = call->op.as()) { + if (const auto* extern_func = call->op.as()) { if (extern_func->global_symbol == "vm.builtin.distributed.attention_kv_cache_view") { - ffi::ObjectPtr new_call_node = ffi::make_object(*call.get()); + ffi::ObjectPtr new_call_node = ffi::make_object(*call.get()); Type new_dtensor_ty = DTensorType(call->ty_args[0].as_or_throw(), device_mesh, placements[0]); new_call_node->ty_args = {new_dtensor_ty}; @@ -501,12 +501,12 @@ class DistributedIRBuilder : public ExprMutator { } else if (call->op.same_as(call_tir_op)) { TVM_FFI_ICHECK(call->ty_args.size() == 1); if (!TypeCompatibleWithDistIR(call->ty_args)) { - ffi::ObjectPtr new_call_node = ffi::make_object(*call.get()); + ffi::ObjectPtr new_call_node = ffi::make_object(*call.get()); if (placements.size() == 1) { new_call_node->ty_args = { DTensorType(call->ty_args[0].as_or_throw(), device_mesh, placements[0])}; } else { - const auto* tuple_ty = call->ty_args[0].as(); + const auto* tuple_ty = call->ty_args[0].as(); TVM_FFI_ICHECK(placements.size() == tuple_ty->fields.size()); ffi::Array new_tuple_ty_fields; for (int i = 0; i < static_cast(placements.size()); i++) { @@ -522,11 +522,11 @@ class DistributedIRBuilder : public ExprMutator { return new_call; } - void VisitBinding_(const VarBindingNode* binding, const CallNode* val) { + void VisitBinding_(const VarBindingObj* binding, const CallObj* val) { ffi::Array orig_output_tys; - if (const auto* tensor_ty = GetTypeAs(binding->var)) { + if (const auto* tensor_ty = GetTypeAs(binding->var)) { orig_output_tys.push_back(ffi::GetRef(tensor_ty)); - } else if (const auto* tuple_ty = GetTypeAs(binding->var)) { + } else if (const auto* tuple_ty = GetTypeAs(binding->var)) { for (const auto& field_ty : tuple_ty->fields) { orig_output_tys.push_back(field_ty.as_or_throw()); } @@ -558,20 +558,20 @@ class DistributedIRBuilder : public ExprMutator { new_call = builder_->Normalize(RewriteOutType(new_call, device_mesh, placements)).as_or_throw(); - if (const auto* inferred_dtensor_ty = new_call->ty.as()) { + if (const auto* inferred_dtensor_ty = new_call->ty.as()) { Expr new_value = RemoveAnnotateSharding(new_call); if (!ffi::StructuralEqual()( DTensorType(inferred_dtensor_ty->tensor_ty, device_mesh, placements[0]), new_call->ty)) { new_value = InsertRedistribute(new_value, device_mesh, placements[0]); } - if (const auto* var = new_value.as()) { + if (const auto* var = new_value.as()) { var_remap_.insert_or_assign(binding->var, ffi::GetRef(var)); } else { ReEmitBinding(binding, builder_->Normalize(new_value)); } } else { - const auto* inferred_tuple_ty = new_call->ty.as(); + const auto* inferred_tuple_ty = new_call->ty.as(); TVM_FFI_ICHECK(inferred_tuple_ty) << new_call; Var new_var = builder_->Emit(new_call); var_remap_.insert_or_assign(binding->var, new_var); @@ -588,7 +588,7 @@ class DistributedIRBuilder : public ExprMutator { } } - void VisitBinding_(const VarBindingNode* binding, const TupleGetItemNode* val) { + void VisitBinding_(const VarBindingObj* binding, const TupleGetItemObj* val) { if (tuple_getitem_remap_.count(ffi::GetRef(val))) { var_remap_.insert_or_assign(binding->var, tuple_getitem_remap_.at(ffi::GetRef(val))); @@ -597,7 +597,7 @@ class DistributedIRBuilder : public ExprMutator { } } - Expr VisitExpr_(const VarNode* var) final { + Expr VisitExpr_(const VarObj* var) final { Var var_ref = ffi::GetRef(var); auto it = input_tensor_remap_.find(var_ref); if (it != input_tensor_remap_.end()) { diff --git a/src/relax/distributed/transform/utils.cc b/src/relax/distributed/transform/utils.cc index df93811ca14b..8ed5644214e5 100644 --- a/src/relax/distributed/transform/utils.cc +++ b/src/relax/distributed/transform/utils.cc @@ -25,10 +25,10 @@ namespace distributed { bool TypeCompatibleWithDistIR(ffi::Array tys) { bool compatible = true; for (const auto& ty : tys) { - if (const auto* tuple_ty = ty.as()) { + if (const auto* tuple_ty = ty.as()) { compatible &= TypeCompatibleWithDistIR(tuple_ty->fields); } else { - compatible &= !ty->IsInstance(); + compatible &= !ty->IsInstance(); } } return compatible; @@ -37,10 +37,10 @@ bool TypeCompatibleWithDistIR(ffi::Array tys) { bool TypeCompatibleWithRelax(ffi::Array tys) { bool compatible = true; for (const auto& ty : tys) { - if (const auto* tuple_ty = ty.as()) { + if (const auto* tuple_ty = ty.as()) { compatible &= TypeCompatibleWithRelax(tuple_ty->fields); } else { - compatible &= !ty->IsInstance(); + compatible &= !ty->IsInstance(); } } return compatible; @@ -65,7 +65,7 @@ bool IsDistIRFunc(Function func) { bool IsShardingAnnotatedFunc(Function func) { bool has_annotate_sharding = false; PostOrderVisit(func, [&has_annotate_sharding](const Expr& e) { - const CallNode* call = e.as(); + const CallObj* call = e.as(); if (!call) { return; } diff --git a/src/relax/distributed/transform/utils.h b/src/relax/distributed/transform/utils.h index d34b5b95e683..516f9c2e2758 100644 --- a/src/relax/distributed/transform/utils.h +++ b/src/relax/distributed/transform/utils.h @@ -38,7 +38,7 @@ inline ffi::Optional MatchFunction(const IRModule& mod_, const E const GlobalVar& global_var = op.as_or_throw(); // NOTE: as check works for nullptr(returns null) ffi::Optional base_func = mod_->functions.Get(global_var); - if (auto* pfunc = base_func.as()) { + if (auto* pfunc = base_func.as()) { return ffi::GetRef(pfunc); } return std::nullopt; diff --git a/src/relax/distributed/type.cc b/src/relax/distributed/type.cc index 60c9f0584bb0..6318059ceccc 100644 --- a/src/relax/distributed/type.cc +++ b/src/relax/distributed/type.cc @@ -34,8 +34,8 @@ namespace { TVM_FFI_INLINE ffi::Expected> DTensorTypeVisit( ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept { - const DTensorTypeNode* self = - ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value); + const DTensorTypeObj* self = + ffi::detail::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value); TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->device_mesh)); TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->placement)); TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->tensor_ty)); @@ -44,8 +44,8 @@ TVM_FFI_INLINE ffi::Expected> DTensorTypeVisi TVM_FFI_INLINE ffi::Expected> DTensorTypeMutate( ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept { - const DTensorTypeNode* self = - ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value); + const DTensorTypeObj* self = + ffi::detail::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value); TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, mapped_device_mesh, mutator->MutateExpected(self->device_mesh)); TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, mapped_placement, @@ -57,7 +57,7 @@ TVM_FFI_INLINE ffi::Expected> DTensorTypeMutate( mapped_tensor_ty.UnchangedOrSameAs(self->tensor_ty)) { return ffi::Unchanged(); } - ffi::ObjectPtr copy = ffi::make_object(*self); + ffi::ObjectPtr copy = ffi::make_object(*self); copy->device_mesh = std::move(mapped_device_mesh).ValueOrUnchanged(std::move(copy->device_mesh)); copy->placement = std::move(mapped_placement).ValueOrUnchanged(std::move(copy->placement)); copy->tensor_ty = std::move(mapped_tensor_ty).ValueOrUnchanged(std::move(copy->tensor_ty)); @@ -66,8 +66,8 @@ TVM_FFI_INLINE ffi::Expected> DTensorTypeMutate( TVM_FFI_INLINE ffi::Expected> DTensorTypeMaybeInplaceMutate( ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept { - DTensorTypeNode* self = const_cast( - ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value)); + DTensorTypeObj* self = const_cast( + ffi::detail::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value)); TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN( ffi::UnchangedOr, mapped_device_mesh, mutator->MutateExpected(self->device_mesh, ffi::InplaceMode::kAllow)); @@ -92,19 +92,19 @@ TVM_FFI_INLINE ffi::Expected> DTensorTypeMaybeInplace } // namespace TVM_FFI_STATIC_INIT_BLOCK() { - PlacementNode::RegisterReflection(); - PlacementSpecNode::RegisterReflection(); + PlacementObj::RegisterReflection(); + PlacementSpecObj::RegisterReflection(); } PlacementSpec PlacementSpec::Sharding(int axis) { - ffi::ObjectPtr n = ffi::make_object(); + ffi::ObjectPtr n = ffi::make_object(); n->axis = axis; n->kind = PlacementSpecKind::kSharding; return PlacementSpec(ffi::UnsafeInit{}, std::move(n)); } PlacementSpec PlacementSpec::Replica() { - ffi::ObjectPtr n = ffi::make_object(); + ffi::ObjectPtr n = ffi::make_object(); n->axis = -1; n->kind = PlacementSpecKind::kReplica; return PlacementSpec(ffi::UnsafeInit{}, std::move(n)); @@ -117,7 +117,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { .def("relax.distributed.Replica", []() { return PlacementSpec::Replica(); }); } -ffi::String PlacementNode::ToString() const { +ffi::String PlacementObj::ToString() const { std::stringstream ss; for (size_t i = 0; i < dim_specs.size(); ++i) { if (i != 0) { @@ -133,7 +133,7 @@ ffi::String PlacementNode::ToString() const { } Placement::Placement(ffi::Array dim_specs) { - ffi::ObjectPtr n = ffi::make_object(); + ffi::ObjectPtr n = ffi::make_object(); n->dim_specs = std::move(dim_specs); data_ = std::move(n); } @@ -193,7 +193,7 @@ DTensorType::DTensorType(TensorType tensor_ty, DeviceMesh device_mesh, Placement TVM_FFI_CHECK_LT(spec->axis, tensor_ty->ndim, ValueError) << "Sharding dimension should be smaller than tensor ndim"; } - ffi::ObjectPtr n = ffi::make_object( + ffi::ObjectPtr n = ffi::make_object( std::move(tensor_ty), std::move(device_mesh), std::move(placement)); n->loc = loc; data_ = std::move(n); @@ -201,8 +201,8 @@ DTensorType::DTensorType(TensorType tensor_ty, DeviceMesh device_mesh, Placement TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - DTensorTypeNode::RegisterReflection(); - refl::TypeAttrDef() + DTensorTypeObj::RegisterReflection(); + refl::TypeAttrDef() .attr(ffi::type_attr::kStructuralVisit, ffi::FStructuralVisit::FromNative<&DTensorTypeVisit>()) .attr(ffi::type_attr::kStructuralMutate, diff --git a/src/relax/ir/binding_rewrite.cc b/src/relax/ir/binding_rewrite.cc index 5ee1dbeaaec2..e83338621514 100644 --- a/src/relax/ir/binding_rewrite.cc +++ b/src/relax/ir/binding_rewrite.cc @@ -71,21 +71,21 @@ TVM_FFI_STATIC_INIT_BLOCK() { void DataflowBlockRewriteObj::ReplaceAllUses(Var old_var, Var new_var) { class ReplaceAllUsePass : public ExprMutator { Var old_var, new_var; - const DataflowBlockNode* const to_catch; + const DataflowBlockObj* const to_catch; public: ffi::Optional caught; - ReplaceAllUsePass(Var old_var, Var new_var, const DataflowBlockNode* to_catch) + ReplaceAllUsePass(Var old_var, Var new_var, const DataflowBlockObj* to_catch) : old_var(old_var), new_var(new_var), to_catch(to_catch) {} using ExprMutator::VisitExpr_; - Expr VisitExpr_(const VarNode* op) override { + Expr VisitExpr_(const VarObj* op) override { return (op == old_var.get()) ? new_var : ffi::GetRef(op); } - BindingBlock VisitBindingBlock_(const DataflowBlockNode* op) override { + BindingBlock VisitBindingBlock_(const DataflowBlockObj* op) override { BindingBlock res = ExprMutator::VisitBindingBlock_(op); if (op == to_catch) caught = res.as_or_throw(); return res; @@ -144,16 +144,16 @@ class UpdateDFB : public ExprMutator { UpdateDFB(DataflowBlock old_dfb, DataflowBlock new_dfb) : old_dfb(std::move(old_dfb)), new_dfb(std::move(new_dfb)) {} - BindingBlock VisitBindingBlock_(const DataflowBlockNode* op) override { + BindingBlock VisitBindingBlock_(const DataflowBlockObj* op) override { return old_dfb.get() == op ? new_dfb : old_dfb; } }; void DataflowBlockRewriteObj::Add(Binding binding) { auto [var, val] = [binding] { - if (auto vb = binding.as()) { + if (auto vb = binding.as()) { return std::make_pair(vb->var, vb->value); - } else if (auto mc = binding.as()) { + } else if (auto mc = binding.as()) { return std::make_pair(mc->var, mc->value); } TVM_FFI_THROW(InternalError) << "Unsupported binding type"; @@ -180,7 +180,7 @@ void DataflowBlockRewriteObj::Add(Binding binding) { root_fn_ = updater.VisitExpr(root_fn_.value()).as_or_throw(); } - for (const VarNode* v : used_vars) { + for (const VarObj* v : used_vars) { auto var = ffi::GetRef(v); if (auto users = to_users_.Get(var)) { users.value().push_back(var); @@ -258,7 +258,7 @@ class RemoveUnusedVars : public ExprMutator { RemoveUnusedVars(ffi::Map> users, ffi::Array fn_outputs) : RemoveUnusedVars(GetUnusedVars(users, fn_outputs)) {} - void VisitBinding_(const VarBindingNode* binding) override { + void VisitBinding_(const VarBindingObj* binding) override { bool can_remove = unused_vars.count(binding->var) && (in_dataflow_block_ || !ContainsImpureCall(binding->value)); if (!can_remove) { @@ -266,7 +266,7 @@ class RemoveUnusedVars : public ExprMutator { } } - BindingBlock VisitBindingBlock_(const DataflowBlockNode* block) override { + BindingBlock VisitBindingBlock_(const DataflowBlockObj* block) override { bool capture_output = caught_rewrite && block == caught_rewrite.value().get(); bool cache = in_dataflow_block_; diff --git a/src/relax/ir/block_builder.cc b/src/relax/ir/block_builder.cc index 13405a90be5e..72888c73b771 100644 --- a/src/relax/ir/block_builder.cc +++ b/src/relax/ir/block_builder.cc @@ -205,7 +205,9 @@ class BlockBuilderImpl : public BlockBuilderObj { // of shape inference. In many cases, knowning that the // shape variable is non-negative allows for simpler // expressions for dynamic shapes. - analyzer_->MarkGlobalNonNegValue(shape_var); + auto context = analyzer_->GetConstraintContext(); + context->Append(sym::PredicateConstraint(shape_var >= 0)); + analyzer_->SyncConstraintContext(context); } else { const PrimExpr& old_shape_expr = (*it).second; TVM_FFI_ICHECK(old_shape_expr.same_as(shape_expr) || @@ -265,9 +267,9 @@ class BlockBuilderImpl : public BlockBuilderObj { void EmitNormalized(Binding binding) final { BindingBlockFrame* cur_frame = CurrentBindingBlockFrame(); - if (const auto* var_binding = binding.as()) { + if (const auto* var_binding = binding.as()) { if (!cur_frame->is_dataflow) { - TVM_FFI_ICHECK(!var_binding->var.as()) + TVM_FFI_ICHECK(!var_binding->var.as()) << "Cannot emit dataflow var in non-dataflow block"; } // normalized check @@ -275,9 +277,9 @@ class BlockBuilderImpl : public BlockBuilderObj { TVM_FFI_ICHECK(!var_binding->value->ty.as().has_value()); cur_frame->bindings.push_back(binding); binding_table_.insert_or_assign(var_binding->var, var_binding->value); - } else if (const auto* match_cast = binding.as()) { + } else if (const auto* match_cast = binding.as()) { if (!cur_frame->is_dataflow) { - TVM_FFI_ICHECK(!match_cast->var.as()) + TVM_FFI_ICHECK(!match_cast->var.as()) << "Cannot emit dataflow var in non-dataflow block"; } // normalized check @@ -408,7 +410,7 @@ class BlockBuilderImpl : public BlockBuilderObj { } name_hint = GetUniqueName(name_hint); return is_dataflow ? DataflowVar(name_hint, /*ty_annotation=*/std::nullopt) - : Var(name_hint, /*ty_annotation=*/std::nullopt); + : Var(name_hint, Type::Missing()); } private: @@ -476,8 +478,8 @@ class BlockBuilderImpl : public BlockBuilderObj { } private: - void VisitType_(const TensorTypeNode* op) final { - if (const auto* shape_expr = op->shape.as()) { + void VisitType_(const TensorTypeObj* op) final { + if (const auto* shape_expr = op->shape.as()) { for (const PrimExpr& s : shape_expr->values) { // Only collect single var defined shape. Ignore something like `R.Tensor((m + 1, n + 1)) if (auto var = s.as()) { @@ -487,7 +489,7 @@ class BlockBuilderImpl : public BlockBuilderObj { } } - void VisitType_(const ShapeTypeNode* op) final { + void VisitType_(const ShapeTypeObj* op) final { for (const PrimExpr& s : op->values.value_or(ffi::Array())) { // Only collect single var defined shape. Ignore something like `R.Shape((m + 1, n + 1)) if (auto var = s.as()) { @@ -527,9 +529,9 @@ class Normalizer : public BlockBuilderImpl, private ExprFunctorIsInstance()) { + if (!normalized->IsInstance()) { TVM_FFI_ICHECK(!normalized->ty.as().has_value()) - << "The ty of an Expr except OpNode after " + << "The ty of an Expr except OpObj after " "normalization must not be missing. However, this Expr does not have ty: " << normalized; } @@ -568,16 +570,16 @@ class Normalizer : public BlockBuilderImpl, private ExprFunctor(static_cast(op)); + Expr expr = ffi::GetRef(static_cast(op)); if (expr.as() || expr.as()) return expr; return ExprFunctor::VisitExprDefault_(op); } @@ -591,7 +593,7 @@ class Normalizer : public BlockBuilderImpl, private ExprFunctor(var); } - Expr VisitExpr_(const VarNode* var_ptr) final { + Expr VisitExpr_(const VarObj* var_ptr) final { auto var = VisitVar_(var_ptr); if (HasVoidType(var)) { return VisitExpr(Tuple(ffi::Array{})); @@ -600,7 +602,7 @@ class Normalizer : public BlockBuilderImpl, private ExprFunctor(var); } + Expr VisitExpr_(const DataflowVarObj* var) final { return VisitVar_(var); } Expr VisitExpr(const Expr& expr) final { // lookup normalize map @@ -614,7 +616,7 @@ class Normalizer : public BlockBuilderImpl, private ExprFunctor new_fields; @@ -636,7 +638,7 @@ class Normalizer : public BlockBuilderImpl, private ExprFunctorVisitWithNewScope(op->body, op->params); if (new_body.same_as(op->body)) { @@ -646,8 +648,8 @@ class Normalizer : public BlockBuilderImpl, private ExprFunctorop.as() || !op->op->ty.as(), TypeError) + Expr VisitExpr_(const CallObj* op) final { + TVM_FFI_CHECK(op->op.as() || !op->op->ty.as(), TypeError) << "Ordinary Relax calls cannot invoke a native TIRx function; " << "use R.call_tir for destination passing or R.call_tir_packed for a direct result"; Expr new_op = this->NormalizeArgument(op->op); @@ -689,7 +691,7 @@ class Normalizer : public BlockBuilderImpl, private ExprFunctor new_blocks; for (BindingBlock block : op->blocks) { @@ -738,7 +740,7 @@ class Normalizer : public BlockBuilderImpl, private ExprFunctorNormalizeArgument(op->cond); Expr new_true = this->VisitWithNewScope(op->true_branch); Expr new_false = this->VisitWithNewScope(op->false_branch); @@ -759,7 +761,7 @@ class Normalizer : public BlockBuilderImpl, private ExprFunctorNormalizeArgument(op->tuple); TupleGetItem node = new_tuple.same_as(op->tuple) ? ffi::GetRef(op) @@ -776,10 +778,10 @@ class Normalizer : public BlockBuilderImpl, private ExprFunctor()) { + if (auto* var_binding = binding.as()) { return this->VisitVarBinding(ffi::GetRef(var_binding)); } else { - auto* match_cast = binding.as(); + auto* match_cast = binding.as(); TVM_FFI_ICHECK(match_cast) << "Unsupported binding type: " << binding->GetTypeKey(); return this->VisitMatchCast(ffi::GetRef(match_cast)); } @@ -808,7 +810,7 @@ class Normalizer : public BlockBuilderImpl, private ExprFunctor()) { + if (block.as()) { this->BeginDataflowBlock(); } else { this->BeginBindingBlock(); @@ -832,19 +834,19 @@ class Normalizer : public BlockBuilderImpl, private ExprFunctorop.as()) { + if (auto* op_ptr = call->op.as()) { // Case 1: the op field is a primitive op, look up FInferTypeWithBuilder attribute Op op = ffi::GetRef(op_ptr); bool is_dist_op = false; for (const auto& arg : call->args) { - if (arg->ty.as()) { + if (arg->ty.as()) { is_dist_op = true; break; } } if (is_dist_op) { for (const auto& arg : call->args) { - TVM_FFI_ICHECK(!arg->ty.as()) + TVM_FFI_ICHECK(!arg->ty.as()) << "Distributed operator must take DTensor instead of Tensor as input"; } TVM_FFI_ICHECK(op_map_dist_infer_ty.count(op)) @@ -900,7 +902,7 @@ class Normalizer : public BlockBuilderImpl, private ExprFunctor Expr { // SeqExpr do not need to prepare for normalization. - if (expr.as()) { + if (expr.as()) { return this->VisitExpr(expr); } else { this->BeginBindingBlock(); @@ -908,7 +910,7 @@ class Normalizer : public BlockBuilderImpl, private ExprFunctorEndBlock(); // "New scopes" (function bodies, if/else clauses) must be wrapped in seq exprs. // Don't wrap if it's already a seq and there are no bindings to add - if (post.as() && prologue->bindings.empty()) { + if (post.as() && prologue->bindings.empty()) { return post; } ffi::Array bindings; @@ -933,13 +935,13 @@ class Normalizer : public BlockBuilderImpl, private ExprFunctor ret; bool changed = false; for (const BindingBlock& block : blocks) { - bool is_dataflow = block->IsInstance(); + bool is_dataflow = block->IsInstance(); ffi::Array current; for (const Binding& binding : block->bindings) { Expr value = [&]() -> Expr { - if (const auto* var_binding = binding.as()) { + if (const auto* var_binding = binding.as()) { return var_binding->value; - } else if (const auto* match_cast = binding.as()) { + } else if (const auto* match_cast = binding.as()) { return match_cast->value; } else { TVM_FFI_THROW(InternalError) << "Unknown binding type: " << binding->GetTypeKey(); @@ -952,14 +954,14 @@ class Normalizer : public BlockBuilderImpl, private ExprFunctor()) { + if (auto seq = value.as()) { changed = true; ret.push_back(is_dataflow ? DataflowBlock(current) : BindingBlock(current)); current = {}; // We do not need to flatten recursively because the normalizer will have normalized // and thus flattened the inner SeqExprs already for (const BindingBlock& block : seq->blocks) { - if (is_dataflow && !block->IsInstance()) { + if (is_dataflow && !block->IsInstance()) { // A DataflowBlock occurring within a non-DataflowBlock // usually is an error, resulting from return of a // `BindingBlock`. However, it may still be well-formed @@ -993,9 +995,9 @@ class Normalizer : public BlockBuilderImpl, private ExprFunctor()) { + if (const auto* var_binding = binding.as()) { current.push_back(VarBinding(var_binding->var, seq->body)); - } else if (const auto* match_cast = binding.as()) { + } else if (const auto* match_cast = binding.as()) { current.push_back(MatchCast(match_cast->var, seq->body, match_cast->ty)); } else { TVM_FFI_THROW(InternalError) << "Unknown binding type: " << binding->GetTypeKey(); @@ -1024,13 +1026,13 @@ class Normalizer : public BlockBuilderImpl, private ExprFunctortype_index() == block->type_index()) { // Case 2. Merge with previous block if possible BindingBlock merged(ffi::UnsafeInit{}); - // NOTE: should check DataflowBlockNode first. - if (const auto* dataflow_block = ret.back().as()) { - auto n = ffi::make_object(*dataflow_block); + // NOTE: should check DataflowBlockObj first. + if (const auto* dataflow_block = ret.back().as()) { + auto n = ffi::make_object(*dataflow_block); n->bindings.insert(n->bindings.end(), block->bindings.begin(), block->bindings.end()); merged = DataflowBlock(ffi::UnsafeInit{}, std::move(n)); - } else if (const auto* binding_block = ret.back().as()) { - auto n = ffi::make_object(*binding_block); + } else if (const auto* binding_block = ret.back().as()) { + auto n = ffi::make_object(*binding_block); n->bindings.insert(n->bindings.end(), block->bindings.begin(), block->bindings.end()); merged = BindingBlock(ffi::UnsafeInit{}, std::move(n)); } else { diff --git a/src/relax/ir/dataflow_block_rewriter.cc b/src/relax/ir/dataflow_block_rewriter.cc index 373ad33b3587..9f7b9caf4b92 100644 --- a/src/relax/ir/dataflow_block_rewriter.cc +++ b/src/relax/ir/dataflow_block_rewriter.cc @@ -46,14 +46,14 @@ namespace relax { class MatcherUseDefAnalysis : public relax::ExprVisitor { public: - std::vector vars; - std::map> def2use; + std::vector vars; + std::map> def2use; // caller -> callee table. - std::map> caller2callees; + std::map> caller2callees; - const VarNode* cur_user_ = nullptr; + const VarObj* cur_user_ = nullptr; - void VisitBinding_(const VarBindingNode* binding) override { + void VisitBinding_(const VarBindingObj* binding) override { // init cur_user_ = binding->var.get(); this->VisitVarDef(binding->var); @@ -61,10 +61,10 @@ class MatcherUseDefAnalysis : public relax::ExprVisitor { cur_user_ = nullptr; } - void VisitExpr_(const VarNode* op) override { + void VisitExpr_(const VarObj* op) override { if (nullptr == cur_user_) return; - auto check_and_push = [](std::vector& vec, const VarNode* var) { + auto check_and_push = [](std::vector& vec, const VarObj* var) { if (std::find(vec.begin(), vec.end(), var) == vec.end()) { vec.push_back(var); } @@ -78,13 +78,13 @@ class MatcherUseDefAnalysis : public relax::ExprVisitor { }; struct PNode { - const DFPatternNode* ptr; + const DFPatternObj* ptr; std::vector&>> children; std::vector&>> parents; }; struct RNode { - const VarNode* ptr; + const VarObj* ptr; std::vector children; std::vector parents; }; @@ -95,7 +95,7 @@ struct MatchState { match_r_p[r] = p; } - void add(const DFConstraintNode* constraint) { validated_constraints_.insert(constraint); } + void add(const DFConstraintObj* constraint) { validated_constraints_.insert(constraint); } void add(MatchState&& other) { match_p_r.merge(std::move(other.match_p_r)); @@ -103,31 +103,31 @@ struct MatchState { validated_constraints_.merge(other.validated_constraints_); } - const VarNode* matched(const PNode* p) const { + const VarObj* matched(const PNode* p) const { if (auto it = match_p_r.find(p); it != match_p_r.end()) { return it->second->ptr; } return nullptr; } - const DFPatternNode* matched(const RNode* r) const { + const DFPatternObj* matched(const RNode* r) const { if (auto it = match_r_p.find(r); it != match_r_p.end()) { return it->second->ptr; } return nullptr; } - const VarNode* matched(const PNode& p) const { return matched(&p); } - const DFPatternNode* matched(const RNode& r) const { return matched(&r); } + const VarObj* matched(const PNode& p) const { return matched(&p); } + const DFPatternObj* matched(const RNode& r) const { return matched(&r); } - bool is_validated(const DFConstraintNode* constraint) const { + bool is_validated(const DFConstraintObj* constraint) const { return validated_constraints_.count(constraint); } private: std::unordered_map match_p_r; std::unordered_map match_r_p; - std::unordered_set validated_constraints_; + std::unordered_set validated_constraints_; }; /** @@ -189,12 +189,12 @@ 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::unordered_map& pattern2node, const std::vector& validation_constraints, sym::AnalyzerObj* analyzer) { MatchState new_match; - std::function(const DFPatternNode*)> query_match_state = - [&pattern2node, ¤t_match](const DFPatternNode* pattern) -> ffi::Optional { + std::function(const DFPatternObj*)> query_match_state = + [&pattern2node, ¤t_match](const DFPatternObj* pattern) -> ffi::Optional { auto it = pattern2node.find(pattern); TVM_FFI_ICHECK(it != pattern2node.end()) << "DFConstraint attempted to access DFPattern " << ffi::GetRef(pattern) @@ -212,7 +212,7 @@ static std::optional TryValidate( auto [necessary_condition, is_sufficient] = constraint->AsCondition(query_match_state); necessary_condition = analyzer->Simplify(necessary_condition); - const auto* known = necessary_condition.as(); + const auto* known = necessary_condition.as(); if (known && known->value != 0 && is_sufficient) { // The condition passes, and the expression provided is both @@ -241,8 +241,8 @@ static std::optional TryValidate( static std::optional MatchTree( const MatchState& current_match, size_t current_root_idx, - const std::unordered_map& pattern2node, - const std::unordered_map& var2node, DFPatternMatcher* matcher, + 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, sym::AnalyzerObj* analyzer) { auto get_next_root = [&](size_t root_idx) -> const PNode* { @@ -303,14 +303,14 @@ ffi::Optional> MatchGraph(const PatternContext& ctx, } // First construct a graph of PNode and RNode. - std::unordered_map var2node; + std::unordered_map var2node; var2node.reserve(bindings.size()); - for (const VarNode* cur_var : ud_analysis.vars) { + for (const VarObj* cur_var : ud_analysis.vars) { const auto& uses = ud_analysis.def2use.at(cur_var); RNode& cur_node = var2node[cur_var]; cur_node.ptr = cur_var; - for (const VarNode* use : uses) { + for (const VarObj* use : uses) { auto& use_node = var2node[use]; use_node.ptr = use; cur_node.children.push_back(&use_node); @@ -318,7 +318,7 @@ ffi::Optional> MatchGraph(const PatternContext& ctx, } } - std::unordered_map pattern2node; + std::unordered_map pattern2node; pattern2node.reserve(ctx->edge_constraints.size()); for (const auto& def_pattern : ctx->src_ordered) { diff --git a/src/relax/ir/dataflow_expr_rewriter.cc b/src/relax/ir/dataflow_expr_rewriter.cc index 6446ccff265d..161247d5ca53 100644 --- a/src/relax/ir/dataflow_expr_rewriter.cc +++ b/src/relax/ir/dataflow_expr_rewriter.cc @@ -51,7 +51,7 @@ class GlobalVarReplacer : public ExprMutator { explicit GlobalVarReplacer(ffi::Map gvar_map) : gvar_map_(gvar_map) {} using ExprMutator::VisitExpr_; - Expr VisitExpr_(const GlobalVarNode* op) override { + Expr VisitExpr_(const GlobalVarObj* op) override { auto gvar = ffi::GetRef(op); if (auto opt = gvar_map_.Get(gvar)) { gvar = opt.value(); @@ -597,7 +597,7 @@ std::optional> TupleRewriterObj::TryMatchByBindingIndex( } std::vector rewrites; - if (auto inline_tuple = rewritten.as()) { + if (auto inline_tuple = rewritten.as()) { const auto& fields = inline_tuple->fields; TVM_FFI_ICHECK_EQ(fields.size(), indices.size()) << "Expected to receive " << indices.size() << " values to replace TuplePattern with " @@ -640,14 +640,14 @@ PatternMatchingRewriter PatternMatchingRewriter::FromPattern( ffi::TypedFunction(Expr, ffi::Map)> func, ffi::Optional> additional_bindings, ffi::Map new_subroutines) { - if (auto or_pattern = pattern.as()) { + if (auto or_pattern = pattern.as()) { auto new_additional_bindings = additional_bindings.value_or({}); new_additional_bindings.push_back(pattern); return OrRewriter(PatternMatchingRewriter::FromPattern( or_pattern->left, func, new_additional_bindings, new_subroutines), PatternMatchingRewriter::FromPattern( or_pattern->right, func, new_additional_bindings, new_subroutines)); - } else if (auto tuple_pattern = pattern.as()) { + } else if (auto tuple_pattern = pattern.as()) { auto new_additional_bindings = additional_bindings.value_or({}); new_additional_bindings.push_back(pattern); // If the Tuple appears as a Relax binding, apply it first. As a @@ -667,7 +667,7 @@ PatternMatchingRewriter PatternMatchingRewriter::FromModule(IRModule mod) { << "a Relax function defining the pattern to be matched, " << "but the module did not contain a 'pattern' function."; auto base_func = mod->Lookup("pattern"); - TVM_FFI_CHECK(base_func->IsInstance(), TypeError) + TVM_FFI_CHECK(base_func->IsInstance(), TypeError) << "Expected module to contain 'pattern', " << "a Relax function defining the pattern to be matched, " << "but the 'pattern' function was of type " << base_func->GetTypeKey() << "."; @@ -679,7 +679,7 @@ PatternMatchingRewriter PatternMatchingRewriter::FromModule(IRModule mod) { << "a Relax function defining the replacement to be matched, " << "but the module did not contain a 'replacement' function."; auto base_func = mod->Lookup("replacement"); - TVM_FFI_CHECK(base_func->IsInstance(), TypeError) + TVM_FFI_CHECK(base_func->IsInstance(), TypeError) << "Expected module to contain 'replacement', " << "a Relax function defining the replacement to be made on a successful match, " << "but the 'replacement' function was of type " << base_func->GetTypeKey() << "."; @@ -718,23 +718,23 @@ PatternMatchingRewriter PatternMatchingRewriter::FromModule(IRModule mod) { if (auto var = expr.as()) { return pattern_lookup[var.value()]; - } else if (auto call = expr.as()) { + } else if (auto call = expr.as()) { auto op = make_pattern(call->op); auto args = call->args.Map(make_pattern); return CallPattern(op, args); - } else if (auto tuple = expr.as()) { + } else if (auto tuple = expr.as()) { auto fields = tuple->fields.Map(make_pattern); return TuplePattern(fields); - } else if (auto tuple_get_item = expr.as()) { + } else if (auto tuple_get_item = expr.as()) { auto tuple = make_pattern(tuple_get_item->tuple); return TupleGetItemPattern(tuple, tuple_get_item->index); } else if (auto op = expr.as()) { return ExprPattern(op.value()); - } else if (auto func = expr.as()) { + } else if (auto func = expr.as()) { return ExternFuncPattern(func->global_symbol); } else if (auto prim = expr.as()) { @@ -749,7 +749,7 @@ PatternMatchingRewriter PatternMatchingRewriter::FromModule(IRModule mod) { for (const auto& block : func_pattern->body->blocks) { for (const auto& binding : block->bindings) { auto value_pattern = make_pattern(GetBoundValue(binding)); - if (auto match_cast = binding.as()) { + if (auto match_cast = binding.as()) { value_pattern = TypePattern(value_pattern, match_cast->ty); } pattern_lookup.Set(binding->var, value_pattern); @@ -832,7 +832,7 @@ class PatternMatchingMutator : public ExprMutator { ffi::Map GetNewSubroutines() const { return new_subroutines_; } - Expr VisitExpr_(const SeqExprNode* seq) override { + Expr VisitExpr_(const SeqExprObj* seq) override { SeqExpr prev = ExprMutator::VisitExpr_(seq).as_or_throw(); ffi::StructuralEqual struct_equal; @@ -875,7 +875,7 @@ class PatternMatchingMutator : public ExprMutator { // were the last variable binding of the last block. This // simplifies the special handling of the SeqExpr's body. ffi::Optional dummy_output_var = std::nullopt; - if (!seq->body->IsInstance()) { + if (!seq->body->IsInstance()) { dummy_output_var = Var("dummy_output_var", GetType(seq->body)); VarBinding dummy_binding(dummy_output_var.value(), seq->body); @@ -903,7 +903,7 @@ class PatternMatchingMutator : public ExprMutator { auto bindings = orig_bindings.Map([&](Binding binding) -> Binding { if (auto new_expr = rewrites.variable_rewrites.Get(binding->var)) { - if (auto match_cast = binding.as()) { + if (auto match_cast = binding.as()) { return MatchCast(binding->var, new_expr.value(), match_cast->ty); } else { return VarBinding(binding->var, new_expr.value()); @@ -929,7 +929,7 @@ class PatternMatchingMutator : public ExprMutator { // Utility function to return the rewrites that should be applied // to a given block. auto get_rewrites = [&](BindingBlock block) -> ffi::Array { - if (block.as()) { + if (block.as()) { // Early return for DataflowBlock. Since neither control flow // nor impure functions are allowed within the dataflow block, // all bindings may be considered at the same time. @@ -957,8 +957,8 @@ class PatternMatchingMutator : public ExprMutator { for (const auto& binding : block->bindings) { auto value = GetBoundValue(binding); - bool is_dataflow = (!value.as()) && - (!(value.as() && IsImpureCall(value.as_or_throw()))); + bool is_dataflow = (!value.as()) && + (!(value.as() && IsImpureCall(value.as_or_throw()))); if (is_dataflow) { // This binding satisfies the dataflow constraints. collected_bindings.push_back(binding); @@ -986,7 +986,7 @@ class PatternMatchingMutator : public ExprMutator { return old_block; } - if (old_block.as()) { + if (old_block.as()) { builder_->BeginDataflowBlock(); } else { builder_->BeginBindingBlock(); @@ -995,9 +995,9 @@ class PatternMatchingMutator : public ExprMutator { for (const auto& binding : new_bindings) { auto value = builder_->Normalize(GetBoundValue(binding)); - if (binding.as()) { + if (binding.as()) { builder_->EmitNormalized(VarBinding(binding->var, value)); - } else if (auto match_cast = binding.as()) { + } else if (auto match_cast = binding.as()) { builder_->EmitNormalized(MatchCast(binding->var, value, match_cast->ty)); } else { TVM_FFI_THROW(InternalError) << "Binding must be either VarBinding or MatchCast"; diff --git a/src/relax/ir/dataflow_matcher.cc b/src/relax/ir/dataflow_matcher.cc index ce5f59ab4c3d..b0fb3caed227 100644 --- a/src/relax/ir/dataflow_matcher.cc +++ b/src/relax/ir/dataflow_matcher.cc @@ -67,7 +67,7 @@ using tvm::sym::AnalyzerObj; */ bool MatchAttrs(const Any& attrs, const ffi::Map& attributes) { // TODO(tqchen): consider lift to common utils - if (auto* dict_attrs = attrs.as()) { + if (auto* dict_attrs = attrs.as()) { for (auto kv : attributes) { auto attr_name = kv.first; auto attr_value = kv.second; @@ -117,7 +117,7 @@ Expr DFPatternMatcher::UnwrapBindings(Expr expr, const ffi::Map& var2 auto unwrap = [&](Expr expr) -> ffi::Optional { // Unwrap variables into the value to which they are bound. if (var2val.size()) { - if (const VarNode* var = expr.as()) { + if (const VarObj* var = expr.as()) { if (auto may = var2val.Get(ffi::GetRef(var))) { return may.value(); } @@ -126,7 +126,7 @@ Expr DFPatternMatcher::UnwrapBindings(Expr expr, const ffi::Map& var2 // Unwrap SeqExpr with no bindings. These can occur due to Relax // IR constraints for the bodies of Function and IfExpr nodes. - if (auto seq = expr.as()) { + if (auto seq = expr.as()) { if (seq->blocks.empty()) { return seq->body; } @@ -170,28 +170,28 @@ bool DFPatternMatcher::VisitDFPattern(const DFPattern& pattern, const Expr& expr } } -bool DFPatternMatcher::VisitDFPattern_(const OrPatternNode* op, const Expr& expr0) { +bool DFPatternMatcher::VisitDFPattern_(const OrPatternObj* op, const Expr& expr0) { auto expr = UnwrapBindings(expr0, var2val_); return VisitDFPattern(op->left, expr) || VisitDFPattern(op->right, expr); } -bool DFPatternMatcher::VisitDFPattern_(const AndPatternNode* op, const Expr& expr0) { +bool DFPatternMatcher::VisitDFPattern_(const AndPatternObj* op, const Expr& expr0) { auto expr = UnwrapBindings(expr0, var2val_); return VisitDFPattern(op->left, expr) && VisitDFPattern(op->right, expr); } -bool DFPatternMatcher::VisitDFPattern_(const NotPatternNode* op, const Expr& expr0) { +bool DFPatternMatcher::VisitDFPattern_(const NotPatternObj* op, const Expr& expr0) { auto expr = UnwrapBindings(expr0, var2val_); return !VisitDFPattern(op->reject, expr); } -bool DFPatternMatcher::VisitDFPattern_(const AttrPatternNode* attr_pattern, const Expr& expr0) { +bool DFPatternMatcher::VisitDFPattern_(const AttrPatternObj* attr_pattern, const Expr& expr0) { auto expr = UnwrapBindings(expr0, var2val_); bool matches = VisitDFPattern(attr_pattern->pattern, expr); if (!matches) return matches; - VLOG(1) << "considering AttrPatternNode at:\n" << expr; - auto attributes = attr_pattern->attrs.as()->dict; - if (const auto* op_node = expr.as()) { + VLOG(1) << "considering AttrPatternObj at:\n" << expr; + auto attributes = attr_pattern->attrs.as()->dict; + if (const auto* op_node = expr.as()) { Op op = ffi::GetRef(op_node); for (auto kv : attributes) { auto attr_name = kv.first; @@ -207,9 +207,9 @@ bool DFPatternMatcher::VisitDFPattern_(const AttrPatternNode* attr_pattern, cons matches = false; } } - } else if (auto* op = expr.as()) { + } else if (auto* op = expr.as()) { matches &= MatchAttrs(op->attrs, attributes); - } else if (auto* op = expr.as()) { + } else if (auto* op = expr.as()) { matches = true; for (auto kv : attributes) { if (matches && op->attrs->dict.count(kv.first)) { @@ -225,18 +225,18 @@ bool DFPatternMatcher::VisitDFPattern_(const AttrPatternNode* attr_pattern, cons return matches; } -bool DFPatternMatcher::VisitDFPattern_(const CallPatternNode* op, const Expr& expr0) { +bool DFPatternMatcher::VisitDFPattern_(const CallPatternObj* op, const Expr& expr0) { auto expr = UnwrapBindings(expr0, var2val_); // utilities - auto get_op_node = [](const CallPatternNode* op) -> const tvm::OpNode* { + auto get_op_node = [](const CallPatternObj* op) -> const tvm::OpObj* { if (op) { - if (auto* expr_pattern = op->op.as()) { - return expr_pattern->expr.as(); + if (auto* expr_pattern = op->op.as()) { + return expr_pattern->expr.as(); } } return nullptr; }; - auto is_pattern_op = [&get_op_node](const CallPatternNode* op, std::string op_type) { + auto is_pattern_op = [&get_op_node](const CallPatternObj* op, std::string op_type) { if (const auto* op_node = get_op_node(op)) { if (op_node->name == op_type) { return true; @@ -245,8 +245,8 @@ bool DFPatternMatcher::VisitDFPattern_(const CallPatternNode* op, const Expr& ex return false; }; auto is_expr_op = [](const Expr& expr, std::string op_type) { - if (const auto* call_node = expr.as()) { - if (const auto* op_node = call_node->op.as()) { + if (const auto* call_node = expr.as()) { + if (const auto* op_node = call_node->op.as()) { if (op_node->name == op_type) { return true; } @@ -257,7 +257,7 @@ bool DFPatternMatcher::VisitDFPattern_(const CallPatternNode* op, const Expr& ex // logic auto watermark = matched_nodes_.size(); - if (const auto* call_node = expr.as()) { + if (const auto* call_node = expr.as()) { auto matches_op = VisitDFPattern(op->op, call_node->op); if (matches_op) { auto watermark2 = matched_nodes_.size(); @@ -286,7 +286,7 @@ bool DFPatternMatcher::VisitDFPattern_(const CallPatternNode* op, const Expr& ex if (match_args(op->args, call_node->args.begin(), call_node->args.end())) return true; // Commutative Matching. - if (const OpNode* op_node = call_node->op.as()) { + if (const OpObj* op_node = call_node->op.as()) { if ((op_node->name == "relax.add") || (op_node->name == "relax.multiply")) { if (match_args(op->args, call_node->args.rbegin(), call_node->args.rend())) { return true; @@ -297,7 +297,7 @@ bool DFPatternMatcher::VisitDFPattern_(const CallPatternNode* op, const Expr& ex ClearMap(watermark); // associate divide/multiply if (is_pattern_op(op, "relax.divide")) { - if (const auto* arg_node = op->args[0].as()) { + if (const auto* arg_node = op->args[0].as()) { if (is_pattern_op(arg_node, "relax.multiply") && is_expr_op(expr, "relax.multiply") && (is_expr_op(call_node->args[0], "relax.divide") || is_expr_op(call_node->args[1], "relax.divide"))) { @@ -319,7 +319,7 @@ bool DFPatternMatcher::VisitDFPattern_(const CallPatternNode* op, const Expr& ex if (is_pattern_op(op, "relax.multiply")) { // associate multiply/divide for (size_t arg_id = 0; arg_id < 2; ++arg_id) { - if (auto* arg_node = op->args[arg_id].as()) { + if (auto* arg_node = op->args[arg_id].as()) { if (is_pattern_op(arg_node, "relax.divide") && is_expr_op(expr, "relax.divide") && (is_expr_op(call_node->args[0], "relax.multiply") || is_expr_op(call_node->args[1], "relax.multiply"))) { @@ -335,15 +335,15 @@ bool DFPatternMatcher::VisitDFPattern_(const CallPatternNode* op, const Expr& ex return false; } -bool DFPatternMatcher::VisitDFPattern_(const ExprPatternNode* op, const Expr& expr0) { +bool DFPatternMatcher::VisitDFPattern_(const ExprPatternObj* op, const Expr& expr0) { auto expr = UnwrapBindings(expr0, var2val_); return ffi::StructuralEqual()(op->expr, expr); } -bool DFPatternMatcher::VisitDFPattern_(const FunctionPatternNode* op, const Expr& expr0) { +bool DFPatternMatcher::VisitDFPattern_(const FunctionPatternObj* op, const Expr& expr0) { auto expr = UnwrapBindings(expr0, var2val_); bool matches = false; - if (const auto* func = expr.as()) { + if (const auto* func = expr.as()) { matches = true; if (op->params.defined()) { size_t i = 0; @@ -363,19 +363,19 @@ bool DFPatternMatcher::VisitDFPattern_(const FunctionPatternNode* op, const Expr return matches; } -bool DFPatternMatcher::VisitDFPattern_(const TupleGetItemPatternNode* op, const Expr& expr0) { +bool DFPatternMatcher::VisitDFPattern_(const TupleGetItemPatternObj* op, const Expr& expr0) { auto expr = UnwrapBindings(expr0, var2val_); - if (const auto* tuple_get_item_node = expr.as()) { + if (const auto* tuple_get_item_node = expr.as()) { return (op->index == -1 || op->index == tuple_get_item_node->index) && VisitDFPattern(op->tuple, tuple_get_item_node->tuple); } return false; } -bool DFPatternMatcher::VisitDFPattern_(const TuplePatternNode* op, const Expr& expr0) { +bool DFPatternMatcher::VisitDFPattern_(const TuplePatternObj* op, const Expr& expr0) { auto expr = UnwrapBindings(expr0, var2val_); bool matches = false; - if (const auto* tuple_node = expr.as()) { + if (const auto* tuple_node = expr.as()) { matches = true; if (op->fields.size() == tuple_node->fields.size()) { size_t i = 0; @@ -413,10 +413,10 @@ bool DFPatternMatcher::TryUnorderedMatch(size_t idx, const tvm::ffi::Array()) { + if (const auto* tuple_node = expr.as()) { if (op->fields.size() == tuple_node->fields.size()) { constexpr int8_t kUnknown = -1; TVM_FFI_ICHECK_LE(op->fields.size(), std::numeric_limits::max()) @@ -435,7 +435,7 @@ bool DFPatternMatcher::VisitDFPattern_(const UnorderedTuplePatternNode* op, cons static bool ShapeConstraintsConsistent(const PrimExpr& condition) { std::vector constraints{condition}; for (size_t i = 0; i < constraints.size();) { - if (const auto* conjunction = constraints[i].as()) { + if (const auto* conjunction = constraints[i].as()) { PrimExpr lhs = conjunction->a; PrimExpr rhs = conjunction->b; constraints[i] = std::move(lhs); @@ -456,16 +456,18 @@ static bool ShapeConstraintsConsistent(const PrimExpr& condition) { for (const PrimExpr& constraint : constraints) { PrimExpr simplified = analyzer->Simplify(constraint, simplify_steps); if (prim::IsZero(simplified)) return false; - if (const auto* equality = simplified.as()) { + if (const auto* equality = simplified.as()) { PrimExpr lhs = equality->a; PrimExpr rhs = equality->b; - if (!lhs.as()) std::swap(lhs, rhs); - if (const auto* var = lhs.as(); + if (!lhs.as()) std::swap(lhs, rhs); + if (const auto* var = lhs.as(); var && SideEffect(rhs) <= CallEffectKind::kPure) { auto variables = tvm::UndefinedVars(rhs); if (std::none_of(variables.begin(), variables.end(), [var](const tvm::Var& used) { return used.get() == var; })) { - analyzer->Bind(ffi::GetRef(var), rhs); + auto context = analyzer->GetConstraintContext(); + context->Append(sym::BindingConstraint(ffi::GetRef(var), rhs)); + analyzer->SyncConstraintContext(context); changed = true; } } @@ -485,13 +487,18 @@ static bool ShapeConstraintsConsistent(const PrimExpr& condition) { for (size_t j = 0; j < constraints.size(); ++j) { if (i != j) others = others && constraints[j]; } - With scope(analyzer, others); - if (prim::IsZero(analyzer->Simplify(constraints[i], simplify_steps))) return false; + auto saved_context = analyzer->GetConstraintContext(); + auto context = saved_context->Snapshot(); + context->Append(sym::PredicateConstraint(others)); + analyzer->SyncConstraintContext(context); + bool inconsistent = prim::IsZero(analyzer->Simplify(constraints[i], simplify_steps)); + analyzer->SyncConstraintContext(saved_context); + if (inconsistent) return false; } return true; } -bool DFPatternMatcher::VisitDFPattern_(const TypePatternNode* op, const Expr& expr0) { +bool DFPatternMatcher::VisitDFPattern_(const TypePatternObj* op, const Expr& expr0) { if (!VisitDFPattern(op->pattern, expr0)) { return false; } @@ -500,13 +507,13 @@ bool DFPatternMatcher::VisitDFPattern_(const TypePatternNode* op, const Expr& ex auto expr_ty = GetType(expr); PrimExpr new_constraint = TypeBaseCheckPrecondition(op->ty, expr_ty); - if (auto* as_int = new_constraint.as()) { + if (auto* as_int = new_constraint.as()) { return static_cast(as_int->value); } symbolic_expr_condition_ = analyzer_->Simplify(symbolic_expr_condition_ && new_constraint); - if (auto* as_int = symbolic_expr_condition_.as()) { + if (auto* as_int = symbolic_expr_condition_.as()) { return static_cast(as_int->value); } else { return ShapeConstraintsConsistent(symbolic_expr_condition_); @@ -521,10 +528,10 @@ static bool ShapeEqual(AnalyzerObj* analyzer, const ffi::Array& lhs, return true; } -bool DFPatternMatcher::VisitDFPattern_(const ShapePatternNode* op, const Expr& expr) { +bool DFPatternMatcher::VisitDFPattern_(const ShapePatternObj* op, const Expr& expr) { // no need to jump, as var.shape == value.shape - if (const auto* tinfo = GetTypeAs(expr)) { - if (const ShapeExprNode* shape_expr = tinfo->shape.as()) { + if (const auto* tinfo = GetTypeAs(expr)) { + if (const ShapeExprObj* shape_expr = tinfo->shape.as()) { return ShapeEqual(analyzer_.get(), op->shape, shape_expr->values) && VisitDFPattern(op->pattern, expr); } @@ -532,8 +539,8 @@ bool DFPatternMatcher::VisitDFPattern_(const ShapePatternNode* op, const Expr& e return false; } -std::tuple SameShapeConstraintNode::AsCondition( - std::function(const DFPatternNode*)> match_state) const { +std::tuple SameShapeConstraintObj::AsCondition( + std::function(const DFPatternObj*)> match_state) const { ffi::Optional> expected_shape; bool all_shapes_defined = true; @@ -545,9 +552,9 @@ std::tuple SameShapeConstraintNode::AsCondition( auto var = opt_var.value(); auto opt_var_shape = [&]() -> ffi::Optional> { auto ty = GetType(var); - if (auto tensor = ty.as()) { + if (auto tensor = ty.as()) { return tensor->GetShape(); - } else if (auto shape_expr = ty.as()) { + } else if (auto shape_expr = ty.as()) { return shape_expr->values; } else { return std::nullopt; @@ -595,55 +602,55 @@ std::tuple SameShapeConstraintNode::AsCondition( return {all_dimensions_equal, all_shapes_defined}; } -bool DFPatternMatcher::VisitDFPattern_(const PrimArrPatternNode* op, const Expr& expr0) { +bool DFPatternMatcher::VisitDFPattern_(const PrimArrPatternObj* op, const Expr& expr0) { auto expr = UnwrapBindings(expr0, var2val_); - if (const ShapeExprNode* shape_expr = expr.as()) + if (const ShapeExprObj* shape_expr = expr.as()) return ShapeEqual(analyzer_.get(), op->fields, shape_expr->values); return false; } -bool DFPatternMatcher::VisitDFPattern_(const DataTypePatternNode* op, const Expr& expr) { +bool DFPatternMatcher::VisitDFPattern_(const DataTypePatternObj* op, const Expr& expr) { // no need to jump, as var.dtype == value.dtype - auto expr_ty = expr.as()->ty; - if (const TensorTypeNode* tensor_ty = expr_ty.as()) { + auto expr_ty = expr.as()->ty; + if (const TensorTypeObj* tensor_ty = expr_ty.as()) { return op->dtype == tensor_ty->dtype.value()->dtype && VisitDFPattern(op->pattern, expr); } return false; } -bool DFPatternMatcher::VisitDFPattern_(const VarPatternNode* op, const Expr& expr) { +bool DFPatternMatcher::VisitDFPattern_(const VarPatternObj* op, const Expr& expr) { // We don't jump for var pattern, as there's no need to access its value to judge it. - if (const auto* var_node = expr.as()) { + if (const auto* var_node = expr.as()) { // "" means any name. return "" == op->name_hint() || op->name_hint() == var_node->name; } return false; } -bool DFPatternMatcher::VisitDFPattern_(const ExternFuncPatternNode* op, const Expr& expr0) { +bool DFPatternMatcher::VisitDFPattern_(const ExternFuncPatternObj* op, const Expr& expr0) { auto expr = UnwrapBindings(expr0, var2val_); - if (const auto* extern_fn = expr.as()) { + if (const auto* extern_fn = expr.as()) { return "" == op->global_symbol() || op->global_symbol() == extern_fn->global_symbol; } return false; } -bool DFPatternMatcher::VisitDFPattern_(const ConstantPatternNode* op, const Expr& expr0) { +bool DFPatternMatcher::VisitDFPattern_(const ConstantPatternObj* op, const Expr& expr0) { // constants can be binded to relax.Var as well. auto expr = UnwrapBindings(expr0, var2val_); - auto* constant = expr.as(); + auto* constant = expr.as(); return constant && constant->value.as(); } -bool DFPatternMatcher::VisitDFPattern_(const DataflowVarPatternNode* op, const Expr& expr) { +bool DFPatternMatcher::VisitDFPattern_(const DataflowVarPatternObj* op, const Expr& expr) { // DataflowVar is inherented from Var, so dispatch it to VarPattern. - return expr->IsInstance() && - VisitDFPattern_(static_cast(op), expr); + return expr->IsInstance() && + VisitDFPattern_(static_cast(op), expr); } -bool DFPatternMatcher::VisitDFPattern_(const GlobalVarPatternNode* op, const Expr& expr) { +bool DFPatternMatcher::VisitDFPattern_(const GlobalVarPatternObj* op, const Expr& expr) { // GlobalVarPattern is not inherited from Var, so we need to handle it separately. - if (const auto* var_node = expr.as()) { + if (const auto* var_node = expr.as()) { std::string pat = std::string(op->name_hint()); std::string var_name = std::string(var_node->name_hint); return pat.empty() || var_name.find(pat) != std::string::npos; @@ -651,7 +658,7 @@ bool DFPatternMatcher::VisitDFPattern_(const GlobalVarPatternNode* op, const Exp return false; } -bool DFPatternMatcher::VisitDFPattern_(const WildcardPatternNode* op, const Expr& expr) { +bool DFPatternMatcher::VisitDFPattern_(const WildcardPatternObj* op, const Expr& expr) { return true; } diff --git a/src/relax/ir/dataflow_matcher.h b/src/relax/ir/dataflow_matcher.h index 98a4adab749a..c27dad457a0a 100644 --- a/src/relax/ir/dataflow_matcher.h +++ b/src/relax/ir/dataflow_matcher.h @@ -51,27 +51,27 @@ class DFPatternMatcher : public DFPatternFunctor patterns, diff --git a/src/relax/ir/dataflow_pattern.cc b/src/relax/ir/dataflow_pattern.cc index d67781babc79..9a3d4ab83588 100644 --- a/src/relax/ir/dataflow_pattern.cc +++ b/src/relax/ir/dataflow_pattern.cc @@ -61,56 +61,56 @@ struct PatternReprPrinterHelper { } ExternFuncPattern::ExternFuncPattern(ffi::String global_symbol) : DFPattern(ffi::UnsafeInit{}) { - ffi::ObjectPtr n = ffi::make_object(); + ffi::ObjectPtr n = ffi::make_object(); n->global_symbol_ = std::move(global_symbol); data_ = std::move(n); } TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - ExternFuncPatternNode::RegisterReflection(); + ExternFuncPatternObj::RegisterReflection(); refl::GlobalDef().def("relax.dpl.ExternFuncPattern", [](ffi::String global_symbol) { return ExternFuncPattern(global_symbol); }); } -RELAX_PATTERN_PRINTER_DEF(ExternFuncPatternNode, [](auto p, auto node) { +RELAX_PATTERN_PRINTER_DEF(ExternFuncPatternObj, [](auto p, auto node) { p->stream << "ExternFuncPattern(" << node->global_symbol() << ")"; }); VarPattern::VarPattern(ffi::String name_hint) : DFPattern(ffi::UnsafeInit{}) { - ffi::ObjectPtr n = ffi::make_object(); + ffi::ObjectPtr n = ffi::make_object(); n->name = std::move(name_hint); data_ = std::move(n); } TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - VarPatternNode::RegisterReflection(); + VarPatternObj::RegisterReflection(); refl::GlobalDef().def("relax.dpl.VarPattern", [](ffi::String name_hint) { return VarPattern(name_hint); }); } -RELAX_PATTERN_PRINTER_DEF(VarPatternNode, [](auto p, auto node) { +RELAX_PATTERN_PRINTER_DEF(VarPatternObj, [](auto p, auto node) { p->stream << "VarPattern(" << node->name_hint() << ")"; }); DataflowVarPattern::DataflowVarPattern(ffi::String name_hint) : DFPattern(ffi::UnsafeInit{}) { - ffi::ObjectPtr n = ffi::make_object(); + ffi::ObjectPtr n = ffi::make_object(); n->name = std::move(name_hint); data_ = std::move(n); } TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - DataflowVarPatternNode::RegisterReflection(); + DataflowVarPatternObj::RegisterReflection(); refl::GlobalDef().def("relax.dpl.DataflowVarPattern", [](ffi::String name_hint) { return DataflowVarPattern(name_hint); }); } -RELAX_PATTERN_PRINTER_DEF(DataflowVarPatternNode, [](auto p, auto node) { +RELAX_PATTERN_PRINTER_DEF(DataflowVarPatternObj, [](auto p, auto node) { p->stream << "DataflowVarPattern(" << node->name_hint() << ")"; }); GlobalVarPattern::GlobalVarPattern(ffi::String name_hint) : DFPattern(ffi::UnsafeInit{}) { - ffi::ObjectPtr n = ffi::make_object(); + ffi::ObjectPtr n = ffi::make_object(); n->name = std::move(name_hint); data_ = std::move(n); } @@ -120,36 +120,36 @@ TVM_FFI_STATIC_INIT_BLOCK() { refl::GlobalDef().def("relax.dpl.GlobalVarPattern", [](ffi::String name_hint) { return GlobalVarPattern(name_hint); }); } -RELAX_PATTERN_PRINTER_DEF(GlobalVarPatternNode, [](auto p, auto node) { +RELAX_PATTERN_PRINTER_DEF(GlobalVarPatternObj, [](auto p, auto node) { p->stream << "GlobalVarPattern(" << node->name_hint() << ")"; }); ExprPattern::ExprPattern(Expr expr) : DFPattern(ffi::UnsafeInit{}) { - ffi::ObjectPtr n = ffi::make_object(std::move(expr)); + ffi::ObjectPtr n = ffi::make_object(std::move(expr)); data_ = std::move(n); } TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - ExprPatternNode::RegisterReflection(); + ExprPatternObj::RegisterReflection(); refl::GlobalDef().def("relax.dpl.ExprPattern", [](Expr e) { return ExprPattern(e); }); } -RELAX_PATTERN_PRINTER_DEF(ExprPatternNode, [](auto p, auto node) { p->Print(node->expr); }); +RELAX_PATTERN_PRINTER_DEF(ExprPatternObj, [](auto p, auto node) { p->Print(node->expr); }); TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - ConstantPatternNode::RegisterReflection(); + ConstantPatternObj::RegisterReflection(); refl::GlobalDef().def("relax.dpl.ConstantPattern", []() { - auto c = ConstantPattern(ffi::UnsafeInit{}, ffi::make_object()); + auto c = ConstantPattern(ffi::UnsafeInit{}, ffi::make_object()); return c; }); } -RELAX_PATTERN_PRINTER_DEF(ConstantPatternNode, +RELAX_PATTERN_PRINTER_DEF(ConstantPatternObj, [](auto p, auto node) { p->stream << "ConstantPattern()"; }); CallPattern::CallPattern(DFPattern op, ffi::Array args, bool varg_default_wildcard) : DFPattern(ffi::UnsafeInit{}) { - ffi::ObjectPtr n = ffi::make_object(std::move(op)); + ffi::ObjectPtr n = ffi::make_object(std::move(op)); n->args = std::move(args); n->varg_default_wildcard = varg_default_wildcard; data_ = std::move(n); @@ -157,13 +157,13 @@ CallPattern::CallPattern(DFPattern op, ffi::Array args, bool varg_def TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - CallPatternNode::RegisterReflection(); + CallPatternObj::RegisterReflection(); refl::GlobalDef().def("relax.dpl.CallPattern", [](DFPattern op, ffi::Array args, bool varg_default_wildcard) { return CallPattern(op, args, varg_default_wildcard); }); } -RELAX_PATTERN_PRINTER_DEF(CallPatternNode, [](auto p, auto node) { +RELAX_PATTERN_PRINTER_DEF(CallPatternObj, [](auto p, auto node) { p->stream << node->op << "("; for (size_t i = 0; i < node->args.size(); ++i) { if (i != 0) p->stream << ", "; @@ -177,186 +177,186 @@ RELAX_PATTERN_PRINTER_DEF(CallPatternNode, [](auto p, auto node) { }); PrimArrPattern::PrimArrPattern(ffi::Array arr) : DFPattern(ffi::UnsafeInit{}) { - ffi::ObjectPtr n = ffi::make_object(); + ffi::ObjectPtr n = ffi::make_object(); n->fields = std::move(arr); data_ = std::move(n); } TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - PrimArrPatternNode::RegisterReflection(); + PrimArrPatternObj::RegisterReflection(); refl::GlobalDef().def("relax.dpl.PrimArrPattern", [](ffi::Array arr) { return PrimArrPattern(std::move(arr)); }); } -RELAX_PATTERN_PRINTER_DEF(PrimArrPatternNode, [](auto p, auto node) { +RELAX_PATTERN_PRINTER_DEF(PrimArrPatternObj, [](auto p, auto node) { p->stream << "PrimArrPattern(" << node->fields << ")"; }); FunctionPattern::FunctionPattern(ffi::Array params, DFPattern body) : DFPattern(ffi::UnsafeInit{}) { - ffi::ObjectPtr n = ffi::make_object(std::move(body)); + ffi::ObjectPtr n = ffi::make_object(std::move(body)); n->params = std::move(params); data_ = std::move(n); } TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - FunctionPatternNode::RegisterReflection(); + FunctionPatternObj::RegisterReflection(); refl::GlobalDef().def( "relax.dpl.FunctionPattern", [](ffi::Array params, DFPattern body) { return FunctionPattern(params, body); }); } -RELAX_PATTERN_PRINTER_DEF(FunctionPatternNode, [](auto p, auto node) { +RELAX_PATTERN_PRINTER_DEF(FunctionPatternObj, [](auto p, auto node) { p->stream << "FunctionPattern(" << node->params << ", " << node->body << ")"; }); TuplePattern::TuplePattern(tvm::ffi::Array fields) : DFPattern(ffi::UnsafeInit{}) { - ffi::ObjectPtr n = ffi::make_object(); + ffi::ObjectPtr n = ffi::make_object(); n->fields = std::move(fields); data_ = std::move(n); } TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - TuplePatternNode::RegisterReflection(); + TuplePatternObj::RegisterReflection(); refl::GlobalDef().def("relax.dpl.TuplePattern", [](tvm::ffi::Array fields) { return TuplePattern(fields); }); } -RELAX_PATTERN_PRINTER_DEF(TuplePatternNode, [](auto p, auto node) { +RELAX_PATTERN_PRINTER_DEF(TuplePatternObj, [](auto p, auto node) { p->stream << "TuplePattern(" << node->fields << ")"; }); UnorderedTuplePattern::UnorderedTuplePattern(tvm::ffi::Array fields) : DFPattern(ffi::UnsafeInit{}) { - ffi::ObjectPtr n = ffi::make_object(); + ffi::ObjectPtr n = ffi::make_object(); n->fields = std::move(fields); data_ = std::move(n); } TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - UnorderedTuplePatternNode::RegisterReflection(); + UnorderedTuplePatternObj::RegisterReflection(); refl::GlobalDef().def("relax.dpl.UnorderedTuplePattern", [](tvm::ffi::Array fields) { return UnorderedTuplePattern(fields); }); } -RELAX_PATTERN_PRINTER_DEF(UnorderedTuplePatternNode, [](auto p, auto node) { +RELAX_PATTERN_PRINTER_DEF(UnorderedTuplePatternObj, [](auto p, auto node) { p->stream << "UnorderedTuplePattern(" << node->fields << ")"; }); TupleGetItemPattern::TupleGetItemPattern(DFPattern tuple, int index) : DFPattern(ffi::UnsafeInit{}) { - ffi::ObjectPtr n = - ffi::make_object(std::move(tuple)); + ffi::ObjectPtr n = + ffi::make_object(std::move(tuple)); n->index = index; data_ = std::move(n); } TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - TupleGetItemPatternNode::RegisterReflection(); + TupleGetItemPatternObj::RegisterReflection(); refl::GlobalDef().def("relax.dpl.TupleGetItemPattern", [](DFPattern tuple, int index) { return TupleGetItemPattern(tuple, index); }); } -RELAX_PATTERN_PRINTER_DEF(TupleGetItemPatternNode, [](auto p, auto node) { +RELAX_PATTERN_PRINTER_DEF(TupleGetItemPatternObj, [](auto p, auto node) { p->stream << "TupleGetItemPattern(" << node->tuple << ", " << node->index << ")"; }); AndPattern::AndPattern(DFPattern left, DFPattern right) : DFPattern(ffi::UnsafeInit{}) { - ffi::ObjectPtr n = - ffi::make_object(std::move(left), std::move(right)); + ffi::ObjectPtr n = + ffi::make_object(std::move(left), std::move(right)); data_ = std::move(n); } TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - AndPatternNode::RegisterReflection(); + AndPatternObj::RegisterReflection(); refl::GlobalDef().def("relax.dpl.AndPattern", [](DFPattern left, DFPattern right) { return AndPattern(left, right); }); } -RELAX_PATTERN_PRINTER_DEF(AndPatternNode, [](auto p, auto node) { +RELAX_PATTERN_PRINTER_DEF(AndPatternObj, [](auto p, auto node) { p->stream << "AndPattern(" << node->left << " & " << node->right << ")"; }); OrPattern::OrPattern(DFPattern left, DFPattern right) : DFPattern(ffi::UnsafeInit{}) { - ffi::ObjectPtr n = - ffi::make_object(std::move(left), std::move(right)); + ffi::ObjectPtr n = + ffi::make_object(std::move(left), std::move(right)); data_ = std::move(n); } TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - OrPatternNode::RegisterReflection(); + OrPatternObj::RegisterReflection(); refl::GlobalDef().def("relax.dpl.OrPattern", [](DFPattern left, DFPattern right) { return OrPattern(left, right); }); } -RELAX_PATTERN_PRINTER_DEF(OrPatternNode, [](auto p, auto node) { +RELAX_PATTERN_PRINTER_DEF(OrPatternObj, [](auto p, auto node) { p->stream << "OrPattern(" << node->left << " | " << node->right << ")"; }); NotPattern::NotPattern(DFPattern reject) : DFPattern(ffi::UnsafeInit{}) { - ffi::ObjectPtr n = ffi::make_object(std::move(reject)); + ffi::ObjectPtr n = ffi::make_object(std::move(reject)); data_ = std::move(n); } TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - NotPatternNode::RegisterReflection(); + NotPatternObj::RegisterReflection(); refl::GlobalDef().def("relax.dpl.NotPattern", [](DFPattern reject) { return NotPattern(reject); }); } -RELAX_PATTERN_PRINTER_DEF(NotPatternNode, +RELAX_PATTERN_PRINTER_DEF(NotPatternObj, [](auto p, auto node) { p->stream << "!(" << node->reject << ")"; }); WildcardPattern::WildcardPattern() : DFPattern(ffi::UnsafeInit{}) { - data_ = ffi::make_object(); + data_ = ffi::make_object(); } TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - WildcardPatternNode::RegisterReflection(); + WildcardPatternObj::RegisterReflection(); refl::GlobalDef().def("relax.dpl.WildcardPattern", []() { return WildcardPattern(); }); } -RELAX_PATTERN_PRINTER_DEF(WildcardPatternNode, [](auto p, auto node) { p->stream << "*"; }); +RELAX_PATTERN_PRINTER_DEF(WildcardPatternObj, [](auto p, auto node) { p->stream << "*"; }); TypePattern::TypePattern(DFPattern pattern, Type ty) : DFPattern(ffi::UnsafeInit{}) { - ffi::ObjectPtr n = ffi::make_object(std::move(pattern)); + ffi::ObjectPtr n = ffi::make_object(std::move(pattern)); n->ty = std::move(ty); data_ = std::move(n); } TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - TypePatternNode::RegisterReflection(); + TypePatternObj::RegisterReflection(); refl::GlobalDef().def("relax.dpl.TypePattern", [](DFPattern pattern, Type ty) { return TypePattern(pattern, ty); }); } -RELAX_PATTERN_PRINTER_DEF(TypePatternNode, [](auto p, auto node) { +RELAX_PATTERN_PRINTER_DEF(TypePatternObj, [](auto p, auto node) { p->stream << "TypePattern(" << node->pattern << " has relax Type " << node->ty << ")"; }); ShapePattern::ShapePattern(DFPattern pattern, ffi::Array shape) : DFPattern(ffi::UnsafeInit{}) { - ffi::ObjectPtr n = ffi::make_object(std::move(pattern)); + ffi::ObjectPtr n = ffi::make_object(std::move(pattern)); n->shape = std::move(shape); data_ = std::move(n); } TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - ShapePatternNode::RegisterReflection(); + ShapePatternObj::RegisterReflection(); refl::GlobalDef().def( "relax.dpl.ShapePattern", [](DFPattern pattern, ffi::Array shape) { return ShapePattern(pattern, shape); }); } -RELAX_PATTERN_PRINTER_DEF(ShapePatternNode, [](auto p, auto node) { +RELAX_PATTERN_PRINTER_DEF(ShapePatternObj, [](auto p, auto node) { p->stream << "ShapePattern(" << node->pattern << " has shape " << node->shape << ")"; }); SameShapeConstraint::SameShapeConstraint(ffi::Array args) : DFConstraint(ffi::UnsafeInit{}) { - ffi::ObjectPtr n = ffi::make_object(); + ffi::ObjectPtr n = ffi::make_object(); n->args = std::move(args); data_ = std::move(n); @@ -367,11 +367,11 @@ SameShapeConstraint::SameShapeConstraint(ffi::Array args) TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - SameShapeConstraintNode::RegisterReflection(); + SameShapeConstraintObj::RegisterReflection(); refl::GlobalDef().def("relax.dpl.SameShapeConstraint", [](ffi::Array args) { return SameShapeConstraint(args); }); } -RELAX_PATTERN_PRINTER_DEF(SameShapeConstraintNode, [](auto p, auto node) { +RELAX_PATTERN_PRINTER_DEF(SameShapeConstraintObj, [](auto p, auto node) { p->stream << "SameShapeConstraint("; for (size_t i = 0; i < node->args.size(); i++) { if (i) { @@ -384,36 +384,36 @@ RELAX_PATTERN_PRINTER_DEF(SameShapeConstraintNode, [](auto p, auto node) { DataTypePattern::DataTypePattern(DFPattern pattern, DLDataType dtype) : DFPattern(ffi::UnsafeInit{}) { - ffi::ObjectPtr n = ffi::make_object(std::move(pattern)); + ffi::ObjectPtr n = ffi::make_object(std::move(pattern)); n->dtype = dtype; data_ = std::move(n); } TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - DataTypePatternNode::RegisterReflection(); + DataTypePatternObj::RegisterReflection(); refl::GlobalDef().def("relax.dpl.DataTypePattern", [](DFPattern pattern, DLDataType dtype) { return DataTypePattern(pattern, dtype); }); } -RELAX_PATTERN_PRINTER_DEF(DataTypePatternNode, [](auto p, auto node) { +RELAX_PATTERN_PRINTER_DEF(DataTypePatternObj, [](auto p, auto node) { p->stream << "DataTypePattern(" << node->pattern << " has dtype " << node->dtype << ")"; }); AttrPattern::AttrPattern(DFPattern pattern, DictAttrs attrs) : DFPattern(ffi::UnsafeInit{}) { - ffi::ObjectPtr n = ffi::make_object(std::move(pattern)); + ffi::ObjectPtr n = ffi::make_object(std::move(pattern)); n->attrs = std::move(attrs); data_ = std::move(n); } TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - AttrPatternNode::RegisterReflection(); + AttrPatternObj::RegisterReflection(); refl::GlobalDef().def("relax.dpl.AttrPattern", [](DFPattern pattern, DictAttrs attrs) { return AttrPattern(pattern, attrs); }); } -RELAX_PATTERN_PRINTER_DEF(AttrPatternNode, [](auto p, auto node) { +RELAX_PATTERN_PRINTER_DEF(AttrPatternObj, [](auto p, auto node) { p->stream << "AttrPattern(" << node->pattern << " has attributes " << node->attrs << ")"; }); @@ -422,56 +422,54 @@ class DFPatternDuplicator : public DFPatternFunctor DFPattern VisitDFPattern(const DFPattern& pattern) override { return DFPatternFunctor::VisitDFPattern(pattern); } - DFPattern VisitDFPattern_(const OrPatternNode* op) override { + DFPattern VisitDFPattern_(const OrPatternObj* op) override { return OrPattern(op->left, op->right); } - DFPattern VisitDFPattern_(const AndPatternNode* op) override { + DFPattern VisitDFPattern_(const AndPatternObj* op) override { return AndPattern(op->left, op->right); } - DFPattern VisitDFPattern_(const NotPatternNode* op) override { return NotPattern(op->reject); } - DFPattern VisitDFPattern_(const VarPatternNode* op) override { return VarPattern(op->name); } - DFPattern VisitDFPattern_(const ConstantPatternNode* op) override { - return ConstantPattern(ffi::UnsafeInit{}, ffi::make_object()); + DFPattern VisitDFPattern_(const NotPatternObj* op) override { return NotPattern(op->reject); } + DFPattern VisitDFPattern_(const VarPatternObj* op) override { return VarPattern(op->name); } + DFPattern VisitDFPattern_(const ConstantPatternObj* op) override { + return ConstantPattern(ffi::UnsafeInit{}, ffi::make_object()); } - DFPattern VisitDFPattern_(const WildcardPatternNode* op) override { - return WildcardPattern(ffi::make_object()); + DFPattern VisitDFPattern_(const WildcardPatternObj* op) override { + return WildcardPattern(ffi::make_object()); } - DFPattern VisitDFPattern_(const ExprPatternNode* op) override { return ExprPattern(op->expr); } - DFPattern VisitDFPattern_(const GlobalVarPatternNode* op) override { + DFPattern VisitDFPattern_(const ExprPatternObj* op) override { return ExprPattern(op->expr); } + DFPattern VisitDFPattern_(const GlobalVarPatternObj* op) override { return GlobalVarPattern(op->name); } - DFPattern VisitDFPattern_(const TuplePatternNode* op) override { - return TuplePattern(op->fields); - } - DFPattern VisitDFPattern_(const UnorderedTuplePatternNode* op) override { + DFPattern VisitDFPattern_(const TuplePatternObj* op) override { return TuplePattern(op->fields); } + DFPattern VisitDFPattern_(const UnorderedTuplePatternObj* op) override { return UnorderedTuplePattern(op->fields); } - DFPattern VisitDFPattern_(const TupleGetItemPatternNode* op) override { + DFPattern VisitDFPattern_(const TupleGetItemPatternObj* op) override { return TupleGetItemPattern(op->tuple, op->index); } - DFPattern VisitDFPattern_(const CallPatternNode* op) override { + DFPattern VisitDFPattern_(const CallPatternObj* op) override { return CallPattern(op->op, op->args); } - DFPattern VisitDFPattern_(const DataTypePatternNode* op) override { + DFPattern VisitDFPattern_(const DataTypePatternObj* op) override { return DataTypePattern(op->pattern, op->dtype); } - DFPattern VisitDFPattern_(const FunctionPatternNode* op) override { + DFPattern VisitDFPattern_(const FunctionPatternObj* op) override { return FunctionPattern(op->params, op->body); } - DFPattern VisitDFPattern_(const ShapePatternNode* op) override { + DFPattern VisitDFPattern_(const ShapePatternObj* op) override { return ShapePattern(op->pattern, op->shape); } - DFPattern VisitDFPattern_(const TypePatternNode* op) override { + DFPattern VisitDFPattern_(const TypePatternObj* op) override { return TypePattern(op->pattern, op->ty); } - DFPattern VisitDFPattern_(const DataflowVarPatternNode* op) override { + DFPattern VisitDFPattern_(const DataflowVarPatternObj* op) override { return DataflowVarPattern(op->name); } - DFPattern VisitDFPattern_(const ExternFuncPatternNode* op) override { + DFPattern VisitDFPattern_(const ExternFuncPatternObj* op) override { return ExternFuncPattern(op->global_symbol()); } - DFPattern VisitDFPattern_(const PrimArrPatternNode* op) override { + DFPattern VisitDFPattern_(const PrimArrPatternObj* op) override { return PrimArrPattern(op->fields); } }; @@ -539,7 +537,7 @@ static void sync_graph_constraints(const DFPattern& lhs, const DFPattern& rhs, P } PatternSeq::PatternSeq(DFPattern init_pattern) { - ffi::ObjectPtr n = ffi::make_object(); + ffi::ObjectPtr n = ffi::make_object(); n->patterns = {init_pattern}; n->pair_constraints = {}; data_ = std::move(n); @@ -548,7 +546,7 @@ PatternSeq::PatternSeq(tvm::ffi::Array patterns, bool only_used_by) { TVM_FFI_ICHECK_GE(patterns.size(), 1) << "PatternSeq must have at least one pattern"; const auto cons = PairCons(only_used_by ? PairCons::kOnlyUsedBy : PairCons::kUsedBy); - ffi::ObjectPtr n = ffi::make_object(); + ffi::ObjectPtr n = ffi::make_object(); n->patterns = std::move(patterns); n->pair_constraints = std::vector(n->patterns.size() - 1, cons); data_ = std::move(n); @@ -565,7 +563,7 @@ PatternSeq PatternSeq::OnlyUsedBy(PatternSeq other, int index) const { PatternSeq PatternSeq::dup() const { PatternSeq ret(ffi::UnsafeInit{}); - ffi::ObjectPtr n = ffi::make_object(); + ffi::ObjectPtr n = ffi::make_object(); n->patterns = ffi::Array{}; n->patterns.reserve(get()->patterns.size()); n->pair_constraints = this->get()->pair_constraints; @@ -583,13 +581,13 @@ PatternSeq PatternSeq::dup() const { TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - PatternSeqNode::RegisterReflection(); + PatternSeqObj::RegisterReflection(); refl::GlobalDef().def("relax.dpl.PatternSeq", [](ffi::Array patterns, bool only_used_by) { return PatternSeq(std::move(patterns), only_used_by); }); } -RELAX_PATTERN_PRINTER_DEF(PatternSeqNode, [](auto p, auto node) { +RELAX_PATTERN_PRINTER_DEF(PatternSeqObj, [](auto p, auto node) { p->stream << "["; for (size_t i = 0; i < node->patterns.size(); ++i) { if (i != 0) @@ -627,7 +625,7 @@ PatternSeq UsedBy(const PatternSeq& lhs, const PatternSeq& rhs, int index) { pair_constraints.insert(pair_constraints.end(), rhs->pair_constraints.begin(), rhs->pair_constraints.end()); - ffi::ObjectPtr n = ffi::make_object(); + ffi::ObjectPtr n = ffi::make_object(); n->patterns = std::move(patterns); n->pair_constraints = std::move(pair_constraints); ret.data_ = std::move(n); @@ -654,7 +652,7 @@ PatternSeq OnlyUsedBy(const PatternSeq& lhs, const PatternSeq& rhs, int index) { pair_constraints.insert(pair_constraints.end(), rhs->pair_constraints.begin(), rhs->pair_constraints.end()); - ffi::ObjectPtr n = ffi::make_object(); + ffi::ObjectPtr n = ffi::make_object(); n->patterns = std::move(patterns); n->pair_constraints = std::move(pair_constraints); ret.data_ = std::move(n); @@ -665,9 +663,9 @@ PatternSeq operator>>(const PatternSeq& lhs, const PatternSeq& rhs) { return lhs VarPattern IsVar(const ffi::String& name) { return VarPattern(name); } ConstantPattern IsConst() { - return ConstantPattern(ffi::UnsafeInit{}, ffi::make_object()); + return ConstantPattern(ffi::UnsafeInit{}, ffi::make_object()); } -WildcardPattern Wildcard() { return WildcardPattern(ffi::make_object()); } +WildcardPattern Wildcard() { return WildcardPattern(ffi::make_object()); } ExprPattern IsExpr(const Expr& expr) { return ExprPattern(expr); } ExprPattern IsOp(const ffi::String& op_name) { return IsExpr(Op::Get(op_name)); } CallPattern IsCallTIR(const ffi::String& name, ffi::Optional var_args) { diff --git a/src/relax/ir/dataflow_pattern_functor.cc b/src/relax/ir/dataflow_pattern_functor.cc index d269196dc32f..3d34eab53645 100644 --- a/src/relax/ir/dataflow_pattern_functor.cc +++ b/src/relax/ir/dataflow_pattern_functor.cc @@ -36,21 +36,21 @@ void DFPatternVisitor::VisitDFPattern(const DFPattern& pattern) { } } -void DFPatternVisitor::VisitDFPattern_(const OrPatternNode* op) { +void DFPatternVisitor::VisitDFPattern_(const OrPatternObj* op) { VisitDFPattern(op->left); VisitDFPattern(op->right); } -void DFPatternVisitor::VisitDFPattern_(const AndPatternNode* op) { +void DFPatternVisitor::VisitDFPattern_(const AndPatternObj* op) { VisitDFPattern(op->left); VisitDFPattern(op->right); } -void DFPatternVisitor::VisitDFPattern_(const NotPatternNode* op) { VisitDFPattern(op->reject); } +void DFPatternVisitor::VisitDFPattern_(const NotPatternObj* op) { VisitDFPattern(op->reject); } -void DFPatternVisitor::VisitDFPattern_(const AttrPatternNode* op) { VisitDFPattern(op->pattern); } +void DFPatternVisitor::VisitDFPattern_(const AttrPatternObj* op) { VisitDFPattern(op->pattern); } -void DFPatternVisitor::VisitDFPattern_(const CallPatternNode* op) { +void DFPatternVisitor::VisitDFPattern_(const CallPatternObj* op) { VisitDFPattern(op->op); if (op->args.defined()) { for (auto arg : op->args) { @@ -59,13 +59,13 @@ void DFPatternVisitor::VisitDFPattern_(const CallPatternNode* op) { } } -void DFPatternVisitor::VisitDFPattern_(const DataTypePatternNode* op) { +void DFPatternVisitor::VisitDFPattern_(const DataTypePatternObj* op) { VisitDFPattern(op->pattern); } -void DFPatternVisitor::VisitDFPattern_(const ExprPatternNode* op) {} +void DFPatternVisitor::VisitDFPattern_(const ExprPatternObj* op) {} -void DFPatternVisitor::VisitDFPattern_(const FunctionPatternNode* op) { +void DFPatternVisitor::VisitDFPattern_(const FunctionPatternObj* op) { if (op->params.defined()) { for (auto param : op->params) { VisitDFPattern(param); @@ -74,13 +74,13 @@ void DFPatternVisitor::VisitDFPattern_(const FunctionPatternNode* op) { VisitDFPattern(op->body); } -void DFPatternVisitor::VisitDFPattern_(const ShapePatternNode* op) { VisitDFPattern(op->pattern); } +void DFPatternVisitor::VisitDFPattern_(const ShapePatternObj* op) { VisitDFPattern(op->pattern); } -void DFPatternVisitor::VisitDFPattern_(const TupleGetItemPatternNode* op) { +void DFPatternVisitor::VisitDFPattern_(const TupleGetItemPatternObj* op) { VisitDFPattern(op->tuple); } -void DFPatternVisitor::VisitDFPattern_(const TuplePatternNode* op) { +void DFPatternVisitor::VisitDFPattern_(const TuplePatternObj* op) { if (op->fields.defined()) { for (auto field : op->fields) { VisitDFPattern(field); @@ -88,7 +88,7 @@ void DFPatternVisitor::VisitDFPattern_(const TuplePatternNode* op) { } } -void DFPatternVisitor::VisitDFPattern_(const UnorderedTuplePatternNode* op) { +void DFPatternVisitor::VisitDFPattern_(const UnorderedTuplePatternObj* op) { if (op->fields.defined()) { for (auto field : op->fields) { VisitDFPattern(field); @@ -96,16 +96,16 @@ void DFPatternVisitor::VisitDFPattern_(const UnorderedTuplePatternNode* op) { } } -void DFPatternVisitor::VisitDFPattern_(const TypePatternNode* op) { VisitDFPattern(op->pattern); } +void DFPatternVisitor::VisitDFPattern_(const TypePatternObj* op) { VisitDFPattern(op->pattern); } // leaf nodes. -void DFPatternVisitor::VisitDFPattern_(const PrimArrPatternNode* op) {} -void DFPatternVisitor::VisitDFPattern_(const VarPatternNode* op) {} -void DFPatternVisitor::VisitDFPattern_(const ConstantPatternNode* op) {} -void DFPatternVisitor::VisitDFPattern_(const DataflowVarPatternNode* op) {} -void DFPatternVisitor::VisitDFPattern_(const GlobalVarPatternNode* op) {} -void DFPatternVisitor::VisitDFPattern_(const ExternFuncPatternNode* op) {} -void DFPatternVisitor::VisitDFPattern_(const WildcardPatternNode* op) {} +void DFPatternVisitor::VisitDFPattern_(const PrimArrPatternObj* op) {} +void DFPatternVisitor::VisitDFPattern_(const VarPatternObj* op) {} +void DFPatternVisitor::VisitDFPattern_(const ConstantPatternObj* op) {} +void DFPatternVisitor::VisitDFPattern_(const DataflowVarPatternObj* op) {} +void DFPatternVisitor::VisitDFPattern_(const GlobalVarPatternObj* op) {} +void DFPatternVisitor::VisitDFPattern_(const ExternFuncPatternObj* op) {} +void DFPatternVisitor::VisitDFPattern_(const WildcardPatternObj* op) {} } // namespace relax } // namespace tvm diff --git a/src/relax/ir/dependent_type.cc b/src/relax/ir/dependent_type.cc index 2134064bc38e..fce1eee6def1 100644 --- a/src/relax/ir/dependent_type.cc +++ b/src/relax/ir/dependent_type.cc @@ -39,8 +39,8 @@ namespace { TVM_FFI_INLINE ffi::Expected> ShapeTypeVisit( ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept { // skips: ndim (scalar) - const ShapeTypeNode* self = - ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value); + const ShapeTypeObj* self = + ffi::detail::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value); TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->values)); return std::nullopt; } @@ -48,12 +48,12 @@ TVM_FFI_INLINE ffi::Expected> ShapeTypeVisit( TVM_FFI_INLINE ffi::Expected> ShapeTypeMutate( ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept { // skips: ndim (scalar) - const ShapeTypeNode* self = - ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value); + const ShapeTypeObj* self = + ffi::detail::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value); TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr>>, mapped_values, mutator->MutateExpected(self->values)); if (mapped_values.UnchangedOrSameAs(self->values)) return ffi::Unchanged(); - ffi::ObjectPtr copy = ffi::make_object(*self); + ffi::ObjectPtr copy = ffi::make_object(*self); copy->values = std::move(mapped_values).ValueOrUnchanged(std::move(copy->values)); return ffi::Any(std::move(copy)); } @@ -61,8 +61,8 @@ TVM_FFI_INLINE ffi::Expected> ShapeTypeMutate( TVM_FFI_INLINE ffi::Expected> ShapeTypeMaybeInplaceMutate( ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept { // skips: ndim (scalar) - ShapeTypeNode* self = const_cast( - ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value)); + ShapeTypeObj* self = const_cast( + ffi::detail::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value)); TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN( ffi::UnchangedOr>>, mapped_values, mutator->MutateExpected(self->values, ffi::InplaceMode::kAllow)); @@ -75,8 +75,8 @@ TVM_FFI_INLINE ffi::Expected> ShapeTypeMaybeInplaceMu TVM_FFI_INLINE ffi::Expected> TensorTypeVisit( ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept { // skips: ndim (scalar) - const TensorTypeNode* self = - ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value); + const TensorTypeObj* self = + ffi::detail::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value); TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->shape)); TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->dtype)); TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->vdevice)); @@ -86,8 +86,8 @@ TVM_FFI_INLINE ffi::Expected> TensorTypeVisit TVM_FFI_INLINE ffi::Expected> TensorTypeMutate( ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept { // skips: ndim (scalar) - const TensorTypeNode* self = - ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value); + const TensorTypeObj* self = + ffi::detail::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value); TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr>, mapped_shape, mutator->MutateExpected(self->shape)); TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr>, mapped_dtype, @@ -98,7 +98,7 @@ TVM_FFI_INLINE ffi::Expected> TensorTypeMutate( mapped_vdevice.UnchangedOrSameAs(self->vdevice)) { return ffi::Unchanged(); } - ffi::ObjectPtr copy = ffi::make_object(*self); + ffi::ObjectPtr copy = ffi::make_object(*self); copy->shape = std::move(mapped_shape).ValueOrUnchanged(std::move(copy->shape)); copy->dtype = std::move(mapped_dtype).ValueOrUnchanged(std::move(copy->dtype)); copy->vdevice = std::move(mapped_vdevice).ValueOrUnchanged(std::move(copy->vdevice)); @@ -108,8 +108,8 @@ TVM_FFI_INLINE ffi::Expected> TensorTypeMutate( TVM_FFI_INLINE ffi::Expected> TensorTypeMaybeInplaceMutate( ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept { // skips: ndim (scalar) - TensorTypeNode* self = const_cast( - ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value)); + TensorTypeObj* self = const_cast( + ffi::detail::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value)); TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr>, mapped_shape, mutator->MutateExpected(self->shape, ffi::InplaceMode::kAllow)); TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr>, mapped_dtype, @@ -132,8 +132,8 @@ TVM_FFI_INLINE ffi::Expected> TensorTypeMaybeInplaceM TVM_FFI_INLINE ffi::Expected> FuncTypeVisit( ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept { // skips: derive_func (environment-backed callable metadata), purity (scalar) - const FuncTypeNode* self = - ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value); + const FuncTypeObj* self = + ffi::detail::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value); TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->WithDefRegionKind( kTVMFFIDefRegionKindPattern, [&]() { return visitor->VisitExpected(self->params); })); TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->ret)); @@ -143,8 +143,8 @@ TVM_FFI_INLINE ffi::Expected> FuncTypeVisit( TVM_FFI_INLINE ffi::Expected> FuncTypeMutate( ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept { // skips: derive_func (environment-backed callable metadata), purity (scalar) - const FuncTypeNode* self = - ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value); + const FuncTypeObj* self = + ffi::detail::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value); TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN( ffi::UnchangedOr>>, mapped_params, mutator->WithDefRegionKind(kTVMFFIDefRegionKindPattern, @@ -154,7 +154,7 @@ TVM_FFI_INLINE ffi::Expected> FuncTypeMutate( if (mapped_params.UnchangedOrSameAs(self->params) && mapped_ret.UnchangedOrSameAs(self->ret)) { return ffi::Unchanged(); } - ffi::ObjectPtr copy = ffi::make_object(*self); + ffi::ObjectPtr copy = ffi::make_object(*self); copy->params = std::move(mapped_params).ValueOrUnchanged(std::move(copy->params)); copy->ret = std::move(mapped_ret).ValueOrUnchanged(std::move(copy->ret)); return ffi::Any(std::move(copy)); @@ -163,8 +163,8 @@ TVM_FFI_INLINE ffi::Expected> FuncTypeMutate( TVM_FFI_INLINE ffi::Expected> FuncTypeMaybeInplaceMutate( ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept { // skips: derive_func (environment-backed callable metadata), purity (scalar) - FuncTypeNode* self = const_cast( - ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value)); + FuncTypeObj* self = const_cast( + ffi::detail::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value)); TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN( ffi::UnchangedOr>>, mapped_params, mutator->WithDefRegionKind(kTVMFFIDefRegionKindPattern, [&]() { @@ -185,10 +185,10 @@ TVM_FFI_INLINE ffi::Expected> FuncTypeMaybeInplaceMut // Shape ShapeType::ShapeType(ffi::Array values, Location loc) : Type(ffi::UnsafeInit{}) { - ffi::ObjectPtr n = ffi::make_object(); + ffi::ObjectPtr n = ffi::make_object(); n->ndim = static_cast(values.size()); n->values = values.Map([](PrimExpr value) { - if (value->IsInstance()) { + if (value->IsInstance()) { return tvm::prim::cast(value, PrimType::Int(64)); } TVM_FFI_ICHECK(value.ty().MatchesElementType(DLDataTypeCode::kDLInt, 64)) @@ -200,7 +200,7 @@ ShapeType::ShapeType(ffi::Array values, Location loc) : Type(ffi::Unsa } ShapeType::ShapeType(int ndim, Location loc) : Type(ffi::UnsafeInit{}) { - ffi::ObjectPtr n = ffi::make_object(); + ffi::ObjectPtr n = ffi::make_object(); TVM_FFI_ICHECK(ndim >= -1) << "ndim of ShapeType must be >= -1, but got " << ndim; n->ndim = ndim; n->loc = loc; @@ -209,8 +209,8 @@ ShapeType::ShapeType(int ndim, Location loc) : Type(ffi::UnsafeInit{}) { TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - ShapeTypeNode::RegisterReflection(); - refl::TypeAttrDef() + ShapeTypeObj::RegisterReflection(); + refl::TypeAttrDef() .attr(ffi::type_attr::kStructuralVisit, ffi::FStructuralVisit::FromNative<&ShapeTypeVisit>()) .attr(ffi::type_attr::kStructuralMutate, ffi::FStructuralMutate::FromNative<&ShapeTypeMutate>()) @@ -235,12 +235,12 @@ TVM_FFI_STATIC_INIT_BLOCK() { TensorType::TensorType(Expr shape, ffi::Optional dtype, ffi::Optional vdevice, Location loc) : Type(ffi::UnsafeInit{}) { - ffi::ObjectPtr n = ffi::make_object(); + ffi::ObjectPtr n = ffi::make_object(); // assign ndim before move TVM_FFI_ICHECK(shape.defined()) << "Must provide a shape in this constructor"; ffi::Optional shape_ty = MatchType(shape); TVM_FFI_ICHECK(shape_ty) << "We expect shape to contain pre-set shape type"; - TVM_FFI_ICHECK(shape->IsInstance() || shape->IsInstance()) + TVM_FFI_ICHECK(shape->IsInstance() || shape->IsInstance()) << "We require shape to be normalized when constructing TensorType"; n->ndim = shape_ty.value()->ndim; // assign rest of the fields. @@ -254,7 +254,7 @@ TensorType::TensorType(Expr shape, ffi::Optional dtype, ffi::Optional< TensorType::TensorType(ffi::Optional dtype, int ndim, ffi::Optional vdevice, Location loc) : Type(ffi::UnsafeInit{}) { - ffi::ObjectPtr n = ffi::make_object(); + ffi::ObjectPtr n = ffi::make_object(); TVM_FFI_ICHECK(ndim >= -1) << "ndim of TensorType must be >= -1, but got " << ndim; n->ndim = ndim; n->dtype = dtype; @@ -265,8 +265,8 @@ TensorType::TensorType(ffi::Optional dtype, int ndim, ffi::Optional() + TensorTypeObj::RegisterReflection(); + refl::TypeAttrDef() .attr(ffi::type_attr::kStructuralVisit, ffi::FStructuralVisit::FromNative<&TensorTypeVisit>()) .attr(ffi::type_attr::kStructuralMutate, ffi::FStructuralMutate::FromNative<&TensorTypeMutate>()) @@ -279,7 +279,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { auto index = slice[0].as(); TVM_FFI_CHECK(index.has_value(), TypeError) << "A Relax expression requires a point index"; - const auto* imm = index.value().as(); + const auto* imm = index.value().as(); TVM_FFI_CHECK(imm != nullptr, TypeError) << "A Relax expression requires a constant integer index"; return TupleGetItem(value, imm->value.as().value(), loc); @@ -303,7 +303,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { // Func FuncType::FuncType(ffi::Array params, Type ret, bool purity, Location loc) : Type(ffi::UnsafeInit{}) { - ffi::ObjectPtr n = ffi::make_object(); + ffi::ObjectPtr n = ffi::make_object(); n->params = std::move(params); n->ret = std::move(ret); n->purity = std::move(purity); @@ -313,8 +313,8 @@ FuncType::FuncType(ffi::Array params, Type ret, bool purity, Location loc) TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - FuncTypeNode::RegisterReflection(); - refl::TypeAttrDef() + FuncTypeObj::RegisterReflection(); + refl::TypeAttrDef() .attr(ffi::type_attr::kStructuralVisit, ffi::FStructuralVisit::FromNative<&FuncTypeVisit>()) .attr(ffi::type_attr::kStructuralMutate, ffi::FStructuralMutate::FromNative<&FuncTypeMutate>()) @@ -328,7 +328,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { } FuncType FuncType::OpaqueFunc(TypeDeriveFunc derive_func, bool purity, Location loc) { - ffi::ObjectPtr n = ffi::make_object(); + ffi::ObjectPtr n = ffi::make_object(); n->derive_func = std::move(derive_func); n->ret = AnyType(); n->purity = std::move(purity); @@ -337,7 +337,7 @@ FuncType FuncType::OpaqueFunc(TypeDeriveFunc derive_func, bool purity, Location } FuncType FuncType::OpaqueFunc(Type ret, bool purity, Location loc) { - ffi::ObjectPtr n = ffi::make_object(); + ffi::ObjectPtr n = ffi::make_object(); n->ret = std::move(ret); n->purity = std::move(purity); n->loc = loc; diff --git a/src/relax/ir/emit_te.cc b/src/relax/ir/emit_te.cc index 8f160aca0cca..37352f27263e 100644 --- a/src/relax/ir/emit_te.cc +++ b/src/relax/ir/emit_te.cc @@ -30,18 +30,18 @@ namespace tvm { namespace relax { -// Pattern A (RM): auto-default repr from reflection for RXPlaceholderOpNode. +// Pattern A (RM): auto-default repr from reflection for RXPlaceholderOpObj. -TVM_FFI_STATIC_INIT_BLOCK() { RXPlaceholderOpNode::RegisterReflection(); } +TVM_FFI_STATIC_INIT_BLOCK() { RXPlaceholderOpObj::RegisterReflection(); } te::Tensor TETensor(Expr value, ffi::Map tir_var_map, std::string name) { - auto n = ffi::make_object(value); + auto n = ffi::make_object(value); n->name = name; // If the value is a constant, it might come as an argument of EmitTE and thus its shape and // checked-type might not be properly set. In this case we set the shape and dtype of the returned // TE tensor. - if (const auto* constant = value.as()) { + if (const auto* constant = value.as()) { n->dtype = PrimType(constant->value.cast().dtype()); int ndim = constant->value.cast().ndim(); @@ -56,9 +56,9 @@ te::Tensor TETensor(Expr value, ffi::Map tir_var_map, std::s } TVM_FFI_ICHECK(!value->ty.as().has_value()) << "value must be normalized and contain Type"; - auto* tensor_ty = GetTypeAs(value); + auto* tensor_ty = GetTypeAs(value); TVM_FFI_ICHECK(tensor_ty) << "Value must be a tensor"; - auto* shape_expr = tensor_ty->shape.as(); + auto* shape_expr = tensor_ty->shape.as(); TVM_FFI_CHECK(shape_expr, ValueError) << "Expression does not have an known symbolic shape, please consider use " "match_cast " diff --git a/src/relax/ir/emit_te.h b/src/relax/ir/emit_te.h index 38377900c002..ea48100b84d6 100644 --- a/src/relax/ir/emit_te.h +++ b/src/relax/ir/emit_te.h @@ -36,24 +36,24 @@ namespace relax { /*! * \brief A placeholder op that represents a relax expression. */ -class RXPlaceholderOpNode : public te::PlaceholderOpNode { +class RXPlaceholderOpObj : public te::PlaceholderOpObj { public: - explicit RXPlaceholderOpNode(Expr value) : value(std::move(value)) {} - explicit RXPlaceholderOpNode(ffi::UnsafeInit) : value(ffi::UnsafeInit{}) {} + explicit RXPlaceholderOpObj(Expr value) : value(std::move(value)) {} + explicit RXPlaceholderOpObj(ffi::UnsafeInit) : value(ffi::UnsafeInit{}) {} /*! \brief The relax expression. */ Expr value; static void RegisterReflection() { namespace refl = tvm::ffi::reflection; - refl::ObjectDef().def_ro("value", &RXPlaceholderOpNode::value); + refl::ObjectDef().def_ro("value", &RXPlaceholderOpObj::value); } // FFI system configuration for structural equality and hashing static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode; - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.TEPlaceholderOp", RXPlaceholderOpNode, - te::PlaceholderOpNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.TEPlaceholderOp", RXPlaceholderOpObj, + te::PlaceholderOpObj); }; /*! diff --git a/src/relax/ir/expr.cc b/src/relax/ir/expr.cc index 716b2fac9ce9..0ca6f9d3a621 100644 --- a/src/relax/ir/expr.cc +++ b/src/relax/ir/expr.cc @@ -31,13 +31,12 @@ namespace relax { namespace { -// Traverses only ExprNode::ty. The per-node payload is an intentional constant leaf: -// ExternFuncNode::global_symbol is scalar and BaseFuncNode::attrs is metadata. +// Traverses only ExprObj::ty. The per-node payload is an intentional constant leaf: +// ExternFuncObj::global_symbol is scalar and BaseFuncObj::attrs is metadata. template TVM_FFI_INLINE ffi::Expected> TypeOnlyExprVisit( ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept { - const TNode* self = - ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value); + const TNode* self = ffi::detail::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value); TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->ty)); return std::nullopt; } @@ -45,8 +44,7 @@ TVM_FFI_INLINE ffi::Expected> TypeOnlyExprVis template TVM_FFI_INLINE ffi::Expected> TypeOnlyExprMutate( ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept { - const TNode* self = - ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value); + const TNode* self = ffi::detail::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value); TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, mapped_ty, mutator->MutateExpected(self->ty)); if (mapped_ty.UnchangedOrSameAs(self->ty)) return ffi::Unchanged(); @@ -59,7 +57,7 @@ template TVM_FFI_INLINE ffi::Expected> TypeOnlyExprMaybeInplaceMutate( ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept { TNode* self = const_cast( - ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value)); + ffi::detail::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value)); TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, mapped_ty, mutator->MutateExpected(self->ty, ffi::InplaceMode::kAllow)); if (!mapped_ty.UnchangedOrSameAs(self->ty)) { @@ -70,8 +68,8 @@ TVM_FFI_INLINE ffi::Expected> TypeOnlyExprMaybeInplac TVM_FFI_INLINE ffi::Expected> IfExprVisit( ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept { - const IfExprNode* self = - ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value); + const IfExprObj* self = + ffi::detail::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value); TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->ty)); TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->cond)); TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->true_branch)); @@ -81,8 +79,8 @@ TVM_FFI_INLINE ffi::Expected> IfExprVisit( TVM_FFI_INLINE ffi::Expected> IfExprMutate( ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept { - const IfExprNode* self = - ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value); + const IfExprObj* self = + ffi::detail::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value); TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, mapped_ty, mutator->MutateExpected(self->ty)); TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, mapped_cond, @@ -96,7 +94,7 @@ TVM_FFI_INLINE ffi::Expected> IfExprMutate( mapped_false_branch.UnchangedOrSameAs(self->false_branch)) { return ffi::Unchanged(); } - ffi::ObjectPtr copy = ffi::make_object(*self); + ffi::ObjectPtr copy = ffi::make_object(*self); copy->ty = std::move(mapped_ty).ValueOrUnchanged(std::move(copy->ty)); copy->cond = std::move(mapped_cond).ValueOrUnchanged(std::move(copy->cond)); copy->true_branch = std::move(mapped_true_branch).ValueOrUnchanged(std::move(copy->true_branch)); @@ -107,8 +105,8 @@ TVM_FFI_INLINE ffi::Expected> IfExprMutate( TVM_FFI_INLINE ffi::Expected> IfExprMaybeInplaceMutate( ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept { - IfExprNode* self = const_cast( - ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value)); + IfExprObj* self = const_cast( + ffi::detail::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value)); TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, mapped_ty, mutator->MutateExpected(self->ty, ffi::InplaceMode::kAllow)); TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, mapped_cond, @@ -134,8 +132,8 @@ TVM_FFI_INLINE ffi::Expected> IfExprMaybeInplaceMutat TVM_FFI_INLINE ffi::Expected> ShapeExprVisit( ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept { - const ShapeExprNode* self = - ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value); + const ShapeExprObj* self = + ffi::detail::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value); TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->ty)); TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->values)); return std::nullopt; @@ -143,8 +141,8 @@ TVM_FFI_INLINE ffi::Expected> ShapeExprVisit( TVM_FFI_INLINE ffi::Expected> ShapeExprMutate( ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept { - const ShapeExprNode* self = - ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value); + const ShapeExprObj* self = + ffi::detail::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value); TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, mapped_ty, mutator->MutateExpected(self->ty)); TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr>, mapped_values, @@ -152,7 +150,7 @@ TVM_FFI_INLINE ffi::Expected> ShapeExprMutate( if (mapped_ty.UnchangedOrSameAs(self->ty) && mapped_values.UnchangedOrSameAs(self->values)) { return ffi::Unchanged(); } - ffi::ObjectPtr copy = ffi::make_object(*self); + ffi::ObjectPtr copy = ffi::make_object(*self); copy->ty = std::move(mapped_ty).ValueOrUnchanged(std::move(copy->ty)); copy->values = std::move(mapped_values).ValueOrUnchanged(std::move(copy->values)); return ffi::Any(std::move(copy)); @@ -160,8 +158,8 @@ TVM_FFI_INLINE ffi::Expected> ShapeExprMutate( TVM_FFI_INLINE ffi::Expected> ShapeExprMaybeInplaceMutate( ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept { - ShapeExprNode* self = const_cast( - ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value)); + ShapeExprObj* self = const_cast( + ffi::detail::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value)); TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, mapped_ty, mutator->MutateExpected(self->ty, ffi::InplaceMode::kAllow)); TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN( @@ -174,13 +172,13 @@ TVM_FFI_INLINE ffi::Expected> ShapeExprMaybeInplaceMu return ffi::Unchanged(); } -// Hooks do not inherit, so DataflowVar must mirror the base VarNode remap, PrimType-skip, and -// Simple-to-None definition-region protocol. Keep this hook triple in lockstep with VarNode. +// Hooks do not inherit, so DataflowVar must mirror the base VarObj remap, PrimType-skip, and +// Simple-to-None definition-region protocol. Keep this hook triple in lockstep with VarObj. TVM_FFI_INLINE ffi::Expected> DataflowVarVisit( ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept { - const DataflowVarNode* self = - ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value); - if (!self->ty.as()) { + const DataflowVarObj* self = + ffi::detail::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value); + if (!self->ty.as()) { if (visitor->def_region_kind() == kTVMFFIDefRegionKindSimple) { TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->WithDefRegionKind( kTVMFFIDefRegionKindNone, [&]() { return visitor->VisitExpected(self->ty); })); @@ -193,11 +191,11 @@ TVM_FFI_INLINE ffi::Expected> DataflowVarVisi TVM_FFI_INLINE ffi::Expected> DataflowVarMutate( ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept { - const DataflowVarNode* self = - ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value); + const DataflowVarObj* self = + ffi::detail::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value); ffi::Expected remap_result = mutator->VarRemapGetExpected(value); TVM_FFI_S_MUTATE_MAYBE_EARLY_RETURN(remap_result); - if (ffi::details::ExpectedUnsafe::GetData(remap_result).type_index() != + if (ffi::detail::ExpectedUnsafe::GetData(remap_result).type_index() != ffi::TypeIndex::kTVMFFINone) { return std::move(remap_result); } @@ -206,7 +204,7 @@ TVM_FFI_INLINE ffi::Expected> DataflowVarMutate( } ffi::UnchangedOr result = ffi::Unchanged(); ffi::Any mapped_value; - if (!self->ty.as()) { + if (!self->ty.as()) { ffi::Expected> mapped_ty_result = mutator->def_region_kind() == kTVMFFIDefRegionKindSimple ? mutator->WithDefRegionKind(kTVMFFIDefRegionKindNone, @@ -215,7 +213,7 @@ TVM_FFI_INLINE ffi::Expected> DataflowVarMutate( TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, mapped_ty, std::move(mapped_ty_result)); if (!mapped_ty.UnchangedOrSameAs(self->ty)) { - ffi::ObjectPtr copy = ffi::make_object(*self); + ffi::ObjectPtr copy = ffi::make_object(*self); copy->ty = std::move(mapped_ty).ValueUnchecked(); mapped_value = ffi::Any(std::move(copy)); result = mapped_value; @@ -233,11 +231,11 @@ TVM_FFI_INLINE ffi::Expected> DataflowVarMutate( TVM_FFI_INLINE ffi::Expected> DataflowVarMaybeInplaceMutate( ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept { - DataflowVarNode* self = const_cast( - ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value)); + DataflowVarObj* self = const_cast( + ffi::detail::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value)); ffi::Expected remap_result = mutator->VarRemapGetExpected(value); TVM_FFI_S_MUTATE_MAYBE_EARLY_RETURN(remap_result); - if (ffi::details::ExpectedUnsafe::GetData(remap_result).type_index() != + if (ffi::detail::ExpectedUnsafe::GetData(remap_result).type_index() != ffi::TypeIndex::kTVMFFINone) { return std::move(remap_result); } @@ -246,7 +244,7 @@ TVM_FFI_INLINE ffi::Expected> DataflowVarMaybeInplace } ffi::UnchangedOr result = ffi::Unchanged(); ffi::Any mapped_value; - if (!self->ty.as()) { + if (!self->ty.as()) { ffi::Expected> mapped_ty_result = mutator->def_region_kind() == kTVMFFIDefRegionKindSimple ? mutator->WithDefRegionKind( @@ -273,8 +271,8 @@ TVM_FFI_INLINE ffi::Expected> DataflowVarMaybeInplace TVM_FFI_INLINE ffi::Expected> SeqExprVisit( ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept { - const SeqExprNode* self = - ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value); + const SeqExprObj* self = + ffi::detail::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value); TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->blocks)); TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->ty)); TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->body)); @@ -283,8 +281,8 @@ TVM_FFI_INLINE ffi::Expected> SeqExprVisit( TVM_FFI_INLINE ffi::Expected> SeqExprMutate( ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept { - const SeqExprNode* self = - ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value); + const SeqExprObj* self = + ffi::detail::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value); TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr>, mapped_blocks, mutator->MutateExpected(self->blocks)); TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, mapped_ty, @@ -295,7 +293,7 @@ TVM_FFI_INLINE ffi::Expected> SeqExprMutate( mapped_body.UnchangedOrSameAs(self->body)) { return ffi::Unchanged(); } - ffi::ObjectPtr copy = ffi::make_object(*self); + ffi::ObjectPtr copy = ffi::make_object(*self); copy->blocks = std::move(mapped_blocks).ValueOrUnchanged(std::move(copy->blocks)); copy->ty = std::move(mapped_ty).ValueOrUnchanged(std::move(copy->ty)); copy->body = std::move(mapped_body).ValueOrUnchanged(std::move(copy->body)); @@ -304,8 +302,8 @@ TVM_FFI_INLINE ffi::Expected> SeqExprMutate( TVM_FFI_INLINE ffi::Expected> SeqExprMaybeInplaceMutate( ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept { - SeqExprNode* self = const_cast( - ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value)); + SeqExprObj* self = const_cast( + ffi::detail::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value)); TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN( ffi::UnchangedOr>, mapped_blocks, mutator->MutateExpected(self->blocks, ffi::InplaceMode::kAllow)); @@ -328,8 +326,8 @@ TVM_FFI_INLINE ffi::Expected> SeqExprMaybeInplaceMuta TVM_FFI_INLINE ffi::Expected> FunctionVisit( ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept { // skips: attrs (metadata), is_pure (scalar) - const FunctionNode* self = - ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value); + const FunctionObj* self = + ffi::detail::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value); TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->WithDefRegionKind( kTVMFFIDefRegionKindPattern, [&]() { return visitor->VisitExpected(self->params); })); TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->ty)); @@ -341,8 +339,8 @@ TVM_FFI_INLINE ffi::Expected> FunctionVisit( TVM_FFI_INLINE ffi::Expected> FunctionMutate( ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept { // skips: attrs (metadata), is_pure (scalar) - const FunctionNode* self = - ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value); + const FunctionObj* self = + ffi::detail::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value); TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr>, mapped_params, mutator->WithDefRegionKind(kTVMFFIDefRegionKindPattern, [&]() { return mutator->MutateExpected(self->params); @@ -357,7 +355,7 @@ TVM_FFI_INLINE ffi::Expected> FunctionMutate( mapped_body.UnchangedOrSameAs(self->body) && mapped_ret_ty.UnchangedOrSameAs(self->ret_ty)) { return ffi::Unchanged(); } - ffi::ObjectPtr copy = ffi::make_object(*self); + ffi::ObjectPtr copy = ffi::make_object(*self); copy->params = std::move(mapped_params).ValueOrUnchanged(std::move(copy->params)); copy->ty = std::move(mapped_ty).ValueOrUnchanged(std::move(copy->ty)); copy->body = std::move(mapped_body).ValueOrUnchanged(std::move(copy->body)); @@ -368,8 +366,8 @@ TVM_FFI_INLINE ffi::Expected> FunctionMutate( TVM_FFI_INLINE ffi::Expected> FunctionMaybeInplaceMutate( ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept { // skips: attrs (metadata), is_pure (scalar) - FunctionNode* self = const_cast( - ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value)); + FunctionObj* self = const_cast( + ffi::detail::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value)); TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr>, mapped_params, mutator->WithDefRegionKind(kTVMFFIDefRegionKindPattern, [&]() { return mutator->MutateExpected(self->params, @@ -397,20 +395,20 @@ TVM_FFI_INLINE ffi::Expected> FunctionMaybeInplaceMut } // namespace -TVM_FFI_STATIC_INIT_BLOCK() { BindingNode::RegisterReflection(); } +TVM_FFI_STATIC_INIT_BLOCK() { BindingObj::RegisterReflection(); } IfExpr::IfExpr(Expr cond, Expr true_branch, Expr false_branch, Location loc) : Expr(ffi::UnsafeInit{}) { - ffi::ObjectPtr n = ffi::make_object( - std::move(cond), std::move(true_branch), std::move(false_branch)); + ffi::ObjectPtr n = + ffi::make_object(std::move(cond), std::move(true_branch), std::move(false_branch)); n->loc = loc; data_ = std::move(n); } TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - IfExprNode::RegisterReflection(); - refl::TypeAttrDef() + IfExprObj::RegisterReflection(); + refl::TypeAttrDef() .attr(ffi::type_attr::kStructuralVisit, ffi::FStructuralVisit::FromNative<&IfExprVisit>()) .attr(ffi::type_attr::kStructuralMutate, ffi::FStructuralMutate::FromNative<&IfExprMutate>()) .attr(ffi::type_attr::kStructuralMaybeInplaceMutate, @@ -423,10 +421,10 @@ TVM_FFI_STATIC_INIT_BLOCK() { } ShapeExpr::ShapeExpr(ffi::Array values, Location loc) : Expr(ffi::UnsafeInit{}) { - ffi::ObjectPtr n = ffi::make_object(); + ffi::ObjectPtr n = ffi::make_object(); n->values = values.Map([](PrimExpr value) { - if (value->IsInstance()) { + if (value->IsInstance()) { return tvm::prim::cast(value, PrimType::Int(64)); } TVM_FFI_ICHECK(value.ty().MatchesElementType(DLDataTypeCode::kDLInt, 64)) @@ -440,8 +438,8 @@ ShapeExpr::ShapeExpr(ffi::Array values, Location loc) : Expr(ffi::Unsa TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - ShapeExprNode::RegisterReflection(); - refl::TypeAttrDef() + ShapeExprObj::RegisterReflection(); + refl::TypeAttrDef() .attr(ffi::type_attr::kStructuralVisit, ffi::FStructuralVisit::FromNative<&ShapeExprVisit>()) .attr(ffi::type_attr::kStructuralMutate, ffi::FStructuralMutate::FromNative<&ShapeExprMutate>()) @@ -455,7 +453,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { DataflowVar::DataflowVar(ffi::String name, ffi::Optional ty_annotation, Location loc) : Var(ffi::UnsafeInit{}) { - ffi::ObjectPtr n = ffi::make_object(); + ffi::ObjectPtr n = ffi::make_object(); n->name = std::move(name); if (ty_annotation.has_value()) { n->ty = ty_annotation.value(); @@ -466,8 +464,8 @@ DataflowVar::DataflowVar(ffi::String name, ffi::Optional ty_annotation, Lo TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - DataflowVarNode::RegisterReflection(); - refl::TypeAttrDef() + DataflowVarObj::RegisterReflection(); + refl::TypeAttrDef() .attr(ffi::type_attr::kStructuralVisit, ffi::FStructuralVisit::FromNative<&DataflowVarVisit>()) .attr(ffi::type_attr::kStructuralMutate, @@ -500,8 +498,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { MatchCast::MatchCast(Var var, Expr value, Type ty, Location loc) { TVM_FFI_ICHECK(var.defined()) << "MatchCast requires var to be defined"; - ffi::ObjectPtr n = - ffi::make_object(std::move(var), std::move(value)); + ffi::ObjectPtr n = ffi::make_object(std::move(var), std::move(value)); n->ty = std::move(ty); n->loc = loc; data_ = std::move(n); @@ -509,7 +506,7 @@ MatchCast::MatchCast(Var var, Expr value, Type ty, Location loc) { TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - MatchCastNode::RegisterReflection(); + MatchCastObj::RegisterReflection(); refl::GlobalDef().def("relax.MatchCast", [](Var var, Expr value, Type ty, Location loc) { return MatchCast(var, value, ty, loc); @@ -517,15 +514,15 @@ TVM_FFI_STATIC_INIT_BLOCK() { } VarBinding::VarBinding(Var var, Expr value, Location loc) { - ffi::ObjectPtr n = - ffi::make_object(std::move(var), std::move(value)); + ffi::ObjectPtr n = + ffi::make_object(std::move(var), std::move(value)); n->loc = loc; data_ = std::move(n); } -bool VarBindingNode::SEqual(const VarBindingNode* other, - ffi::TypedFunction equal) const { - if (value->IsInstance()) { +bool VarBindingObj::SEqual(const VarBindingObj* other, + ffi::TypedFunction equal) const { + if (value->IsInstance()) { // Recursive function definitions may reference the bound variable // within the value being bound. In these cases, the // var comparison must occur first to define the var, to ensure it is @@ -538,10 +535,10 @@ bool VarBindingNode::SEqual(const VarBindingNode* other, } } -int64_t VarBindingNode::SHash(int64_t init_hash, - ffi::TypedFunction hash) const { +int64_t VarBindingObj::SHash(int64_t init_hash, + ffi::TypedFunction hash) const { int64_t hash_value = init_hash; - if (value->IsInstance()) { + if (value->IsInstance()) { hash_value = hash(var, hash_value, true); hash_value = hash(value, hash_value, false); } else { @@ -553,7 +550,7 @@ int64_t VarBindingNode::SHash(int64_t init_hash, TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - VarBindingNode::RegisterReflection(); + VarBindingObj::RegisterReflection(); refl::GlobalDef().def("relax.VarBinding", [](Var var, Expr value, Location loc) { return VarBinding(var, value, loc); @@ -561,7 +558,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { } BindingBlock::BindingBlock(ffi::Array bindings, Location loc) { - ffi::ObjectPtr n = ffi::make_object(); + ffi::ObjectPtr n = ffi::make_object(); n->bindings = std::move(bindings); n->loc = loc; data_ = std::move(n); @@ -569,14 +566,14 @@ BindingBlock::BindingBlock(ffi::Array bindings, Location loc) { TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - BindingBlockNode::RegisterReflection(); + BindingBlockObj::RegisterReflection(); refl::GlobalDef().def("relax.BindingBlock", [](ffi::Array bindings, Location loc) { return BindingBlock(bindings, loc); }); } -BindingBlockNode* BindingBlock::CopyOnWrite() { +BindingBlockObj* BindingBlock::CopyOnWrite() { // The `TVM_DEFINE_OBJECT_REF_COW_METHOD` cannot be used for // BindingBlock, because it is the base class for `DataflowBlock`. // If the `TVM_DEFINE_OBJECT_REF_COW_METHOD` were used, the @@ -584,20 +581,20 @@ BindingBlockNode* BindingBlock::CopyOnWrite() { // `DataflowBlock` to a `BindingBlock`. TVM_FFI_ICHECK(data_ != nullptr); if (!data_.unique()) { - ffi::ObjectPtr node; - if (auto dataflow_block = as()) { - node = ffi::make_object(*dataflow_block); + ffi::ObjectPtr node; + if (auto dataflow_block = as()) { + node = ffi::make_object(*dataflow_block); } else { - node = ffi::make_object(*(operator->())); + node = ffi::make_object(*(operator->())); } ffi::ObjectPtr(std::move(node)).swap(data_); } - return static_cast(data_.get()); + return static_cast(data_.get()); } DataflowBlock::DataflowBlock(ffi::Array bindings, Location loc) : BindingBlock(ffi::UnsafeInit{}) { - ffi::ObjectPtr n = ffi::make_object(); + ffi::ObjectPtr n = ffi::make_object(); n->bindings = std::move(bindings); n->loc = loc; data_ = std::move(n); @@ -605,7 +602,7 @@ DataflowBlock::DataflowBlock(ffi::Array bindings, Location loc) TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - DataflowBlockNode::RegisterReflection(); + DataflowBlockObj::RegisterReflection(); refl::GlobalDef().def("relax.DataflowBlock", [](ffi::Array bindings, Location loc) { return DataflowBlock(bindings, loc); @@ -622,7 +619,7 @@ SeqExpr::SeqExpr(Expr body) : Expr(ffi::UnsafeInit{}) { SeqExpr::SeqExpr(ffi::Array blocks, Expr body, Location loc) : Expr(ffi::UnsafeInit{}) { - ffi::ObjectPtr n = ffi::make_object(std::move(body)); + ffi::ObjectPtr n = ffi::make_object(std::move(body)); n->blocks = std::move(blocks); n->loc = loc; data_ = std::move(n); @@ -630,8 +627,8 @@ SeqExpr::SeqExpr(ffi::Array blocks, Expr body, Location loc) TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - SeqExprNode::RegisterReflection(); - refl::TypeAttrDef() + SeqExprObj::RegisterReflection(); + refl::TypeAttrDef() .attr(ffi::type_attr::kStructuralVisit, ffi::FStructuralVisit::FromNative<&SeqExprVisit>()) .attr(ffi::type_attr::kStructuralMutate, ffi::FStructuralMutate::FromNative<&SeqExprMutate>()) .attr(ffi::type_attr::kStructuralMaybeInplaceMutate, @@ -693,7 +690,7 @@ Function::Function(ffi::Array params, Expr body, ffi::Optional ret_ty FuncType func_ty(param_ty, ret_ty.value(), is_pure); // set the fields - ffi::ObjectPtr n = ffi::make_object(std::move(body)); + ffi::ObjectPtr n = ffi::make_object(std::move(body)); n->params = std::move(params); n->ret_ty = ret_ty.value(); n->is_pure = is_pure; @@ -705,8 +702,8 @@ Function::Function(ffi::Array params, Expr body, ffi::Optional ret_ty TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - FunctionNode::RegisterReflection(); - refl::TypeAttrDef() + FunctionObj::RegisterReflection(); + refl::TypeAttrDef() .attr(ffi::type_attr::kStructuralVisit, ffi::FStructuralVisit::FromNative<&FunctionVisit>()) .attr(ffi::type_attr::kStructuralMutate, ffi::FStructuralMutate::FromNative<&FunctionMutate>()) @@ -740,7 +737,7 @@ Function Function::CreateEmpty(ffi::Array params, Type ret_ty, bool is_pure }(); // set the fields - ffi::ObjectPtr n = ffi::make_object(std::move(body)); + ffi::ObjectPtr n = ffi::make_object(std::move(body)); n->params = std::move(params); n->is_pure = is_pure; n->ty = std::move(finfo); @@ -764,7 +761,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; auto infer_by_ty_args = [](const Call& call, const BlockBuilder& ctx) -> Type { - TVM_FFI_ICHECK(call->ty_args.defined()) << "ty_args field of CallNode should always be defined"; + TVM_FFI_ICHECK(call->ty_args.defined()) << "ty_args field of CallObj should always be defined"; if (call->ty_args.empty()) { return AnyType(); } else if (call->ty_args.size() == 1) { @@ -788,11 +785,11 @@ ExternFunc::ExternFunc(ffi::String global_symbol, Location loc) ExternFunc::ExternFunc(ffi::String global_symbol, Type ty, Location loc) : BaseFunc(ffi::UnsafeInit{}) { - TVM_FFI_ICHECK(ty.as()) + TVM_FFI_ICHECK(ty.as()) << "ExternFunc must have FuncType, " << "but declaration of '" << global_symbol << "' received " << ty; - ffi::ObjectPtr n = ffi::make_object(); + ffi::ObjectPtr n = ffi::make_object(); n->global_symbol = std::move(global_symbol); n->loc = loc; n->ty = ty; @@ -801,14 +798,14 @@ ExternFunc::ExternFunc(ffi::String global_symbol, Type ty, Location loc) TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - ExternFuncNode::RegisterReflection(); - refl::TypeAttrDef() + ExternFuncObj::RegisterReflection(); + refl::TypeAttrDef() .attr(ffi::type_attr::kStructuralVisit, - ffi::FStructuralVisit::FromNative<&TypeOnlyExprVisit>()) + ffi::FStructuralVisit::FromNative<&TypeOnlyExprVisit>()) .attr(ffi::type_attr::kStructuralMutate, - ffi::FStructuralMutate::FromNative<&TypeOnlyExprMutate>()) + ffi::FStructuralMutate::FromNative<&TypeOnlyExprMutate>()) .attr(ffi::type_attr::kStructuralMaybeInplaceMutate, - ffi::FStructuralMutate::FromNative<&TypeOnlyExprMaybeInplaceMutate>()); + ffi::FStructuralMutate::FromNative<&TypeOnlyExprMaybeInplaceMutate>()); refl::GlobalDef().def("relax.ExternFunc", [](ffi::String global_symbol, ffi::Optional ty, Location loc) { @@ -824,7 +821,7 @@ Expr GetShapeOf(const Expr& expr) { // default case, to be normalized. TVM_FFI_ICHECK(!expr->ty.as().has_value()) << "GetShapeOf can only be applied to normalized expr"; - auto* tinfo = GetTypeAs(expr); + auto* tinfo = GetTypeAs(expr); TVM_FFI_ICHECK(tinfo != nullptr) << "ShapeOf can only be applied to expr with TensorType"; if (tinfo->shape.has_value()) return tinfo->shape.value(); @@ -842,20 +839,20 @@ TVM_FFI_STATIC_INIT_BLOCK() { .def("relax.GetShapeOf", [](const Expr& expr) { return GetShapeOf(expr); }) .def("relax.FuncWithAttr", [](BaseFunc func, ffi::String key, ffi::ObjectRef value) -> ffi::Optional { - if (func->IsInstance()) { + if (func->IsInstance()) { return WithAttr(std::move(func).as_or_throw(), key, value); } return std::nullopt; }) .def("relax.FuncWithAttrs", [](BaseFunc func, ffi::Map attr_map) -> ffi::Optional { - if (func->IsInstance()) { + if (func->IsInstance()) { return WithAttrs(std::move(func).as_or_throw(), attr_map); } return std::nullopt; }) .def("relax.FuncWithoutAttr", [](BaseFunc func, ffi::String key) -> ffi::Optional { - if (func->IsInstance()) { + if (func->IsInstance()) { return WithoutAttr(std::move(func).as_or_throw(), key); } return std::nullopt; diff --git a/src/relax/ir/expr_functor.cc b/src/relax/ir/expr_functor.cc index 059bfd31358f..7f92d507f33f 100644 --- a/src/relax/ir/expr_functor.cc +++ b/src/relax/ir/expr_functor.cc @@ -32,33 +32,33 @@ #include // functions to be overriden. -#define RELAX_VISIT_BINDING_DISPATCH(OP) \ - vtable.template SetDispatch( \ - [](const ffi::ObjectRef& n, TSelf* self, const VarBindingNode* binding) { \ - self->VisitBinding_(binding, static_cast(n.get())); \ +#define RELAX_VISIT_BINDING_DISPATCH(OP) \ + vtable.template SetDispatch( \ + [](const ffi::ObjectRef& n, TSelf* self, const VarBindingObj* binding) { \ + self->VisitBinding_(binding, static_cast(n.get())); \ }); #define RELAX_VAR_BINDING_DISPATCH_IMPL(Type) \ Type::VisitBindingVTable Type::InitVisitBindingVTable() { \ VisitBindingVTable vtable; \ - RELAX_VISIT_BINDING_DISPATCH(GenericConstNode); \ - RELAX_VISIT_BINDING_DISPATCH(TupleNode); \ - RELAX_VISIT_BINDING_DISPATCH(VarNode); \ - RELAX_VISIT_BINDING_DISPATCH(DataflowVarNode); \ - RELAX_VISIT_BINDING_DISPATCH(ShapeExprNode); \ - RELAX_VISIT_BINDING_DISPATCH(ExternFuncNode); \ - RELAX_VISIT_BINDING_DISPATCH(GlobalVarNode); \ - RELAX_VISIT_BINDING_DISPATCH(FunctionNode); \ - RELAX_VISIT_BINDING_DISPATCH(CallNode); \ - RELAX_VISIT_BINDING_DISPATCH(SeqExprNode); \ - RELAX_VISIT_BINDING_DISPATCH(IfExprNode); \ - RELAX_VISIT_BINDING_DISPATCH(OpNode); \ - RELAX_VISIT_BINDING_DISPATCH(TupleGetItemNode); \ - RELAX_VISIT_BINDING_DISPATCH(StringImmNode); \ - RELAX_VISIT_BINDING_DISPATCH(DataTypeImmNode); \ + RELAX_VISIT_BINDING_DISPATCH(GenericConstObj); \ + RELAX_VISIT_BINDING_DISPATCH(TupleObj); \ + RELAX_VISIT_BINDING_DISPATCH(VarObj); \ + RELAX_VISIT_BINDING_DISPATCH(DataflowVarObj); \ + RELAX_VISIT_BINDING_DISPATCH(ShapeExprObj); \ + RELAX_VISIT_BINDING_DISPATCH(ExternFuncObj); \ + RELAX_VISIT_BINDING_DISPATCH(GlobalVarObj); \ + RELAX_VISIT_BINDING_DISPATCH(FunctionObj); \ + RELAX_VISIT_BINDING_DISPATCH(CallObj); \ + RELAX_VISIT_BINDING_DISPATCH(SeqExprObj); \ + RELAX_VISIT_BINDING_DISPATCH(IfExprObj); \ + RELAX_VISIT_BINDING_DISPATCH(OpObj); \ + RELAX_VISIT_BINDING_DISPATCH(TupleGetItemObj); \ + RELAX_VISIT_BINDING_DISPATCH(StringImmObj); \ + RELAX_VISIT_BINDING_DISPATCH(DataTypeImmObj); \ return vtable; \ } \ - void Type::VisitBinding_(const VarBindingNode* binding) { \ + void Type::VisitBinding_(const VarBindingObj* binding) { \ static VisitBindingVTable vtable = InitVisitBindingVTable(); \ const Expr& value = binding->value; \ TVM_FFI_ICHECK(value.defined()) << "Found null pointer node while traversing AST."; \ @@ -70,17 +70,17 @@ } // functions to be overriden. -#define RELAX_EXPR_VISITOR_VISIT_BINDING_IMPL(OP) \ - void ExprVisitor::VisitBinding_(const VarBindingNode* binding, const OP* value) { \ - this->VisitExpr(binding->value); \ - this->VisitVarDef(binding->var); \ +#define RELAX_EXPR_VISITOR_VISIT_BINDING_IMPL(OP) \ + void ExprVisitor::VisitBinding_(const VarBindingObj* binding, const OP* value) { \ + this->VisitExpr(binding->value); \ + this->VisitVarDef(binding->var); \ } // functions to be overriden. -#define RELAX_EXPR_MUTATOR_VISIT_BINDING_IMPL(OP) \ - void ExprMutator::VisitBinding_(const VarBindingNode* binding, const OP* value) { \ - Expr new_value = this->VisitExpr(binding->value); \ - this->ReEmitBinding(binding, new_value); \ +#define RELAX_EXPR_MUTATOR_VISIT_BINDING_IMPL(OP) \ + void ExprMutator::VisitBinding_(const VarBindingObj* binding, const OP* value) { \ + Expr new_value = this->VisitExpr(binding->value); \ + this->ReEmitBinding(binding, new_value); \ } namespace tvm { @@ -107,14 +107,14 @@ void ExprVisitor::DefaultTypeFieldVisitor::VisitTypeExprField(const PrimExpr& ex parent_->VisitTypePrimExprField(expr); } -void ExprVisitor::DefaultTypeFieldVisitor::VisitType_(const FuncTypeNode* op) { +void ExprVisitor::DefaultTypeFieldVisitor::VisitType_(const FuncTypeObj* op) { // Do not recurse into function type // as they won't contain ref to values in current scope. } void VisitExprDepTypeFieldIfNeeded(ExprVisitor* visitor, const Type& ty) { if (!ty.as().has_value()) { - auto* ty_node = ty.as(); + auto* ty_node = ty.as(); TVM_FFI_DCHECK(ty_node != nullptr); visitor->VisitExprDepTypeField(ffi::GetRef(ty_node)); } @@ -122,17 +122,17 @@ void VisitExprDepTypeFieldIfNeeded(ExprVisitor* visitor, const Type& ty) { void ExprVisitor::VisitExpr(const Expr& expr) { ExprFunctor::VisitExpr(expr); } -void ExprVisitor::VisitExpr_(const GenericConstNode* op) { +void ExprVisitor::VisitExpr_(const GenericConstObj* op) { this->VisitLoc(op->loc); // GenericConst's Type does not depend on Expr. } -void ExprVisitor::VisitExpr_(const GlobalVarNode* op) { +void ExprVisitor::VisitExpr_(const GlobalVarObj* op) { this->VisitLoc(op->loc); // FuncType is not value-dep } -void ExprVisitor::VisitExpr_(const TupleNode* op) { +void ExprVisitor::VisitExpr_(const TupleObj* op) { this->VisitLoc(op->loc); for (Expr field : op->fields) { this->VisitExpr(field); @@ -141,17 +141,17 @@ void ExprVisitor::VisitExpr_(const TupleNode* op) { } // Visit the use-site of a defined Var -void ExprVisitor::VisitExpr_(const VarNode* op) { +void ExprVisitor::VisitExpr_(const VarObj* op) { this->VisitLoc(op->loc); VisitExprDepTypeFieldIfNeeded(this, op->ty); } // Visit the use-site of a defined DataflowVar -void ExprVisitor::VisitExpr_(const DataflowVarNode* op) { - VisitExpr_(static_cast(op)); +void ExprVisitor::VisitExpr_(const DataflowVarObj* op) { + VisitExpr_(static_cast(op)); } -void ExprVisitor::VisitExpr_(const FunctionNode* op) { +void ExprVisitor::VisitExpr_(const FunctionObj* op) { this->VisitLoc(op->loc); for (Var param : op->params) { this->VisitVarDef(param); @@ -161,7 +161,7 @@ void ExprVisitor::VisitExpr_(const FunctionNode* op) { // FuncType does not depend on Expr. } -void ExprVisitor::VisitExpr_(const CallNode* op) { +void ExprVisitor::VisitExpr_(const CallObj* op) { this->VisitLoc(op->loc); this->VisitExpr(op->op); @@ -176,7 +176,7 @@ void ExprVisitor::VisitExpr_(const CallNode* op) { VisitExprDepTypeFieldIfNeeded(this, op->ty); } -void ExprVisitor::VisitExpr_(const TensorLoadNode* op) { +void ExprVisitor::VisitExpr_(const TensorLoadObj* op) { this->VisitLoc(op->loc); for (const PrimExpr& index : op->indices) { this->VisitExpr(index); @@ -192,50 +192,50 @@ void ExprVisitor::VisitExpr_(const TensorLoadNode* op) { VisitExprDepTypeFieldIfNeeded(this, op->ty); \ } -RELAX_VISIT_TIRX_BINOP(AddNode); -RELAX_VISIT_TIRX_BINOP(LShiftNode); -RELAX_VISIT_TIRX_BINOP(RShiftNode); -RELAX_VISIT_TIRX_BINOP(BitwiseAndNode); -RELAX_VISIT_TIRX_BINOP(BitwiseOrNode); -RELAX_VISIT_TIRX_BINOP(BitwiseXorNode); -RELAX_VISIT_TIRX_BINOP(SubNode); -RELAX_VISIT_TIRX_BINOP(MulNode); -RELAX_VISIT_TIRX_BINOP(DivNode); -RELAX_VISIT_TIRX_BINOP(ModNode); -RELAX_VISIT_TIRX_BINOP(FloorDivNode); -RELAX_VISIT_TIRX_BINOP(FloorModNode); -RELAX_VISIT_TIRX_BINOP(MinNode); -RELAX_VISIT_TIRX_BINOP(MaxNode); -RELAX_VISIT_TIRX_BINOP(EQNode); -RELAX_VISIT_TIRX_BINOP(NENode); -RELAX_VISIT_TIRX_BINOP(LTNode); -RELAX_VISIT_TIRX_BINOP(LENode); -RELAX_VISIT_TIRX_BINOP(GTNode); -RELAX_VISIT_TIRX_BINOP(GENode); -RELAX_VISIT_TIRX_BINOP(AndNode); -RELAX_VISIT_TIRX_BINOP(OrNode); +RELAX_VISIT_TIRX_BINOP(AddObj); +RELAX_VISIT_TIRX_BINOP(LShiftObj); +RELAX_VISIT_TIRX_BINOP(RShiftObj); +RELAX_VISIT_TIRX_BINOP(BitwiseAndObj); +RELAX_VISIT_TIRX_BINOP(BitwiseOrObj); +RELAX_VISIT_TIRX_BINOP(BitwiseXorObj); +RELAX_VISIT_TIRX_BINOP(SubObj); +RELAX_VISIT_TIRX_BINOP(MulObj); +RELAX_VISIT_TIRX_BINOP(DivObj); +RELAX_VISIT_TIRX_BINOP(ModObj); +RELAX_VISIT_TIRX_BINOP(FloorDivObj); +RELAX_VISIT_TIRX_BINOP(FloorModObj); +RELAX_VISIT_TIRX_BINOP(MinObj); +RELAX_VISIT_TIRX_BINOP(MaxObj); +RELAX_VISIT_TIRX_BINOP(EQObj); +RELAX_VISIT_TIRX_BINOP(NEObj); +RELAX_VISIT_TIRX_BINOP(LTObj); +RELAX_VISIT_TIRX_BINOP(LEObj); +RELAX_VISIT_TIRX_BINOP(GTObj); +RELAX_VISIT_TIRX_BINOP(GEObj); +RELAX_VISIT_TIRX_BINOP(AndObj); +RELAX_VISIT_TIRX_BINOP(OrObj); #undef RELAX_VISIT_TIRX_BINOP -void ExprVisitor::VisitExpr_(const prim::CastNode* op) { +void ExprVisitor::VisitExpr_(const prim::CastObj* op) { this->VisitLoc(op->loc); this->VisitExpr(op->value); VisitExprDepTypeFieldIfNeeded(this, op->ty); } -void ExprVisitor::VisitExpr_(const prim::NotNode* op) { +void ExprVisitor::VisitExpr_(const prim::NotObj* op) { this->VisitLoc(op->loc); this->VisitExpr(op->a); VisitExprDepTypeFieldIfNeeded(this, op->ty); } -void ExprVisitor::VisitExpr_(const prim::BitwiseNotNode* op) { +void ExprVisitor::VisitExpr_(const prim::BitwiseNotObj* op) { this->VisitLoc(op->loc); this->VisitExpr(op->a); VisitExprDepTypeFieldIfNeeded(this, op->ty); } -void ExprVisitor::VisitExpr_(const prim::SelectNode* op) { +void ExprVisitor::VisitExpr_(const prim::SelectObj* op) { this->VisitLoc(op->loc); this->VisitExpr(op->condition); this->VisitExpr(op->true_value); @@ -243,7 +243,7 @@ void ExprVisitor::VisitExpr_(const prim::SelectNode* op) { VisitExprDepTypeFieldIfNeeded(this, op->ty); } -void ExprVisitor::VisitExpr_(const prim::RampNode* op) { +void ExprVisitor::VisitExpr_(const prim::RampObj* op) { this->VisitLoc(op->loc); this->VisitExpr(op->base); this->VisitExpr(op->stride); @@ -251,14 +251,14 @@ void ExprVisitor::VisitExpr_(const prim::RampNode* op) { VisitExprDepTypeFieldIfNeeded(this, op->ty); } -void ExprVisitor::VisitExpr_(const prim::BroadcastNode* op) { +void ExprVisitor::VisitExpr_(const prim::BroadcastObj* op) { this->VisitLoc(op->loc); this->VisitExpr(op->value); this->VisitExpr(op->lanes); VisitExprDepTypeFieldIfNeeded(this, op->ty); } -void ExprVisitor::VisitExpr_(const prim::ShuffleNode* op) { +void ExprVisitor::VisitExpr_(const prim::ShuffleObj* op) { this->VisitLoc(op->loc); for (const PrimExpr& vector : op->vectors) { this->VisitExpr(vector); @@ -269,11 +269,11 @@ void ExprVisitor::VisitExpr_(const prim::ShuffleNode* op) { VisitExprDepTypeFieldIfNeeded(this, op->ty); } -void ExprVisitor::VisitExpr_(const tvm::IntImmNode* op) { this->VisitLoc(op->loc); } +void ExprVisitor::VisitExpr_(const tvm::IntImmObj* op) { this->VisitLoc(op->loc); } -void ExprVisitor::VisitExpr_(const tvm::FloatImmNode* op) { this->VisitLoc(op->loc); } +void ExprVisitor::VisitExpr_(const tvm::FloatImmObj* op) { this->VisitLoc(op->loc); } -void ExprVisitor::VisitExpr_(const IfExprNode* op) { +void ExprVisitor::VisitExpr_(const IfExprObj* op) { this->VisitLoc(op->loc); this->VisitExpr(op->cond); this->VisitExpr(op->true_branch); @@ -282,28 +282,28 @@ void ExprVisitor::VisitExpr_(const IfExprNode* op) { VisitExprDepTypeFieldIfNeeded(this, op->ty); } -void ExprVisitor::VisitExpr_(const OpNode* op) { this->VisitLoc(op->loc); } +void ExprVisitor::VisitExpr_(const OpObj* op) { this->VisitLoc(op->loc); } -void ExprVisitor::VisitExpr_(const TupleGetItemNode* op) { +void ExprVisitor::VisitExpr_(const TupleGetItemObj* op) { this->VisitLoc(op->loc); this->VisitExpr(op->tuple); VisitExprDepTypeFieldIfNeeded(this, op->ty); } -void ExprVisitor::VisitExpr_(const ShapeExprNode* op) { +void ExprVisitor::VisitExpr_(const ShapeExprObj* op) { for (PrimExpr val : op->values) { this->VisitExpr(val); } this->VisitLoc(op->loc); } -void ExprVisitor::VisitExpr_(const ExternFuncNode* op) { +void ExprVisitor::VisitExpr_(const ExternFuncObj* op) { this->VisitLoc(op->loc); // FuncType does not depend on Expr. } -void ExprVisitor::VisitExpr_(const SeqExprNode* op) { +void ExprVisitor::VisitExpr_(const SeqExprObj* op) { this->VisitLoc(op->loc); for (BindingBlock block : op->blocks) { this->VisitBindingBlock(block); @@ -313,15 +313,15 @@ void ExprVisitor::VisitExpr_(const SeqExprNode* op) { VisitExprDepTypeFieldIfNeeded(this, op->ty); } -void ExprVisitor::VisitExprFallback_(const ExprNode* op) { - if (op->IsInstance()) { +void ExprVisitor::VisitExprFallback_(const ExprObj* op) { + if (op->IsInstance()) { return; } this->VisitExprDefault_(op); } -void ExprVisitor::VisitExpr_(const StringImmNode* op) { this->VisitLoc(op->loc); } -void ExprVisitor::VisitExpr_(const DataTypeImmNode* op) { this->VisitLoc(op->loc); } +void ExprVisitor::VisitExpr_(const StringImmObj* op) { this->VisitLoc(op->loc); } +void ExprVisitor::VisitExpr_(const DataTypeImmObj* op) { this->VisitLoc(op->loc); } void ExprVisitor::VisitLoc(const Location& loc) {} @@ -329,51 +329,51 @@ void ExprVisitor::VisitTypePrimExprField(const PrimExpr& expr) { this->VisitExpr // implementations of binding visitor dispatch RELAX_VAR_BINDING_DISPATCH_IMPL(ExprVisitor); -RELAX_EXPR_VISITOR_VISIT_BINDING_IMPL(GenericConstNode); -RELAX_EXPR_VISITOR_VISIT_BINDING_IMPL(TupleNode); -RELAX_EXPR_VISITOR_VISIT_BINDING_IMPL(VarNode); -RELAX_EXPR_VISITOR_VISIT_BINDING_IMPL(DataflowVarNode); -RELAX_EXPR_VISITOR_VISIT_BINDING_IMPL(ShapeExprNode); -RELAX_EXPR_VISITOR_VISIT_BINDING_IMPL(ExternFuncNode); -RELAX_EXPR_VISITOR_VISIT_BINDING_IMPL(GlobalVarNode); -RELAX_EXPR_VISITOR_VISIT_BINDING_IMPL(FunctionNode); -RELAX_EXPR_VISITOR_VISIT_BINDING_IMPL(CallNode); -RELAX_EXPR_VISITOR_VISIT_BINDING_IMPL(SeqExprNode); -RELAX_EXPR_VISITOR_VISIT_BINDING_IMPL(IfExprNode); -RELAX_EXPR_VISITOR_VISIT_BINDING_IMPL(OpNode); -RELAX_EXPR_VISITOR_VISIT_BINDING_IMPL(TupleGetItemNode); -RELAX_EXPR_VISITOR_VISIT_BINDING_IMPL(ExprNode); -RELAX_EXPR_VISITOR_VISIT_BINDING_IMPL(StringImmNode); -RELAX_EXPR_VISITOR_VISIT_BINDING_IMPL(DataTypeImmNode); - -void ExprVisitor::VisitBinding_(const MatchCastNode* binding) { +RELAX_EXPR_VISITOR_VISIT_BINDING_IMPL(GenericConstObj); +RELAX_EXPR_VISITOR_VISIT_BINDING_IMPL(TupleObj); +RELAX_EXPR_VISITOR_VISIT_BINDING_IMPL(VarObj); +RELAX_EXPR_VISITOR_VISIT_BINDING_IMPL(DataflowVarObj); +RELAX_EXPR_VISITOR_VISIT_BINDING_IMPL(ShapeExprObj); +RELAX_EXPR_VISITOR_VISIT_BINDING_IMPL(ExternFuncObj); +RELAX_EXPR_VISITOR_VISIT_BINDING_IMPL(GlobalVarObj); +RELAX_EXPR_VISITOR_VISIT_BINDING_IMPL(FunctionObj); +RELAX_EXPR_VISITOR_VISIT_BINDING_IMPL(CallObj); +RELAX_EXPR_VISITOR_VISIT_BINDING_IMPL(SeqExprObj); +RELAX_EXPR_VISITOR_VISIT_BINDING_IMPL(IfExprObj); +RELAX_EXPR_VISITOR_VISIT_BINDING_IMPL(OpObj); +RELAX_EXPR_VISITOR_VISIT_BINDING_IMPL(TupleGetItemObj); +RELAX_EXPR_VISITOR_VISIT_BINDING_IMPL(ExprObj); +RELAX_EXPR_VISITOR_VISIT_BINDING_IMPL(StringImmObj); +RELAX_EXPR_VISITOR_VISIT_BINDING_IMPL(DataTypeImmObj); + +void ExprVisitor::VisitBinding_(const MatchCastObj* binding) { this->VisitExpr(binding->value); this->VisitExprDepTypeField(binding->ty); this->VisitVarDef(binding->var); } -void ExprVisitor::VisitBindingBlock_(const BindingBlockNode* block) { +void ExprVisitor::VisitBindingBlock_(const BindingBlockObj* block) { for (Binding binding : block->bindings) { this->VisitBinding(binding); } } -void ExprVisitor::VisitBindingBlock_(const DataflowBlockNode* block) { +void ExprVisitor::VisitBindingBlock_(const DataflowBlockObj* block) { for (Binding binding : block->bindings) { this->VisitBinding(binding); } } -void ExprVisitor::VisitVarDef_(const DataflowVarNode* var) { - VisitVarDef_(static_cast(var)); +void ExprVisitor::VisitVarDef_(const DataflowVarObj* var) { + VisitVarDef_(static_cast(var)); } -void ExprVisitor::VisitVarDef_(const VarNode* var) { this->VisitLoc(var->loc); } +void ExprVisitor::VisitVarDef_(const VarObj* var) { this->VisitLoc(var->loc); } void ExprVisitor::VisitBinding(const Binding& binding) { - if (const auto* node = binding.as()) { + if (const auto* node = binding.as()) { VisitBinding_(node); - } else if (const auto* node = binding.as()) { + } else if (const auto* node = binding.as()) { VisitBinding_(node); } else { TVM_FFI_THROW(TypeError) << "Invalid type: " << binding->GetTypeKey(); @@ -381,9 +381,9 @@ void ExprVisitor::VisitBinding(const Binding& binding) { } void ExprVisitor::VisitBindingBlock(const BindingBlock& block) { - if (const auto* node = block.as()) { + if (const auto* node = block.as()) { VisitBindingBlock_(node); - } else if (const auto* node = block.as()) { + } else if (const auto* node = block.as()) { VisitBindingBlock_(node); } else { TVM_FFI_THROW(TypeError) << "Invalid type: " << block->GetTypeKey(); @@ -391,9 +391,9 @@ void ExprVisitor::VisitBindingBlock(const BindingBlock& block) { } void ExprVisitor::VisitVarDef(const Var& var) { - if (const auto* node = var.as()) { + if (const auto* node = var.as()) { VisitVarDef_(node); - } else if (const auto* node = var.as()) { + } else if (const auto* node = var.as()) { VisitVarDef_(node); } else { TVM_FFI_THROW(TypeError) << "Invalid type: " << var->GetTypeKey(); @@ -444,7 +444,7 @@ PrimExpr ExprMutatorBase::DefaultTypeFieldMutator::VisitTypeExprField(const Prim return parent_->VisitTypePrimExprField(expr); } -Type ExprMutatorBase::DefaultTypeFieldMutator::VisitType_(const FuncTypeNode* op) { +Type ExprMutatorBase::DefaultTypeFieldMutator::VisitType_(const FuncTypeObj* op) { // Do not recurse into function type // as they won't contain ref to values in current scope. return ffi::GetRef(op); @@ -452,17 +452,17 @@ Type ExprMutatorBase::DefaultTypeFieldMutator::VisitType_(const FuncTypeNode* op Expr ExprMutatorBase::VisitExpr(const Expr& expr) { return ExprFunctor::VisitExpr(expr); } -Expr ExprMutatorBase::VisitExpr_(const GenericConstNode* op) { +Expr ExprMutatorBase::VisitExpr_(const GenericConstObj* op) { // GenericConst' type won't be affected by Expr/PrimExpr change. return ffi::GetRef(op); } -Expr ExprMutatorBase::VisitExpr_(const GlobalVarNode* op) { +Expr ExprMutatorBase::VisitExpr_(const GlobalVarObj* op) { // FuncType won't be affected by Expr/PrimExpr change. return ffi::GetRef(op); } -Expr ExprMutatorBase::VisitExpr_(const TupleNode* op) { +Expr ExprMutatorBase::VisitExpr_(const TupleObj* op) { bool unchanged = true; tvm::ffi::Array fields; for (Expr field : op->fields) { @@ -483,14 +483,14 @@ Expr ExprMutatorBase::VisitExpr_(const TupleNode* op) { } // Visit the use-site of a defined Var -Expr ExprMutatorBase::VisitExpr_(const VarNode* op) { return ffi::GetRef(op); } +Expr ExprMutatorBase::VisitExpr_(const VarObj* op) { return ffi::GetRef(op); } // Visit the use-site of a defined DataflowVar -Expr ExprMutatorBase::VisitExpr_(const DataflowVarNode* op) { - return VisitExpr_(static_cast(op)); +Expr ExprMutatorBase::VisitExpr_(const DataflowVarObj* op) { + return VisitExpr_(static_cast(op)); } -Expr ExprMutatorBase::VisitExpr_(const FunctionNode* op) { +Expr ExprMutatorBase::VisitExpr_(const FunctionObj* op) { // type of function is not value dependent // so no need to check ty field Expr body = this->VisitExpr(op->body); @@ -502,7 +502,7 @@ Expr ExprMutatorBase::VisitExpr_(const FunctionNode* op) { } } -Expr ExprMutatorBase::VisitExpr_(const CallNode* call_node) { +Expr ExprMutatorBase::VisitExpr_(const CallObj* call_node) { Expr new_op = this->VisitExpr(call_node->op); bool unchanged = call_node->op.same_as(new_op); @@ -539,7 +539,7 @@ Expr ExprMutatorBase::VisitExpr_(const CallNode* call_node) { } } -Expr ExprMutatorBase::VisitExpr_(const TensorLoadNode* op) { +Expr ExprMutatorBase::VisitExpr_(const TensorLoadObj* op) { ffi::Array indices = op->indices.Map( [this](const PrimExpr& e) { return this->VisitExpr(e).as_or_throw(); }); if (indices.same_as(op->indices)) { @@ -549,7 +549,7 @@ Expr ExprMutatorBase::VisitExpr_(const TensorLoadNode* op) { } #define RELAX_MUTATE_TIRX_BINOP(OP) \ - Expr ExprMutatorBase::VisitExpr_(const prim::OP##Node* op) { \ + Expr ExprMutatorBase::VisitExpr_(const prim::OP##Obj* op) { \ PrimExpr a = this->VisitExpr(op->a).as_or_throw(); \ PrimExpr b = this->VisitExpr(op->b).as_or_throw(); \ if (a.same_as(op->a) && b.same_as(op->b)) { \ @@ -583,24 +583,24 @@ RELAX_MUTATE_TIRX_BINOP(Or); #undef RELAX_MUTATE_TIRX_BINOP -Expr ExprMutatorBase::VisitExpr_(const prim::CastNode* op) { +Expr ExprMutatorBase::VisitExpr_(const prim::CastObj* op) { PrimExpr value = this->VisitExpr(op->value).as_or_throw(); return value.same_as(op->value) ? ffi::GetRef(op) : Expr(prim::Cast(op->ty.as_or_throw(), value, op->loc)); } -Expr ExprMutatorBase::VisitExpr_(const prim::NotNode* op) { +Expr ExprMutatorBase::VisitExpr_(const prim::NotObj* op) { PrimExpr a = this->VisitExpr(op->a).as_or_throw(); return a.same_as(op->a) ? ffi::GetRef(op) : Expr(prim::Not(a, op->loc)); } -Expr ExprMutatorBase::VisitExpr_(const prim::BitwiseNotNode* op) { +Expr ExprMutatorBase::VisitExpr_(const prim::BitwiseNotObj* op) { PrimExpr a = this->VisitExpr(op->a).as_or_throw(); return a.same_as(op->a) ? ffi::GetRef(op) : Expr(prim::BitwiseNot(a, op->loc)); } -Expr ExprMutatorBase::VisitExpr_(const prim::SelectNode* op) { +Expr ExprMutatorBase::VisitExpr_(const prim::SelectObj* op) { PrimExpr condition = this->VisitExpr(op->condition).as_or_throw(); PrimExpr true_value = this->VisitExpr(op->true_value).as_or_throw(); PrimExpr false_value = this->VisitExpr(op->false_value).as_or_throw(); @@ -611,7 +611,7 @@ Expr ExprMutatorBase::VisitExpr_(const prim::SelectNode* op) { return prim::Select(condition, true_value, false_value, op->loc); } -Expr ExprMutatorBase::VisitExpr_(const prim::RampNode* op) { +Expr ExprMutatorBase::VisitExpr_(const prim::RampObj* op) { PrimExpr base = this->VisitExpr(op->base).as_or_throw(); PrimExpr stride = this->VisitExpr(op->stride).as_or_throw(); PrimExpr lanes = this->VisitExpr(op->lanes).as_or_throw(); @@ -621,7 +621,7 @@ Expr ExprMutatorBase::VisitExpr_(const prim::RampNode* op) { return prim::Ramp(base, stride, lanes, op->loc); } -Expr ExprMutatorBase::VisitExpr_(const prim::BroadcastNode* op) { +Expr ExprMutatorBase::VisitExpr_(const prim::BroadcastObj* op) { PrimExpr value = this->VisitExpr(op->value).as_or_throw(); PrimExpr lanes = this->VisitExpr(op->lanes).as_or_throw(); if (value.same_as(op->value) && lanes.same_as(op->lanes)) { @@ -630,7 +630,7 @@ Expr ExprMutatorBase::VisitExpr_(const prim::BroadcastNode* op) { return prim::Broadcast(value, lanes, op->loc); } -Expr ExprMutatorBase::VisitExpr_(const prim::ShuffleNode* op) { +Expr ExprMutatorBase::VisitExpr_(const prim::ShuffleObj* op) { ffi::Array vectors = op->vectors.Map( [this](const PrimExpr& e) { return this->VisitExpr(e).as_or_throw(); }); ffi::Array indices = op->indices.Map( @@ -641,11 +641,11 @@ Expr ExprMutatorBase::VisitExpr_(const prim::ShuffleNode* op) { return prim::Shuffle(vectors, indices, op->loc); } -Expr ExprMutatorBase::VisitExpr_(const tvm::IntImmNode* op) { return ffi::GetRef(op); } +Expr ExprMutatorBase::VisitExpr_(const tvm::IntImmObj* op) { return ffi::GetRef(op); } -Expr ExprMutatorBase::VisitExpr_(const tvm::FloatImmNode* op) { return ffi::GetRef(op); } +Expr ExprMutatorBase::VisitExpr_(const tvm::FloatImmObj* op) { return ffi::GetRef(op); } -Expr ExprMutatorBase::VisitExpr_(const IfExprNode* op) { +Expr ExprMutatorBase::VisitExpr_(const IfExprObj* op) { Expr guard = this->VisitExpr(op->cond); Expr true_b = this->VisitExpr(op->true_branch); Expr false_b = this->VisitExpr(op->false_branch); @@ -657,9 +657,9 @@ Expr ExprMutatorBase::VisitExpr_(const IfExprNode* op) { } } -Expr ExprMutatorBase::VisitExpr_(const OpNode* op) { return ffi::GetRef(op); } +Expr ExprMutatorBase::VisitExpr_(const OpObj* op) { return ffi::GetRef(op); } -Expr ExprMutatorBase::VisitExpr_(const TupleGetItemNode* op) { +Expr ExprMutatorBase::VisitExpr_(const TupleGetItemObj* op) { auto t = this->VisitExpr(op->tuple); if (op->tuple.same_as(t)) { // type can be deterministically derived by tuple and index @@ -670,17 +670,17 @@ Expr ExprMutatorBase::VisitExpr_(const TupleGetItemNode* op) { } } -Expr ExprMutatorBase::VisitExprFallback_(const ExprNode* op) { - if (op->IsInstance()) { +Expr ExprMutatorBase::VisitExprFallback_(const ExprObj* op) { + if (op->IsInstance()) { return ffi::GetRef(op); } return this->VisitExprDefault_(op); } -Expr ExprMutatorBase::VisitExpr_(const StringImmNode* op) { return ffi::GetRef(op); } -Expr ExprMutatorBase::VisitExpr_(const DataTypeImmNode* op) { return ffi::GetRef(op); } +Expr ExprMutatorBase::VisitExpr_(const StringImmObj* op) { return ffi::GetRef(op); } +Expr ExprMutatorBase::VisitExpr_(const DataTypeImmObj* op) { return ffi::GetRef(op); } -Expr ExprMutatorBase::VisitExpr_(const ShapeExprNode* op) { +Expr ExprMutatorBase::VisitExpr_(const ShapeExprObj* op) { auto values = op->values.Map( [this](const PrimExpr& e) { return this->VisitExpr(e).as_or_throw(); }); @@ -692,12 +692,12 @@ Expr ExprMutatorBase::VisitExpr_(const ShapeExprNode* op) { } } -Expr ExprMutatorBase::VisitExpr_(const ExternFuncNode* op) { +Expr ExprMutatorBase::VisitExpr_(const ExternFuncObj* op) { // Type of function remains value independent. return ffi::GetRef(op); } -Expr ExprMutatorBase::VisitExpr_(const SeqExprNode* op) { +Expr ExprMutatorBase::VisitExpr_(const SeqExprObj* op) { bool all_blocks_unchanged = true; ffi::Array blocks; for (auto block : op->blocks) { @@ -718,12 +718,12 @@ Expr ExprMutatorBase::VisitExpr_(const SeqExprNode* op) { BindingBlock ExprMutatorBase::VisitBindingBlock(const BindingBlock& block) { ffi::Array bindings; - if (const auto* node = block.as()) { + if (const auto* node = block.as()) { for (auto binding : node->bindings) { - if (auto var_binding = binding.as()) { + if (auto var_binding = binding.as()) { Expr new_value = this->VisitExpr(var_binding->value); bindings.push_back(VarBinding(var_binding->var, new_value)); - } else if (auto match_cast = binding.as()) { + } else if (auto match_cast = binding.as()) { Expr new_value = this->VisitExpr(match_cast->value); bindings.push_back(MatchCast(match_cast->var, new_value, match_cast->ty)); } else { @@ -734,7 +734,7 @@ BindingBlock ExprMutatorBase::VisitBindingBlock(const BindingBlock& block) { TVM_FFI_THROW(TypeError) << "Invalid type: " << block->GetTypeKey(); } - if (block.as()) { + if (block.as()) { return DataflowBlock(bindings); } else { return BindingBlock(bindings); @@ -753,7 +753,7 @@ Expr ExprMutator::VisitExpr(const Expr& expr) { } // Visit the use-site of a defined Var -Expr ExprMutator::VisitExpr_(const VarNode* op) { +Expr ExprMutator::VisitExpr_(const VarObj* op) { Var var = ffi::GetRef(op); auto it = var_remap_.find(var); if (it != var_remap_.end()) { @@ -765,11 +765,11 @@ Expr ExprMutator::VisitExpr_(const VarNode* op) { } // Visit the use-site of a defined DataflowVar -Expr ExprMutator::VisitExpr_(const DataflowVarNode* op) { - return VisitExpr_(static_cast(op)); +Expr ExprMutator::VisitExpr_(const DataflowVarObj* op) { + return VisitExpr_(static_cast(op)); } -Expr ExprMutator::VisitExpr_(const FunctionNode* op) { +Expr ExprMutator::VisitExpr_(const FunctionObj* op) { tvm::ffi::Array params; bool all_params_unchanged = true; for (Var param : op->params) { @@ -809,7 +809,7 @@ Expr ExprMutator::VisitExpr_(const FunctionNode* op) { } } -Expr ExprMutator::VisitExpr_(const IfExprNode* op) { +Expr ExprMutator::VisitExpr_(const IfExprObj* op) { Expr guard = this->VisitExpr(op->cond); Expr true_b = this->VisitWithInnerScope(op->true_branch); Expr false_b = this->VisitWithInnerScope(op->false_branch); @@ -821,7 +821,7 @@ Expr ExprMutator::VisitExpr_(const IfExprNode* op) { } } -Expr ExprMutator::VisitExpr_(const SeqExprNode* op) { +Expr ExprMutator::VisitExpr_(const SeqExprObj* op) { bool all_blocks_unchanged = true; ffi::Array blocks; for (auto block : op->blocks) { @@ -848,24 +848,24 @@ Expr ExprMutator::VisitExpr_(const SeqExprNode* op) { } RELAX_VAR_BINDING_DISPATCH_IMPL(ExprMutator); -RELAX_EXPR_MUTATOR_VISIT_BINDING_IMPL(GenericConstNode); -RELAX_EXPR_MUTATOR_VISIT_BINDING_IMPL(TupleNode); -RELAX_EXPR_MUTATOR_VISIT_BINDING_IMPL(VarNode); -RELAX_EXPR_MUTATOR_VISIT_BINDING_IMPL(DataflowVarNode); -RELAX_EXPR_MUTATOR_VISIT_BINDING_IMPL(ShapeExprNode); -RELAX_EXPR_MUTATOR_VISIT_BINDING_IMPL(ExternFuncNode); -RELAX_EXPR_MUTATOR_VISIT_BINDING_IMPL(GlobalVarNode); -RELAX_EXPR_MUTATOR_VISIT_BINDING_IMPL(FunctionNode); -RELAX_EXPR_MUTATOR_VISIT_BINDING_IMPL(CallNode); -RELAX_EXPR_MUTATOR_VISIT_BINDING_IMPL(SeqExprNode); -RELAX_EXPR_MUTATOR_VISIT_BINDING_IMPL(IfExprNode); -RELAX_EXPR_MUTATOR_VISIT_BINDING_IMPL(OpNode); -RELAX_EXPR_MUTATOR_VISIT_BINDING_IMPL(TupleGetItemNode); -RELAX_EXPR_MUTATOR_VISIT_BINDING_IMPL(ExprNode); -RELAX_EXPR_MUTATOR_VISIT_BINDING_IMPL(StringImmNode); -RELAX_EXPR_MUTATOR_VISIT_BINDING_IMPL(DataTypeImmNode); - -void ExprMutator::ReEmitBinding(const VarBindingNode* binding, Expr new_value) { +RELAX_EXPR_MUTATOR_VISIT_BINDING_IMPL(GenericConstObj); +RELAX_EXPR_MUTATOR_VISIT_BINDING_IMPL(TupleObj); +RELAX_EXPR_MUTATOR_VISIT_BINDING_IMPL(VarObj); +RELAX_EXPR_MUTATOR_VISIT_BINDING_IMPL(DataflowVarObj); +RELAX_EXPR_MUTATOR_VISIT_BINDING_IMPL(ShapeExprObj); +RELAX_EXPR_MUTATOR_VISIT_BINDING_IMPL(ExternFuncObj); +RELAX_EXPR_MUTATOR_VISIT_BINDING_IMPL(GlobalVarObj); +RELAX_EXPR_MUTATOR_VISIT_BINDING_IMPL(FunctionObj); +RELAX_EXPR_MUTATOR_VISIT_BINDING_IMPL(CallObj); +RELAX_EXPR_MUTATOR_VISIT_BINDING_IMPL(SeqExprObj); +RELAX_EXPR_MUTATOR_VISIT_BINDING_IMPL(IfExprObj); +RELAX_EXPR_MUTATOR_VISIT_BINDING_IMPL(OpObj); +RELAX_EXPR_MUTATOR_VISIT_BINDING_IMPL(TupleGetItemObj); +RELAX_EXPR_MUTATOR_VISIT_BINDING_IMPL(ExprObj); +RELAX_EXPR_MUTATOR_VISIT_BINDING_IMPL(StringImmObj); +RELAX_EXPR_MUTATOR_VISIT_BINDING_IMPL(DataTypeImmObj); + +void ExprMutator::ReEmitBinding(const VarBindingObj* binding, Expr new_value) { Var new_var = this->VisitVarDef(binding->var); Var visited_var = new_var; @@ -894,7 +894,7 @@ void ExprMutator::ReEmitBinding(const VarBindingNode* binding, Expr new_value) { builder_->EmitNormalized(VarBinding(new_var, new_value)); } -void ExprMutator::VisitBinding_(const MatchCastNode* binding) { +void ExprMutator::VisitBinding_(const MatchCastObj* binding) { Expr new_value = this->VisitExpr(binding->value); Type new_ty = this->VisitExprDepTypeField(binding->ty); @@ -922,7 +922,7 @@ void ExprMutator::VisitBinding_(const MatchCastNode* binding) { builder_->AddDefinitionToScope(new_binding->var); } -BindingBlock ExprMutator::VisitBindingBlock_(const BindingBlockNode* block) { +BindingBlock ExprMutator::VisitBindingBlock_(const BindingBlockObj* block) { builder_->BeginBindingBlock(); for (Binding binding : block->bindings) { this->VisitBinding(binding); @@ -930,7 +930,7 @@ BindingBlock ExprMutator::VisitBindingBlock_(const BindingBlockNode* block) { return builder_->EndBlock(); } -BindingBlock ExprMutator::VisitBindingBlock_(const DataflowBlockNode* block) { +BindingBlock ExprMutator::VisitBindingBlock_(const DataflowBlockObj* block) { builder_->BeginDataflowBlock(); for (auto binding : block->bindings) { this->VisitBinding(binding); @@ -938,12 +938,12 @@ BindingBlock ExprMutator::VisitBindingBlock_(const DataflowBlockNode* block) { return builder_->EndBlock(); } -Var ExprMutator::VisitVarDef_(const DataflowVarNode* var) { - Var output = VisitVarDef_(static_cast(var)); +Var ExprMutator::VisitVarDef_(const DataflowVarObj* var) { + Var output = VisitVarDef_(static_cast(var)); // Because we delegate from DataflowVar visitor to Var visitor to // provide default behavior in subclasses, we may produce a Var // where we should produce a DataflowVar. - if (!output->IsInstance()) { + if (!output->IsInstance()) { Var delegated_output = output; output = DataflowVar(output->name, GetType(output), output->loc); var_remap_.insert_or_assign(delegated_output, output); @@ -951,8 +951,8 @@ Var ExprMutator::VisitVarDef_(const DataflowVarNode* var) { return output; } -Var ExprMutator::VisitVarDef_(const VarNode* var) { - if (auto* ty_node = var->ty.as()) { +Var ExprMutator::VisitVarDef_(const VarObj* var) { + if (auto* ty_node = var->ty.as()) { Type ty = this->VisitExprDepTypeField(ffi::GetRef(ty_node)); if (ty.same_as(var->ty)) { return ffi::GetRef(var); @@ -965,9 +965,9 @@ Var ExprMutator::VisitVarDef_(const VarNode* var) { } void ExprMutator::VisitBinding(const Binding& binding) { - if (const auto* node = binding.as()) { + if (const auto* node = binding.as()) { VisitBinding_(node); - } else if (const auto* node = binding.as()) { + } else if (const auto* node = binding.as()) { VisitBinding_(node); } else { TVM_FFI_THROW(TypeError) << "Invalid type: " << binding->GetTypeKey(); @@ -976,9 +976,9 @@ void ExprMutator::VisitBinding(const Binding& binding) { BindingBlock ExprMutator::VisitBindingBlock(const BindingBlock& block) { BindingBlock ret(ffi::UnsafeInit{}); - if (const auto* node = block.as()) { + if (const auto* node = block.as()) { ret = VisitBindingBlock_(node); - } else if (const auto* node = block.as()) { + } else if (const auto* node = block.as()) { ret = VisitBindingBlock_(node); } else { TVM_FFI_THROW(TypeError) << "Invalid type: " << block->GetTypeKey(); @@ -987,9 +987,9 @@ BindingBlock ExprMutator::VisitBindingBlock(const BindingBlock& block) { } Var ExprMutator::VisitVarDef(const Var& var) { - if (const auto* node = var.as()) { + if (const auto* node = var.as()) { return VisitVarDef_(node); - } else if (const auto* node = var.as()) { + } else if (const auto* node = var.as()) { return VisitVarDef_(node); } TVM_FFI_THROW(TypeError) << "Invalid type: " << var->GetTypeKey(); @@ -997,7 +997,7 @@ Var ExprMutator::VisitVarDef(const Var& var) { } Expr ExprMutator::VisitWithNewScope(const Expr& expr, ffi::Optional> params) { - TVM_FFI_ICHECK(expr->IsInstance()) + TVM_FFI_ICHECK(expr->IsInstance()) << "Normal form requires all new scope is stored as SeqExpr"; PrimExpr constraint = IntImm::Bool(true); @@ -1012,7 +1012,11 @@ 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); + auto analyzer = builder_->GetAnalyzer(); + auto saved_context = analyzer->GetConstraintContext(); + auto context = saved_context->Snapshot(); + context->Append(sym::PredicateConstraint(constraint)); + analyzer->SyncConstraintContext(context); builder_->BeginInnerScope(); // Inner scope also includes any TIR variables that are defined by // MatchCast nodes, and are internal to the scope. @@ -1027,11 +1031,12 @@ Expr ExprMutator::VisitWithNewScope(const Expr& expr, ffi::OptionalNormalize(ret); builder_->EndScope(); + analyzer->SyncConstraintContext(saved_context); return ret; } Expr ExprMutator::VisitWithInnerScope(const Expr& expr) { - TVM_FFI_ICHECK(expr->IsInstance()) + TVM_FFI_ICHECK(expr->IsInstance()) << "Normal form requires all new scope is stored as SeqExpr"; builder_->BeginInnerScope(); @@ -1056,8 +1061,8 @@ Var ExprMutator::WithType(Var var, Type ty) { if (var->ty.same_as(ty) || ffi::StructuralEqual()(var->ty, ty)) { return var; } else { - Var new_var = var.as() ? DataflowVar(var->name, ty, var->loc) - : Var(var->name, ty, var->loc); + Var new_var = var.as() ? DataflowVar(var->name, ty, var->loc) + : Var(var->name, ty, var->loc); return new_var; } } else { diff --git a/src/relax/ir/global_info.cc b/src/relax/ir/global_info.cc index b37bdfd4390d..c064267a6479 100644 --- a/src/relax/ir/global_info.cc +++ b/src/relax/ir/global_info.cc @@ -28,21 +28,21 @@ namespace tvm { namespace relax { TVM_FFI_STATIC_INIT_BLOCK() { - VDeviceNode::RegisterReflection(); - DummyGlobalInfoNode::RegisterReflection(); + VDeviceObj::RegisterReflection(); + DummyGlobalInfoObj::RegisterReflection(); } TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; auto make_dummy = []() { - return DummyGlobalInfo(ffi::UnsafeInit{}, ffi::make_object()); + return DummyGlobalInfo(ffi::UnsafeInit{}, ffi::make_object()); }; refl::GlobalDef().def("relax.DummyGlobalInfo", make_dummy); } VDevice::VDevice(ffi::Optional tgt, int dev_id, MemoryScope mem_scope) : GlobalInfo(ffi::UnsafeInit{}) { - ffi::ObjectPtr n = ffi::make_object(); + ffi::ObjectPtr n = ffi::make_object(); n->target = std::move(tgt); n->vdevice_id = std::move(dev_id); n->memory_scope = std::move(mem_scope); diff --git a/src/relax/ir/py_expr_functor.cc b/src/relax/ir/py_expr_functor.cc index c3910fbeb862..e2fd33255e06 100644 --- a/src/relax/ir/py_expr_functor.cc +++ b/src/relax/ir/py_expr_functor.cc @@ -40,60 +40,60 @@ class PyExprVisitorObj : public ffi::Object, public ExprVisitor { public: /*! \brief The packed function to the `VisitExpr(const Expr& expr)` function. */ ffi::Function f_visit_expr{nullptr}; - /*! \brief The packed function to the `VisitExpr_(const GenericConstNode* op)` function. */ + /*! \brief The packed function to the `VisitExpr_(const GenericConstObj* op)` function. */ ffi::Function f_visit_generic_const_{nullptr}; - /*! \brief The packed function to the `VisitExpr_(const TupleNode* op)` function. */ + /*! \brief The packed function to the `VisitExpr_(const TupleObj* op)` function. */ ffi::Function f_visit_tuple_{nullptr}; - /*! \brief The packed function to the `VisitExpr_(const VarNode* op)` function. */ + /*! \brief The packed function to the `VisitExpr_(const VarObj* op)` function. */ ffi::Function f_visit_var_{nullptr}; - /*! \brief The packed function to the `VisitExpr_(const DataflowVarNode* op)` function. */ + /*! \brief The packed function to the `VisitExpr_(const DataflowVarObj* op)` function. */ ffi::Function f_visit_dataflow_var_{nullptr}; - /*! \brief The packed function to the `VisitExpr_(const ShapeExprNode* op)` function. */ + /*! \brief The packed function to the `VisitExpr_(const ShapeExprObj* op)` function. */ ffi::Function f_visit_shape_expr_{nullptr}; - /*! \brief The packed function to the `VisitExpr_(const ExternFuncNode* op)` function. */ + /*! \brief The packed function to the `VisitExpr_(const ExternFuncObj* op)` function. */ ffi::Function f_visit_extern_func_{nullptr}; - /*! \brief The packed function to the `VisitExpr_(const GlobalVarNode* op)` function. */ + /*! \brief The packed function to the `VisitExpr_(const GlobalVarObj* op)` function. */ ffi::Function f_visit_global_var_{nullptr}; - /*! \brief The packed function to the `VisitExpr_(const FunctionNode* op)` function. */ + /*! \brief The packed function to the `VisitExpr_(const FunctionObj* op)` function. */ ffi::Function f_visit_function_{nullptr}; - /*! \brief The packed function to the `VisitExpr_(const CallNode* op)` function. */ + /*! \brief The packed function to the `VisitExpr_(const CallObj* op)` function. */ ffi::Function f_visit_call_{nullptr}; - /*! \brief The packed function to the `VisitExpr_(const SeqExprNode* op)` function. */ + /*! \brief The packed function to the `VisitExpr_(const SeqExprObj* op)` function. */ ffi::Function f_visit_seq_expr_{nullptr}; - /*! \brief The packed function to the `VisitExpr_(const IfExprNode* op)` function. */ + /*! \brief The packed function to the `VisitExpr_(const IfExprObj* op)` function. */ ffi::Function f_visit_if_expr_{nullptr}; - /*! \brief The packed function to the `VisitExpr_(const OpNode* op)` function. */ + /*! \brief The packed function to the `VisitExpr_(const OpObj* op)` function. */ ffi::Function f_visit_op_{nullptr}; - /*! \brief The packed function to the `VisitExpr_(const TupleGetItemNode* op)` function. */ + /*! \brief The packed function to the `VisitExpr_(const TupleGetItemObj* op)` function. */ ffi::Function f_visit_tuple_getitem_{nullptr}; /*! \brief The packed function to the generic expression fallback. */ ffi::Function f_visit_expr_fallback_{nullptr}; - /*! \brief The packed function to the `VisitExpr_(const StringImmNode* op)` function. */ + /*! \brief The packed function to the `VisitExpr_(const StringImmObj* op)` function. */ ffi::Function f_visit_string_imm_{nullptr}; - /*! \brief The packed function to the `VisitExpr_(const DataTypeImmNode* op)` function. */ + /*! \brief The packed function to the `VisitExpr_(const DataTypeImmObj* op)` function. */ ffi::Function f_visit_data_type_imm_{nullptr}; /*! \brief The packed function to the `VisitBinding(const Binding& binding)` function. */ ffi::Function f_visit_binding{nullptr}; - /*! \brief The packed function to the `VisitBinding_(const VarBindingNode* binding)` + /*! \brief The packed function to the `VisitBinding_(const VarBindingObj* binding)` * function. */ ffi::Function f_visit_var_binding_{nullptr}; - /*! \brief The packed function to the `VisitBinding_(const MatchCastNode* binding)` + /*! \brief The packed function to the `VisitBinding_(const MatchCastObj* binding)` * function. */ ffi::Function f_visit_match_cast_{nullptr}; /*! \brief The packed function to the `VisitBindingBlock(const BindingBlock& block)` * function. */ ffi::Function f_visit_binding_block{nullptr}; - /*! \brief The packed function to the `VisitBindingBlock_(const BindingBlockNode* block)` + /*! \brief The packed function to the `VisitBindingBlock_(const BindingBlockObj* block)` * function. */ ffi::Function f_visit_binding_block_{nullptr}; - /*! \brief The packed function to the `VisitBindingBlock_(const DataflowBlockNode* block)` + /*! \brief The packed function to the `VisitBindingBlock_(const DataflowBlockObj* block)` * function. */ ffi::Function f_visit_dataflow_block_{nullptr}; /*! \brief The packed function to the `VisitVarDef(const Var& var)` function. */ ffi::Function f_visit_var_def{nullptr}; - /*! \brief The packed function to the `VisitVarDef_(const VarNode* var)` function. */ + /*! \brief The packed function to the `VisitVarDef_(const VarObj* var)` function. */ ffi::Function f_visit_var_def_{nullptr}; - /*! \brief The packed function to the `VisitVarDef_(const DataflowVarNode* var)` function. */ + /*! \brief The packed function to the `VisitVarDef_(const DataflowVarObj* var)` function. */ ffi::Function f_visit_dataflow_var_def_{nullptr}; /*! \brief The packed function to the `VisitLoc(const Location& loc)` function. */ ffi::Function f_visit_loc{nullptr}; @@ -112,36 +112,36 @@ class PyExprVisitorObj : public ffi::Object, public ExprVisitor { } } - void VisitExprFallback_(const ExprNode* op) override + void VisitExprFallback_(const ExprObj* op) override PY_EXPR_VISITOR_DEFAULT(ffi::GetRef(op), f_visit_expr_fallback_, ExprVisitor::VisitExprFallback_(op)); void VisitBinding(const Binding& binding) override PY_EXPR_VISITOR_DEFAULT(binding, f_visit_binding, ExprVisitor::VisitBinding(binding)); - void VisitBinding_(const VarBindingNode* binding) override + void VisitBinding_(const VarBindingObj* binding) override PY_EXPR_VISITOR_DEFAULT(ffi::GetRef(binding), f_visit_var_binding_, ExprVisitor::VisitBinding_(binding)); - void VisitBinding_(const MatchCastNode* binding) override + void VisitBinding_(const MatchCastObj* binding) override PY_EXPR_VISITOR_DEFAULT(ffi::GetRef(binding), f_visit_match_cast_, ExprVisitor::VisitBinding_(binding)); void VisitBindingBlock(const BindingBlock& block) override PY_EXPR_VISITOR_DEFAULT(block, f_visit_binding_block, ExprVisitor::VisitBindingBlock(block)); - void VisitBindingBlock_(const BindingBlockNode* block) override + void VisitBindingBlock_(const BindingBlockObj* block) override PY_EXPR_VISITOR_DEFAULT(ffi::GetRef(block), f_visit_binding_block_, ExprVisitor::VisitBindingBlock_(block)); - void VisitBindingBlock_(const DataflowBlockNode* block) override + void VisitBindingBlock_(const DataflowBlockObj* block) override PY_EXPR_VISITOR_DEFAULT(ffi::GetRef(block), f_visit_dataflow_block_, ExprVisitor::VisitBindingBlock_(block)); void VisitVarDef(const Var& var) override PY_EXPR_VISITOR_DEFAULT(var, f_visit_var_def, ExprVisitor::VisitVarDef(var)); - void VisitVarDef_(const VarNode* var) override + void VisitVarDef_(const VarObj* var) override PY_EXPR_VISITOR_DEFAULT(ffi::GetRef(var), f_visit_var_def_, ExprVisitor::VisitVarDef_(var)); - void VisitVarDef_(const DataflowVarNode* var) override + void VisitVarDef_(const DataflowVarObj* var) override PY_EXPR_VISITOR_DEFAULT(ffi::GetRef(var), f_visit_dataflow_var_def_, ExprVisitor::VisitVarDef_(var)); @@ -160,21 +160,21 @@ class PyExprVisitorObj : public ffi::Object, public ExprVisitor { static FType InitVTable() { FType vtable; // Set dispatch - PY_EXPR_VISITOR_DISPATCH(GenericConstNode, f_visit_generic_const_); - PY_EXPR_VISITOR_DISPATCH(TupleNode, f_visit_tuple_); - PY_EXPR_VISITOR_DISPATCH(VarNode, f_visit_var_); - PY_EXPR_VISITOR_DISPATCH(DataflowVarNode, f_visit_dataflow_var_); - PY_EXPR_VISITOR_DISPATCH(ShapeExprNode, f_visit_shape_expr_); - PY_EXPR_VISITOR_DISPATCH(ExternFuncNode, f_visit_extern_func_); - PY_EXPR_VISITOR_DISPATCH(GlobalVarNode, f_visit_global_var_); - PY_EXPR_VISITOR_DISPATCH(FunctionNode, f_visit_function_); - PY_EXPR_VISITOR_DISPATCH(CallNode, f_visit_call_); - PY_EXPR_VISITOR_DISPATCH(SeqExprNode, f_visit_seq_expr_); - PY_EXPR_VISITOR_DISPATCH(IfExprNode, f_visit_if_expr_); - PY_EXPR_VISITOR_DISPATCH(OpNode, f_visit_op_); - PY_EXPR_VISITOR_DISPATCH(TupleGetItemNode, f_visit_tuple_getitem_); - PY_EXPR_VISITOR_DISPATCH(StringImmNode, f_visit_string_imm_); - PY_EXPR_VISITOR_DISPATCH(DataTypeImmNode, f_visit_data_type_imm_); + PY_EXPR_VISITOR_DISPATCH(GenericConstObj, f_visit_generic_const_); + PY_EXPR_VISITOR_DISPATCH(TupleObj, f_visit_tuple_); + PY_EXPR_VISITOR_DISPATCH(VarObj, f_visit_var_); + PY_EXPR_VISITOR_DISPATCH(DataflowVarObj, f_visit_dataflow_var_); + PY_EXPR_VISITOR_DISPATCH(ShapeExprObj, f_visit_shape_expr_); + PY_EXPR_VISITOR_DISPATCH(ExternFuncObj, f_visit_extern_func_); + PY_EXPR_VISITOR_DISPATCH(GlobalVarObj, f_visit_global_var_); + PY_EXPR_VISITOR_DISPATCH(FunctionObj, f_visit_function_); + PY_EXPR_VISITOR_DISPATCH(CallObj, f_visit_call_); + PY_EXPR_VISITOR_DISPATCH(SeqExprObj, f_visit_seq_expr_); + PY_EXPR_VISITOR_DISPATCH(IfExprObj, f_visit_if_expr_); + PY_EXPR_VISITOR_DISPATCH(OpObj, f_visit_op_); + PY_EXPR_VISITOR_DISPATCH(TupleGetItemObj, f_visit_tuple_getitem_); + PY_EXPR_VISITOR_DISPATCH(StringImmObj, f_visit_string_imm_); + PY_EXPR_VISITOR_DISPATCH(DataTypeImmObj, f_visit_data_type_imm_); vtable.Finalize(); return vtable; } @@ -192,36 +192,36 @@ class PyExprVisitor : public ffi::ObjectRef { /*! * \brief Create a PyExprVisitor with customized methods on the python-side. * \param f_visit_expr The packed function of `VisitExpr(const Expr& expr)`. - * \param f_visit_generic_const_ The packed function of `VisitExpr_(const GenericConstNode* op)`. - * \param f_visit_tuple_ The packed function of `VisitExpr_(const TupleNode* op)`. - * \param f_visit_var_ The packed function of `VisitExpr_(const VarNode* op)`. - * \param f_visit_dataflow_var_ The packed function of `VisitExpr_(const DataflowVarNode* op)`. - * \param f_visit_shape_expr_ The packed function of `VisitExpr_(const ShapeExprNode* op)`. - * \param f_visit_extern_func_ The packed function of `VisitExpr_(const ExternFuncNode* op)`. - * \param f_visit_global_var_ The packed function of `VisitExpr_(const GlobalVarNode* op)`. - * \param f_visit_function_ The packed function of `VisitExpr_(const FunctionNode* op)`. - * \param f_visit_call_ The packed function of `VisitExpr_(const CallNode* op)`. - * \param f_visit_seq_expr_ The packed function of `VisitExpr_(const SeqExprNode* op)`. - * \param f_visit_if_expr_ The packed function of `VisitExpr_(const IfExprNode* op)`. - * \param f_visit_op_ The packed function of `VisitExpr_(const OpNode* op)`. - * \param f_visit_tuple_getitem_ The packed function of `VisitExpr_(const TupleGetItemNode* op)`. + * \param f_visit_generic_const_ The packed function of `VisitExpr_(const GenericConstObj* op)`. + * \param f_visit_tuple_ The packed function of `VisitExpr_(const TupleObj* op)`. + * \param f_visit_var_ The packed function of `VisitExpr_(const VarObj* op)`. + * \param f_visit_dataflow_var_ The packed function of `VisitExpr_(const DataflowVarObj* op)`. + * \param f_visit_shape_expr_ The packed function of `VisitExpr_(const ShapeExprObj* op)`. + * \param f_visit_extern_func_ The packed function of `VisitExpr_(const ExternFuncObj* op)`. + * \param f_visit_global_var_ The packed function of `VisitExpr_(const GlobalVarObj* op)`. + * \param f_visit_function_ The packed function of `VisitExpr_(const FunctionObj* op)`. + * \param f_visit_call_ The packed function of `VisitExpr_(const CallObj* op)`. + * \param f_visit_seq_expr_ The packed function of `VisitExpr_(const SeqExprObj* op)`. + * \param f_visit_if_expr_ The packed function of `VisitExpr_(const IfExprObj* op)`. + * \param f_visit_op_ The packed function of `VisitExpr_(const OpObj* op)`. + * \param f_visit_tuple_getitem_ The packed function of `VisitExpr_(const TupleGetItemObj* op)`. * \param f_visit_expr_fallback_ The packed function of the generic expression fallback. - * \param f_visit_string_imm_ The packed function of `VisitExpr_(const StringImmNode* op)`. - * \param f_visit_data_type_imm_ The packed function of `VisitExpr_(const DataTypeImmNode* op)`. + * \param f_visit_string_imm_ The packed function of `VisitExpr_(const StringImmObj* op)`. + * \param f_visit_data_type_imm_ The packed function of `VisitExpr_(const DataTypeImmObj* op)`. * \param f_visit_binding The packed function of `VisitBinding(const Binding& binding)`. - * \param f_visit_var_binding_ The packed function of `VisitBinding_(const VarBindingNode* + * \param f_visit_var_binding_ The packed function of `VisitBinding_(const VarBindingObj* * binding)`. - * \param f_visit_match_cast_ The packed function of `VisitBinding_(const MatchCastNode* + * \param f_visit_match_cast_ The packed function of `VisitBinding_(const MatchCastObj* * binding)`. * \param f_visit_binding_block The packed function of `VisitBindingBlock(const BindingBlock& * block)`. * \param f_visit_binding_block_ The packed function of `VisitBindingBlock_(const - * BindingBlockNode* block)`. + * BindingBlockObj* block)`. * \param f_visit_dataflow_block_ The packed function of `VisitBindingBlock_(const - * DataflowBlockNode* block)`. + * DataflowBlockObj* block)`. * \param f_visit_var_def The packed function of `VisitVarDef(const Var& var)`. - * \param f_visit_var_def_ The packed function of `VisitVarDef_(const VarNode* var)`. - * \param f_visit_dataflow_var_def_ The packed function of `VisitVarDef_(const DataflowVarNode* + * \param f_visit_var_def_ The packed function of `VisitVarDef_(const VarObj* var)`. + * \param f_visit_dataflow_var_def_ The packed function of `VisitVarDef_(const DataflowVarObj* * var)`. * \param f_visit_loc The packed function of `VisitLoc(const Location& loc)`. * \return The PyVisitor created. @@ -285,60 +285,60 @@ class PyExprMutatorObj : public ffi::Object, public ExprMutator { public: /*! \brief The packed function to the `VisitExpr(const Expr& expr)` function. */ ffi::Function f_visit_expr{nullptr}; - /*! \brief The packed function to the `VisitExpr_(const GenericConstNode* op)` function. */ + /*! \brief The packed function to the `VisitExpr_(const GenericConstObj* op)` function. */ ffi::Function f_visit_generic_const_{nullptr}; - /*! \brief The packed function to the `VisitExpr_(const TupleNode* op)` function. */ + /*! \brief The packed function to the `VisitExpr_(const TupleObj* op)` function. */ ffi::Function f_visit_tuple_{nullptr}; - /*! \brief The packed function to the `VisitExpr_(const VarNode* op)` function. */ + /*! \brief The packed function to the `VisitExpr_(const VarObj* op)` function. */ ffi::Function f_visit_var_{nullptr}; - /*! \brief The packed function to the `VisitExpr_(const DataflowVarNode* op)` function. */ + /*! \brief The packed function to the `VisitExpr_(const DataflowVarObj* op)` function. */ ffi::Function f_visit_dataflow_var_{nullptr}; - /*! \brief The packed function to the `VisitExpr_(const ShapeExprNode* op)` function. */ + /*! \brief The packed function to the `VisitExpr_(const ShapeExprObj* op)` function. */ ffi::Function f_visit_shape_expr_{nullptr}; - /*! \brief The packed function to the `VisitExpr_(const ExternFuncNode* op)` function. */ + /*! \brief The packed function to the `VisitExpr_(const ExternFuncObj* op)` function. */ ffi::Function f_visit_extern_func_{nullptr}; - /*! \brief The packed function to the `VisitExpr_(const GlobalVarNode* op)` function. */ + /*! \brief The packed function to the `VisitExpr_(const GlobalVarObj* op)` function. */ ffi::Function f_visit_global_var_{nullptr}; - /*! \brief The packed function to the `VisitExpr_(const FunctionNode* op)` function. */ + /*! \brief The packed function to the `VisitExpr_(const FunctionObj* op)` function. */ ffi::Function f_visit_function_{nullptr}; - /*! \brief The packed function to the `VisitExpr_(const CallNode* op)` function. */ + /*! \brief The packed function to the `VisitExpr_(const CallObj* op)` function. */ ffi::Function f_visit_call_{nullptr}; - /*! \brief The packed function to the `VisitExpr_(const SeqExprNode* op)` function. */ + /*! \brief The packed function to the `VisitExpr_(const SeqExprObj* op)` function. */ ffi::Function f_visit_seq_expr_{nullptr}; - /*! \brief The packed function to the `VisitExpr_(const IfExprNode* op)` function. */ + /*! \brief The packed function to the `VisitExpr_(const IfExprObj* op)` function. */ ffi::Function f_visit_if_expr_{nullptr}; - /*! \brief The packed function to the `VisitExpr_(const OpNode* op)` function. */ + /*! \brief The packed function to the `VisitExpr_(const OpObj* op)` function. */ ffi::Function f_visit_op_{nullptr}; - /*! \brief The packed function to the `VisitExpr_(const TupleGetItemNode* op)` function. */ + /*! \brief The packed function to the `VisitExpr_(const TupleGetItemObj* op)` function. */ ffi::Function f_visit_tuple_getitem_{nullptr}; /*! \brief The packed function to the generic expression fallback. */ ffi::Function f_visit_expr_fallback_{nullptr}; - /*! \brief The packed function to the `VisitExpr_(const StringImmNode* op)` function. */ + /*! \brief The packed function to the `VisitExpr_(const StringImmObj* op)` function. */ ffi::Function f_visit_string_imm_{nullptr}; - /*! \brief The packed function to the `VisitExpr_(const DataTypeImmNode* op)` function. */ + /*! \brief The packed function to the `VisitExpr_(const DataTypeImmObj* op)` function. */ ffi::Function f_visit_data_type_imm_{nullptr}; /*! \brief The packed function to the `VisitBinding(const Binding& binding)` function. */ ffi::Function f_visit_binding{nullptr}; - /*! \brief The packed function to the `VisitBinding_(const VarBindingNode* binding)` + /*! \brief The packed function to the `VisitBinding_(const VarBindingObj* binding)` * function. */ ffi::Function f_visit_var_binding_{nullptr}; - /*! \brief The packed function to the `VisitBinding_(const MatchCastNode* binding)` + /*! \brief The packed function to the `VisitBinding_(const MatchCastObj* binding)` * function. */ ffi::Function f_visit_match_cast_{nullptr}; /*! \brief The packed function to the `VisitBindingBlock(const BindingBlock& block)` * function. */ ffi::Function f_visit_binding_block{nullptr}; - /*! \brief The packed function to the `VisitBindingBlock_(const BindingBlockNode* block)` + /*! \brief The packed function to the `VisitBindingBlock_(const BindingBlockObj* block)` * function. */ ffi::Function f_visit_binding_block_{nullptr}; - /*! \brief The packed function to the `VisitBindingBlock_(const DataflowBlockNode* block)` + /*! \brief The packed function to the `VisitBindingBlock_(const DataflowBlockObj* block)` * function. */ ffi::Function f_visit_dataflow_block_{nullptr}; /*! \brief The packed function to the `VisitVarDef(const Var& var)` function. */ ffi::Function f_visit_var_def{nullptr}; - /*! \brief The packed function to the `VisitVarDef_(const VarNode* var)` function. */ + /*! \brief The packed function to the `VisitVarDef_(const VarObj* var)` function. */ ffi::Function f_visit_var_def_{nullptr}; - /*! \brief The packed function to the `VisitVarDef_(const DataflowVarNode* var)` function. */ + /*! \brief The packed function to the `VisitVarDef_(const DataflowVarObj* var)` function. */ ffi::Function f_visit_dataflow_var_def_{nullptr}; /*! \brief The packed function to the `VisitLoc(const Location& loc)` function. */ ffi::Function f_visit_loc{nullptr}; @@ -355,7 +355,7 @@ class PyExprMutatorObj : public ffi::Object, public ExprMutator { } } - Expr VisitExprFallback_(const ExprNode* op) override + Expr VisitExprFallback_(const ExprObj* op) override PY_EXPR_MUTATOR_DEFAULT(ffi::GetRef(op), f_visit_expr_fallback_, ExprMutator::VisitExprFallback_(op), Expr); @@ -366,14 +366,14 @@ class PyExprMutatorObj : public ffi::Object, public ExprMutator { ExprMutator::VisitBinding(binding); } - void VisitBinding_(const VarBindingNode* binding) override { + void VisitBinding_(const VarBindingObj* binding) override { if (f_visit_var_binding_ != nullptr) f_visit_var_binding_(ffi::GetRef(binding)); else ExprMutator::VisitBinding_(binding); } - void VisitBinding_(const MatchCastNode* binding) override { + void VisitBinding_(const MatchCastObj* binding) override { if (f_visit_match_cast_ != nullptr) f_visit_match_cast_(ffi::GetRef(binding)); else @@ -384,19 +384,19 @@ class PyExprMutatorObj : public ffi::Object, public ExprMutator { PY_EXPR_MUTATOR_DEFAULT(block, f_visit_binding_block, ExprMutator::VisitBindingBlock(block), BindingBlock); - BindingBlock VisitBindingBlock_(const BindingBlockNode* block) override + BindingBlock VisitBindingBlock_(const BindingBlockObj* block) override PY_EXPR_MUTATOR_DEFAULT(ffi::GetRef(block), f_visit_binding_block_, ExprMutator::VisitBindingBlock_(block), BindingBlock); - BindingBlock VisitBindingBlock_(const DataflowBlockNode* block) override + BindingBlock VisitBindingBlock_(const DataflowBlockObj* block) override PY_EXPR_MUTATOR_DEFAULT(ffi::GetRef(block), f_visit_dataflow_block_, ExprMutator::VisitBindingBlock_(block), BindingBlock); Var VisitVarDef(const Var& var) override PY_EXPR_MUTATOR_DEFAULT(var, f_visit_var_def, ExprMutator::VisitVarDef(var), Var); - Var VisitVarDef_(const VarNode* var) override + Var VisitVarDef_(const VarObj* var) override PY_EXPR_MUTATOR_DEFAULT(ffi::GetRef(var), f_visit_var_def_, ExprMutator::VisitVarDef_(var), Var); - Var VisitVarDef_(const DataflowVarNode* var) override + Var VisitVarDef_(const DataflowVarObj* var) override PY_EXPR_MUTATOR_DEFAULT(ffi::GetRef(var), f_visit_dataflow_var_def_, ExprMutator::VisitVarDef_(var), Var); @@ -432,21 +432,21 @@ class PyExprMutatorObj : public ffi::Object, public ExprMutator { static FType InitVTable() { FType vtable; // Set dispatch - PY_EXPR_MUTATOR_DISPATCH(GenericConstNode, f_visit_generic_const_); - PY_EXPR_MUTATOR_DISPATCH(TupleNode, f_visit_tuple_); - PY_EXPR_MUTATOR_DISPATCH(VarNode, f_visit_var_); - PY_EXPR_MUTATOR_DISPATCH(DataflowVarNode, f_visit_dataflow_var_); - PY_EXPR_MUTATOR_DISPATCH(ShapeExprNode, f_visit_shape_expr_); - PY_EXPR_MUTATOR_DISPATCH(ExternFuncNode, f_visit_extern_func_); - PY_EXPR_MUTATOR_DISPATCH(GlobalVarNode, f_visit_global_var_); - PY_EXPR_MUTATOR_DISPATCH(FunctionNode, f_visit_function_); - PY_EXPR_MUTATOR_DISPATCH(CallNode, f_visit_call_); - PY_EXPR_MUTATOR_DISPATCH(SeqExprNode, f_visit_seq_expr_); - PY_EXPR_MUTATOR_DISPATCH(IfExprNode, f_visit_if_expr_); - PY_EXPR_MUTATOR_DISPATCH(OpNode, f_visit_op_); - PY_EXPR_MUTATOR_DISPATCH(TupleGetItemNode, f_visit_tuple_getitem_); - PY_EXPR_MUTATOR_DISPATCH(StringImmNode, f_visit_string_imm_); - PY_EXPR_MUTATOR_DISPATCH(DataTypeImmNode, f_visit_data_type_imm_); + PY_EXPR_MUTATOR_DISPATCH(GenericConstObj, f_visit_generic_const_); + PY_EXPR_MUTATOR_DISPATCH(TupleObj, f_visit_tuple_); + PY_EXPR_MUTATOR_DISPATCH(VarObj, f_visit_var_); + PY_EXPR_MUTATOR_DISPATCH(DataflowVarObj, f_visit_dataflow_var_); + PY_EXPR_MUTATOR_DISPATCH(ShapeExprObj, f_visit_shape_expr_); + PY_EXPR_MUTATOR_DISPATCH(ExternFuncObj, f_visit_extern_func_); + PY_EXPR_MUTATOR_DISPATCH(GlobalVarObj, f_visit_global_var_); + PY_EXPR_MUTATOR_DISPATCH(FunctionObj, f_visit_function_); + PY_EXPR_MUTATOR_DISPATCH(CallObj, f_visit_call_); + PY_EXPR_MUTATOR_DISPATCH(SeqExprObj, f_visit_seq_expr_); + PY_EXPR_MUTATOR_DISPATCH(IfExprObj, f_visit_if_expr_); + PY_EXPR_MUTATOR_DISPATCH(OpObj, f_visit_op_); + PY_EXPR_MUTATOR_DISPATCH(TupleGetItemObj, f_visit_tuple_getitem_); + PY_EXPR_MUTATOR_DISPATCH(StringImmObj, f_visit_string_imm_); + PY_EXPR_MUTATOR_DISPATCH(DataTypeImmObj, f_visit_data_type_imm_); vtable.Finalize(); return vtable; } @@ -455,21 +455,21 @@ class PyExprMutatorObj : public ffi::Object, public ExprMutator { static FType InitPostOrderVTable() { FType post_order_vtable; // Set dispatch - PY_EXPR_MUTATOR_VISIT_EXPR_POST_ORDER_DISPATCH(GenericConstNode); - PY_EXPR_MUTATOR_VISIT_EXPR_POST_ORDER_DISPATCH(TupleNode); - PY_EXPR_MUTATOR_VISIT_EXPR_POST_ORDER_DISPATCH(VarNode); - PY_EXPR_MUTATOR_VISIT_EXPR_POST_ORDER_DISPATCH(DataflowVarNode); - PY_EXPR_MUTATOR_VISIT_EXPR_POST_ORDER_DISPATCH(ShapeExprNode); - PY_EXPR_MUTATOR_VISIT_EXPR_POST_ORDER_DISPATCH(ExternFuncNode); - PY_EXPR_MUTATOR_VISIT_EXPR_POST_ORDER_DISPATCH(GlobalVarNode); - PY_EXPR_MUTATOR_VISIT_EXPR_POST_ORDER_DISPATCH(FunctionNode); - PY_EXPR_MUTATOR_VISIT_EXPR_POST_ORDER_DISPATCH(CallNode); - PY_EXPR_MUTATOR_VISIT_EXPR_POST_ORDER_DISPATCH(SeqExprNode); - PY_EXPR_MUTATOR_VISIT_EXPR_POST_ORDER_DISPATCH(IfExprNode); - PY_EXPR_MUTATOR_VISIT_EXPR_POST_ORDER_DISPATCH(OpNode); - PY_EXPR_MUTATOR_VISIT_EXPR_POST_ORDER_DISPATCH(TupleGetItemNode); - PY_EXPR_MUTATOR_VISIT_EXPR_POST_ORDER_DISPATCH(StringImmNode); - PY_EXPR_MUTATOR_VISIT_EXPR_POST_ORDER_DISPATCH(DataTypeImmNode); + PY_EXPR_MUTATOR_VISIT_EXPR_POST_ORDER_DISPATCH(GenericConstObj); + PY_EXPR_MUTATOR_VISIT_EXPR_POST_ORDER_DISPATCH(TupleObj); + PY_EXPR_MUTATOR_VISIT_EXPR_POST_ORDER_DISPATCH(VarObj); + PY_EXPR_MUTATOR_VISIT_EXPR_POST_ORDER_DISPATCH(DataflowVarObj); + PY_EXPR_MUTATOR_VISIT_EXPR_POST_ORDER_DISPATCH(ShapeExprObj); + PY_EXPR_MUTATOR_VISIT_EXPR_POST_ORDER_DISPATCH(ExternFuncObj); + PY_EXPR_MUTATOR_VISIT_EXPR_POST_ORDER_DISPATCH(GlobalVarObj); + PY_EXPR_MUTATOR_VISIT_EXPR_POST_ORDER_DISPATCH(FunctionObj); + PY_EXPR_MUTATOR_VISIT_EXPR_POST_ORDER_DISPATCH(CallObj); + PY_EXPR_MUTATOR_VISIT_EXPR_POST_ORDER_DISPATCH(SeqExprObj); + PY_EXPR_MUTATOR_VISIT_EXPR_POST_ORDER_DISPATCH(IfExprObj); + PY_EXPR_MUTATOR_VISIT_EXPR_POST_ORDER_DISPATCH(OpObj); + PY_EXPR_MUTATOR_VISIT_EXPR_POST_ORDER_DISPATCH(TupleGetItemObj); + PY_EXPR_MUTATOR_VISIT_EXPR_POST_ORDER_DISPATCH(StringImmObj); + PY_EXPR_MUTATOR_VISIT_EXPR_POST_ORDER_DISPATCH(DataTypeImmObj); post_order_vtable.Finalize(); return post_order_vtable; } @@ -487,36 +487,36 @@ class PyExprMutator : public ffi::ObjectRef { /*! * \brief Create a PyExprMutator with customized methods on the python-side. * \param f_visit_expr The packed function of `VisitExpr(const Expr& expr)`. - * \param f_visit_generic_const_ The packed function of `VisitExpr_(const GenericConstNode* op)`. - * \param f_visit_tuple_ The packed function of `VisitExpr_(const TupleNode* op)`. - * \param f_visit_var_ The packed function of `VisitExpr_(const VarNode* op)`. - * \param f_visit_dataflow_var_ The packed function of `VisitExpr_(const DataflowVarNode* op)`. - * \param f_visit_shape_expr_ The packed function of `VisitExpr_(const ShapeExprNode* op)`. - * \param f_visit_extern_func_ The packed function of `VisitExpr_(const ExternFuncNode* op)`. - * \param f_visit_global_var_ The packed function of `VisitExpr_(const GlobalVarNode* op)`. - * \param f_visit_function_ The packed function of `VisitExpr_(const FunctionNode* op)`. - * \param f_visit_call_ The packed function of `VisitExpr_(const CallNode* op)`. - * \param f_visit_seq_expr_ The packed function of `VisitExpr_(const SeqExprNode* op)`. - * \param f_visit_if_expr_ The packed function of `VisitExpr_(const IfExprNode* op)`. - * \param f_visit_op_ The packed function of `VisitExpr_(const OpNode* op)`. - * \param f_visit_tuple_getitem_ The packed function of `VisitExpr_(const TupleGetItemNode* op)`. + * \param f_visit_generic_const_ The packed function of `VisitExpr_(const GenericConstObj* op)`. + * \param f_visit_tuple_ The packed function of `VisitExpr_(const TupleObj* op)`. + * \param f_visit_var_ The packed function of `VisitExpr_(const VarObj* op)`. + * \param f_visit_dataflow_var_ The packed function of `VisitExpr_(const DataflowVarObj* op)`. + * \param f_visit_shape_expr_ The packed function of `VisitExpr_(const ShapeExprObj* op)`. + * \param f_visit_extern_func_ The packed function of `VisitExpr_(const ExternFuncObj* op)`. + * \param f_visit_global_var_ The packed function of `VisitExpr_(const GlobalVarObj* op)`. + * \param f_visit_function_ The packed function of `VisitExpr_(const FunctionObj* op)`. + * \param f_visit_call_ The packed function of `VisitExpr_(const CallObj* op)`. + * \param f_visit_seq_expr_ The packed function of `VisitExpr_(const SeqExprObj* op)`. + * \param f_visit_if_expr_ The packed function of `VisitExpr_(const IfExprObj* op)`. + * \param f_visit_op_ The packed function of `VisitExpr_(const OpObj* op)`. + * \param f_visit_tuple_getitem_ The packed function of `VisitExpr_(const TupleGetItemObj* op)`. * \param f_visit_expr_fallback_ The packed function of the generic expression fallback. - * \param f_visit_string_imm_ The packed function of `VisitExpr_(const StringImmNode* op)`. - * \param f_visit_data_type_imm_ The packed function of `VisitExpr_(const DataTypeImmNode* op)`. + * \param f_visit_string_imm_ The packed function of `VisitExpr_(const StringImmObj* op)`. + * \param f_visit_data_type_imm_ The packed function of `VisitExpr_(const DataTypeImmObj* op)`. * \param f_visit_binding The packed function of `VisitBinding(const Binding& binding)`. - * \param f_visit_var_binding_ The packed function of `VisitBinding_(const VarBindingNode* + * \param f_visit_var_binding_ The packed function of `VisitBinding_(const VarBindingObj* * binding)`. - * \param f_visit_match_cast_ The packed function of `VisitBinding_(const MatchCastNode* + * \param f_visit_match_cast_ The packed function of `VisitBinding_(const MatchCastObj* * binding)`. * \param f_visit_binding_block The packed function of `VisitBindingBlock(const BindingBlock& * block)`. * \param f_visit_binding_block_ The packed function of `VisitBindingBlock_(const - * BindingBlockNode* block)`. + * BindingBlockObj* block)`. * \param f_visit_dataflow_block_ The packed function of `VisitBindingBlock_(const - * DataflowBlockNode* block)`. + * DataflowBlockObj* block)`. * \param f_visit_var_def The packed function of `VisitVarDef(const Var& var)`. - * \param f_visit_var_def_ The packed function of `VisitVarDef_(const VarNode* var)`. - * \param f_visit_dataflow_var_def_ The packed function of `VisitVarDef_(const DataflowVarNode* + * \param f_visit_var_def_ The packed function of `VisitVarDef_(const VarObj* var)`. + * \param f_visit_dataflow_var_def_ The packed function of `VisitVarDef_(const DataflowVarObj* * var)`. * \param f_visit_loc The packed function of `VisitLoc(const Location& loc)`. * \return The PyExprMutator created. @@ -588,9 +588,9 @@ TVM_FFI_STATIC_INIT_BLOCK() { [](PyExprVisitor visitor, const Expr& expr) { visitor->ExprVisitor::VisitExpr(expr); }) .def("relax.ExprVisitorVisitBinding", [](PyExprVisitor visitor, const Binding& binding) { - if (const auto* ptr = binding.as()) { + if (const auto* ptr = binding.as()) { visitor->ExprVisitor::VisitBinding_(ptr); - } else if (const auto* ptr = binding.as()) { + } else if (const auto* ptr = binding.as()) { visitor->ExprVisitor::VisitBinding_(ptr); } else { TVM_FFI_THROW(InternalError) << "unreachable"; @@ -598,9 +598,9 @@ TVM_FFI_STATIC_INIT_BLOCK() { }) .def("relax.ExprVisitorVisitBindingBlock", [](PyExprVisitor visitor, const BindingBlock& block) { - if (const auto* ptr = block.as()) { + if (const auto* ptr = block.as()) { visitor->ExprVisitor::VisitBindingBlock_(ptr); - } else if (const auto* ptr = block.as()) { + } else if (const auto* ptr = block.as()) { visitor->ExprVisitor::VisitBindingBlock_(ptr); } else { TVM_FFI_THROW(TypeError) << "Invalid type: " << block->GetTypeKey(); @@ -608,9 +608,9 @@ TVM_FFI_STATIC_INIT_BLOCK() { }) .def("relax.ExprVisitorVisitVarDef", [](PyExprVisitor visitor, const Var& var) { - if (const auto* node = var.as()) { + if (const auto* node = var.as()) { visitor->ExprVisitor::VisitVarDef_(node); - } else if (const auto* node = var.as()) { + } else if (const auto* node = var.as()) { visitor->ExprVisitor::VisitVarDef_(node); } else { TVM_FFI_THROW(TypeError) << "Invalid type: " << var->GetTypeKey(); @@ -639,9 +639,9 @@ TVM_FFI_STATIC_INIT_BLOCK() { }) .def("relax.ExprMutatorVisitBinding", [](PyExprMutator mutator, const Binding& binding) { - if (const auto* ptr = binding.as()) { + if (const auto* ptr = binding.as()) { return mutator->ExprMutator::VisitBinding_(ptr); - } else if (const auto* ptr = binding.as()) { + } else if (const auto* ptr = binding.as()) { return mutator->ExprMutator::VisitBinding_(ptr); } else { TVM_FFI_THROW(InternalError) << "unreachable"; @@ -649,9 +649,9 @@ TVM_FFI_STATIC_INIT_BLOCK() { }) .def("relax.ExprMutatorVisitBindingBlock", [](PyExprMutator mutator, const BindingBlock& block) { - if (const auto* node = block.as()) { + if (const auto* node = block.as()) { return mutator->ExprMutator::VisitBindingBlock_(node); - } else if (const auto* node = block.as()) { + } else if (const auto* node = block.as()) { return mutator->ExprMutator::VisitBindingBlock_(node); } else { TVM_FFI_THROW(TypeError) << "Invalid type: " << block->GetTypeKey(); @@ -659,9 +659,9 @@ TVM_FFI_STATIC_INIT_BLOCK() { }) .def("relax.ExprMutatorVisitVarDef", [](PyExprMutator mutator, const Var& var) { - if (const auto* node = var.as()) { + if (const auto* node = var.as()) { return mutator->ExprMutator::VisitVarDef_(node); - } else if (const auto* node = var.as()) { + } else if (const auto* node = var.as()) { return mutator->ExprMutator::VisitVarDef_(node); } else { TVM_FFI_THROW(TypeError) << "Invalid type: " << var->GetTypeKey(); diff --git a/src/relax/ir/transform.cc b/src/relax/ir/transform.cc index a20564f5c919..b6daedefa17f 100644 --- a/src/relax/ir/transform.cc +++ b/src/relax/ir/transform.cc @@ -127,7 +127,7 @@ IRModule FunctionPassObj::operator()(IRModule mod, const PassContext& pass_ctx) std::vector > updates; for (const auto& it : updated_mod->functions) { // only picks up relax::Function - if (auto* n = it.second.as()) { + if (auto* n = it.second.as()) { Function func = ffi::GetRef(n); // Enrich at this leaf executor, rendering the location local to the // currently-processed function (the access path is function-rooted). @@ -216,12 +216,12 @@ class DataflowBlockMutator : public ExprMutator { : pass_func_(pass_func), mod_(mod), pass_ctx_(pass_ctx) {} /*! - * \brief Rewrite the DataflowBlockNode with pass_func_ + * \brief Rewrite the DataflowBlockObj with pass_func_ * * This function will check that there are no rewrites of the global-scope * Vars or symbolic Vars defined inside the dataflow block. */ - BindingBlock VisitBindingBlock_(const DataflowBlockNode* n) final { + BindingBlock VisitBindingBlock_(const DataflowBlockObj* n) final { PreservedVars preserved_vars = CollectPreservedVars(n->bindings); // apply pass_func_ to the DataflowBlock @@ -229,11 +229,11 @@ class DataflowBlockMutator : public ExprMutator { DataflowBlock updated_block = pass_func_(block, mod_, pass_ctx_); PreservedVars updated_vars = CollectPreservedVars(updated_block->bindings); - for (const VarNode* var : preserved_vars.binding_vars) { + for (const VarObj* var : preserved_vars.binding_vars) { TVM_FFI_ICHECK(updated_vars.binding_vars.count(var)) << "Error: DataflowBlock Pass should not rewrite or delete any global-scope Var."; } - for (const VarNode* var : preserved_vars.match_cast_symbolic_vars) { + for (const VarObj* var : preserved_vars.match_cast_symbolic_vars) { TVM_FFI_ICHECK(updated_vars.match_cast_symbolic_vars.count(var)) << "Error: DataflowBlock Pass should not rewrite or delete any symbolic Var declared " "by a MatchCast."; @@ -243,7 +243,7 @@ class DataflowBlockMutator : public ExprMutator { } private: - using VarSet = std::unordered_set; + using VarSet = std::unordered_set; struct PreservedVars { VarSet binding_vars; @@ -253,12 +253,12 @@ class DataflowBlockMutator : public ExprMutator { static PreservedVars CollectPreservedVars(const ffi::Array& bindings) { PreservedVars vars; for (const Binding& binding : bindings) { - if (const auto* match_cast = binding.as()) { + if (const auto* match_cast = binding.as()) { for (const tvm::Var& var : DefinableTIRVarsInType(match_cast->ty)) { vars.match_cast_symbolic_vars.insert(var.get()); } } - if (!binding->var.as()) { + if (!binding->var.as()) { vars.binding_vars.insert(binding->var.get()); } } @@ -309,7 +309,7 @@ IRModule DataflowBlockPassObj::operator()(IRModule mod, const PassContext& pass_ std::vector > updates; for (const auto& it : updated_mod->functions) { // only picks up relax::Function - if (auto* n = it.second.as()) { + if (auto* n = it.second.as()) { Function func = ffi::GetRef(n); // Enrich at this leaf executor, rendering the location local to the // currently-processed function. diff --git a/src/relax/ir/type.cc b/src/relax/ir/type.cc index 9f33f6f701de..2198e1a52acb 100644 --- a/src/relax/ir/type.cc +++ b/src/relax/ir/type.cc @@ -50,15 +50,15 @@ TVM_FFI_INLINE ffi::Expected> PackedFuncTypeMaybeInpl } // namespace PackedFuncType::PackedFuncType(Location loc) : Type(ffi::UnsafeInit{}) { - ffi::ObjectPtr n = ffi::make_object(); + ffi::ObjectPtr n = ffi::make_object(); n->loc = loc; data_ = std::move(n); } TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - PackedFuncTypeNode::RegisterReflection(); - refl::TypeAttrDef() + PackedFuncTypeObj::RegisterReflection(); + refl::TypeAttrDef() .attr(ffi::type_attr::kStructuralVisit, ffi::FStructuralVisit::FromNative<&PackedFuncTypeVisit>()) .attr(ffi::type_attr::kStructuralMutate, diff --git a/src/relax/ir/type_functor.cc b/src/relax/ir/type_functor.cc index 38d7cf49c650..0b158e6e1a94 100644 --- a/src/relax/ir/type_functor.cc +++ b/src/relax/ir/type_functor.cc @@ -27,13 +27,13 @@ namespace tvm { namespace relax { -void TypeVisitor::VisitType_(const AnyTypeNode* op) {} +void TypeVisitor::VisitType_(const AnyTypeObj* op) {} -void TypeVisitor::VisitType_(const PrimTypeNode* op) {} +void TypeVisitor::VisitType_(const PrimTypeObj* op) {} -void TypeVisitor::VisitType_(const StringTypeNode* op) {} +void TypeVisitor::VisitType_(const StringTypeObj* op) {} -void TypeVisitor::VisitType_(const ShapeTypeNode* op) { +void TypeVisitor::VisitType_(const ShapeTypeObj* op) { if (op->values.has_value()) { for (PrimExpr value : op->values.value()) { this->VisitTypeExprField(value); @@ -41,28 +41,28 @@ void TypeVisitor::VisitType_(const ShapeTypeNode* op) { } } -void TypeVisitor::VisitType_(const TensorTypeNode* op) { +void TypeVisitor::VisitType_(const TensorTypeObj* op) { if (op->shape.has_value()) { this->VisitTypeExprField(op->shape.value()); } } -void TypeVisitor::VisitType_(const distributed::DTensorTypeNode* op) { +void TypeVisitor::VisitType_(const distributed::DTensorTypeObj* op) { this->VisitType(op->tensor_ty); } -void TypeVisitor::VisitType_(const TupleTypeNode* op) { +void TypeVisitor::VisitType_(const TupleTypeObj* op) { for (Type field : op->fields) { this->VisitType(field); } } -void TypeVisitor::VisitType_(const tvm::FuncTypeNode* op) { +void TypeVisitor::VisitType_(const tvm::FuncTypeObj* op) { // Native callbacks may appear inside Relax tuples. Their signature symbols // and tensor contracts belong to the native function, not the Relax scope. } -void TypeVisitor::VisitType_(const FuncTypeNode* op) { +void TypeVisitor::VisitType_(const FuncTypeObj* op) { if (op->params.has_value()) { for (Type param : op->params.value()) { this->VisitType(param); @@ -71,13 +71,13 @@ void TypeVisitor::VisitType_(const FuncTypeNode* op) { this->VisitType(op->ret); } -Type TypeMutator::VisitType_(const AnyTypeNode* op) { return ffi::GetRef(op); } +Type TypeMutator::VisitType_(const AnyTypeObj* op) { return ffi::GetRef(op); } -Type TypeMutator::VisitType_(const PrimTypeNode* op) { return ffi::GetRef(op); } +Type TypeMutator::VisitType_(const PrimTypeObj* op) { return ffi::GetRef(op); } -Type TypeMutator::VisitType_(const StringTypeNode* op) { return ffi::GetRef(op); } +Type TypeMutator::VisitType_(const StringTypeObj* op) { return ffi::GetRef(op); } -Type TypeMutator::VisitType_(const ShapeTypeNode* op) { +Type TypeMutator::VisitType_(const ShapeTypeObj* op) { if (!op->values.has_value()) { return ffi::GetRef(op); } @@ -93,7 +93,7 @@ Type TypeMutator::VisitType_(const ShapeTypeNode* op) { } } -Type TypeMutator::VisitType_(const TensorTypeNode* op) { +Type TypeMutator::VisitType_(const TensorTypeObj* op) { if (!op->shape.has_value()) { return ffi::GetRef(op); } @@ -108,12 +108,12 @@ Type TypeMutator::VisitType_(const TensorTypeNode* op) { } } -Type TypeMutator::VisitType_(const distributed::DTensorTypeNode* op) { +Type TypeMutator::VisitType_(const distributed::DTensorTypeObj* op) { TensorType tensor_ty = this->VisitType(op->tensor_ty).as_or_throw(); return distributed::DTensorType(tensor_ty, op->device_mesh, op->placement); } -Type TypeMutator::VisitType_(const TupleTypeNode* op) { +Type TypeMutator::VisitType_(const TupleTypeObj* op) { ffi::Array fields = op->fields.Map([this](const Type& ty) { return this->VisitType(ty); }); if (fields.same_as(op->fields)) { @@ -123,12 +123,12 @@ Type TypeMutator::VisitType_(const TupleTypeNode* op) { } } -Type TypeMutator::VisitType_(const tvm::FuncTypeNode* op) { +Type TypeMutator::VisitType_(const tvm::FuncTypeObj* op) { // Relax type rewrites preserve the native callback's signature. return ffi::GetRef(op); } -Type TypeMutator::VisitType_(const FuncTypeNode* op) { +Type TypeMutator::VisitType_(const FuncTypeObj* op) { ffi::Optional> params; if (op->params.has_value()) { diff --git a/src/relax/op/ccl/ccl.cc b/src/relax/op/ccl/ccl.cc index 5efe5c0d4d06..055bd48e5d09 100644 --- a/src/relax/op/ccl/ccl.cc +++ b/src/relax/op/ccl/ccl.cc @@ -90,7 +90,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { refl::GlobalDef().def("relax.op.ccl.allreduce", allreduce); } -Type InferTypeAllReduce(const CallNode* call_node) { +Type InferTypeAllReduce(const CallObj* call_node) { const Call call = ffi::GetRef(call_node); if (RequiresTensorInputNormalization(call)) return Type::Missing(); TensorType input_ty = GetUnaryInputTensorType(call); @@ -122,7 +122,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { refl::GlobalDef().def("relax.op.ccl.allgather", allgather); } -Type InferTypeAllGather(const CallNode* call_node) { +Type InferTypeAllGather(const CallObj* call_node) { const Call call = ffi::GetRef(call_node); TensorType input_ty = GetUnaryInputTensorType(call); @@ -159,7 +159,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { refl::GlobalDef().def("relax.op.ccl.broadcast_from_worker0", broadcast_from_worker0); } -Type InferTypeBroadcastFromZero(const CallNode* call_node) { +Type InferTypeBroadcastFromZero(const CallObj* call_node) { const Call call = ffi::GetRef(call_node); TensorType input_ty = GetUnaryInputTensorType(call); return input_ty; diff --git a/src/relax/op/distributed/binary.h b/src/relax/op/distributed/binary.h index f68f2bf1538f..a685488a4b16 100644 --- a/src/relax/op/distributed/binary.h +++ b/src/relax/op/distributed/binary.h @@ -49,8 +49,8 @@ Type InferDistTypeBroadcast(const Call& call, const BlockBuilder& ctx, FType f_c << "Unknown ndim is not supported for distributed operators."; int output_ndim = std::max(x1_ty->ndim, x2_ty->ndim); - const auto* x1_shape = x1_ty->shape.as(); - const auto* x2_shape = x2_ty->shape.as(); + const auto* x1_shape = x1_ty->shape.as(); + const auto* x2_shape = x2_ty->shape.as(); Type output_tensor_ty = Type::Missing(); // Shapes and ndims if (x1_shape && x2_shape) { diff --git a/src/relax/op/distributed/distributed.cc b/src/relax/op/distributed/distributed.cc index 5791451d2083..81d4b8eb41b5 100644 --- a/src/relax/op/distributed/distributed.cc +++ b/src/relax/op/distributed/distributed.cc @@ -97,7 +97,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { Type InferDistTypeRedistribute(const Call& call, const BlockBuilder& ctx) { const auto* attrs = call->attrs.as(); - const auto* ty = GetTypeAs(call->args[0]); + const auto* ty = GetTypeAs(call->args[0]); TVM_FFI_ICHECK(ty); return distributed::DTensorType(ty->tensor_ty, attrs->device_mesh, attrs->placement); } @@ -113,7 +113,7 @@ Type InferTypeCallTIRLocalView(const Call& call, const BlockBuilder& ctx) { if (call->ty_args.size() != 1) { TVM_FFI_VISIT_THROW(InternalError, call) << "ty_args should have exactly 1 output type."; } - TVM_FFI_ICHECK(call->args[0]->IsInstance()) + TVM_FFI_ICHECK(call->args[0]->IsInstance()) << "call_tir_local_view expects the first argument to be a GlobalVar referring to a TIR " "tirx::Function. " << "However, gets " << call->args[0]; @@ -131,7 +131,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { Expr MakeCallTIRLocalView(Expr func, Tuple args, ffi::Array out_ty_list) { for (const distributed::DTensorType& ty : out_ty_list) { - const auto* shape = ty->tensor_ty->shape.as(); + const auto* shape = ty->tensor_ty->shape.as(); TVM_FFI_ICHECK(shape != nullptr) << "out_ty of call_tir_local_view should have defined ShapeExpr as shape. " "However, one given type information is " diff --git a/src/relax/op/distributed/linear_algebra.cc b/src/relax/op/distributed/linear_algebra.cc index 6c81a855a8cb..a1c4a1cc6754 100644 --- a/src/relax/op/distributed/linear_algebra.cc +++ b/src/relax/op/distributed/linear_algebra.cc @@ -63,8 +63,8 @@ Type InferDistTypeMatmul(const Call& call, const BlockBuilder& ctx) { } int output_ndim = std::max(x1_ndim, x2_ndim) - x1_prepended - x2_appended; - const auto* x1_shape = x1_ty->shape.as(); - const auto* x2_shape = x2_ty->shape.as(); + const auto* x1_shape = x1_ty->shape.as(); + const auto* x2_shape = x2_ty->shape.as(); if (x1_shape == nullptr || x2_shape == nullptr) { TVM_FFI_VISIT_THROW(ValueError, call) << "input of distributed operator must have shape"; } diff --git a/src/relax/op/distributed/manipulate.cc b/src/relax/op/distributed/manipulate.cc index 58d954313974..b179b155c050 100644 --- a/src/relax/op/distributed/manipulate.cc +++ b/src/relax/op/distributed/manipulate.cc @@ -65,7 +65,7 @@ Type InferDistTypePermuteDims(const Call& call, const BlockBuilder& ctx) { return input_dtensor_tys[0]; } - const auto* data_shape = data_ty->shape.as(); + const auto* data_shape = data_ty->shape.as(); if (data_shape == nullptr) { TVM_FFI_VISIT_THROW(ValueError, call) << "Input of distributed operator must have known shape"; } @@ -92,7 +92,7 @@ Type InferDistTypeReshape(const Call& call, const BlockBuilder& ctx) { ffi::Array input_dtensor_tys = GetInputDTensorType(call, ctx); TensorType data_ty = input_dtensor_tys[0]->tensor_ty; - const auto* new_shape_ty = GetTypeAs(call->args[1]); + const auto* new_shape_ty = GetTypeAs(call->args[1]); if (!data_ty.defined()) { TVM_FFI_VISIT_THROW(TypeError, call) << "Reshape requires the input data to be Tensor. However, the given one is " @@ -106,7 +106,7 @@ Type InferDistTypeReshape(const Call& call, const BlockBuilder& ctx) { ffi::Optional> old_shape_values; if (data_ty->shape.has_value()) { - const auto* old_shape_ty = GetTypeAs(data_ty->shape.value()); + const auto* old_shape_ty = GetTypeAs(data_ty->shape.value()); TVM_FFI_ICHECK_NOTNULL(old_shape_ty); old_shape_values = old_shape_ty->values; } @@ -125,7 +125,7 @@ Type InferDistTypeReshape(const Call& call, const BlockBuilder& ctx) { Expr target_shape = call->args[1]; Type output_tensor_ty = Type::Missing(); // If shape values are defined, use them - if (target_shape->IsInstance() && new_shape_ty->values.has_value()) { + if (target_shape->IsInstance() && new_shape_ty->values.has_value()) { output_tensor_ty = TensorType(ShapeExpr(new_shape_ty->values.value()), data_ty->dtype); } else { output_tensor_ty = TensorType(target_shape, data_ty->dtype); diff --git a/src/relax/op/distributed/op.cc b/src/relax/op/distributed/op.cc index 2b13f20a7b2b..f2dbfbaa7938 100644 --- a/src/relax/op/distributed/op.cc +++ b/src/relax/op/distributed/op.cc @@ -30,7 +30,7 @@ Type InferDistTypeCallTIR(const Call& call, const BlockBuilder& ctx) { if (call->ty_args.size() != 1) { TVM_FFI_VISIT_THROW(InternalError, call) << "ty_args should have exact 1 output type."; } - TVM_FFI_ICHECK(call->args[0]->IsInstance()) + TVM_FFI_ICHECK(call->args[0]->IsInstance()) << "call_tir expects the first argument to be a GlobalVar referring to a TIR tirx::Function. " << "However, gets " << call->args[0]; return call->ty_args[0]; diff --git a/src/relax/op/distributed/statistical.cc b/src/relax/op/distributed/statistical.cc index 3bdd322d3502..d7b5b169691d 100644 --- a/src/relax/op/distributed/statistical.cc +++ b/src/relax/op/distributed/statistical.cc @@ -56,7 +56,7 @@ Type InferDistTypeStatistical(const Call& call, const BlockBuilder& ctx) { // - axes is not None, keepdims is false -> the returned shape does not contain the input axes. // - axes is not None, keepdims is true -> the returned shape has value 1 at the positions of the // input axes - const auto* data_shape = data_ty->shape.as(); + const auto* data_shape = data_ty->shape.as(); if (data_shape == nullptr) { TVM_FFI_VISIT_THROW(ValueError, call) << "Input of distributed operator must be known shape"; diff --git a/src/relax/op/distributed/unary.h b/src/relax/op/distributed/unary.h index a5dc0d6747d6..6c8f9a28eda1 100644 --- a/src/relax/op/distributed/unary.h +++ b/src/relax/op/distributed/unary.h @@ -49,7 +49,7 @@ Type InferDistTypeUnary(const Call& call, const BlockBuilder& ctx, FType f_compu << " requires the input tensor to have float dtype. However, the given input dtype is " << input_tensor_ty->dtype; } - auto output_ty = ffi::make_object(*input_tensor_ty.get()); + auto output_ty = ffi::make_object(*input_tensor_ty.get()); ffi::Optional computed_dtype = f_compute_out_dtype(input_tensor_ty); output_ty->dtype = computed_dtype; TensorType out_tensor_ty(output_ty); diff --git a/src/relax/op/distributed/utils.cc b/src/relax/op/distributed/utils.cc index b02fa9df040a..8e2a434d12ba 100644 --- a/src/relax/op/distributed/utils.cc +++ b/src/relax/op/distributed/utils.cc @@ -33,7 +33,7 @@ ffi::Array GetInputDTensorType(const Call& call, ffi::Array input_tensor_ty; input_tensor_ty.reserve(args.size()); for (const Expr& arg : args) { - const auto* ty = GetTypeAs(arg); + const auto* ty = GetTypeAs(arg); if (ty != nullptr) { input_tensor_ty.push_back(ffi::GetRef(ty)); } @@ -68,10 +68,10 @@ Type InferShardingSpec(const Call& call, const BlockBuilder& ctx, const Type& or } axis_group_graph.PropagateShardingSpec(); ffi::Array orig_output_tensor_tys; - if (const auto* tensor_ty = orig_output_ty.as()) { + if (const auto* tensor_ty = orig_output_ty.as()) { orig_output_tensor_tys.push_back(ffi::GetRef(tensor_ty)); } else { - const auto* tuple_ty = orig_output_ty.as(); + const auto* tuple_ty = orig_output_ty.as(); TVM_FFI_ICHECK(tuple_ty); for (const auto& ty : tuple_ty->fields) { orig_output_tensor_tys.push_back(ty.as_or_throw()); diff --git a/src/relax/op/image/resize.cc b/src/relax/op/image/resize.cc index ae683aee09a5..0a15523c65fb 100644 --- a/src/relax/op/image/resize.cc +++ b/src/relax/op/image/resize.cc @@ -173,7 +173,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { refl::GlobalDef().def("relax.op.image.resize2d", resize2d); } -Type InferTypeResize2D(const CallNode* call_node) { +Type InferTypeResize2D(const CallObj* call_node) { const Call call = ffi::GetRef(call_node); if (call->args.size() != 2) { TVM_FFI_VISIT_THROW(ValueError, call) @@ -181,9 +181,9 @@ Type InferTypeResize2D(const CallNode* call_node) { << call->args.size(); } - const auto* data_ty = GetTypeAs(call->args[0]); - const auto* size_ty = GetTypeAs(call->args[1]); - const auto* size_value = call->args[1].as(); + const auto* data_ty = GetTypeAs(call->args[0]); + const auto* size_ty = GetTypeAs(call->args[1]); + const auto* size_value = call->args[1].as(); if (data_ty == nullptr) { TVM_FFI_VISIT_THROW(TypeError, call) << "Resize2D expects the input data to be a Tensor, while the given data is " @@ -293,7 +293,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { refl::GlobalDef().def("relax.op.image.resize3d", resize3d); } -Type InferTypeResize3D(const CallNode* call_node) { +Type InferTypeResize3D(const CallObj* call_node) { const Call call = ffi::GetRef(call_node); if (call->args.size() != 2) { TVM_FFI_VISIT_THROW(ValueError, call) @@ -301,9 +301,9 @@ Type InferTypeResize3D(const CallNode* call_node) { << call->args.size(); } - const auto* data_ty = GetTypeAs(call->args[0]); - const auto* size_ty = GetTypeAs(call->args[1]); - const auto* size_value = call->args[1].as(); + const auto* data_ty = GetTypeAs(call->args[0]); + const auto* size_ty = GetTypeAs(call->args[1]); + const auto* size_value = call->args[1].as(); if (data_ty == nullptr) { TVM_FFI_VISIT_THROW(TypeError, call) << "Resize3D expects the input data to be a Tensor, while the given data is " @@ -406,7 +406,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { refl::GlobalDef().def("relax.op.image.grid_sample", grid_sample); } -Type InferTypeGridSample(const CallNode* call_node) { +Type InferTypeGridSample(const CallObj* call_node) { const Call call = ffi::GetRef(call_node); if (call->args.size() != 2) { TVM_FFI_VISIT_THROW(ValueError, call) @@ -414,8 +414,8 @@ Type InferTypeGridSample(const CallNode* call_node) { << call->args.size(); } - const auto* data_ty = GetTypeAs(call->args[0]); - const auto* grid_ty = GetTypeAs(call->args[1]); + const auto* data_ty = GetTypeAs(call->args[0]); + const auto* grid_ty = GetTypeAs(call->args[1]); if (data_ty == nullptr) { TVM_FFI_VISIT_THROW(TypeError, call) @@ -443,7 +443,7 @@ Type InferTypeGridSample(const CallNode* call_node) { ffi::Optional data_shape = CheckNdimPerLayoutAndGetShape(call, ffi::GetRef(data_ty), data_layout); - const auto* grid_shape = grid_ty->shape.as(); + const auto* grid_shape = grid_ty->shape.as(); if (!data_shape.has_value() || grid_shape == nullptr) { return TensorType(out_dtype, data_layout.ndim(), data_ty->vdevice); @@ -493,7 +493,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { refl::GlobalDef().def("relax.op.image.affine_grid", affine_grid); } -Type InferTypeAffineGrid(const CallNode* call_node) { +Type InferTypeAffineGrid(const CallObj* call_node) { const Call call = ffi::GetRef(call_node); if (call->args.size() != 2) { TVM_FFI_VISIT_THROW(ValueError, call) @@ -501,9 +501,9 @@ Type InferTypeAffineGrid(const CallNode* call_node) { << call->args.size(); } - const auto* data_ty = GetTypeAs(call->args[0]); - const auto* size_ty = GetTypeAs(call->args[1]); - const auto* size_value = call->args[1].as(); + const auto* data_ty = GetTypeAs(call->args[0]); + const auto* size_ty = GetTypeAs(call->args[1]); + const auto* size_value = call->args[1].as(); if (data_ty == nullptr) { TVM_FFI_VISIT_THROW(TypeError, call) @@ -531,10 +531,10 @@ Type InferTypeAffineGrid(const CallNode* call_node) { << data_ty->ndim; } - const auto* data_shape = data_ty->shape.as(); + const auto* data_shape = data_ty->shape.as(); if (data_shape != nullptr) { if (data_shape->values.size() >= 2) { - auto* dim1 = data_shape->values[1].as(); + auto* dim1 = data_shape->values[1].as(); if (dim1 != nullptr && dim1->value != spatial) { TVM_FFI_VISIT_THROW(ValueError, call) << "AffineGrid expects the second dimension of input to be " << spatial << ", but got " @@ -542,7 +542,7 @@ Type InferTypeAffineGrid(const CallNode* call_node) { } } if (data_shape->values.size() >= 3) { - auto* dim2 = data_shape->values[2].as(); + auto* dim2 = data_shape->values[2].as(); if (dim2 != nullptr && dim2->value != spatial + 1) { TVM_FFI_VISIT_THROW(ValueError, call) << "AffineGrid expects the third dimension of input to be " << spatial + 1 diff --git a/src/relax/op/memory/view.cc b/src/relax/op/memory/view.cc index 4da2756c2dd8..80116c5785f8 100644 --- a/src/relax/op/memory/view.cc +++ b/src/relax/op/memory/view.cc @@ -74,13 +74,13 @@ Type InferTypeView(const Call& call, const BlockBuilder& ctx) { << "but received " << arg_data << " with type " << ty; } }(); - auto view_shape_ty = [&]() -> const ShapeTypeNode* { + auto view_shape_ty = [&]() -> const ShapeTypeObj* { Type ty = GetType(arg_shape); if (HasVoidType(arg_shape)) { // No shape change is applied. The input tensor's shape is // kept as-is. return nullptr; - } else if (auto ptr = ty.as()) { + } else if (auto ptr = ty.as()) { // The `R.view` operation returns a different shape. return ptr; } else { @@ -111,7 +111,7 @@ Type InferTypeView(const Call& call, const BlockBuilder& ctx) { } static const Op null_value_op = Op::Get("relax.null_value"); - if (const CallNode* call_node = arg_value.as()) { + if (const CallObj* call_node = arg_value.as()) { if (call_node->op.same_as(null_value_op)) { // No datatype change is applied. This is the non-void // representation used for missing optional dtype arguments. @@ -126,14 +126,14 @@ Type InferTypeView(const Call& call, const BlockBuilder& ctx) { // themselves. However, dtype DataTypeImm uses // `AnyType`, so we need to inspect the argument itself // in this case. - if (auto dtype_imm = arg_value.as()) { + if (auto dtype_imm = arg_value.as()) { // We know the datatype for the view. TVM_FFI_CHECK(!PrimType(dtype_imm->value).IsVoid(), TypeError) << "Operator " << call->op << " expects the dtype argument to be a concrete tensor datatype, " << "but received void. Use relax.null_value() if no dtype change is requested."; return PrimType(dtype_imm->value); - } else if (ty.as()) { + } else if (ty.as()) { // The view changes the datatype, but we don't know what it is // being changed into. return std::nullopt; @@ -150,12 +150,12 @@ Type InferTypeView(const Call& call, const BlockBuilder& ctx) { if (HasVoidType(arg_relative_byte_offset)) { // No byte offset is specified, so no change is applied. return IntImm::Int64(0); - } else if (auto prim_ty = ty.as()) { + } else if (auto prim_ty = ty.as()) { TVM_FFI_CHECK_EQ(prim_ty->dtype, (DLDataType{kDLInt, 64, 1}), TypeError) << "Operator " << call->op << " expects the relative_byte_offset to be a 64-bit integer, but received " << arg_relative_byte_offset << ", which has type " << ty; - if (arg_relative_byte_offset.as()) { + if (arg_relative_byte_offset.as()) { // A primitive DataflowVar is a runtime value, not a symbolic value // that can be analyzed. return std::nullopt; @@ -341,7 +341,7 @@ Expr LowerBuiltinView(const BlockBuilder& bb, const Call& call) { } static const Op null_value_op = Op::Get("relax.null_value"); - if (const CallNode* call_node = arg.as()) { + if (const CallObj* call_node = arg.as()) { return call_node->op.same_as(null_value_op); } return false; @@ -415,7 +415,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { refl::GlobalDef().def("relax.op.memory.ensure_zero_offset", ensure_zero_offset); } -Type InferTypeEnsureZeroOffset(const CallNode* call_node) { +Type InferTypeEnsureZeroOffset(const CallObj* call_node) { const Call call = ffi::GetRef(call_node); if (call->args.size() != 1) { TVM_FFI_VISIT_THROW(ValueError, call) diff --git a/src/relax/op/nn/attention.cc b/src/relax/op/nn/attention.cc index 91145d6f4ad7..330a6f991e0e 100644 --- a/src/relax/op/nn/attention.cc +++ b/src/relax/op/nn/attention.cc @@ -97,9 +97,9 @@ Type InferTypeAttention(const Call& call, const BlockBuilder& ctx) { diag_dim(q_ty, "query"); diag_dim(k_ty, "key"); diag_dim(v_ty, "value"); - const ShapeExprNode* q_shape = q_ty->shape.as(); - const ShapeExprNode* k_shape = k_ty->shape.as(); - const ShapeExprNode* v_shape = v_ty->shape.as(); + const ShapeExprObj* q_shape = q_ty->shape.as(); + const ShapeExprObj* k_shape = k_ty->shape.as(); + const ShapeExprObj* v_shape = v_ty->shape.as(); PrimExpr num_batches = q_shape->values[0]; PrimExpr num_queries = q_shape->values[1]; PrimExpr num_heads = q_shape->values[2]; @@ -134,7 +134,7 @@ Type InferTypeAttention(const Call& call, const BlockBuilder& ctx) { if (input_ty.size() == 4) { TensorType bias_ty = input_ty[3]; - const ShapeExprNode* bias_shape = bias_ty->shape.as(); + const ShapeExprObj* bias_shape = bias_ty->shape.as(); if (bias_ty->ndim != 4) { TVM_FFI_VISIT_THROW(ValueError, call) << "The bias should have 4 dimensions." diff --git a/src/relax/op/nn/convolution.cc b/src/relax/op/nn/convolution.cc index 497f3924008d..7e2b27466b48 100644 --- a/src/relax/op/nn/convolution.cc +++ b/src/relax/op/nn/convolution.cc @@ -626,9 +626,9 @@ InferLayoutOutput InferLayoutConv2d( TensorType data_ty = data_si.as().value(); TensorType kernel_ty = kernel_si.as().value(); ffi::Optional data_shape = - ffi::GetRef(data_ty->shape.as()); + ffi::GetRef(data_ty->shape.as()); ffi::Optional kernel_shape = - ffi::GetRef(kernel_ty->shape.as()); + ffi::GetRef(kernel_ty->shape.as()); bool can_data_proved = CanProveLayoutTransform(input_layout, desired_data_layout, data_shape.value()->values); @@ -1247,9 +1247,9 @@ InferLayoutOutput InferLayoutConv2dTranspose( TensorType data_ty = data_si.as().value(); TensorType kernel_ty = kernel_si.as().value(); ffi::Optional data_shape = - ffi::GetRef(data_ty->shape.as()); + ffi::GetRef(data_ty->shape.as()); ffi::Optional kernel_shape = - ffi::GetRef(kernel_ty->shape.as()); + ffi::GetRef(kernel_ty->shape.as()); bool can_data_proved = CanProveLayoutTransform(input_layout, desired_data_layout, data_shape.value()->values); @@ -1495,9 +1495,9 @@ InferLayoutOutput InferLayoutConv3dTranspose( TensorType data_ty = data_si.as().value(); TensorType kernel_ty = kernel_si.as().value(); ffi::Optional data_shape = - ffi::GetRef(data_ty->shape.as()); + ffi::GetRef(data_ty->shape.as()); ffi::Optional kernel_shape = - ffi::GetRef(kernel_ty->shape.as()); + ffi::GetRef(kernel_ty->shape.as()); bool can_data_proved = CanProveLayoutTransform(input_layout, desired_data_layout, data_shape.value()->values); diff --git a/src/relax/op/nn/nn.cc b/src/relax/op/nn/nn.cc index b566cba643e8..d0701830ba21 100644 --- a/src/relax/op/nn/nn.cc +++ b/src/relax/op/nn/nn.cc @@ -345,7 +345,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { refl::GlobalDef().def("relax.op.nn.prelu", prelu); } -Type InferTypePRelu(const CallNode* call_node) { +Type InferTypePRelu(const CallObj* call_node) { const Call call = ffi::GetRef(call_node); if (RequiresTensorInputNormalization(call)) return Type::Missing(); TensorType data_ty = GetUnaryInputTensorType(call); @@ -376,7 +376,7 @@ InferLayoutOutput InferLayoutPRelu( // TODO(Siva): We could handle if the axis is not the sub indexed one. if ((layout->layout.has_value() ? layout->layout.value().ndim() : 0) != (layout->layout.has_value() ? layout->layout.value().ndim_primal() : 0)) { - const auto* tensor_ty = GetTypeAs(call->args[0]); + const auto* tensor_ty = GetTypeAs(call->args[0]); TVM_FFI_ICHECK(tensor_ty != nullptr) << "Invalid Call"; TVM_FFI_ICHECK(!tensor_ty->IsUnknownNdim()) << "Only support static ndim for now"; int ndim = tensor_ty->ndim; @@ -414,7 +414,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { refl::GlobalDef().def("relax.op.nn.softmax", softmax); } -Type InferTypeSoftmax(const CallNode* call_node) { +Type InferTypeSoftmax(const CallObj* call_node) { const Call call = ffi::GetRef(call_node); if (RequiresTensorInputNormalization(call)) return Type::Missing(); TensorType data_ty = GetUnaryInputTensorType(call); @@ -448,7 +448,7 @@ InferLayoutOutput InferLayoutSoftmax( // TODO(Siva): We could handle if the axis is not the sub indexed one. if ((layout->layout.has_value() ? layout->layout.value().ndim() : 0) != (layout->layout.has_value() ? layout->layout.value().ndim_primal() : 0)) { - const auto* tensor_ty = GetTypeAs(call->args[0]); + const auto* tensor_ty = GetTypeAs(call->args[0]); TVM_FFI_ICHECK(tensor_ty != nullptr) << "Invalid Call"; TVM_FFI_ICHECK(!tensor_ty->IsUnknownNdim()) << "Only support static ndim for now"; int ndim = tensor_ty->ndim; @@ -502,7 +502,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { refl::GlobalDef().def("relax.op.nn.pad", pad); } -Type InferTypePad(const CallNode* call_node) { +Type InferTypePad(const CallObj* call_node) { const Call call = ffi::GetRef(call_node); ffi::Array input_ty = GetInputTensorType(call); const auto* attrs = call->attrs.as(); @@ -513,7 +513,7 @@ Type InferTypePad(const CallNode* call_node) { ffi::Array out_shape; if (input_ty[0]->shape.has_value()) { // Compute output shape by adding corresponding pad width to each axis. - const auto* data_shape = input_ty[0]->shape.as(); + const auto* data_shape = input_ty[0]->shape.as(); for (int i = 0; i < ndim; i++) { // Sum pad width for this axis. PrimExpr added_width = IntImm::Int64(pad_width[2 * i] + pad_width[(2 * i) + 1]); @@ -548,7 +548,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { refl::GlobalDef().def("relax.op.nn.pixel_shuffle", pixel_shuffle); } -Type InferTypePixelShuffle(const CallNode* call_node) { +Type InferTypePixelShuffle(const CallObj* call_node) { const Call call = ffi::GetRef(call_node); ffi::Array input_ty = GetInputTensorType(call); const auto* attrs = call->attrs.as(); @@ -563,7 +563,7 @@ Type InferTypePixelShuffle(const CallNode* call_node) { return TensorType(input->dtype, ndim); } - const auto* shape = input->shape.as(); + const auto* shape = input->shape.as(); ffi::Array in_shape = shape->values; int channel_idx = ndim - 3; @@ -577,8 +577,8 @@ Type InferTypePixelShuffle(const CallNode* call_node) { PrimExpr r_expr = IntImm::Int32(r); PrimExpr r_squared = r_expr * r_expr; - const auto* c_in_imm = c_in.as(); - const auto* r2_imm = r_squared.as(); + const auto* c_in_imm = c_in.as(); + const auto* r2_imm = r_squared.as(); TVM_FFI_ICHECK_EQ(c_in_imm->value % r2_imm->value, 0) << "Number of input channels must be divisible by the square of the upscale factor"; @@ -646,7 +646,7 @@ bool NormCheckDtypeAndShape(const Call& call, const BlockBuilder& ctx, std::vector> axis_lengths; axis_lengths.reserve(n_input); - if (const auto* data_shape = data_ty->shape.as()) { + if (const auto* data_shape = data_ty->shape.as()) { std::vector lengths; lengths.reserve(n_axis); for (int d = 0; d < n_axis; ++d) { @@ -655,7 +655,7 @@ bool NormCheckDtypeAndShape(const Call& call, const BlockBuilder& ctx, axis_lengths.push_back(lengths); } for (int i = 1; i < n_input; ++i) { - if (const auto* shape = input_ty[i]->shape.as()) { + if (const auto* shape = input_ty[i]->shape.as()) { axis_lengths.push_back(shape->values); } } @@ -722,7 +722,7 @@ InferLayoutOutput InferLayoutBatchNorm( TVM_FFI_ICHECK(NoDesiredLayout(call, desired_layouts)); std::vector initial_layouts; for (size_t i = 0; i < 5; ++i) { - const auto* tensor_ty = GetTypeAs(call->args[i]); + const auto* tensor_ty = GetTypeAs(call->args[i]); TVM_FFI_ICHECK(tensor_ty != nullptr) << "Invalid Call"; TVM_FFI_ICHECK(!tensor_ty->IsUnknownNdim()) << "Only support known ndim"; initial_layouts.push_back(InitialLayoutDecision(tensor_ty->ndim)); @@ -734,7 +734,7 @@ InferLayoutOutput InferLayoutBatchNorm( // While dealing with sub layouts, its adviced to deal with batchnorm // on other ways like decomposing or fusion methods. // This handling is fail safe fallback. - const auto* input_ty = GetTypeAs(call->args[0]); + const auto* input_ty = GetTypeAs(call->args[0]); int ndim = input_ty->ndim; if ((layout->layout.has_value() ? layout->layout.value().ndim() : 0) != (layout->layout.has_value() ? layout->layout.value().ndim_primal() : 0)) { @@ -796,7 +796,7 @@ InferLayoutOutput InferLayoutLayerNorm( TVM_FFI_ICHECK(NoDesiredLayout(call, desired_layouts)); std::vector initial_layouts; for (size_t i = 0; i < 3; ++i) { - const auto* tensor_ty = GetTypeAs(call->args[i]); + const auto* tensor_ty = GetTypeAs(call->args[i]); TVM_FFI_ICHECK(tensor_ty != nullptr) << "Invalid Call"; TVM_FFI_ICHECK(!tensor_ty->IsUnknownNdim()) << "Only support known ndim"; initial_layouts.push_back(InitialLayoutDecision(tensor_ty->ndim)); @@ -806,7 +806,7 @@ InferLayoutOutput InferLayoutLayerNorm( LayoutDecision layout = GetLayoutDecision(var_layout_map, call->args[0]); ffi::ObjectPtr new_attrs = ffi::make_object(*attrs); - const auto* input_ty = GetTypeAs(call->args[0]); + const auto* input_ty = GetTypeAs(call->args[0]); int ndim = input_ty->ndim; std::vector new_axis; for (int64_t axis : attrs->axes) { @@ -873,7 +873,7 @@ Type InferTypeGroupNorm(const Call& call, const BlockBuilder& ctx) { << op << " expects that data must be float, but got " << data_ty->dtype; } sym::Analyzer analyzer = ctx->GetAnalyzer(); - const auto* data_shape = data_ty->shape.as(); + 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)) { TVM_FFI_VISIT_THROW(ValueError, call) @@ -889,7 +889,7 @@ Type InferTypeGroupNorm(const Call& call, const BlockBuilder& ctx) { TVM_FFI_VISIT_THROW(ValueError, call) << op << " expects that all inputs must have ndim=1, but got " << input_ty[i]->ndim; } else if (channel_axis != -1) { - const auto* shape = input_ty[i]->shape.as(); + const auto* shape = input_ty[i]->shape.as(); if (shape != nullptr && data_shape != nullptr) { PrimExpr channel_size = data_shape->values[channel_axis]; PrimExpr input_size = shape->values[0]; @@ -911,7 +911,7 @@ InferLayoutOutput InferLayoutGroupNorm( TVM_FFI_ICHECK(NoDesiredLayout(call, desired_layouts)); std::vector initial_layouts; for (size_t i = 0; i < 3; ++i) { - const auto* tensor_ty = GetTypeAs(call->args[i]); + const auto* tensor_ty = GetTypeAs(call->args[i]); TVM_FFI_ICHECK(tensor_ty != nullptr) << "Invalid Call"; TVM_FFI_ICHECK(!tensor_ty->IsUnknownNdim()) << "Only support known ndim"; initial_layouts.push_back(InitialLayoutDecision(tensor_ty->ndim)); @@ -981,7 +981,7 @@ Type InferTypeInstanceNorm(const Call& call, const BlockBuilder& ctx) { << channel_axis << ", axes: " << attrs->axes; } } - const auto* data_shape = data_ty->shape.as(); + const auto* data_shape = data_ty->shape.as(); sym::Analyzer analyzer = ctx->GetAnalyzer(); for (int i = 1; i < static_cast(op->args_info.size()); ++i) { if (input_ty[i]->dtype != data_ty->dtype) { @@ -992,7 +992,7 @@ Type InferTypeInstanceNorm(const Call& call, const BlockBuilder& ctx) { TVM_FFI_VISIT_THROW(ValueError, call) << op << " expects that all inputs must have ndim=1, but got " << input_ty[i]->ndim; } - const auto* shape = input_ty[i]->shape.as(); + const auto* shape = input_ty[i]->shape.as(); if (shape != nullptr && data_shape != nullptr) { PrimExpr channel_size = data_shape->values[channel_axis]; PrimExpr input_size = shape->values[0]; @@ -1013,7 +1013,7 @@ InferLayoutOutput InferLayoutInstanceNorm( TVM_FFI_ICHECK(NoDesiredLayout(call, desired_layouts)); std::vector initial_layouts; for (size_t i = 0; i < 3; ++i) { - const auto* tensor_ty = GetTypeAs(call->args[i]); + const auto* tensor_ty = GetTypeAs(call->args[i]); TVM_FFI_ICHECK(tensor_ty != nullptr) << "Invalid Call"; TVM_FFI_ICHECK(!tensor_ty->IsUnknownNdim()) << "Only support known ndim"; initial_layouts.push_back(InitialLayoutDecision(tensor_ty->ndim)); @@ -1075,7 +1075,7 @@ InferLayoutOutput InferLayoutRMSNorm( TVM_FFI_ICHECK(NoDesiredLayout(call, desired_layouts)); std::vector initial_layouts; for (size_t i = 0; i < 2; ++i) { - const auto* tensor_ty = GetTypeAs(call->args[i]); + const auto* tensor_ty = GetTypeAs(call->args[i]); TVM_FFI_ICHECK(tensor_ty != nullptr) << "Invalid Call"; TVM_FFI_ICHECK(!tensor_ty->IsUnknownNdim()) << "Only support known ndim"; initial_layouts.push_back(InitialLayoutDecision(tensor_ty->ndim)); @@ -1118,7 +1118,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { refl::GlobalDef().def("relax.op.nn.dropout", dropout); } -Type InferTypeDropout(const CallNode* call_node) { +Type InferTypeDropout(const CallObj* call_node) { const Call call = ffi::GetRef(call_node); TensorType data_ty = GetUnaryInputTensorType(call); return TupleType({data_ty, data_ty}); @@ -1156,12 +1156,12 @@ Type InferTypeCrossEntropy(const Call& call, const BlockBuilder& ctx) { ffi::Optional> pred_shape_value; if (pred_ty->shape.has_value()) { - pred_shape_value = GetTypeAs(pred_ty->shape.value())->values; + pred_shape_value = GetTypeAs(pred_ty->shape.value())->values; } ffi::Optional> label_shape_value; if (label_ty->shape.has_value()) { - label_shape_value = GetTypeAs(label_ty->shape.value())->values; + label_shape_value = GetTypeAs(label_ty->shape.value())->values; } if (pred_shape_value.has_value() && label_shape_value.has_value()) { @@ -1230,11 +1230,11 @@ Type InferTypeNLLLoss(const Call& call, const BlockBuilder& ctx) { TVM_FFI_VISIT_THROW(ValueError, call) << "NLLLoss op should take 2 or 3 arguments"; } - const auto* pred_ty = GetTypeAs(call->args[0]); - const auto* tgt_ty = GetTypeAs(call->args[1]); - const TensorTypeNode* wgt_ty = nullptr; + const auto* pred_ty = GetTypeAs(call->args[0]); + const auto* tgt_ty = GetTypeAs(call->args[1]); + const TensorTypeObj* wgt_ty = nullptr; if (call->args.size() == 3) { - wgt_ty = GetTypeAs(call->args[2]); + wgt_ty = GetTypeAs(call->args[2]); if (wgt_ty == nullptr) { TVM_FFI_VISIT_THROW(TypeError, call) << "NLLLoss requires the argument weights to be Tensor. However, the given one is " @@ -1307,7 +1307,7 @@ Type InferTypeNLLLoss(const Call& call, const BlockBuilder& ctx) { ffi::Optional> pred_shape_value; if (pred_ty->shape.has_value()) { - pred_shape_value = GetTypeAs(pred_ty->shape.value())->values; + pred_shape_value = GetTypeAs(pred_ty->shape.value())->values; } if (pred_shape_value.has_value()) { if (pred_shape_value.value().size() == 1) { @@ -1330,7 +1330,7 @@ Type InferTypeNLLLoss(const Call& call, const BlockBuilder& ctx) { ffi::Optional> tgt_shape_value; if (tgt_ty->shape.has_value()) { - tgt_shape_value = GetTypeAs(tgt_ty->shape.value())->values; + tgt_shape_value = GetTypeAs(tgt_ty->shape.value())->values; } if (tgt_shape_value.has_value()) { if (tgt_shape_value.value().empty()) { @@ -1383,7 +1383,7 @@ Type InferTypeNLLLoss(const Call& call, const BlockBuilder& ctx) { if (wgt_ty != nullptr) { ffi::Optional> wgt_shape_value; if (wgt_ty->shape.has_value()) { - wgt_shape_value = GetTypeAs(wgt_ty->shape.value())->values; + wgt_shape_value = GetTypeAs(wgt_ty->shape.value())->values; } if (wgt_shape_value.has_value()) { TVM_FFI_ICHECK(wgt_shape_value.value().size() == 1); @@ -1403,7 +1403,7 @@ Type InferTypeNLLLoss(const Call& call, const BlockBuilder& ctx) { if (reduction == "none") { // () or (N,) or (N, d1, d2, ..., dk) - if (pred_ty->shape.as()) { + if (pred_ty->shape.as()) { return TensorType(ShapeExpr(output_shape), output_dtype, vdevice); } else { int output_ndim = pred_ty->ndim == kUnknownNDim ? kUnknownNDim : pred_ty->ndim - 1; @@ -1438,7 +1438,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { refl::GlobalDef().def("relax.op.nn.batch_flatten", batch_flatten); } -Type InferTypeBatchFlatten(const CallNode* call_node) { +Type InferTypeBatchFlatten(const CallObj* call_node) { const Call call = ffi::GetRef(call_node); TensorType data_ty = GetUnaryInputTensorType(call); @@ -1456,7 +1456,7 @@ Type InferTypeBatchFlatten(const CallNode* call_node) { return data_ty; } - const auto* data_shape = data_ty->shape.as(); + const auto* data_shape = data_ty->shape.as(); if (data_shape == nullptr) { return TensorType(data_ty->dtype, /*ndim=*/2, data_ty->vdevice); } diff --git a/src/relax/op/nn/pooling.cc b/src/relax/op/nn/pooling.cc index 2d9c8fc70363..bcb47fb864d6 100644 --- a/src/relax/op/nn/pooling.cc +++ b/src/relax/op/nn/pooling.cc @@ -293,7 +293,7 @@ InferLayoutOutput InferLayoutPool1d( const Call& call, const ffi::Map>& desired_layouts, const VarLayoutMap& var_layout_map) { TVM_FFI_ICHECK(NoDesiredLayout(call, desired_layouts)); - const auto* tensor_ty = GetTypeAs(call); + const auto* tensor_ty = GetTypeAs(call); TVM_FFI_ICHECK(tensor_ty != nullptr) << "Invalid Call"; TVM_FFI_ICHECK_EQ(tensor_ty->ndim, 3) << "Unsupported initial layout"; const auto* attrs = call->attrs.as(); @@ -432,7 +432,7 @@ InferLayoutOutput InferLayoutPool2d( const Call& call, const ffi::Map>& desired_layouts, const VarLayoutMap& var_layout_map) { TVM_FFI_ICHECK(NoDesiredLayout(call, desired_layouts)); - const auto* tensor_ty = GetTypeAs(call); + const auto* tensor_ty = GetTypeAs(call); TVM_FFI_ICHECK(tensor_ty != nullptr) << "Invalid Call"; TVM_FFI_ICHECK_EQ(tensor_ty->ndim, 4) << "Unsupported initial layout"; const auto* attrs = call->attrs.as(); @@ -449,8 +449,7 @@ InferLayoutOutput InferLayoutPool2d( TransposeSubLayoutLike(attrs->layout, InitialLayout(4), layout->layout.value()); auto data_si = GetType(call->args[0]); TensorType data_ty = data_si.as().value(); - ffi::Optional data_shape = - ffi::GetRef(data_ty->shape.as()); + ffi::Optional data_shape = ffi::GetRef(data_ty->shape.as()); if (CanProveLayoutTransform(in_layout, desired_layout, data_shape.value()->values)) { // Not handling out_layout being different from in_layout now. Any use case ? new_attrs->layout = desired_layout.value().name(); @@ -609,7 +608,7 @@ InferLayoutOutput InferLayoutPool3d( const Call& call, const ffi::Map>& desired_layouts, const VarLayoutMap& var_layout_map) { TVM_FFI_ICHECK(NoDesiredLayout(call, desired_layouts)); - const auto* tensor_ty = GetTypeAs(call); + const auto* tensor_ty = GetTypeAs(call); TVM_FFI_ICHECK(tensor_ty != nullptr) << "Invalid Call"; TVM_FFI_ICHECK_EQ(tensor_ty->ndim, 5) << "Unsupported initial layout"; const auto* attrs = call->attrs.as(); @@ -717,7 +716,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { refl::GlobalDef().def("relax.op.nn.adaptive_avg_pool1d", adaptive_avg_pool1d); } -Type InferTypeAdaptiveAvgPool1D(const CallNode* call_node) { +Type InferTypeAdaptiveAvgPool1D(const CallObj* call_node) { const Call call = ffi::GetRef(call_node); TensorType data_ty = GetUnaryInputTensorType(call); @@ -753,7 +752,7 @@ InferLayoutOutput InferLayoutAdaptiveAvgPool1D( const Call& call, const ffi::Map>& desired_layouts, const VarLayoutMap& var_layout_map) { TVM_FFI_ICHECK(NoDesiredLayout(call, desired_layouts)); - const auto* tensor_ty = GetTypeAs(call); + const auto* tensor_ty = GetTypeAs(call); TVM_FFI_ICHECK(tensor_ty != nullptr) << "Invalid Call"; TVM_FFI_ICHECK_EQ(tensor_ty->ndim, 3) << "Unsupported initial layout"; const auto* attrs = call->attrs.as(); @@ -804,7 +803,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { refl::GlobalDef().def("relax.op.nn.adaptive_avg_pool2d", adaptive_avg_pool2d); } -Type InferTypeAdaptiveAvgPool2D(const CallNode* call_node) { +Type InferTypeAdaptiveAvgPool2D(const CallObj* call_node) { const Call call = ffi::GetRef(call_node); TensorType data_ty = GetUnaryInputTensorType(call); @@ -841,7 +840,7 @@ InferLayoutOutput InferLayoutAdaptiveAvgPool2D( const Call& call, const ffi::Map>& desired_layouts, const VarLayoutMap& var_layout_map) { TVM_FFI_ICHECK(NoDesiredLayout(call, desired_layouts)); - const auto* tensor_ty = GetTypeAs(call); + const auto* tensor_ty = GetTypeAs(call); TVM_FFI_ICHECK(tensor_ty != nullptr) << "Invalid Call"; TVM_FFI_ICHECK_EQ(tensor_ty->ndim, 4) << "Unsupported initial layout"; const auto* attrs = call->attrs.as(); @@ -857,8 +856,7 @@ InferLayoutOutput InferLayoutAdaptiveAvgPool2D( TransposeSubLayoutLike(attrs->layout, InitialLayout(4), layout->layout.value()); auto data_si = GetType(call->args[0]); TensorType data_ty = data_si.as().value(); - ffi::Optional data_shape = - ffi::GetRef(data_ty->shape.as()); + ffi::Optional data_shape = ffi::GetRef(data_ty->shape.as()); if (CanProveLayoutTransform(in_layout, desired_layout, data_shape.value()->values)) { // Not handling out_layout being different from in_layout now. Any use case ? new_attrs->layout = desired_layout.value().name(); @@ -913,7 +911,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { refl::GlobalDef().def("relax.op.nn.adaptive_avg_pool3d", adaptive_avg_pool3d); } -Type InferTypeAdaptiveAvgPool3D(const CallNode* call_node) { +Type InferTypeAdaptiveAvgPool3D(const CallObj* call_node) { const Call call = ffi::GetRef(call_node); TensorType data_ty = GetUnaryInputTensorType(call); @@ -951,7 +949,7 @@ InferLayoutOutput InferLayoutAdaptiveAvgPool3D( const Call& call, const ffi::Map>& desired_layouts, const VarLayoutMap& var_layout_map) { TVM_FFI_ICHECK(NoDesiredLayout(call, desired_layouts)); - const auto* tensor_ty = GetTypeAs(call); + const auto* tensor_ty = GetTypeAs(call); TVM_FFI_ICHECK(tensor_ty != nullptr) << "Invalid Call"; TVM_FFI_ICHECK_EQ(tensor_ty->ndim, 5) << "Unsupported initial layout"; const auto* attrs = call->attrs.as(); diff --git a/src/relax/op/op.cc b/src/relax/op/op.cc index 657379a64bdf..d3b02cffbfae 100644 --- a/src/relax/op/op.cc +++ b/src/relax/op/op.cc @@ -82,7 +82,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { } bool EqualConstInt(const PrimExpr& lhs, int64_t value) { - if (const auto* pvalue = lhs.as()) { + if (const auto* pvalue = lhs.as()) { return pvalue->value == value; } return false; @@ -90,37 +90,37 @@ bool EqualConstInt(const PrimExpr& lhs, int64_t value) { bool EqualCheck(const PrimExpr& lhs, const PrimExpr& rhs) { PrimExpr diff = lhs - rhs; - if (const auto* pdiff = diff.as()) { + if (const auto* pdiff = diff.as()) { return pdiff->value == 0; } tvm::sym::Analyzer ana; diff = ana->Simplify(diff); - if (const auto* pdiff = diff.as()) { + if (const auto* pdiff = diff.as()) { return pdiff->value == 0; } return false; } -Type ReturnVoidType(const CallNode*) { return TupleType(ffi::Array()); } +Type ReturnVoidType(const CallObj*) { return TupleType(ffi::Array()); } -Type ReturnAnyType(const CallNode*) { return AnyType(); } +Type ReturnAnyType(const CallObj*) { return AnyType(); } -Type InferTypeShapeOf(const CallNode* call_node) { +Type InferTypeShapeOf(const CallObj* call_node) { const Call call = ffi::GetRef(call_node); // use the Type of the argument auto arg_ty = GetType(call->args[0]); - auto* tensor_ty = GetType(call->args[0]).as(); + auto* tensor_ty = GetType(call->args[0]).as(); TVM_FFI_ICHECK(tensor_ty) << "shape_of expects a tensor input, but received " << arg_ty << "; use MatchCast if necessary"; if (tensor_ty->ndim == kUnknownNDim) { return ShapeType(kUnknownNDim); } // if the tensor shape is a Relax var or omitted, do not try to construct a shape expr from it - if (!tensor_ty->shape.has_value() || tensor_ty->shape.as()) { + if (!tensor_ty->shape.has_value() || tensor_ty->shape.as()) { return ShapeType(tensor_ty->ndim); } // otherwise, copy over the values from the tensor shape - auto* tensor_shape = tensor_ty->shape.as(); + auto* tensor_shape = tensor_ty->shape.as(); TVM_FFI_ICHECK(tensor_shape); return ShapeType(tensor_shape->values); } @@ -141,7 +141,7 @@ Type InferTypeCallPurePacked(const Call& call, const BlockBuilder& ctx) { return Call::ReinferType(inner.get()); } // Other callees must remain opaque packed functions. - TVM_FFI_ICHECK(!callee.as()) << "call_pure_packed cannot be used with an op node"; + TVM_FFI_ICHECK(!callee.as()) << "call_pure_packed cannot be used with an op node"; auto opt = MatchType(callee); TVM_FFI_ICHECK(opt) << "Callee must have a function type"; FuncType finfo = opt.value(); @@ -159,20 +159,20 @@ Type InferTypeCallPurePacked(const Call& call, const BlockBuilder& ctx) { } } -void ValidateCallPurePacked(const CallNode* call) { +void ValidateCallPurePacked(const CallObj* call) { TVM_FFI_CHECK(call->args.size() >= 1, TypeError) << "call_pure_packed expects a function argument"; for (const ffi::Any* arg = call->args.GetArrayObj()->begin(); arg != call->args.GetArrayObj()->end(); ++arg) { TVM_FFI_CHECK( - *arg != nullptr && ffi::details::AnyUnsafe::CheckAnyViewStrict(ffi::AnyView(*arg)), + *arg != nullptr && ffi::detail::AnyUnsafe::CheckAnyViewStrict(ffi::AnyView(*arg)), TypeError) << "call_pure_packed has an invalid value argument"; } for (const ffi::Any* ty_arg = call->ty_args.GetArrayObj()->begin(); ty_arg != call->ty_args.GetArrayObj()->end(); ++ty_arg) { TVM_FFI_CHECK(*ty_arg != nullptr && - ffi::details::AnyUnsafe::CheckAnyViewStrict(ffi::AnyView(*ty_arg)), + ffi::detail::AnyUnsafe::CheckAnyViewStrict(ffi::AnyView(*ty_arg)), TypeError) << "call_pure_packed has an invalid type argument"; } @@ -203,7 +203,7 @@ Expr NormalizeCallPurePacked(const BlockBuilder& ctx, Call call) { TVM_FFI_STATIC_INIT_BLOCK() { OpDef("relax.call_pure_packed") .set_validator(ffi::reflection::NativeFunctionView::FromNative<&ValidateCallPurePacked>()) + const CallObj*)>::FromNative<&ValidateCallPurePacked>()) .signature(sig::arg("func", "The first argument is the function being called. The rest are the " "arguments to that function."), @@ -242,7 +242,7 @@ Type InferTypeCallInplacePacked(const Call& call, const BlockBuilder& ctx) { // the callee must be an opaque function auto callee = call->args[0]; - TVM_FFI_ICHECK(!callee.as()) << "call_pure_packed cannot be used with an op node"; + TVM_FFI_ICHECK(!callee.as()) << "call_pure_packed cannot be used with an op node"; auto opt = MatchType(callee); TVM_FFI_ICHECK(opt) << "Callee must have a function type"; FuncType finfo = opt.value(); @@ -295,7 +295,7 @@ Type InferTypeCallInplacePacked(const Call& call, const BlockBuilder& ctx) { << "the in-place argument at index " << (arg_idx - 1) << ": " << ret << " vs " << arg_ty; } } else { - auto* tup_info = ret.as(); + auto* tup_info = ret.as(); if (!tup_info) { TVM_FFI_VISIT_THROW(ValueError, call) << "Multiple outputs given via the inplace indices " "but the derived Type is not a tuple"; @@ -369,7 +369,7 @@ bool IsGlobalMemoryScope(runtime::MemoryScope scope) { Type TIRxPackedValueType(const Type& type, bool is_result) { if (is_result && IsVoidType(type)) return VoidType(); - if (const auto* scalar = type.as()) { + if (const auto* scalar = type.as()) { DLDataType dtype = scalar->dtype; bool supported = dtype.lanes == 1 && @@ -380,15 +380,16 @@ Type TIRxPackedValueType(const Type& type, bool is_result) { << "R.call_tir_packed does not support packed scalar type " << type; return type; } - if (type.as()) return AnyType(); + if (type.as()) return AnyType(); if (!is_result) { - if (const auto* tensor = type.as()) { + if (const auto* tensor = type.as()) { bool regular = tirx::MatchShapeStridePattern(ffi::GetRef(tensor)).has_value(); - const auto* tile = tensor->layout.as(); + const auto* tile = tensor->layout.as(); // Runtime origin matching is handled separately by the native call boundary. bool regular_contiguous = - regular && (!tensor->layout.has_value() || (tile && tile->IsContiguous())); + regular && (!tensor->layout.has_value() || + (tile && ffi::GetRef(tile).IsContiguous())); TVM_FFI_CHECK(regular_contiguous, TypeError) << "A Relax-to-TIRx call cannot implicitly convert tensor layout or allocated address: " << type; @@ -404,10 +405,10 @@ Type TIRxPackedValueType(const Type& type, bool is_result) { TVM_FFI_UNREACHABLE(); } -const tvm::FuncTypeNode* NativeTIRxSignature(const Expr& callee) { - TVM_FFI_CHECK(callee.as(), TypeError) +const tvm::FuncTypeObj* NativeTIRxSignature(const Expr& callee) { + TVM_FFI_CHECK(callee.as(), TypeError) << "R.call_tir_packed requires a GlobalVar referring to a native TIRx function"; - const auto* signature = callee->ty.as(); + const auto* signature = callee->ty.as(); TVM_FFI_CHECK(signature, TypeError) << "R.call_tir_packed requires a native TIRx function signature, but received " << callee->ty; return signature; @@ -428,7 +429,7 @@ Tuple InlineTIRxArguments(const BlockBuilder& ctx, const Expr& argument) { value = bound.value(); if (auto tuple = value.as()) return tuple.value(); } - const auto* type = argument->ty.as(); + const auto* type = argument->ty.as(); TVM_FFI_CHECK(type, TypeError) << "R.call_tir_packed expects an argument tuple"; ffi::Array fields; for (size_t i = 0; i < type->fields.size(); ++i) { @@ -445,14 +446,14 @@ tvm::FuncType ContextualTIRxSignature(const BlockBuilder& ctx, const Expr& calle IRModule mod = ctx->GetContextIRModule(); auto it = mod->functions.find(callee.as_or_throw()); if (it == mod->functions.end()) return signature; - const auto* function = (*it).second.as(); + const auto* function = (*it).second.as(); TVM_FFI_CHECK(function, TypeError) << "A Relax-to-TIRx call must refer to a native TIRx function"; // Specialization may replace the function while retaining its typed GlobalVar. // Use the available definition at this boundary without modifying the symbol. signature = function->ty.as_or_throw(); auto substitutor = ffi::make_object(); for (size_t i = 0; i < arguments->fields.size() && i < function->params.size(); ++i) { - if (!function->params[i]->ty.as()) continue; + if (!function->params[i]->ty.as()) continue; Expr value = arguments->fields[i]; while (auto var = value.as()) { auto binding = ctx->LookupBinding(var.value()); @@ -470,14 +471,14 @@ tvm::FuncType ContextualTIRxSignature(const BlockBuilder& ctx, const Expr& calle void CheckTIRxCarrier(const Type& native, const Type& actual, bool allow_storage_specialization = false) { - if (native.as()) { - TVM_FFI_CHECK(!actual.as() && !actual.as(), TypeError) + if (native.as()) { + TVM_FFI_CHECK(!actual.as() && !actual.as(), TypeError) << "A Relax-to-TIRx pointer argument requires a handle-compatible object or Any, received " << actual; } - if (const auto* tensor = native.as()) { + if (const auto* tensor = native.as()) { auto native_scope = tensor->scope; - const auto* argument = actual.as(); + const auto* argument = actual.as(); ffi::Optional actual_scope; if (argument && argument->vdevice.has_value()) { const auto& scope = argument->vdevice.value()->memory_scope; @@ -499,11 +500,11 @@ void CheckTIRxArguments(const tvm::FuncType& signature, const ffi::Array& FuncType adapted = TIRxToRelaxFuncType(signature); auto params = adapted->params.value(); for (size_t i = 0; i < params.size() && i < arguments.size(); ++i) { - if (!defer_storage_check || !signature->arg_types[i].as()) { + if (!defer_storage_check || !signature->arg_types[i].as()) { CheckTIRxCarrier(signature->arg_types[i], arguments[i]->ty); } - const auto* expected = params[i].as(); - const auto* actual = arguments[i]->ty.as(); + const auto* expected = params[i].as(); + const auto* actual = arguments[i]->ty.as(); if (!expected || !actual) continue; // Unknown Relax metadata defers to the native packed checks. Keep every // known constraint, so partial information never hides a contradiction. @@ -518,7 +519,7 @@ void CheckTIRxArguments(const tvm::FuncType& signature, const ffi::Array& DeriveCallRetType(adapted, Call(Type::Missing(), Var("callee", adapted), arguments), ctx); } -void CheckFreshTIRxDestination(const tirx::TensorTypeNode* tensor) { +void CheckFreshTIRxDestination(const tirx::TensorTypeObj* tensor) { if (!tensor->layout.has_value()) return; sym::Analyzer analyzer; auto pattern = tirx::MatchShapeStridePattern(ffi::GetRef(tensor)); @@ -537,24 +538,24 @@ void CheckFreshTIRxDestination(const tirx::TensorTypeNode* tensor) { } } -void ValidateCallTIRPacked(const CallNode* call) { +void ValidateCallTIRPacked(const CallObj* call) { TVM_FFI_CHECK(call->args.size() == 2 && call->ty_args.empty(), TypeError) << "R.call_tir_packed expects a callee and an argument tuple, without type arguments"; TVM_FFI_CHECK(!call->attrs.has_value(), TypeError) << "R.call_tir_packed does not accept attributes"; const auto* native = NativeTIRxSignature(call->args[0]); - const auto* arguments = call->args[1]->ty.as(); + const auto* arguments = call->args[1]->ty.as(); TVM_FFI_CHECK(arguments, TypeError) << "R.call_tir_packed expects an argument tuple"; TVM_FFI_CHECK_EQ(native->arg_types.size(), arguments->fields.size(), TypeError) << "R.call_tir_packed requires one argument for each native parameter"; for (size_t i = 0; i < native->arg_types.size(); ++i) { - if (native->arg_types[i].as()) { + if (native->arg_types[i].as()) { const Type& argument = arguments->fields[i]; - TVM_FFI_CHECK(!argument.as() && !argument.as(), TypeError) + TVM_FFI_CHECK(!argument.as() && !argument.as(), TypeError) << "R.call_tir_packed pointer parameter " << i << " requires a handle-compatible object or Any, but received " << argument; } - TVM_FFI_CHECK(!arguments->fields[i].as(), TypeError) + TVM_FFI_CHECK(!arguments->fields[i].as(), TypeError) << "R.call_tir_packed requires distributed tensors to be explicitly lowered"; } ffi::Array args = @@ -565,7 +566,7 @@ void ValidateCallTIRPacked(const CallNode* call) { true); } -Type InferTypeCallTIRPacked(const CallNode* call) { +Type InferTypeCallTIRPacked(const CallObj* call) { TVM_FFI_CHECK_EQ(call->args.size(), 2, TypeError) << "R.call_tir_packed expects a callee and an argument tuple"; return TIRxPackedValueType(NativeTIRxSignature(call->args[0])->ret_type, true); @@ -589,7 +590,7 @@ Expr NormalizeCallTIRPacked(const BlockBuilder& ctx, Call call) { TVM_FFI_STATIC_INIT_BLOCK() { OpDef("relax.call_tir_packed") .set_validator(ffi::reflection::NativeFunctionView::FromNative<&ValidateCallTIRPacked>()) + const CallObj*)>::FromNative<&ValidateCallTIRPacked>()) .signature(sig::arg("func", "The native TIRx function."), sig::arg("args", "Every native argument, in parameter order.")) .set_native_fn_attr<&InferTypeCallTIRPacked>(tvm::op_attr::kInferType) @@ -645,7 +646,7 @@ ffi::Optional InferCallTIROutputTypeFromArguments( << "However, the first argument was of type " << callee_ty; auto callee_params = callee_ty->params.value(); - const TupleTypeNode* args = arg_ty.as(); + const TupleTypeObj* args = arg_ty.as(); TVM_FFI_CHECK(args, TypeError) << "The second argument to `R.call_tir` must be a tuple, " << "but instead received expression of type " << arg_ty; @@ -659,12 +660,12 @@ ffi::Optional InferCallTIROutputTypeFromArguments( size_t num_input_arguments = args->fields.size(); const auto& native_params = opt_native_ty.value()->arg_types; for (size_t i = 0; i < num_input_arguments && i < native_params.size(); ++i) { - if (!args->fields[i].as()) { + if (!args->fields[i].as()) { CheckTIRxCarrier(native_params[i], args->fields[i], true); } } for (size_t i = num_input_arguments; i < native_params.size(); ++i) { - const auto* destination = native_params[i].as(); + const auto* destination = native_params[i].as(); TVM_FFI_CHECK(destination, TypeError) << "R.call_tir destination parameter " << i << " must have a native tensor type"; CheckFreshTIRxDestination(destination); @@ -680,9 +681,9 @@ ffi::Optional InferCallTIROutputTypeFromArguments( // shape for `R.dist.call_tir` calls, as it depends on the lowering // of DistIR into regular Relax. std::function contains_dtensor = [&contains_dtensor](Type ty) -> bool { - if (ty.as()) { + if (ty.as()) { return true; - } else if (auto tuple = ty.as()) { + } else if (auto tuple = ty.as()) { return std::any_of(tuple->fields.begin(), tuple->fields.end(), contains_dtensor); } else { return false; @@ -696,8 +697,8 @@ ffi::Optional InferCallTIROutputTypeFromArguments( // in `callee_params` may contain dynamic shape parameters that are // not present in the caller's scope. The `DeriveCallRetType` // utility can infer the value of dynamic parameters in - // `FuncTypeNode::ret` based on definitions in - // `FuncTypeNode::params`, inferring the correct values in the + // `FuncTypeObj::ret` based on definitions in + // `FuncTypeObj::params`, inferring the correct values in the // caller's scope. // // Since the callee of `R.call_tir` is provided with output @@ -753,7 +754,7 @@ void CheckTIRxDestinations(const tvm::FuncType& signature, const Type& arguments const Type& result, const ffi::Optional>& inplace_indices) { ffi::Array outputs; - if (const auto* tuple = result.as()) { + if (const auto* tuple = result.as()) { outputs = tuple->fields; } else { outputs.push_back(result); @@ -770,19 +771,19 @@ void CheckTIRxDestinations(const tvm::FuncType& signature, const Type& arguments for (size_t i = 0; i < outputs.size(); ++i) { if (inplace_indices && inplace_indices.value()[i] >= 0) continue; if (destination < signature->arg_types.size() && - !outputs[i].as()) { + !outputs[i].as()) { CheckTIRxCarrier(signature->arg_types[destination], outputs[i], true); } ++destination; } } -Type InferTypeCallTIR(const CallNode* call_node) { +Type InferTypeCallTIR(const CallObj* call_node) { const Call call = ffi::GetRef(call_node); if (call->ty_args.size() != 1) { TVM_FFI_VISIT_THROW(InternalError, call) << "ty_args should have exactly 1 output type."; } - TVM_FFI_ICHECK(call->args[0]->IsInstance()) + TVM_FFI_ICHECK(call->args[0]->IsInstance()) << "R.call_tir expects the first argument to be a GlobalVar referring to a TIR " "tirx::Function. " << "However, the argument " << call->args[0] << " instead has type " @@ -804,17 +805,17 @@ Expr NormalizeCallTIR(const BlockBuilder& ctx, Call call) { << "but " << call << " has " << call->args.size() << " arguments."; auto callee = call->args[0]; - TVM_FFI_ICHECK(callee->ty.as()) + TVM_FFI_ICHECK(callee->ty.as()) << "Operation " << call->op << " expects the first argument to be a TIR callee. " << "However, the first argument " << callee << " has type " << callee->ty; Expr arg_tuple = call->args[1]; - TVM_FFI_ICHECK(arg_tuple->ty.as()) + TVM_FFI_ICHECK(arg_tuple->ty.as()) << "Operation " << call->op << " expects the second argument to be a tuple of relax Expr. " << "However, the second argument " << arg_tuple << " has type " << arg_tuple->ty << "."; - TVM_FFI_ICHECK(arg_tuple.as() || arg_tuple.as()) + TVM_FFI_ICHECK(arg_tuple.as() || arg_tuple.as()) << "Operation " << call->op << " must hold its arguments as an in-line tuple. " << "However, " << call << " has arguments " << arg_tuple << ", which is neither an in-line tuple, " @@ -885,7 +886,7 @@ Expr NormalizeCallTIR(const BlockBuilder& ctx, Call call) { return call; } -void ValidateCallTIR(const CallNode* call) { +void ValidateCallTIR(const CallObj* call) { // This function is used for validation of `relax.call_tir`, // along with the variants `relax.call_tir_with_grad` and // `relax.call_tir_inplace`. Therefore, all error messages should @@ -918,8 +919,8 @@ void ValidateCallTIR(const CallNode* call) { TVM_FFI_STATIC_INIT_BLOCK() { OpDef("relax.call_tir") - .set_validator(ffi::reflection::NativeFunctionView::FromNative<&ValidateCallTIR>()) + .set_validator( + ffi::reflection::NativeFunctionView::FromNative<&ValidateCallTIR>()) .signature(sig::arg("func", "The destination-passing-style function."), sig::arg("args", "The input arguments."), sig::ty_arg("out_type", "The output type.")) @@ -930,7 +931,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { Expr MakeCallTIR(Expr func, Tuple args, ffi::Array out_ty_list) { for (const TensorType& ty : out_ty_list) { - const auto* shape = ty->shape.as(); + const auto* shape = ty->shape.as(); TVM_FFI_ICHECK(shape != nullptr) << "out_ty of call_tir should have defined ShapeExpr as shape. " "However, one given type information is " @@ -955,8 +956,8 @@ TVM_FFI_STATIC_INIT_BLOCK() { // call_tir_with_grad OpDef("relax.call_tir_with_grad") - .set_validator(ffi::reflection::NativeFunctionView::FromNative<&ValidateCallTIR>()) + .set_validator( + ffi::reflection::NativeFunctionView::FromNative<&ValidateCallTIR>()) .signature(sig::arg("func", "The destination-passing-style function."), sig::arg("args", "The input arguments."), sig::ty_arg("out_type", "The output type."), @@ -969,7 +970,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { Expr MakeCallTIRWithGrad(Expr func, Tuple args, ffi::Array out_ty_list, ffi::String te_grad_name, ffi::Map te_grad_kwargs) { for (const TensorType& ty : out_ty_list) { - const auto* shape = ty->shape.as(); + const auto* shape = ty->shape.as(); TVM_FFI_ICHECK(shape != nullptr) << "out_ty of call_tir_with_grad should have defined ShapeExpr as shape. " "However, one given type information is " @@ -1006,7 +1007,7 @@ Expr NormalizeCallTIRInPlace(const BlockBuilder& ctx, Call call) { ffi::Array ty_outputs = [&]() -> ffi::Array { auto out_ty = call->ty_args[0]; - if (auto* tuple_output = out_ty.as()) { + if (auto* tuple_output = out_ty.as()) { return tuple_output->fields; } else { return {out_ty}; @@ -1056,7 +1057,7 @@ Expr NormalizeCallTIRInPlace(const BlockBuilder& ctx, Call call) { } auto ty_output = ty_outputs[i_output]; - auto tinfo_output = ty_output.as(); + auto tinfo_output = ty_output.as(); if (!tinfo_output || !tinfo_output->shape.has_value() || tinfo_output->IsUnknownDtype()) { TVM_FFI_VISIT_THROW(ValueError, call) @@ -1066,7 +1067,7 @@ Expr NormalizeCallTIRInPlace(const BlockBuilder& ctx, Call call) { } auto ty_input = GetType(call_args->fields[i_input]); - auto tinfo_input = ty_input.as(); + auto tinfo_input = ty_input.as(); if (!tinfo_input || (tinfo_output->IsUnknownDtype() || tinfo_output->dtype != tinfo_input->dtype) || @@ -1087,8 +1088,8 @@ Expr NormalizeCallTIRInPlace(const BlockBuilder& ctx, Call call) { TVM_FFI_STATIC_INIT_BLOCK() { OpDef("relax.call_tir_inplace") - .set_validator(ffi::reflection::NativeFunctionView::FromNative<&ValidateCallTIR>()) + .set_validator( + ffi::reflection::NativeFunctionView::FromNative<&ValidateCallTIR>()) .signature(sig::arg("func", "The destination-passing-style function."), sig::arg("args", "The input arguments."), sig::ty_arg("out_type", "The output type."), @@ -1104,7 +1105,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { Expr MakeCallTIRInplace(Expr func, Tuple args, ffi::Array inplace_indices, ffi::Array out_ty_list) { for (const TensorType& ty : out_ty_list) { - const auto* shape = ty->shape.as(); + const auto* shape = ty->shape.as(); TVM_FFI_ICHECK(shape != nullptr) << "out_ty of call_tir should have defined ShapeExpr as shape. " "However, one given type information is " @@ -1132,7 +1133,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { // call_dps_packed -Type InferTypeCallDPSPacked(const CallNode* call_node) { +Type InferTypeCallDPSPacked(const CallObj* call_node) { const Call call = ffi::GetRef(call_node); if (call->ty_args.size() != 1) { TVM_FFI_VISIT_THROW(InternalError, call) << "ty_args should have exact 1 output type."; @@ -1153,7 +1154,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { Expr MakeCallDPSPacked(Expr func, Tuple args, ffi::Array out_ty_list) { for (const TensorType& ty : out_ty_list) { - const auto* shape = ty->shape.as(); + const auto* shape = ty->shape.as(); TVM_FFI_ICHECK(shape != nullptr) << "out_ty of call_dps_packed should have defined ShapeExpr as shape. " "However, one given type information is " @@ -1178,7 +1179,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { // call_py_func -Type InferTypeCallPyFunc(const CallNode* call_node) { +Type InferTypeCallPyFunc(const CallObj* call_node) { const Call call = ffi::GetRef(call_node); if (call->ty_args.size() != 1) { TVM_FFI_VISIT_THROW(InternalError, call) << "ty_args should have exact 1 output type."; @@ -1186,21 +1187,21 @@ Type InferTypeCallPyFunc(const CallNode* call_node) { return call->ty_args[0]; } -void ValidateCallPyFunc(const CallNode* call) { +void ValidateCallPyFunc(const CallObj* call) { // Validate that the function name is a string literal auto func_name = call->args[0]; - TVM_FFI_ICHECK(func_name->IsInstance()) + TVM_FFI_ICHECK(func_name->IsInstance()) << "Operation " << call->op << " expects the first argument to be a string literal " << "specifying the Python function name. However, the first argument " << func_name << " is not a string literal."; // Validate that args is a tuple Expr arg_tuple = call->args[1]; - TVM_FFI_ICHECK(arg_tuple->ty.as()) + TVM_FFI_ICHECK(arg_tuple->ty.as()) << "Operation " << call->op << " expects the second argument to be a tuple of relax Expr. " << "However, the second argument " << arg_tuple << " has type " << arg_tuple->ty << "."; - TVM_FFI_ICHECK(arg_tuple.as() || arg_tuple.as()) + TVM_FFI_ICHECK(arg_tuple.as() || arg_tuple.as()) << "Operation " << call->op << " must hold its arguments as an in-line tuple. " << "However, the argument tuple is " << arg_tuple << ", which is neither an in-line tuple, " << "nor a variable binding that may be normalized to an in-line tuple."; @@ -1209,7 +1210,7 @@ void ValidateCallPyFunc(const CallNode* call) { TVM_FFI_STATIC_INIT_BLOCK() { OpDef("relax.call_py_func") .set_validator(ffi::reflection::NativeFunctionView::FromNative<&ValidateCallPyFunc>()) + const CallObj*)>::FromNative<&ValidateCallPyFunc>()) .signature(sig::arg("func_name", "The name of the Python function to call."), sig::arg("args", "The input arguments."), sig::ty_arg("out_type", "The output type.")) @@ -1219,7 +1220,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { Expr MakeCallPyFunc(StringImm func_name, Tuple args, ffi::Array out_ty_list) { for (const TensorType& ty : out_ty_list) { - const auto* shape = ty->shape.as(); + const auto* shape = ty->shape.as(); TVM_FFI_ICHECK(shape != nullptr) << "out_ty of call_py_func should have defined ShapeExpr as shape. " "However, one given type information is " @@ -1243,7 +1244,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { } // call builtin -Type InferTypeCallBuiltinWithCtx(const CallNode* call_node) { +Type InferTypeCallBuiltinWithCtx(const CallObj* call_node) { const Call call = ffi::GetRef(call_node); if (call->ty_args.size() == 0) { // by default return void. @@ -1319,14 +1320,14 @@ TVM_FFI_STATIC_INIT_BLOCK() { // can't actually name it assert or else Python will consider it a syntax error -Type InferAssertType(const CallNode* call_node) { +Type InferAssertType(const CallObj* call_node) { const Call call = ffi::GetRef(call_node); // Ensure that the condition argument is a boolean scalar. // Also permitted is a tensor with unknown shape and unknown dtype // (checked dynamically in that case). Returns void. CheckNumArguments(call); Type arg_ty = call->args[0]->ty; - if (arg_ty.as()) return Type::Missing(); + if (arg_ty.as()) return Type::Missing(); if (!IsBoolType(arg_ty)) { TVM_FFI_VISIT_THROW(TypeError, call) << "The argument to assert must be a boolean scalar, but received " << arg_ty; @@ -1365,7 +1366,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { // invoke_closure -Type InferTypeInvokeClosure(const CallNode* call_node) { +Type InferTypeInvokeClosure(const CallObj* call_node) { const Call call = ffi::GetRef(call_node); if (call->ty_args.empty()) { return AnyType(); @@ -1434,10 +1435,10 @@ TVM_FFI_STATIC_INIT_BLOCK() { // size -Type InferTypeSize(const CallNode* call_node) { +Type InferTypeSize(const CallObj* call_node) { const Call call = ffi::GetRef(call_node); auto arg_ty = GetType(call->args[0]); - auto* tensor_ty = GetType(call->args[0]).as(); + auto* tensor_ty = GetType(call->args[0]).as(); TVM_FFI_ICHECK(tensor_ty) << "size expects a tensor input, but received " << arg_ty << "; use MatchCast if necessary"; return TensorType(ShapeExpr(ffi::Array{}), PrimType::Int(64)); @@ -1462,11 +1463,11 @@ TVM_FFI_STATIC_INIT_BLOCK() { // tensor_to_shape -Type ReturnTensorToShapeType(const CallNode* call_node) { +Type ReturnTensorToShapeType(const CallObj* call_node) { const Call call = ffi::GetRef(call_node); TVM_FFI_ICHECK(call->args.size() == 1); TVM_FFI_ICHECK(!call->args[0]->ty.as().has_value()); - const auto* tensor_ty = GetTypeAs(call->args[0]); + const auto* tensor_ty = GetTypeAs(call->args[0]); TVM_FFI_ICHECK(tensor_ty); TVM_FFI_ICHECK_EQ(tensor_ty->ndim, 1) << "relax.tensor_to_shape expected argument to be 1-d, " @@ -1474,7 +1475,7 @@ Type ReturnTensorToShapeType(const CallNode* call_node) { if (tensor_ty->shape.has_value()) { ShapeExpr shape_expr = tensor_ty->shape.value().as_or_throw(); - const IntImmNode* ndim = shape_expr->values[0].as(); + const IntImmObj* ndim = shape_expr->values[0].as(); if (ndim) { return ShapeType(ndim->value.as().value()); } @@ -1500,11 +1501,11 @@ TVM_FFI_STATIC_INIT_BLOCK() { } // shape_to_tensor -Type ReturnShapeToTensorType(const CallNode* call_node) { +Type ReturnShapeToTensorType(const CallObj* call_node) { const Call call = ffi::GetRef(call_node); TVM_FFI_ICHECK(call->args.size() == 1); TVM_FFI_ICHECK(!call->args[0]->ty.as().has_value()); - const auto* ty = GetTypeAs(call->args[0]); + const auto* ty = GetTypeAs(call->args[0]); TVM_FFI_ICHECK(ty); int32_t ndim = ty->ndim; return TensorType(ShapeExpr({PrimExpr(ndim)}), PrimType::Int(64)); @@ -1531,17 +1532,17 @@ TVM_FFI_STATIC_INIT_BLOCK() { // alloc_tensor Type InferTypeAllocateTensor(const Call& call, const BlockBuilder& ctx) { - TVM_FFI_ICHECK(call->args[0].as()) + TVM_FFI_ICHECK(call->args[0].as()) << "must be ShapeExpr, but got " << call->args[0]->GetTypeKey(); - TVM_FFI_ICHECK(call->args[1].as()) + TVM_FFI_ICHECK(call->args[1].as()) << "must be DataTypeImm, but got " << call->args[1]->GetTypeKey(); PrimType out_dtype = PrimType::Void(); - if (const auto* dtype_node = call->args[1].as()) { + if (const auto* dtype_node = call->args[1].as()) { const DataTypeImm dtype_imm = ffi::GetRef(dtype_node); out_dtype = PrimType(dtype_imm->value); } int64_t vdevice_index = -1; - if (const auto* int_imm = call->args[2].as()) { + if (const auto* int_imm = call->args[2].as()) { vdevice_index = int_imm->value.as().value(); } auto vdevice = GetGlobalVDevice(ctx->GetContextIRModule(), vdevice_index); @@ -1613,17 +1614,17 @@ TVM_FFI_STATIC_INIT_BLOCK() { // memory planning alloc_tensor Type InferTypeMemAllocTensor(const Call& call, const BlockBuilder& ctx) { - TVM_FFI_ICHECK(GetTypeAs(call->args[2])) + TVM_FFI_ICHECK(GetTypeAs(call->args[2])) << "must be a Expr of ShapeType, but got " << call->args[1]->GetTypeKey(); PrimType out_dtype = PrimType::Void(); - if (const auto* dtype_node = call->args[3].as()) { + if (const auto* dtype_node = call->args[3].as()) { const DataTypeImm dtype_imm = ffi::GetRef(dtype_node); out_dtype = PrimType(dtype_imm->value); } if (call->args.size() == 5) { int64_t vdevice_index = -1; - if (const auto* int_imm = call->args[4].as()) { + if (const auto* int_imm = call->args[4].as()) { vdevice_index = int_imm->value.as().value(); } auto vdevice = GetGlobalVDevice(ctx->GetContextIRModule(), vdevice_index); @@ -1741,19 +1742,19 @@ TVM_FFI_STATIC_INIT_BLOCK() { Type InferTypeVMAllocTensor(const Call& call, const BlockBuilder& ctx) { PrimType out_dtype = PrimType::Void(); - if (const auto* dtype_node = call->args[3].as()) { + if (const auto* dtype_node = call->args[3].as()) { const DataTypeImm dtype_imm = ffi::GetRef(dtype_node); out_dtype = PrimType(dtype_imm->value); } int64_t vdevice_index = -1; - if (const auto* int_imm = call->args[4].as()) { + if (const auto* int_imm = call->args[4].as()) { vdevice_index = int_imm->value.as().value(); } auto vdevice = GetGlobalVDevice(ctx->GetContextIRModule(), vdevice_index); - if (const auto* output_shape = call->args[2].as()) { + if (const auto* output_shape = call->args[2].as()) { return TensorType(ffi::GetRef(output_shape), out_dtype, vdevice); - } else if (const auto* shape_ty = GetTypeAs(call->args[2])) { + } else if (const auto* shape_ty = GetTypeAs(call->args[2])) { if (shape_ty->values.has_value()) { return TensorType(ShapeExpr(shape_ty->values.value()), out_dtype, vdevice); } else { @@ -1837,7 +1838,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { } // builtin stop_lift_params -Type InferTypeStopLiftParams(const CallNode* call) { +Type InferTypeStopLiftParams(const CallObj* call) { return InferTypeUnaryArithContextFree(call); } @@ -1861,7 +1862,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { // to_vdevice -Type InferToVDeviceType(const CallNode* call_node) { +Type InferToVDeviceType(const CallObj* call_node) { const Call call = ffi::GetRef(call_node); TVM_FFI_ICHECK(call->args.size() == 1); TVM_FFI_ICHECK(!call->args[0]->ty.as().has_value()); @@ -1895,7 +1896,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { // hint_on_device -Type InferHintOnDeviceType(const CallNode* call_node) { +Type InferHintOnDeviceType(const CallObj* call_node) { const Call call = ffi::GetRef(call_node); TVM_FFI_ICHECK(call->args.size() == 1); TVM_FFI_ICHECK(!call->args[0]->ty.as().has_value()); diff --git a/src/relax/op/op_common.cc b/src/relax/op/op_common.cc index accce4fe925b..18757be6ea0d 100644 --- a/src/relax/op/op_common.cc +++ b/src/relax/op/op_common.cc @@ -89,7 +89,7 @@ ffi::Array GetInputTensorType(const Call& call) { } ffi::Array GetTensorTypeFromTuple(const Call& call, const Expr& tup) { - const auto* tuple_ty = GetTypeAs(tup); + const auto* tuple_ty = GetTypeAs(tup); if (tuple_ty == nullptr) { TVM_FFI_VISIT_THROW(TypeError, call) << call->op << " expects the input to be a Tuple of Tensors. However, the given input is " @@ -99,7 +99,7 @@ ffi::Array GetTensorTypeFromTuple(const Call& call, const Expr& tup) ffi::Array tensor_ty; tensor_ty.reserve(tuple_ty->fields.size()); for (Type field_ty : tuple_ty->fields) { - const auto* field_tensor_ty = field_ty.as(); + const auto* field_tensor_ty = field_ty.as(); if (field_tensor_ty == nullptr) { TVM_FFI_VISIT_THROW(TypeError, call) << call->op << " expects the input to be a Tuple of Tensors. However, the given input is " @@ -125,8 +125,8 @@ BinaryBroadcastShapeInferResult InferBinaryBroadcastShape(sym::AnalyzerObj* anal for (; i <= std::min(x1_ndim, x2_ndim); ++i) { const PrimExpr& dim0 = x1_shape[x1_ndim - i]; const PrimExpr& dim1 = x2_shape[x2_ndim - i]; - const auto* int_dim0 = dim0.as(); - const auto* int_dim1 = dim1.as(); + const auto* int_dim0 = dim0.as(); + const auto* int_dim1 = dim1.as(); if (int_dim0 != nullptr && int_dim0->value == 1) { output_shape.push_back(dim1); } else if (int_dim1 != nullptr && int_dim1->value == 1) { diff --git a/src/relax/op/op_common.h b/src/relax/op/op_common.h index ad3e68091d47..52977bafa362 100644 --- a/src/relax/op/op_common.h +++ b/src/relax/op/op_common.h @@ -100,7 +100,7 @@ inline TensorType GetUnaryInputTensorType(const Call& call) { return GetInputTen inline bool RequiresTensorInputNormalization(const Call& call) { CheckNumArguments(call); for (const Expr& arg : call->args) { - if (arg->ty.as() || arg->ty.as()) return true; + if (arg->ty.as() || arg->ty.as()) return true; } return false; } @@ -214,10 +214,10 @@ inline Type InferTypeUnary(const Call& call, FType f_compute_out_dtype) { computed_dtype.value()->dtype == input_ty->dtype.value()->dtype); TensorType inferred = input_ty; if (!call->ty_args.empty() || !same_dtype) { - auto output_ty = ffi::make_object(*input_ty.get()); + auto output_ty = ffi::make_object(*input_ty.get()); output_ty->dtype = computed_dtype; if (call->ty_args.size() > 0) { - auto defined_ty = call->ty_args[0].as(); + auto defined_ty = call->ty_args[0].as(); TVM_FFI_ICHECK(defined_ty); auto shape = output_ty->GetShape(); TVM_FFI_ICHECK(shape.has_value()); @@ -228,7 +228,7 @@ inline Type InferTypeUnary(const Call& call, FType f_compute_out_dtype) { inferred = TensorType(output_ty); } } - if (const auto* old = call->ty.as()) { + if (const auto* old = call->ty.as()) { auto same_optional_ref = [](const auto& lhs, const auto& rhs) { return lhs.has_value() == rhs.has_value() && (!lhs.has_value() || lhs.value().same_as(rhs.value())); @@ -278,7 +278,7 @@ Type ReturnTypeFromArg(const Call& call, const BlockBuilder& ctx) { } template -Type ReturnTypeFromArgContextFree(const CallNode* call_node) { +Type ReturnTypeFromArgContextFree(const CallObj* call_node) { Call call = ffi::GetRef(call_node); Op op = call->op.as_or_throw(); CheckNumArguments(call); @@ -306,7 +306,7 @@ Type InferTypeUnaryArith(const Call& call, const BlockBuilder& ctx) { } template -Type InferTypeUnaryArithContextFree(const CallNode* call_node) { +Type InferTypeUnaryArithContextFree(const CallObj* call_node) { Call call = ffi::GetRef(call_node); return InferTypeUnary( call, [](const TensorType& input_ty) { return input_ty->dtype; }); @@ -331,9 +331,9 @@ InferLayoutOutput InferLayoutUnaryEwise( * \throw Throw exception if the Type doesn't have an element type. */ inline std::optional GetElementDType(const Type& ty) { - if (const auto* prim = ty.as()) { + if (const auto* prim = ty.as()) { return ffi::GetRef(prim); - } else if (const auto* tensor = ty.as()) { + } else if (const auto* tensor = ty.as()) { if (tensor->dtype.has_value()) { return tensor->dtype.value(); } else { @@ -362,7 +362,7 @@ inline ffi::Optional InferBinaryArithOpOutDtype(const Call& call, cons // and reenter the same inference hook before the exception is thrown. auto opt_lhs_dtype = GetElementDType(lhs_ty); if (!opt_lhs_dtype) { - if (const auto* lhs_tensor = lhs_ty.as()) { + if (const auto* lhs_tensor = lhs_ty.as()) { if (lhs_tensor->IsUnknownDtype()) return std::nullopt; } TVM_FFI_VISIT_THROW(TypeError, call) @@ -374,7 +374,7 @@ inline ffi::Optional InferBinaryArithOpOutDtype(const Call& call, cons auto opt_rhs_dtype = GetElementDType(rhs_ty); if (!opt_rhs_dtype) { - if (const auto* rhs_tensor = rhs_ty.as()) { + if (const auto* rhs_tensor = rhs_ty.as()) { if (rhs_tensor->IsUnknownDtype()) return std::nullopt; } TVM_FFI_VISIT_THROW(TypeError, call) @@ -411,7 +411,7 @@ inline ffi::Optional InferBinaryArithOpOutDtype(const Call& call, cons inline ffi::Optional InferBinaryArithOpOutVDevice(const Call& call, const Type& lhs_ty, const Type& rhs_ty) { auto get_vdevice = [&](const Type& ty) -> ffi::Optional { - if (const auto* tensor = ty.as()) { + if (const auto* tensor = ty.as()) { return tensor->vdevice; } else { return std::nullopt; @@ -676,7 +676,7 @@ inline ffi::Optional CheckNdimPerLayoutAndGetShape(const Call& call, << "In " << call->op << ", layout " << layout << " requires the input to be " << layout.ndim() << "-dim tensor. However, the given input has ndim " << ty->ndim; } - if (const auto* shape_expr = ty->shape.as()) { + if (const auto* shape_expr = ty->shape.as()) { return ffi::GetRef(shape_expr); } return std::nullopt; diff --git a/src/relax/op/tensor/binary.cc b/src/relax/op/tensor/binary.cc index f071ab74aaff..06a4b01ef16a 100644 --- a/src/relax/op/tensor/binary.cc +++ b/src/relax/op/tensor/binary.cc @@ -40,17 +40,17 @@ Type InferTypeBroadcast(const Call& call, const BlockBuilder& ctx, FType f_compu auto lhs_ty = GetType(call->args[0]); auto rhs_ty = GetType(call->args[1]); - TVM_FFI_CHECK(lhs_ty.as() || lhs_ty.as(), TypeError) + TVM_FFI_CHECK(lhs_ty.as() || lhs_ty.as(), TypeError) << "Arguments to binary operators must be either R.Tensor or R.Prim types, " << "but " << call->op << " has LHS type " << lhs_ty; - TVM_FFI_CHECK(rhs_ty.as() || rhs_ty.as(), TypeError) + TVM_FFI_CHECK(rhs_ty.as() || rhs_ty.as(), TypeError) << "Arguments to binary operators must be either R.Tensor or R.Prim types, " << "but " << call->op << " has RHS type " << rhs_ty; // Dtype ffi::Optional output_dtype = f_compute_out_dtype(call, ctx, lhs_ty, rhs_ty); - if (lhs_ty.as() && rhs_ty.as()) { + if (lhs_ty.as() && rhs_ty.as()) { TVM_FFI_ICHECK(output_dtype.has_value()); return output_dtype.value(); } @@ -59,9 +59,9 @@ Type InferTypeBroadcast(const Call& call, const BlockBuilder& ctx, FType f_compu ffi::Optional vdevice = InferBinaryArithOpOutVDevice(call, ctx, lhs_ty, rhs_ty); auto get_ndim = [&](const Type& ty) -> int { - if (ty.as()) { + if (ty.as()) { return 1; - } else if (const auto* tensor = ty.as()) { + } else if (const auto* tensor = ty.as()) { return tensor->ndim; } else { return kUnknownNDim; @@ -82,9 +82,9 @@ Type InferTypeBroadcast(const Call& call, const BlockBuilder& ctx, FType f_compu // Shapes auto get_shape = [](const Type& ty) -> ffi::Optional> { - if (ty.as()) { + if (ty.as()) { return ffi::Array{IntImm::Int64(1)}; - } else if (const auto* tensor = ty.as()) { + } else if (const auto* tensor = ty.as()) { return tensor->GetShape(); } else { return std::nullopt; @@ -105,7 +105,7 @@ Type InferTypeBroadcast(const Call& call, const BlockBuilder& ctx, FType f_compu } auto get_shape_expr = [](const Type& ty) -> ffi::Optional { - if (const auto* tensor = ty.as()) { + if (const auto* tensor = ty.as()) { return tensor->shape; } else { return std::nullopt; @@ -146,8 +146,8 @@ InferLayoutOutput InferLayoutBinaryEwise( LayoutDecision layout1 = GetLayoutDecision(var_layout_map, call->args[0]); LayoutDecision layout2 = GetLayoutDecision(var_layout_map, call->args[1]); - auto* x1_ty = GetTypeAs(call->args[0]); - auto* x2_ty = GetTypeAs(call->args[1]); + auto* x1_ty = GetTypeAs(call->args[0]); + auto* x2_ty = GetTypeAs(call->args[1]); TVM_FFI_ICHECK(!x1_ty->IsUnknownNdim() && !x2_ty->IsUnknownNdim()) << "Unknown dim tensors should not be handled by this function"; diff --git a/src/relax/op/tensor/create.cc b/src/relax/op/tensor/create.cc index a278839d113c..4ad38135af5d 100644 --- a/src/relax/op/tensor/create.cc +++ b/src/relax/op/tensor/create.cc @@ -65,7 +65,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { Expr full(ffi::Variant, Expr> shape, Expr fill_value, ffi::Optional dtype) { Expr shape_in_expr = [&]() -> Expr { - if (const auto* expr = shape.as()) { + if (const auto* expr = shape.as()) { return ffi::GetRef(expr); } return ShapeExpr(shape.get>()); @@ -84,13 +84,13 @@ TVM_FFI_STATIC_INIT_BLOCK() { refl::GlobalDef().def("relax.op.full", full); } -Type InferTypeFull(const CallNode* call_node) { +Type InferTypeFull(const CallObj* call_node) { const Call call = ffi::GetRef(call_node); if (call->args.size() != 2) { TVM_FFI_VISIT_THROW(ValueError, call) << "Full op should have 2 arguments"; } - const auto* shape_ty = GetTypeAs(call->args[0]); - const auto* fill_value_ty = GetTypeAs(call->args[1]); + const auto* shape_ty = GetTypeAs(call->args[0]); + const auto* fill_value_ty = GetTypeAs(call->args[1]); if (shape_ty == nullptr) { TVM_FFI_VISIT_THROW(TypeError, call) << "Full requires the input shape to be a Shape. However, the given one is " @@ -134,7 +134,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { refl::GlobalDef().def("relax.op.full_like", full_like); } -Type InferTypeFullLike(const CallNode* call_node) { +Type InferTypeFullLike(const CallObj* call_node) { const Call call = ffi::GetRef(call_node); ffi::Array input_ty = GetInputTensorType(call); TensorType data_ty = input_ty[0]; @@ -149,7 +149,7 @@ Type InferTypeFullLike(const CallNode* call_node) { if (!attrs->dtype.has_value()) { return data_ty; } else { - auto output_ty = ffi::make_object(*data_ty.get()); + auto output_ty = ffi::make_object(*data_ty.get()); output_ty->dtype = PrimType(attrs->dtype.value()); return TensorType(output_ty); } @@ -165,13 +165,13 @@ TVM_FFI_STATIC_INIT_BLOCK() { } // Structure info inference for ones and zeros -Type InferTypeOnesZeros(const CallNode* call_node) { +Type InferTypeOnesZeros(const CallObj* call_node) { const Call call = ffi::GetRef(call_node); if (call->args.size() != 1) { TVM_FFI_VISIT_THROW(ValueError, call) << "Ones/Zeros should have 1 argument"; } - const auto* shape_ty = GetTypeAs(call->args[0]); + const auto* shape_ty = GetTypeAs(call->args[0]); if (shape_ty == nullptr) { TVM_FFI_VISIT_THROW(TypeError, call) << "Ones/Zeros requires the input shape to be a Shape. However, the given one is " @@ -183,14 +183,14 @@ Type InferTypeOnesZeros(const CallNode* call_node) { } // Structure info inference for ones_like and zeros_like -Type InferTypeOnesLikeZerosLike(const CallNode* call_node) { +Type InferTypeOnesLikeZerosLike(const CallObj* call_node) { const Call call = ffi::GetRef(call_node); TensorType data_ty = GetUnaryInputTensorType(call); const auto* attrs = call->attrs.as(); if (!attrs->dtype.has_value()) { return data_ty; } else { - auto output_ty = ffi::make_object(*data_ty.get()); + auto output_ty = ffi::make_object(*data_ty.get()); output_ty->dtype = PrimType(attrs->dtype.value()); return TensorType(output_ty); } @@ -282,7 +282,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { refl::GlobalDef().def("relax.op.eye", eye).def("relax.op.eye_like", eye_like); } -Type InferTypeEye(const CallNode* call_node) { +Type InferTypeEye(const CallObj* call_node) { const Call call = ffi::GetRef(call_node); if (call->args.size() != 3) { TVM_FFI_VISIT_THROW(ValueError, call) << "Eye op should have 3 arguments: n, m, and k, but got " @@ -307,7 +307,7 @@ Type InferTypeEye(const CallNode* call_node) { return TensorType(ShapeExpr({n, m}), PrimType(dtype)); } -Type InferTypeEyeLike(const CallNode* call_node) { +Type InferTypeEyeLike(const CallObj* call_node) { const Call call = ffi::GetRef(call_node); if (call->args.size() != 2) { TVM_FFI_VISIT_THROW(ValueError, call) @@ -315,7 +315,7 @@ Type InferTypeEyeLike(const CallNode* call_node) { << " arguments"; } - const auto* x_ty = GetTypeAs(call->args[0]); + const auto* x_ty = GetTypeAs(call->args[0]); if (x_ty == nullptr) { TVM_FFI_VISIT_THROW(TypeError, call) << "Eye_like expects the input `x` to be a Tensor, but got " @@ -365,7 +365,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { refl::GlobalDef().def("relax.op.arange", arange); } -Type InferTypeArange(const CallNode* call_node) { +Type InferTypeArange(const CallObj* call_node) { const Call call = ffi::GetRef(call_node); if (call->args.size() != 3) { TVM_FFI_VISIT_THROW(ValueError, call) @@ -429,7 +429,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { refl::GlobalDef().def("relax.op.hamming_window", hamming_window); } -Type InferTypeHammingWindow(const CallNode* call_node) { +Type InferTypeHammingWindow(const CallObj* call_node) { const Call call = ffi::GetRef(call_node); const auto* attrs = call->attrs.as(); TVM_FFI_ICHECK(attrs->dtype.has_value()); @@ -491,7 +491,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { .def("relax.op.triu", static_cast(triu)); } -Type InferTypeTrilTriu(const CallNode* call_node) { +Type InferTypeTrilTriu(const CallObj* call_node) { const Call call = ffi::GetRef(call_node); auto [data_ty, offset] = GetArgType(call); diff --git a/src/relax/op/tensor/datatype.cc b/src/relax/op/tensor/datatype.cc index c33efce461b6..4c14368cad7c 100644 --- a/src/relax/op/tensor/datatype.cc +++ b/src/relax/op/tensor/datatype.cc @@ -64,12 +64,12 @@ TVM_FFI_STATIC_INIT_BLOCK() { refl::GlobalDef().def("relax.op.astype", astype); } -Type InferTypeAstype(const CallNode* call_node) { +Type InferTypeAstype(const CallObj* call_node) { const Call call = ffi::GetRef(call_node); if (RequiresTensorInputNormalization(call)) return Type::Missing(); TensorType ty = GetUnaryInputTensorType(call); const auto* attrs = call->attrs.as(); - ffi::ObjectPtr new_ty = ffi::make_object(*ty.get()); + ffi::ObjectPtr new_ty = ffi::make_object(*ty.get()); new_ty->dtype = PrimType(attrs->dtype); return TensorType(new_ty); } @@ -98,11 +98,11 @@ TVM_FFI_STATIC_INIT_BLOCK() { refl::GlobalDef().def("relax.op.wrap_param", MakeWrapParam); } -Type InferTypeWrapParam(const CallNode* call_node) { +Type InferTypeWrapParam(const CallObj* call_node) { const Call call = ffi::GetRef(call_node); TensorType ty = GetUnaryInputTensorType(call); const auto* attrs = call->attrs.as(); - ffi::ObjectPtr new_ty = ffi::make_object(*ty.get()); + ffi::ObjectPtr new_ty = ffi::make_object(*ty.get()); new_ty->dtype = PrimType(attrs->dtype); return TensorType(new_ty); } diff --git a/src/relax/op/tensor/grad.cc b/src/relax/op/tensor/grad.cc index 1c7c4c5de4cf..84271f6c7df4 100644 --- a/src/relax/op/tensor/grad.cc +++ b/src/relax/op/tensor/grad.cc @@ -52,7 +52,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { refl::GlobalDef().def("relax.op.grad.no_grad", no_grad); } -Type InferTypeNoGrad(const CallNode* call_node) { +Type InferTypeNoGrad(const CallObj* call_node) { const Call call = ffi::GetRef(call_node); return GetType(call->args[0]); } @@ -76,7 +76,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { } Type InferTypeStartCheckpoint(const Call& call, const BlockBuilder& ctx) { - if (!call->args[0].as()) { + if (!call->args[0].as()) { TVM_FFI_VISIT_THROW(TypeError, call) << "The argument of relax.op.grad.start_checkpoint should be a Var."; } @@ -102,7 +102,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { } Type InferTypeEndCheckpoint(const Call& call, const BlockBuilder& ctx) { - if (!call->args[0].as()) { + if (!call->args[0].as()) { TVM_FFI_VISIT_THROW(TypeError, call) << "The argument of relax.op.grad.end_checkpoint should be a Var."; } @@ -142,7 +142,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { refl::GlobalDef().def("relax.op.grad.nll_loss_backward", nll_loss_backward); } -Type InferTypeNLLLossBackward(const CallNode* call_node) { +Type InferTypeNLLLossBackward(const CallObj* call_node) { const Call call = ffi::GetRef(call_node); return GetType(call->args[1]); } @@ -180,7 +180,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { refl::GlobalDef().def("relax.op.grad.max_pool2d_backward", max_pool2d_backward); } -Type InferTypeMaxPool2DBackward(const CallNode* call_node) { +Type InferTypeMaxPool2DBackward(const CallObj* call_node) { const Call call = ffi::GetRef(call_node); return GetType(call->args[1]); } @@ -216,7 +216,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { refl::GlobalDef().def("relax.op.grad.avg_pool2d_backward", avg_pool2d_backward); } -Type InferTypeAvgPool2DBackward(const CallNode* call_node) { +Type InferTypeAvgPool2DBackward(const CallObj* call_node) { const Call call = ffi::GetRef(call_node); return GetType(call->args[1]); } @@ -246,7 +246,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { refl::GlobalDef().def("relax.op.grad.take_backward", take_backward); } -Type InferTypeTakeBackward(const CallNode* call_node) { +Type InferTypeTakeBackward(const CallObj* call_node) { const Call call = ffi::GetRef(call_node); return GetType(call->args[1]); } diff --git a/src/relax/op/tensor/index.cc b/src/relax/op/tensor/index.cc index 8a58194d3f46..479e311eb1b6 100644 --- a/src/relax/op/tensor/index.cc +++ b/src/relax/op/tensor/index.cc @@ -81,7 +81,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { refl::GlobalDef().def("relax.op.take", take); } -Type InferTypeTake(const CallNode* call_node) { +Type InferTypeTake(const CallObj* call_node) { const Call call = ffi::GetRef(call_node); if (RequiresTensorInputNormalization(call)) return Type::Missing(); CheckNumArguments(call); @@ -94,7 +94,7 @@ Type InferTypeTake(const CallNode* call_node) { auto ty = GetType(arg); if (auto tensor_ty = ty.as()) { return tensor_ty.value(); - } else if (auto prim_ty = ty.as()) { + } else if (auto prim_ty = ty.as()) { return TensorType(ShapeExpr(ffi::Array{}), ffi::GetRef(prim_ty)); } else { TVM_FFI_VISIT_THROW(TypeError, call) @@ -132,8 +132,8 @@ Type InferTypeTake(const CallNode* call_node) { if (attrs->axis.has_value()) { axis = NormalizeAxis(call, data_ty->ndim, attrs->axis.value()); } - const auto* data_shape = data_ty->shape.as(); - const auto* indices_shape = indices_ty->shape.as(); + const auto* data_shape = data_ty->shape.as(); + const auto* indices_shape = indices_ty->shape.as(); if (data_shape == nullptr || indices_shape == nullptr) { return TensorType(data_ty->dtype, indices_ty->ndim + data_ty->ndim - 1, data_ty->vdevice); } @@ -168,7 +168,7 @@ Expr strided_slice(Expr x, Expr axes, Expr begin, Expr end, ffi::Optional std::optional> known_length; auto check_tuple = [&known_length](const char* name, Expr expr) { - if (const auto* tuple = expr.as()) { + if (const auto* tuple = expr.as()) { size_t length = tuple->fields.size(); if (known_length.has_value()) { const auto& prev = known_length.value(); @@ -235,9 +235,9 @@ ffi::Optional> UnpackTupleOfPrimExpr(ffi::Optional ty // An AnyType may contain a tuple of the desired type, but // it isn't yet known whether it does. Return early, as we cannot // provide a known `ffi::Array` to the caller. - if (ty.as()) return std::nullopt; + if (ty.as()) return std::nullopt; - auto tuple = ty.as(); + auto tuple = ty.as(); TVM_FFI_CHECK(tuple, TypeError) << "The type " << ty << " cannot contain a tuple whose elements are " << PrimType::ContainerType::_type_key; @@ -246,9 +246,9 @@ ffi::Optional> UnpackTupleOfPrimExpr(ffi::Optional ty for (size_t i = 0; i < tuple->fields.size(); i++) { auto field = tuple->fields[i]; - if (field.as()) return std::nullopt; + if (field.as()) return std::nullopt; - auto prim_ty = field.as(); + auto prim_ty = field.as(); TVM_FFI_CHECK(prim_ty, TypeError) << "The type " << ty << " cannot contain a tuple whose elements are " << PrimType::ContainerType::_type_key << ", because element " << i << " has type " << field; @@ -284,7 +284,7 @@ ffi::Optional> UnpackTupleOfPrimExpr(ffi::Optional ex if (!expr) return std::nullopt; const Expr& value = expr.value(); - if (const auto* tuple = value.as()) { + if (const auto* tuple = value.as()) { ffi::Array output; for (size_t i = 0; i < tuple->fields.size(); i++) { const Expr& field = tuple->fields[i]; @@ -346,11 +346,11 @@ Type InferTypeStridedSlice(const Call& call, const BlockBuilder& ctx) { // so will require a way to represent a `relax::TupleType` of // unknown length, where each element has the same `Type`. auto is_base_of_tuple_of_int64 = [&](const Type& ty) -> bool { - if (ty.as()) { + if (ty.as()) { return true; } - const auto* tuple = ty.as(); + const auto* tuple = ty.as(); if (!tuple) return false; return std::all_of(tuple->fields.begin(), tuple->fields.end(), @@ -372,7 +372,7 @@ Type InferTypeStridedSlice(const Call& call, const BlockBuilder& ctx) { check_tuple("strides", call->args[4]); } - const auto* data_ty = data->ty.as(); + const auto* data_ty = data->ty.as(); ffi::Optional dtype = std::nullopt; ffi::Optional vdevice = std::nullopt; @@ -453,12 +453,16 @@ Type InferTypeStridedSlice(const Call& call, const BlockBuilder& ctx) { topi::GetLength(begin, end, strides_tuple[i], input_dim, attrs->assume_inbound); sym::Analyzer analyzer = ctx->GetAnalyzer(); - std::optional> context; + auto saved_context = analyzer->GetConstraintContext(); if (attrs->assume_inbound) { - context.emplace(analyzer, 0 <= begin && begin <= input_dim && 0 <= end && end <= input_dim); + auto context = saved_context->Snapshot(); + context->Append(sym::PredicateConstraint(0 <= begin && begin <= input_dim && 0 <= end && + end <= input_dim)); + analyzer->SyncConstraintContext(context); } output_dim = analyzer->Simplify(output_dim); + analyzer->SyncConstraintContext(saved_context); output_shape.Set(axis, output_dim); } @@ -480,7 +484,7 @@ InferLayoutOutput InferLayoutStridedSlice( const auto* attrs = call->attrs.as(); TVM_FFI_ICHECK(attrs != nullptr) << "Invalid Call"; - const auto* tensor_ty = GetTypeAs(call->args[0]); + const auto* tensor_ty = GetTypeAs(call->args[0]); TVM_FFI_ICHECK(tensor_ty) << "Invalid Call"; TVM_FFI_ICHECK(!tensor_ty->IsUnknownNdim()) << "Layout inference only supports known dimensionality, " @@ -536,13 +540,13 @@ TVM_FFI_STATIC_INIT_BLOCK() { refl::GlobalDef().def("relax.op.dynamic_strided_slice", dynamic_strided_slice); } -Type InferTypeDynStridedSlice(const CallNode* call_node) { +Type InferTypeDynStridedSlice(const CallObj* call_node) { const Call call = ffi::GetRef(call_node); if (RequiresTensorInputNormalization(call)) return Type::Missing(); - const auto* data_ty = GetTypeAs(call->args[0]); - const auto* begin_ty = GetTypeAs(call->args[1]); - const auto* end_ty = GetTypeAs(call->args[2]); - const auto* strides_ty = GetTypeAs(call->args[3]); + const auto* data_ty = GetTypeAs(call->args[0]); + const auto* begin_ty = GetTypeAs(call->args[1]); + const auto* end_ty = GetTypeAs(call->args[2]); + const auto* strides_ty = GetTypeAs(call->args[3]); TVM_FFI_ICHECK(data_ty); if (data_ty->IsUnknownNdim()) { @@ -558,17 +562,17 @@ Type InferTypeDynStridedSlice(const CallNode* call_node) { } int n_axis = data_ty->ndim; - auto diag_def = [&](const TensorTypeNode* ty, ffi::String name) { + auto diag_def = [&](const TensorTypeObj* ty, ffi::String name) { TVM_FFI_ICHECK(ty) << "Dynamic strided slice requires the input " << name << " to be have the type. Please try normalizing the inputs."; TVM_FFI_ICHECK_EQ(ty->ndim, 1) << "Dynamic strided slice requires " << name << " to be 1d tensor (list of values)."; - const auto* shape = ty->shape.as(); + const auto* shape = ty->shape.as(); TVM_FFI_ICHECK(shape) << "Dynamic strided slice requires the input " << name << " to have well-defined shape."; // NOTE(tvm-team): This strong restriction seems necessary for now until we have a generic // solution in converting 1d Tensor with unknown num_elem to ffi::Array. - const auto* num_elem = shape->values[0].as(); + const auto* num_elem = shape->values[0].as(); TVM_FFI_ICHECK(num_elem) << "Dynamic strided slice requires the input " << name << " to have a known integer shape value."; TVM_FFI_ICHECK_EQ(num_elem->value, n_axis) @@ -598,7 +602,7 @@ InferLayoutOutput InferLayoutDynStridedSlice( const VarLayoutMap& var_layout_map) { TVM_FFI_ICHECK(NoDesiredLayout(call, desired_layouts)); - const auto* tensor_ty = GetTypeAs(call->args[0]); + const auto* tensor_ty = GetTypeAs(call->args[0]); TVM_FFI_ICHECK(tensor_ty) << "Invalid Call"; TVM_FFI_ICHECK(!tensor_ty->IsUnknownNdim()) << "Layout inference only supports known dimensionality, " diff --git a/src/relax/op/tensor/inspect.cc b/src/relax/op/tensor/inspect.cc index 5ee3330f7772..91f439c94e70 100644 --- a/src/relax/op/tensor/inspect.cc +++ b/src/relax/op/tensor/inspect.cc @@ -62,19 +62,19 @@ std::tuple> GetTensorArgInfoWithIndex(const C const auto& arg = call->args[0]; const auto& axis = call->args[1]; - auto tensor_ty = arg->ty.as(); + auto tensor_ty = arg->ty.as(); TVM_FFI_CHECK(tensor_ty, TypeError) << "Operator " << call->op << " expects arguments (tensor, axis), " << "but the first argument " << arg << " in expression " << call << " has type " << arg->ty; - auto axis_ty = axis->ty.as(); + auto axis_ty = axis->ty.as(); TVM_FFI_CHECK(axis_ty, TypeError) << "Operator " << call->op << " expects arguments (tensor, axis), " << "but the second argument " << arg << " in expression " << call << " has type " << axis->ty; ffi::Optional int_imm_axis = std::nullopt; if (auto prim_value = axis.as()) { - if (const auto* int_imm = prim_value->as()) { + if (const auto* int_imm = prim_value->as()) { int_imm_axis = static_cast(int_imm->value); } } diff --git a/src/relax/op/tensor/linear_algebra.cc b/src/relax/op/tensor/linear_algebra.cc index 671e870334b0..d62922723efa 100644 --- a/src/relax/op/tensor/linear_algebra.cc +++ b/src/relax/op/tensor/linear_algebra.cc @@ -126,8 +126,8 @@ Type InferTypeMatmul(const Call& call, const BlockBuilder& ctx) { } int output_ndim = std::max(x1_ndim, x2_ndim) - x1_prepended - x2_appended; - const auto* x1_shape = x1_ty->shape.as(); - const auto* x2_shape = x2_ty->shape.as(); + const auto* x1_shape = x1_ty->shape.as(); + const auto* x2_shape = x2_ty->shape.as(); if (x1_shape == nullptr || x2_shape == nullptr) { if (vdev.has_value()) { return TensorType(out_dtype, output_ndim, vdev); @@ -202,7 +202,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { refl::GlobalDef().def("relax.op.einsum", einsum); } -Type InferTypeEinsum(const CallNode* call_node) { +Type InferTypeEinsum(const CallObj* call_node) { const Call call = ffi::GetRef(call_node); if (call->args.size() != 1) { TVM_FFI_VISIT_THROW(ValueError, call) << "Einsum op should take 1 argument"; @@ -249,7 +249,7 @@ Type InferTypeEinsum(const CallNode* call_node) { } // Get input shapes - const auto* shape_expr = tensor_ty->shape.as(); + const auto* shape_expr = tensor_ty->shape.as(); if (shape_expr != nullptr) { input_shapes.push_back(shape_expr->values); } else { @@ -287,7 +287,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { refl::GlobalDef().def("relax.op.outer", outer); } -Type InferTypeOuter(const CallNode* call_node) { +Type InferTypeOuter(const CallObj* call_node) { const Call call = ffi::GetRef(call_node); auto input_ty = GetInputTensorType(call); auto x1_ty = input_ty[0]; @@ -299,8 +299,8 @@ Type InferTypeOuter(const CallNode* call_node) { } // Determine output shape - auto x1_shape = x1_ty->shape.as(); - auto x2_shape = x2_ty->shape.as(); + auto x1_shape = x1_ty->shape.as(); + auto x2_shape = x2_ty->shape.as(); if (!x1_shape || !x2_shape) { return TensorType(x1_ty->dtype, 2); } diff --git a/src/relax/op/tensor/manipulate.cc b/src/relax/op/tensor/manipulate.cc index 92b2950632d1..fb3141b14721 100644 --- a/src/relax/op/tensor/manipulate.cc +++ b/src/relax/op/tensor/manipulate.cc @@ -236,8 +236,8 @@ Type InferTypeBroadcastTo(const Call& call, const BlockBuilder& ctx) { if (call->args.size() != 2) { TVM_FFI_VISIT_THROW(ValueError, call) << "broadcast_to should take 2 arguments."; } - const auto* data_ty = GetTypeAs(call->args[0]); - const auto* tgt_shape_ty = GetTypeAs(call->args[1]); + const auto* data_ty = GetTypeAs(call->args[0]); + const auto* tgt_shape_ty = GetTypeAs(call->args[1]); if (data_ty == nullptr) { TVM_FFI_VISIT_THROW(TypeError, call) << "broadcast_to requires the input data to be Tensor. However, the given one is " @@ -274,7 +274,7 @@ Type InferTypeBroadcastTo(const Call& call, const BlockBuilder& ctx) { for (int i = 0; i < old_ndim; ++i) { PrimExpr old_len = old_shape_value[old_ndim - i - 1]; PrimExpr tgt_len = tgt_shape_value[tgt_ndim - i - 1]; - const auto* old_len_int = old_len.as(); + const auto* old_len_int = old_len.as(); if (old_len_int != nullptr && old_len_int->value == 1) { continue; } else if (analyzer->CanProve(old_len != tgt_len)) { @@ -426,7 +426,7 @@ Type InferTypeConcat(const Call& call, const BlockBuilder& ctx) { } // Update the shape values for best effort check. - const auto* shape_expr = ty->shape.as(); + const auto* shape_expr = ty->shape.as(); if (shape_expr != nullptr) { shape_values.push_back(shape_expr->values); continue; @@ -508,19 +508,19 @@ InferLayoutOutput InferLayoutConcat( LayoutDecision in_layout = n_layout.LeafValue(); if ((in_layout->layout.has_value() ? in_layout->layout.value().ndim() : 0) != (in_layout->layout.has_value() ? in_layout->layout.value().ndim_primal() : 0)) { - const auto* tuple_ty = GetTypeAs(call->args[0]); + const auto* tuple_ty = GetTypeAs(call->args[0]); TVM_FFI_ICHECK(tuple_ty != nullptr) << " expects the input to be a Tuple of Tensors. However, the given input is " << call->args[0]->ty->GetTypeKey(); for (size_t i = 0; i < tuple_ty->fields.size(); ++i) { Type field_ty = tuple_ty->fields[i]; - const auto* field_tensor_ty = field_ty.as(); + const auto* field_tensor_ty = field_ty.as(); TVM_FFI_ICHECK(field_tensor_ty != nullptr) << call->op << " expects the input to be a Tuple of Tensors. However, the given input is " << call->args[0]->ty; auto t_ty = ffi::GetRef(field_tensor_ty); - ffi::Optional t_shape = ffi::GetRef(t_ty->shape.as()); + ffi::Optional t_shape = ffi::GetRef(t_ty->shape.as()); LayoutDecision curr_layout = nlayout_array[i].LeafValue(); if (!CanProveLayoutTransform(curr_layout->layout, in_layout->layout, t_shape.value()->values)) { @@ -577,7 +577,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { refl::GlobalDef().def("relax.op.expand_dims", expand_dims); } -Type InferTypeExpandDims(const CallNode* call_node) { +Type InferTypeExpandDims(const CallObj* call_node) { const Call call = ffi::GetRef(call_node); if (RequiresTensorInputNormalization(call)) return Type::Missing(); TensorType data_ty = GetUnaryInputTensorType(call); @@ -594,7 +594,7 @@ Type InferTypeExpandDims(const CallNode* call_node) { int output_ndim = data_ty->ndim + n_new_dim; std::vector axes = NormalizeAxes(call, output_ndim, attrs->axis); - const auto* data_shape = data_ty->shape.as(); + const auto* data_shape = data_ty->shape.as(); if (data_shape == nullptr) { return TensorType(data_ty->dtype, output_ndim, data_ty->vdevice); } @@ -620,7 +620,7 @@ InferLayoutOutput InferLayoutExpandDims( TVM_FFI_ICHECK(NoDesiredLayout(call, desired_layouts)); const auto* attrs = call->attrs.as(); TVM_FFI_ICHECK(attrs != nullptr) << "Invalid Call"; - const auto* tensor_ty = GetTypeAs(call->args[0]); + const auto* tensor_ty = GetTypeAs(call->args[0]); TVM_FFI_ICHECK(tensor_ty != nullptr) << "Invalid Call"; TVM_FFI_ICHECK(!tensor_ty->IsUnknownNdim()) << "Only support static ndim for now"; @@ -685,7 +685,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { refl::GlobalDef().def("relax.op.flatten", flatten); } -Type InferTypeFlatten(const CallNode* call_node) { +Type InferTypeFlatten(const CallObj* call_node) { const Call call = ffi::GetRef(call_node); TensorType data_ty = GetUnaryInputTensorType(call); if (data_ty->IsUnknownNdim()) { @@ -696,7 +696,7 @@ Type InferTypeFlatten(const CallNode* call_node) { return data_ty; } - const auto* data_shape = data_ty->shape.as(); + const auto* data_shape = data_ty->shape.as(); if (data_shape == nullptr) { return TensorType(data_ty->dtype, /*ndim=*/1, data_ty->vdevice); } @@ -765,7 +765,7 @@ Type InferTypeIndexTensor(const Call& call, const BlockBuilder& ctx) { int max_index_ndim = 0; for (const auto& s : indices_ty) { - const auto* shp = s->shape.as(); + const auto* shp = s->shape.as(); if (!shp) { all_index_have_shape_value = false; } else { @@ -797,8 +797,8 @@ Type InferTypeIndexTensor(const Call& call, const BlockBuilder& ctx) { PrimExpr lhs_dim = out_shape[lhs_axis]; PrimExpr rhs_dim = ishape[rhs_axis]; - const auto* lhs_int = lhs_dim.as(); - const auto* rhs_int = rhs_dim.as(); + const auto* lhs_int = lhs_dim.as(); + const auto* rhs_int = rhs_dim.as(); // Case 1: current broadcast slot is 1 -> always replace if (lhs_int && lhs_int->value == 1) { @@ -840,7 +840,7 @@ Type InferTypeIndexTensor(const Call& call, const BlockBuilder& ctx) { // Derive output shape if (broadcast_shape.has_value()) { - const auto* data_shape_expr = data_ty->shape.as(); + const auto* data_shape_expr = data_ty->shape.as(); if (data_shape_expr) { ffi::Array result_shape = broadcast_shape.value(); for (int i = n_indices; i < data_ty->ndim; ++i) { @@ -878,7 +878,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { refl::GlobalDef().def("relax.op.layout_transform", layout_transform); } -Type InferTypeLayoutTransform(const CallNode* call_node) { +Type InferTypeLayoutTransform(const CallObj* call_node) { const Call call = ffi::GetRef(call_node); TensorType data_ty = GetUnaryInputTensorType(call); const auto* attrs = call->attrs.as(); @@ -955,7 +955,7 @@ bool IsIdentityPermutation(const std::vector& permutation) { return true; } -Type InferTypePermuteDims(const CallNode* call_node) { +Type InferTypePermuteDims(const CallObj* call_node) { const Call call = ffi::GetRef(call_node); if (RequiresTensorInputNormalization(call)) return Type::Missing(); TensorType data_ty = GetUnaryInputTensorType(call); @@ -990,7 +990,7 @@ Type InferTypePermuteDims(const CallNode* call_node) { return data_ty; } - const auto* data_shape = data_ty->shape.as(); + const auto* data_shape = data_ty->shape.as(); if (data_shape == nullptr) { return TensorType(data_ty->dtype, data_ty->ndim, data_ty->vdevice); } @@ -1009,7 +1009,7 @@ InferLayoutOutput InferLayoutPermuteDims( const auto* attrs = call->attrs.as(); TVM_FFI_ICHECK(attrs != nullptr) << "Invalid Call"; - const auto* tensor_ty = GetTypeAs(call->args[0]); + const auto* tensor_ty = GetTypeAs(call->args[0]); TVM_FFI_ICHECK(tensor_ty != nullptr) << "Invalid Call"; TVM_FFI_ICHECK(!tensor_ty->IsUnknownNdim()) << "Only support static ndim for now"; int ndim = tensor_ty->ndim; @@ -1060,10 +1060,10 @@ Expr ConvertNewShapeToExpr(const Expr& data, const ffi::Variant, Expr>& shape) { const ffi::ArrayObj* array; // Treat shape expressions as constant arrays to handle special values. - if (const auto* e = shape.as()) { + if (const auto* e = shape.as()) { array = e->values.as(); // Other non-shape expressions are used directly. - } else if (const auto* e = shape.as()) { + } else if (const auto* e = shape.as()) { return ffi::GetRef(e); // Process special values in constants and produce an expression. } else { @@ -1083,13 +1083,13 @@ Expr ConvertNewShapeToExpr(const Expr& data, "Array of PrimExprs. However, the given new shape is " << shape; PrimExpr len = prim_len.value(); - TVM_FFI_CHECK(!len.as(), TypeError) + TVM_FFI_CHECK(!len.as(), TypeError) << "Reshape shape dimensions must be symbolic primitive expressions, but received " << len; TVM_FFI_ICHECK(len.ty().code() == DLDataTypeCode::kDLInt) << "Reshape requires the new shape values to be all " "integers. However, the give new shape is " << shape; - const auto* int_len = len.as(); + const auto* int_len = len.as(); if (int_len != nullptr && int_len->value == 0) { // Note that this dimension should be copied from the original shape. zero_dims.push_back(i); @@ -1115,14 +1115,14 @@ Expr ConvertNewShapeToExpr(const Expr& data, } // Otherwise, we require the input tensor to have known shape value for inference. - const auto* data_ty = GetTypeAs(data); + const auto* data_ty = GetTypeAs(data); TVM_FFI_ICHECK(data_ty != nullptr) << "Reshape expects the input data to be a Tensor. However, the given input is " << data->ty->GetTypeKey(); TVM_FFI_ICHECK(data_ty->shape.has_value()) << "Reshape expects the input tensor to have known shape when there is some dimension length " "to infer. However, the given input has no shape."; - const auto* shape_ty = GetTypeAs(data_ty->shape.value()); + const auto* shape_ty = GetTypeAs(data_ty->shape.value()); TVM_FFI_ICHECK(shape_ty != nullptr && shape_ty->values.has_value()) << "Reshape expects the input tensor to have known shape when there is some dimension length " "to infer. However, the given input shape is " @@ -1140,7 +1140,7 @@ Expr ConvertNewShapeToExpr(const Expr& data, PrimExpr new_shape_prod = IntImm::Int64(1); for (int i = 0; i < static_cast(array_ref.size()); ++i) { PrimExpr new_dim = array_ref[i]; - const auto* int_dim = new_dim.as(); + const auto* int_dim = new_dim.as(); // We expect any symbolic not to signal the intent of -1, and therefore do no check for // symbolic value here. if (int_dim == nullptr || int_dim->value > 0) { @@ -1175,8 +1175,8 @@ Type InferTypeReshape(const Call& call, const BlockBuilder& ctx) { if (call->args.size() != 2) { TVM_FFI_VISIT_THROW(ValueError, call) << "Reshape op should take 2 arguments"; } - const auto* data_ty = GetTypeAs(call->args[0]); - const auto* new_shape_ty = GetTypeAs(call->args[1]); + const auto* data_ty = GetTypeAs(call->args[0]); + const auto* new_shape_ty = GetTypeAs(call->args[1]); if (data_ty == nullptr) { TVM_FFI_VISIT_THROW(TypeError, call) << "Reshape requires the input data to be Tensor. However, the given one is " @@ -1190,7 +1190,7 @@ Type InferTypeReshape(const Call& call, const BlockBuilder& ctx) { ffi::Optional> old_shape_values; if (data_ty->shape.has_value()) { - const auto* old_shape_ty = GetTypeAs(data_ty->shape.value()); + const auto* old_shape_ty = GetTypeAs(data_ty->shape.value()); TVM_FFI_ICHECK_NOTNULL(old_shape_ty); old_shape_values = old_shape_ty->values; } @@ -1208,7 +1208,7 @@ Type InferTypeReshape(const Call& call, const BlockBuilder& ctx) { } Expr target_shape = call->args[1]; // If shape values are defined, use them - if (target_shape->IsInstance() && new_shape_ty->values.has_value()) { + if (target_shape->IsInstance() && new_shape_ty->values.has_value()) { return TensorType(ShapeExpr(new_shape_ty->values.value()), data_ty->dtype, data_ty->vdevice); } return TensorType(target_shape, data_ty->dtype, data_ty->vdevice); @@ -1230,7 +1230,7 @@ Expr split(Expr x, ffi::Variant> indices_or_sections, if (const auto* indices = indices_or_sections.as()) { for (int i = 0; i < static_cast(indices->size()); ++i) { - const auto* idx = indices->at(i).as(); + const auto* idx = indices->at(i).as(); TVM_FFI_ICHECK(idx != nullptr) << "Split op only accepts an array of integers as the indices. " "However, the given indices " @@ -1238,7 +1238,7 @@ Expr split(Expr x, ffi::Variant> indices_or_sections, } indices_or_sections_obj = ConvertIntImmToInt64( ffi::GetRef(indices).as_or_throw>()); - } else if (const auto* n_section = indices_or_sections.as()) { + } else if (const auto* n_section = indices_or_sections.as()) { TVM_FFI_ICHECK_GT(n_section->value, 0) << "Split op expects the input number of sections to be a " "positive integer. However, the given number of sections is " @@ -1264,7 +1264,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { Type InferTypeSplit(const Call& call, const BlockBuilder& ctx) { TensorType data_ty = GetUnaryInputTensorType(call, ctx); const auto* attrs = call->attrs.as(); - const auto* data_shape = data_ty->shape.as(); + const auto* data_shape = data_ty->shape.as(); int axis = data_ty->IsUnknownNdim() ? -1 : NormalizeAxis(call, ctx, data_ty->ndim, attrs->axis); if (auto opt_indices = attrs->indices_or_sections.as>()) { @@ -1301,7 +1301,7 @@ Type InferTypeSplit(const Call& call, const BlockBuilder& ctx) { output_ty.push_back(TensorType(ShapeExpr(shape), data_ty->dtype, data_ty->vdevice)); } return TupleType(output_ty); - } else if (const auto* p_n_section = attrs->indices_or_sections.as()) { + } else if (const auto* p_n_section = attrs->indices_or_sections.as()) { TVM_FFI_ICHECK_GT(p_n_section->value, 0); int n_section = p_n_section->value.as().value(); // When the number of section is one, return the input tensor's type. @@ -1341,13 +1341,13 @@ InferLayoutOutput InferLayoutSplit( const auto* attrs = call->attrs.as(); TVM_FFI_ICHECK(attrs != nullptr) << "Invalid Call"; - const auto* tensor_ty = GetTypeAs(call->args[0]); + const auto* tensor_ty = GetTypeAs(call->args[0]); TVM_FFI_ICHECK(tensor_ty != nullptr) << "Invalid Call"; TVM_FFI_ICHECK(!tensor_ty->IsUnknownNdim()) << "Only support known ndim"; LayoutDecision existing_layout = GetLayoutDecision(var_layout_map, call->args[0]); Type out_ty = InferTypeSplit(call, BlockBuilder::Create(IRModule())); - const auto* out_tuple = out_ty.as(); + const auto* out_tuple = out_ty.as(); /* * Fallback if the outputs can't be represented in input sub indexed layout @@ -1356,11 +1356,11 @@ InferLayoutOutput InferLayoutSplit( if ((existing_layout->layout.has_value() ? existing_layout->layout.value().ndim() : 0) != (existing_layout->layout.has_value() ? existing_layout->layout.value().ndim_primal() : 0)) { for (const auto& si : out_tuple->fields) { - TVM_FFI_ICHECK(si->IsInstance()) << "Fields of TupleType must be TensorType" - "output structinfo, but got " - << si; + TVM_FFI_ICHECK(si->IsInstance()) << "Fields of TupleType must be TensorType" + "output structinfo, but got " + << si; auto ty = si.as_or_throw(); - ffi::Optional shape_expr = ffi::GetRef(ty->shape.as()); + ffi::Optional shape_expr = ffi::GetRef(ty->shape.as()); TVM_FFI_ICHECK(shape_expr.has_value()); auto shape_arr = shape_expr.value(); if (!CanProveLayoutTransform(InitialLayout(tensor_ty->ndim), existing_layout->layout, @@ -1402,7 +1402,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { refl::GlobalDef().def("relax.op.squeeze", squeeze); } -Type InferTypeSqueeze(const CallNode* call_node) { +Type InferTypeSqueeze(const CallObj* call_node) { const Call call = ffi::GetRef(call_node); if (RequiresTensorInputNormalization(call)) return Type::Missing(); TensorType data_ty = GetUnaryInputTensorType(call); @@ -1432,7 +1432,7 @@ Type InferTypeSqueeze(const CallNode* call_node) { for (int i = 0; i < static_cast(axes.size()); ++i) { // Todo(relax-team): revisit here for better check on if the axis being squeezed has length 1. // When `axis` is given, the dim lengths at the axes must be integer 1 when it is not symbolic - const auto* int_len = shape_value.value()[axes[i]].as(); + const auto* int_len = shape_value.value()[axes[i]].as(); // If a dimension is not 1, silently skip it (no-op), matching PyTorch behavior. if ((int_len != nullptr && int_len->value == 1) || int_len == nullptr) { axis_removal_mask[axes[i]] = true; @@ -1448,7 +1448,7 @@ Type InferTypeSqueeze(const CallNode* call_node) { } for (int i = 0; i < data_ty->ndim; ++i) { // Whenever a dimension length is symbolic, fall back to unknown ndim. - const auto* int_len = shape_value.value()[i].as(); + const auto* int_len = shape_value.value()[i].as(); if (int_len == nullptr) { return TensorType(data_ty->dtype, kUnknownNDim, data_ty->vdevice); } @@ -1466,7 +1466,7 @@ Type InferTypeSqueeze(const CallNode* call_node) { } } - if (data_ty->shape.value()->IsInstance()) { + if (data_ty->shape.value()->IsInstance()) { if (static_cast(output_shape.size()) == data_ty->ndim) { return data_ty; } else if (attrs->axis.has_value()) { @@ -1486,12 +1486,12 @@ InferLayoutOutput InferLayoutSqueeze( const auto* attrs = call->attrs.as(); TVM_FFI_ICHECK(attrs != nullptr) << "Invalid Call"; - const auto* tensor_ty = GetTypeAs(call->args[0]); + const auto* tensor_ty = GetTypeAs(call->args[0]); TVM_FFI_ICHECK(tensor_ty != nullptr) << "Invalid Call"; TVM_FFI_ICHECK(!tensor_ty->IsUnknownNdim()) << "Only support static ndim for now"; TVM_FFI_ICHECK(tensor_ty->shape.has_value()) << "Only support static shape for now"; int ndim = tensor_ty->ndim; - const auto* shape = tensor_ty->shape.as(); + const auto* shape = tensor_ty->shape.as(); TVM_FFI_ICHECK(shape != nullptr) << "Only support static shape for now"; ffi::Array axis; @@ -1565,8 +1565,8 @@ void CheckCollapseShape(const Call& call, const BlockBuilder& ctx, } const PrimExpr& dim0 = data_shape[data_ax]; const PrimExpr& dim1 = target_shape[target_ax]; - const auto* int_dim0 = dim0.as(); - const auto* int_dim1 = dim1.as(); + const auto* int_dim0 = dim0.as(); + const auto* int_dim1 = dim1.as(); if (analyzer->CanProveEqual(dim0, dim1) || (int_dim1 != nullptr && int_dim1->value == 1)) { --target_ax; @@ -1695,7 +1695,7 @@ Type InferTypeStack(const Call& call, const BlockBuilder& ctx) { } // Collect shape information - const auto* shape_expr = ty->shape.as(); + const auto* shape_expr = ty->shape.as(); if (shape_expr != nullptr) { shape_values.push_back(shape_expr->values); continue; @@ -1823,12 +1823,12 @@ Type InferTypeCollapseSumLike(const Call& call, const BlockBuilder& ctx) { ffi::Optional> data_shape_value; if (data_ty->shape.has_value()) { - data_shape_value = GetTypeAs(data_ty->shape.value())->values; + data_shape_value = GetTypeAs(data_ty->shape.value())->values; } ffi::Optional> collapse_target_shape_value; if (collapse_target_ty->shape.has_value()) { collapse_target_shape_value = - GetTypeAs(collapse_target_ty->shape.value())->values; + GetTypeAs(collapse_target_ty->shape.value())->values; } if (data_shape_value.has_value() && collapse_target_shape_value.has_value()) { @@ -1866,8 +1866,8 @@ Type InferTypeCollapseSumTo(const Call& call, const BlockBuilder& ctx) { TVM_FFI_VISIT_THROW(ValueError, call) << "CollapseSumTo should have 2 arguments"; } - const auto* data_ty = GetTypeAs(call->args[0]); - const auto* shape_ty = GetTypeAs(call->args[1]); + const auto* data_ty = GetTypeAs(call->args[0]); + const auto* shape_ty = GetTypeAs(call->args[1]); if (data_ty == nullptr) { TVM_FFI_VISIT_THROW(TypeError, call) @@ -1884,7 +1884,7 @@ Type InferTypeCollapseSumTo(const Call& call, const BlockBuilder& ctx) { ffi::Optional> data_shape_value; if (data_ty->shape.has_value()) { - data_shape_value = GetTypeAs(data_ty->shape.value())->values; + data_shape_value = GetTypeAs(data_ty->shape.value())->values; } if (data_shape_value.has_value() && shape_ty->values.has_value()) { @@ -1921,7 +1921,7 @@ Type InferTypeRepeat(const Call& call, const BlockBuilder& ctx) { 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(); + const auto* data_shape = data_ty->shape.as(); if (attrs->axis.has_value() && !data_ty->IsUnknownNdim()) { int axis = attrs->axis.value(); @@ -1967,7 +1967,7 @@ InferLayoutOutput InferLayoutRepeat( const auto* attrs = call->attrs.as(); TVM_FFI_ICHECK(attrs != nullptr) << "Invalid Call"; - const auto* tensor_ty = GetTypeAs(call->args[0]); + const auto* tensor_ty = GetTypeAs(call->args[0]); TVM_FFI_ICHECK(tensor_ty != nullptr) << "Invalid Call"; TVM_FFI_ICHECK(!tensor_ty->IsUnknownNdim()) << "Only support static ndim for now"; @@ -2046,7 +2046,7 @@ Type InferTypeTile(const Call& call, const BlockBuilder& ctx) { 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(); + const auto* data_shape = data_ty->shape.as(); int l = attrs->repeats.size(); int ndim = data_ty->ndim; @@ -2092,7 +2092,7 @@ InferLayoutOutput InferLayoutTile( const auto* attrs = call->attrs.as(); TVM_FFI_ICHECK(attrs != nullptr) << "Invalid Call"; - const auto* tensor_ty = GetTypeAs(call->args[0]); + const auto* tensor_ty = GetTypeAs(call->args[0]); TVM_FFI_ICHECK(tensor_ty != nullptr) << "Invalid Call"; TVM_FFI_ICHECK(!tensor_ty->IsUnknownNdim()) << "Only support static ndim for now"; @@ -2186,7 +2186,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { refl::GlobalDef().def("relax.op.flip", flip); } -Type InferTypeFlip(const CallNode* call_node) { +Type InferTypeFlip(const CallObj* call_node) { const Call call = ffi::GetRef(call_node); if (call->args.size() != 1) { TVM_FFI_VISIT_THROW(ValueError, call) << "Flip op should take 1 argument"; @@ -2212,7 +2212,7 @@ InferLayoutOutput InferLayoutFlip( const auto* attrs = call->attrs.as(); TVM_FFI_ICHECK(attrs != nullptr) << "Invalid Call"; - const auto* tensor_ty = GetTypeAs(call->args[0]); + const auto* tensor_ty = GetTypeAs(call->args[0]); TVM_FFI_ICHECK(tensor_ty != nullptr) << "Invalid Call"; TVM_FFI_ICHECK(!tensor_ty->IsUnknownNdim()) << "Only support static ndim for now"; @@ -2311,8 +2311,8 @@ Type InferTypeReverseSequence(const Call& call, const BlockBuilder& ctx) { } if (data_ty->shape.has_value() && seq_lengths_ty->shape.has_value()) { - const auto* data_shape_ty = GetTypeAs(data_ty->shape.value()); - const auto* seq_lengths_shape_ty = GetTypeAs(seq_lengths_ty->shape.value()); + const auto* data_shape_ty = GetTypeAs(data_ty->shape.value()); + const auto* seq_lengths_shape_ty = GetTypeAs(seq_lengths_ty->shape.value()); if (data_shape_ty != nullptr && seq_lengths_shape_ty != nullptr && data_shape_ty->values.has_value() && seq_lengths_shape_ty->values.has_value()) { PrimExpr batch_extent = data_shape_ty->values.value()[batch_axis]; @@ -2354,11 +2354,11 @@ TVM_FFI_STATIC_INIT_BLOCK() { refl::GlobalDef().def("relax.op.gather_elements", gather_elements); } -Type InferTypeGatherElements(const CallNode* call_node) { +Type InferTypeGatherElements(const CallObj* call_node) { const Call call = ffi::GetRef(call_node); if (RequiresTensorInputNormalization(call)) return Type::Missing(); - const auto* data_ty = GetTypeAs(call->args[0]); - const auto* indices_ty = GetTypeAs(call->args[1]); + const auto* data_ty = GetTypeAs(call->args[0]); + const auto* indices_ty = GetTypeAs(call->args[1]); const auto* attrs = call->attrs.as(); if (data_ty == nullptr) { @@ -2429,7 +2429,7 @@ InferLayoutOutput InferLayoutGatherElements( if ((layout->layout.has_value() ? layout->layout.value().ndim() : 0) != (layout->layout.has_value() ? layout->layout.value().ndim_primal() : 0)) { - const auto* tensor_ty = GetTypeAs(call->args[0]); + const auto* tensor_ty = GetTypeAs(call->args[0]); TVM_FFI_ICHECK(tensor_ty != nullptr) << "Invalid Call"; TVM_FFI_ICHECK(!tensor_ty->IsUnknownNdim()) << "Only support static ndim for now"; int ndim = tensor_ty->ndim; @@ -2464,11 +2464,11 @@ TVM_FFI_STATIC_INIT_BLOCK() { refl::GlobalDef().def("relax.op.gather_nd", gather_nd); } -Type InferTypeGatherND(const CallNode* call_node) { +Type InferTypeGatherND(const CallObj* call_node) { const Call call = ffi::GetRef(call_node); if (RequiresTensorInputNormalization(call)) return Type::Missing(); - const auto* data_ty = GetTypeAs(call->args[0]); - const auto* indices_ty = GetTypeAs(call->args[1]); + const auto* data_ty = GetTypeAs(call->args[0]); + const auto* indices_ty = GetTypeAs(call->args[1]); const auto* attrs = call->attrs.as(); if (data_ty == nullptr) { @@ -2507,12 +2507,12 @@ Type InferTypeGatherND(const CallNode* call_node) { } // Check if indices shape is known - const auto* indices_shape = indices_ty->shape.as(); - const auto* data_shape = data_ty->shape.as(); - if (!indices_shape || !indices_shape->values.back()->IsInstance()) { + const auto* indices_shape = indices_ty->shape.as(); + const auto* data_shape = data_ty->shape.as(); + if (!indices_shape || !indices_shape->values.back()->IsInstance()) { return TensorType(data_ty->dtype, kUnknownNDim, data_ty->vdevice); } - int l = indices_shape->values.back().as()->value.as().value(); + int l = indices_shape->values.back().as()->value.as().value(); int output_ndim = indices_ty->ndim + input_dims - l - 1 - batch_dims; if (!data_shape) { return TensorType(data_ty->dtype, output_ndim, data_ty->vdevice); @@ -2559,12 +2559,12 @@ TVM_FFI_STATIC_INIT_BLOCK() { refl::GlobalDef().def("relax.op.index_put", index_put); } -Type InferTypeIndexPut(const CallNode* call_node) { +Type InferTypeIndexPut(const CallObj* call_node) { const Call call = ffi::GetRef(call_node); - const auto* data_ty = GetTypeAs(call->args[0]); - const auto* values_ty = GetTypeAs(call->args[2]); + const auto* data_ty = GetTypeAs(call->args[0]); + const auto* values_ty = GetTypeAs(call->args[2]); - auto diag_def = [&](const TensorTypeNode* ty, ffi::String name, ffi::String type_key) { + auto diag_def = [&](const TensorTypeObj* ty, ffi::String name, ffi::String type_key) { if (ty == nullptr) { TVM_FFI_VISIT_THROW(TypeError, call) << "IndexPut requires the input " << name << " to be a Tensor. However, the given one is " @@ -2578,10 +2578,10 @@ Type InferTypeIndexPut(const CallNode* call_node) { // Handle indices: either a single tensor or a tuple of tensors ffi::Array indices_tensors; - if (const auto* tuple_ty = GetTypeAs(call->args[1])) { + if (const auto* tuple_ty = GetTypeAs(call->args[1])) { // Indices is a tuple of tensors for (size_t i = 0; i < tuple_ty->fields.size(); ++i) { - const auto* tensor_ty = tuple_ty->fields[i].as(); + const auto* tensor_ty = tuple_ty->fields[i].as(); if (tensor_ty == nullptr) { TVM_FFI_VISIT_THROW(TypeError, call) << "IndexPut requires each index in the indices tuple to be a Tensor. " @@ -2589,7 +2589,7 @@ Type InferTypeIndexPut(const CallNode* call_node) { } indices_tensors.push_back(ffi::GetRef(tensor_ty)); } - } else if (const auto* tensor_ty = GetTypeAs(call->args[1])) { + } else if (const auto* tensor_ty = GetTypeAs(call->args[1])) { // Indices is a single tensor indices_tensors.push_back(ffi::GetRef(tensor_ty)); } else { @@ -2637,7 +2637,7 @@ Type InferTypeIndexPut(const CallNode* call_node) { const auto& tensor_ty = indices_tensors[i]; if (!tensor_ty->IsUnknownNdim() && tensor_ty->ndim > 1) { // Check that multi-dimensional indices are broadcastable - const auto* shape = tensor_ty->shape.as(); + const auto* shape = tensor_ty->shape.as(); if (shape) { // Verify trailing dimensions can broadcast // For now, we accept any multi-dimensional index and rely on runtime validation @@ -2657,7 +2657,7 @@ Type InferTypeIndexPut(const CallNode* call_node) { // Check data and values dtype compatibility if (data_ty->IsUnknownDtype() || values_ty->IsUnknownDtype()) { - auto diag_dtype = [&](const TensorTypeNode* ty, ffi::String name) { + auto diag_dtype = [&](const TensorTypeObj* ty, ffi::String name) { if (ty->IsUnknownDtype()) { LOG(WARNING) << "Data type of " << name << " has not been specified. Assume it has an integer type."; @@ -2673,7 +2673,7 @@ Type InferTypeIndexPut(const CallNode* call_node) { } // Check values shape compatibility - const auto* values_shape = values_ty->shape.as(); + const auto* values_shape = values_ty->shape.as(); if (values_shape) { if (values_ty->ndim != 1) { LOG(WARNING) << "IndexPut typically expects values to be 1D, but got ndim=" @@ -2681,7 +2681,7 @@ Type InferTypeIndexPut(const CallNode* call_node) { } } - const auto* data_shape = data_ty->shape.as(); + const auto* data_shape = data_ty->shape.as(); if (data_shape) { return TensorType(ShapeExpr(data_shape->values), data_ty->dtype, data_ty->vdevice); } @@ -2711,7 +2711,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { refl::GlobalDef().def("relax.op.meshgrid", meshgrid); } -Type InferTypeMeshgrid(const CallNode* call_node) { +Type InferTypeMeshgrid(const CallObj* call_node) { const Call call = ffi::GetRef(call_node); if (call->args.size() != 1) { TVM_FFI_VISIT_THROW(ValueError, call) << "meshgrid op expects 1 Tuple input argument."; @@ -2750,7 +2750,7 @@ Type InferTypeMeshgrid(const CallNode* call_node) { << " and " << common_dtype; } - const auto* shape_expr = ty->shape.as(); + const auto* shape_expr = ty->shape.as(); if (shape_expr && shape_expr->values.size() == 1) { lengths.push_back(shape_expr->values[0]); } else { @@ -2821,11 +2821,11 @@ TVM_FFI_STATIC_INIT_BLOCK() { Type InferTypeScatterElements(const Call& call, const BlockBuilder& ctx) { 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]); + const auto* data_ty = GetTypeAs(call->args[0]); + const auto* indices_ty = GetTypeAs(call->args[1]); + const auto* updates_ty = GetTypeAs(call->args[2]); - auto diag_def = [&](const TensorTypeNode* ty, ffi::String name, ffi::String type_key) { + auto diag_def = [&](const TensorTypeObj* ty, ffi::String name, ffi::String type_key) { if (ty == nullptr) { TVM_FFI_VISIT_THROW(TypeError, call) << "ScatterElements requires the input " << name @@ -2860,7 +2860,7 @@ Type InferTypeScatterElements(const Call& call, const BlockBuilder& ctx) { } if (data_ty->IsUnknownDtype() || updates_ty->IsUnknownDtype()) { - auto diag_dtype = [&](const TensorTypeNode* ty, ffi::String name) { + auto diag_dtype = [&](const TensorTypeObj* ty, ffi::String name) { if (ty->IsUnknownDtype()) { LOG(WARNING) << "Data type of " << name << " has not been specified. Assume it has an integer type."; @@ -2890,8 +2890,8 @@ Type InferTypeScatterElements(const Call& call, const BlockBuilder& ctx) { } } - const auto* indices_shape = indices_ty->shape.as(); - const auto* updates_shape = updates_ty->shape.as(); + const auto* indices_shape = indices_ty->shape.as(); + const auto* updates_shape = updates_ty->shape.as(); if (indices_shape && updates_shape) { for (int i = 0; i < indices_ty->ndim; i++) { if (analyzer->CanProve(indices_shape->values[i] != updates_shape->values[i])) { @@ -2903,7 +2903,7 @@ Type InferTypeScatterElements(const Call& call, const BlockBuilder& ctx) { } } } - const auto* data_shape = data_ty->shape.as(); + const auto* data_shape = data_ty->shape.as(); if (data_shape) { return TensorType(ShapeExpr(data_shape->values), data_ty->dtype, data_ty->vdevice); } @@ -2928,7 +2928,7 @@ InferLayoutOutput InferLayoutScatterElements( if ((layout->layout.has_value() ? layout->layout.value().ndim() : 0) != (layout->layout.has_value() ? layout->layout.value().ndim_primal() : 0)) { - const auto* tensor_ty = GetTypeAs(call->args[0]); + const auto* tensor_ty = GetTypeAs(call->args[0]); TVM_FFI_ICHECK(tensor_ty != nullptr) << "Invalid Call"; TVM_FFI_ICHECK(!tensor_ty->IsUnknownNdim()) << "Only support static ndim for now"; int ndim = tensor_ty->ndim; @@ -2968,9 +2968,9 @@ Type InferTypeScatterND(const Call& call, const BlockBuilder& ctx) { // `call->args` contains: [data, indices, updates] 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]); - const auto* updates_ty = GetTypeAs(call->args[2]); + const auto* data_ty = GetTypeAs(call->args[0]); + const auto* indices_ty = GetTypeAs(call->args[1]); + const auto* updates_ty = GetTypeAs(call->args[2]); if (data_ty == nullptr) { TVM_FFI_VISIT_THROW(TypeError, call) @@ -3015,12 +3015,12 @@ Type InferTypeScatterND(const Call& call, const BlockBuilder& ctx) { } } - const auto* data_shape = data_ty->shape.as(); - const auto* indices_shape = indices_ty->shape.as(); - const auto* updates_shape = updates_ty->shape.as(); + const auto* data_shape = data_ty->shape.as(); + const auto* indices_shape = indices_ty->shape.as(); + const auto* updates_shape = updates_ty->shape.as(); if (data_shape && indices_shape && updates_shape) { - const IntImmNode* k_dim = indices_shape->values[indices_ty->ndim - 1].as(); + const IntImmObj* k_dim = indices_shape->values[indices_ty->ndim - 1].as(); if (!k_dim) { TVM_FFI_VISIT_THROW(ValueError, call) << "ScatterND needs a static shape for the last axis of indices, got " @@ -3088,8 +3088,8 @@ InferLayoutOutput InferLayoutScatterND( LayoutDecision indices_layout = GetLayoutDecision(var_layout_map, call->args[1]); LayoutDecision updates_layout = GetLayoutDecision(var_layout_map, call->args[2]); - const auto* data_ty = GetTypeAs(call->args[0]); - const auto* updates_ty = GetTypeAs(call->args[2]); + const auto* data_ty = GetTypeAs(call->args[0]); + const auto* updates_ty = GetTypeAs(call->args[2]); TVM_FFI_ICHECK(data_ty != nullptr) << "Invalid Call"; TVM_FFI_ICHECK(updates_ty != nullptr) << "Invalid Call"; TVM_FFI_ICHECK(!data_ty->IsUnknownNdim()) << "Only support static ndim for now"; @@ -3145,11 +3145,11 @@ TVM_FFI_STATIC_INIT_BLOCK() { Type InferTypeSliceScatter(const Call& call, const BlockBuilder& ctx) { sym::Analyzer analyzer = ctx->GetAnalyzer(); - const auto* data_ty = GetTypeAs(call->args[0]); - const auto* src_ty = GetTypeAs(call->args[1]); + const auto* data_ty = GetTypeAs(call->args[0]); + const auto* src_ty = GetTypeAs(call->args[1]); auto* attrs = call->attrs.as(); - auto diag_tensor_check = [&](const TensorTypeNode* ty, const Expr& arg_expr, ffi::String name) { + auto diag_tensor_check = [&](const TensorTypeObj* ty, const Expr& arg_expr, ffi::String name) { if (ty == nullptr) { TVM_FFI_VISIT_THROW(TypeError, call) << "SliceScatter requires the input " << name @@ -3183,7 +3183,7 @@ Type InferTypeSliceScatter(const Call& call, const BlockBuilder& ctx) { } if (data_ty->IsUnknownDtype() || src_ty->IsUnknownDtype()) { - auto diag_dtype_warn = [&](const TensorTypeNode* ty, ffi::String name) { + auto diag_dtype_warn = [&](const TensorTypeObj* ty, ffi::String name) { if (ty->IsUnknownDtype()) { LOG(WARNING) << "SliceScatter: Data type of " << name << " has not been specified for call node " << call @@ -3234,8 +3234,8 @@ Type InferTypeSliceScatter(const Call& call, const BlockBuilder& ctx) { int axis = NormalizeAxis(call, ctx, ndim, attrs->axis); - const auto* data_shape_node = data_ty->shape.as(); - const auto* src_shape_node = src_ty->shape.as(); + const auto* data_shape_node = data_ty->shape.as(); + const auto* src_shape_node = src_ty->shape.as(); if (data_shape_node && src_shape_node && !src_ty->IsUnknownNdim()) { TVM_FFI_ICHECK_EQ(data_shape_node->values.size(), static_cast(ndim)) @@ -3309,7 +3309,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { refl::GlobalDef().def("relax.op.one_hot", one_hot); } -Type InferTypeOneHot(const CallNode* call_node) { +Type InferTypeOneHot(const CallObj* call_node) { const Call call = ffi::GetRef(call_node); TensorType indices_ty = GetInputTensorType(call, 0); const auto* attrs = call->attrs.as(); @@ -3341,7 +3341,7 @@ Type InferTypeOneHot(const CallNode* call_node) { return TensorType(dtype, kUnknownNDim, indices_ty->vdevice); } // Get the shape of indices - const auto* indices_shape = indices_ty->shape.as(); + const auto* indices_shape = indices_ty->shape.as(); if (indices_shape == nullptr) { return TensorType(dtype, indices_ty->ndim + 1, indices_ty->vdevice); } diff --git a/src/relax/op/tensor/qdq.cc b/src/relax/op/tensor/qdq.cc index d188d5536fc2..f81dc51dbb91 100644 --- a/src/relax/op/tensor/qdq.cc +++ b/src/relax/op/tensor/qdq.cc @@ -65,7 +65,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { } TensorType WithUnknownDtype(const TensorType& input_ty) { - auto output_ty = ffi::make_object(*input_ty.get()); + auto output_ty = ffi::make_object(*input_ty.get()); output_ty->dtype = ffi::Optional(); return TensorType(output_ty); } @@ -137,8 +137,8 @@ Type InferTypeQuantize(const Call& call, const BlockBuilder& ctx) { auto is_scalar_or_singleton_vector = [&](const TensorType& param_ty) { if (IsScalarTensor(param_ty)) return true; - if (param_ty->shape.has_value() && param_ty->shape.value()->IsInstance()) { - const auto& values = param_ty->shape.as()->values; + if (param_ty->shape.has_value() && param_ty->shape.value()->IsInstance()) { + const auto& values = param_ty->shape.as()->values; if (!values.empty()) { return std::all_of(values.begin(), values.end(), [&](const PrimExpr& dim) { return ctx->GetAnalyzer()->CanProveEqual(dim, 1); @@ -161,7 +161,7 @@ Type InferTypeQuantize(const Call& call, const BlockBuilder& ctx) { if (!is_scalar_or_singleton_vector(zp_ty)) check_param_size(zp_ty, input_ty, "zero_point"); } - auto output_ty = ffi::make_object(*input_ty.get()); + auto output_ty = ffi::make_object(*input_ty.get()); output_ty->dtype = PrimType(attrs->out_dtype); return TensorType(output_ty); } @@ -260,8 +260,8 @@ Type InferTypeDequantize(const Call& call, const BlockBuilder& ctx) { auto is_scalar_or_singleton_vector = [&](const TensorType& param_ty) { if (IsScalarTensor(param_ty)) return true; - if (param_ty->shape.has_value() && param_ty->shape.value()->IsInstance()) { - const auto& values = param_ty->shape.as()->values; + if (param_ty->shape.has_value() && param_ty->shape.value()->IsInstance()) { + const auto& values = param_ty->shape.as()->values; if (!values.empty()) { return std::all_of(values.begin(), values.end(), [&](const PrimExpr& dim) { return ctx->GetAnalyzer()->CanProveEqual(dim, 1); @@ -284,7 +284,7 @@ Type InferTypeDequantize(const Call& call, const BlockBuilder& ctx) { if (!is_scalar_or_singleton_vector(zp_ty)) check_param_size(zp_ty, input_ty, "zero_point"); } - auto output_ty = ffi::make_object(*input_ty.get()); + auto output_ty = ffi::make_object(*input_ty.get()); output_ty->dtype = PrimType(attrs->out_dtype); return TensorType(output_ty); } diff --git a/src/relax/op/tensor/sampling.cc b/src/relax/op/tensor/sampling.cc index d552f89895c8..4447a4cace65 100644 --- a/src/relax/op/tensor/sampling.cc +++ b/src/relax/op/tensor/sampling.cc @@ -63,7 +63,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { refl::GlobalDef().def("relax.op.multinomial_from_uniform", multinomial_from_uniform); } -Type InferTypeMultinomialFromUniform(const CallNode* call_node) { +Type InferTypeMultinomialFromUniform(const CallObj* call_node) { const Call call = ffi::GetRef(call_node); if (RequiresTensorInputNormalization(call)) return Type::Missing(); CheckNumArguments(call); @@ -121,11 +121,11 @@ Type InferTypeMultinomialFromUniform(const CallNode* call_node) { } // Expected to be `(batch, vocab_size)` - const auto* prob_shape = prob_ty->shape.as(); + const auto* prob_shape = prob_ty->shape.as(); // Expected to be `(n, 1)` - const auto* uniform_sample_shape = uniform_sample_ty->shape.as(); + const auto* uniform_sample_shape = uniform_sample_ty->shape.as(); // Expected to be `(n, 1)` - const auto* sample_indices_shape = sample_indices_ty->shape.as(); + const auto* sample_indices_shape = sample_indices_ty->shape.as(); // The output shape is expected to be `(n, 1)` if (prob_shape == nullptr || uniform_sample_shape == nullptr || sample_indices_shape == nullptr) { diff --git a/src/relax/op/tensor/search.cc b/src/relax/op/tensor/search.cc index ba3a982f9b6b..ac988ff707bc 100644 --- a/src/relax/op/tensor/search.cc +++ b/src/relax/op/tensor/search.cc @@ -79,7 +79,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { refl::GlobalDef().def("relax.op.bucketize", bucketize); } -Type InferTypeBucketize(const CallNode* call_node) { +Type InferTypeBucketize(const CallObj* call_node) { const Call call = ffi::GetRef(call_node); ffi::Array input_ty = GetInputTensorType(call); TensorType input_tensor_info = input_ty[0]; @@ -96,7 +96,7 @@ Type InferTypeBucketize(const CallNode* call_node) { out_dtype = PrimType::Int(32); } - const auto* data_shape = input_tensor_info->shape.as(); + const auto* data_shape = input_tensor_info->shape.as(); if (data_shape) { return TensorType(ShapeExpr(data_shape->values), out_dtype, input_tensor_info->vdevice); } @@ -165,9 +165,9 @@ Type InferTypeWhere(const Call& call, const BlockBuilder& ctx) { output_ndim = std::max(cond_ty->ndim, std::max(x1_ty->ndim, x2_ty->ndim)); } - const auto* cond_shape = cond_ty->shape.as(); - const auto* x1_shape = x1_ty->shape.as(); - const auto* x2_shape = x2_ty->shape.as(); + const auto* cond_shape = cond_ty->shape.as(); + const auto* x1_shape = x1_ty->shape.as(); + const auto* x2_shape = x2_ty->shape.as(); if (cond_shape && x1_shape && x2_shape) { // Step 1. Compute the broadcasted shape of x1's and x2's ffi::Optional> broadcasted_shape = @@ -220,7 +220,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { /* relax.argmax & relax.argmin */ -Type InferTypeArgmaxArgmin(const CallNode* call_node) { +Type InferTypeArgmaxArgmin(const CallObj* call_node) { const Call call = ffi::GetRef(call_node); TensorType data_ty = GetUnaryInputTensorType(call); const auto* attrs = call->attrs.as(); @@ -250,7 +250,7 @@ Type InferTypeArgmaxArgmin(const CallNode* call_node) { // - axes is not None, keepdims is false -> the returned shape does not contain the input axes. // - axes is not None, keepdims is true -> the returned shape has value 1 at the positions of the // input axes - const auto* data_shape = data_ty->shape.as(); + const auto* data_shape = data_ty->shape.as(); if (data_shape == nullptr) { if (!attrs->axis.has_value() && attrs->keepdims && out_ndim != kUnknownNDim) { return TensorType(ShapeExpr(ffi::Array(out_ndim, IntImm(out_dtype, /*value=*/1))), diff --git a/src/relax/op/tensor/set.cc b/src/relax/op/tensor/set.cc index 36f1bf1617c8..73166c64e170 100644 --- a/src/relax/op/tensor/set.cc +++ b/src/relax/op/tensor/set.cc @@ -56,7 +56,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { refl::GlobalDef().def("relax.op.unique", unique); } -Type InferTypeUnique(const CallNode* call_node) { +Type InferTypeUnique(const CallObj* call_node) { const Call call = ffi::GetRef(call_node); TensorType data_ty = call->args[0]->ty.as_or_throw(); ffi::Optional axis; @@ -67,7 +67,7 @@ Type InferTypeUnique(const CallNode* call_node) { } if (!data_ty->IsUnknownNdim() && axis.has_value()) { // Normalize the axis for sanity check purpose. - if (const auto* axis_int = axis.as()) { + if (const auto* axis_int = axis.as()) { NormalizeAxis(call, data_ty->ndim, axis_int->value.as().value()); } } @@ -80,9 +80,9 @@ Type InferTypeUnique(const CallNode* call_node) { PrimExpr return_counts = call->args[4].as_or_throw(); auto f_convert_to_int64 = [](const PrimExpr& value) { - TVM_FFI_ICHECK(value->IsInstance()) + TVM_FFI_ICHECK(value->IsInstance()) << value << " expects to be IntImm, but gets " << value->GetTypeKey(); - const auto* val_node = value.as(); + const auto* val_node = value.as(); auto val_imm = ffi::GetRef(val_node); return static_cast(val_imm->value); }; @@ -176,7 +176,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { refl::GlobalDef().def("relax.op.nonzero", nonzero); } -Type InferTypeNonzero(const CallNode* call_node) { +Type InferTypeNonzero(const CallObj* call_node) { const Call call = ffi::GetRef(call_node); TensorType data_ty = GetInputTensorType(call, 0); return TensorType(PrimType::Int(64), 2, data_ty->vdevice); diff --git a/src/relax/op/tensor/sorting.cc b/src/relax/op/tensor/sorting.cc index 57921126486c..ddfc15f02429 100644 --- a/src/relax/op/tensor/sorting.cc +++ b/src/relax/op/tensor/sorting.cc @@ -103,7 +103,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { refl::GlobalDef().def("relax.op.sort", sort); } -Type InferTypeSort(const CallNode* call_node) { +Type InferTypeSort(const CallObj* call_node) { const Call call = ffi::GetRef(call_node); return GetUnaryInputTensorType(call); } @@ -132,7 +132,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { refl::GlobalDef().def("relax.op.argsort", argsort); } -Type InferTypeArgsort(const CallNode* call_node) { +Type InferTypeArgsort(const CallObj* call_node) { const Call call = ffi::GetRef(call_node); TensorType data_ty = GetUnaryInputTensorType(call); const auto* attrs = call->attrs.as(); @@ -172,11 +172,11 @@ TVM_FFI_STATIC_INIT_BLOCK() { refl::GlobalDef().def("relax.op.topk", topk); } -Type InferTypeTopK(const CallNode* call_node) { +Type InferTypeTopK(const CallObj* call_node) { const Call call = ffi::GetRef(call_node); if (RequiresTensorInputNormalization(call)) return Type::Missing(); TensorType data_ty = GetUnaryInputTensorType(call); - const auto* data_shape = data_ty->shape.as(); + const auto* data_shape = data_ty->shape.as(); const auto* attrs = call->attrs.as(); ffi::Optional indices_type = attrs->dtype.has_value() ? ffi::Optional(PrimType(attrs->dtype.value())) @@ -196,7 +196,7 @@ Type InferTypeTopK(const CallNode* call_node) { output_tys.push_back(TensorType(indices_type, data_ty->ndim, data_ty->vdevice)); } else { ffi::Array out_shape = data_shape->values; - const auto* int_dim = out_shape[axis].as(); + const auto* int_dim = out_shape[axis].as(); if (k > 0 && (int_dim == nullptr || k < int_dim->value)) { out_shape.Set(axis, k); } diff --git a/src/relax/op/tensor/statistical.cc b/src/relax/op/tensor/statistical.cc index dc4ce4b6ea62..dda9acc2fe86 100644 --- a/src/relax/op/tensor/statistical.cc +++ b/src/relax/op/tensor/statistical.cc @@ -67,7 +67,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { ScanopAttrs::RegisterReflection(); } -Type InferTypeStatistical(const CallNode* call_node) { +Type InferTypeStatistical(const CallObj* call_node) { const Call call = ffi::GetRef(call_node); if (RequiresTensorInputNormalization(call)) return Type::Missing(); TensorType data_ty = GetUnaryInputTensorType(call); @@ -97,7 +97,7 @@ Type InferTypeStatistical(const CallNode* call_node) { // - axes is not None, keepdims is false -> the returned shape does not contain the input axes. // - axes is not None, keepdims is true -> the returned shape has value 1 at the positions of the // input axes - const auto* data_shape = data_ty->shape.as(); + const auto* data_shape = data_ty->shape.as(); if (data_shape == nullptr) { if (!attrs->axis.has_value() && attrs->keepdims && out_ndim != kUnknownNDim) { return TensorType(ShapeExpr(ffi::Array(out_ndim, IntImm::Int64(/*value=*/1))), @@ -129,7 +129,7 @@ InferLayoutOutput InferLayoutStatistical( const auto* attrs = call->attrs.as(); TVM_FFI_ICHECK(attrs != nullptr) << "Invalid Call"; - const auto* tensor_ty = GetTypeAs(call->args[0]); + const auto* tensor_ty = GetTypeAs(call->args[0]); TVM_FFI_ICHECK(tensor_ty != nullptr) << "Invalid Call"; TVM_FFI_ICHECK(!tensor_ty->IsUnknownNdim()) << "Only support known ndim"; int ndim = tensor_ty->ndim; @@ -186,7 +186,7 @@ InferLayoutOutput InferLayoutStatistical( Attrs(new_attrs)); } -Type InferTypeScan(const CallNode* call_node) { +Type InferTypeScan(const CallObj* call_node) { const Call call = ffi::GetRef(call_node); TensorType data_ty = GetUnaryInputTensorType(call); const auto* attrs = call->attrs.as(); @@ -197,7 +197,7 @@ Type InferTypeScan(const CallNode* call_node) { if (!attrs->axis.has_value()) { // flattened - const auto* data_shape = data_ty->shape.as(); + const auto* data_shape = data_ty->shape.as(); if (data_shape == nullptr) { return TensorType(out_type, data_ty->ndim, data_ty->vdevice); } else { @@ -216,7 +216,7 @@ Type InferTypeScan(const CallNode* call_node) { } } -Type InferTypeStatisticalExtension(const CallNode* call_node) { +Type InferTypeStatisticalExtension(const CallObj* call_node) { const Call call = ffi::GetRef(call_node); TensorType data_ty = GetUnaryInputTensorType(call); const auto* attrs = call->attrs.as(); @@ -245,7 +245,7 @@ Type InferTypeStatisticalExtension(const CallNode* call_node) { // - len(axes) == 1, keepdims is false -> the returned shape does not contain the input axis. // - len(axes) == 1, keepdims is true -> the returned shape has value 1 at the positions of the // input axis - const auto* data_shape = data_ty->shape.as(); + const auto* data_shape = data_ty->shape.as(); if (data_shape == nullptr) { if (!attrs->axis.has_value() && attrs->keepdims && out_ndim != kUnknownNDim) { return TensorType(ShapeExpr(ffi::Array(out_ndim, IntImm::Int64(/*value=*/1))), diff --git a/src/relax/op/tensor/ternary.cc b/src/relax/op/tensor/ternary.cc index 94a3d508a74a..c1e6093a97e9 100644 --- a/src/relax/op/tensor/ternary.cc +++ b/src/relax/op/tensor/ternary.cc @@ -83,9 +83,9 @@ Type InferTypeEwiseFMA(const Call& call, const BlockBuilder& ctx) { } } - auto* s1 = t1->shape.as(); - auto* s2 = t2->shape.as(); - auto* s3 = t3->shape.as(); + auto* s1 = t1->shape.as(); + auto* s2 = t2->shape.as(); + auto* s3 = t3->shape.as(); sym::Analyzer analyzer = ctx->GetAnalyzer(); if (s1 && s2 && s3) { ffi::Array output_shape; diff --git a/src/relax/op/tensor/unary.cc b/src/relax/op/tensor/unary.cc index ab82f0b7d1ea..689c5714088a 100644 --- a/src/relax/op/tensor/unary.cc +++ b/src/relax/op/tensor/unary.cc @@ -33,7 +33,7 @@ namespace tvm { namespace relax { -Type InferTypeUnaryCheck(const CallNode* call_node) { +Type InferTypeUnaryCheck(const CallObj* call_node) { const Call call = ffi::GetRef(call_node); return InferTypeUnary(call, [](const TensorType& input_ty) { return PrimType::Bool(); }); } diff --git a/src/relax/op/vision/multibox_transform_loc.cc b/src/relax/op/vision/multibox_transform_loc.cc index 06ee3c8fcca0..00d630e46c19 100644 --- a/src/relax/op/vision/multibox_transform_loc.cc +++ b/src/relax/op/vision/multibox_transform_loc.cc @@ -88,7 +88,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { * skips those N-based relations; other checks (ndim, dtype, loc dim divisible by 4, etc.) * still apply when their inputs are known. */ -Type InferTypeMultiboxTransformLoc(const CallNode* call_node) { +Type InferTypeMultiboxTransformLoc(const CallObj* call_node) { const Call call = ffi::GetRef(call_node); if (RequiresTensorInputNormalization(call)) return Type::Missing(); if (call->args.size() != 3) { @@ -132,12 +132,12 @@ Type InferTypeMultiboxTransformLoc(const CallNode* call_node) { } auto vdev = cls_ty->vdevice; - const auto* cls_shape = cls_ty->shape.as(); - const auto* loc_shape = loc_ty->shape.as(); - const auto* anchor_shape = anchor_ty->shape.as(); + const auto* cls_shape = cls_ty->shape.as(); + const auto* loc_shape = loc_ty->shape.as(); + const auto* anchor_shape = anchor_ty->shape.as(); if (loc_shape != nullptr) { - const auto* loc_dim1 = loc_shape->values[1].as(); + const auto* loc_dim1 = loc_shape->values[1].as(); if (loc_dim1 != nullptr && loc_dim1->value % 4 != 0) { TVM_FFI_VISIT_THROW(ValueError, call) << "multibox_transform_loc: loc_pred.shape[1] must be divisible by 4, got " @@ -146,8 +146,8 @@ Type InferTypeMultiboxTransformLoc(const CallNode* call_node) { } if (cls_shape != nullptr && loc_shape != nullptr) { - const auto* cls_b = cls_shape->values[0].as(); - const auto* loc_b = loc_shape->values[0].as(); + const auto* cls_b = cls_shape->values[0].as(); + const auto* loc_b = loc_shape->values[0].as(); if (cls_b != nullptr && loc_b != nullptr && cls_b->value != loc_b->value) { TVM_FFI_VISIT_THROW(ValueError, call) << "multibox_transform_loc: cls_pred.shape[0] must match loc_pred.shape[0], " @@ -157,12 +157,12 @@ Type InferTypeMultiboxTransformLoc(const CallNode* call_node) { } if (anchor_shape != nullptr) { - const auto* anchor_batch = anchor_shape->values[0].as(); + const auto* anchor_batch = anchor_shape->values[0].as(); if (anchor_batch != nullptr && anchor_batch->value != 1) { TVM_FFI_VISIT_THROW(ValueError, call) << "multibox_transform_loc: anchor.shape[0] must be 1, got " << anchor_batch->value; } - const auto* anchor_last = anchor_shape->values[2].as(); + const auto* anchor_last = anchor_shape->values[2].as(); if (anchor_last != nullptr && anchor_last->value != 4) { TVM_FFI_VISIT_THROW(ValueError, call) << "multibox_transform_loc: anchor.shape[2] must be 4 (ltrb), got " << anchor_last->value; @@ -180,8 +180,8 @@ Type InferTypeMultiboxTransformLoc(const CallNode* call_node) { const auto& num_anchors = cls_shape->values[2]; if (loc_shape != nullptr) { - const auto* num_anchors_imm = num_anchors.as(); - const auto* loc_dim1 = loc_shape->values[1].as(); + const auto* num_anchors_imm = num_anchors.as(); + const auto* loc_dim1 = loc_shape->values[1].as(); if (num_anchors_imm != nullptr && loc_dim1 != nullptr && loc_dim1->value != num_anchors_imm->value * 4) { TVM_FFI_VISIT_THROW(ValueError, call) @@ -191,8 +191,8 @@ Type InferTypeMultiboxTransformLoc(const CallNode* call_node) { } } if (anchor_shape != nullptr) { - const auto* num_anchors_imm = num_anchors.as(); - const auto* anchor_num_anchors = anchor_shape->values[1].as(); + const auto* num_anchors_imm = num_anchors.as(); + const auto* anchor_num_anchors = anchor_shape->values[1].as(); if (num_anchors_imm != nullptr && anchor_num_anchors != nullptr && anchor_num_anchors->value != num_anchors_imm->value) { TVM_FFI_VISIT_THROW(ValueError, call) diff --git a/src/relax/op/vision/nms.cc b/src/relax/op/vision/nms.cc index 290061a066c1..b8fce1852752 100644 --- a/src/relax/op/vision/nms.cc +++ b/src/relax/op/vision/nms.cc @@ -112,7 +112,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { all_class_non_max_suppression); } -Type InferTypeAllClassNMS(const CallNode* call_node) { +Type InferTypeAllClassNMS(const CallObj* call_node) { const Call call = ffi::GetRef(call_node); tvm::ffi::Array input_ty = GetInputTensorType(call); const auto boxes_ty = input_ty[0]; @@ -122,9 +122,9 @@ Type InferTypeAllClassNMS(const CallNode* call_node) { TVM_FFI_ICHECK_EQ(boxes_ty->ndim, 3) << "AllClassNMS input boxes should be 3-D."; TVM_FFI_ICHECK_EQ(scores_ty->ndim, 3) << "AllClassNMS input scores count should be 3-D."; - const auto batch = boxes_ty->shape.as()->values[0]; - const auto num_classes = scores_ty->shape.as()->values[1]; - const auto num_boxes = boxes_ty->shape.as()->values[1]; + const auto batch = boxes_ty->shape.as()->values[0]; + const auto num_classes = scores_ty->shape.as()->values[1]; + const auto num_boxes = boxes_ty->shape.as()->values[1]; auto vdev = input_ty[0]->vdevice; const auto* attrs = call->attrs.as(); @@ -184,7 +184,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { refl::GlobalDef().def("relax.op.vision.get_valid_counts", get_valid_counts); } -Type InferTypeGetValidCounts(const CallNode* call_node) { +Type InferTypeGetValidCounts(const CallObj* call_node) { const Call call = ffi::GetRef(call_node); if (RequiresTensorInputNormalization(call)) return Type::Missing(); if (call->args.size() != 1) { @@ -192,7 +192,7 @@ Type InferTypeGetValidCounts(const CallNode* call_node) { << "get_valid_counts expects 1 argument, got " << call->args.size(); } - const auto* data_ty = GetTypeAs(call->args[0]); + const auto* data_ty = GetTypeAs(call->args[0]); if (data_ty == nullptr) { TVM_FFI_VISIT_THROW(TypeError, call) << "get_valid_counts expects input data to be a Tensor."; } @@ -204,7 +204,7 @@ Type InferTypeGetValidCounts(const CallNode* call_node) { const auto* attrs = call->attrs.as(); TVM_FFI_ICHECK(attrs != nullptr) << "Invalid get_valid_counts attrs"; auto vdev = data_ty->vdevice; - const auto* data_shape = data_ty->shape.as(); + const auto* data_shape = data_ty->shape.as(); if (data_shape == nullptr) { tvm::ffi::Array fields = {TensorType(PrimType::Int(32), /*ndim=*/1, vdev), TensorType(data_ty->dtype, /*ndim=*/3, vdev), @@ -215,7 +215,7 @@ Type InferTypeGetValidCounts(const CallNode* call_node) { auto batch = data_shape->values[0]; auto num_anchors = data_shape->values[1]; auto elem_length = data_shape->values[2]; - const auto* elem_length_imm = elem_length.as(); + const auto* elem_length_imm = elem_length.as(); if (elem_length_imm != nullptr) { if (attrs->score_index < 0 || attrs->score_index >= elem_length_imm->value) { TVM_FFI_VISIT_THROW(ValueError, call) @@ -279,9 +279,9 @@ Type InferTypeNMS(const Call& call, const BlockBuilder& ctx) { << "non_max_suppression expects 3 arguments, got " << call->args.size(); } - const auto* data_ty = GetTypeAs(call->args[0]); - const auto* valid_count_ty = GetTypeAs(call->args[1]); - const auto* indices_ty = GetTypeAs(call->args[2]); + const auto* data_ty = GetTypeAs(call->args[0]); + const auto* valid_count_ty = GetTypeAs(call->args[1]); + const auto* indices_ty = GetTypeAs(call->args[2]); if (data_ty == nullptr) { TVM_FFI_VISIT_THROW(TypeError, call) << "non_max_suppression expects input data to be a Tensor."; @@ -315,9 +315,9 @@ Type InferTypeNMS(const Call& call, const BlockBuilder& ctx) { << "non_max_suppression expects indices to have dtype int32, got " << indices_ty->dtype; } - const auto* data_shape = data_ty->shape.as(); - const auto* valid_count_shape = valid_count_ty->shape.as(); - const auto* indices_shape = indices_ty->shape.as(); + const auto* data_shape = data_ty->shape.as(); + const auto* valid_count_shape = valid_count_ty->shape.as(); + const auto* indices_shape = indices_ty->shape.as(); if (data_shape != nullptr) { sym::Analyzer analyzer = ctx->GetAnalyzer(); PrimExpr batch = data_shape->values[0]; @@ -344,7 +344,7 @@ Type InferTypeNMS(const Call& call, const BlockBuilder& ctx) { TVM_FFI_ICHECK(attrs != nullptr) << "Invalid non_max_suppression attrs"; auto vdev = data_ty->vdevice; if (data_shape != nullptr) { - const auto* elem_length_imm = data_shape->values[2].as(); + const auto* elem_length_imm = data_shape->values[2].as(); if (elem_length_imm != nullptr) { const ffi::BigInt& elem_length = elem_length_imm->value; if (attrs->score_index < 0 || attrs->score_index >= elem_length) { @@ -401,7 +401,7 @@ Type InferTypeNMS(const Call& call, const BlockBuilder& ctx) { } // Returns modified data tensor with the same shape as input. - if (const auto* data_shape = data_ty->shape.as()) { + if (const auto* data_shape = data_ty->shape.as()) { return TensorType(ffi::GetRef(data_shape), data_ty->dtype, vdev); } return TensorType(data_ty->dtype, /*ndim=*/3, vdev); diff --git a/src/relax/op/vision/roi_align.cc b/src/relax/op/vision/roi_align.cc index a2e3205b2d15..544756e24b07 100644 --- a/src/relax/op/vision/roi_align.cc +++ b/src/relax/op/vision/roi_align.cc @@ -80,7 +80,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { refl::GlobalDef().def("relax.op.vision.roi_align", roi_align); } -Type InferTypeROIAlign(const CallNode* call_node) { +Type InferTypeROIAlign(const CallObj* call_node) { const Call call = ffi::GetRef(call_node); if (RequiresTensorInputNormalization(call)) return Type::Missing(); if (call->args.size() != 2) { @@ -89,8 +89,8 @@ Type InferTypeROIAlign(const CallNode* call_node) { << call->args.size(); } - const auto* data_ty = GetTypeAs(call->args[0]); - const auto* rois_ty = GetTypeAs(call->args[1]); + const auto* data_ty = GetTypeAs(call->args[0]); + const auto* rois_ty = GetTypeAs(call->args[1]); if (data_ty == nullptr) { TVM_FFI_VISIT_THROW(TypeError, call) << "ROIAlign expects the input data to be a Tensor, while the given data is " @@ -123,9 +123,9 @@ Type InferTypeROIAlign(const CallNode* call_node) { << "ROIAlign only supports avg and max mode, but got " << attrs->mode; } - const auto* rois_shape = rois_ty->shape.as(); + const auto* rois_shape = rois_ty->shape.as(); if (rois_shape != nullptr) { - const auto* last_dim = rois_shape->values[1].as(); + const auto* last_dim = rois_shape->values[1].as(); if (last_dim != nullptr && last_dim->value != 5) { TVM_FFI_VISIT_THROW(ValueError, call) << "ROIAlign expects rois to have shape (num_roi, 5), but got last " @@ -134,11 +134,11 @@ Type InferTypeROIAlign(const CallNode* call_node) { } } - if (data_ty->shape.as() == nullptr || rois_shape == nullptr) { + if (data_ty->shape.as() == nullptr || rois_shape == nullptr) { return TensorType(data_ty->dtype, 4, data_ty->vdevice); } - ffi::Array data_shape = data_ty->shape.as()->values; + ffi::Array data_shape = data_ty->shape.as()->values; ffi::Array out_shape; if (attrs->layout == "NCHW") { out_shape = {rois_shape->values[0], data_shape[1], IntImm::Int32(attrs->pooled_size[0]), diff --git a/src/relax/op/vision/roi_pool.cc b/src/relax/op/vision/roi_pool.cc index 6c322b4b0eb9..b807a504f2d4 100644 --- a/src/relax/op/vision/roi_pool.cc +++ b/src/relax/op/vision/roi_pool.cc @@ -69,7 +69,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { refl::GlobalDef().def("relax.op.vision.roi_pool", roi_pool); } -Type InferTypeROIPool(const CallNode* call_node) { +Type InferTypeROIPool(const CallObj* call_node) { const Call call = ffi::GetRef(call_node); if (call->args.size() != 2) { TVM_FFI_VISIT_THROW(ValueError, call) @@ -77,8 +77,8 @@ Type InferTypeROIPool(const CallNode* call_node) { << call->args.size(); } - const auto* data_ty = GetTypeAs(call->args[0]); - const auto* rois_ty = GetTypeAs(call->args[1]); + const auto* data_ty = GetTypeAs(call->args[0]); + const auto* rois_ty = GetTypeAs(call->args[1]); if (data_ty == nullptr) { TVM_FFI_VISIT_THROW(TypeError, call) << "ROIPool expects the input data to be a Tensor, while the given data is " @@ -107,9 +107,9 @@ Type InferTypeROIPool(const CallNode* call_node) { << "ROIPool only supports NCHW layout, but got " << attrs->layout; } - const auto* rois_shape = rois_ty->shape.as(); + const auto* rois_shape = rois_ty->shape.as(); if (rois_shape != nullptr) { - const auto* last_dim = rois_shape->values[1].as(); + const auto* last_dim = rois_shape->values[1].as(); if (last_dim != nullptr && last_dim->value != 5) { TVM_FFI_VISIT_THROW(ValueError, call) << "ROIPool expects rois to have shape (num_roi, 5), but got last " @@ -118,11 +118,11 @@ Type InferTypeROIPool(const CallNode* call_node) { } } - if (data_ty->shape.as() == nullptr || rois_shape == nullptr) { + if (data_ty->shape.as() == nullptr || rois_shape == nullptr) { return TensorType(data_ty->dtype, 4, data_ty->vdevice); } - ffi::Array data_shape = data_ty->shape.as()->values; + ffi::Array data_shape = data_ty->shape.as()->values; ffi::Array out_shape = {rois_shape->values[0], data_shape[1], IntImm::Int32(attrs->pooled_size[0]), IntImm::Int32(attrs->pooled_size[1])}; diff --git a/src/relax/script/ir_builder/distributed.cc b/src/relax/script/ir_builder/distributed.cc index 642217512ce9..77c52d0e67bc 100644 --- a/src/relax/script/ir_builder/distributed.cc +++ b/src/relax/script/ir_builder/distributed.cc @@ -30,7 +30,7 @@ namespace tvm { namespace relax { Expr MakeCallTIRDist(Expr func, Tuple args, ffi::Array out_ty_list) { for (const distributed::DTensorType& ty : out_ty_list) { - const auto* shape = ty->tensor_ty->shape.as(); + const auto* shape = ty->tensor_ty->shape.as(); TVM_FFI_ICHECK(shape != nullptr) << "out_ty of call_tir should have defined ShapeExpr as shape. " "However, one given type information is " diff --git a/src/relax/script/ir_builder/frame.cc b/src/relax/script/ir_builder/frame.cc index d43c8e6a5466..8c789930d353 100644 --- a/src/relax/script/ir_builder/frame.cc +++ b/src/relax/script/ir_builder/frame.cc @@ -148,7 +148,7 @@ void FunctionFrameObj::ExitWithScope() { tvm::relax::FuncType(signature->params.value().Map(bind_param), bind_param(signature->ret), signature->purity, signature->loc); } - if (local_var.value()->IsInstance()) { + if (local_var.value()->IsInstance()) { // The self-reference is captured by the nested function, so it must // survive the dataflow region. Region finalization rewrites every use // together, including recursive calls, to the same ordinary Var. @@ -327,7 +327,7 @@ void IfFrameObj::ExitWithScope() { auto binding_var = Emit(body, std::nullopt); var = binding_var; // Finalization uses the frame's already-composed location, never the exit context. - if (!loc.as()) { + if (!loc.as()) { CheckBindingBlockFrameExistAndUnended()->binding_locs.Set(binding_var, loc); } binding_var->loc = loc; diff --git a/src/relax/script/ir_builder/ir.cc b/src/relax/script/ir_builder/ir.cc index b9847813ab02..0fe2cc23bcae 100644 --- a/src/relax/script/ir_builder/ir.cc +++ b/src/relax/script/ir_builder/ir.cc @@ -70,13 +70,13 @@ TVM_FFI_STATIC_INIT_BLOCK() { ///////////////////////////////// Vars ////////////////////////////////// -using tvm::script::ir_builder::details::Namer; +using tvm::script::ir_builder::detail::Namer; TVM_FFI_STATIC_INIT_BLOCK() { - Namer::vtable().SetDispatch( + Namer::vtable().SetDispatch( [](const ffi::ObjectRef& node, ffi::String name) -> void { - using tvm::relax::DataflowVarNode; - DataflowVarNode* var = const_cast(node.as()); + using tvm::relax::DataflowVarObj; + DataflowVarObj* var = const_cast(node.as()); var->name = name; }); } @@ -178,7 +178,7 @@ void FuncRetValue(const tvm::relax::Expr& value) { if (auto opt = ir_builder->GetLastFrame()) { auto block_frame = opt.value(); for (const auto& var : tvm::relax::FreeVars(normalized_value)) { - if (var->IsInstance()) { + if (var->IsInstance()) { block_frame->output_vars.push_back(var); } } @@ -302,10 +302,10 @@ namespace { tvm::Var RecordBindingLoc(tvm::Var var, const Location& name_loc, const Location& statement_loc) { Location loc = IRBuilder::Current()->ResolveLoc(statement_loc); - if (!loc.as()) { + if (!loc.as()) { CheckBindingBlockFrameExistAndUnended()->binding_locs.Set(var, loc); } - var->loc = name_loc.as() ? loc : name_loc; + var->loc = name_loc.as() ? loc : name_loc; return var; } diff --git a/src/relax/script/ir_builder/utils.h b/src/relax/script/ir_builder/utils.h index 1c738473d16e..cfb23f3fcbd3 100644 --- a/src/relax/script/ir_builder/utils.h +++ b/src/relax/script/ir_builder/utils.h @@ -96,7 +96,7 @@ inline tvm::relax::SeqExpr GetSeqExprForBranch(const SeqExprFrame& frame, ffi::S // variable name. tvm::relax::Binding last_binding = last_block->bindings.back(); - TVM_FFI_ICHECK(!last_binding->var->IsInstance()) + TVM_FFI_ICHECK(!last_binding->var->IsInstance()) << "A non-dataflow var is expected in the last binding of '" << method << "'."; *var_name = last_binding->var->name; @@ -111,21 +111,21 @@ inline tvm::relax::SeqExpr GetSeqExprForBranch(const SeqExprFrame& frame, ffi::S tvm::relax::GetType(last_binding->var), last_binding->var->loc); tvm::relax::Expr body = new_var; - const auto* var_binding = last_binding.as(); + const auto* var_binding = last_binding.as(); if (var_binding && tvm::relax::IsLeafOrTuple(var_binding->value)) { body = var_binding->value; } else if (var_binding) { last_block_bindings.push_back( tvm::relax::VarBinding(new_var, var_binding->value, var_binding->loc)); - } else if (const auto* match_cast = last_binding.as()) { + } else if (const auto* match_cast = last_binding.as()) { last_block_bindings.push_back( tvm::relax::MatchCast(new_var, match_cast->value, match_cast->ty, match_cast->loc)); } else { TVM_FFI_CHECK(false, TypeError) << "Unsupported binding type: " << last_binding->GetTypeKey(); } - new_blocks.push_back(last_block->IsInstance() + new_blocks.push_back(last_block->IsInstance() ? tvm::relax::DataflowBlock(last_block_bindings, last_block->loc) : tvm::relax::BindingBlock(last_block_bindings, last_block->loc)); diff --git a/src/relax/script/printer/binding.cc b/src/relax/script/printer/binding.cc index cb6113d16dd0..c4e6a888670c 100644 --- a/src/relax/script/printer/binding.cc +++ b/src/relax/script/printer/binding.cc @@ -36,7 +36,7 @@ namespace { ffi::Optional MatchCastDocTranslate(DocTranslatorObj* d, ffi::AnyView input, const ffi::Object* destination) { const auto* binding = - ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(input); + ffi::detail::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(input); TVM_FFI_CHECK(destination == nullptr, TypeError) << "printer statement-only node cannot fulfill a destination"; ExprDoc rhs = @@ -53,7 +53,7 @@ ffi::Optional MatchCastDocTranslate(DocTranslatorObj* d, ffi::AnyView i } TVM_FFI_STATIC_INIT_BLOCK() { - ffi::reflection::TypeAttrDef().attr( + ffi::reflection::TypeAttrDef().attr( tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&MatchCastDocTranslate>()); } @@ -61,11 +61,10 @@ TVM_FFI_STATIC_INIT_BLOCK() { ffi::Optional VarBindingDocTranslate(DocTranslatorObj* d, ffi::AnyView input, const ffi::Object* destination) { const auto* binding = - ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck( - input); + ffi::detail::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(input); TVM_FFI_CHECK(destination == nullptr, TypeError) << "printer statement-only node cannot fulfill a destination"; - if (auto func = binding->value.as()) { + if (auto func = binding->value.as()) { IdDoc lhs = VarDoc(d, binding->var); d->Translate(binding->value); FunctionDoc function = d->CurrentScopeDocs().back().as_or_throw(); @@ -88,14 +87,14 @@ ffi::Optional VarBindingDocTranslate(DocTranslatorObj* d, ffi::AnyView if (!rhs.has_value()) return std::nullopt; ffi::Optional annotation = std::nullopt; bool infer_vdevice = false; - if (auto tensor = binding->var->ty.as(); + if (auto tensor = binding->var->ty.as(); tensor && tensor->vdevice.has_value()) { - if (auto call = binding->value.as()) { + if (auto call = binding->value.as()) { if (auto op = call->op.as()) infer_vdevice = op.value()->name == "relax.to_vdevice"; } } bool show_all_ty = d->GetExtraConfig("relax.show_all_ty", true); - const auto* call = binding->value.as(); + const auto* call = binding->value.as(); bool output_type_argument = call && call->ty_args.size() == 1 && ffi::StructuralEqual()(binding->var->ty, call->ty_args[0]); ffi::Optional inferred = std::nullopt; @@ -116,7 +115,7 @@ ffi::Optional VarBindingDocTranslate(DocTranslatorObj* d, ffi::AnyView // Primitive aliases need their annotation. Without context-free inference, // keep the binding type and let the parser's deferred inference use it. bool elide_annotation = infer_vdevice || explicit_output_type || - (!show_all_ty && !binding->var->ty.as() && inferable); + (!show_all_ty && !binding->var->ty.as() && inferable); if (!binding->var->ty.as().has_value() && !elide_annotation) { annotation = d->Translate(binding->var->ty).value(); } @@ -125,7 +124,7 @@ ffi::Optional VarBindingDocTranslate(DocTranslatorObj* d, ffi::AnyView } TVM_FFI_STATIC_INIT_BLOCK() { - ffi::reflection::TypeAttrDef().attr( + ffi::reflection::TypeAttrDef().attr( tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&VarBindingDocTranslate>()); } @@ -133,7 +132,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { ffi::Optional IfDocTranslate(DocTranslatorObj* d, ffi::AnyView input, const ffi::Object* destination) { const auto* branch = - ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(input); + ffi::detail::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(input); // A value context needs a usable expression after the conditional. A // destination supplied by a binding is completed directly on both arms. ffi::Optional temporary = std::nullopt; @@ -142,9 +141,9 @@ ffi::Optional IfDocTranslate(DocTranslatorObj* d, ffi::AnyView input, VarDoc(d, temporary.value()); destination = temporary.value().get(); } - TVM_FFI_CHECK(destination->IsInstance(), TypeError) + TVM_FFI_CHECK(destination->IsInstance(), TypeError) << "printer Relax IfExpr destination must be a Var"; - Var var = ffi::GetRef(static_cast(destination)); + Var var = ffi::GetRef(static_cast(destination)); ffi::Optional lhs = VarDoc(d, var); ffi::Optional annotation = std::nullopt; if (!var->ty.as().has_value()) annotation = d->Translate(var->ty).value(); @@ -156,7 +155,7 @@ ffi::Optional IfDocTranslate(DocTranslatorObj* d, ffi::AnyView input, } TVM_FFI_STATIC_INIT_BLOCK() { - ffi::reflection::TypeAttrDef().attr( + ffi::reflection::TypeAttrDef().attr( tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&IfDocTranslate>()); } diff --git a/src/relax/script/printer/call.cc b/src/relax/script/printer/call.cc index 64bd57ebe790..cbfb8a7d69ce 100644 --- a/src/relax/script/printer/call.cc +++ b/src/relax/script/printer/call.cc @@ -35,16 +35,16 @@ namespace { // These semantic wrappers convert bare Python scalar tuple fields to 64-bit // values. Use raw construction when ordinary tuple syntax would change the IR. bool InlineTupleNeedsRawCall(DocTranslatorObj* d, const Expr& expr) { - if (const auto* tuple = expr.as()) { + if (const auto* tuple = expr.as()) { for (const Expr& field : tuple->fields) { if (InlineTupleNeedsRawCall(d, field)) return true; } - } else if (const auto* imm = expr.as()) { + } else if (const auto* imm = expr.as()) { DLDataType dtype = imm->ty.as_or_throw()->dtype; auto implicit = ffi::StringToDLDataType(d->GetExtraConfig("ir.int_dtype", "int32")); return dtype == implicit && dtype != ffi::StringToDLDataType("int64"); - } else if (const auto* imm = expr.as()) { + } else if (const auto* imm = expr.as()) { DLDataType dtype = imm->ty.as_or_throw()->dtype; auto implicit = ffi::StringToDLDataType(d->GetExtraConfig("ir.float_dtype", "void")); @@ -56,7 +56,7 @@ bool InlineTupleNeedsRawCall(DocTranslatorObj* d, const Expr& expr) { ffi::Optional InlineTupleCallDocTranslate(DocTranslatorObj* d, ffi::AnyView input, const ffi::Object*) { const auto* call = - ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(input); + ffi::detail::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(input); if (call->args.size() > 1 && InlineTupleNeedsRawCall(d, call->args[1])) return RawCallDocTranslate(d, call); if (auto doc = StandardCallDocTranslate(d, call)) return doc; @@ -76,7 +76,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { ffi::Optional LogicalCallDocTranslate(DocTranslatorObj* d, ffi::AnyView input, const ffi::Object*) { const auto* call = - ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(input); + ffi::detail::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(input); for (const Expr& arg : call->args) { if (arg->ty.as()) return RawCallDocTranslate(d, call); } @@ -89,10 +89,10 @@ ffi::Optional LogicalCallDocTranslate(DocTranslatorObj* d, ffi::AnyView ffi::Optional VMAllocCallDocTranslate(DocTranslatorObj* d, ffi::AnyView input, const ffi::Object*) { const auto* call = - ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(input); + ffi::detail::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(input); auto needs_raw = [&](size_t index) { if (index >= call->args.size()) return false; - const auto* imm = call->args[index].as(); + const auto* imm = call->args[index].as(); if (!imm) return false; DLDataType dtype = imm->ty.as_or_throw()->dtype; auto implicit = @@ -119,18 +119,18 @@ TVM_FFI_STATIC_INIT_BLOCK() { ffi::Optional CallDPSPackedDocTranslate(DocTranslatorObj* d, ffi::AnyView input, const ffi::Object*) { const auto* call = - ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(input); - if (call->attrs.has_value() || call->args.size() != 2 || !call->args[1].as() || + ffi::detail::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(input); + if (call->attrs.has_value() || call->args.size() != 2 || !call->args[1].as() || call->ty_args.size() != 1 || InlineTupleNeedsRawCall(d, call->args[1])) { return RawCallDocTranslate(d, call); } ExprDoc callee = d->Translate(call->args[0]).value(); - if (const auto* packed = call->args[0].as()) { + if (const auto* packed = call->args[0].as()) { if (!ffi::StructuralEqual()(packed->ty, relax::ExternFunc(packed->global_symbol)->ty)) { return RawCallDocTranslate(d, call); } callee = LiteralDoc::Str(packed->global_symbol, std::nullopt); - } else if (call->args[0].as()) { + } else if (call->args[0].as()) { // A bare Python string now constructs ExternFunc, so retain this explicit IR value. callee = d->WithExtraStateScope("ir.explicit_value_type", true, [&] { return d->Translate(call->args[0]).value(); }); diff --git a/src/relax/script/printer/dependent_type.cc b/src/relax/script/printer/dependent_type.cc index b12683dbd815..6ac2aee33d7f 100644 --- a/src/relax/script/printer/dependent_type.cc +++ b/src/relax/script/printer/dependent_type.cc @@ -35,7 +35,7 @@ namespace { ffi::Optional ShapeTypeDocTranslate(DocTranslatorObj* d, ffi::AnyView input, const ffi::Object*) { const auto* ty = - ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(input); + ffi::detail::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(input); if (ty->values.has_value()) { ffi::Array values; for (const PrimExpr& value : ty->values.value()) values.push_back(RelaxShapeDim(d, value)); @@ -46,7 +46,7 @@ ffi::Optional ShapeTypeDocTranslate(DocTranslatorObj* d, ffi::AnyView i } TVM_FFI_STATIC_INIT_BLOCK() { - ffi::reflection::TypeAttrDef().attr( + ffi::reflection::TypeAttrDef().attr( tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&ShapeTypeDocTranslate>()); } @@ -54,13 +54,12 @@ TVM_FFI_STATIC_INIT_BLOCK() { ffi::Optional TensorTypeDocTranslate(DocTranslatorObj* d, ffi::AnyView input, const ffi::Object*) { const auto* ty = - ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck( - input); + ffi::detail::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(input); ffi::Array args; ffi::Array keys; ffi::Array values; if (ty->shape.has_value()) { - if (auto shape = ty->shape.value().as()) { + if (auto shape = ty->shape.value().as()) { ffi::Array dimensions; for (const PrimExpr& dim : shape->values) dimensions.push_back(RelaxShapeDim(d, dim)); args.push_back(TupleDoc(dimensions)); @@ -88,7 +87,7 @@ ffi::Optional TensorTypeDocTranslate(DocTranslatorObj* d, ffi::AnyView } TVM_FFI_STATIC_INIT_BLOCK() { - ffi::reflection::TypeAttrDef().attr( + ffi::reflection::TypeAttrDef().attr( tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&TensorTypeDocTranslate>()); } @@ -96,11 +95,11 @@ TVM_FFI_STATIC_INIT_BLOCK() { ffi::Optional RelaxFuncTypeDocTranslate(DocTranslatorObj* d, ffi::AnyView input, const ffi::Object*) { const auto* ty = - ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(input); + ffi::detail::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(input); if (!ty->params.has_value()) { ffi::Array keys; ffi::Array values; - if (!ty->ret.as()) { + if (!ty->ret.as()) { keys.push_back("ret"); values.push_back(d->Translate(ty->ret).value()); } @@ -124,7 +123,7 @@ ffi::Optional RelaxFuncTypeDocTranslate(DocTranslatorObj* d, ffi::AnyVi } TVM_FFI_STATIC_INIT_BLOCK() { - ffi::reflection::TypeAttrDef().attr( + ffi::reflection::TypeAttrDef().attr( tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&RelaxFuncTypeDocTranslate>()); } diff --git a/src/relax/script/printer/distributed.cc b/src/relax/script/printer/distributed.cc index d3a3b58246eb..a79ed5c2c43e 100644 --- a/src/relax/script/printer/distributed.cc +++ b/src/relax/script/printer/distributed.cc @@ -32,28 +32,28 @@ namespace { ffi::Optional PlacementDocTranslate(DocTranslatorObj*, ffi::AnyView input, const ffi::Object*) { - const auto* placement = ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck< - const relax::distributed::PlacementNode>(input); + const auto* placement = ffi::detail::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck< + const relax::distributed::PlacementObj>(input); return LiteralDoc::Str(placement->ToString(), std::nullopt); } TVM_FFI_STATIC_INIT_BLOCK() { - ffi::reflection::TypeAttrDef().attr( + ffi::reflection::TypeAttrDef().attr( tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&PlacementDocTranslate>()); } ffi::Optional DTensorTypeDocTranslate(DocTranslatorObj* d, ffi::AnyView input, const ffi::Object*) { - const auto* ty = ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck< - const relax::distributed::DTensorTypeNode>(input); + const auto* ty = ffi::detail::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck< + const relax::distributed::DTensorTypeObj>(input); ffi::Array args; ffi::Array keys; ffi::Array values; const relax::TensorType& tensor = ty->tensor_ty; bool require_keywords = !tensor->shape.has_value(); if (tensor->shape.has_value()) { - if (auto shape = tensor->shape.value().as()) { + if (auto shape = tensor->shape.value().as()) { ffi::Array dimensions; for (const PrimExpr& dim : shape->values) dimensions.push_back(RelaxShapeDim(d, dim)); args.push_back(TupleDoc(dimensions)); @@ -94,15 +94,15 @@ ffi::Optional DTensorTypeDocTranslate(DocTranslatorObj* d, ffi::AnyView } TVM_FFI_STATIC_INIT_BLOCK() { - ffi::reflection::TypeAttrDef().attr( + ffi::reflection::TypeAttrDef().attr( tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&DTensorTypeDocTranslate>()); } ffi::Optional DeviceMeshDocTranslate(DocTranslatorObj* d, ffi::AnyView input, const ffi::Object*) { - const auto* mesh = ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck< - const relax::distributed::DeviceMeshNode>(input); + const auto* mesh = ffi::detail::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck< + const relax::distributed::DeviceMeshObj>(input); ffi::Array dimensions; for (int64_t value : mesh->shape) { dimensions.push_back(LiteralDoc::Int(value, std::nullopt)); @@ -119,7 +119,7 @@ ffi::Optional DeviceMeshDocTranslate(DocTranslatorObj* d, ffi::AnyView } TVM_FFI_STATIC_INIT_BLOCK() { - ffi::reflection::TypeAttrDef().attr( + ffi::reflection::TypeAttrDef().attr( tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&DeviceMeshDocTranslate>()); } diff --git a/src/relax/script/printer/expr.cc b/src/relax/script/printer/expr.cc index dc1c5178d718..f47dc9deec6f 100644 --- a/src/relax/script/printer/expr.cc +++ b/src/relax/script/printer/expr.cc @@ -42,7 +42,7 @@ namespace { ffi::Optional GenericConstDocTranslate(DocTranslatorObj* d, ffi::AnyView input, const ffi::Object*) { const auto* constant = - ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(input); + ffi::detail::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(input); if (auto dtype = constant->value.as()) { return LiteralDoc::DataType(dtype.value(), std::nullopt); } @@ -51,7 +51,7 @@ ffi::Optional GenericConstDocTranslate(DocTranslatorObj* d, ffi::AnyVie const ffi::Tensor& data = tensor.value(); DLDataType dtype = data.dtype(); auto print_const = [&](ExprDoc value) -> ExprDoc { - if (constant->ty.as()) { + if (constant->ty.as()) { return NamespaceDoc("relax")->Attr("dist")->Attr("const")->Call( {value, d->Translate(constant->ty).value()}); } @@ -104,7 +104,7 @@ ffi::Optional GenericConstDocTranslate(DocTranslatorObj* d, ffi::AnyVie } TVM_FFI_STATIC_INIT_BLOCK() { - ffi::reflection::TypeAttrDef().attr( + ffi::reflection::TypeAttrDef().attr( tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&GenericConstDocTranslate>()); } @@ -112,20 +112,20 @@ TVM_FFI_STATIC_INIT_BLOCK() { ffi::Optional ShapeExprDocTranslate(DocTranslatorObj* d, ffi::AnyView input, const ffi::Object*) { const auto* shape = - ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(input); + ffi::detail::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(input); ffi::Array dimensions; for (const PrimExpr& dim : shape->values) dimensions.push_back(RelaxShapeDim(d, dim)); return NamespaceDoc("relax")->Attr("shape")->Call({ListDoc(dimensions)}); } TVM_FFI_STATIC_INIT_BLOCK() { - ffi::reflection::TypeAttrDef().attr( + ffi::reflection::TypeAttrDef().attr( tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&ShapeExprDocTranslate>()); } TVM_FFI_STATIC_INIT_BLOCK() { - ffi::reflection::TypeAttrDef().attr( + ffi::reflection::TypeAttrDef().attr( tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&VarDocTranslate>()); } @@ -142,23 +142,23 @@ TVM_FFI_STATIC_INIT_BLOCK() { [](const ffi::ObjectRef& obj, const PrinterConfig& config) { return tvm::Script(obj, config); }); - RegisterScriptRepr(); - RegisterScriptRepr(); - RegisterScriptRepr(); - RegisterScriptRepr(); - RegisterScriptRepr(); - RegisterScriptRepr(); - RegisterScriptRepr(); - RegisterScriptRepr(); - RegisterScriptRepr(); - RegisterScriptRepr(); - RegisterScriptRepr(); - RegisterScriptRepr(); - RegisterScriptRepr(); - RegisterScriptRepr(); - RegisterScriptRepr(); - RegisterScriptRepr(); - RegisterScriptRepr(); + RegisterScriptRepr(); + RegisterScriptRepr(); + RegisterScriptRepr(); + RegisterScriptRepr(); + RegisterScriptRepr(); + RegisterScriptRepr(); + RegisterScriptRepr(); + RegisterScriptRepr(); + RegisterScriptRepr(); + RegisterScriptRepr(); + RegisterScriptRepr(); + RegisterScriptRepr(); + RegisterScriptRepr(); + RegisterScriptRepr(); + RegisterScriptRepr(); + RegisterScriptRepr(); + RegisterScriptRepr(); } } // namespace tvm diff --git a/src/relax/script/printer/function.cc b/src/relax/script/printer/function.cc index 78014339e799..d6735ab48709 100644 --- a/src/relax/script/printer/function.cc +++ b/src/relax/script/printer/function.cc @@ -40,7 +40,7 @@ namespace { ffi::Optional FunctionDocTranslate(DocTranslatorObj* d, ffi::AnyView input, const ffi::Object* destination) { const auto* func = - ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(input); + ffi::detail::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(input); TVM_FFI_CHECK(destination == nullptr, TypeError) << "printer statement-only node cannot fulfill a destination"; @@ -105,7 +105,7 @@ ffi::Optional FunctionDocTranslate(DocTranslatorObj* d, ffi::AnyView in } for (const relax::BindingBlock& block : func->body->blocks) d->Translate(block); ExprDoc result = d->Translate(func->body->body).value(); - if (auto integer = func->body->body.as(); + if (auto integer = func->body->body.as(); integer && ffi::StructuralEqual()(integer->ty, func->ret_ty)) { // The declared primitive return type reconstructs this literal's dtype. result = LiteralDoc::Int(ffi::GetRef(integer), std::nullopt); @@ -127,7 +127,7 @@ ffi::Optional FunctionDocTranslate(DocTranslatorObj* d, ffi::AnyView in auto pending = d->GetImplicitDefs(); if (!module.has_value()) { for (const auto& [var, signature_id] : signature_candidates) { - if (!pending.count(var) || !var->ty.as()) continue; + if (!pending.count(var) || !var->ty.as()) continue; IdDoc id = VarDoc(d, var); if (var->ty.as_or_throw()->dtype == (DLDataType{kDLInt, 32, 1})) { function->type_params.push_back(id); @@ -146,7 +146,7 @@ ffi::Optional FunctionDocTranslate(DocTranslatorObj* d, ffi::AnyView in TVM_FFI_STATIC_INIT_BLOCK() { RegisterNamespaceAlias("relax.prefix", "R"); - ffi::reflection::TypeAttrDef().attr( + ffi::reflection::TypeAttrDef().attr( tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&FunctionDocTranslate>()); } @@ -154,17 +154,16 @@ TVM_FFI_STATIC_INIT_BLOCK() { ffi::Optional ExternFuncDocTranslate(DocTranslatorObj* d, ffi::AnyView input, const ffi::Object*) { const auto* func = - ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck( - input); - if (!func->ty.as()) { + ffi::detail::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(input); + if (!func->ty.as()) { // The ExternFunc constructor requires a FuncType; retain diagnostic IR in metadata. return AddMetadata(d, ffi::GetRef(func)); } ffi::Array args = {LiteralDoc::Str(func->global_symbol, std::nullopt)}; if (!ffi::StructuralEqual()(func->ty, relax::ExternFunc(func->global_symbol)->ty)) { ExprDoc type = d->Translate(func->ty).value(); - if (auto opaque = func->ty.as(); - opaque && !opaque->params.has_value() && opaque->ret.as() && !opaque->purity && + if (auto opaque = func->ty.as(); + opaque && !opaque->params.has_value() && opaque->ret.as() && !opaque->purity && !opaque->derive_func.has_value()) { type = type->Call({}); } @@ -174,7 +173,7 @@ ffi::Optional ExternFuncDocTranslate(DocTranslatorObj* d, ffi::AnyView } TVM_FFI_STATIC_INIT_BLOCK() { - ffi::reflection::TypeAttrDef().attr( + ffi::reflection::TypeAttrDef().attr( tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&ExternFuncDocTranslate>()); } diff --git a/src/relax/script/printer/global_info.cc b/src/relax/script/printer/global_info.cc index a60e729ea2e5..21055d284d84 100644 --- a/src/relax/script/printer/global_info.cc +++ b/src/relax/script/printer/global_info.cc @@ -68,7 +68,7 @@ namespace { ffi::Optional VDeviceDocTranslate(DocTranslatorObj* d, ffi::AnyView input, const ffi::Object*) { const auto* device = - ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(input); + ffi::detail::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(input); return NamespaceDoc("relax")->Attr("vdevice")->Call( {d->Translate(device->target).value()}, {"vdevice_id", "memory_scope"}, {LiteralDoc::Int(device->vdevice_id, std::nullopt), @@ -76,7 +76,7 @@ ffi::Optional VDeviceDocTranslate(DocTranslatorObj* d, ffi::AnyView inp } TVM_FFI_STATIC_INIT_BLOCK() { - ffi::reflection::TypeAttrDef().attr( + ffi::reflection::TypeAttrDef().attr( tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&VDeviceDocTranslate>()); } @@ -87,7 +87,7 @@ ffi::Optional DummyGlobalInfoDocTranslate(DocTranslatorObj*, ffi::AnyVi } TVM_FFI_STATIC_INIT_BLOCK() { - ffi::reflection::TypeAttrDef().attr( + ffi::reflection::TypeAttrDef().attr( tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&DummyGlobalInfoDocTranslate>()); } diff --git a/src/relax/script/printer/region.cc b/src/relax/script/printer/region.cc index 57e22c41016b..9a97c509618b 100644 --- a/src/relax/script/printer/region.cc +++ b/src/relax/script/printer/region.cc @@ -33,14 +33,14 @@ namespace { ffi::Optional BindingVar(const ffi::Object* destination) { if (!destination) return std::nullopt; - TVM_FFI_CHECK(destination->IsInstance(), TypeError) + TVM_FFI_CHECK(destination->IsInstance(), TypeError) << "printer binding destination must be a Var"; - return ffi::GetRef(static_cast(destination)); + return ffi::GetRef(static_cast(destination)); } } // namespace -ffi::Array RelaxSeqBody(DocTranslatorObj* d, const relax::SeqExprNode* seq, +ffi::Array RelaxSeqBody(DocTranslatorObj* d, const relax::SeqExprObj* seq, ffi::Optional destination, ffi::Optional annotation, const ffi::Object* destination_object) { @@ -69,12 +69,12 @@ namespace { ffi::Optional SeqExprDocTranslate(DocTranslatorObj* d, ffi::AnyView input, const ffi::Object* destination) { const auto* seq = - ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(input); + ffi::detail::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(input); for (const relax::BindingBlock& block : seq->blocks) d->Translate(block); if (seq->blocks.empty()) { if (auto var = seq->body.as()) { d->VarGetOrAllocId(var.value(), false); - if (d->GetImplicitDefs().count(var.value()) && !var.value()->ty.as()) { + if (d->GetImplicitDefs().count(var.value()) && !var.value()->ty.as()) { IdDoc id = VarDoc(d, var.value()); d->Emit(AssignDoc(id, std::nullopt, d->Translate(var.value()->ty).value()), ffi::GetRef(var.value().get())); @@ -85,7 +85,7 @@ ffi::Optional SeqExprDocTranslate(DocTranslatorObj* d, ffi::AnyView inp } TVM_FFI_STATIC_INIT_BLOCK() { - ffi::reflection::TypeAttrDef().attr( + ffi::reflection::TypeAttrDef().attr( tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&SeqExprDocTranslate>()); } @@ -93,7 +93,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { ffi::Optional BindingBlockDocTranslate(DocTranslatorObj* d, ffi::AnyView input, const ffi::Object* destination) { const auto* block = - ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck( + ffi::detail::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck( input); TVM_FFI_CHECK(destination == nullptr, TypeError) << "printer statement-only node cannot fulfill a destination"; @@ -102,7 +102,7 @@ ffi::Optional BindingBlockDocTranslate(DocTranslatorObj* d, ffi::AnyVie } TVM_FFI_STATIC_INIT_BLOCK() { - ffi::reflection::TypeAttrDef().attr( + ffi::reflection::TypeAttrDef().attr( tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&BindingBlockDocTranslate>()); } @@ -110,7 +110,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { ffi::Optional DataflowBlockDocTranslate(DocTranslatorObj* d, ffi::AnyView input, const ffi::Object* destination) { const auto* block = - ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck( + ffi::detail::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck( input); TVM_FFI_CHECK(destination == nullptr, TypeError) << "printer statement-only node cannot fulfill a destination"; @@ -118,7 +118,7 @@ ffi::Optional DataflowBlockDocTranslate(DocTranslatorObj* d, ffi::AnyVi ffi::Array outputs; for (const relax::Binding& binding : block->bindings) { d->Translate(binding); - if (!binding->var.as()) { + if (!binding->var.as()) { outputs.push_back(d->Translate(binding->var).value()); } } @@ -132,7 +132,7 @@ ffi::Optional DataflowBlockDocTranslate(DocTranslatorObj* d, ffi::AnyVi } TVM_FFI_STATIC_INIT_BLOCK() { - ffi::reflection::TypeAttrDef().attr( + ffi::reflection::TypeAttrDef().attr( tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&DataflowBlockDocTranslate>()); } diff --git a/src/relax/script/printer/tir.cc b/src/relax/script/printer/tir.cc index 9ec3794c6466..b6c6bbce3f15 100644 --- a/src/relax/script/printer/tir.cc +++ b/src/relax/script/printer/tir.cc @@ -31,15 +31,15 @@ namespace { ExprDoc RelaxShapeDim(DocTranslatorObj* d, const PrimExpr& dim, const ExprDoc& translated) { auto operation = translated.as(); - if (auto mul = dim.as(); - mul && operation && operation.value()->kind == OperationDocNode::Kind::kMult) { - ExprDoc compact = OperationDoc(OperationDocNode::Kind::kMult, + if (auto mul = dim.as(); + mul && operation && operation.value()->kind == OperationDocObj::Kind::kMult) { + ExprDoc compact = OperationDoc(OperationDocObj::Kind::kMult, {RelaxShapeDim(d, mul->a, operation.value()->operands[0]), RelaxShapeDim(d, mul->b, operation.value()->operands[1])}); d->RecordOrigin(compact, dim); return compact; } - if (auto integer = dim.as(); + if (auto integer = dim.as(); integer && integer->ty.as_or_throw()->dtype == (DLDataType{kDLInt, 64, 1})) { ExprDoc compact = LiteralDoc::Int(ffi::GetRef(integer), std::nullopt); d->RecordOrigin(compact, dim); diff --git a/src/relax/script/printer/utils.h b/src/relax/script/printer/utils.h index 4ba0c793c7e2..52a0f415f08f 100644 --- a/src/relax/script/printer/utils.h +++ b/src/relax/script/printer/utils.h @@ -33,7 +33,7 @@ namespace details { ffi::Optional GlobalInfoSelector(DocTranslatorObj* d, const GlobalInfo& info); ExprDoc RelaxShapeDim(DocTranslatorObj* d, const PrimExpr& dim); -ffi::Array RelaxSeqBody(DocTranslatorObj* d, const relax::SeqExprNode* seq, +ffi::Array RelaxSeqBody(DocTranslatorObj* d, const relax::SeqExprObj* seq, ffi::Optional destination = std::nullopt, ffi::Optional annotation = std::nullopt, const ffi::Object* destination_object = nullptr); diff --git a/src/relax/training/utils.cc b/src/relax/training/utils.cc index 0fa7b29d4fd0..5f4f1e8de827 100644 --- a/src/relax/training/utils.cc +++ b/src/relax/training/utils.cc @@ -43,7 +43,7 @@ class AppendLossMutator : private ExprMutator { public: static IRModule Transform(IRModule mod, ffi::String func_name, Function loss_function, int num_backbone_outputs, ffi::Optional new_func_name) { - auto* old_func = mod->Lookup(func_name).as(); + auto* old_func = mod->Lookup(func_name).as(); TVM_FFI_ICHECK(old_func) << func_name << "is not a Relax Function"; // functions should be copied to satisfy the well-formed check @@ -67,7 +67,7 @@ class AppendLossMutator : private ExprMutator { loss_function_(loss_function), num_backbone_outputs_(num_backbone_outputs) {} - Expr VisitExpr_(const FunctionNode* func) final { + Expr VisitExpr_(const FunctionObj* func) final { // Well-formed checks and setting up class members loss_body_ = loss_function_->body; CheckLossBody(); @@ -83,9 +83,9 @@ class AppendLossMutator : private ExprMutator { return Function(new_params, new_body, std::nullopt, func->is_pure, func->attrs); } - Expr VisitExpr_(const SeqExprNode* seq_expr) final { + Expr VisitExpr_(const SeqExprObj* seq_expr) final { TVM_FFI_ICHECK(seq_expr->blocks.size() == 1 && - seq_expr->blocks[0]->IsInstance()) + seq_expr->blocks[0]->IsInstance()) << "Backbone should have only one DataflowBlock"; auto new_blocks = ffi::Array({this->VisitBindingBlock(seq_expr->blocks[0])}); @@ -95,7 +95,7 @@ class AppendLossMutator : private ExprMutator { return SeqExpr(new_blocks, ret.size() == 1 ? ret[0] : Tuple(ret)); } - BindingBlock VisitBindingBlock_(const DataflowBlockNode* block) final { + BindingBlock VisitBindingBlock_(const DataflowBlockObj* block) final { builder_->BeginDataflowBlock(); // Emit original bindings. for (const auto& binding : block->bindings) { @@ -119,9 +119,9 @@ class AppendLossMutator : private ExprMutator { /*! \brief Checks the loss function have only one DataflowBlock, and returns a scalar Var. */ void CheckLossBody() { TVM_FFI_ICHECK(loss_body_->blocks.size() == 1 && - loss_body_->blocks[0]->IsInstance()) + loss_body_->blocks[0]->IsInstance()) << "The loss function should have only one DataflowBlock"; - auto var_node = loss_body_->body.as(); + auto var_node = loss_body_->body.as(); TVM_FFI_ICHECK(var_node && IsScalarTensor(ffi::GetRef(var_node))) << "The loss function must return a scalar(0-dim Tensor) Var"; } @@ -131,11 +131,11 @@ class AppendLossMutator : private ExprMutator { * one or a tuple of Vars. */ void BackboneReturnToArr(const Expr& backbone_return) { - if (auto* var = backbone_return.as()) { + if (auto* var = backbone_return.as()) { backbone_return_arr_.push_back(ffi::GetRef(var)); - } else if (auto* tuple = backbone_return.as()) { + } else if (auto* tuple = backbone_return.as()) { for (auto i : tuple->fields) { - auto var = i.as(); + auto var = i.as(); TVM_FFI_ICHECK(var) << "The return value of the backbone should be either a Var or a Tuple of Vars"; backbone_return_arr_.push_back(ffi::GetRef(var)); @@ -192,7 +192,7 @@ class AppendLossMutator : private ExprMutator { backbone_return_arr_.begin() + num_backbone_outputs_, backbone_return_arr_.end()); for (int i = 0; i < num_backbone_outputs_; ++i) { auto var = backbone_return_arr_[i]; - if (other_outputs_var.count(var) == 0 && !var->ty.as()) { + if (other_outputs_var.count(var) == 0 && !var->ty.as()) { auto new_var = DataflowVar(var->name, GetType(var), var->loc); this->var_remap_.insert_or_assign(var, new_var); backbone_return_arr_.Set(i, new_var); diff --git a/src/relax/transform/allocate_workspace.cc b/src/relax/transform/allocate_workspace.cc index af3a34d50cf4..9dffbb562c0e 100644 --- a/src/relax/transform/allocate_workspace.cc +++ b/src/relax/transform/allocate_workspace.cc @@ -43,8 +43,8 @@ class ExternFunctionRewriter : ExprMutator { ExternFunctionRewriter(IRModule mod, size_t max_workspace_size) : ExprMutator(mod), max_workspace_size_(max_workspace_size) {} - std::unordered_map Run() { - std::unordered_map ret; + std::unordered_map Run() { + std::unordered_map ret; for (const auto& [gvar, f] : builder_->GetContextIRModule()->functions) { if (f->GetAttr(tvm::relax::attr::kWorkspaceSize)) { ret.insert_or_assign(gvar.get(), VisitExpr(f).as_or_throw()); @@ -53,7 +53,7 @@ class ExternFunctionRewriter : ExprMutator { return ret; } - Expr VisitExpr_(const FunctionNode* func_node) override { + Expr VisitExpr_(const FunctionObj* func_node) override { if (!func_node->GetAttr(tvm::relax::attr::kCodegen) && !func_node->GetAttr(tvm::relax::attr::kComposite)) { return ExprMutator::VisitExpr_(func_node); @@ -78,11 +78,11 @@ class ExternFunctionRewriter : ExprMutator { return ExprMutator::VisitExpr_(func_node); } - Expr VisitExpr_(const CallNode* call_node) override { + Expr VisitExpr_(const CallObj* call_node) override { auto new_op = VisitExpr(call_node->op); if (auto var = new_op.as()) { if (auto callee = builder_->LookupBinding(var.value()); - callee && callee.value()->IsInstance() && + callee && callee.value()->IsInstance() && callee.value().as_or_throw()->GetAttr( tvm::relax::attr::kComposite.name)) { // Append the workspace argument to this call. The callee should have been updated to accept @@ -136,7 +136,7 @@ class WorkspaceProvider : ExprMutator { for (const auto& [gvar, f] : mod_->functions) { workspace_var_main_ = std::nullopt; - if (!f->IsInstance() || f->GetAttr(tvm::relax::attr::kCodegen) || + if (!f->IsInstance() || f->GetAttr(tvm::relax::attr::kCodegen) || f->GetAttr(tvm::relax::attr::kComposite)) { continue; } @@ -148,7 +148,7 @@ class WorkspaceProvider : ExprMutator { return builder_->GetContextIRModule(); } - BindingBlock VisitBindingBlock_(const DataflowBlockNode* block_node) final { + BindingBlock VisitBindingBlock_(const DataflowBlockObj* block_node) final { builder_->BeginDataflowBlock(); if (!workspace_var_main_.has_value()) { auto shape = ShapeExpr({IntImm::Int32(max_workspace_size_)}); @@ -162,14 +162,14 @@ class WorkspaceProvider : ExprMutator { return builder_->EndBlock(); } - Expr VisitExpr_(const GlobalVarNode* gvar_node) override { + Expr VisitExpr_(const GlobalVarObj* gvar_node) override { if (gvar_map_.count(gvar_node)) { return gvar_map_.at(gvar_node); } return ExprMutator::VisitExpr_(gvar_node); } - Expr VisitExpr_(const CallNode* call_node) override { + Expr VisitExpr_(const CallObj* call_node) override { auto new_op = VisitExpr(call_node->op); if (auto gv = new_op.as()) { @@ -193,7 +193,7 @@ class WorkspaceProvider : ExprMutator { /*! \brief A map from old global variables representing a function with workspace requirement to * the new ones that are transformed to take an additional workspace parameter. This is only * needed since the type of the global variables changes between transformation. */ - std::unordered_map gvar_map_; + std::unordered_map gvar_map_; std::unordered_set new_gvars_; }; diff --git a/src/relax/transform/alter_op_impl.cc b/src/relax/transform/alter_op_impl.cc index 39a8ebf4c19a..34d599732337 100644 --- a/src/relax/transform/alter_op_impl.cc +++ b/src/relax/transform/alter_op_impl.cc @@ -93,7 +93,7 @@ class AlterOpImplMutator : public ExprMutator { IRModule Run() { for (const auto& gv : mod_->GetGlobalVars()) { const auto& func = mod_->Lookup(gv); - if (func->IsInstance()) { + if (func->IsInstance()) { relax::Function update_func = VisitExpr(func).as_or_throw(); builder_->UpdateFunction(gv, update_func); } @@ -104,7 +104,7 @@ class AlterOpImplMutator : public ExprMutator { private: using ExprMutator::VisitExpr_; - Expr VisitExpr_(const CallNode* op) final { + Expr VisitExpr_(const CallObj* op) final { auto call = ExprMutator::VisitExpr_(op).as_or_throw(); // TODO(@tvm-team): When we differentiate the call for tirx function and packed function, @@ -112,10 +112,10 @@ class AlterOpImplMutator : public ExprMutator { if (!call->op.same_as(call_tir_op_)) return call; // Do not do anything for external function - if (call->args[0].as()) return call; + if (call->args[0].as()) return call; // Get operator name from callee - TVM_FFI_ICHECK(call->args[0]->IsInstance()); + TVM_FFI_ICHECK(call->args[0]->IsInstance()); const tirx::Function& old_func = mod_->Lookup(call->args[0].as_or_throw()).as_or_throw(); ffi::Optional maybe_op_kind = @@ -140,7 +140,7 @@ class AlterOpImplMutator : public ExprMutator { GlobalVar replacement_gv = GetOrCreateGlobalVarForFunc(replacement_func, op_kind); - auto call_tir_inputs_tuple = ffi::GetRef(call->args[1].as()); + auto call_tir_inputs_tuple = ffi::GetRef(call->args[1].as()); Tuple updated_inputs = UpdateInputs(call_tir_inputs_tuple, tensor_transforms); TVM_FFI_ICHECK_EQ(call->ty_args.size(), 1) << "call_tir ty_args.size() is expected to be 1"; @@ -154,15 +154,15 @@ class AlterOpImplMutator : public ExprMutator { } ffi::Array GetTensorTypePerOutput(const Type& output_ty) { - if (const auto* tensor_ty = output_ty.as()) + if (const auto* tensor_ty = output_ty.as()) return {ffi::GetRef(tensor_ty)}; - const auto* tuple_ty = output_ty.as(); + const auto* tuple_ty = output_ty.as(); TVM_FFI_ICHECK(tuple_ty); ffi::Array tensor_tys; tensor_tys.reserve(tuple_ty->fields.size()); for (const auto& ty : tuple_ty->fields) { - const auto* tensor_ty = ty.as(); + const auto* tensor_ty = ty.as(); TVM_FFI_ICHECK(tensor_ty) << "Nested tuples in output of call_tir is not supported yet"; tensor_tys.push_back(ffi::GetRef(tensor_ty)); } @@ -170,9 +170,9 @@ class AlterOpImplMutator : public ExprMutator { } bool IsScalarConstant(const Expr& expr) { - if (expr->IsInstance() && - expr.as()->value.as() && - expr.as()->value.cast().ndim() == 0) { + if (expr->IsInstance() && + expr.as()->value.as() && + expr.as()->value.cast().ndim() == 0) { return true; } return false; @@ -296,11 +296,11 @@ class AlterOpImplMutator : public ExprMutator { const ffi::Array>& tensor_transforms) { if (tensor_transforms.empty()) return out_ty; - if (out_ty->IsInstance()) + if (out_ty->IsInstance()) return UpdateOutputType(out_ty.as_or_throw(), tensor_transforms[tensor_transforms.size() - 1]); - TVM_FFI_ICHECK(out_ty->IsInstance()) + TVM_FFI_ICHECK(out_ty->IsInstance()) << "Expect output type of call_tir to be either TupleType or " "TensorType, but got " << out_ty; @@ -310,7 +310,7 @@ class AlterOpImplMutator : public ExprMutator { size_t first_output_index = tensor_transforms.size() - tuple_ty->fields.size(); size_t i = 0; for (const auto& si : tuple_ty->fields) { - TVM_FFI_ICHECK(si->IsInstance()) + TVM_FFI_ICHECK(si->IsInstance()) << "Fields of TupleType must be TensorType for call_tir " "output structinfo, but got " << si; diff --git a/src/relax/transform/attach_attr_layout_free_tensors.cc b/src/relax/transform/attach_attr_layout_free_tensors.cc index 38750cc01a51..dbf5ada9559e 100644 --- a/src/relax/transform/attach_attr_layout_free_tensors.cc +++ b/src/relax/transform/attach_attr_layout_free_tensors.cc @@ -37,7 +37,7 @@ class AttrAttacher : public ExprMutator { static IRModule Transform(const IRModule& mod) { AttrAttacher mutator(mod); for (auto [gvar, func] : mod->functions) { - if (func->IsInstance()) { + if (func->IsInstance()) { // clear the layout_free_exprs_ for each function mutator.layout_free_exprs_.clear(); mutator.builder_->UpdateFunction(gvar, mutator.VisitExpr(func).as_or_throw()); @@ -50,7 +50,7 @@ class AttrAttacher : public ExprMutator { explicit AttrAttacher(IRModule mod) : ExprMutator(mod), mod_(mod) {} using ExprMutator::VisitExpr_; - Expr VisitExpr_(const FunctionNode* op) final { + Expr VisitExpr_(const FunctionObj* op) final { if (auto opt_num_input = op->attrs.GetAttr(tvm::relax::attr::kNumInput)) { TVM_FFI_ICHECK(layout_free_exprs_.empty()) << "meet a non-global function with num_input attr"; @@ -62,12 +62,12 @@ class AttrAttacher : public ExprMutator { return ExprMutator::VisitExpr_(op); } - Expr VisitExpr_(const GenericConstNode* op) final { + Expr VisitExpr_(const GenericConstObj* op) final { if (op->value.as()) layout_free_exprs_.insert(op); return ExprMutator::VisitExpr_(op); } - Expr VisitExpr_(const CallNode* op) final { + Expr VisitExpr_(const CallObj* op) final { static const Op call_tir_op_ = Op::Get("relax.call_tir"); Call call = ExprMutator::VisitExpr_(op).as_or_throw(); if (!call->op.same_as(call_tir_op_)) { @@ -92,14 +92,14 @@ class AttrAttacher : public ExprMutator { // So we don't need to worry about the duplicate insertion GlobalVar new_gv = builder_->AddFunction(func, gv->name_hint); // Create a new call node with the updated tirx::Function - auto n = ffi::make_object(*op); + auto n = ffi::make_object(*op); n->args = {new_gv, Tuple(call_tir_args)}; return Call(n); } private: IRModule mod_; - std::unordered_set layout_free_exprs_; + std::unordered_set layout_free_exprs_; }; namespace transform { diff --git a/src/relax/transform/attach_global_symbol.cc b/src/relax/transform/attach_global_symbol.cc index 09d0bfdc4e74..155503844ac1 100644 --- a/src/relax/transform/attach_global_symbol.cc +++ b/src/relax/transform/attach_global_symbol.cc @@ -46,7 +46,7 @@ struct RelaxGvarMutator : ExprMutator { : replacements(replacements) {} using ExprMutator::VisitExpr_; - Expr VisitExpr_(const GlobalVarNode* node) override { + Expr VisitExpr_(const GlobalVarObj* node) override { auto gvar = ffi::GetRef(node); return replacements.Get(gvar).value_or(gvar); } @@ -59,7 +59,7 @@ struct TirxGvarMutator : tirx::StmtExprMutator { : replacements(replacements) {} using tirx::StmtExprMutator::Mutate_; - UnchangedOr Mutate_(const CallNode* node, InplaceMode inplace_mode) override { + UnchangedOr Mutate_(const CallObj* node, InplaceMode inplace_mode) override { auto call = tirx::StmtExprMutator::Mutate_(node, inplace_mode) .ValueOrUnchanged(ffi::GetRef(node)) .as_or_throw(); @@ -87,7 +87,7 @@ IRModule ReplaceGlobalVarsInModule(IRModule mod, ffi::Map auto new_gvar = replacements.Get(old_gvar).value_or(old_gvar); BaseFunc new_func = old_func; - if (auto* function_node = old_func.as()) { + if (auto* function_node = old_func.as()) { auto func = ffi::GetRef(function_node); auto mutator = ffi::make_object(replacements); auto new_body = @@ -102,7 +102,7 @@ IRModule ReplaceGlobalVarsInModule(IRModule mod, ffi::Map } } new_func = func; - } else if (auto* relax_func_node = old_func.as()) { + } else if (auto* relax_func_node = old_func.as()) { RelaxGvarMutator mutator(replacements); auto new_relax_func = mutator(ffi::GetRef(relax_func_node)).as_or_throw(); // Update kGlobalSymbol if the function is externally exposed and being renamed. @@ -113,7 +113,7 @@ IRModule ReplaceGlobalVarsInModule(IRModule mod, ffi::Map } } new_func = new_relax_func; - } else if (old_func.as()) { + } else if (old_func.as()) { // ExternFunc: no internal GlobalVar references to update. new_func = old_func; } else { @@ -155,11 +155,11 @@ Pass AttachGlobalSymbol() { ffi::Optional new_name; BaseFunc new_func = func; - if (auto* function = func.as()) { + if (auto* function = func.as()) { new_name = c_prefix + gvar->name_hint; new_func = WithAttr(ffi::GetRef(function), tvm::attr::kGlobalSymbol.name, new_name); - } else if (auto* relax_func = func.as()) { + } else if (auto* relax_func = func.as()) { new_name = gvar->name_hint; new_func = WithAttr(ffi::GetRef(relax_func), tvm::attr::kGlobalSymbol.name, new_name); diff --git a/src/relax/transform/bind_params.cc b/src/relax/transform/bind_params.cc index e9fbab70bd14..e7daebe438e5 100644 --- a/src/relax/transform/bind_params.cc +++ b/src/relax/transform/bind_params.cc @@ -37,7 +37,7 @@ ffi::Map NormalizeBindings(const Function& func, // Map from string to the variable(s) with that name. std::unordered_map> string_lookup; - std::unordered_set var_set; + std::unordered_set var_set; for (const auto& param : func->params) { string_lookup[param->name].push_back(param); var_set.insert(param.get()); @@ -115,10 +115,10 @@ Function FunctionBindParams(Function func, const ffi::Map& * \return The module after binding params. */ IRModule BindParam(IRModule m, ffi::String func_name, ffi::Map bind_params) { - IRModuleNode* new_module = m.CopyOnWrite(); + IRModuleObj* new_module = m.CopyOnWrite(); ffi::Map functions = m->functions; for (const auto& func_pr : functions) { - if (const auto* relax_f = func_pr.second.as()) { + if (const auto* relax_f = func_pr.second.as()) { if (relax_f->GetLinkageType() == LinkageType::kExternal) { // Use global_symbol if it's external linkage ffi::Optional gsymbol = relax_f->GetAttr(tvm::attr::kGlobalSymbol); diff --git a/src/relax/transform/bind_symbolic_vars.cc b/src/relax/transform/bind_symbolic_vars.cc index ebb4310bfd48..1d0404a6068c 100644 --- a/src/relax/transform/bind_symbolic_vars.cc +++ b/src/relax/transform/bind_symbolic_vars.cc @@ -41,7 +41,7 @@ Function FunctionBindSymbolicVars( // Map from string to the variable(s) with that name. std::unordered_map> string_lookup; - std::unordered_set symbolic_var_set; + std::unordered_set symbolic_var_set; for (const auto& var : old_symbolic_vars) { string_lookup[var->name].push_back(var); symbolic_var_set.insert(var.get()); @@ -104,7 +104,7 @@ IRModule ModuleBindSymbolicVars( // Collect bindings that are used by this function. auto func_binding_map = [&]() -> ffi::Map, PrimExpr> { std::unordered_set var_names; - std::unordered_set vars; + std::unordered_set vars; for (const auto& var : DefinedSymbolicVars(func)) { var_names.insert(var->name); vars.insert(var.get()); diff --git a/src/relax/transform/bundle_model_params.cc b/src/relax/transform/bundle_model_params.cc index 5222ffd8486d..dffcd572e87f 100644 --- a/src/relax/transform/bundle_model_params.cc +++ b/src/relax/transform/bundle_model_params.cc @@ -42,7 +42,7 @@ class ModelParamBundler : public ExprMutator { explicit ModelParamBundler(ffi::Optional param_tuple_name) : param_tuple_name_(param_tuple_name) {} - Expr VisitExpr_(const FunctionNode* op) override { + Expr VisitExpr_(const FunctionObj* op) override { Function func = ffi::GetRef(op); auto opt_num_input = func->attrs.GetAttr(tvm::relax::attr::kNumInput); if (!opt_num_input) return func; @@ -59,12 +59,12 @@ class ModelParamBundler : public ExprMutator { params.push_back(func->params[i]); } - std::unordered_set signature_vars; + std::unordered_set signature_vars; for (const tvm::Var& var : DefinableTIRVarsInType(TupleType(params.Map(GetType)))) { signature_vars.insert(var.get()); } - std::unordered_set bundled_prim_params; + std::unordered_set bundled_prim_params; ffi::Array bundled_prim_params_in_order; for (size_t i = num_input; i < func->params.size(); i++) { if (func->params[i].as()) { @@ -107,7 +107,7 @@ class ModelParamBundler : public ExprMutator { return Function(params, body, ret_ty, func->is_pure, func->attrs); } - Expr VisitExpr_(const SeqExprNode* op) override { + Expr VisitExpr_(const SeqExprObj* op) override { if (pending_prim_params_.empty()) { return ExprMutator::VisitExpr_(op); } @@ -132,7 +132,7 @@ class ModelParamBundler : public ExprMutator { return SeqExpr(blocks, body->body); } - Expr VisitExpr_(const VarNode* op) override { + Expr VisitExpr_(const VarObj* op) override { auto var = ffi::GetRef(op); if (!rewrite_model_params_) { return ExprMutator::VisitExpr_(op); diff --git a/src/relax/transform/call_tir_rewrite.cc b/src/relax/transform/call_tir_rewrite.cc index 6fd17681d1b1..a667b37df596 100644 --- a/src/relax/transform/call_tir_rewrite.cc +++ b/src/relax/transform/call_tir_rewrite.cc @@ -52,7 +52,7 @@ class CallTIRMutator : public ExprMutator { IRModule Run() { for (const auto& [gv, func] : mod_->functions) { - if (func->IsInstance()) { + if (func->IsInstance()) { auto updated_func = this->VisitExpr(func).as_or_throw(); builder_->UpdateFunction(gv, updated_func); } @@ -62,10 +62,10 @@ class CallTIRMutator : public ExprMutator { private: using ExprMutator::VisitExpr_; - Expr VisitExpr_(const CallNode* call) override { + Expr VisitExpr_(const CallObj* call) override { // post-order mutation Expr expr = VisitExprPostOrder_(call); - call = expr.as(); + call = expr.as(); static const Op call_tir_op = Op::Get("relax.call_tir"); static const Op call_tir_inplace_op = Op::Get("relax.call_tir_inplace"); @@ -96,7 +96,7 @@ class CallTIRMutator : public ExprMutator { tuple_output_type = output_ty; bool has_nested_tuple = std::any_of(output_ty->fields.begin(), output_ty->fields.end(), - [](const Type& field) { return field->IsInstance(); }); + [](const Type& field) { return field->IsInstance(); }); if (has_nested_tuple) { TVM_FFI_ICHECK(!is_inplace) @@ -106,7 +106,7 @@ class CallTIRMutator : public ExprMutator { for (size_t i = 0; i < output_ty->fields.size(); ++i) { const auto& field = output_ty->fields[i]; - TVM_FFI_ICHECK(field->IsInstance()) + TVM_FFI_ICHECK(field->IsInstance()) << "call_tir expects Tuple of TensorType, but got " << field << " as an element of TupleType"; const auto& field_tensor = field.as_or_throw(); @@ -126,7 +126,7 @@ class CallTIRMutator : public ExprMutator { } ffi::Array args; - if (call->args[1].as()) { + if (call->args[1].as()) { args = call->args[1].as_or_throw()->fields; // for call_tir_inplace, don't reinsert in-place args, only the newly allocated ones if (!is_inplace) { @@ -193,11 +193,11 @@ class CallTIRMutator : public ExprMutator { void FlattenAndAllocateOutputs(const Type& type, const Op& alloc_tensor_op, ffi::Array* outs) { - if (const auto* tensor_ty = type.as()) { + if (const auto* tensor_ty = type.as()) { outs->push_back(AllocateOutputTensor(ffi::GetRef(tensor_ty), alloc_tensor_op)); return; } - if (const auto* tuple_ty = type.as()) { + if (const auto* tuple_ty = type.as()) { for (const Type& field : tuple_ty->fields) { FlattenAndAllocateOutputs(field, alloc_tensor_op, outs); } @@ -209,11 +209,11 @@ class CallTIRMutator : public ExprMutator { } Expr RebuildOutputTuple(const Type& type, const ffi::Array& outs, size_t* index) { - if (type->IsInstance()) { + if (type->IsInstance()) { TVM_FFI_ICHECK_LT(*index, outs.size()); return outs[(*index)++]; } - if (const auto* tuple_ty = type.as()) { + if (const auto* tuple_ty = type.as()) { ffi::Array fields; fields.reserve(tuple_ty->fields.size()); for (const Type& field : tuple_ty->fields) { diff --git a/src/relax/transform/canonicalize_bindings.cc b/src/relax/transform/canonicalize_bindings.cc index 3dafe80d9aec..9249544faa6e 100644 --- a/src/relax/transform/canonicalize_bindings.cc +++ b/src/relax/transform/canonicalize_bindings.cc @@ -55,7 +55,7 @@ class SymbolicVarCanonicalizer : public ExprMutator { return CanonicalizeShapeValue(expr); } - Expr VisitExpr_(const CallNode* op) final { + Expr VisitExpr_(const CallObj* op) final { static const Op& call_tir_op = Op::Get("relax.call_tir"); static const Op& call_tir_inplace_op = Op::Get("relax.call_tir_inplace"); static const Op& call_tir_with_grad_op = Op::Get("relax.call_tir_with_grad"); @@ -99,8 +99,8 @@ class SymbolicVarCanonicalizer : public ExprMutator { if (IsBaseOf(implied, out_ty)) { return out_ty; } - const auto* implied_tuple = implied.as(); - const auto* out_tuple = out_ty.as(); + const auto* implied_tuple = implied.as(); + const auto* out_tuple = out_ty.as(); if (implied_tuple && out_tuple && implied_tuple->fields.size() == out_tuple->fields.size()) { ffi::Array fields; for (size_t i = 0; i < out_tuple->fields.size(); ++i) { @@ -108,8 +108,8 @@ class SymbolicVarCanonicalizer : public ExprMutator { } return TupleType(fields, out_tuple->loc); } - const auto* implied_tensor = implied.as(); - const auto* out_tensor = out_ty.as(); + const auto* implied_tensor = implied.as(); + const auto* out_tensor = out_ty.as(); if (implied_tensor && out_tensor && implied_tensor->shape.has_value()) { return TensorType(implied_tensor->shape.value(), out_tensor->dtype, out_tensor->vdevice, out_tensor->loc); @@ -117,7 +117,7 @@ class SymbolicVarCanonicalizer : public ExprMutator { return implied; } - Expr VisitExpr_(const ShapeExprNode* op) final { + Expr VisitExpr_(const ShapeExprObj* op) final { if (!canonicalize_shape_values_) return ffi::GetRef(op); ffi::Array values = op->values.Map([this](const PrimExpr& value) { return CanonicalizeShapeValue(value); }); @@ -127,7 +127,7 @@ class SymbolicVarCanonicalizer : public ExprMutator { return ShapeExpr(values, op->loc); } - Expr VisitExpr_(const FunctionNode* func) override { + Expr VisitExpr_(const FunctionObj* func) override { auto cached = known_values_; auto cached_runtime_uses = runtime_prim_var_uses_; runtime_prim_var_uses_ = RuntimePrimVarCollector::Collect(func->body); @@ -137,7 +137,7 @@ class SymbolicVarCanonicalizer : public ExprMutator { return output; } - void VisitBinding_(const MatchCastNode* binding) override { + void VisitBinding_(const MatchCastObj* binding) override { auto tir_var_map = InferSymbolicVarMap({{binding->var, binding->value}}, builder_->GetAnalyzer()); bool has_runtime_use = false; @@ -167,7 +167,7 @@ class SymbolicVarCanonicalizer : public ExprMutator { canonicalize_shape_values_ = cached; } - Expr VisitExpr_(const IfExprNode* op) override { + Expr VisitExpr_(const IfExprObj* op) override { Expr guard = this->VisitExpr(op->cond); auto cached = known_values_; @@ -229,9 +229,9 @@ class SymbolicVarCanonicalizer : public ExprMutator { void VisitExprDepTypeField(const Type&) final {} - void VisitExpr_(const ShapeExprNode*) final {} + void VisitExpr_(const ShapeExprObj*) final {} - void VisitExpr_(const VarNode* op) final { + void VisitExpr_(const VarObj* op) final { Var var = ffi::GetRef(op); if (auto prim_var = var.as()) { uses_.insert(*prim_var); @@ -267,9 +267,9 @@ class SymbolicVarCanonicalizer : public ExprMutator { Expr SymbolicVarCanonicalizer::Apply(Expr expr) { SymbolicVarCanonicalizer mutator; - if (!expr->IsInstance()) { + if (!expr->IsInstance()) { // Dataflow rewriting canonicalizes a SeqExpr directly, without entering - // VisitExpr_(FunctionNode), so initialize the same runtime-use context. + // VisitExpr_(FunctionObj), so initialize the same runtime-use context. mutator.runtime_prim_var_uses_ = RuntimePrimVarCollector::Collect(expr); } return mutator(std::move(expr)); @@ -302,7 +302,7 @@ class CanonicalizePlanner : public ExprVisitor { // within a DataflowBlock, and is not already handled by removal // of trivial bindings, then we can replace it with a DataflowVar. for (auto var : visitor.defined_inside_dataflow_) { - if (!var.as() && !visitor.used_outside_home_dataflow_.count(var)) { + if (!var.as() && !visitor.used_outside_home_dataflow_.count(var)) { DataflowVar new_var(var->name, GetType(var)); plan.replace_binding.Set(var, new_var); @@ -331,7 +331,7 @@ class CanonicalizePlanner : public ExprVisitor { bound_to = opt.value(); } - if (bound_var.as() || !bound_to.as() || + if (bound_var.as() || !bound_to.as() || !visitor.used_outside_home_dataflow_.count(bound_var)) { // Case 1: Var = Var // Case 2: DataflowVar = Var @@ -361,7 +361,7 @@ class CanonicalizePlanner : public ExprVisitor { } private: - void VisitExpr_(const FunctionNode* func) override { + void VisitExpr_(const FunctionObj* func) override { // for functions, treat any free vars as used outside their home DF block auto cache = current_block_; current_block_ = ffi::Optional(); @@ -373,7 +373,7 @@ class CanonicalizePlanner : public ExprVisitor { current_block_ = cache; } - void VisitExpr_(const SeqExprNode* seq) override { + void VisitExpr_(const SeqExprObj* seq) override { // need to reset current_block_ for nested seq exprs (such as in If nodes) auto cache = current_block_; current_block_ = ffi::Optional(); @@ -381,14 +381,14 @@ class CanonicalizePlanner : public ExprVisitor { current_block_ = cache; } - void VisitBindingBlock_(const BindingBlockNode* block) override { + void VisitBindingBlock_(const BindingBlockObj* block) override { TVM_FFI_ICHECK(!current_block_.has_value()) << "Forgetting to unset current block"; current_block_ = ffi::GetRef(block); ExprVisitor::VisitBindingBlock_(block); current_block_ = ffi::Optional(); } - void VisitBindingBlock_(const DataflowBlockNode* block) override { + void VisitBindingBlock_(const DataflowBlockObj* block) override { TVM_FFI_ICHECK(!current_block_.has_value()) << "Forgetting to unset current block"; current_block_ = ffi::GetRef(block); ExprVisitor::VisitBindingBlock_(block); @@ -421,9 +421,9 @@ class CanonicalizePlanner : public ExprVisitor { // If the expression is a TupleGetItem, which accesses a field of // a known tuple, then it can be unwrapped into a direct access of // that field. - if (auto tuple_get_item = expr.as()) { + if (auto tuple_get_item = expr.as()) { Expr tuple = recursively_unwrap_var(tuple_get_item->tuple); - if (auto ptr = tuple.as()) { + if (auto ptr = tuple.as()) { return ptr->fields[tuple_get_item->index]; } } @@ -432,7 +432,7 @@ class CanonicalizePlanner : public ExprVisitor { // `TupleGetItem(earlier_tuple, i)`, then this is just a copy of // `earlier_tuple`. auto earlier_tuple = [&]() -> ffi::Optional { - auto expr_tuple = expr.as(); + auto expr_tuple = expr.as(); if (!expr_tuple) { return std::nullopt; } @@ -441,7 +441,7 @@ class CanonicalizePlanner : public ExprVisitor { return std::nullopt; } - auto first_element = recursively_unwrap_var(expr_tuple->fields[0]).as(); + auto first_element = recursively_unwrap_var(expr_tuple->fields[0]).as(); if (!first_element) { return std::nullopt; } @@ -455,7 +455,7 @@ class CanonicalizePlanner : public ExprVisitor { Expr earlier_tuple = recursively_unwrap_var(first_element->tuple); for (size_t i = 0; i < expr_tuple->fields.size(); i++) { - auto element = recursively_unwrap_var(expr_tuple->fields[i]).as(); + auto element = recursively_unwrap_var(expr_tuple->fields[i]).as(); if (!element) { return std::nullopt; } @@ -481,9 +481,9 @@ class CanonicalizePlanner : public ExprVisitor { void VisitBinding(const Binding& binding) override { bool has_same_ty = [&]() { - if (binding.as()) { + if (binding.as()) { return true; - } else if (auto match_cast = binding.as()) { + } else if (auto match_cast = binding.as()) { return ffi::StructuralEqual()(GetType(binding->var), GetType(match_cast->value)); } else { TVM_FFI_THROW(InternalError) << "Invalid binding type: " << binding->GetTypeKey(); @@ -516,7 +516,7 @@ class CanonicalizePlanner : public ExprVisitor { } } - void VisitExpr_(const VarNode* var) override { + void VisitExpr_(const VarObj* var) override { auto var_ref = ffi::GetRef(var); // if a var is used in a dataflow block but *not* the one // where it was defined, it also needs to be exposed, so also we treat that as @@ -529,7 +529,7 @@ class CanonicalizePlanner : public ExprVisitor { } inline bool inside_dataflow() { - return current_block_.has_value() && current_block_.value().as(); + return current_block_.has_value() && current_block_.value().as(); } ffi::Optional current_block_; @@ -571,7 +571,7 @@ class BindingCanonicalizer : public ExprMutator { return ExprMutator::VisitVarDef(new_var); } - Expr VisitExpr_(const VarNode* var) override { + Expr VisitExpr_(const VarObj* var) override { Var new_var = ffi::GetRef(var); while (auto opt = plan_.replace_usage.Get(new_var)) { new_var = opt.value(); @@ -586,7 +586,7 @@ class BindingCanonicalizer : public ExprMutator { // Special case: for dataflow blocks, we will check for dataflow vars that solely exist // to be bound to the output. In this case, we will get rid of those bindings and // use the dataflow var's definition directly - BindingBlock VisitBindingBlock_(const DataflowBlockNode* block) override { + BindingBlock VisitBindingBlock_(const DataflowBlockObj* block) override { auto new_block = ExprMutator::VisitBindingBlock_(block).as_or_throw(); std::unordered_set disqualified_set; std::unordered_set output_vars; @@ -597,20 +597,20 @@ class BindingCanonicalizer : public ExprMutator { auto var = binding->var; auto value = GetBoundValue(binding); - if (var->IsInstance()) { + if (var->IsInstance()) { auto df_var = var.as_or_throw(); // disqualify any vars that appear in the RHS // (for a function literal, consider only free vars) ffi::Array rhs_vars; - if (!value->IsInstance()) { + if (!value->IsInstance()) { rhs_vars = FreeVars(value); } else { rhs_vars = AllVars(value); } for (auto rhs_var : rhs_vars) { - if (rhs_var->IsInstance()) { + if (rhs_var->IsInstance()) { disqualified_set.insert(rhs_var.as_or_throw()); } } @@ -625,7 +625,7 @@ class BindingCanonicalizer : public ExprMutator { // We are looking for cases where the RHS is a single dataflow var; // disqualify if the RHS is not a single dataflow var // or if the var has been output before - if (const auto* rhs_var = value.as()) { + if (const auto* rhs_var = value.as()) { if (output_vars.count(ffi::GetRef(rhs_var))) { disqualified_set.insert(ffi::GetRef(rhs_var)); } @@ -633,14 +633,14 @@ class BindingCanonicalizer : public ExprMutator { } else { ffi::Array disqualified; // for function literal, consider only free vars - if (value->IsInstance()) { + if (value->IsInstance()) { disqualified = FreeVars(value); } else { disqualified = AllVars(value); } for (auto rhs_var : disqualified) { - if (rhs_var->IsInstance()) { + if (rhs_var->IsInstance()) { disqualified_set.insert(rhs_var.as_or_throw()); } } @@ -653,20 +653,20 @@ class BindingCanonicalizer : public ExprMutator { ffi::Array new_bindings; bool changed = false; for (auto binding : new_block->bindings) { - if (binding->var->IsInstance() && + if (binding->var->IsInstance() && candidates.count(binding->var.as_or_throw())) { changed = true; continue; - } else if (!binding->var->IsInstance() && - GetBoundValue(binding)->IsInstance() && + } else if (!binding->var->IsInstance() && + GetBoundValue(binding)->IsInstance() && candidates.count(GetBoundValue(binding).as_or_throw())) { changed = true; - if (auto* match_binding = binding.as()) { + if (auto* match_binding = binding.as()) { auto new_binding = MatchCast( binding->var, candidates.at(match_binding->value.as_or_throw()), match_binding->ty); new_bindings.push_back(new_binding); - } else if (auto* var_binding = binding.as()) { + } else if (auto* var_binding = binding.as()) { auto new_binding = VarBinding( binding->var, candidates.at(var_binding->value.as_or_throw())); new_bindings.push_back(new_binding); diff --git a/src/relax/transform/combine_parallel_matmul.cc b/src/relax/transform/combine_parallel_matmul.cc index 5479895712a9..bfeb9158bfd4 100644 --- a/src/relax/transform/combine_parallel_matmul.cc +++ b/src/relax/transform/combine_parallel_matmul.cc @@ -175,7 +175,7 @@ ffi::TypedFunction(ffi::Map, ffi::Map bool { - return !split.split_size->IsInstance(); + return !split.split_size->IsInstance(); }; std::stable_sort(splits.begin(), splits.end(), [&is_dynamic_split](const auto& a, const auto& b) { @@ -237,7 +237,7 @@ ffi::TypedFunction(ffi::Map, ffi::Map sections; for (size_t i = 0; i + 1 < splits.size(); i++) { - auto width = splits[i].split_size.as(); + auto width = splits[i].split_size.as(); TVM_FFI_CHECK(width, InternalError) << "All splits except the last one must have a static shape"; split_index = (split_index + width->value).as().value(); @@ -286,10 +286,10 @@ std::vector GetBranchInfo(Function f) { auto create_group = [&](DFPattern pat) { // Maps a LHS matrix to consumer parallel matmuls - std::unordered_map groups; + std::unordered_map groups; PostOrderVisit(f, [&](const Expr& e) { - if (!e->IsInstance()) return; + if (!e->IsInstance()) return; auto match = ExtractMatchedExpr(pat, e, bindings); if (!match) return; @@ -333,7 +333,7 @@ std::vector GetBranchInfo(Function f) { return groups; }; - std::unordered_map groups_activation; + std::unordered_map groups_activation; for (size_t i = 0; i < activations.size(); ++i) { auto groups = create_group(bias_activation_pat[i]); groups_activation.merge(std::move(groups)); diff --git a/src/relax/transform/compute_prim_value.cc b/src/relax/transform/compute_prim_value.cc index 1bb1edb00016..ed31d1557f7b 100644 --- a/src/relax/transform/compute_prim_value.cc +++ b/src/relax/transform/compute_prim_value.cc @@ -35,7 +35,7 @@ namespace relax { namespace { -bool HasRelaxCallCapabilities(const CallNode* call) { +bool HasRelaxCallCapabilities(const CallObj* call) { auto op = call->op.as(); if (!op) return true; static auto infer_type_map = Op::GetAttrMap(tvm::op_attr::kInferType); @@ -53,7 +53,7 @@ class PrimExprComputeInjector : public ExprMutator { private: using ExprMutator::VisitExpr_; - Expr VisitExpr_(const CallNode* op) final { + Expr VisitExpr_(const CallObj* op) final { Call call = ffi::GetRef(op); if (auto prim_expr = call.as()) { if (call->op.as() && !HasRelaxCallCapabilities(op)) { @@ -68,47 +68,47 @@ class PrimExprComputeInjector : public ExprMutator { return LiftPrimValue(ffi::GetRef(op).as_or_throw()); \ } - RELAX_LIFT_PRIM_EXPR(TensorLoadNode); - RELAX_LIFT_PRIM_EXPR(prim::AddNode); - RELAX_LIFT_PRIM_EXPR(prim::LShiftNode); - RELAX_LIFT_PRIM_EXPR(prim::RShiftNode); - RELAX_LIFT_PRIM_EXPR(prim::BitwiseAndNode); - RELAX_LIFT_PRIM_EXPR(prim::BitwiseOrNode); - RELAX_LIFT_PRIM_EXPR(prim::BitwiseXorNode); - RELAX_LIFT_PRIM_EXPR(prim::BitwiseNotNode); - RELAX_LIFT_PRIM_EXPR(prim::SubNode); - RELAX_LIFT_PRIM_EXPR(prim::MulNode); - RELAX_LIFT_PRIM_EXPR(prim::DivNode); - RELAX_LIFT_PRIM_EXPR(prim::ModNode); - RELAX_LIFT_PRIM_EXPR(prim::FloorDivNode); - RELAX_LIFT_PRIM_EXPR(prim::FloorModNode); - RELAX_LIFT_PRIM_EXPR(prim::MinNode); - RELAX_LIFT_PRIM_EXPR(prim::MaxNode); - RELAX_LIFT_PRIM_EXPR(prim::EQNode); - RELAX_LIFT_PRIM_EXPR(prim::NENode); - RELAX_LIFT_PRIM_EXPR(prim::LTNode); - RELAX_LIFT_PRIM_EXPR(prim::LENode); - RELAX_LIFT_PRIM_EXPR(prim::GTNode); - RELAX_LIFT_PRIM_EXPR(prim::GENode); - RELAX_LIFT_PRIM_EXPR(prim::AndNode); - RELAX_LIFT_PRIM_EXPR(prim::OrNode); - RELAX_LIFT_PRIM_EXPR(prim::CastNode); - RELAX_LIFT_PRIM_EXPR(prim::NotNode); - RELAX_LIFT_PRIM_EXPR(prim::SelectNode); - RELAX_LIFT_PRIM_EXPR(prim::RampNode); - RELAX_LIFT_PRIM_EXPR(prim::BroadcastNode); - RELAX_LIFT_PRIM_EXPR(prim::ShuffleNode); - RELAX_LIFT_PRIM_EXPR(tvm::IntImmNode); - RELAX_LIFT_PRIM_EXPR(tvm::FloatImmNode); + RELAX_LIFT_PRIM_EXPR(TensorLoadObj); + RELAX_LIFT_PRIM_EXPR(prim::AddObj); + RELAX_LIFT_PRIM_EXPR(prim::LShiftObj); + RELAX_LIFT_PRIM_EXPR(prim::RShiftObj); + RELAX_LIFT_PRIM_EXPR(prim::BitwiseAndObj); + RELAX_LIFT_PRIM_EXPR(prim::BitwiseOrObj); + RELAX_LIFT_PRIM_EXPR(prim::BitwiseXorObj); + RELAX_LIFT_PRIM_EXPR(prim::BitwiseNotObj); + RELAX_LIFT_PRIM_EXPR(prim::SubObj); + RELAX_LIFT_PRIM_EXPR(prim::MulObj); + RELAX_LIFT_PRIM_EXPR(prim::DivObj); + RELAX_LIFT_PRIM_EXPR(prim::ModObj); + RELAX_LIFT_PRIM_EXPR(prim::FloorDivObj); + RELAX_LIFT_PRIM_EXPR(prim::FloorModObj); + RELAX_LIFT_PRIM_EXPR(prim::MinObj); + RELAX_LIFT_PRIM_EXPR(prim::MaxObj); + RELAX_LIFT_PRIM_EXPR(prim::EQObj); + RELAX_LIFT_PRIM_EXPR(prim::NEObj); + RELAX_LIFT_PRIM_EXPR(prim::LTObj); + RELAX_LIFT_PRIM_EXPR(prim::LEObj); + RELAX_LIFT_PRIM_EXPR(prim::GTObj); + RELAX_LIFT_PRIM_EXPR(prim::GEObj); + RELAX_LIFT_PRIM_EXPR(prim::AndObj); + RELAX_LIFT_PRIM_EXPR(prim::OrObj); + RELAX_LIFT_PRIM_EXPR(prim::CastObj); + RELAX_LIFT_PRIM_EXPR(prim::NotObj); + RELAX_LIFT_PRIM_EXPR(prim::SelectObj); + RELAX_LIFT_PRIM_EXPR(prim::RampObj); + RELAX_LIFT_PRIM_EXPR(prim::BroadcastObj); + RELAX_LIFT_PRIM_EXPR(prim::ShuffleObj); + RELAX_LIFT_PRIM_EXPR(tvm::IntImmObj); + RELAX_LIFT_PRIM_EXPR(tvm::FloatImmObj); #undef RELAX_LIFT_PRIM_EXPR - Expr VisitExpr_(const ShapeExprNode* op) final { return ffi::GetRef(op); } + Expr VisitExpr_(const ShapeExprObj* op) final { return ffi::GetRef(op); } PrimExpr VisitTypePrimExprField(const PrimExpr& expr) final { return expr; } Expr LiftPrimValue(const PrimExpr& node) { - if (node->IsInstance() || node->IsInstance()) { + if (node->IsInstance() || node->IsInstance()) { return node; } diff --git a/src/relax/transform/convert_dataflow.cc b/src/relax/transform/convert_dataflow.cc index 9713a0ddfd6f..c0e3bbe79c8d 100644 --- a/src/relax/transform/convert_dataflow.cc +++ b/src/relax/transform/convert_dataflow.cc @@ -39,7 +39,7 @@ class DataflowBlockExtractor : public ExprMutator { public: explicit DataflowBlockExtractor(size_t min_size) : ExprMutator(), min_size_(min_size) {} - Expr VisitExpr_(const SeqExprNode* seq) override { + Expr VisitExpr_(const SeqExprObj* seq) override { ffi::Array new_blocks; Expr new_body = VisitExpr(seq->body); bool changed = !new_body.same_as(seq->body); @@ -115,8 +115,8 @@ class DataflowBlockExtractor : public ExprMutator { for (const auto& binding : binding_block->bindings) { Expr value = GetBoundValue(binding); // dataflow values: not an if node and not an impure call - bool is_dataflow = (!value.as()) && - (!(value.as() && IsImpureCall(value.as_or_throw()))); + bool is_dataflow = (!value.as()) && + (!(value.as() && IsImpureCall(value.as_or_throw()))); if (is_dataflow) { // extend the streak dataflow_bindings.push_back(binding); diff --git a/src/relax/transform/convert_layout.cc b/src/relax/transform/convert_layout.cc index be920b838a12..efb2532cf709 100644 --- a/src/relax/transform/convert_layout.cc +++ b/src/relax/transform/convert_layout.cc @@ -123,7 +123,7 @@ class LayoutConvertMutator : public ExprMutator { TVM_FFI_ICHECK(!NLayoutEqual()(from, LayoutDecision::InitUnknownDim()) && !NLayoutEqual()(to, LayoutDecision::InitUnknownDim())) << "Cannot convert when exactly one of the layouts is unknown"; - const auto* tensor = GetTypeAs(expr); + const auto* tensor = GetTypeAs(expr); TVM_FFI_ICHECK(tensor != nullptr) << "Expect a tensor, but got: " << expr; if ((from.LeafValue()->layout.has_value() ? from.LeafValue()->layout.value().ndim() : 0) == @@ -182,7 +182,7 @@ class LayoutConvertMutator : public ExprMutator { return RewriteExpr(var, InitialNLayout(var)); } - Expr VisitExpr_(const VarNode* op) final { return VisitVars_(ffi::GetRef(op)); } + Expr VisitExpr_(const VarObj* op) final { return VisitVars_(ffi::GetRef(op)); } bool HasUnknownDimTensor(const NLayout& nlayout) { bool find = false; @@ -205,10 +205,10 @@ class LayoutConvertMutator : public ExprMutator { } ffi::Optional GetInferLayoutInfo( - const CallNode* call_node, + const CallObj* call_node, const ffi::Map>& desired_layouts, const LayoutCb& layout_cb, const VarLayoutMap& var_layout_map) { - const OpNode* op_node = call_node->op.as(); + const OpObj* op_node = call_node->op.as(); if (op_node == nullptr) return std::nullopt; Op op = ffi::GetRef(op_node).as_or_throw(); const auto attr_map = Op::GetAttrMap(tvm::relax::op_attr::kInferLayout); @@ -228,13 +228,13 @@ class LayoutConvertMutator : public ExprMutator { } } - void VisitBinding_(const VarBindingNode* binding, const CallNode* call_node) final { + void VisitBinding_(const VarBindingObj* binding, const CallObj* call_node) final { ffi::Optional res = GetInferLayoutInfo(call_node, desired_layouts_, layout_cb_, var_layout_map_); - ffi::ObjectPtr new_call = ffi::make_object(*call_node); + ffi::ObjectPtr new_call = ffi::make_object(*call_node); new_call->ty = Type::Missing(); if (!res.has_value() || - (!IsNestedTensor(binding->var) && !binding->var->IsInstance())) { + (!IsNestedTensor(binding->var) && !binding->var->IsInstance())) { // Default policy: use the initial layout. // When we don't have the infer layout info, or it's a non-tensor global var binding. std::vector input_layout; @@ -256,7 +256,7 @@ class LayoutConvertMutator : public ExprMutator { new_call->attrs = std::move(res.value()->new_attrs); Expr cur_call = builder_->Normalize(Call(new_call)); - if (binding->var->IsInstance()) { + if (binding->var->IsInstance()) { // Dataflow var, we emit the rewritten call. ReEmitBinding(binding, cur_call); // update the layout map @@ -276,10 +276,10 @@ class LayoutConvertMutator : public ExprMutator { } } - void VisitBinding_(const VarBindingNode* binding, const TupleNode* val) final { + void VisitBinding_(const VarBindingObj* binding, const TupleObj* val) final { std::vector input_layout; for (const auto& field : val->fields) { - if (binding->var->IsInstance()) { + if (binding->var->IsInstance()) { // Df var: Use the current realized layout to group the tuple; input_layout.push_back(GetNLayout(var_layout_map_, field)); } else { @@ -294,8 +294,8 @@ class LayoutConvertMutator : public ExprMutator { } } - void VisitBinding_(const VarBindingNode* binding, const TupleGetItemNode* val) final { - NLayout input_layout = binding->var->IsInstance() + void VisitBinding_(const VarBindingObj* binding, const TupleGetItemObj* val) final { + NLayout input_layout = binding->var->IsInstance() ? GetNLayout(var_layout_map_, val->tuple) : InitialNLayout(val->tuple); ReEmitBinding(binding, builder_->Normalize( @@ -304,8 +304,8 @@ class LayoutConvertMutator : public ExprMutator { var_layout_map_[binding->var] = input_layout.NestedArray()[val->index]; } - void VisitBinding_(const MatchCastNode* binding) final { - if (!binding->var->IsInstance()) { + void VisitBinding_(const MatchCastObj* binding) final { + if (!binding->var->IsInstance()) { ExprMutator::VisitBinding_(binding); return; } @@ -318,10 +318,10 @@ class LayoutConvertMutator : public ExprMutator { TVM_FFI_ICHECK(!NLayoutEqual()(from, LayoutDecision::InitUnknownDim()) && !NLayoutEqual()(to, LayoutDecision::InitUnknownDim())) << "Cannot convert when exactly one of the layouts is unknown"; - const TensorTypeNode* tensor_ty = ty.as(); + const TensorTypeObj* tensor_ty = ty.as(); TVM_FFI_ICHECK(tensor_ty != nullptr) << "We can not set layout for non-tensor struct"; if (!tensor_ty->shape.has_value()) return ty; - const ShapeExprNode* shape = tensor_ty->shape.value().as(); + const ShapeExprObj* shape = tensor_ty->shape.value().as(); if (shape == nullptr) return ty; TVM_FFI_ICHECK_EQ( shape->values.size(), diff --git a/src/relax/transform/dataflow_inplace.cc b/src/relax/transform/dataflow_inplace.cc index ce3e7743c55e..cfb420ebdea4 100644 --- a/src/relax/transform/dataflow_inplace.cc +++ b/src/relax/transform/dataflow_inplace.cc @@ -46,7 +46,7 @@ using namespace tvm::prim; // Ops that may return a tensor sharing storage with the first argument. // These ops has been verified to share storage with the first argument in // tests/python/relax/test_dataflow_inplace.py. -bool IsViewMemoryOp(const OpNode* op_node) { +bool IsViewMemoryOp(const OpObj* op_node) { // TODO: Consider to add more ops that may return a tensor sharing storage with // the first argument in the future. static const std::unordered_set kViewOps = { @@ -61,7 +61,7 @@ bool IsViewMemoryOp(const OpNode* op_node) { // Look up alias ids for a call argument (only Var args are expected in dataflow blocks). std::unordered_set GetVarAliasSetFromExpr( const Expr& arg, const std::unordered_map>& alias_sets) { - if (auto* var_node = arg.as()) { + if (auto* var_node = arg.as()) { Var var = ffi::GetRef(var_node); if (!alias_sets.count(var)) { return {-1}; @@ -76,9 +76,9 @@ std::unordered_set GetVarAliasSetFromExpr( // alias id; -1 in the other operand's set does not skip checking other ids. Same var twice // (e.g. add(z, z)) is allowed. bool InplaceArgDisjointFromOtherCallArgs( - const CallNode* call_node, int candidate, + const CallObj* call_node, int candidate, const std::unordered_map>& alias_sets) { - const auto* cand_var_node = call_node->args[candidate].as(); + const auto* cand_var_node = call_node->args[candidate].as(); if (!cand_var_node) { return false; } @@ -117,9 +117,9 @@ std::unordered_map> AnalyzeLiveness(const DataflowBlock ffi::Array used_vars; // for a function literal, we consider only the free vars // (those captured from the outer scope) - if (value.as()) { + if (value.as()) { used_vars = FreeVars(value); - } else if (value.as()) { + } else if (value.as()) { // Special case: we do not consider a tuple index to be a "use." // This is a bit of a hack but allows us to do operations that // create tuples to be done in-place (otherwise, any index of the tuple @@ -133,7 +133,7 @@ std::unordered_map> AnalyzeLiveness(const DataflowBlock int range_end = i; // if the var is not a dataflow var, then it is live // after the block (we are not checking later blocks) - if (!var.as()) { + if (!var.as()) { range_end = block->bindings.size(); } if (!ret.count(var)) { @@ -143,7 +143,7 @@ std::unordered_map> AnalyzeLiveness(const DataflowBlock if (!ret.count(defined_var)) { // if it's an output, then it lives past the end of the block - if (!defined_var.as()) { + if (!defined_var.as()) { ret[defined_var] = {i, block->bindings.size()}; } else { // otherwise, it's live only here @@ -175,7 +175,7 @@ class AliasAnalyzer { for (auto input : inputs) { int curr_idx = get_fresh_idx(); alias_map_[input] = {curr_idx}; - if (auto* tup_info = GetTypeAs(input)) { + if (auto* tup_info = GetTypeAs(input)) { InsertFreshTuple(curr_idx, tup_info); } } @@ -197,12 +197,12 @@ class AliasAnalyzer { } // Fresh tuple = each element is assumed to be a unique allocation - void InsertFreshTuple(int tup_idx, const TupleTypeNode* tup_info) { + void InsertFreshTuple(int tup_idx, const TupleTypeObj* tup_info) { std::vector> tuple_set; for (int i = 0; i < static_cast(tup_info->fields.size()); i++) { int curr_field = get_fresh_idx(); tuple_set.push_back({curr_field}); - if (auto* nested_tup_info = tup_info->fields[i].as()) { + if (auto* nested_tup_info = tup_info->fields[i].as()) { InsertFreshTuple(curr_field, nested_tup_info); } } @@ -249,13 +249,13 @@ class AliasAnalyzer { // For tuples, assume all members are aliased. Yeah, it's bad. // (Skip first arg is for handling call_pure_packed, where the first arg is an ExternFunc that we // should ignore) - std::unordered_set HandleMysteryCall(const CallNode* call_node, const Var& bound_var, + std::unordered_set HandleMysteryCall(const CallObj* call_node, const Var& bound_var, bool skip_first_arg = false) { // the result may or may not be newly allocated std::unordered_set ret; int res_idx = get_fresh_idx(); // the result may be a tuple - if (auto* tup_info_node = GetTypeAs(bound_var)) { + if (auto* tup_info_node = GetTypeAs(bound_var)) { InsertFreshTuple(res_idx, tup_info_node); } AddCapturedIndices(&ret, res_idx); @@ -289,10 +289,10 @@ class AliasAnalyzer { // function constant: give them a fresh index (TODO: we can handle in more detail if this is a // case we need to support) prim value: fresh index if node: should not happen inside dataflow // block - if (value.as() || value.as()) { + if (value.as() || value.as()) { // TODO(@slyubomirsky): We will probably want special handling for closures ret.insert(get_fresh_idx()); - } else if (auto* target_var_node = value.as()) { + } else if (auto* target_var_node = value.as()) { auto target_var = ffi::GetRef(target_var_node); if (alias_map_.count(target_var)) { ret.insert(alias_map_[target_var].begin(), alias_map_[target_var].end()); @@ -301,7 +301,7 @@ class AliasAnalyzer { } } else if (value.as()) { ret.insert(get_fresh_idx()); - } else if (auto* target_tuple = value.as()) { + } else if (auto* target_tuple = value.as()) { // fresh idx but we update the tuple map int tup_idx = get_fresh_idx(); ret.insert(tup_idx); @@ -310,7 +310,7 @@ class AliasAnalyzer { new_tuple_map.push_back(GetAliasSet(field, bound_var)); } tuple_map_[tup_idx] = new_tuple_map; - } else if (auto* target_tgi = value.as()) { + } else if (auto* target_tgi = value.as()) { std::unordered_set tuple_set = GetAliasSet(target_tgi->tuple, bound_var); // if -1 is a member of the tuple set, then we have to assume the result is -1 if (tuple_set.count(-1)) { @@ -327,14 +327,14 @@ class AliasAnalyzer { } } } - } else if (auto* call_node = value.as()) { - if (auto* op_node = call_node->op.as()) { + } else if (auto* call_node = value.as()) { + if (auto* op_node = call_node->op.as()) { // call_pure_packed: treat as non-op call if (op_node->name == "relax.call_pure_packed") { return HandleMysteryCall(call_node, bound_var, true); } else if (op_node->name == "relax.call_tir") { // call_tir: can potentially return a tuple - if (auto* tuple_ty = call_node->ty_args[0].as()) { + if (auto* tuple_ty = call_node->ty_args[0].as()) { int tup_idx = get_fresh_idx(); ret.insert(tup_idx); InsertFreshTuple(tup_idx, tuple_ty); @@ -350,7 +350,7 @@ class AliasAnalyzer { // If the returned value is a tuple, we'll assume it's a fresh tuple // (there may be exceptions to this too) - if (auto* tup_info = GetTypeAs(bound_var)) { + if (auto* tup_info = GetTypeAs(bound_var)) { int tup_idx = get_fresh_idx(); ret.insert(tup_idx); InsertFreshTuple(tup_idx, tup_info); @@ -387,19 +387,19 @@ PrimExpr NumElements(const ShapeExpr& shape) { // all members are eligible though we can consider only individual members separately) std::unordered_set GatherCandidateType( const Type& result_ty) { - if (auto* tensor_info = result_ty.as()) { + if (auto* tensor_info = result_ty.as()) { // don't consider void dtype (don't know the size at compile time) if (tensor_info->IsUnknownDtype()) { return {}; } // don't consider cases where we don't know the shape at compile time // (we will use the analyzer to do best-effort analysis where there are vars) - if (tensor_info->shape.as()) { + if (tensor_info->shape.as()) { return {ffi::GetRef(tensor_info)}; } else { return {}; } - } else if (auto* tuple_info = result_ty.as()) { + } else if (auto* tuple_info = result_ty.as()) { // we can see if the whole tuple matches or go for any of the components std::unordered_set ret; for (auto field : tuple_info->fields) { @@ -424,11 +424,11 @@ std::unordered_set GatherCandidat // if a single pair of corresponding tensors does not meet the condition). std::pair SizeMatches(const Type& target_info, const Type& arg_info, const BlockBuilder& ctx) { - if (target_info.as() && arg_info.as()) { + if (target_info.as() && arg_info.as()) { auto target_tensor = target_info.as_or_throw(); auto arg_tensor = arg_info.as_or_throw(); - if (target_tensor->shape.has_value() && target_tensor->shape.as() && - arg_tensor->shape.has_value() && arg_tensor->shape.as()) { + if (target_tensor->shape.has_value() && target_tensor->shape.as() && + arg_tensor->shape.has_value() && arg_tensor->shape.as()) { if (target_tensor->dtype != arg_tensor->dtype) { return {false, false}; } @@ -452,7 +452,7 @@ std::pair SizeMatches(const Type& target_info, const Type& arg_info, } else { return {false, false}; } - } else if (target_info.as() && arg_info.as()) { + } else if (target_info.as() && arg_info.as()) { auto target_tup = target_info.as_or_throw(); auto arg_tup = arg_info.as_or_throw(); if (target_tup->fields.size() != arg_tup->fields.size()) { @@ -462,9 +462,9 @@ std::pair SizeMatches(const Type& target_info, const Type& arg_info, for (size_t i = 0; i < target_tup->fields.size(); i++) { // if members aren't either tuples or tensors, simply skip them, // since they don't matter for in-place computations - if (!(target_tup->fields[i].as() || - target_tup->fields[i].as()) && - !(arg_tup->fields[i].as() || arg_tup->fields[i].as())) { + if (!(target_tup->fields[i].as() || + target_tup->fields[i].as()) && + !(arg_tup->fields[i].as() || arg_tup->fields[i].as())) { continue; } auto [field_size_match, field_exact_match] = @@ -516,7 +516,7 @@ bool InplaceConditionsMet( const std::unordered_map>& alias_sets, const std::unordered_map>>& tuple_map, const std::unordered_set& currently_live, const Expr& target, int binding_idx) { - if (auto* var_node = target.as()) { + if (auto* var_node = target.as()) { auto current_var = ffi::GetRef(var_node); // if the var is live past this point, we can't use it for in-place computations anyway if (live_ranges.count(current_var)) { @@ -567,7 +567,7 @@ bool InplaceConditionsMet( } } return true; - } else if (auto* tup_node = target.as()) { + } else if (auto* tup_node = target.as()) { for (auto field : tup_node->fields) { if (!InplaceConditionsMet(live_ranges, alias_sets, tuple_map, currently_live, field, binding_idx)) { @@ -686,8 +686,8 @@ FindInplaceOpportunities(const DataflowBlock& block, const ffi::Array& inpu Var defined_var = b->var; Expr value = GetBoundValue(b); - if (auto* call_node = value.as()) { - if (auto* op_node = call_node->op.as()) { + if (auto* call_node = value.as()) { + if (auto* op_node = call_node->op.as()) { if (!OpSupportsInplace(ffi::GetRef(op_node))) { continue; } @@ -789,7 +789,7 @@ class ModuleInplaceTransformer : public ExprMutator { IRModule Transform() { // visit every Relax function in the module for (auto kv : mod_->functions) { - if (auto* func_node = kv.second.as()) { + if (auto* func_node = kv.second.as()) { auto gv = kv.first; auto func_params = func_node->params; auto function = VisitExpr(ffi::GetRef(func_node)).as_or_throw(); @@ -805,7 +805,7 @@ class ModuleInplaceTransformer : public ExprMutator { return ret; } - Expr VisitExpr_(const FunctionNode* op) override { + Expr VisitExpr_(const FunctionObj* op) override { auto old_func_params = func_params; func_params = op->params; auto ret = ExprMutator::VisitExpr_(op); @@ -815,7 +815,7 @@ class ModuleInplaceTransformer : public ExprMutator { // the only case we will override: we will visit all binding blocks // and replace any valid calls in them - BindingBlock VisitBindingBlock_(const DataflowBlockNode* op) override { + BindingBlock VisitBindingBlock_(const DataflowBlockObj* op) override { auto block = ffi::GetRef(op); auto old_idxs = inplace_idxs; @@ -843,7 +843,7 @@ class ModuleInplaceTransformer : public ExprMutator { return builder_->Normalize(new_call); } - void VisitBinding_(const VarBindingNode* binding) override { + void VisitBinding_(const VarBindingObj* binding) override { auto binding_ref = ffi::GetRef(binding); if (!inplace_idxs.count(binding_ref)) { ExprMutator::VisitBinding_(binding); @@ -853,7 +853,7 @@ class ModuleInplaceTransformer : public ExprMutator { builder_->EmitNormalized(VarBinding(binding->var, new_value, binding->loc)); } - void VisitBinding_(const MatchCastNode* binding) override { + void VisitBinding_(const MatchCastObj* binding) override { auto binding_ref = ffi::GetRef(binding); if (!inplace_idxs.count(binding_ref)) { ExprMutator::VisitBinding_(binding); diff --git a/src/relax/transform/dead_code_elimination.cc b/src/relax/transform/dead_code_elimination.cc index 6bf82d6905f7..7cfb7f70b1d5 100644 --- a/src/relax/transform/dead_code_elimination.cc +++ b/src/relax/transform/dead_code_elimination.cc @@ -53,7 +53,7 @@ namespace { struct RelaxCalleeCollector : relax::ExprVisitor { std::vector* callees; explicit RelaxCalleeCollector(std::vector* out) : callees(out) {} - void VisitExpr_(const GlobalVarNode* node) final { + void VisitExpr_(const GlobalVarObj* node) final { callees->push_back(ffi::GetRef(node)); } }; @@ -62,7 +62,7 @@ struct TIRxCalleeCollector : tirx::StmtExprVisitor { std::vector* callees; explicit TIRxCalleeCollector(std::vector* out) : callees(out) {} using tirx::StmtExprVisitor::Visit_; - ffi::Optional Visit_(const CallNode* node) final { + ffi::Optional Visit_(const CallObj* node) final { TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(tirx::StmtExprVisitor::Visit_(node)); if (auto opt_gvar = node->op.as()) { callees->push_back(opt_gvar.value()); @@ -162,7 +162,7 @@ IRModule DeadCodeElimination(const IRModule& arg_mod, } for (const auto& gv : mod->GetGlobalVars()) { const auto& func = mod->Lookup(gv); - if (func.as() || func->GetLinkageType() == LinkageType::kExternal) { + if (func.as() || func->GetLinkageType() == LinkageType::kExternal) { entry_functions.insert(gv); } } diff --git a/src/relax/transform/decompose_ops.cc b/src/relax/transform/decompose_ops.cc index ef720fdc7f2d..2e0684e21683 100644 --- a/src/relax/transform/decompose_ops.cc +++ b/src/relax/transform/decompose_ops.cc @@ -147,7 +147,7 @@ Expr DecomposeLayerNorm(const Call& call) { Expr TensorToShape(const Call& call_node, const BlockBuilder& builder) { TVM_FFI_ICHECK(!call_node->ty.as().has_value()); Expr expr = call_node->args[0]; - const ShapeTypeNode* ty = GetTypeAs(call_node); + const ShapeTypeObj* ty = GetTypeAs(call_node); TVM_FFI_ICHECK(ty); // call builtin function that converts tensor to shape tuple // TODO(@sunggg): Register operator for "vm.builtin.tensor_to_shape" @@ -178,7 +178,7 @@ class TrainingOperatorMutator : public ExprMutator { private: using ExprMutator::VisitExpr_; - Expr VisitExpr_(const CallNode* call_node) final { + Expr VisitExpr_(const CallObj* call_node) final { Call call = VisitExprPostOrder_(call_node).as_or_throw(); if (call->op.same_as(batch_norm_op_)) { return MutateBatchNormForTraining(call); @@ -200,7 +200,7 @@ class OpDecomposer : public ExprMutator { private: using ExprMutator::VisitExpr_; - Expr VisitExpr_(const CallNode* call_node) final { + Expr VisitExpr_(const CallObj* call_node) final { Call call = VisitExprPostOrder_(call_node).as_or_throw(); if (call->op.same_as(batch_norm_op_)) { return DecomposeBatchNorm(call); diff --git a/src/relax/transform/eliminate_common_subexpr.cc b/src/relax/transform/eliminate_common_subexpr.cc index 1a0e603df6ce..b146b3d67b4c 100644 --- a/src/relax/transform/eliminate_common_subexpr.cc +++ b/src/relax/transform/eliminate_common_subexpr.cc @@ -24,7 +24,7 @@ * * Currently it removes common subexpressions within a Function. */ -#include +#include #include #include #include @@ -56,7 +56,7 @@ struct ReplacementKey { explicit ReplacementKey(const tvm::relax::Binding& binding) : bound_value(GetBoundValue(binding)) { - if (const auto* ptr = binding.as()) { + if (const auto* ptr = binding.as()) { match_cast = ptr->ty; } } @@ -81,7 +81,7 @@ template <> struct std::hash { std::size_t operator()(const tvm::relax::ReplacementKey& key) const { tvm::ffi::StructuralHash hasher; - return tvm::ffi::details::StableHashCombine(hasher(key.bound_value), hasher(key.match_cast)); + return tvm::ffi::detail::StableHashCombine(hasher(key.bound_value), hasher(key.match_cast)); } }; @@ -94,14 +94,14 @@ class CommonSubexprEliminator : public ExprMutator { public: explicit CommonSubexprEliminator(bool call_only = false) : call_only_(call_only) {} - BindingBlock VisitBindingBlock_(const DataflowBlockNode* block) override { + BindingBlock VisitBindingBlock_(const DataflowBlockObj* block) override { auto cache_vars = var_remap_; auto output = ExprMutator::VisitBindingBlock_(block); for (auto& [key, replacements] : expr_replacements_) { replacements.erase( std::remove_if(replacements.begin(), replacements.end(), - [](const Var& var) -> bool { return var->IsInstance(); }), + [](const Var& var) -> bool { return var->IsInstance(); }), replacements.end()); } @@ -113,9 +113,9 @@ class CommonSubexprEliminator : public ExprMutator { Expr bound_value = VisitExpr(GetBoundValue(binding)); Binding output_binding = [&]() -> Binding { - if (binding.as()) { + if (binding.as()) { return VarBinding(binding->var, bound_value); - } else if (auto match_cast = binding.as()) { + } else if (auto match_cast = binding.as()) { return MatchCast(binding->var, bound_value, match_cast->ty); } else { TVM_FFI_THROW(InternalError) << "Binding must be either VarBinding or MatchCast, " @@ -125,7 +125,7 @@ class CommonSubexprEliminator : public ExprMutator { ReplacementKey lookup_key(output_binding); - if (call_only_ && !bound_value->IsInstance()) { + if (call_only_ && !bound_value->IsInstance()) { VLOG(1) << "Since call_only_ is true, it is forbidden to de-duplicate " << bound_value; } else if (ContainsImpureCall(bound_value)) { @@ -154,7 +154,7 @@ class CommonSubexprEliminator : public ExprMutator { builder_->EmitNormalized(output_binding); } - Expr VisitExpr_(const FunctionNode* op) override { + Expr VisitExpr_(const FunctionObj* op) override { // If we have accumulated any state, visit the function in a fresh // copy of the mutator, to avoid replacing a child-scope // expression with a parent-scope binding, or vice versa. @@ -165,7 +165,7 @@ class CommonSubexprEliminator : public ExprMutator { } } - Expr VisitExpr_(const IfExprNode* op) override { + Expr VisitExpr_(const IfExprObj* op) override { Expr cond = VisitExpr(op->cond); Expr true_branch = VisitWithInnerScope(op->true_branch); Expr false_branch = VisitWithInnerScope(op->false_branch); @@ -194,8 +194,8 @@ class CommonSubexprEliminator : public ExprMutator { bool IsAllocatorCall(const Expr& expr) { static const auto& allocator_attr_map = Op::GetAttrMap(tvm::relax::op_attr::kAllocator); - if (const auto* call = expr.as()) { - if (const auto* op = call->op.as()) { + if (const auto* call = expr.as()) { + if (const auto* op = call->op.as()) { bool is_allocator = allocator_attr_map.get(ffi::GetRef(op), false); if (is_allocator) { return true; diff --git a/src/relax/transform/expand_tuple_arguments.cc b/src/relax/transform/expand_tuple_arguments.cc index d25fa4bb88df..09089a15d027 100644 --- a/src/relax/transform/expand_tuple_arguments.cc +++ b/src/relax/transform/expand_tuple_arguments.cc @@ -39,7 +39,7 @@ ffi::Optional ExpandParams(Function func) { bool has_tuple_param = std::any_of(func->params.begin(), func->params.end(), - [](const Var& param) -> bool { return param->ty.as(); }); + [](const Var& param) -> bool { return param->ty.as(); }); if (!has_tuple_param) return std::nullopt; @@ -47,7 +47,7 @@ ffi::Optional ExpandParams(Function func) { ffi::Array bindings; std::function expand_param = [&](const Var& param) { - if (auto ty = param->ty.as()) { + if (auto ty = param->ty.as()) { ffi::Array internal_tuple; for (size_t i = 0; i < ty->fields.size(); i++) { auto name = @@ -84,7 +84,7 @@ class TupleExpander : public ExprMutator { using ExprMutator::VisitExpr_; - Expr VisitExpr_(const CallNode* op) override { + Expr VisitExpr_(const CallObj* op) override { auto node = ExprMutator::VisitExpr_(op).as_or_throw(); if (auto gvar = node->op.as()) { @@ -92,7 +92,7 @@ class TupleExpander : public ExprMutator { ffi::Array new_args; std::function expand_arg = [&](const Expr& arg) { - if (auto ty = arg->ty.as()) { + if (auto ty = arg->ty.as()) { for (size_t i = 0; i < ty->fields.size(); i++) { expand_arg(TupleGetItem(arg, i)); } diff --git a/src/relax/transform/fold_constant.cc b/src/relax/transform/fold_constant.cc index 3fb85faa7c9d..78f5b3f07b65 100644 --- a/src/relax/transform/fold_constant.cc +++ b/src/relax/transform/fold_constant.cc @@ -55,17 +55,17 @@ class ConstantFolder : public ExprMutator { * if the input type is not TensorType. */ static ffi::Optional MatchConstShape(const Type& ty) { - const auto* tensor_ty = ty.as(); + const auto* tensor_ty = ty.as(); if (tensor_ty == nullptr) { return std::nullopt; } - const auto* shape = tensor_ty->shape.as(); + const auto* shape = tensor_ty->shape.as(); TVM_FFI_ICHECK(shape != nullptr) << "type given by call_tir should have ShapeExpr shape"; std::vector shape_values; for (const auto v : shape->values) { - auto* ptr = v.as(); + auto* ptr = v.as(); if (!ptr) return std::nullopt; auto value = ptr->value.as(); if (!value.has_value()) return std::nullopt; @@ -81,7 +81,7 @@ class ConstantFolder : public ExprMutator { static ffi::Optional> MatchConstArrayArgs(const ffi::Array& args) { ffi::Array res; for (auto arg : args) { - auto* ptr = arg.as<::tvm::GenericConstNode>(); + auto* ptr = arg.as<::tvm::GenericConstObj>(); if (!ptr || !ptr->value.as()) return std::nullopt; res.push_back(ptr->value.cast()); } @@ -96,7 +96,7 @@ class ConstantFolder : public ExprMutator { const GlobalVar& global_var = op.as_or_throw(); // NOTE: as check works for nullptr(returns null) ffi::Optional base_func = builder_->GetContextIRModule()->functions.Get(global_var); - if (auto* pfunc = base_func.as()) { + if (auto* pfunc = base_func.as()) { return ffi::GetRef(pfunc); } return std::nullopt; @@ -142,10 +142,10 @@ class ConstantFolder : public ExprMutator { * of the program. */ static bool ExprContainsTensor(const Expr& expr) { - if (GetType(expr).as()) { + if (GetType(expr).as()) { return true; } - if (const auto* tuple = expr.as()) { + if (const auto* tuple = expr.as()) { for (const auto& field : tuple->fields) { if (ExprContainsTensor(field)) { return true; @@ -161,10 +161,10 @@ class ConstantFolder : public ExprMutator { // and folding them would materialize large constants in the binary. static constexpr int64_t kMaxFoldElements = 1024; - const auto* call = expr.as(); + const auto* call = expr.as(); if (!call) return true; - const auto* tensor_ty = call->ty.as(); + const auto* tensor_ty = call->ty.as(); if (!tensor_ty) return true; auto opt_shape = tensor_ty->GetShape(); @@ -172,7 +172,7 @@ class ConstantFolder : public ExprMutator { int64_t num_elements = 1; for (const auto& dim : opt_shape.value()) { - const auto* int_dim = dim.as(); + const auto* int_dim = dim.as(); if (!int_dim) return true; auto d = int_dim->value.as(); if (int_dim->value <= 0) return true; @@ -232,7 +232,7 @@ class ConstantFolder : public ExprMutator { // Returns std::nullopt on failure. ffi::Optional ConstEvaluateCallTIRTuple(tirx::Function tir_func, ffi::Array arr_args, - const TupleTypeNode* tuple_ty) { + const TupleTypeObj* tuple_ty) { ffi::Optional func = GetCachedBuild(tir_func); if (!func) return std::nullopt; @@ -276,15 +276,15 @@ class ConstantFolder : public ExprMutator { // call_tir needs to have at least two arguments TVM_FFI_ICHECK_GE(call->args.size(), 2); ffi::Optional func = MatchFunction(call->args[0]); - TVM_FFI_ICHECK(call->args[1].as()) << "call_tir.args[1] must be Tuple"; + TVM_FFI_ICHECK(call->args[1].as()) << "call_tir.args[1] must be Tuple"; ffi::Optional> arr_args = - MatchConstArrayArgs(call->args[1].as()->fields); + MatchConstArrayArgs(call->args[1].as()->fields); TVM_FFI_ICHECK_EQ(call->ty_args.size(), 1) << "call_tir should have exactly one ty arg"; if (!func || !arr_args) return {}; // Handle tuple output: ty_args[0] is a TupleType. - if (const auto* tuple_ty = call->ty_args[0].as()) { + if (const auto* tuple_ty = call->ty_args[0].as()) { return ConstEvaluateCallTIRTuple(func.value(), arr_args.value(), tuple_ty); } @@ -304,7 +304,7 @@ class ConstantFolder : public ExprMutator { // Next PR will support fold with ffi::Function and MatchCast // Until then, DecomposeOps() should be applied after // this pass to fold `tensor_to_shape` op. - Expr VisitExpr_(const CallNode* call) final { + Expr VisitExpr_(const CallObj* call) final { // post-order mutation Call post_call = VisitExprPostOrder_(call).as_or_throw(); @@ -316,9 +316,9 @@ class ConstantFolder : public ExprMutator { static const auto& infer_type_with_builder_map = Op::GetAttrMap(tvm::relax::op_attr::kInferTypeWithBuilder); static const auto& legalize_map = Op::GetAttrMap(tvm::relax::op_attr::kLegalize); - auto* op_node = post_call->op.as(); + auto* op_node = post_call->op.as(); - // Not an OpNode + // Not an OpObj if (op_node == nullptr) { return post_call; } @@ -338,9 +338,9 @@ class ConstantFolder : public ExprMutator { // ffi::Array new_args; for (auto arg : post_call->args) { - if (arg->IsInstance()) { + if (arg->IsInstance()) { ffi::Optional val = LookupBinding(arg.as_or_throw()); - if (val.has_value() && val.value()->IsInstance()) { + if (val.has_value() && val.value()->IsInstance()) { new_args.push_back(val.value()); continue; } @@ -348,7 +348,7 @@ class ConstantFolder : public ExprMutator { new_args.push_back(arg); } Type ret_ty = Type::Missing(); - if (post_call->ty.as() && !infer_type_map.count(op) && + if (post_call->ty.as() && !infer_type_map.count(op) && !infer_type_with_builder_map.count(op)) { ret_ty = post_call->ty.as_or_throw(); } @@ -363,7 +363,7 @@ class ConstantFolder : public ExprMutator { Call post_call_normalized = builder_->Normalize(post_call).as_or_throw(); Expr legalized_expr = builder_->Normalize(legalize_map[op](builder_, post_call_normalized)); // If the legalized expression is call_tir, try to fold it. - const CallNode* call = legalized_expr.as(); + const CallObj* call = legalized_expr.as(); if (call && call->op.same_as(call_tir_op)) { return VisitCallTIR(ffi::GetRef(call)).value_or(post_call); } @@ -377,7 +377,7 @@ class ConstantFolder : public ExprMutator { // decomposition map for each op in a similar way we do for legalization. TVM_FFI_ICHECK_EQ(post_call->args.size(), 1); Expr arg = post_call->args[0]; - if (arg->IsInstance()) { + if (arg->IsInstance()) { GenericConst constant = arg.as_or_throw(); ffi::Tensor ndarray = constant->value.cast(); TVM_FFI_ICHECK(ndarray.defined()); @@ -407,7 +407,7 @@ class ConstantFolder : public ExprMutator { is_known = false; break; } - arr.push_back(static_cast(val.as()->value)); + arr.push_back(static_cast(val.as()->value)); } if (is_known) { const auto func = tvm::ffi::Function::GetGlobalRequired("relax.run.shape_to_tensor"); @@ -421,10 +421,10 @@ class ConstantFolder : public ExprMutator { return post_call; } - Expr VisitExpr_(const VarNode* op) final { + Expr VisitExpr_(const VarObj* op) final { ffi::Optional opt = LookupBinding(ffi::GetRef(op)); // `as` check checks if opt is not null and is instance of constant - if (opt.as<::tvm::GenericConstNode>()) { + if (opt.as<::tvm::GenericConstObj>()) { return opt.value(); } return ExprMutator::VisitExpr_(op); diff --git a/src/relax/transform/fuse_ops.cc b/src/relax/transform/fuse_ops.cc index a681756227e5..ec639fb837b8 100644 --- a/src/relax/transform/fuse_ops.cc +++ b/src/relax/transform/fuse_ops.cc @@ -131,7 +131,7 @@ class GraphCreator : public ExprVisitor { // `attr::kCodegen.name` are previously fused functions from // `FuseOpsByPattern`, when the `annotate_codegen` option is // true. - const auto* func = it.second.as(); + const auto* func = it.second.as(); if (func == nullptr || func->HasNonzeroAttr(tvm::relax::attr::kPrimitive.name) || func->GetAttr(tvm::relax::attr::kCodegen).has_value()) { continue; @@ -153,7 +153,7 @@ class GraphCreator : public ExprVisitor { explicit GraphCreator(IRModule mod, support::Arena* arena) : mod_(std::move(mod)), arena_(arena) {} - void VisitExpr_(const FunctionNode* func) final { + void VisitExpr_(const FunctionObj* func) final { for (const Var& param : func->params) { IndexedForwardGraph::Node* param_node = CreateNode(param.get()); // The parameter is passed in from the outside, and thus it's marked as an external reference, @@ -171,25 +171,25 @@ class GraphCreator : public ExprVisitor { ExprVisitor::VisitExpr_(func); } - void VisitBinding_(const MatchCastNode* binding) final { + void VisitBinding_(const MatchCastObj* binding) final { IndexedForwardGraph::Node* node = CreateNode(binding->var.get()); SetNodePattern(node, OpPatternKind::kOpaque); AddToPostDFSOrder(node, binding->var.get()); } - void VisitBinding_(const VarBindingNode* binding) final { + void VisitBinding_(const VarBindingObj* binding) final { IndexedForwardGraph::Node* node = CreateNode(binding->var.get()); // If the variable is not a dataflow variable, it must be the output variable of this dataflow // block - if (!binding->var->IsInstance()) { + if (!binding->var->IsInstance()) { this->MarkAsExternRef(node); } - if (const auto* call = binding->value.as()) { - // Case 1. The expression is a CallNode + if (const auto* call = binding->value.as()) { + // Case 1. The expression is a CallObj VisitCall(call, node); - } else if (const auto* tuple_get_item = binding->value.as()) { - // Case 2. The expression is a TupleGetItemNode + } else if (const auto* tuple_get_item = binding->value.as()) { + // Case 2. The expression is a TupleGetItemObj VisitTupleGetItem(tuple_get_item, node); } else { VisitUnsupportedNode(binding->value, node); @@ -201,7 +201,7 @@ class GraphCreator : public ExprVisitor { /********** Non-Leaf Expression Nodes **********/ - void VisitCall(const CallNode* call, IndexedForwardGraph::Node* binding_var_node) { + void VisitCall(const CallObj* call, IndexedForwardGraph::Node* binding_var_node) { TVM_FFI_ICHECK_NOTNULL(binding_var_node); static const Op call_tir_op_ = Op::Get("relax.call_tir"); @@ -214,7 +214,7 @@ class GraphCreator : public ExprVisitor { // from the function attribute and visit the arguments one by one. // - Otherwise, the pattern of the current binding variable node is set to `kOpaque`, and we // recurse into the call expression. - const auto* op = call->op.as(); + const auto* op = call->op.as(); if (op == call_tir_op_.get() || op == call_tir_inplace_op_.get()) { const GlobalVar& global_var = call->args[0].as_or_throw(); tirx::Function func = mod_->Lookup(global_var).as_or_throw(); @@ -241,12 +241,12 @@ class GraphCreator : public ExprVisitor { } } - void VisitTupleGetItem(const TupleGetItemNode* tuple_item, + void VisitTupleGetItem(const TupleGetItemObj* tuple_item, IndexedForwardGraph::Node* binding_var_node) { TVM_FFI_ICHECK_NOTNULL(binding_var_node); auto pattern = OpPatternKind::kInjective; - if (input_params_.count(tuple_item->tuple.as())) { + if (input_params_.count(tuple_item->tuple.as())) { // TupleGetItem for fetching the parameter from the packed param tuple is treated as opaque // and won't be fused. This prevents the usage of packed param tuple changes the order of the // fusion result as the function usually begins with fetching the parameters. @@ -261,7 +261,7 @@ class GraphCreator : public ExprVisitor { SetNodePattern(binding_var_node, OpPatternKind::kOpaque); auto visit_leaves = [this, &binding_var_node](const Expr& e) { - if (e->IsInstance() || e->IsInstance()) { + if (e->IsInstance() || e->IsInstance()) { VisitLeaf(e, binding_var_node, OpPatternKind::kOpaque); } }; @@ -275,17 +275,17 @@ class GraphCreator : public ExprVisitor { TVM_FFI_ICHECK_NOTNULL(binding_var_node); // Recursive visit if it's Tuple - if (const auto* tuple = leaf_expr.as()) { + if (const auto* tuple = leaf_expr.as()) { for (const Expr& expr : tuple->fields) { VisitLeaf(expr, binding_var_node, pattern); } return; } - if (!leaf_expr.as() && !leaf_expr.as() && - !leaf_expr.as() && !leaf_expr.as() && - !leaf_expr.as() && !leaf_expr.as()) { - // Skip GlobalVar, ExternFunc, OpNode. + if (!leaf_expr.as() && !leaf_expr.as() && + !leaf_expr.as() && !leaf_expr.as() && + !leaf_expr.as() && !leaf_expr.as()) { + // Skip GlobalVar, ExternFunc, OpObj. return; } @@ -386,7 +386,7 @@ class GraphCreator : public ExprVisitor { /*! \brief The graph nodes whose patterns are set */ std::unordered_set initialized_nodes_; /*! \brief The model params in the function input */ - std::unordered_set input_params_; + std::unordered_set input_params_; }; /*! @@ -418,8 +418,8 @@ class FunctionCreator : public ExprMutator { TVM_FFI_ICHECK(!function_.has_value()) << "The `function_` is supposed to be uncreated when adding bindings"; - if (const auto* var_binding = binding.as()) { - if (const auto* call = var_binding->value.as()) { + if (const auto* var_binding = binding.as()) { + if (const auto* call = var_binding->value.as()) { if (call->op.same_as(Op::Get("relax.call_tir")) || call->op.same_as(Op::Get("relax.call_tir_inplace"))) { // Update the name of the function. @@ -428,13 +428,13 @@ class FunctionCreator : public ExprMutator { const Tuple& args = call->args[1].as_or_throw(); for (const Expr& arg : args->fields) { CheckDefAndUpdateParam(arg); - TVM_FFI_ICHECK(GetTypeAs(arg) == nullptr); + TVM_FFI_ICHECK(GetTypeAs(arg) == nullptr); } // TODO(tvm-team): handle shape expr } else { - if (call->op->IsInstance()) { + if (call->op->IsInstance()) { name_hint_ = name_hint_ + "_" + call->op.as_or_throw()->name; - } else if (call->op->IsInstance()) { + } else if (call->op->IsInstance()) { std::string gvar_name = call->op.as_or_throw()->name_hint; if (auto pos = gvar_name.find("fused_"); pos == 0) { name_hint_ = name_hint_ + "_" + gvar_name.substr(std::string("fused_").size()); @@ -444,22 +444,22 @@ class FunctionCreator : public ExprMutator { } for (const Expr& arg : call->args) { - if (auto tuple = arg.as()) { + if (auto tuple = arg.as()) { for (const Expr& tup_arg : tuple->fields) { CheckDefAndUpdateParam(tup_arg); - TVM_FFI_ICHECK(GetTypeAs(tup_arg) == nullptr); + TVM_FFI_ICHECK(GetTypeAs(tup_arg) == nullptr); } } else { CheckDefAndUpdateParam(arg); } - if (GetTypeAs(arg) != nullptr) { + if (GetTypeAs(arg) != nullptr) { // The argument is fully referenced. Thus we remove it from the mapping. partially_used_tuple_params_.erase(arg.get()); } } } - } else if (var_binding->value.as()) { - const auto* tuple_item = var_binding->value.as(); + } else if (var_binding->value.as()) { + const auto* tuple_item = var_binding->value.as(); CheckDefAndUpdateParam(tuple_item->tuple); if (partially_used_tuple_params_.find(tuple_item->tuple.get()) != @@ -472,7 +472,7 @@ class FunctionCreator : public ExprMutator { // Mark the binding variable as defined. defined_vars_.insert(var_binding->var.get()); // Set var as output true if the binding is not a dataflow variable - if (!var_binding->var->IsInstance()) { + if (!var_binding->var->IsInstance()) { AppendOutput(var_binding->var); } } else { @@ -489,7 +489,7 @@ class FunctionCreator : public ExprMutator { } /*! \brief Variables returned from the grouped function, in return-value order. */ - const std::vector& output_vars() const { return output_vars_; } + const std::vector& output_vars() const { return output_vars_; } /*! * \brief Create the grouped function according to the collected bindings and parameters @@ -504,7 +504,7 @@ class FunctionCreator : public ExprMutator { // Step 2. Handing partially used tuple parameters: replacing entire tuple // parameters with the parameters of its fields that are accessed in the // function. - std::unordered_map> tuple_get_item_remap; + std::unordered_map> tuple_get_item_remap; for (auto& [tuple_arg, item_indices] : partially_used_tuple_params_) { TVM_FFI_ICHECK(!item_indices.empty()); int param_idx = tuple_param_idx_[tuple_arg]; @@ -536,8 +536,8 @@ class FunctionCreator : public ExprMutator { std::vector> output_slots(output_vars_.size()); for (const Binding& binding : bindings_) { // Special handing for TupleGetItem. - if (const auto* var_binding = binding.as()) { - if (const auto* tuple_get_item = var_binding->value.as()) { + if (const auto* var_binding = binding.as()) { + if (const auto* tuple_get_item = var_binding->value.as()) { auto it = tuple_get_item_remap.find(tuple_get_item->tuple.get()); if (it != tuple_get_item_remap.end()) { TVM_FFI_ICHECK(it->second.find(tuple_get_item->index) != it->second.end()); @@ -553,7 +553,7 @@ class FunctionCreator : public ExprMutator { if (auto output_idx = GetOutputIndex(binding->var)) { // Case 1. It is an output binding // We only allow VarBinding as output. - const auto* var_binding = binding.as(); + const auto* var_binding = binding.as(); TVM_FFI_ICHECK_NOTNULL(var_binding); Var output_var = builder_->EmitOutput(VisitExpr(var_binding->value)); var_remap_.insert_or_assign(var_binding->var, output_var); @@ -625,8 +625,8 @@ class FunctionCreator : public ExprMutator { Expr ResolveOuterBinding(const Expr& expr) { Expr bound_value = expr; - std::unordered_set visited; - while (const auto* current_var = bound_value.as()) { + std::unordered_set visited; + while (const auto* current_var = bound_value.as()) { if (!visited.insert(current_var).second) break; auto it = outer_bindings_.find(ffi::GetRef(current_var)); if (it == outer_bindings_.end()) break; @@ -637,7 +637,7 @@ class FunctionCreator : public ExprMutator { void InlineShapeDependentPrimArgs() { ffi::Array boundary_types = params_.Map([](const Var& param) { return GetType(param); }); - for (const VarNode* output_var : output_vars_) { + for (const VarObj* output_var : output_vars_) { boundary_types.push_back(GetType(ffi::GetRef(output_var))); } @@ -659,7 +659,7 @@ class FunctionCreator : public ExprMutator { if (!inline_argument && !inline_bound_value) continue; if (inline_bound_value) { - const auto* argument_var = argument.as(); + const auto* argument_var = argument.as(); TVM_FFI_ICHECK(argument_var); inlined_bindings_.insert_or_assign(argument_var, bound_value); } @@ -685,7 +685,7 @@ class FunctionCreator : public ExprMutator { // If the expression is not a variable or is a undefined variable, it should be populated as a // parameter of the relax function. - const auto* var = expr.as(); + const auto* var = expr.as(); if (var != nullptr && defined_vars_.count(var) == 0) { Var bound_var = ffi::GetRef(var); Expr bound_value = ResolveOuterBinding(bound_var); @@ -695,7 +695,7 @@ class FunctionCreator : public ExprMutator { } } if ((var == nullptr || defined_vars_.count(var) == 0) && - (lift_constant_ || !expr->IsInstance())) { + (lift_constant_ || !expr->IsInstance())) { ffi::String name = var != nullptr ? var->name : ffi::String("param_" + std::to_string(n_param_for_const_++)); Type param_ty = GetType(expr); @@ -710,7 +710,7 @@ class FunctionCreator : public ExprMutator { // Mark the tuple parameter is partially referenced in the beginning. // We will remove it from the mapping once we find it is fully referenced. - if (param_ty->IsInstance()) { + if (param_ty->IsInstance()) { partially_used_tuple_params_[expr.get()] = {}; tuple_param_idx_[expr.get()] = static_cast(arguments_.size()) - 1; } @@ -724,7 +724,7 @@ class FunctionCreator : public ExprMutator { if (it != arguments_.end()) { return params_[it - arguments_.begin()]; } - if (const auto* var = expr.as()) { + if (const auto* var = expr.as()) { if (auto inlined = inlined_bindings_.find(var); inlined != inlined_bindings_.end()) { return inlined->second; } @@ -736,7 +736,7 @@ class FunctionCreator : public ExprMutator { // Check if the expression is constant PrimExpr or ShapeExpr or tuple of them that can be // inlined in the composite functions and excluded from args/params. bool IsSymbolicPrimExpr(const Expr& expr) { - if (expr.as()) return false; + if (expr.as()) return false; if (auto prim_value = expr.as()) { return !tvm::UndefinedVars(prim_value.value()).empty(); } @@ -752,14 +752,14 @@ class FunctionCreator : public ExprMutator { } bool IsInlinableConstants(const Expr& expr) { - if (const auto* tuple = expr.as()) { + if (const auto* tuple = expr.as()) { return std::all_of(tuple->fields.begin(), tuple->fields.end(), [this](const Expr& e) { return IsInlinableConstants(e); }); - } else if (expr.as() || expr.as()) { + } else if (expr.as() || expr.as()) { return false; } else if (auto prim_value = expr.as()) { return tvm::UndefinedVars(prim_value.value()).empty(); - } else if (const auto* shape_expr = expr.as()) { + } else if (const auto* shape_expr = expr.as()) { return std::all_of(shape_expr->values.begin(), shape_expr->values.end(), [](const PrimExpr& e) { return tvm::UndefinedVars(e).empty(); }); } @@ -768,26 +768,26 @@ class FunctionCreator : public ExprMutator { private: /*! \brief The variables defined in this function */ - std::unordered_set defined_vars_; + std::unordered_set defined_vars_; /*! \brief Caller variables replaced by statically inlinable bound values. */ - std::unordered_map inlined_bindings_; + std::unordered_map inlined_bindings_; /*! \brief Symbolic primitive arguments classified after the complete boundary is known. */ ffi::Array deferred_prim_args_; /*! \brief The number of parameters reserved for constants */ int n_param_for_const_ = 0; /*! \brief The output vars */ - std::vector output_vars_; + std::vector output_vars_; /*! \brief Bindings in the caller function, used to inline static leaf expressions. */ ffi::Map outer_bindings_; /*! \brief Whether or not to lift bound constants to parameters */ bool lift_constant_; /*! \brief Mapping from tuple parameter of the function to its position index */ - std::unordered_map tuple_param_idx_; + std::unordered_map tuple_param_idx_; /*! * \brief Mapping from partially referenced tuple parameter to the list of * indices that the parameter is referred by TupleGetItem */ - std::unordered_map> partially_used_tuple_params_; + std::unordered_map> partially_used_tuple_params_; }; /*! @@ -849,7 +849,7 @@ class OperatorFusor : public ExprMutator { const auto& func = mod_->Lookup(gv); // Only visit Relax functions with neither attr::kPrimitive.name nor // attr::kCodegen.name. - if (func->IsInstance() && + if (func->IsInstance() && !func->HasNonzeroAttr(tvm::relax::attr::kPrimitive.name) && !func->GetAttr(tvm::relax::attr::kCodegen).has_value()) { outer_bindings_ = AnalyzeVar2Value(func); @@ -874,7 +874,7 @@ class OperatorFusor : public ExprMutator { return obj2group; } - BindingBlock VisitBindingBlock_(const DataflowBlockNode* block) final { + BindingBlock VisitBindingBlock_(const DataflowBlockObj* block) final { group2func_.clear(); // Step 1. Collect the bindings for each grouped function. @@ -937,7 +937,7 @@ class OperatorFusor : public ExprMutator { } // Case 3. The binding is the last binding of the group. - const auto* var_binding = binding.as(); + const auto* var_binding = binding.as(); TVM_FFI_ICHECK(var_binding != nullptr) << "The last binding of a group whose size is larger than 1 " "is supposed to be a variable binding"; @@ -952,7 +952,7 @@ class OperatorFusor : public ExprMutator { TVM_FFI_ICHECK(!output_vars.empty()); Call call_to_emit = Call(Type::Missing(), gv, UpdateArgs(func_info.arguments_)); - Var new_var = output_vars.size() == 1 && !output_vars[0]->IsInstance() + Var new_var = output_vars.size() == 1 && !output_vars[0]->IsInstance() ? builder_->EmitOutput(call_to_emit) : builder_->Emit(call_to_emit); @@ -964,13 +964,13 @@ class OperatorFusor : public ExprMutator { continue; } - std::unordered_set remapped_outputs; + std::unordered_set remapped_outputs; auto remap_output = [&](const Var& output_var) { auto it = std::find(output_vars.begin(), output_vars.end(), output_var.get()); TVM_FFI_ICHECK(it != output_vars.end()); int index = static_cast(std::distance(output_vars.begin(), it)); TupleGetItem tuple_get(new_var, index); - if (output_var->IsInstance()) { + if (output_var->IsInstance()) { var_remap_.insert_or_assign(output_var, builder_->Emit(tuple_get)); } else { var_remap_.insert_or_assign(output_var, builder_->EmitOutput(tuple_get)); @@ -983,7 +983,7 @@ class OperatorFusor : public ExprMutator { remap_output(output_var); } } - for (const VarNode* output_var : output_vars) { + for (const VarObj* output_var : output_vars) { if (!remapped_outputs.count(output_var)) { remap_output(ffi::GetRef(output_var)); } @@ -1023,7 +1023,7 @@ class OperatorFusor : public ExprMutator { // - If the var's group is different with the binding's, the var must be the output from // another group. Mark it to be the group output. auto update_boundary = [this, binding, &cur_group](const Expr& e) { - if (e->IsInstance() && obj2group_.count(e.get())) { + if (e->IsInstance() && obj2group_.count(e.get())) { const Var& used_var = e.as_or_throw(); Group* producer_group = GetGroupFromVar(used_var); // Only check those group defined before. @@ -1044,10 +1044,10 @@ class OperatorFusor : public ExprMutator { } }; - if (const auto* var_binding = binding.as()) { + if (const auto* var_binding = binding.as()) { PostOrderVisit(var_binding->value, update_boundary); } else { - const auto* match_cast = binding.as(); + const auto* match_cast = binding.as(); TVM_FFI_ICHECK_NOTNULL(match_cast); PostOrderVisit(match_cast->value, update_boundary); } @@ -1204,7 +1204,7 @@ class PatternBasedPartitioner : ExprVisitor { arena_(arena), attrs_getter_(attrs_getter) {} - void VisitBindingBlock_(const DataflowBlockNode* block) final { + void VisitBindingBlock_(const DataflowBlockObj* block) final { current_block_use_def_ = DataflowBlockUseDef(ffi::GetRef(block)); ExprVisitor::VisitBindingBlock_(block); current_block_use_def_ = {}; @@ -1216,16 +1216,16 @@ class PatternBasedPartitioner : ExprVisitor { vars_in_group_[g].push_back(var); } - void VisitBinding_(const VarBindingNode* binding) final { + void VisitBinding_(const VarBindingObj* binding) final { bindings_.Set(binding->var, binding->value); value_to_bound_var_.Set(binding->value, binding->var); ExprVisitor::VisitBinding_(binding); } - void VisitExpr_(const GenericConstNode* op) final { group_map_[op] = arena_->make(); } - void VisitExpr_(const DataTypeImmNode* op) final { group_map_[op] = arena_->make(); } + void VisitExpr_(const GenericConstObj* op) final { group_map_[op] = arena_->make(); } + void VisitExpr_(const DataTypeImmObj* op) final { group_map_[op] = arena_->make(); } - void VisitBinding_(const VarBindingNode* binding, const CallNode* call) final { + void VisitBinding_(const VarBindingObj* binding, const CallObj* call) final { VisitVarDef(binding->var); if (auto matches_opt = ExtractMatchedExpr(pat_, ffi::GetRef(call), bindings_)) { const auto& context = CreatePatternCheckContext(call, matches_opt.value()); @@ -1234,8 +1234,8 @@ class PatternBasedPartitioner : ExprVisitor { } for (const auto& [pat, match] : matches_opt.value()) { - if ((pat->IsInstance() && !match.same_as(ffi::GetRef(call))) || - pat->IsInstance()) { + if ((pat->IsInstance() && !match.same_as(ffi::GetRef(call))) || + pat->IsInstance()) { auto g = GetGroup(match); if (g && g->FindRoot()->num_nodes > 1) { // This expression has already been matched to a previous pattern. @@ -1279,8 +1279,8 @@ class PatternBasedPartitioner : ExprVisitor { // the previous group. For example, when there are two back-to-back conv2d ops, the output // of the first conv2d is matched to the input of the second conv2d via a wildcard pattern. // But we must avoid merging the first conv2d into the group of the second conv2d. - if ((pat->IsInstance() && !match.same_as(ffi::GetRef(call))) || - pat->IsInstance()) { + if ((pat->IsInstance() && !match.same_as(ffi::GetRef(call))) || + pat->IsInstance()) { // Put the bound variable on the LHS into the same parent group. AddToGroup(value_to_bound_var_[match], parent_group); } @@ -1310,7 +1310,7 @@ class PatternBasedPartitioner : ExprVisitor { return nullptr; } - PatternCheckContext CreatePatternCheckContext(const CallNode* call, + PatternCheckContext CreatePatternCheckContext(const CallObj* call, const ffi::Map& matched_result) { ffi::Map annotated_expr; for (const auto& it : annotation_pat_) { @@ -1321,7 +1321,7 @@ class PatternBasedPartitioner : ExprVisitor { ffi::Map matched_bindings; for (const auto& [pat, match] : matched_result) { - if (pat->IsInstance() || pat->IsInstance()) { + if (pat->IsInstance() || pat->IsInstance()) { matched_bindings.Set(value_to_bound_var_[match], match); } } @@ -1335,7 +1335,7 @@ class PatternBasedPartitioner : ExprVisitor { const ffi::Map& matched_result) { std::set matched_vars; for (const auto& [pat, match] : matched_result) { - if ((pat->IsInstance() || pat->IsInstance())) + if ((pat->IsInstance() || pat->IsInstance())) matched_vars.insert(value_to_bound_var_[match]); } @@ -1377,7 +1377,7 @@ class CompositeFunctionAnnotator : public ExprMutator { continue; } const auto& base_func = (*it).second; - if (const auto* func = base_func.as()) { + if (const auto* func = base_func.as()) { if (func->GetAttr(tvm::relax::attr::kComposite).has_value() || func->GetAttr(tvm::relax::attr::kCodegen).has_value()) { continue; @@ -1394,8 +1394,8 @@ class CompositeFunctionAnnotator : public ExprMutator { return builder_->GetContextIRModule(); } - Expr VisitExpr_(const CallNode* call_node) final { - if (auto const* gvar = call_node->op.as()) { + Expr VisitExpr_(const CallObj* call_node) final { + if (auto const* gvar = call_node->op.as()) { if (auto it = gvar_map_.find(gvar); it != gvar_map_.end()) { return Call(Type::Missing(), it->second, call_node->args); } @@ -1416,7 +1416,7 @@ class CompositeFunctionAnnotator : public ExprMutator { return ExprMutator::VisitExpr_(call_node); } - Expr VisitExpr_(const FunctionNode* func_node) final { + Expr VisitExpr_(const FunctionObj* func_node) final { Function f_inner = ExprMutator::VisitExpr_(func_node).as_or_throw(); if (!func_node->GetAttr(tvm::relax::attr::kComposite)) { @@ -1436,7 +1436,7 @@ class CompositeFunctionAnnotator : public ExprMutator { params.push_back(new_v); } - // We cannot delegate to `ExprMutator::VisitExpr_(const FunctionNode*)` at this point, as it + // We cannot delegate to `ExprMutator::VisitExpr_(const FunctionObj*)` at this point, as it // would recursively visit the Call node. However, we are still required to generate // well-formed Relax IR. As a result, we need to build the SeqExpr ourselves. Var local_func_var("local_func", GetType(f_inner)); @@ -1453,7 +1453,7 @@ class CompositeFunctionAnnotator : public ExprMutator { private: /*! \brief A map from old global vars to their replacements. */ - std::unordered_map gvar_map_; + std::unordered_map gvar_map_; }; IRModule FuseOpsByPattern(const tvm::ffi::Array& patterns, IRModule mod, @@ -1467,17 +1467,17 @@ IRModule FuseOpsByPattern(const tvm::ffi::Array& patte for (const auto& name : entry_function_names) { auto gv = mod->GetGlobalVar(name); auto func = mod->Lookup(gv); - TVM_FFI_ICHECK(func->IsInstance()) + TVM_FFI_ICHECK(func->IsInstance()) << "Entry function must be a relax function"; entry_functions.push_back(func.as_or_throw()); } } else { for (const auto& gv : mod->GetGlobalVars()) { const auto& base_func = mod->Lookup(gv); - if (base_func->IsInstance()) { + if (base_func->IsInstance()) { continue; } - const FunctionNode* function = base_func.as(); + const FunctionObj* function = base_func.as(); if (function->GetAttr(tvm::relax::attr::kPrimitive).value_or(false) || function->GetAttr(tvm::relax::attr::kComposite).has_value() || function->GetAttr(tvm::relax::attr::kCodegen).has_value()) { diff --git a/src/relax/transform/fuse_tir.cc b/src/relax/transform/fuse_tir.cc index 8d75beece6ef..ffff12c5b8ca 100644 --- a/src/relax/transform/fuse_tir.cc +++ b/src/relax/transform/fuse_tir.cc @@ -90,26 +90,26 @@ class SymbolicMatcher : ExprFunctor } \ } - TVM_DECLARE_SYMBOLIC_MATCHER_BINOP(prim::AddNode); - TVM_DECLARE_SYMBOLIC_MATCHER_BINOP(prim::SubNode); - TVM_DECLARE_SYMBOLIC_MATCHER_BINOP(prim::MulNode); - TVM_DECLARE_SYMBOLIC_MATCHER_BINOP(prim::DivNode); - TVM_DECLARE_SYMBOLIC_MATCHER_BINOP(prim::ModNode); - TVM_DECLARE_SYMBOLIC_MATCHER_BINOP(prim::EQNode); - TVM_DECLARE_SYMBOLIC_MATCHER_BINOP(prim::NENode); - TVM_DECLARE_SYMBOLIC_MATCHER_BINOP(prim::LTNode); - TVM_DECLARE_SYMBOLIC_MATCHER_BINOP(prim::LENode); - TVM_DECLARE_SYMBOLIC_MATCHER_BINOP(prim::GTNode); - TVM_DECLARE_SYMBOLIC_MATCHER_BINOP(prim::GENode); - TVM_DECLARE_SYMBOLIC_MATCHER_BINOP(prim::AndNode); - TVM_DECLARE_SYMBOLIC_MATCHER_BINOP(prim::OrNode); - TVM_DECLARE_SYMBOLIC_MATCHER_BINOP(prim::MinNode); - TVM_DECLARE_SYMBOLIC_MATCHER_BINOP(prim::MaxNode); - TVM_DECLARE_SYMBOLIC_MATCHER_BINOP(prim::FloorDivNode); - TVM_DECLARE_SYMBOLIC_MATCHER_BINOP(prim::FloorModNode); - - void Dispatch_(const IntImmNode* op, const PrimExpr& other) { - const auto* rhs = other.as(); + TVM_DECLARE_SYMBOLIC_MATCHER_BINOP(prim::AddObj); + TVM_DECLARE_SYMBOLIC_MATCHER_BINOP(prim::SubObj); + TVM_DECLARE_SYMBOLIC_MATCHER_BINOP(prim::MulObj); + TVM_DECLARE_SYMBOLIC_MATCHER_BINOP(prim::DivObj); + TVM_DECLARE_SYMBOLIC_MATCHER_BINOP(prim::ModObj); + TVM_DECLARE_SYMBOLIC_MATCHER_BINOP(prim::EQObj); + TVM_DECLARE_SYMBOLIC_MATCHER_BINOP(prim::NEObj); + TVM_DECLARE_SYMBOLIC_MATCHER_BINOP(prim::LTObj); + TVM_DECLARE_SYMBOLIC_MATCHER_BINOP(prim::LEObj); + TVM_DECLARE_SYMBOLIC_MATCHER_BINOP(prim::GTObj); + TVM_DECLARE_SYMBOLIC_MATCHER_BINOP(prim::GEObj); + TVM_DECLARE_SYMBOLIC_MATCHER_BINOP(prim::AndObj); + TVM_DECLARE_SYMBOLIC_MATCHER_BINOP(prim::OrObj); + TVM_DECLARE_SYMBOLIC_MATCHER_BINOP(prim::MinObj); + TVM_DECLARE_SYMBOLIC_MATCHER_BINOP(prim::MaxObj); + TVM_DECLARE_SYMBOLIC_MATCHER_BINOP(prim::FloorDivObj); + TVM_DECLARE_SYMBOLIC_MATCHER_BINOP(prim::FloorModObj); + + void Dispatch_(const IntImmObj* op, const PrimExpr& other) { + const auto* rhs = other.as(); if (!rhs || (op->value != rhs->value)) { TVM_FFI_THROW(InternalError) << "Parameter expression " << ffi::GetRef(op) @@ -118,8 +118,8 @@ class SymbolicMatcher : ExprFunctor } } - void Dispatch_(const FloatImmNode* op, const PrimExpr& other) { - const auto* rhs = other.as(); + void Dispatch_(const FloatImmObj* op, const PrimExpr& other) { + const auto* rhs = other.as(); if (!rhs || (op->value != rhs->value)) { TVM_FFI_THROW(InternalError) << "Parameter expression " << ffi::GetRef(op) << " expected an float argument with value " << op->value << ", " @@ -127,8 +127,8 @@ class SymbolicMatcher : ExprFunctor } } - void Dispatch_(const prim::CastNode* op, const PrimExpr& other) { - const auto* rhs = other.as(); + void Dispatch_(const prim::CastObj* op, const PrimExpr& other) { + const auto* rhs = other.as(); if (!rhs) { TVM_FFI_THROW(InternalError) << "Parameter expression " << ffi::GetRef(op) << " expected an cast to " @@ -138,7 +138,7 @@ class SymbolicMatcher : ExprFunctor Dispatch(op->value, rhs->value); } - void Dispatch_(const VarNode* op, const PrimExpr& rhs) { + void Dispatch_(const VarObj* op, const PrimExpr& rhs) { auto lhs = ffi::GetRef(op); PrimType lhs_ty = op->ty.as_or_throw(); @@ -155,8 +155,8 @@ class SymbolicMatcher : ExprFunctor } } - void Dispatch_(const prim::SelectNode* op, const PrimExpr& other) { - const auto* rhs = other.as(); + void Dispatch_(const prim::SelectObj* op, const PrimExpr& other) { + const auto* rhs = other.as(); if (rhs) { Dispatch(op->true_value, rhs->true_value); Dispatch(op->false_value, rhs->false_value); @@ -193,7 +193,7 @@ class FuseTIRTensorSubstitutor : public s_tir::StmtExprMutator { } private: - UnchangedOr Mutate_(const s_tir::SBlockNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const s_tir::SBlockObj* op, InplaceMode inplace_mode) final { s_tir::SBlock block = s_tir::StmtExprMutator::Mutate_(op, inplace_mode) .ValueOrUnchanged(ffi::GetRef(op)) .as_or_throw(); @@ -213,7 +213,7 @@ class FuseTIRTensorSubstitutor : public s_tir::StmtExprMutator { // However, `A[vi, vj], A[vi, vj + 1]` is not allow for now. // Note: the order of return region should remain the same as the first occurrence of the region ffi::Array ret; - std::unordered_map> tensor_region_set; + std::unordered_map> tensor_region_set; for (const TensorRegion& region : regions) { auto it = tensor_region_set.find(region->source.as_or_throw().get()); @@ -237,12 +237,12 @@ class SBlockNameDeduplicator : public s_tir::StmtExprMutator { public: using s_tir::StmtExprMutator::Mutate; UnchangedOr Mutate(ffi::AnyView value, InplaceMode inplace_mode) final { - if (value.as()) return ffi::Unchanged(); + if (value.as()) return ffi::Unchanged(); return s_tir::StmtExprMutator::Mutate(value, inplace_mode); } private: - UnchangedOr Mutate_(const s_tir::SBlockNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const s_tir::SBlockObj* op, InplaceMode inplace_mode) final { s_tir::SBlock block = s_tir::StmtExprMutator::Mutate_(op, inplace_mode) .ValueOrUnchanged(ffi::GetRef(op)) .as_or_throw(); @@ -342,12 +342,12 @@ class RelaxToTIRVarMapCollector : public ExprVisitor { } private: - void VisitBinding_(const VarBindingNode* binding) final { + void VisitBinding_(const VarBindingObj* binding) final { current_var_ = binding->var; ExprVisitor::VisitBinding_(binding); } - void VisitExpr_(const CallNode* call) { + void VisitExpr_(const CallObj* call) { static const Op call_tir_op_ = Op::Get("relax.call_tir"); static const Op call_tir_inplace_op_ = Op::Get("relax.call_tir_inplace"); @@ -357,7 +357,7 @@ class RelaxToTIRVarMapCollector : public ExprVisitor { CollectVarMapping(call, current_var_, call->op.same_as(call_tir_inplace_op_)); } - void CollectVarMapping(const CallNode* call, const Expr& lhs_var, bool in_place) { + void CollectVarMapping(const CallObj* call, const Expr& lhs_var, bool in_place) { GlobalVar gv = call->args[0].as_or_throw(); tirx::Function function_ = mod_->Lookup(gv).as_or_throw(); const auto& tir_args = function_->params; @@ -365,10 +365,10 @@ class RelaxToTIRVarMapCollector : public ExprVisitor { const auto& relax_args = call->args[1].as_or_throw()->fields; ffi::Array relax_results; - if (lhs_var->IsInstance()) { + if (lhs_var->IsInstance()) { relax_results = lhs_var.as_or_throw()->fields; } else { - TVM_FFI_ICHECK(lhs_var->IsInstance()) + TVM_FFI_ICHECK(lhs_var->IsInstance()) << "The lhs_var is expected to be either tuple or var"; relax_results = {lhs_var.as_or_throw()}; } @@ -434,7 +434,7 @@ class FusedTIRConstructor : public ExprVisitor { const GlobalVar& gv) { FusedTIRConstructor visitor(mod, gv->name_hint); BaseFunc f = mod->Lookup(gv); - TVM_FFI_ICHECK(f->IsInstance()) + TVM_FFI_ICHECK(f->IsInstance()) << "Expected relax functions, but got: " << f->GetTypeKey(); TVM_FFI_ICHECK(f->HasNonzeroAttr(tvm::relax::attr::kPrimitive.name)) << "Expected a function with attr `kPrimitive`"; @@ -450,7 +450,7 @@ class FusedTIRConstructor : public ExprVisitor { explicit FusedTIRConstructor(const IRModule& mod, const ffi::String& func_name) : mod_(mod), func_name_(func_name) {} - void VisitExpr_(const FunctionNode* func) final { + void VisitExpr_(const FunctionObj* func) final { auto relax_to_tir_var_map = RelaxToTIRVarMapCollector::Collect(mod_, ffi::GetRef(func)); std::vector> function_params; @@ -541,7 +541,7 @@ class FusedTIRConstructor : public ExprVisitor { fused_tir_ = ConstructFunc(); } - void VisitBinding_(const VarBindingNode* binding) final { + void VisitBinding_(const VarBindingObj* binding) final { // Update expr2tensors by visiting values. this->VisitExpr(binding->value); auto it = func_info_.expr2tensors.find(binding->value); @@ -553,11 +553,11 @@ class FusedTIRConstructor : public ExprVisitor { } } - void VisitBinding_(const MatchCastNode* match_cast) final { + void VisitBinding_(const MatchCastObj* match_cast) final { TVM_FFI_THROW(InternalError) << "MatchCast is unsupported in primitive functions"; } - void VisitExpr_(const CallNode* call) final { + void VisitExpr_(const CallObj* call) final { ExprVisitor::VisitExpr_(call); static const Op call_tir_op_ = Op::Get("relax.call_tir"); static const Op call_tir_inplace_op_ = Op::Get("relax.call_tir_inplace"); @@ -576,7 +576,7 @@ class FusedTIRConstructor : public ExprVisitor { // Step 3. Check functions are all schedulable funcs. i.e. the body of func is root block // TODO(Siyuan): support un-schedulable functions. TVM_FFI_ICHECK(function->body.has_value() && function->body.value()->size() == 1 && - function->body.value()->seq[0].as()) + function->body.value()->seq[0].as()) << "Only schedulable functions (whose body is the root block) can be fused"; s_tir::SBlockRealize root_realize = function->body.value()->seq[0].as_or_throw(); @@ -598,7 +598,7 @@ class FusedTIRConstructor : public ExprVisitor { func_info_.global_name += "_" + gv->name_hint; } - void VisitExpr_(const TupleGetItemNode* tuple_get_item) final { + void VisitExpr_(const TupleGetItemObj* tuple_get_item) final { ExprVisitor::VisitExpr_(tuple_get_item); auto it = func_info_.expr2tensors.find(tuple_get_item->tuple); if (it != func_info_.expr2tensors.end()) { @@ -615,7 +615,7 @@ class FusedTIRConstructor : public ExprVisitor { } } - void VisitExpr_(const TupleNode* tuple) final { + void VisitExpr_(const TupleObj* tuple) final { ExprVisitor::VisitExpr_(tuple); ffi::Array tensors; for (const Expr& expr : tuple->fields) { @@ -629,7 +629,7 @@ class FusedTIRConstructor : public ExprVisitor { } } - void VisitExpr_(const GenericConstNode* op) final { + void VisitExpr_(const GenericConstObj* op) final { if (!op->value.as()) return; TVM_FFI_THROW(InternalError) << "Tensor constants are not supported in primitive functions."; } @@ -638,29 +638,29 @@ class FusedTIRConstructor : public ExprVisitor { * \brief Get the number of outputs for a call_tir node. * \return The number of outputs. */ - static ffi::Array> GetCallTIROutputShapes(const CallNode* call) { + static ffi::Array> GetCallTIROutputShapes(const CallObj* call) { static const Op call_tir_op_ = Op::Get("relax.call_tir"); static const Op call_tir_inplace_op_ = Op::Get("relax.call_tir_inplace"); TVM_FFI_ICHECK(call->op.same_as(call_tir_op_) || call->op.same_as(call_tir_inplace_op_)); TVM_FFI_ICHECK_EQ(call->ty_args.size(), 1); auto get_tensor_shape = - [](const TensorTypeNode* ty) { - const auto* shape_expr = ty->shape.as(); + [](const TensorTypeObj* ty) { + const auto* shape_expr = ty->shape.as(); TVM_FFI_ICHECK(shape_expr) << "FuseTIR expects all parameters are Tensors with symbolic shape."; return shape_expr->values; }; - if (const auto* tuple_ty = call->ty_args[0].as()) { + if (const auto* tuple_ty = call->ty_args[0].as()) { ffi::Array> shapes; for (const Type& field : tuple_ty->fields) { - const auto* tensor_ty = field.as(); + const auto* tensor_ty = field.as(); TVM_FFI_ICHECK(tensor_ty) << "CallTIR ty_args are expected to be TensorType or Tuple of " "TensorType, but got " << call->ty_args[0]; shapes.push_back(get_tensor_shape(tensor_ty)); } return shapes; - } else if (const auto* tensor_ty = call->ty_args[0].as()) { + } else if (const auto* tensor_ty = call->ty_args[0].as()) { return {get_tensor_shape(tensor_ty)}; } else { TVM_FFI_ICHECK(tensor_ty) << "CallTIR ty_args are expected to be TensorType or Tuple of " @@ -674,7 +674,7 @@ class FusedTIRConstructor : public ExprVisitor { void MapArgsToTensor(const ffi::Array args, const ffi::Array& tensors) { size_t tensor_idx = 0; for (const Expr& arg : args) { - if (const auto* v = arg.as()) { + if (const auto* v = arg.as()) { auto it = func_info_.expr2tensors.find(ffi::GetRef(v)); // Substitute the tensor with the already allocated one if it is an intermediate var if (it != func_info_.expr2tensors.end()) { @@ -748,7 +748,7 @@ class FusedTIRConstructor : public ExprVisitor { * \param func The old TIR tirx::Function * \param output_shapes The shape of output params. */ - void AllocateIntermediateTensor(const CallNode* call, const tirx::Function& func, + void AllocateIntermediateTensor(const CallObj* call, const tirx::Function& func, const ffi::Array>& output_shapes) { bool is_inplace = call->op.same_as(Op::Get("relax.call_tir_inplace")); @@ -797,7 +797,7 @@ class FusedTIRConstructor : public ExprVisitor { return unique_name; }; // Update tensor with new symbolic shape according to the ty - auto new_type_node = ffi::make_object(*tensor.ty().get()); + auto new_type_node = ffi::make_object(*tensor.ty().get()); new_type_node->shape = output_shapes[i]; tirx::TensorType new_type(std::move(new_type_node)); tirx::TensorVar new_tensor(unify_name_hints(), std::move(new_type), tensor->loc); @@ -823,15 +823,15 @@ class FusedTIRConstructor : public ExprVisitor { const ffi::Optional& tir_tensor_param) { auto ty = GetType(relax_param); - TVM_FFI_CHECK(!ty.as(), InternalError) + TVM_FFI_CHECK(!ty.as(), InternalError) << "All tuple parameters should be expanded before this point in FuseTIR. " << "However, parameter " << relax_param << " has type " << ty; auto name_hint = relax_param->name; - if (const auto* tensor_type = ty.as()) { + if (const auto* tensor_type = ty.as()) { // Case 1. The relax param is a Tensor, we directly create a tirx var and tensor - const auto* shape_expr = tensor_type->shape.as(); + const auto* shape_expr = tensor_type->shape.as(); TVM_FFI_ICHECK(shape_expr) << "FuseTIR expects all Tensor parameters have a known shape."; PrimType dtype = tensor_type->dtype.value(); tirx::TensorVar tensor = @@ -841,11 +841,11 @@ class FusedTIRConstructor : public ExprVisitor { : tirx::TensorVar(name_hint, tirx::TensorType(shape_expr->values, dtype)); out->push_back(std::move(tensor)); - } else if (ty.as()) { + } else if (ty.as()) { // Case 2. The relax param is a scalar, so its canonical Var is a TIR parameter. out->push_back(relax_param.as_or_throw()); - } else if (const auto* shape_expr = ty.as()) { + } else if (const auto* shape_expr = ty.as()) { // Case 3. The relax param is a tuple of scalars, each represented as a tirx var for (const auto& var : shape_expr->values.value()) { auto prim_var = var.as(); @@ -896,9 +896,9 @@ class FusedTIRConstructor : public ExprVisitor { /*! \brief Get DynTensor numbers from recursive Tuples. */ static size_t GetTotalTensorSize(const Type& ty) { - if (ty.as()) { + if (ty.as()) { return 1; - } else if (const auto* tuple_ty = ty.as()) { + } else if (const auto* tuple_ty = ty.as()) { size_t num = 0; for (const Type& ty : tuple_ty->fields) { num += GetTotalTensorSize(ty); @@ -935,7 +935,7 @@ class FusedTIRConstructor : public ExprVisitor { /*! \brief Tensor annotations keyed by their placeholder parameters. */ ffi::Map tensor_map; /*! \brief The output tensors among the function parameters. */ - std::unordered_set output_tensors; + std::unordered_set output_tensors; /*! \brief The name of the fused function */ std::string global_name = "fused"; @@ -972,11 +972,11 @@ class FusedTIRConstructor : public ExprVisitor { std::unordered_set inplace_indices_; }; -std::vector GetTupleAccessedIndices(const FunctionNode* func, const Var& tuple_var) { +std::vector GetTupleAccessedIndices(const FunctionObj* func, const Var& tuple_var) { // Need to be ordered std::vector indices; PostOrderVisit(func->body, [&indices, tuple_var](Expr e) { - if (auto tup_get = e.as(); tup_get && tup_get->tuple.same_as(tuple_var)) { + if (auto tup_get = e.as(); tup_get && tup_get->tuple.same_as(tuple_var)) { if (std::find(indices.begin(), indices.end(), tup_get->index) == indices.end()) { indices.push_back(tup_get->index); } @@ -1032,7 +1032,7 @@ class TIRFuseMutator : public ExprMutator { // into the new IRModule for (const auto& [gv, func] : mod->functions) { - if (func->IsInstance()) { + if (func->IsInstance()) { TVM_FFI_ICHECK(!func->HasNonzeroAttr(tvm::relax::attr::kPrimitive.name)) << "Module should not contain any primitive relax functions at this point"; relax::Function update_func = mutator.VisitExpr(func).as_or_throw(); @@ -1061,19 +1061,19 @@ class TIRFuseMutator : public ExprMutator { // Get shape from call tirx static Expr GetCallTIRShape(Type ty) { - if (auto* tuple = ty.as()) { + if (auto* tuple = ty.as()) { ffi::Array fields = tuple->fields.Map([&](Type x) { return GetCallTIRShape(x); }); return Tuple(fields); } else { - auto* tensor = ty.as(); + auto* tensor = ty.as(); TVM_FFI_ICHECK(tensor) << "FuseTIR can only take tensor or tuple type"; - auto* shape_expr = tensor->shape.as(); + auto* shape_expr = tensor->shape.as(); TVM_FFI_ICHECK(shape_expr) << "FuseTIR requires all intermediate values have shape"; return ffi::GetRef(shape_expr); } } - Expr VisitExpr_(const CallNode* op) final { + Expr VisitExpr_(const CallObj* op) final { static const Op call_tir_op_ = Op::Get("relax.call_tir"); static const Op call_tir_inplace_op_ = Op::Get("relax.call_tir_inplace"); @@ -1108,14 +1108,14 @@ class TIRFuseMutator : public ExprMutator { auto ty = GetType(arg); TVM_FFI_CHECK( - !relax_func->params[i]->ty->IsInstance() && !ty.as(), + !relax_func->params[i]->ty->IsInstance() && !ty.as(), InternalError) << "All tuple parameters should be expanded before this point in FuseTIR. " << "However, argument " << arg << " with type " << arg->ty << " is passed as argument " << i << " to Primitive Relax function " << old_gvar << ", which expects parameter " << relax_func->params[i] << " to have type " << relax_func->params[i]->ty; - if (const auto* shape = ty.as()) { + if (const auto* shape = ty.as()) { TVM_FFI_ICHECK(shape->values.has_value()) << "FuseTIR requires all shape input has ty value."; for (const PrimExpr& prim_value : shape->values.value()) { @@ -1123,7 +1123,7 @@ class TIRFuseMutator : public ExprMutator { << "All shape inputs are expected to be single tirx var."; arg_list.push_back(prim_value); } - } else if (ty.as()) { + } else if (ty.as()) { if (auto literal = arg.as()) { arg_list.push_back(literal.value()); } else { diff --git a/src/relax/transform/gradient.cc b/src/relax/transform/gradient.cc index 23f310d61ece..cfbab7268f83 100644 --- a/src/relax/transform/gradient.cc +++ b/src/relax/transform/gradient.cc @@ -74,7 +74,7 @@ class CallTIRWithGradEliminator : private ExprMutator { private: using ExprMutator::VisitExpr_; - Expr VisitExpr_(const CallNode* call_node) final { + Expr VisitExpr_(const CallObj* call_node) final { if (!call_node->op.same_as(Op::Get("relax.call_tir_with_grad"))) { return ExprMutator::VisitExpr_(call_node); } @@ -116,7 +116,7 @@ class CheckpointCollector : private ExprMutator { std::unordered_map var_mapping; private: - Expr VisitExpr_(const FunctionNode* func) final { + Expr VisitExpr_(const FunctionObj* func) final { for (auto var : func->params) { checkpoints.insert(var); } @@ -131,14 +131,14 @@ class CheckpointCollector : private ExprMutator { // If every variable that the variable of binding relies on is either // 1) the output of end_checkpoint; 2) checkpointed // then the variable of binding will be checkpointed - auto var_binding = binding.as(); + auto var_binding = binding.as(); TVM_FFI_ICHECK(var_binding); - auto value_call = var_binding->value.as(); + auto value_call = var_binding->value.as(); if (!value_call || (!value_call->op.same_as(s_cp) && !value_call->op.same_as(e_cp))) { bool all_inner_var_checkpointed = true; PostOrderVisit(var_binding->value, [this, &all_inner_var_checkpointed](const Expr& expr) { - if (auto var = expr.as()) { + if (auto var = expr.as()) { Var var_ref = ffi::GetRef(var); all_inner_var_checkpointed &= (checkpoints.count(var_ref) != 0 || e_vars_.count(var_ref) != 0); @@ -154,19 +154,19 @@ class CheckpointCollector : private ExprMutator { } // mark vars to be checkpointed, and eliminate bindings with checkpoint calls - void VisitBinding_(const VarBindingNode* binding, const CallNode* value) final { + void VisitBinding_(const VarBindingObj* binding, const CallObj* value) final { static const auto s_cp = Op::Get("relax.grad.start_checkpoint"); static const auto e_cp = Op::Get("relax.grad.end_checkpoint"); if (value->op.same_as(s_cp) || value->op.same_as(e_cp)) { // Eliminate the binding - auto var = value->args[0].as(); + auto var = value->args[0].as(); TVM_FFI_ICHECK(var) << "The first argument of relax.grad.start_checkpoint and " "relax.grad.end_checkpoint should be a Var"; // var might already be remapped. Find the original var auto orig_var = ExprMutator::VisitExpr(ffi::GetRef(var)).as_or_throw(); // Add remapping from binding->var to new_var - if (!binding->var.as() && var->IsInstance()) { + if (!binding->var.as() && var->IsInstance()) { // For output binding, emit a dummy binding this->var_remap_.insert_or_assign(binding->var, builder_->EmitOutput(orig_var, orig_var->name)); @@ -217,8 +217,8 @@ class CheckpointGenerator : private ExprMutator { } for (auto binding : forward_block->bindings) { - auto* var_binding = binding.as(); - TVM_FFI_ICHECK(var_binding) << "Now only support VarBindingNode"; + auto* var_binding = binding.as(); + TVM_FFI_ICHECK(var_binding) << "Now only support VarBindingObj"; auto var = var_binding->var; binding_map_.Set(var, var_binding->value); if (checkpoints.count(var)) { @@ -243,10 +243,10 @@ class CheckpointGenerator : private ExprMutator { using ExprMutator::VisitExpr_; // Visit the use-site of a defined Var - Expr VisitExpr_(const VarNode* op) final { return VisitVar(ffi::GetRef(op)); } + Expr VisitExpr_(const VarObj* op) final { return VisitVar(ffi::GetRef(op)); } // Visit the use-site of a defined DataflowVar - Expr VisitExpr_(const DataflowVarNode* op) final { return VisitVar(ffi::GetRef(op)); } + Expr VisitExpr_(const DataflowVarObj* op) final { return VisitVar(ffi::GetRef(op)); } Expr VisitVar(const Var& var) { auto it = checkpoint_map_.find(var); @@ -260,7 +260,7 @@ class CheckpointGenerator : private ExprMutator { // The only purpose of this function is create a new expr for Call node // to pass the structual equal check - Expr VisitExpr_(const CallNode* call_node) final { + Expr VisitExpr_(const CallObj* call_node) final { Expr new_op = this->VisitExpr(call_node->op); tvm::ffi::Array call_args; @@ -269,7 +269,7 @@ class CheckpointGenerator : private ExprMutator { call_args.push_back(new_arg); } Type ret_ty = Type::Missing(); - if (call_node->ty.as()) { + if (call_node->ty.as()) { if (auto op = call_node->op.as()) { static auto infer_type_map = Op::GetAttrMap(tvm::op_attr::kInferType); static auto infer_type_with_builder_map = @@ -318,7 +318,7 @@ class BackwardBindingGenerator : private ExprVisitor { BackwardBindingGenerator generator(builder, cp_collector, checkpoint_generator); // Initialize the adjoint of target_var as ones op. We have already checked the target. - auto* target_ty = GetTypeAs(target_var); + auto* target_ty = GetTypeAs(target_var); generator.UpdateAdjoint(target_var, ones(target_ty->shape.value(), target_ty->dtype.value()->dtype)); @@ -340,8 +340,8 @@ class BackwardBindingGenerator : private ExprVisitor { void VisitBinding(const Binding& binding) final { // TODO(chaofan, yixin): support other types of bindings - TVM_FFI_ICHECK(binding->IsInstance()) << "Now only support VarBindingNode"; - auto* var_binding = binding.as(); + TVM_FFI_ICHECK(binding->IsInstance()) << "Now only support VarBindingObj"; + auto* var_binding = binding.as(); if (adjoint_var_map_.count(var_binding->var) == 0) { // Optimization: this var is not used in the following bindings @@ -350,9 +350,9 @@ class BackwardBindingGenerator : private ExprVisitor { Expr value = var_binding->value; // TODO(chaofan, yixin): support other types of binding values - TVM_FFI_ICHECK(value->IsInstance() || value->IsInstance() || - value->IsInstance() || value->IsInstance() || - value->IsInstance()) + TVM_FFI_ICHECK(value->IsInstance() || value->IsInstance() || + value->IsInstance() || value->IsInstance() || + value->IsInstance()) << "Now does not support the type of binding value: " << value; ExprVisitor::VisitBinding_(var_binding); @@ -360,9 +360,9 @@ class BackwardBindingGenerator : private ExprVisitor { // The following functions will handle the adjoint expr of the inputs of binding // For call node, we would call the registered gradient functions - void VisitBinding_(const VarBindingNode* binding, const CallNode* call) final { + void VisitBinding_(const VarBindingObj* binding, const CallObj* call) final { // Skip if it is not an Op - if (!call->op->IsInstance()) { + if (!call->op->IsInstance()) { return; } @@ -389,7 +389,7 @@ class BackwardBindingGenerator : private ExprVisitor { grad_func(checkpoint_var, checkpoint_call.as_or_throw(), adjoint_var, builder_) .cast(); Tuple args = call->args[1].as_or_throw(); - auto* tuple_ty = GetTypeAs(partials); + auto* tuple_ty = GetTypeAs(partials); if (!tuple_ty) { // result_var is a tensor TVM_FFI_ICHECK(args->fields.size() == 1); @@ -422,7 +422,7 @@ class BackwardBindingGenerator : private ExprVisitor { // a = ((b, c), d) --> // b_adjoint += a_adjoint_var[0][0], c_adjoint += a_adjoint_var[0][1], // d_adjoint += a_adjoint_var[1] - void VisitBinding_(const VarBindingNode* binding, const TupleNode* tuple) final { + void VisitBinding_(const VarBindingObj* binding, const TupleObj* tuple) final { UpdateAdjoint(ffi::GetRef(tuple), adjoint_var_map_[binding->var]); } @@ -431,10 +431,10 @@ class BackwardBindingGenerator : private ExprVisitor { // b = a[0] --> // a_adjoint[0] += b_adjoint_var // If a_adjoint does not exist, we would create a zeros tuple as a_adjoint first, and then add - void VisitBinding_(const VarBindingNode* binding, const TupleGetItemNode* tuple_get_item) final { - TVM_FFI_ICHECK(tuple_get_item->tuple->IsInstance()) + void VisitBinding_(const VarBindingObj* binding, const TupleGetItemObj* tuple_get_item) final { + TVM_FFI_ICHECK(tuple_get_item->tuple->IsInstance()) << "The tuple field of a TupleGetItem is not bound to a Var"; - auto* tuple_ty = GetTypeAs(tuple_get_item->tuple); + auto* tuple_ty = GetTypeAs(tuple_get_item->tuple); TVM_FFI_ICHECK(tuple_ty) << "The tuple field of a TupleGetItem must has a TupleType"; const Var& tuple_var = tuple_get_item->tuple.as_or_throw(); @@ -451,16 +451,16 @@ class BackwardBindingGenerator : private ExprVisitor { } // For assign nodes, we add the adjoint of output to the adjoint of input - void VisitBinding_(const VarBindingNode* binding, const DataflowVarNode* var) final { + void VisitBinding_(const VarBindingObj* binding, const DataflowVarObj* var) final { UpdateAdjoint(ffi::GetRef(var), adjoint_var_map_[binding->var]); } - void VisitBinding_(const VarBindingNode* binding, const VarNode* var) final { + void VisitBinding_(const VarBindingObj* binding, const VarObj* var) final { UpdateAdjoint(ffi::GetRef(var), adjoint_var_map_[binding->var]); } // For constant nodes, we do not have to handle it because it does not contribute to the adjoint - void VisitBinding_(const VarBindingNode* binding, const GenericConstNode* var) final { return; } + void VisitBinding_(const VarBindingObj* binding, const GenericConstObj* var) final { return; } // Add partial to the adjoint of expr // expr may be a argument of a func call / tuple definition. Its type can be @@ -472,7 +472,7 @@ class BackwardBindingGenerator : private ExprVisitor { void UpdateAdjoint(const Expr& expr, const Expr& partial) { AdjointMsg partial_msg = ExprToAdjointMsg(builder_->Normalize(partial)); DecomposeNestedMsg(expr, partial_msg, [&](Expr leaf, AdjointMsg msg) { - if (leaf->IsInstance()) { + if (leaf->IsInstance()) { const Var& v = leaf.as_or_throw(); Expr updated_adjoint_expr = builder_->Normalize(AdjointMsgToExpr(msg)); auto it = adjoint_var_map_.find(v); @@ -480,9 +480,9 @@ class BackwardBindingGenerator : private ExprVisitor { updated_adjoint_expr = TupleAwareAdd((*it).second, updated_adjoint_expr); } EmitAdjoint(v, updated_adjoint_expr, false); - } else if (leaf->IsInstance()) { + } else if (leaf->IsInstance()) { // nothing to do - } else if (leaf->IsInstance()) { + } else if (leaf->IsInstance()) { // must be no grad TVM_FFI_ICHECK(IsCallNoGrad(partial)); } else { @@ -531,7 +531,7 @@ class BackwardBindingGenerator : private ExprVisitor { } static bool IsCallNoGrad(const Expr& expr) { - return expr->IsInstance() && + return expr->IsInstance() && expr.as_or_throw()->op.same_as(Op::Get("relax.grad.no_grad")); } @@ -546,7 +546,7 @@ class BackwardBindingGenerator : private ExprVisitor { static AdjointMsg ExprToAdjointMsg(Expr expr) { return MapToNestedMsgByType(expr, [](Expr leaf) { - TVM_FFI_ICHECK(GetTypeAs(leaf)) + TVM_FFI_ICHECK(GetTypeAs(leaf)) << "The leaf of adjoint: " << leaf << " should have Type and be a Tensor."; return AdjointMsg(leaf); }); @@ -556,7 +556,7 @@ class BackwardBindingGenerator : private ExprVisitor { // When ty is TupleType, we would create a (nested) Tuple containing zeros static Expr NestedZeros(const Type& ty) { AdjointMsg msg = MapToNestedMsg(ty, [](Type ty) { - auto* tensor_ty = ty.as(); + auto* tensor_ty = ty.as(); TVM_FFI_ICHECK(tensor_ty) << "The leaf of adjoint should be a Tensor."; TVM_FFI_ICHECK(tensor_ty->shape.has_value()) << "Missing shape when building zeros tuple."; const Expr& init = zeros(tensor_ty->shape.value(), tensor_ty->dtype.value()->dtype); @@ -570,9 +570,9 @@ class BackwardBindingGenerator : private ExprVisitor { static Expr TupleAwareAdd(const Expr& lhs, const Expr& rhs) { AdjointMsg res = CombineNestedMsg( ExprToAdjointMsg(lhs), ExprToAdjointMsg(rhs), [](Expr l_leaf, Expr r_leaf) { - auto* ty = GetTypeAs(l_leaf); + auto* ty = GetTypeAs(l_leaf); TVM_FFI_ICHECK(ty) << "The leaf of adjoint should have Type and be a Tensor."; - TVM_FFI_ICHECK(GetTypeAs(r_leaf)) + TVM_FFI_ICHECK(GetTypeAs(r_leaf)) << "The leaf of adjoint should have Type and be a Tensor."; Expr res = add(l_leaf, r_leaf); UpdateType(res, ffi::GetRef(ty)); @@ -588,13 +588,13 @@ class BackwardBindingGenerator : private ExprVisitor { // Step 2）t2_new = t2 + increment (TupleAwareAdd) // Step 3) tuple_new = (t1, t2_new, t3) static Expr AddInTuple(const Expr& tuple, int index, const Expr& increment) { - auto* ty = GetTypeAs(tuple); + auto* ty = GetTypeAs(tuple); TVM_FFI_ICHECK(ty) << "The first argument of AddInTuple should have tuple type."; TVM_FFI_ICHECK(index >= 0 && index < static_cast(ty->fields.size())); ffi::Array res; for (size_t i = 0; i < ty->fields.size(); ++i) { Expr field = [&]() -> Expr { - if (const auto* expr_tuple = tuple.as()) { + if (const auto* expr_tuple = tuple.as()) { return expr_tuple->fields[i]; } return TupleGetItem(tuple, i); @@ -622,7 +622,7 @@ class GradientMutator : private ExprMutator { static IRModule Transform(IRModule mod, ffi::String func_name, ffi::Optional> require_grads, int target_index) { // Step 1. Copy function - auto* old_func = mod->Lookup(func_name).as(); + auto* old_func = mod->Lookup(func_name).as(); TVM_FFI_ICHECK(old_func) << func_name << "is not a Relax Function"; auto copier = FunctionCopier(); auto new_func = copier.Copy(ffi::GetRef(old_func)); @@ -681,21 +681,21 @@ class GradientMutator : private ExprMutator { return builder_->GetContextIRModule(); } - Expr VisitExpr_(const FunctionNode* func) final { + Expr VisitExpr_(const FunctionObj* func) final { orig_params_ = func->params; Expr new_body = this->VisitExpr(func->body); return Function(func->params, new_body, std::nullopt, func->is_pure, func->attrs); } - Expr VisitExpr_(const SeqExprNode* seq_expr) final { + Expr VisitExpr_(const SeqExprObj* seq_expr) final { // TODO(chaofan, yixin): multiple blocks AD TVM_FFI_ICHECK(seq_expr->blocks.size() == 1) << "now only support one dataflow block"; // TODO(chaofan, yixin): AD in non-dataflow block. - TVM_FFI_ICHECK(seq_expr->blocks[0]->IsInstance()) + TVM_FFI_ICHECK(seq_expr->blocks[0]->IsInstance()) << "now only support one dataflow block"; - // the return value should be a VarNode, and a scalar + // the return value should be a VarObj, and a scalar orig_return_expr_ = seq_expr->body; CheckAndSetTarget(seq_expr->body, target_index_); @@ -703,7 +703,7 @@ class GradientMutator : private ExprMutator { return SeqExpr({new_block}, return_expr_); } - BindingBlock VisitBindingBlock_(const DataflowBlockNode* block) final { + BindingBlock VisitBindingBlock_(const DataflowBlockObj* block) final { builder_->BeginDataflowBlock(); // accept bindings in the original block for (const auto& binding : block->bindings) { @@ -719,7 +719,7 @@ class GradientMutator : private ExprMutator { } static bool IsFloatTensorType(const Type& ty) { - auto* tensor_ty = ty.as(); + auto* tensor_ty = ty.as(); // Gradient eligibility preserves the old float-kind check; lanes do not affect this policy. return tensor_ty && !tensor_ty->IsUnknownDtype() && tensor_ty->dtype.value().MatchesCode(DLDataTypeCode::kDLFloat); @@ -729,19 +729,19 @@ class GradientMutator : private ExprMutator { // when the return value is a Tuple, the target is the target_index-th field of the return value // Check that the target should be a Var of scalar tensor ty void CheckAndSetTarget(const Expr& e, int target_index) { - if (auto* var = e.as()) { + if (auto* var = e.as()) { TVM_FFI_ICHECK_EQ(target_index, 0) << "When the function has only one return value, target_index can " "only be 0. But the target_index specified is " << target_index; target_var_ = ffi::GetRef(var); - } else if (auto* tuple = e.as()) { + } else if (auto* tuple = e.as()) { TVM_FFI_ICHECK(target_index >= 0 && target_index < static_cast(tuple->fields.size())) << "target_index should be in the range of the number of return values of the " "function. " "But the specified target_index is " << target_index << ", while the number of return values is " << tuple->fields.size(); - auto* var = tuple->fields[target_index].as(); + auto* var = tuple->fields[target_index].as(); TVM_FFI_ICHECK(var) << "Target must be a Var, but the specified target is " << tuple->fields[target_index]; target_var_ = ffi::GetRef(var); diff --git a/src/relax/transform/gradient_simplifier.cc b/src/relax/transform/gradient_simplifier.cc index b6075ef3810b..92ba4fdf7659 100644 --- a/src/relax/transform/gradient_simplifier.cc +++ b/src/relax/transform/gradient_simplifier.cc @@ -80,7 +80,7 @@ class GradientSimplifier : private ExprMutator { } private: - static bool IsTransposeOp(const CallNode* call_node) { + static bool IsTransposeOp(const CallObj* call_node) { if (!call_node->op.same_as(Op::Get("relax.permute_dims"))) { return false; } @@ -126,23 +126,23 @@ class GradientSimplifier : private ExprMutator { // permute_dims op // Else, return permute_dims(expr) Expr GetTransposeAccordingToCtx(const Expr& expr) { - if (!expr->IsInstance()) { + if (!expr->IsInstance()) { return GetTransposeOf(expr); } auto prev_expr = builder_->LookupBinding(expr.as_or_throw()); - if (!prev_expr || !prev_expr.value()->IsInstance()) { + if (!prev_expr || !prev_expr.value()->IsInstance()) { return GetTransposeOf(expr); } - auto prev_call_node = prev_expr.as(); + auto prev_call_node = prev_expr.as(); if (!IsTransposeOp(prev_call_node)) { return GetTransposeOf(expr); } return prev_call_node->args[0]; } - void VisitBinding_(const VarBindingNode* binding, const CallNode* call_node) { + void VisitBinding_(const VarBindingObj* binding, const CallObj* call_node) { auto result = ExprMutator::VisitExpr(ffi::GetRef(call_node)); - auto new_call_node = result.as(); + auto new_call_node = result.as(); auto reemit_and_return = [&]() { ReEmitBinding(binding, result); return; @@ -153,19 +153,19 @@ class GradientSimplifier : private ExprMutator { } auto arg = new_call_node->args[0]; - if (!arg->IsInstance()) { + if (!arg->IsInstance()) { return reemit_and_return(); } auto prev_expr = builder_->LookupBinding(arg.as_or_throw()); - if (!prev_expr || !prev_expr.value()->IsInstance()) { + if (!prev_expr || !prev_expr.value()->IsInstance()) { return reemit_and_return(); } - auto prev_call_node = prev_expr.as(); + auto prev_call_node = prev_expr.as(); if (IsTransposeOp(prev_call_node)) { // rewrite rule #1: permute_dims(permute_dims(a)) -> a - if (prev_call_node->args[0]->IsInstance()) { + if (prev_call_node->args[0]->IsInstance()) { var_remap_.insert_or_assign(binding->var, prev_call_node->args[0].as_or_throw()); return; } else { diff --git a/src/relax/transform/infer_amp_utils.cc b/src/relax/transform/infer_amp_utils.cc index f9f67d13dba2..7b2608aff690 100644 --- a/src/relax/transform/infer_amp_utils.cc +++ b/src/relax/transform/infer_amp_utils.cc @@ -24,7 +24,7 @@ namespace relax { NType NTypeFrom(const Type& ty, DLDataType dtype) { auto fmapleaf = [&](const Type& ty) -> NType { - const auto* tensor = ty.as(); + const auto* tensor = ty.as(); TVM_FFI_ICHECK(tensor) << "Expected TensorType, but got " << ty; if (dtype == DLDataType{kDLOpaqueHandle, 0, 0}) { if (tensor->IsUnknownDtype()) return NType(ffi::String("")); diff --git a/src/relax/transform/infer_layout_utils.cc b/src/relax/transform/infer_layout_utils.cc index cb97b47fe75f..3974a0c8ec5a 100644 --- a/src/relax/transform/infer_layout_utils.cc +++ b/src/relax/transform/infer_layout_utils.cc @@ -112,7 +112,7 @@ LayoutDecision InitialLayoutDecision(int ndim) { NLayout InitialNLayout(const Type& ty) { auto fmapleaf = [&](const Type& ty) -> NLayout { - if (const auto* tensor_ty = ty.as()) { + if (const auto* tensor_ty = ty.as()) { return NLayout(InitialLayoutDecision(tensor_ty->ndim)); } return LayoutDecision::InitUnknownDim(); @@ -130,14 +130,14 @@ LayoutDecision GetLayoutDecision(const VarLayoutMap& var_layout_map, const Expr& NLayout GetNLayout(const VarLayoutMap& var_layout_map, const Expr& arg) { auto fmapleaf = [&](const Expr& expr) -> NLayout { - if (const auto* var = expr.as()) { + if (const auto* var = expr.as()) { auto it = var_layout_map.find(ffi::GetRef(var)); if (it != var_layout_map.end()) { return (*it).second; } else { return InitialNLayout(expr); } - } else if (const auto* constant = expr.as(); + } else if (const auto* constant = expr.as(); constant && constant->value.as()) { return InitialLayoutDecision(constant->value.cast().shape().size()); } @@ -148,7 +148,7 @@ NLayout GetNLayout(const VarLayoutMap& var_layout_map, const Expr& arg) { bool NoDesiredLayout(const Call& call, const ffi::Map>& desired_layouts) { - const OpNode* op_node = call->op.as(); + const OpObj* op_node = call->op.as(); if (op_node == nullptr) return false; const auto& it = desired_layouts.find(op_node->name); return it == desired_layouts.end(); diff --git a/src/relax/transform/inline_functions.cc b/src/relax/transform/inline_functions.cc index f5203c658fb2..4c41a81b5584 100644 --- a/src/relax/transform/inline_functions.cc +++ b/src/relax/transform/inline_functions.cc @@ -43,7 +43,7 @@ class FunctionInliner : public ExprMutator { using ExprMutator::VisitExpr_; - Expr VisitExpr_(const FunctionNode* op) override { + Expr VisitExpr_(const FunctionObj* op) override { auto node = ExprMutator::VisitExpr_(op); if (node.get() != op) { node = CanonicalizeBindings(node); @@ -52,7 +52,7 @@ class FunctionInliner : public ExprMutator { return node; } - Expr VisitExpr_(const CallNode* op) override { + Expr VisitExpr_(const CallObj* op) override { auto node = ExprMutator::VisitExpr_(op).as_or_throw(); if (auto opt = node->op.as()) { @@ -130,13 +130,13 @@ class FunctionInliner : public ExprMutator { // // This implementation uses Option 4. - Var param_var(func->params[i]->name, args[i]->ty.as()); + Var param_var(func->params[i]->name, args[i]->ty); param_bindings.push_back(VarBinding(param_var, args[i])); param_map.Set(func->params[i], param_var); } DataflowBlock binding_block(param_bindings); - Expr body = Bind(func, param_map).as()->body; + Expr body = Bind(func, param_map).as()->body; return SeqExpr({binding_block}, body); } @@ -155,7 +155,7 @@ class FunctionInliner : public ExprMutator { Function FunctionInlineFunctions( Function func, const ffi::Map, Function>& replacements) { for (const auto& [key, func] : replacements) { - if (auto ptr = key.as()) { + if (auto ptr = key.as()) { TVM_FFI_CHECK(!replacements.count(ptr->name_hint), ValueError) << "Map of functions to inline must be unambiguous. " << "However, the map provided contains both the GlobalVar " << key << " and the string \'" diff --git a/src/relax/transform/kill_after_last_use.cc b/src/relax/transform/kill_after_last_use.cc index 8f0e6096bab9..6a82277d8d3d 100644 --- a/src/relax/transform/kill_after_last_use.cc +++ b/src/relax/transform/kill_after_last_use.cc @@ -43,28 +43,28 @@ class UnusedTrivialBindingRemover : public ExprMutator { public: static Expr Apply(Expr expr) { struct UsedCollector : ExprVisitor { - void VisitExpr_(const VarNode* val) override { used.insert(val); } - void VisitExpr_(const DataflowVarNode* val) override { - VisitExpr_(static_cast(val)); + void VisitExpr_(const VarObj* val) override { used.insert(val); } + void VisitExpr_(const DataflowVarObj* val) override { + VisitExpr_(static_cast(val)); } - void VisitBinding_(const VarBindingNode* binding, const VarNode* val) override { + void VisitBinding_(const VarBindingObj* binding, const VarObj* val) override { has_trivial_binding.insert(binding->var.get()); ExprVisitor::VisitBinding_(binding, val); } - void VisitBinding_(const MatchCastNode* binding) override { - if (binding->value.as() && + void VisitBinding_(const MatchCastObj* binding) override { + if (binding->value.as() && ffi::StructuralEqual()(GetType(binding->var), GetType(binding->value))) { has_trivial_binding.insert(binding->var.get()); } ExprVisitor::VisitBinding_(binding); } - void VisitBinding_(const VarBindingNode* binding, const DataflowVarNode* val) override { - VisitBinding_(binding, static_cast(val)); + void VisitBinding_(const VarBindingObj* binding, const DataflowVarObj* val) override { + VisitBinding_(binding, static_cast(val)); } - std::unordered_set used; - std::unordered_set has_trivial_binding; + std::unordered_set used; + std::unordered_set has_trivial_binding; }; UsedCollector collector; @@ -80,7 +80,7 @@ class UnusedTrivialBindingRemover : public ExprMutator { } private: - explicit UnusedTrivialBindingRemover(std::unordered_set to_remove) + explicit UnusedTrivialBindingRemover(std::unordered_set to_remove) : to_remove_(std::move(to_remove)) {} void VisitBinding(const Binding& binding) override { @@ -89,17 +89,17 @@ class UnusedTrivialBindingRemover : public ExprMutator { } } - std::unordered_set to_remove_; + std::unordered_set to_remove_; }; class CollectLastUsage : public ExprVisitor { public: struct LastUsage { - std::vector tensors; - std::vector storage; - std::vector objects; + std::vector tensors; + std::vector storage; + std::vector objects; }; - using Result = std::unordered_map; + using Result = std::unordered_map; static Result Collect(const Expr& expr) { CollectLastUsage visitor; @@ -124,13 +124,12 @@ class CollectLastUsage : public ExprVisitor { // codegen special case should be reflected here as well. bool stored_in_vm_register = !(visitor.constant_tensors_.count(var) || visitor.null_value_objects_.count(var) || - var->ty.as() || var->ty.as() || - var->ty.as()); + var->ty.as() || var->ty.as() || var->ty.as()); if (!is_output && !already_killed) { if (visitor.storage_objects_.count(var)) { output[last_usage_point].storage.push_back(var); - } else if (var->ty.as() && stored_in_vm_register) { + } else if (var->ty.as() && stored_in_vm_register) { output[last_usage_point].tensors.push_back(var); } else if (stored_in_vm_register) { output[last_usage_point].objects.push_back(var); @@ -150,7 +149,7 @@ class CollectLastUsage : public ExprVisitor { current_binding_ = cache; } - void VisitExpr_(const VarNode* op) override { + void VisitExpr_(const VarObj* op) override { ExprVisitor::VisitExpr_(op); // Overwrite any previous usage, such that after the visitor // completes, last_usage_of_ contains the last usage point. If @@ -159,7 +158,7 @@ class CollectLastUsage : public ExprVisitor { last_usage_of_[op] = current_binding_; } - void VisitBinding_(const VarBindingNode* binding, const CallNode* val) override { + void VisitBinding_(const VarBindingObj* binding, const CallObj* val) override { static const Op vm_alloc_storage = Op::Get("relax.vm.alloc_storage"); static const Op mem_alloc_storage = Op::Get("relax.memory.alloc_storage"); static const Op null_value_op = Op::Get("relax.null_value"); @@ -177,7 +176,7 @@ class CollectLastUsage : public ExprVisitor { TVM_FFI_ICHECK_EQ(val->args.size(), 1) << "Operator " << val->op << " should have one argument, " << "but instead found " << val->args.size() << " arguments: " << val->args; - auto killed_object = val->args[0].as(); + auto killed_object = val->args[0].as(); TVM_FFI_ICHECK(killed_object) << "Internal error: non-normalized expression " << ffi::GetRef(val); killed_objects_.insert(killed_object); @@ -187,54 +186,54 @@ class CollectLastUsage : public ExprVisitor { } } - void VisitBinding_(const VarBindingNode* binding, const GenericConstNode* val) override { + void VisitBinding_(const VarBindingObj* binding, const GenericConstObj* val) override { if (val->value.as()) constant_tensors_.insert(binding->var.get()); } private: // The current binding being visited, or nullptr if no binding is // being visited. - const VarNode* current_binding_{nullptr}; + const VarObj* current_binding_{nullptr}; // Order of bindings, to ensure consistent order of destruction, in // case a Binding is the last usage for more than one variable. - std::vector binding_order_; + std::vector binding_order_; // Map from a variable to the last variable binding that makes use // of it. - std::unordered_map last_usage_of_; + std::unordered_map last_usage_of_; // Storage objects, eligible for R.vm.kill_object. This cannot be // determined solely from the Type, because the // `R.*.alloc_storage` operators return AnyType - std::unordered_set storage_objects_; + std::unordered_set storage_objects_; // Constants, which do not have a VM register, and may *not* have // R.builtin.kill_tensor called on them. - std::unordered_set constant_tensors_; + std::unordered_set constant_tensors_; // Variables bound to `relax.null_value`, which do not occupy a VM // register in CodeGenVM, and therefore must // never be passed to R.vm.kill_object. - std::unordered_set null_value_objects_; + std::unordered_set null_value_objects_; // Set of objects that already have a call node to kill them. Should not have a duplicate - std::unordered_set killed_objects_; + std::unordered_set killed_objects_; // Trivial var-to-var bindings. - std::unordered_map trivial_bindings_; + std::unordered_map trivial_bindings_; }; class KillInserter : public ExprMutator { private: - Expr VisitExpr_(const FunctionNode* op) override { + Expr VisitExpr_(const FunctionObj* op) override { last_usage_ = CollectLastUsage::Collect(ffi::GetRef(op)); auto mutated = ExprMutator::VisitExpr_(op); last_usage_.clear(); return mutated; } - Expr VisitExpr_(const SeqExprNode* op) override { + Expr VisitExpr_(const SeqExprObj* op) override { last_usage_ = CollectLastUsage::Collect(ffi::GetRef(op)); auto mutated = ExprMutator::VisitExpr_(op); last_usage_.clear(); @@ -251,13 +250,13 @@ class KillInserter : public ExprMutator { } static const Op mem_kill_storage = Op::Get("relax.memory.kill_storage"); - for (const VarNode* storage_obj : it->second.storage) { + for (const VarObj* storage_obj : it->second.storage) { builder_->Emit(Call(Type::Missing(), mem_kill_storage, {ffi::GetRef(storage_obj)}), /*name_hint=*/"_"); } static const Op vm_kill_object = Op::Get("relax.vm.kill_object"); - for (const VarNode* obj : it->second.objects) { + for (const VarObj* obj : it->second.objects) { builder_->Emit(Call(Type::Missing(), vm_kill_object, {ffi::GetRef(obj)}), /*name_hint=*/"_"); } diff --git a/src/relax/transform/lambda_lift.cc b/src/relax/transform/lambda_lift.cc index 75debf420147..258228ecd493 100644 --- a/src/relax/transform/lambda_lift.cc +++ b/src/relax/transform/lambda_lift.cc @@ -43,7 +43,7 @@ namespace { /* \brief Collect names of functions to be lifted out */ class LambdaNameCollector : ExprVisitor { public: - static std::unordered_map Collect(const IRModule& mod) { + static std::unordered_map Collect(const IRModule& mod) { LambdaNameCollector visitor; for (const auto& [gvar, base_func] : mod->functions) { @@ -62,7 +62,7 @@ class LambdaNameCollector : ExprVisitor { } private: - void VisitBinding_(const VarBindingNode* binding, const FunctionNode* func) override { + void VisitBinding_(const VarBindingObj* binding, const FunctionObj* func) override { if (auto opt = func->GetAttr(tvm::attr::kGlobalSymbol)) { ffi::String public_name = opt.value(); @@ -105,13 +105,13 @@ class LambdaNameCollector : ExprVisitor { } // De-duplication of collected names - std::unordered_map Finalize() const { + std::unordered_map Finalize() const { // The functions which still must be assigned a name - std::unordered_map> remaining_to_name = + std::unordered_map> remaining_to_name = lambda_location_; // Collecting the functions that now have a name. - std::unordered_map lifted_names; + std::unordered_map lifted_names; // A lookup for names that are unavailable for use. std::unordered_set unavailable_names = previous_global_vars_; @@ -120,7 +120,7 @@ class LambdaNameCollector : ExprVisitor { // `proposed_name_generation_func` should be a function with // signature: // - // ffi::Optional func(const FunctionNode*, const ffi::Array&) + // ffi::Optional func(const FunctionObj*, const ffi::Array&) // // The first argument will be the lambda function being lifted. // The second argument will be the nested location where that @@ -139,7 +139,7 @@ class LambdaNameCollector : ExprVisitor { return; } - std::unordered_map new_names; + std::unordered_map new_names; for (const auto& [func, location] : remaining_to_name) { if (ffi::Optional opt_proposed_name = proposed_name_generation_func(func, location)) { @@ -168,7 +168,7 @@ class LambdaNameCollector : ExprVisitor { }; // 1. Start with any publicly explosed names from kGlobalSymbol - attempt_name_generation([&](const FunctionNode* func, + attempt_name_generation([&](const FunctionObj* func, const auto&) -> ffi::Optional { if (auto it = lifted_with_global_symbol_.find(func); it != lifted_with_global_symbol_.end()) { return it->second; @@ -179,7 +179,7 @@ class LambdaNameCollector : ExprVisitor { // 2. Try concatenating the name of the relax variable with the // name of the function that contains it. - attempt_name_generation([&](const FunctionNode*, const auto& location) -> ffi::String { + attempt_name_generation([&](const FunctionObj*, const auto& location) -> ffi::String { std::stringstream stream; stream << location.front() << "_" << location.back(); return stream.str(); @@ -188,7 +188,7 @@ class LambdaNameCollector : ExprVisitor { // 3. Try concatenating the entire path together. Don't include // paths of length 2, as they would already be attempted earlier. attempt_name_generation( - [&](const FunctionNode*, const auto& location) -> ffi::Optional { + [&](const FunctionObj*, const auto& location) -> ffi::Optional { if (location.size() == 2) return std::nullopt; std::stringstream stream; @@ -207,7 +207,7 @@ class LambdaNameCollector : ExprVisitor { // 4. Fallback. Count the number of times a relax variable with // that name was used. std::unordered_map usage_count; - attempt_name_generation([&](const FunctionNode*, const auto& location) -> ffi::String { + attempt_name_generation([&](const FunctionObj*, const auto& location) -> ffi::String { std::stringstream stream; stream << location.front() << "_" << location.back(); int usage = usage_count[stream.str()]++; @@ -225,8 +225,8 @@ class LambdaNameCollector : ExprVisitor { ffi::Array name_stack_; std::unordered_set previous_global_vars_; std::unordered_map> new_public_names_; - std::unordered_map lifted_with_global_symbol_; - std::unordered_map> lambda_location_; + std::unordered_map lifted_with_global_symbol_; + std::unordered_map> lambda_location_; }; } // namespace @@ -244,7 +244,7 @@ class LambdaLifter : public ExprMutator { using ExprMutator::VisitExpr_; - void VisitBinding_(const VarBindingNode* binding, const FunctionNode* func_node) final { + void VisitBinding_(const VarBindingObj* binding, const FunctionObj* func_node) final { auto cache = current_lambda_var_; current_lambda_var_ = binding->var; @@ -256,7 +256,7 @@ class LambdaLifter : public ExprMutator { current_lambda_var_ = cache; } - Expr VisitExpr_(const FunctionNode* func_node) final { + Expr VisitExpr_(const FunctionObj* func_node) final { if (!current_lambda_var_) { // Early bail-out for top-level functions return ExprMutator::VisitExpr_(func_node); @@ -358,7 +358,7 @@ class LambdaLifter : public ExprMutator { return callable_value; } - Expr VisitExpr_(const CallNode* call_node) final { + Expr VisitExpr_(const CallObj* call_node) final { auto call = ffi::GetRef(call_node); auto orig_ty = call->ty.as_or_throw(); @@ -369,14 +369,14 @@ class LambdaLifter : public ExprMutator { // Call "relax.invoke_closure" to invoke closure auto bound_value = LookupBinding(var); - if (IsClosure(var) && bound_value.as()) { + if (IsClosure(var) && bound_value.as()) { // if the original op was pure, we should use invoke_pure_closure Call orig_call = bound_value.value().as_or_throw(); bool is_pure = [&]() -> bool { if (auto op = orig_call->op.as()) { static const auto& purity_map = Op::GetAttrMap(tvm::relax::op_attr::kPurity); return purity_map.get(op.value(), false); - } else if (const auto* func_ty = orig_call->op->ty.as()) { + } else if (const auto* func_ty = orig_call->op->ty.as()) { return func_ty->purity; } else { TVM_FFI_THROW(InternalError) @@ -411,7 +411,7 @@ class LambdaLifter : public ExprMutator { return ExprMutator::VisitExpr_(call.get()); } - Expr VisitExpr_(const VarNode* op) override { + Expr VisitExpr_(const VarObj* op) override { auto var = ffi::GetRef(op); if (auto it = rebind_map_.find(var); it != rebind_map_.end()) { return it->second; @@ -429,7 +429,7 @@ class LambdaLifter : public ExprMutator { } } - if (const auto* call_node = val.as()) { + if (const auto* call_node = val.as()) { // recursive call auto op = call_node->op; if (auto local_var = op.as()) { @@ -440,24 +440,24 @@ class LambdaLifter : public ExprMutator { return make_closure_op_.same_as(op); } - } else if (const auto* global_var = val.as()) { + } else if (const auto* global_var = val.as()) { if (closures_.count(ffi::GetRef(global_var))) { return true; } IRModule ctx_mod = builder_->GetContextIRModule(); TVM_FFI_ICHECK(ctx_mod->functions.size() > 0); BaseFunc func = ctx_mod->Lookup(ffi::GetRef(global_var)); - const auto* func_node = func.as(); + const auto* func_node = func.as(); if (func_node) { return IsClosure(func_node->body); } else { return false; } - } else if (const auto* func_node = val.as()) { + } else if (const auto* func_node = val.as()) { return IsClosure(func_node->body); - } else if (const auto* seq_node = val.as()) { + } else if (const auto* seq_node = val.as()) { return IsClosure(seq_node->body); } else { @@ -487,7 +487,7 @@ class LambdaLifter : public ExprMutator { ffi::Optional current_lambda_var_ = std::nullopt; IRModule mod_; - std::unordered_map lifted_names_; + std::unordered_map lifted_names_; /*! \brief Cache ops that would be used later to reduce lookup overhead. */ const Op make_closure_op_ = Op::Get("relax.make_closure"); diff --git a/src/relax/transform/lazy_transform_params.cc b/src/relax/transform/lazy_transform_params.cc index bb51c7438dc8..755293def379 100644 --- a/src/relax/transform/lazy_transform_params.cc +++ b/src/relax/transform/lazy_transform_params.cc @@ -37,7 +37,7 @@ namespace relax { using namespace tvm::prim; namespace { -std::optional GetNumInputParams(const FunctionNode* func) { +std::optional GetNumInputParams(const FunctionObj* func) { if (auto opt_int_imm = func->GetAttr(tvm::relax::attr::kNumInput.name)) { int64_t num_input_params = static_cast(opt_int_imm.value()->value); TVM_FFI_CHECK_GE(num_input_params, 0, ValueError) @@ -55,7 +55,7 @@ std::optional GetNumInputParams(const FunctionNode* func) { class LazyInputMutator : public ExprMutator { public: - Expr VisitExpr_(const FunctionNode* func) override { + Expr VisitExpr_(const FunctionObj* func) override { if (plan_.has_value()) { return ExprMutator::VisitExpr_(func); } @@ -94,7 +94,7 @@ class LazyInputMutator : public ExprMutator { return output; } - Expr VisitExpr_(const VarNode* op) override { + Expr VisitExpr_(const VarObj* op) override { if (plan_) { Var var = ffi::GetRef(op); if (auto it = plan_->param_lookup.find(var); it != plan_->param_lookup.end()) { @@ -121,7 +121,7 @@ class LazyInputMutator : public ExprMutator { class LazyOutputMutator : public ExprMutator { public: - Expr VisitExpr_(const FunctionNode* func) override { + Expr VisitExpr_(const FunctionObj* func) override { if (plan_.has_value()) { return ExprMutator::VisitExpr_(func); } @@ -137,7 +137,7 @@ class LazyOutputMutator : public ExprMutator { }; auto func_body = func->body.as_or_throw(); - if (auto tuple_output = func_body->body.as()) { + if (auto tuple_output = func_body->body.as()) { for (size_t i = 0; i < tuple_output->fields.size(); i++) { define_lookup(i, tuple_output->fields[i]); } diff --git a/src/relax/transform/legalize_ops.cc b/src/relax/transform/legalize_ops.cc index 014f77e6c6fe..3341b700a59d 100644 --- a/src/relax/transform/legalize_ops.cc +++ b/src/relax/transform/legalize_ops.cc @@ -55,14 +55,14 @@ TVM_FFI_STATIC_INIT_BLOCK() { * \return A boolean indicating the given type contains shape values that are all known. */ bool KnowAllShapeValues(const Type& ty) { - if (const auto* tensor_ty = ty.as()) { - return tensor_ty->shape.has_value() && tensor_ty->shape.value()->IsInstance(); - } else if (const auto* shape_ty = ty.as()) { + if (const auto* tensor_ty = ty.as()) { + return tensor_ty->shape.has_value() && tensor_ty->shape.value()->IsInstance(); + } else if (const auto* shape_ty = ty.as()) { return shape_ty->values.has_value(); - } else if (const auto* tuple_ty = ty.as()) { + } else if (const auto* tuple_ty = ty.as()) { return std::all_of(tuple_ty->fields.begin(), tuple_ty->fields.end(), [](Type field_ty) { return KnowAllShapeValues(field_ty); }); - } else if (ty.as()) { + } else if (ty.as()) { return true; } else { return false; @@ -89,7 +89,7 @@ class LegalizeMutator : public ExprMutator { IRModule Transform() { for (const auto& gv : mod_->GetGlobalVars()) { const auto& func = mod_->Lookup(gv); - if (func->IsInstance()) { + if (func->IsInstance()) { auto updated_func = this->VisitExpr(func).as_or_throw(); builder_->UpdateFunction(gv, updated_func.as_or_throw()); } @@ -118,7 +118,7 @@ class LegalizeMutator : public ExprMutator { bool WrapPureCondition(const Op& op, const Expr& legalized) { static const auto& purity_map = Op::GetAttrMap(tvm::relax::op_attr::kPurity); - const CallNode* call = legalized.as(); + const CallObj* call = legalized.as(); if (!call) { // Unlikely for this condition to be met, but it is possible. @@ -131,7 +131,7 @@ class LegalizeMutator : public ExprMutator { bool pure_legalized_op = [&]() -> bool { if (auto legalized_op = call->op.as()) { return purity_map.get(legalized_op.value(), false); - } else if (auto func_ty = call->op->ty.as()) { + } else if (auto func_ty = call->op->ty.as()) { return func_ty->purity; } else { return false; @@ -157,14 +157,14 @@ class LegalizeMutator : public ExprMutator { ffi::Optional GetTarget(const ffi::Array& types) { for (auto ty : types) { - if (const auto* tinfo = ty.as()) { + if (const auto* tinfo = ty.as()) { if (tinfo->vdevice.has_value()) { auto vdevice = tinfo->vdevice.value(); if (vdevice->target.has_value()) { return vdevice->target; } } - } else if (const auto* tup_ty = ty.as()) { + } else if (const auto* tup_ty = ty.as()) { return GetTarget(tup_ty->fields); } } @@ -172,7 +172,7 @@ class LegalizeMutator : public ExprMutator { } Expr BindTarget(Expr expr) { - if (!expr->IsInstance()) { + if (!expr->IsInstance()) { // FLegalize returned something other than a relax::Call. This // post-processing only handles cases where legalization // produces a lowered call node. In principle, this @@ -240,7 +240,7 @@ class LegalizeMutator : public ExprMutator { return call; } - Expr VisitExpr_(const CallNode* call) final { + Expr VisitExpr_(const CallObj* call) final { Call visited_call = this->VisitExprPostOrder_(call).as_or_throw(); static const auto& legalize_map = Op::GetAttrMap(tvm::relax::op_attr::kLegalize); static const auto& call_packed_map = Op::GetAttrMap(tvm::relax::op_attr::kCallPacked); @@ -249,9 +249,9 @@ class LegalizeMutator : public ExprMutator { static const Op call_pure_packed_op = Op::Get("relax.call_pure_packed"); static const Op call_tir_op = Op::Get("relax.call_tir"); static const Op call_dps_packed_op = Op::Get("relax.call_dps_packed"); - auto* op_node = visited_call->op.as(); + auto* op_node = visited_call->op.as(); - // Not an OpNode + // Not an OpObj if (op_node == nullptr) { return visited_call; } @@ -409,7 +409,7 @@ class LegalizeMutator : public ExprMutator { /*! \brief If VDevice annotations produced at least one tirx::Function with a Target attr*/ bool generated_tir_with_target_attr_{false}; /*! - * \brief A boolean value indicating if to print warnings for CallNode whose op's + * \brief A boolean value indicating if to print warnings for CallObj whose op's * legalization function is not registered. */ bool enable_warning_; diff --git a/src/relax/transform/lift_transform_params.cc b/src/relax/transform/lift_transform_params.cc index cab95099c072..94d69828f88f 100644 --- a/src/relax/transform/lift_transform_params.cc +++ b/src/relax/transform/lift_transform_params.cc @@ -305,7 +305,7 @@ struct LocalCollectInfo : public BaseCollectInfo { } using ExprMutator::VisitExpr_; - Expr VisitExpr_(const CallNode* call) override { + Expr VisitExpr_(const CallObj* call) override { static const Op stop_lift_params_op = Op::Get("relax.builtin.stop_lift_params"); if (call->op.same_as(stop_lift_params_op)) { return VisitExpr(call->args[0]); @@ -329,7 +329,7 @@ struct LocalCollectInfo : public BaseCollectInfo { class BaseLiftableBindingCollector : public ExprVisitor { protected: - void VisitBindingBlock_(const DataflowBlockNode* block) final { + void VisitBindingBlock_(const DataflowBlockObj* block) final { bool cache = is_in_dataflow_block_; is_in_dataflow_block_ = true; ExprVisitor::VisitBindingBlock_(block); @@ -345,7 +345,7 @@ class BaseLiftableBindingCollector : public ExprVisitor { } // Cond 2. Do not lift regarding the "builtin.stop_lift_params" op. - if (const auto* call = value.as()) { + if (const auto* call = value.as()) { static const Op stop_lift_params_op = Op::Get("relax.builtin.stop_lift_params"); if (call->op.same_as(stop_lift_params_op)) { return false; @@ -368,7 +368,7 @@ class BaseLiftableBindingCollector : public ExprVisitor { } // Cond 5. Do not lift declarations of external functions - if (value.as()) { + if (value.as()) { return false; } @@ -412,7 +412,7 @@ class LocalLiftableBindingCollector : public BaseLiftableBindingCollector { : info_(std::move(func)) { info_.global_info = global_info; } - void VisitExpr_(const FunctionNode* func) override { + void VisitExpr_(const FunctionObj* func) override { size_t num_runtime_params = func->params.size(); if (auto opt = func->attrs.GetAttr(tvm::relax::attr::kNumInput)) { num_runtime_params = opt.value(); @@ -515,7 +515,7 @@ class ParamRemapper : private ExprFunctor { } private: - void VisitExpr_(const VarNode* lhs_var, const Expr& rhs_expr) final { + void VisitExpr_(const VarObj* lhs_var, const Expr& rhs_expr) final { auto rhs_var = rhs_expr.as_or_throw(); if (auto it = var_remap_.find(ffi::GetRef(lhs_var)); it != var_remap_.end()) { TVM_FFI_ICHECK((*it).second.same_as(rhs_var)); @@ -592,7 +592,7 @@ class GlobalLiftableBindingCollector : public BaseLiftableBindingCollector { explicit GlobalLiftableBindingCollector(const ffi::Map& var_remap) : var_remap_(var_remap) {} void VisitBinding(const Binding& binding) override { - TVM_FFI_ICHECK(!binding->IsInstance()) + TVM_FFI_ICHECK(!binding->IsInstance()) << "MatchCast is not supported in global lifting"; if (CanLiftBinding(binding)) { liftable_vars_.insert(binding->var); @@ -639,7 +639,7 @@ inline bool ends_with(const std::string& value, const std::string& ending) { */ class ConsumeBundledParams : public ExprMutator { public: - void VisitBinding_(const VarBindingNode* binding, const TupleGetItemNode* tuple_get_item) final { + void VisitBinding_(const VarBindingObj* binding, const TupleGetItemObj* tuple_get_item) final { static const auto call_pure_packed = Op::Get("relax.call_pure_packed"); static const auto& builtin_tuple_reset_item = ExternFunc("vm.builtin.tuple_reset_item"); if (tuple_get_item->tuple.same_as(params_)) { @@ -659,13 +659,13 @@ class ConsumeBundledParams : public ExprMutator { } } - Expr VisitExpr_(const FunctionNode* func) final { + Expr VisitExpr_(const FunctionObj* func) final { auto opt_num_input = func->GetAttr(tvm::relax::attr::kNumInput); TVM_FFI_ICHECK(opt_num_input.has_value()); auto num_input = opt_num_input.value(); TVM_FFI_ICHECK_EQ(func->params.size(), num_input + 1); params_ = func->params.back(); - TVM_FFI_ICHECK(params_->ty.as()); + TVM_FFI_ICHECK(params_->ty.as()); return ExprMutator::VisitExpr_(func); } @@ -709,7 +709,7 @@ std::vector> GetTargetFunctions( // Get all the functions that have the `num_input` attribute, and // are not already the result of `LiftTransformParams`. for (const auto& [gvar, func] : mod->functions) { - if (func->IsInstance()) { + if (func->IsInstance()) { auto opt_num_input = func->GetAttr(tvm::relax::attr::kNumInput); if (opt_num_input && !ends_with(gvar->name_hint, "transform_params")) { target_functions.emplace_back(gvar, func.as_or_throw()); diff --git a/src/relax/transform/lower_alloc_tensor.cc b/src/relax/transform/lower_alloc_tensor.cc index 11726005ce49..83292153bdfc 100644 --- a/src/relax/transform/lower_alloc_tensor.cc +++ b/src/relax/transform/lower_alloc_tensor.cc @@ -38,7 +38,7 @@ class Mutator : public ExprMutator { explicit Mutator(IRModule mod) : ctx_mod_(mod) {} using ExprMutator::VisitExpr_; - Expr VisitExpr_(const CallNode* op) override { + Expr VisitExpr_(const CallObj* op) override { static const Op alloc_tensor_op = Op::Get("relax.builtin.alloc_tensor"); static const Op mem_alloc_storage_op = Op::Get("relax.memory.alloc_storage"); static const Op mem_alloc_tensor_op = Op::Get("relax.memory.alloc_tensor"); @@ -55,12 +55,12 @@ class Mutator : public ExprMutator { StringImm storage_scope = op->args[3].as_or_throw(); auto shape = [&]() -> ffi::Array { - if (auto ptr = shape_arg.as()) { + if (auto ptr = shape_arg.as()) { return ptr->values; } auto ty = GetType(shape_arg); - if (auto ptr = ty.as()) { + if (auto ptr = ty.as()) { if (ptr->values) { return ptr->values.value(); } @@ -87,7 +87,7 @@ class Mutator : public ExprMutator { ShapeExpr size({nbytes}); int64_t vdevice_index = -1; - if (const auto* int_imm = op->args[2].as()) { + if (const auto* int_imm = op->args[2].as()) { vdevice_index = int_imm->value.as().value(); } ffi::Optional vdevice = GetGlobalVDevice(ctx_mod_, vdevice_index); diff --git a/src/relax/transform/merge_composite_functions.cc b/src/relax/transform/merge_composite_functions.cc index 8b8b80fead5b..bf4dc734b748 100644 --- a/src/relax/transform/merge_composite_functions.cc +++ b/src/relax/transform/merge_composite_functions.cc @@ -95,8 +95,8 @@ class CompositeGroupsBuilder : public MemoizedExprTranslator { } PostOrderVisit(func, [this](const Expr& expr) { - if (expr->IsInstance() || expr->IsInstance() || - (!expr->IsInstance() && !expr->IsInstance() && expr.as())) { + if (expr->IsInstance() || expr->IsInstance() || + (!expr->IsInstance() && !expr->IsInstance() && expr.as())) { memo_[expr] = arena_->make(); } }); @@ -112,43 +112,43 @@ class CompositeGroupsBuilder : public MemoizedExprTranslator { } Group* VisitBinding(const Binding& binding) { - if (const auto* node = binding.as()) { + if (const auto* node = binding.as()) { return VisitBinding_(node); } else { TVM_FFI_THROW(TypeError) << "Invalid type: " << binding->GetTypeKey(); } } - void VisitBindingBlock_(const BindingBlockNode* block) { + void VisitBindingBlock_(const BindingBlockObj* block) { for (Binding binding : block->bindings) { VisitBinding(binding); } } - void VisitBindingBlock_(const DataflowBlockNode* block) { + void VisitBindingBlock_(const DataflowBlockObj* block) { for (Binding binding : block->bindings) { VisitBinding(binding); } } void VisitBindingBlock(const BindingBlock& block) { - if (const auto* node = block.as()) { + if (const auto* node = block.as()) { VisitBindingBlock_(node); - } else if (const auto* node = block.as()) { + } else if (const auto* node = block.as()) { VisitBindingBlock_(node); } else { TVM_FFI_THROW(TypeError) << "Invalid type: " << block->GetTypeKey(); } } - Group* VisitExpr_(const SeqExprNode* op) { + Group* VisitExpr_(const SeqExprObj* op) { for (BindingBlock block : op->blocks) { VisitBindingBlock(block); } return VisitExpr(op->body); } - Group* VisitExpr_(const CallNode* call) { + Group* VisitExpr_(const CallObj* call) { for (const Expr& arg : call->args) { EnsureVisited(arg); } @@ -175,7 +175,7 @@ class CompositeGroupsBuilder : public MemoizedExprTranslator { return group; } - Group* VisitExpr_(const TupleNode* tuple) { + Group* VisitExpr_(const TupleObj* tuple) { Expr tuple_expr = ffi::GetRef(tuple); if (!IsFlatTensorTuple(tuple_expr)) return arena_->make(); @@ -209,7 +209,7 @@ class CompositeGroupsBuilder : public MemoizedExprTranslator { return group; } - Group* VisitExpr_(const TupleGetItemNode* tuple_get_item) { + Group* VisitExpr_(const TupleGetItemObj* tuple_get_item) { if (!IsFlatTensorTuple(tuple_get_item->tuple)) return arena_->make(); EnsureVisited(tuple_get_item->tuple); @@ -229,13 +229,13 @@ class CompositeGroupsBuilder : public MemoizedExprTranslator { private: void EnsureVisited(const Expr& expr) { - if (!expr.as() && !memo_.count(expr)) { + if (!expr.as() && !memo_.count(expr)) { VisitExpr(expr); } } ffi::Optional GetCodegenName(const Expr& callee) { - auto const* gvar = callee.as(); + auto const* gvar = callee.as(); if (!gvar) { return std::nullopt; } @@ -259,10 +259,10 @@ class CompositeGroupsBuilder : public MemoizedExprTranslator { bool CanAbsorbTupleNodes(Group* group) { return GetCodegenName(group->FindRoot()).has_value(); } bool IsFlatTensorTuple(const Expr& expr) { - const auto* tuple_type = GetType(expr).as(); + const auto* tuple_type = GetType(expr).as(); if (tuple_type == nullptr || tuple_type->fields.empty()) return false; return std::all_of(tuple_type->fields.begin(), tuple_type->fields.end(), - [](const Type& field) { return field->IsInstance(); }); + [](const Type& field) { return field->IsInstance(); }); } bool HasOnlyTupleGetItemUsers(const Expr& tuple_expr) { @@ -278,7 +278,7 @@ class CompositeGroupsBuilder : public MemoizedExprTranslator { bool ComputeHasOnlyTupleGetItemUsers(const Expr& tuple_expr) { ffi::Optional tuple_var; - if (const auto* var = tuple_expr.as()) { + if (const auto* var = tuple_expr.as()) { tuple_var = ffi::GetRef(var); } else if (auto it = value_to_bound_vars_.find(tuple_expr.get()); it != value_to_bound_vars_.end() && it->second.size() == 1) { @@ -304,7 +304,7 @@ class CompositeGroupsBuilder : public MemoizedExprTranslator { // Walk toward the original tuple value when current is itself an alias. if (auto binding_it = var_usage_.bound_values.find(current); binding_it != var_usage_.bound_values.end()) { - if (const auto* source = (*binding_it).second.as()) { + if (const auto* source = (*binding_it).second.as()) { pending.push_back(ffi::GetRef(source)); } } @@ -315,10 +315,10 @@ class CompositeGroupsBuilder : public MemoizedExprTranslator { auto binding_it = var_usage_.bound_values.find(user); if (binding_it == var_usage_.bound_values.end()) return false; const Expr& bound_value = (*binding_it).second; - if (const auto* tuple_get_item = bound_value.as()) { + if (const auto* tuple_get_item = bound_value.as()) { if (!tuple_get_item->tuple.same_as(current)) return false; has_tuple_get_item = true; - } else if (const auto* source = bound_value.as()) { + } else if (const auto* source = bound_value.as()) { if (!ffi::GetRef(source).same_as(current)) return false; pending.push_back(user); } else { @@ -329,7 +329,7 @@ class CompositeGroupsBuilder : public MemoizedExprTranslator { return has_tuple_get_item; } - Group* CreateNewGroup(const CallNode* call) { + Group* CreateNewGroup(const CallObj* call) { Group* group = arena_->make(); if (ffi::Optional codegen_name = GetCodegenName(call->op)) { group->attrs.Set(tvm::relax::attr::kCodegen.name, codegen_name.value()); @@ -366,7 +366,7 @@ class CompositeGroupsBuilder : public MemoizedExprTranslator { std::unordered_set dependencies; for (const auto& arg : args) { - if (arg.as()) continue; + if (arg.as()) continue; for (auto dep : group_deps_[memo_[arg]->FindRoot()]) { dependencies.insert(dep); } @@ -379,8 +379,8 @@ class CompositeGroupsBuilder : public MemoizedExprTranslator { Group* group_root = group->FindRoot(); std::function visit_expr = [&](Expr expr) { - if (expr.as()) return; - if (auto tuple = expr.as()) { + if (expr.as()) return; + if (auto tuple = expr.as()) { for (const auto& field : tuple->fields) { visit_expr(field); } @@ -410,7 +410,7 @@ class CompositeGroupsBuilder : public MemoizedExprTranslator { } } - std::vector GetGroupsToMerge(const CallNode* call) { + std::vector GetGroupsToMerge(const CallObj* call) { ffi::Optional codegen_name = GetCodegenName(call->op); if (!codegen_name.has_value()) { return {}; @@ -420,7 +420,7 @@ class CompositeGroupsBuilder : public MemoizedExprTranslator { std::unordered_set parent_dependencies = GetParentGroupDependencies(call->args); for (const auto& arg : call->args) { - if (arg.as()) continue; + if (arg.as()) continue; auto arg_group = memo_[arg]; ffi::Optional arg_codegen_name = GetCodegenName(arg_group); if (arg_codegen_name == codegen_name && !parent_dependencies.count(arg_group->FindRoot())) { @@ -461,8 +461,8 @@ class CompositeInliner : public ExprMutator { return new_func; } - Expr VisitExpr_(const CallNode* call) { - if (call->op->IsInstance()) { + Expr VisitExpr_(const CallObj* call) { + if (call->op->IsInstance()) { auto gvar = call->op.as_or_throw(); auto func = mod_->Lookup(gvar).as_or_throw(); if (func->GetAttr(tvm::relax::attr::kComposite)) { @@ -503,8 +503,8 @@ class CompositeFunctionAnnotator : public ExprMutator { return builder_->GetContextIRModule(); } - Expr VisitExpr_(const CallNode* call) { - if (call->op->IsInstance()) { + Expr VisitExpr_(const CallObj* call) { + if (call->op->IsInstance()) { GlobalVar cur_var = call->op.as_or_throw(); auto func = mod_->Lookup(cur_var).as_or_throw(); if (auto codegen_name = func->GetAttr(tvm::relax::attr::kCodegen)) { diff --git a/src/relax/transform/meta_schedule.cc b/src/relax/transform/meta_schedule.cc index 13001c922814..715c6a12ebc3 100644 --- a/src/relax/transform/meta_schedule.cc +++ b/src/relax/transform/meta_schedule.cc @@ -104,7 +104,7 @@ Pass MetaScheduleApplyDatabase(ffi::Optional work_dir, bool enable_ for (const auto& iter : mod->functions) { GlobalVar gv = iter.first; BaseFunc base_func = iter.second; - if (const auto* function_node = base_func.as()) { + if (const auto* function_node = base_func.as()) { tirx::Function function = ffi::GetRef(function_node); IRModule tir_mod = (*normalize_mod_func_)(function).cast(); @@ -126,7 +126,7 @@ Pass MetaScheduleApplyDatabase(ffi::Optional work_dir, bool enable_ IRModule new_mod = sch->mod(); TVM_FFI_ICHECK_EQ(new_mod->functions.size(), 1); BaseFunc new_base_func = (*new_mod->functions.begin()).second; - TVM_FFI_ICHECK(new_base_func->IsInstance()); + TVM_FFI_ICHECK(new_base_func->IsInstance()); tirx::Function tuned_function = new_base_func.as_or_throw(); // maintain the original attributes tirx::Function new_function = tirx::Function(/*params=*/tuned_function->params, diff --git a/src/relax/transform/normalize.cc b/src/relax/transform/normalize.cc index 21e74d5d3481..d5dc81cd73ac 100644 --- a/src/relax/transform/normalize.cc +++ b/src/relax/transform/normalize.cc @@ -44,7 +44,7 @@ class NormalizeMutator : public ExprMutatorBase { return builder_->Normalize(ExprMutatorBase::VisitExpr(expr)); } - Expr VisitExpr_(const FunctionNode* op) final { + Expr VisitExpr_(const FunctionObj* op) final { Expr body = this->VisitWithNewScope(op->body, op->params); if (body.same_as(op->body)) { @@ -54,7 +54,7 @@ class NormalizeMutator : public ExprMutatorBase { } } - Expr VisitExpr_(const IfExprNode* op) final { + Expr VisitExpr_(const IfExprObj* op) final { Expr guard = this->VisitExpr(op->cond); Expr true_b = this->VisitWithNewScope(op->true_branch); Expr false_b = this->VisitWithNewScope(op->false_branch); @@ -82,7 +82,7 @@ class NormalizeMutator : public ExprMutatorBase { return ret; } - Expr VisitExpr_(const SeqExprNode* op) final { + Expr VisitExpr_(const SeqExprObj* op) final { bool all_blocks_unchanged = true; ffi::Array blocks; for (auto block : op->blocks) { @@ -110,9 +110,9 @@ class NormalizeMutator : public ExprMutatorBase { BindingBlock VisitBindingBlock(const BindingBlock& block) final { BindingBlock ret(ffi::UnsafeInit{}); - if (const auto* node = block.as()) { + if (const auto* node = block.as()) { ret = VisitBindingBlock_(node); - } else if (const auto* node = block.as()) { + } else if (const auto* node = block.as()) { ret = VisitBindingBlock_(node); } else { TVM_FFI_THROW(TypeError) << "Invalid type: " << block->GetTypeKey(); @@ -120,7 +120,7 @@ class NormalizeMutator : public ExprMutatorBase { return ret; } - BindingBlock VisitBindingBlock_(const BindingBlockNode* block) { + BindingBlock VisitBindingBlock_(const BindingBlockObj* block) { builder_->BeginBindingBlock(); for (Binding binding : block->bindings) { this->VisitBinding(binding); @@ -128,7 +128,7 @@ class NormalizeMutator : public ExprMutatorBase { return builder_->EndBlock(); } - BindingBlock VisitBindingBlock_(const DataflowBlockNode* block) { + BindingBlock VisitBindingBlock_(const DataflowBlockObj* block) { builder_->BeginDataflowBlock(); for (Binding binding : block->bindings) { this->VisitBinding(binding); @@ -137,16 +137,16 @@ class NormalizeMutator : public ExprMutatorBase { } void VisitBinding(const Binding& binding) { - if (const auto* node = binding.as()) { + if (const auto* node = binding.as()) { VisitBinding_(node); - } else if (const auto* node = binding.as()) { + } else if (const auto* node = binding.as()) { VisitBinding_(node); } else { TVM_FFI_THROW(TypeError) << "Invalid type: " << binding->GetTypeKey(); } } - void VisitBinding_(const VarBindingNode* binding) { + void VisitBinding_(const VarBindingObj* binding) { Expr new_value = this->VisitExpr(binding->value); if (binding->var->ty.as().has_value()) { UpdateType(binding->var, GetType(new_value)); @@ -159,7 +159,7 @@ class NormalizeMutator : public ExprMutatorBase { } } - void VisitBinding_(const MatchCastNode* binding) { + void VisitBinding_(const MatchCastObj* binding) { Expr new_value = this->VisitExpr(binding->value); if (new_value.same_as(binding->value)) { @@ -202,7 +202,7 @@ class GlobalVarNormalizer : private ExprMutator { // Step 3. Substitute global vars in functions for (auto [gvar, func] : module_->functions) { - if (!func->IsInstance()) { + if (!func->IsInstance()) { continue; } auto new_func = this->VisitExpr(func).as_or_throw(); @@ -263,7 +263,7 @@ class GlobalVarNormalizer : private ExprMutator { } } - Expr VisitExpr_(const GlobalVarNode* op) final { + Expr VisitExpr_(const GlobalVarObj* op) final { TVM_FFI_ICHECK(gvar_map_.count(ffi::GetRef(op))); return gvar_map_[ffi::GetRef(op)]; } diff --git a/src/relax/transform/realize_vdevice.cc b/src/relax/transform/realize_vdevice.cc index 92af0558b456..9e1549d5ab76 100644 --- a/src/relax/transform/realize_vdevice.cc +++ b/src/relax/transform/realize_vdevice.cc @@ -100,14 +100,14 @@ class DeviceHintCollector : ExprVisitor { private: explicit DeviceHintCollector(VDeviceLookup vdevice_lookup) : vdevice_lookup_(vdevice_lookup) {} - void VisitExpr_(const FunctionNode* func) override { + void VisitExpr_(const FunctionObj* func) override { ExprVisitor::VisitExpr_(func); std::function check_ret_ty = [this, &check_ret_ty](Expr expr, Type ty) { // If the function is annotated as returning a tensor on a // specific device, then that annotation may be propagated into // the returned variable. - if (auto tensor_info = ty.as(); + if (auto tensor_info = ty.as(); tensor_info && tensor_info->vdevice.has_value()) { if (auto opt_var = expr.as()) { auto var = opt_var.value(); @@ -121,7 +121,7 @@ class DeviceHintCollector : ExprVisitor { // where some elements of the tuple are tensors that exist on a // specific device, then those annotations may be propagated // into the corresponding tensor annotations. - if (auto tuple_info = ty.as()) { + if (auto tuple_info = ty.as()) { // The returned tuple is not necessarily an in-line tuple. In // order to find the variables that are bound to the // individual tuple elements, we may need to unwrap the @@ -138,7 +138,7 @@ class DeviceHintCollector : ExprVisitor { // example, the function may return one of its arguments as an // output, or may return the result of a `relax::Call` that // produces a tuple of outputs. - if (auto tuple = expr.as()) { + if (auto tuple = expr.as()) { TVM_FFI_CHECK_EQ(tuple_info->fields.size(), tuple->fields.size(), ValueError) << "Function returns a tuple with " << tuple->fields.size() << " elements, " << "but is annotated as returning a tuple with " << tuple_info->fields.size() @@ -154,7 +154,7 @@ class DeviceHintCollector : ExprVisitor { } void VisitVarDef(const Var& var) override { - if (auto tinfo = var->ty.as(); tinfo && tinfo->vdevice.has_value()) { + if (auto tinfo = var->ty.as(); tinfo && tinfo->vdevice.has_value()) { known_vdevice_.Set(var, tinfo->vdevice.value()); } ExprVisitor::VisitVarDef(var); @@ -165,7 +165,7 @@ class DeviceHintCollector : ExprVisitor { binding_lookup_.Set(binding->var, GetBoundValue(binding)); } - void VisitBinding_(const VarBindingNode* binding, const CallNode* call) override { + void VisitBinding_(const VarBindingObj* binding, const CallObj* call) override { ExprVisitor::VisitBinding_(binding, call); if (call->op.same_as(hint_on_device_op_)) { auto vdevice = vdevice_lookup_(call->attrs); @@ -241,7 +241,7 @@ class VDeviceSetCollector : ExprVisitor { current_binding_ = cached; } - void VisitExpr_(const CallNode* call) override { + void VisitExpr_(const CallObj* call) override { if (!call->op.same_as(to_vdevice_op_) && !call->op.same_as(hint_on_device_op_)) { ExprVisitor::VisitExpr_(call); } @@ -249,10 +249,10 @@ class VDeviceSetCollector : ExprVisitor { void VisitExprDepTypeField(const Type&) final {} - void VisitExpr_(const VarNode* op) override { - if (current_binding_ && !GetType(current_binding_.value()).as()) { + void VisitExpr_(const VarObj* op) override { + if (current_binding_ && !GetType(current_binding_.value()).as()) { auto var = ffi::GetRef(op); - if (GetType(var).as()) { + if (GetType(var).as()) { return; } var_to_co_located_vars_[current_binding_.value()].push_back(var); @@ -356,7 +356,7 @@ class VDeviceTypeUpdater : ExprMutator { Var VisitVarDef(const Var& old_var) override { auto var = ExprMutator::VisitVarDef(old_var); - if (auto tinfo = var->ty.as()) { + if (auto tinfo = var->ty.as()) { if (auto opt = vdevice_map_.Get(old_var)) { auto vdevice = opt.value(); TensorType new_ty = [&]() { @@ -367,7 +367,7 @@ class VDeviceTypeUpdater : ExprMutator { } }(); - if (var->IsInstance()) { + if (var->IsInstance()) { var = DataflowVar(var->name, new_ty, var->loc); } else { var = Var(var->name, new_ty, var->loc); @@ -380,7 +380,7 @@ class VDeviceTypeUpdater : ExprMutator { using ExprMutator::VisitExpr_; - Expr VisitExpr_(const CallNode* op) override { + Expr VisitExpr_(const CallObj* op) override { auto call = ExprMutator::VisitExpr_(op).as_or_throw(); if (!call->op.same_as(hint_on_device_op_)) { diff --git a/src/relax/transform/remove_purity_checking.cc b/src/relax/transform/remove_purity_checking.cc index f86a368c33d4..766a6b2e7058 100644 --- a/src/relax/transform/remove_purity_checking.cc +++ b/src/relax/transform/remove_purity_checking.cc @@ -48,7 +48,7 @@ class PurityRemover : public ExprMutator { return ret; } - Expr VisitExpr_(const CallNode* call) override { + Expr VisitExpr_(const CallObj* call) override { if (call->op.same_as(call_pure_packed_op_)) { auto ret = Call(Type::Missing(), call->args[0], ffi::Array(call->args.begin() + 1, call->args.end()), call->attrs, @@ -69,7 +69,7 @@ class PurityRemover : public ExprMutator { return ExprMutator::VisitExpr_(call); } - Expr VisitExpr_(const FunctionNode* func) override { + Expr VisitExpr_(const FunctionObj* func) override { // handling inner functions: we will remove purity annotations from them too return RemovePurity(ffi::GetRef(func)); } diff --git a/src/relax/transform/remove_unused_outputs.cc b/src/relax/transform/remove_unused_outputs.cc index 660e87cbd1d4..7a6a72456bd4 100644 --- a/src/relax/transform/remove_unused_outputs.cc +++ b/src/relax/transform/remove_unused_outputs.cc @@ -51,8 +51,8 @@ class PartialTupleUsageCollector : ExprVisitor { bool is_exposed = base_func->attrs.GetAttr(tvm::attr::kGlobalSymbol).has_value(); if (!is_exposed) { - if (auto relax_func = base_func.as()) { - if (auto out_tuple = relax_func->ret_ty.as()) { + if (auto relax_func = base_func.as()) { + if (auto out_tuple = relax_func->ret_ty.as()) { num_outputs[gvar] = out_tuple->fields.size(); } } @@ -96,7 +96,7 @@ class PartialTupleUsageCollector : ExprVisitor { known_bindings_.Set(binding->var, GetBoundValue(binding)); } - void VisitExpr_(const TupleGetItemNode* op) override { + void VisitExpr_(const TupleGetItemObj* op) override { if (auto* usage_mask_ptr = GetCalleeUsageMask(op->tuple)) { auto& used_indices = *usage_mask_ptr; @@ -113,7 +113,7 @@ class PartialTupleUsageCollector : ExprVisitor { } } - void VisitExpr_(const VarNode* op) override { + void VisitExpr_(const VarObj* op) override { if (auto* usage_mask_ptr = GetCalleeUsageMask(ffi::GetRef(op))) { auto& usage_mask = *usage_mask_ptr; for (size_t i = 0; i < usage_mask.size(); i++) { @@ -123,12 +123,12 @@ class PartialTupleUsageCollector : ExprVisitor { } std::vector* GetCalleeUsageMask(Expr expr) { - if (!expr->ty.as()) { + if (!expr->ty.as()) { return nullptr; } expr = UnwrapBindings(expr); - if (auto call = expr.as()) { + if (auto call = expr.as()) { if (auto callee = call->op.as()) { if (auto it = output_usage_mask_.find(callee.value()); it != output_usage_mask_.end()) { return &it->second; @@ -160,9 +160,9 @@ class PartialTupleUsageCollector : ExprVisitor { }; Function UpdateCallee(Function func, const std::vector& usage_mask) { - auto old_func_ty = func->ty.as(); + auto old_func_ty = func->ty.as(); - auto old_ret_ty = func->ret_ty.as(); + auto old_ret_ty = func->ret_ty.as(); TVM_FFI_ICHECK(old_ret_ty) << "All functions returning non-tuple outputs " << "should have been pruned already by PartialTupleUsageCollector"; @@ -202,7 +202,7 @@ class CallSiteMutator : public ExprMutator { using ExprMutator::VisitExpr_; - Expr VisitExpr_(const CallNode* op) override { + Expr VisitExpr_(const CallObj* op) override { auto node = ExprMutator::VisitExpr_(op).as_or_throw(); if (auto gvar = node->op.as()) { @@ -250,7 +250,7 @@ Pass RemoveUnusedOutputs() { << "Updater should be applied to " << old_gvar << ", but was applied to " << call->op; - auto old_call_ty = call->ty.as(); + auto old_call_ty = call->ty.as(); TVM_FFI_CHECK(old_call_ty, InternalError) << "Updater should be applied to Call producing an output tuple, " << "but " << call << " has type " << call->ty; diff --git a/src/relax/transform/remove_unused_parameters.cc b/src/relax/transform/remove_unused_parameters.cc index b105e417800e..c00de8e98d26 100644 --- a/src/relax/transform/remove_unused_parameters.cc +++ b/src/relax/transform/remove_unused_parameters.cc @@ -157,7 +157,7 @@ class CallSiteMutator : public ExprMutator { using ExprMutator::VisitExpr_; - Expr VisitExpr_(const FunctionNode* op) override { + Expr VisitExpr_(const FunctionObj* op) override { auto node = ExprMutator::VisitExpr_(op); // If a function was modified, that means it called into a private @@ -170,7 +170,7 @@ class CallSiteMutator : public ExprMutator { return node; } - Expr VisitExpr_(const CallNode* op) override { + Expr VisitExpr_(const CallObj* op) override { auto node = ExprMutator::VisitExpr_(op).as_or_throw(); if (auto gvar = node->op.as()) { diff --git a/src/relax/transform/reorder_permute_dims_after_concat.cc b/src/relax/transform/reorder_permute_dims_after_concat.cc index efbdb1daf05f..e8f6f6894e39 100644 --- a/src/relax/transform/reorder_permute_dims_after_concat.cc +++ b/src/relax/transform/reorder_permute_dims_after_concat.cc @@ -83,7 +83,7 @@ std::tuple)>> } auto get_permute_dims_optional_axes = [](const Expr& expr) -> ffi::Optional> { - auto call = expr.as(); + auto call = expr.as(); TVM_FFI_ICHECK(call); auto attrs = call->attrs.as(); TVM_FFI_ICHECK(attrs); @@ -99,7 +99,7 @@ std::tuple)>> auto call = expr.as_or_throw(); ffi::Array permutation; - auto arg_ty = call->args[0]->ty.as(); + auto arg_ty = call->args[0]->ty.as(); TVM_FFI_ICHECK(arg_ty) << "Expected permute_dims to have a single tensor argument, " << "but argument " << call->args[0] << " has type " << call->args[0]->ty; diff --git a/src/relax/transform/reorder_take_after_matmul.cc b/src/relax/transform/reorder_take_after_matmul.cc index 0908da4db932..447d0e24a3b8 100644 --- a/src/relax/transform/reorder_take_after_matmul.cc +++ b/src/relax/transform/reorder_take_after_matmul.cc @@ -55,7 +55,7 @@ std::tuple)>> auto weights = matches[pat_weights]; auto indices = matches[pat_indices]; - const auto* take_call = matches[pat_rhs].as(); + const auto* take_call = matches[pat_rhs].as(); TVM_FFI_CHECK(take_call, InternalError) << "Match of relax.take operator should produce Call, " << "but instead produces " << matches[pat_rhs] << " with type " << matches[pat_rhs]->GetTypeKey(); @@ -64,16 +64,16 @@ std::tuple)>> << "Attributes for relax.take operator should be TakeAttrs, " << "but were instead " << take_call->attrs << " with type " << take_call->GetTypeKey(); - const auto* lhs_ty = lhs->ty.as(); + const auto* lhs_ty = lhs->ty.as(); if (!lhs_ty) return expr; - const auto* weights_ty = weights->ty.as(); + const auto* weights_ty = weights->ty.as(); if (!weights_ty) return expr; - const auto* indices_ty = indices->ty.as(); + const auto* indices_ty = indices->ty.as(); if (!indices_ty) return expr; - const auto* matmul_ty = expr->ty.as(); + const auto* matmul_ty = expr->ty.as(); if (!matmul_ty) return expr; if (!attrs->axis.has_value()) return expr; diff --git a/src/relax/transform/rewrite_cuda_graph.cc b/src/relax/transform/rewrite_cuda_graph.cc index a32ba207fa59..edf87047cd11 100644 --- a/src/relax/transform/rewrite_cuda_graph.cc +++ b/src/relax/transform/rewrite_cuda_graph.cc @@ -82,17 +82,17 @@ struct LiftedFunctionRewritePlan { // Whether the lifted function is for allocation or capturing bool is_alloc; // The binding var before which the lifted function should be invoked - const VarNode* launch_point; + const VarObj* launch_point; // Variable remappings between the original function and the lifted function // The bindings in the original function that are lifted - std::vector lifted_bindings; + std::vector lifted_bindings; // The corresponding binding vars in the original function of the outputs of the lifted function // to the index of the element in the output tuple of the lifted function. - std::unordered_map outputs; + std::unordered_map outputs; // The corresponding binding vars in the original function of the inputs of the lifted function - std::vector inputs; + std::vector inputs; // The tirx vars in the original function that are propagated to the lifted function ffi::Optional propogated_tir_vars = std::nullopt; }; @@ -104,13 +104,13 @@ class FuncBuilder : public ExprMutator { * \brief Add a binding to the new function * \param binding The binding to add */ - void AddBinding(const VarBindingNode* binding) { bindings_.push_back(binding); } + void AddBinding(const VarBindingObj* binding) { bindings_.push_back(binding); } /*! * \brief Mark a variable as the input of the new function. * \param var The variable to mark as input */ - void MarkInput(const VarNode* var) { inputs_.push_back(var); } + void MarkInput(const VarObj* var) { inputs_.push_back(var); } /*! * \brief Mark a TIR variable as the ShapeExpr input of the new function. * \param var The variable to mark as input @@ -121,7 +121,7 @@ class FuncBuilder : public ExprMutator { * existing binding in the new function. * \param var The variable to mark as output */ - void MarkOutput(const VarNode* var) { outputs_.push_back(var); } + void MarkOutput(const VarObj* var) { outputs_.push_back(var); } /*! \brief Get the number of bindings in the new function */ auto size() const { return bindings_.size(); } @@ -168,23 +168,23 @@ class FuncBuilder : public ExprMutator { return func; } - support::OrderedSet inputs_; - support::OrderedSet outputs_; + support::OrderedSet inputs_; + support::OrderedSet outputs_; support::OrderedSet shape_expr_inputs_; - std::vector bindings_; + std::vector bindings_; }; // Collect the storage objects that are used as the function output class OutputStorageCollector : public ExprVisitor { public: - static std::unordered_set Collect(const Function& func) { + static std::unordered_set Collect(const Function& func) { OutputStorageCollector collector; collector.VisitExpr(func); return std::move(collector.output_storages_); } private: - void VisitExpr_(const SeqExprNode* seq_expr) final { + void VisitExpr_(const SeqExprObj* seq_expr) final { auto output_vars = FreeVars(seq_expr->body); for (const auto& var : output_vars) { output_vars_.insert(var.get()); @@ -195,36 +195,36 @@ class OutputStorageCollector : public ExprVisitor { } } - void VisitBinding_(const VarBindingNode* binding, const CallNode* call) final { + void VisitBinding_(const VarBindingObj* binding, const CallObj* call) final { static const auto mem_alloc_tensor_op = Op::Get("relax.memory.alloc_tensor"); if (output_vars_.count(binding->var.get()) && call->op.same_as(mem_alloc_tensor_op)) { - output_storages_.insert(call->args[0].as()); + output_storages_.insert(call->args[0].as()); } } - void VisitBindingBlock_(const BindingBlockNode* binding_block) override { + void VisitBindingBlock_(const BindingBlockObj* binding_block) override { // Visit the bindings in reverse order for (auto it = binding_block->bindings.rbegin(); it != binding_block->bindings.rend(); ++it) { VisitBinding(*it); } } - void VisitBinding_(const VarBindingNode* binding, const VarNode* var) final { + void VisitBinding_(const VarBindingObj* binding, const VarObj* var) final { if (output_vars_.count(binding->var.get())) { output_vars_.insert(var); } } - void VisitBinding_(const VarBindingNode* binding, const TupleNode* tuple) final { + void VisitBinding_(const VarBindingObj* binding, const TupleObj* tuple) final { if (output_vars_.count(binding->var.get())) { for (const auto& field : tuple->fields) { - output_vars_.insert(field.as()); + output_vars_.insert(field.as()); } } } - std::unordered_set output_storages_; - std::unordered_set output_vars_; + std::unordered_set output_storages_; + std::unordered_set output_vars_; }; /*! @@ -237,7 +237,7 @@ class CUDAGraphRewritePlanner : public ExprVisitor { std::pair, std::vector> Plan() { for (const auto& pair : mod_->functions) { - if (pair.second->IsInstance()) { + if (pair.second->IsInstance()) { // If a function has the num_input attribute, the last func->params.size() - num_inputs // inputs are assumed to be fixed and thus they can be captured into a cuda graph. // The symbolic variables in the type of the fixed inputs (weights) are also allowed @@ -331,7 +331,7 @@ class CUDAGraphRewritePlanner : public ExprVisitor { current_block_scope_.capture_builder = nullptr; } - void VisitExpr_(const FunctionNode* func) final { + void VisitExpr_(const FunctionObj* func) final { current_function_scope_.alloc_storage_builder = arena_->make(); ExprVisitor::VisitExpr_(func); if (current_function_scope_.alloc_storage_builder->outputs_.size()) { @@ -340,7 +340,7 @@ class CUDAGraphRewritePlanner : public ExprVisitor { current_function_scope_.alloc_storage_builder = nullptr; } - void VisitBindingBlock_(const BindingBlockNode* binding_block) final { + void VisitBindingBlock_(const BindingBlockObj* binding_block) final { BindingSBlockScope new_scope; std::swap(new_scope, current_block_scope_); for (const auto& binding : binding_block->bindings) { @@ -350,7 +350,7 @@ class CUDAGraphRewritePlanner : public ExprVisitor { std::swap(new_scope, current_block_scope_); } - void VisitBinding_(const VarBindingNode* binding, const CallNode* call) final { + void VisitBinding_(const VarBindingObj* binding, const CallObj* call) final { static const auto mem_alloc_storage_op = Op::Get("relax.memory.alloc_storage"); static const auto builtin_alloc_tensor_op = Op::Get("relax.builtin.alloc_tensor"); static const auto call_builtin_with_ctx_op = Op::Get("relax.call_builtin_with_ctx"); @@ -364,14 +364,14 @@ class CUDAGraphRewritePlanner : public ExprVisitor { return; } - const auto* call_gv = call->op.as(); + const auto* call_gv = call->op.as(); bool call_function = - call_gv ? mod_->Lookup(ffi::GetRef(call_gv))->IsInstance() + call_gv ? mod_->Lookup(ffi::GetRef(call_gv))->IsInstance() : false; // Check whether the call can be lifted to the capture function. It requires all the arguments // to be static and the call to be a kernel launch or a pure operation (e.g. memory view). - std::vector args; + std::vector args; std::vector tir_vars; bool is_all_static = [&]() { if (!IsStatic(call->args, &args, &tir_vars)) { @@ -381,7 +381,7 @@ class CUDAGraphRewritePlanner : public ExprVisitor { // calls to other Relax functions are not allowed return false; } - if (const auto* extern_func = call->op.as(); + if (const auto* extern_func = call->op.as(); extern_func != nullptr && extern_func->global_symbol.starts_with("vm.builtin")) { return false; } @@ -395,10 +395,10 @@ class CUDAGraphRewritePlanner : public ExprVisitor { if (call_function) { return true; } - if (call->op.as()) { + if (call->op.as()) { return true; } - if (const auto* op = call->op.as()) { + if (const auto* op = call->op.as()) { return !op->name.starts_with("relax.memory") && !op->name.starts_with("relax.builtin") && op->name != "relax.reshape" && !ffi::GetRef(op).same_as(null_value_op) && !ffi::GetRef(op).same_as(call_builtin_with_ctx_op); @@ -417,12 +417,12 @@ class CUDAGraphRewritePlanner : public ExprVisitor { MarkAsFuncOutput(args); } - void MarkAsFuncInput(const std::vector& vars, + void MarkAsFuncInput(const std::vector& vars, const std::vector& tir_vars = {}) { if (current_block_scope_.capture_builder == nullptr) { return; } - for (const VarNode* var : vars) { + for (const VarObj* var : vars) { auto it = binding_to_region_.find(var); if (it == binding_to_region_.end() || it->second != current_block_scope_.capture_builder) { current_block_scope_.capture_builder->MarkInput(var); @@ -433,8 +433,8 @@ class CUDAGraphRewritePlanner : public ExprVisitor { } } - void MarkAsFuncOutput(const std::vector& vars) { - for (const VarNode* var : vars) { + void MarkAsFuncOutput(const std::vector& vars) { + for (const VarObj* var : vars) { if (auto it = binding_to_region_.find(var); it != binding_to_region_.end() && it->second != current_block_scope_.capture_builder) { it->second->MarkOutput(var); @@ -442,7 +442,7 @@ class CUDAGraphRewritePlanner : public ExprVisitor { } } - void VisitBinding_(const VarBindingNode* binding, const VarNode* var) final { + void VisitBinding_(const VarBindingObj* binding, const VarObj* var) final { if (IsStatic(ffi::GetRef(var))) { AddStaticBinding(binding, false); MarkAsFuncInput({var}); @@ -452,14 +452,14 @@ class CUDAGraphRewritePlanner : public ExprVisitor { MarkAsFuncOutput({var}); } - void VisitBinding_(const VarBindingNode* binding, const GenericConstNode* constant) final { + void VisitBinding_(const VarBindingObj* binding, const GenericConstObj* constant) final { AddStaticBinding(binding, false); } - void VisitBinding_(const MatchCastNode* binding) final { + void VisitBinding_(const MatchCastObj* binding) final { // A match_cast stays in the original function even when its value was lifted into the // current capture region, so the region has to return that value. - if (const auto* var = binding->value.as()) { + if (const auto* var = binding->value.as()) { if (auto it = binding_to_region_.find(var); it != binding_to_region_.end()) { it->second->MarkOutput(var); } @@ -467,8 +467,8 @@ class CUDAGraphRewritePlanner : public ExprVisitor { ExprVisitor::VisitBinding_(binding); } - void VisitBinding_(const VarBindingNode* binding, const TupleNode* tuple) final { - std::vector args; + void VisitBinding_(const VarBindingObj* binding, const TupleObj* tuple) final { + std::vector args; std::vector tir_vars; if (IsStatic(tuple->fields, &args, &tir_vars)) { AddStaticBinding(binding, false); @@ -479,8 +479,8 @@ class CUDAGraphRewritePlanner : public ExprVisitor { MarkAsFuncOutput(args); } - void VisitBinding_(const VarBindingNode* binding, const TupleGetItemNode* tuple_get_item) final { - const VarNode* tuple = tuple_get_item->tuple.as(); + void VisitBinding_(const VarBindingObj* binding, const TupleGetItemObj* tuple_get_item) final { + const VarObj* tuple = tuple_get_item->tuple.as(); TVM_FFI_ICHECK(tuple); if (IsStatic(tuple_get_item->tuple)) { AddStaticBinding(binding, false); @@ -492,7 +492,7 @@ class CUDAGraphRewritePlanner : public ExprVisitor { } bool IsStatic(const PrimExpr& expr, - [[maybe_unused]] std::vector* vars_collector = nullptr, + [[maybe_unused]] std::vector* vars_collector = nullptr, std::vector* tir_vars_collector = nullptr) { bool is_static = true; std::unordered_set visited; @@ -514,33 +514,33 @@ class CUDAGraphRewritePlanner : public ExprVisitor { return is_static; } - bool IsStatic(const Expr& expr, std::vector* vars_collector = nullptr, + bool IsStatic(const Expr& expr, std::vector* vars_collector = nullptr, std::vector* tir_vars_collector = nullptr) { - if (expr->IsInstance() || expr->IsInstance() || - expr->IsInstance() || expr->IsInstance()) { + if (expr->IsInstance() || expr->IsInstance() || + expr->IsInstance() || expr->IsInstance()) { return true; } - if (const auto* var = expr.as(); - var && (expr.as() || static_vars_.count(var))) { + if (const auto* var = expr.as(); + var && (expr.as() || static_vars_.count(var))) { return IsStaticRuntimeVar(var, vars_collector, tir_vars_collector); } if (auto prim_value = expr.as()) { return IsStatic(prim_value.value(), vars_collector, tir_vars_collector); } - if (const auto* var = expr.as()) { + if (const auto* var = expr.as()) { return IsStaticRuntimeVar(var, vars_collector, tir_vars_collector); } - if (const auto* shape = expr.as()) { + if (const auto* shape = expr.as()) { return IsStatic(shape->values, vars_collector, tir_vars_collector); } - if (const auto* tuple = expr.as()) { + if (const auto* tuple = expr.as()) { return IsStatic(tuple->fields, vars_collector, tir_vars_collector); } return false; } - bool IsStaticRuntimeVar(const VarNode* var, std::vector* vars_collector, + bool IsStaticRuntimeVar(const VarObj* var, std::vector* vars_collector, std::vector* tir_vars_collector) { if (vars_collector != nullptr) { vars_collector->push_back(var); @@ -550,7 +550,7 @@ class CUDAGraphRewritePlanner : public ExprVisitor { } template - bool IsStatic(const ffi::Array& exprs, std::vector* vars_collector = nullptr, + bool IsStatic(const ffi::Array& exprs, std::vector* vars_collector = nullptr, std::vector* tir_vars_collector = nullptr) { bool result = true; for (const auto& expr : exprs) { @@ -564,34 +564,34 @@ class CUDAGraphRewritePlanner : public ExprVisitor { return result; } - bool IsStatic(const Type& ty, std::vector* vars_collector = nullptr, + bool IsStatic(const Type& ty, std::vector* vars_collector = nullptr, std::vector* tir_vars_collector = nullptr) { - if (const auto* tensor_ty = ty.as()) { + if (const auto* tensor_ty = ty.as()) { if (auto shape = tensor_ty->GetShape()) { return IsStatic(shape.value(), vars_collector, tir_vars_collector); } - } else if (const auto* shape_ty = ty.as()) { + } else if (const auto* shape_ty = ty.as()) { if (shape_ty->values) { return IsStatic(shape_ty->values.value(), vars_collector, tir_vars_collector); } - } else if (const auto* tuple_ty = ty.as()) { + } else if (const auto* tuple_ty = ty.as()) { return IsStatic(tuple_ty->fields, vars_collector, tir_vars_collector); - } else if (ty.as() || ty.as()) { + } else if (ty.as() || ty.as()) { return true; } return false; } private: - bool IsStaticAllocStorage(const VarBindingNode* binding) { + bool IsStaticAllocStorage(const VarBindingObj* binding) { if (disabled_storage_vars_.count(binding->var.get())) { return false; } // Check if the allocation has constant shape - const auto* alloc_storage_call = binding->value.as(); + const auto* alloc_storage_call = binding->value.as(); auto shape = alloc_storage_call->args[0].as_or_throw(); return std::all_of(shape->values.begin(), shape->values.end(), - [](const PrimExpr& expr) { return expr.as() != nullptr; }); + [](const PrimExpr& expr) { return expr.as() != nullptr; }); } /*! @@ -600,7 +600,7 @@ class CUDAGraphRewritePlanner : public ExprVisitor { * \param binding the binding to add * \param is_alloc_storage whether the binding is call to R.memory.alloc_storage */ - void AddStaticBinding(const VarBindingNode* binding, bool is_alloc_storage) { + void AddStaticBinding(const VarBindingObj* binding, bool is_alloc_storage) { if (is_alloc_storage) { current_function_scope_.alloc_storage_builder->AddBinding(binding); binding_to_region_[binding->var.get()] = current_function_scope_.alloc_storage_builder; @@ -638,19 +638,19 @@ class CUDAGraphRewritePlanner : public ExprVisitor { // States of the current function scope FunctionScope current_function_scope_; // Variables whose buffer address is fixed - std::unordered_set static_vars_; + std::unordered_set static_vars_; // Symbolic variables that are allowed to be captured. This can come from symbolic shapes of // weights or hints in the function annotations. std::unordered_set capture_symbolic_vars_; // Binding to the FuncBuilder if the binding is lifted. This is used to update the inputs/outputs // of the lifted function when its binding is used outside. - std::unordered_map binding_to_region_; + std::unordered_map binding_to_region_; // The regions for capturing. std::vector captured_regions_; // The regions for allocation. std::vector alloc_storages_; // The binding variables that are not allowed to be captured. - std::unordered_set disabled_storage_vars_; + std::unordered_set disabled_storage_vars_; // The arena. support::Arena* arena_; }; @@ -675,7 +675,7 @@ Function MergeAllocationPlans(const std::vector& all // The size of the storage object in bytes int64_t size; // The binding of the storage allocation - const VarBindingNode* binding; + const VarBindingObj* binding; // The source rewriting plan that the storage record is from LiftedFunctionRewritePlan* src; @@ -690,7 +690,7 @@ Function MergeAllocationPlans(const std::vector& all for (int plan_id = 0; plan_id < static_cast(alloc_plans.size()); ++plan_id) { LiftedFunctionRewritePlan* plan = alloc_plans[plan_id]; TVM_FFI_ICHECK(plan->is_alloc); - for (const VarBindingNode* binding : plan->lifted_bindings) { + for (const VarBindingObj* binding : plan->lifted_bindings) { // Extract the stroage record from the Call expr. Call alloc_storage = binding->value.as_or_throw(); TVM_FFI_ICHECK(alloc_storage->op.same_as(mem_alloc_storage_op)); @@ -711,7 +711,7 @@ Function MergeAllocationPlans(const std::vector& all // storage allocations from different functions to the largest allocation that can be shared // among all the functions. // This assumes that multiple functions will not run concurrently. - std::vector merged_allocs; + std::vector merged_allocs; // Merge the storage records within each storage scope. for (auto& [storage_scope, curr_scope_records] : storage_records) { // The number of storages needed for the current storage scope, which is the maximum number of @@ -767,7 +767,7 @@ class CUDAGraphRewriter : public ExprMutator { // Collect the target functions for rewriting before any mutation. std::vector> target_functions; for (const auto& [gv, func] : builder_->GetContextIRModule()->functions) { - if (func->IsInstance()) { + if (func->IsInstance()) { target_functions.emplace_back(gv, func.as_or_throw()); } } @@ -794,7 +794,7 @@ class CUDAGraphRewriter : public ExprMutator { return builder_->GetContextIRModule(); } - void LaunchSubgraph(const VarBindingNode* op, const LiftedFunctionRewritePlan* plan) { + void LaunchSubgraph(const VarBindingObj* op, const LiftedFunctionRewritePlan* plan) { static const auto call_builtin_with_ctx_op = Op::Get("relax.call_builtin_with_ctx"); static const auto& builtin_run_or_capture = ExternFunc("vm.builtin.cuda_graph.run_or_capture"); static const auto& builtin_get_cached_alloc = @@ -853,10 +853,10 @@ class CUDAGraphRewriter : public ExprMutator { } std::transform(plan->lifted_bindings.begin(), plan->lifted_bindings.end(), std::inserter(lifted_binding_vars_, lifted_binding_vars_.end()), - [](const BindingNode* binding) { return binding->var.get(); }); + [](const BindingObj* binding) { return binding->var.get(); }); } - void VisitBinding_(const VarBindingNode* op) final { + void VisitBinding_(const VarBindingObj* op) final { if (subgraph_launches_.count(op->var.get())) { LaunchSubgraph(op, subgraph_launches_[op->var.get()]); } @@ -872,7 +872,7 @@ class CUDAGraphRewriter : public ExprMutator { ExprMutator::VisitBinding_(op); } - Expr VisitExpr_(const VarNode* op) final { + Expr VisitExpr_(const VarObj* op) final { Var var = ffi::GetRef(op); if (auto it = var_remap_.find(var); it != var_remap_.end()) { return it->second; @@ -885,15 +885,15 @@ class CUDAGraphRewriter : public ExprMutator { return ffi::GetRef(op); } - Var EmitRedef(const VarNode* var, const Expr& redef) { + Var EmitRedef(const VarObj* var, const Expr& redef) { auto new_var = builder_->Emit(redef, var->name); var_remap_.insert_or_assign(ffi::GetRef(var), new_var); return new_var; } - std::unordered_map subgraph_launches_; - std::unordered_map var_redef_; - std::unordered_set lifted_binding_vars_; + std::unordered_map subgraph_launches_; + std::unordered_map var_redef_; + std::unordered_set lifted_binding_vars_; int index_alloc_ = 0; int index_capture_ = 0; support::Arena arena_; diff --git a/src/relax/transform/rewrite_dataflow_reshape.cc b/src/relax/transform/rewrite_dataflow_reshape.cc index 8f1f0c5ee0ea..da190c132440 100644 --- a/src/relax/transform/rewrite_dataflow_reshape.cc +++ b/src/relax/transform/rewrite_dataflow_reshape.cc @@ -63,24 +63,24 @@ class DataflowReshapeRewriter : public ExprMutator { BindingBlock VisitBindingBlock(const BindingBlock& block) final { // We only rewrite the bindings inside dataflow blocks. - if (const auto* dataflow_block = block.as()) { + if (const auto* dataflow_block = block.as()) { return VisitBindingBlock_(dataflow_block); } else { return block; } } - void VisitBinding_(const VarBindingNode* binding) final { + void VisitBinding_(const VarBindingObj* binding) final { // We only rewrite the bindings that are not dataflow output (which means they are not // externally referenced) - if (!binding->var->IsInstance()) { + if (!binding->var->IsInstance()) { this->builder_->EmitNormalized(ffi::GetRef(binding)); } else { ExprMutator::VisitBinding_(binding); } } - Expr VisitExpr_(const CallNode* call) final { + Expr VisitExpr_(const CallObj* call) final { static const Op call_tir_op = Op::Get("relax.call_tir"); if (!call->op.same_as(call_tir_op)) { return ffi::GetRef(call); @@ -113,9 +113,9 @@ class DataflowReshapeRewriter : public ExprMutator { return reshape(arg, res_ty->shape.value()); } - bool IsCallingTIRReshape(const CallNode* call, Expr inp) { + bool IsCallingTIRReshape(const CallObj* call, Expr inp) { const GlobalVar& global_var = call->args[0].as_or_throw(); - const auto* func = mod_->functions.Get(global_var).value().as(); + const auto* func = mod_->functions.Get(global_var).value().as(); TVM_FFI_ICHECK_NOTNULL(func); if (!HasReshapePattern(ffi::GetRef(func))) { return false; diff --git a/src/relax/transform/run_codegen.cc b/src/relax/transform/run_codegen.cc index f42576815711..eb5d18dc41d8 100644 --- a/src/relax/transform/run_codegen.cc +++ b/src/relax/transform/run_codegen.cc @@ -65,7 +65,7 @@ class CodeGenRunner : ExprMutator { std::vector attr_entry_functions; for (const auto& [gv, func] : mod->functions) { if (func->GetLinkageType() == LinkageType::kExternal && - !func->GetAttr(tvm::relax::attr::kCodegen) && func->IsInstance()) { + !func->GetAttr(tvm::relax::attr::kCodegen) && func->IsInstance()) { attr_entry_functions.push_back(gv); } } @@ -110,9 +110,9 @@ class CodeGenRunner : ExprMutator { using ExprMutator::VisitExpr_; - Expr VisitExpr_(const CallNode* call_node) override { + Expr VisitExpr_(const CallObj* call_node) override { auto call = ExprMutator::VisitExpr_(call_node).as_or_throw(); - if (auto const* gvar_node = call_node->op.as()) { + if (auto const* gvar_node = call_node->op.as()) { const GlobalVar gvar = ffi::GetRef(gvar_node); auto create_call_dps_packed = [call_node, this](Expr extern_func, Type ret_ty) { @@ -132,7 +132,7 @@ class CodeGenRunner : ExprMutator { Function func = opt_func.value(); Expr new_func = VisitExpr(func); - if (new_func->IsInstance()) { + if (new_func->IsInstance()) { extern_funcs_.insert_or_assign(gvar_node, new_func); // Remove the global symbol and codegen attributes from the function so that it can be // removed the module. @@ -149,7 +149,7 @@ class CodeGenRunner : ExprMutator { } Type ret_ty = Type::Missing(); - if (call_node->ty.as()) { + if (call_node->ty.as()) { if (auto op = call_node->op.as()) { static auto infer_type_map = Op::GetAttrMap(tvm::op_attr::kInferType); static auto infer_type_with_builder_map = @@ -163,7 +163,7 @@ class CodeGenRunner : ExprMutator { call_node->loc); } - Expr VisitExpr_(const FunctionNode* func_node) override { + Expr VisitExpr_(const FunctionObj* func_node) override { Function func = ffi::GetRef(func_node); auto opt_codegen = func->GetAttr(tvm::relax::attr::kCodegen); if (opt_codegen) { @@ -189,11 +189,11 @@ class CodeGenRunner : ExprMutator { std::unordered_map> target_functions; for (const auto& entry : mod->functions) { - if (entry.second->IsInstance()) { + if (entry.second->IsInstance()) { continue; } PostOrderVisit(entry.second, [&target_functions](Expr e) { - if (e->IsInstance()) { + if (e->IsInstance()) { auto f = e.as_or_throw(); if (auto target_opt = f->GetAttr(tvm::relax::attr::kCodegen)) { ffi::String target = target_opt.value(); @@ -224,7 +224,7 @@ class CodeGenRunner : ExprMutator { /*! \brief The names of all constants in the original module. */ ffi::Map constant_names; /*! \brief Extern funcs for each global variable. */ - std::unordered_map extern_funcs_; + std::unordered_map extern_funcs_; }; } // namespace relax diff --git a/src/relax/transform/specialize_function_based_on_callsite.cc b/src/relax/transform/specialize_function_based_on_callsite.cc index 4c7776769c50..227b6f52a500 100644 --- a/src/relax/transform/specialize_function_based_on_callsite.cc +++ b/src/relax/transform/specialize_function_based_on_callsite.cc @@ -54,7 +54,7 @@ class SpecializeTIRCallArgs : ExprMutator { IRModule Run(IRModule mod) { mod_ = mod; for (const auto& [gv, func] : mod->functions) { - if (func->IsInstance()) { + if (func->IsInstance()) { const auto& base_func = mod->Lookup(gv); // Only non primitive relax functions if (base_func->HasNonzeroAttr(tvm::relax::attr::kPrimitive.name)) { @@ -70,7 +70,7 @@ class SpecializeTIRCallArgs : ExprMutator { using ExprMutator::VisitExpr_; - Expr VisitExpr_(const CallNode* call_node) override { + Expr VisitExpr_(const CallObj* call_node) override { auto call = ExprMutator::VisitExpr_(call_node).as_or_throw(); static const Op call_tir_op = Op::Get("relax.call_tir"); if (call->op.same_as(call_tir_op)) { @@ -88,7 +88,7 @@ class SpecializeTIRCallArgs : ExprMutator { for (size_t i = 0; i < args.size(); ++i) { auto ty = GetType(args[i]); - TVM_FFI_ICHECK(ty->IsInstance()) + TVM_FFI_ICHECK(ty->IsInstance()) << "Expected Tensor struct Info for call :" << call->op; auto tensor_ty = ty.as_or_throw(); TVM_FFI_ICHECK(tensor_ty->shape.has_value()) << "Shape undefined for call:" << call->args[0]; @@ -97,7 +97,7 @@ class SpecializeTIRCallArgs : ExprMutator { scope = tensor_ty->vdevice.value()->memory_scope; } ffi::String name; - if (args[i]->IsInstance()) { + if (args[i]->IsInstance()) { name = args[i].as_or_throw()->name; } else { name = std::string({static_cast('A' + i)}); @@ -110,7 +110,7 @@ class SpecializeTIRCallArgs : ExprMutator { } runtime::MemoryScope scope; auto out_ty = call->ty_args[0]; - if (out_ty->IsInstance()) { + if (out_ty->IsInstance()) { auto ty = out_ty.as_or_throw(); if (ty->vdevice.has_value()) { scope = ty->vdevice.value()->memory_scope; @@ -119,7 +119,7 @@ class SpecializeTIRCallArgs : ExprMutator { "ret_val", tirx::TensorType(GetShapeFromTensorType(ty), ty->dtype.value(), scope)); param_map.Set(pfunc->params[pfunc->params.size() - 1], tensor); } else { - TVM_FFI_ICHECK(out_ty->IsInstance()) + TVM_FFI_ICHECK(out_ty->IsInstance()) << "Expect output type of call_tir to be either TupleType or " "TensorType, but got " << out_ty; @@ -128,7 +128,7 @@ class SpecializeTIRCallArgs : ExprMutator { ffi::Array ty_fields; int index = 0; for (const auto& si : tuple_ty->fields) { - TVM_FFI_ICHECK(si->IsInstance()) + TVM_FFI_ICHECK(si->IsInstance()) << "Fields of TupleType must be TensorType for call_tir " "output structinfo, but got " << si; @@ -156,7 +156,9 @@ class SpecializeTIRCallArgs : ExprMutator { TVM_FFI_CHECK(var.has_value() && !replacements.count(var.value()), ValueError) << "Inconsistent call_tir tensor dimension: " << symbolic << " vs. " << concrete; replacements.Set(var.value(), concrete); - analyzer->Bind(var.value(), concrete); + auto context = analyzer->GetConstraintContext(); + context->Append(sym::BindingConstraint(var.value(), concrete)); + analyzer->SyncConstraintContext(context); }; for (const auto& param : pfunc->params) { auto descriptor = param_map.Get(param); @@ -196,7 +198,7 @@ class SpecializeTIRCallArgs : ExprMutator { [](const tvm::Call& call) { if (!call->op.same_as(tirx::tensor_data_ptr_op())) return call; // Pointer projections inherit the newly bound tensor storage scope. - auto copy = ffi::make_object(*call.get()); + auto copy = ffi::make_object(*call.get()); copy->ty = tvm::Call::ReinferType(copy.get()); return tvm::Call(std::move(copy)); }, diff --git a/src/relax/transform/split_call_tir_by_pattern.cc b/src/relax/transform/split_call_tir_by_pattern.cc index 778be8ae423e..d53be0ac578c 100644 --- a/src/relax/transform/split_call_tir_by_pattern.cc +++ b/src/relax/transform/split_call_tir_by_pattern.cc @@ -65,9 +65,9 @@ class ForMatcher : public TensorizeComparator { bool Match(const For& top) { TVM_FFI_ICHECK(pattern_->body.has_value() && pattern_->body.value()->size() == 1) << "Invalid pattern function"; - const auto* root = pattern_->body.value()->seq[0].as(); + const auto* root = pattern_->body.value()->seq[0].as(); TVM_FFI_ICHECK(root && root->block->body->size() == 1) << "Invalid pattern function"; - const ForNode* pattern_top = root->block->body->seq[0].as(); + const ForObj* pattern_top = root->block->body->seq[0].as(); TVM_FFI_ICHECK(pattern_top) << "Invalid pattern function"; if (!Dispatch(top, ffi::GetRef(pattern_top))) { return false; @@ -106,7 +106,7 @@ class ForMatcher : public TensorizeComparator { if (pattern_vars_.count(rhs_var)) { // special case for pattern vars if (!lhs.as()) { - if (lhs->IsInstance() || lhs->IsInstance()) { + if (lhs->IsInstance() || lhs->IsInstance()) { ffi::Optional value = QueryEvaluatedSymbols(rhs_var); if (value.has_value()) { if (!analyzer_->CanProveEqual(lhs, value.value())) return false; @@ -121,7 +121,7 @@ class ForMatcher : public TensorizeComparator { } } // pattern_var * expr - if (const auto* rhs_ptr = rhs.as()) { + if (const auto* rhs_ptr = rhs.as()) { auto operand_a = rhs_ptr->a.as(); auto operand_b = rhs_ptr->b.as(); if (operand_a && pattern_vars_.count(operand_a.value())) { @@ -154,7 +154,7 @@ class ForMatcher : public TensorizeComparator { } } // pattern_Var + expr - if (const auto* rhs_ptr = rhs.as()) { + if (const auto* rhs_ptr = rhs.as()) { auto operand_a = rhs_ptr->a.as(); auto operand_b = rhs_ptr->b.as(); if (operand_a && pattern_vars_.count(operand_a.value())) { @@ -189,8 +189,8 @@ class ForMatcher : public TensorizeComparator { return TensorizeComparator::Dispatch(lhs, rhs); } - bool Dispatch_(const prim::AddNode* add, const PrimExpr& other) final { - const auto* rhs = other.as(); + bool Dispatch_(const prim::AddObj* add, const PrimExpr& other) final { + const auto* rhs = other.as(); if (rhs == nullptr) return false; { this->evaluated_symbols.push_back(SymbolMap()); @@ -215,8 +215,8 @@ class ForMatcher : public TensorizeComparator { return false; } - bool Dispatch_(const prim::MulNode* mul, const PrimExpr& other) final { - const auto* rhs = other.as(); + bool Dispatch_(const prim::MulObj* mul, const PrimExpr& other) final { + const auto* rhs = other.as(); if (rhs == nullptr) return false; { this->evaluated_symbols.push_back(SymbolMap()); @@ -241,11 +241,11 @@ class ForMatcher : public TensorizeComparator { return false; } - bool Dispatch_(const CallNode* call, const PrimExpr& other) final { - const auto* rhs = other.as(); + bool Dispatch_(const CallObj* call, const PrimExpr& other) final { + const auto* rhs = other.as(); if (rhs == nullptr) return false; - const auto* lhs_op = call->op.as(); - const auto* rhs_op = rhs->op.as(); + const auto* lhs_op = call->op.as(); + const auto* rhs_op = rhs->op.as(); if (lhs_op == nullptr || rhs_op == nullptr) return false; if (lhs_op->name != rhs_op->name) return false; if (call->args.size() != rhs->args.size()) return false; @@ -255,17 +255,17 @@ class ForMatcher : public TensorizeComparator { return true; } - bool Dispatch_(const tvm::ForNode* op, const Stmt& other) final { - const auto* rhs = other.as(); + bool Dispatch_(const tvm::ForObj* op, const Stmt& other) final { + const auto* rhs = other.as(); loop_stack_lhs_.push_back(ffi::GetRef(op)); loop_stack_rhs_.push_back(ffi::GetRef(rhs)); // The body of loop must be loop or BlockRealize - if (op->body->size() != 1 || (!op->body->seq[0]->IsInstance() && - !op->body->seq[0]->IsInstance())) { + if (op->body->size() != 1 || (!op->body->seq[0]->IsInstance() && + !op->body->seq[0]->IsInstance())) { return false; } - if (rhs->body->size() != 1 || (!rhs->body->seq[0]->IsInstance() && - !rhs->body->seq[0]->IsInstance())) { + if (rhs->body->size() != 1 || (!rhs->body->seq[0]->IsInstance() && + !rhs->body->seq[0]->IsInstance())) { return false; } // Build mapping between the loop vars @@ -279,8 +279,8 @@ class ForMatcher : public TensorizeComparator { return Dispatch(op->body->seq[0], rhs->body->seq[0]); } - bool Dispatch_(const s_tir::SBlockNode* op, const Stmt& other) final { - const auto* rhs = other.as(); + bool Dispatch_(const s_tir::SBlockObj* op, const Stmt& other) final { + const auto* rhs = other.as(); // Check block equality. // All iter vars and tensor regions including the order should match. // When checking iter vars, DefEqual is used to remap variables. @@ -297,8 +297,8 @@ class ForMatcher : public TensorizeComparator { } // The body of the block has to be TensorStore if (op->body->size() != 1 || rhs->body->size() != 1 || - !op->body->seq[0]->IsInstance() || - !rhs->body->seq[0]->IsInstance()) { + !op->body->seq[0]->IsInstance() || + !rhs->body->seq[0]->IsInstance()) { return false; } // Handle init block @@ -310,8 +310,8 @@ class ForMatcher : public TensorizeComparator { return Dispatch(op->body->seq[0], rhs->body->seq[0]); } - bool Dispatch_(const s_tir::SBlockRealizeNode* op, const Stmt& other) final { - const auto* rhs = other.as(); + bool Dispatch_(const s_tir::SBlockRealizeObj* op, const Stmt& other) final { + const auto* rhs = other.as(); // Only allow trivial bindings for (size_t i = 0; i < op->iter_values.size(); ++i) { if (!op->iter_values[i].same_as(loop_stack_lhs_[i]->var)) return false; @@ -324,13 +324,13 @@ class ForMatcher : public TensorizeComparator { return Dispatch(op->block, rhs->block); } - bool Dispatch_(const TensorStoreNode* op, const Stmt& other) { - const auto* rhs = other.as(); + bool Dispatch_(const TensorStoreObj* op, const Stmt& other) { + const auto* rhs = other.as(); return CompareTensorAccess(op, rhs) && Dispatch(op->value, rhs->value); } - bool Dispatch_(const TensorLoadNode* op, const PrimExpr& other) { - const auto* rhs = other.as(); + bool Dispatch_(const TensorLoadObj* op, const PrimExpr& other) { + const auto* rhs = other.as(); return CompareTensorAccess(op, rhs); } @@ -362,7 +362,7 @@ class ForMatcher : public TensorizeComparator { return CompareArray(lhs->region, rhs->region, &ForMatcher::CompareRange); } - bool CompareTensorAccess(const TensorStoreNode* lhs, const TensorStoreNode* rhs) { + bool CompareTensorAccess(const TensorStoreObj* lhs, const TensorStoreObj* rhs) { if (!CompareTensor(lhs->dest.as_or_throw(), rhs->dest.as_or_throw())) { return false; } @@ -371,7 +371,7 @@ class ForMatcher : public TensorizeComparator { static_cast(&ForMatcher::Dispatch)); } - bool CompareTensorAccess(const TensorLoadNode* lhs, const TensorLoadNode* rhs) { + bool CompareTensorAccess(const TensorLoadObj* lhs, const TensorLoadObj* rhs) { if (rhs == nullptr) return false; if (!CompareTensor(lhs->source.as_or_throw(), rhs->source.as_or_throw())) { @@ -418,7 +418,7 @@ class TIRPatternMatcher { ffi::Array pattern_symbolic_vars; int tensor_count = 0; while (tensor_count < static_cast(pattern_func->params.size()) && - pattern_func->params[tensor_count]->ty.as()) { + pattern_func->params[tensor_count]->ty.as()) { ++tensor_count; } for (int i = tensor_count; i < static_cast(pattern_func->params.size()); i++) { @@ -444,7 +444,7 @@ class TIRPatternMatcher { // For each block in the body, try to find its corresponding pattern one by one void OpMatternMatch(const SeqStmt& body) { for (const Stmt& stmt : body->seq) { - const ForNode* loop = stmt.as(); + const ForObj* loop = stmt.as(); if (loop == nullptr || !BlockPatternMatch(ffi::GetRef(loop))) { break; } @@ -484,7 +484,7 @@ class FunctionPartitioner : public s_tir::StmtExprVisitor { bool fail = false; private: - ffi::Optional Visit_(const s_tir::SBlockNode* op) final { + ffi::Optional Visit_(const s_tir::SBlockObj* op) final { block_counter_++; bool is_matching_ = block_counter_ <= num_matched_ops_; if (block_counter_ == num_matched_ops_ && intermediate_tensor.has_value()) { @@ -538,11 +538,11 @@ class BlockRemover : public s_tir::StmtExprMutator { : block_partition(block_partition), allocs_(allocs), is_library_part_(is_library_part) {} private: - UnchangedOr Mutate_(const s_tir::SBlockNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const s_tir::SBlockObj* op, InplaceMode inplace_mode) final { s_tir::SBlock block = s_tir::StmtExprMutator::Mutate_(op, inplace_mode) .ValueOrUnchanged(ffi::GetRef(op)) .as_or_throw(); - ffi::ObjectPtr n = ffi::make_object(*block.operator->()); + ffi::ObjectPtr n = ffi::make_object(*block.operator->()); if (op->name_hint != "root") { TVM_FFI_ICHECK(block_partition.count(ffi::GetRef(op))); bool block_is_library = block_partition[ffi::GetRef(op)]; @@ -562,7 +562,7 @@ class BlockRemover : public s_tir::StmtExprMutator { return s_tir::SBlock(n); } - UnchangedOr Mutate_(const SeqStmtNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const SeqStmtObj* op, InplaceMode inplace_mode) final { ffi::Array seq; bool erased_all = !op->seq.empty(); for (const Stmt& s : op->seq) { @@ -602,7 +602,7 @@ std::pair> SplitFunctions( if (!func->body.has_value() || func->body.value()->size() != 1) { return {func, std::nullopt}; } - const auto* root = func->body.value()->seq[0].as(); + const auto* root = func->body.value()->seq[0].as(); if (root == nullptr) return {func, std::nullopt}; SeqStmt body = root->block->body; ffi::Array match_results = TIRPatternMatcher::Match(patterns, body); @@ -685,7 +685,7 @@ void StringReplace(std::string* subject, const std::string& search, const std::s } } -tvm::BaseFunc CodegenWithLibrary(const tirx::FunctionNode* pf, ffi::String global_symbol) { +tvm::BaseFunc CodegenWithLibrary(const tirx::FunctionObj* pf, ffi::String global_symbol) { using namespace tvm::tirx; ffi::Optional library_code = pf->attrs.GetAttr(tvm::relax::attr::kLibraryKernel); if (!library_code.has_value()) { @@ -712,7 +712,7 @@ class SplitMutator : public ExprMutator { FCodegen fcodegen) { SplitMutator mutator(mod, patterns, fcodegen); for (auto& kv : mod->functions) { - if (auto* func = kv.second.as()) { + if (auto* func = kv.second.as()) { Function new_func = mutator(ffi::GetRef(func)).as_or_throw(); mutator.builder_->UpdateFunction(kv.first, new_func); } @@ -724,20 +724,20 @@ class SplitMutator : public ExprMutator { using ExprMutator::VisitExpr_; inline ffi::Array GetCallTIRArgs(Expr args) { - if (args.as()) { - return args.as()->fields; + if (args.as()) { + return args.as()->fields; } else { return {args}; } } - Expr VisitExpr_(const CallNode* op) final { + Expr VisitExpr_(const CallObj* op) final { Call call = ExprMutator::VisitExpr_(op).as_or_throw(); static const Op call_tir_op_ = Op::Get("relax.call_tir"); static const Op call_dps_packed_ = Op::Get("relax.call_dps_packed"); if (!call->op.same_as(call_tir_op_)) return call; // the first argument is the function to be called - const auto* gv_ptr = call->args[0].as(); + const auto* gv_ptr = call->args[0].as(); if (gv_ptr == nullptr) return call; GlobalVar gv = ffi::GetRef(gv_ptr); // retrieve the function from the module and split it @@ -749,12 +749,12 @@ class SplitMutator : public ExprMutator { if (!split_funcs.second.has_value()) { // no need to split, the function itself a library kernel tvm::BaseFunc lib_func = CodegenWithLibrary(split_funcs.first.get(), gv->name_hint); - if (lib_func->IsInstance()) return ffi::GetRef(op); + if (lib_func->IsInstance()) return ffi::GetRef(op); // Update the function in the module with the library kernel - TVM_FFI_ICHECK(lib_func->IsInstance()); + TVM_FFI_ICHECK(lib_func->IsInstance()); builder_->UpdateFunction(gv, lib_func); // emit the call to the library kernel - ffi::ObjectPtr new_call = ffi::make_object(*call.operator->()); + ffi::ObjectPtr new_call = ffi::make_object(*call.operator->()); new_call->op = this->call_dps_packed_; new_call->args = {lib_func, call->args[1]}; return Call(new_call); @@ -769,8 +769,8 @@ class SplitMutator : public ExprMutator { } // replace the function in the module with the library kernel tvm::BaseFunc lib_func = CodegenWithLibrary(func1.get(), gv->name_hint); - if (lib_func->IsInstance()) return ffi::GetRef(op); - TVM_FFI_ICHECK(lib_func->IsInstance()); + if (lib_func->IsInstance()) return ffi::GetRef(op); + TVM_FFI_ICHECK(lib_func->IsInstance()); builder_->UpdateFunction(gv, lib_func); tirx::TensorVar intermediate_tensor = func1->params.back().as_or_throw(); PrimType dtype = intermediate_tensor.dtype(); diff --git a/src/relax/transform/split_layout_rewrite_preproc.cc b/src/relax/transform/split_layout_rewrite_preproc.cc index b8acefc2512f..7f63daa7b603 100644 --- a/src/relax/transform/split_layout_rewrite_preproc.cc +++ b/src/relax/transform/split_layout_rewrite_preproc.cc @@ -42,16 +42,16 @@ class SplitFunctionLayoutRewrite : public s_tir::StmtExprMutator { public: using s_tir::StmtExprMutator::Mutate; UnchangedOr Mutate(ffi::AnyView value, InplaceMode inplace_mode) final { - if (value.as()) return ffi::Unchanged(); + if (value.as()) return ffi::Unchanged(); return s_tir::StmtExprMutator::Mutate(value, inplace_mode); } explicit SplitFunctionLayoutRewrite(const tirx::Function& func) : original_func_(func) {} std::tuple, tirx::Function> Transform(const tirx::Function& func) { TVM_FFI_ICHECK(func->body.has_value() && func->body.value()->size() == 1 && - func->body.value()->seq[0].as()) + func->body.value()->seq[0].as()) << "The body of the function should be a root block."; - const auto& block = func->body.value()->seq[0].as()->block; + const auto& block = func->body.value()->seq[0].as()->block; visit_root_block(block.get()); if (layout_rewrite_preproc_stmts_.size() > 0) { return std::make_tuple(create_layout_rewrite_preproc_func(), create_compute_func()); @@ -118,7 +118,7 @@ class SplitFunctionLayoutRewrite : public s_tir::StmtExprMutator { // Step 2: Create the body for the new tirx::Function Stmt body = SeqStmt(compute_stmts_); s_tir::SBlock original_block = - original_func_->body.value()->seq[0].as()->block; + original_func_->body.value()->seq[0].as()->block; ffi::Array alloc_tensors; for (const auto& tensor : original_block->alloc_tensors) { auto it = @@ -152,7 +152,7 @@ class SplitFunctionLayoutRewrite : public s_tir::StmtExprMutator { return tirx::RenewDef(func); } - void visit_root_block(const s_tir::SBlockNode* op) { + void visit_root_block(const s_tir::SBlockObj* op) { for (const auto& stmt : op->body->seq) { current_subtree_ = 0; Stmt new_stmt = this->Mutate(stmt).ValueOrUnchanged(stmt); @@ -166,7 +166,7 @@ class SplitFunctionLayoutRewrite : public s_tir::StmtExprMutator { } } } - UnchangedOr Mutate_(const s_tir::SBlockNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const s_tir::SBlockObj* op, InplaceMode inplace_mode) final { s_tir::SBlock block = s_tir::StmtExprMutator::Mutate_(op, inplace_mode) .ValueOrUnchanged(ffi::GetRef(op)) .as_or_throw(); @@ -210,7 +210,7 @@ class SplitFunctionLayoutRewrite : public s_tir::StmtExprMutator { auto new_annotations = op->annotations; new_annotations.erase(tvm::s_tir::attr::kMetaScheduleLayoutRewritePreproc.name); - auto n = ffi::make_object(*block.get()); + auto n = ffi::make_object(*block.get()); n->annotations = new_annotations; return s_tir::SBlock(n); } @@ -250,7 +250,7 @@ class SplitLayoutRewritePreproc : public ExprMutator { // Step 1: Split the function into preproc and compute for (auto [gv, func] : mod->functions) { - if (func->IsInstance()) { + if (func->IsInstance()) { auto tir_rewriter = ffi::make_object(func.as_or_throw()); auto [preproc_func, compute_func] = @@ -264,7 +264,7 @@ class SplitLayoutRewritePreproc : public ExprMutator { } for (auto [gv, func] : mod->functions) { - if (func->IsInstance()) { + if (func->IsInstance()) { auto relax_func = func.as_or_throw(); mutator.builder_->UpdateFunction(gv, mutator(relax_func).as_or_throw()); } @@ -276,7 +276,7 @@ class SplitLayoutRewritePreproc : public ExprMutator { explicit SplitLayoutRewritePreproc(const IRModule& mod) : ExprMutator(mod) {} using ExprMutator::VisitExpr_; - Expr VisitExpr_(const CallNode* op) final { + Expr VisitExpr_(const CallObj* op) final { static const Op call_tir_op = Op::Get("relax.call_tir"); Call call = ExprMutator::VisitExpr_(op).as_or_throw(); @@ -310,7 +310,7 @@ class SplitLayoutRewritePreproc : public ExprMutator { preproc_args.push_back(call_tir_args[info.tensor_index]); tirx::TensorVar rewritten_tensor = info.post_rewrite_tensor; for (const auto& shape_expr : rewritten_tensor.shape()) { - TVM_FFI_ICHECK(shape_expr.as()) + TVM_FFI_ICHECK(shape_expr.as()) << "Currently does not support rewrite tensor with " "dynamic shape."; } @@ -338,8 +338,8 @@ class SplitLayoutRewritePreproc : public ExprMutator { } private: - std::unordered_map> split_funcs_; - std::unordered_map> split_funcs_; + std::unordered_map> rewrite_infos_; }; diff --git a/src/relax/transform/static_plan_block_memory.cc b/src/relax/transform/static_plan_block_memory.cc index cb54fc5cbf65..c519e984f99a 100644 --- a/src/relax/transform/static_plan_block_memory.cc +++ b/src/relax/transform/static_plan_block_memory.cc @@ -94,8 +94,8 @@ using namespace tvm::prim; * - that is a function parameter, * - that is a function return value, * - one of whose use site is a BindingBlock different from its allocation site, - * - that is used as a condition or branch return of a IfExprNode, - * - that is used as the body of a SeqExprNode, + * - that is used as a condition or branch return of a IfExprObj, + * - that is used as the body of a SeqExprObj, * - that is used as arguments in a Call whose op is not a tirx::Function. * * In practice, we do create a storage token for such tensor at first. But at @@ -123,7 +123,7 @@ class StorageTokenObj : public ffi::Object { /*! \brief Get the constant number of bytes that this token requires, or -1 if the number of bytes * is symbolic */ int64_t const_bytes() const { - const auto* imm = bytes.as(); + const auto* imm = bytes.as(); auto const_val = imm ? imm->value.as() : std::nullopt; if (const_val.has_value()) { return *const_val; @@ -176,7 +176,7 @@ class StorageToken : public ffi::ObjectRef { } if (!size_computed) { for (const PrimExpr& dim_len : shape) { - if (const IntImmNode* const_dim_len = dim_len.as()) { + if (const IntImmObj* const_dim_len = dim_len.as()) { const_coeff *= const_dim_len->value; } else { size *= dim_len; @@ -341,7 +341,7 @@ bool IsInplaceMemoryOp(const Expr& op) { static const Op reshape_op = Op::Get("relax.reshape"); static const Op view_op = Op::Get("relax.memory.view"); static const Op ensure_zero_offset_op = Op::Get("relax.memory.ensure_zero_offset"); - const auto* extern_func = op.as(); + const auto* extern_func = op.as(); bool is_builtin_reshape = extern_func != nullptr && extern_func->global_symbol == "vm.builtin.reshape"; return op.same_as(reshape_op) || op.same_as(view_op) || op.same_as(ensure_zero_offset_op) || @@ -353,7 +353,7 @@ class StorageAllocatorBaseVisitor : public ExprVisitor { protected: using ExprVisitor::VisitExpr_; - void VisitBindingBlock_(const BindingBlockNode* block) override { + void VisitBindingBlock_(const BindingBlockObj* block) override { // We maintain a block stack for token allocation-site and use-site check. block_stack_.push_back(block); ExprVisitor::VisitBindingBlock_(block); @@ -362,19 +362,19 @@ class StorageAllocatorBaseVisitor : public ExprVisitor { block_stack_.pop_back(); } - void VisitBinding_(const VarBindingNode* binding) override { + void VisitBinding_(const VarBindingObj* binding) override { ExprVisitor::VisitBinding_(binding); // The binding var has the same tokens as the binding value. SetTokens(binding->var.get(), token_map_[binding->value.get()]); } - void VisitBinding_(const MatchCastNode* binding) override { + void VisitBinding_(const MatchCastObj* binding) override { ExprVisitor::VisitBinding_(binding); // MatchCast refines the type without changing the underlying storage. SetTokens(binding->var.get(), token_map_[binding->value.get()]); } - void VisitBindingBlock_(const DataflowBlockNode* block) override { + void VisitBindingBlock_(const DataflowBlockObj* block) override { // We maintain a block stack for token allocation-site and use-site check. block_stack_.push_back(block); ExprVisitor::VisitBindingBlock_(block); @@ -383,7 +383,7 @@ class StorageAllocatorBaseVisitor : public ExprVisitor { block_stack_.pop_back(); } - void VisitExpr_(const TupleNode* tuple) final { + void VisitExpr_(const TupleObj* tuple) final { ffi::Array tokens; tokens.reserve(tuple->fields.size()); for (const Expr& field : tuple->fields) { @@ -393,7 +393,7 @@ class StorageAllocatorBaseVisitor : public ExprVisitor { SetTokens(tuple, Tokens(tokens)); } - void VisitExpr_(const TupleGetItemNode* tuple_item) final { + void VisitExpr_(const TupleGetItemObj* tuple_item) final { Tokens tokens = GetTokens(tuple_item->tuple); // If the tuple has no token, every of its field has no token as well. if (tokens.IsNull()) { @@ -414,12 +414,12 @@ class StorageAllocatorBaseVisitor : public ExprVisitor { return token_map_[expr.get()]; } - virtual void SetTokens(const ExprNode* expr, Tokens tokens) { token_map_[expr] = tokens; } + virtual void SetTokens(const ExprObj* expr, Tokens tokens) { token_map_[expr] = tokens; } /*! \brief The mapping from each Expr to its corresponding storage tokens. */ - std::unordered_map token_map_; + std::unordered_map token_map_; /*! \brief The binding block stack. */ - std::vector block_stack_; + std::vector block_stack_; }; /*! @@ -467,12 +467,18 @@ void SetTIRVarRangeConstraints(Function func, sym::AnalyzerObj* ana, const ffi::BigInt& upper = it_upper->second->value; tvm::Range range = tvm::Range::FromMinExtent(tvm::IntImm::Int64(lower), tvm::IntImm::Int64(upper - lower + 1)); - ana->Bind(tir_var, range); + auto context = ana->GetConstraintContext(); + context->Append(sym::RangeConstraint(tir_var, range)); + ana->SyncConstraintContext(context); 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()); + auto context = ana->GetConstraintContext(); + context->Append(sym::PredicateConstraint(tir_var.as_or_throw() >= 0)); + ana->SyncConstraintContext(context); } else if (non_negative_var_attr.count(tir_var->name)) { - ana->MarkGlobalNonNegValue(tir_var.as_or_throw()); + auto context = ana->GetConstraintContext(); + context->Append(sym::PredicateConstraint(tir_var.as_or_throw() >= 0)); + ana->SyncConstraintContext(context); } } } @@ -510,7 +516,7 @@ ffi::Array GetUpperBoundShape(ffi::Array shape, sym::Analyze /*! \brief Check if a shape is static (a.k.a., has no TIR variable). */ bool IsStaticShape(ffi::Array shape) { for (const PrimExpr& dim : shape) { - const auto* int_len = dim.as(); + const auto* int_len = dim.as(); if (!int_len) { return false; } @@ -534,12 +540,12 @@ class StorageAllocatorInit : public StorageAllocatorBaseVisitor { * \param analyzer The arithmetic analyzer. * \return The mapping from each Expr to the token it uses. */ - static std::unordered_map Initialize(const IRModule& mod, - sym::AnalyzerObj* analyzer) { + static std::unordered_map Initialize(const IRModule& mod, + sym::AnalyzerObj* analyzer) { StorageAllocatorInit initializer(mod, analyzer); for (auto it : mod->functions) { - const auto* func = it.second.as(); + const auto* func = it.second.as(); if (func == nullptr) { continue; } @@ -554,7 +560,7 @@ class StorageAllocatorInit : public StorageAllocatorBaseVisitor { explicit StorageAllocatorInit(const IRModule& ctx_mod, sym::AnalyzerObj* analyzer) : ctx_mod_(ctx_mod), analyzer_(analyzer) {} - void VisitExpr_(const FunctionNode* func) final { + void VisitExpr_(const FunctionObj* func) final { // Set the range constraints of TIR variables in the analyzer. SetTIRVarRangeConstraints(ffi::GetRef(func), analyzer_, &dom_map_); // Recurse into the function to get its tokens. @@ -563,7 +569,7 @@ class StorageAllocatorInit : public StorageAllocatorBaseVisitor { DiscardTokensIn(body_tokens); } - void VisitExpr_(const CallNode* call) final { + void VisitExpr_(const CallObj* call) final { static const Op alloc_tensor_op = Op::Get("relax.builtin.alloc_tensor"); static const Op call_tir_dyn_op = Op::Get("relax.vm.call_tir_dyn"); static const Op call_tir_packed_op = Op::Get("relax.call_tir_packed"); @@ -586,8 +592,8 @@ class StorageAllocatorInit : public StorageAllocatorBaseVisitor { // potential external reference. // A native pointer result may alias an argument, with no lifetime tracked by // Relax. Keep such inputs out of the reusable storage pool. - bool reusable_tirx_call = call->op.same_as(call_tir_packed_op) && !call->ty.as(); - if (IsFunctionGlobalVar(call->op) || call->op->IsInstance() || + bool reusable_tirx_call = call->op.same_as(call_tir_packed_op) && !call->ty.as(); + if (IsFunctionGlobalVar(call->op) || call->op->IsInstance() || call->op.same_as(call_tir_dyn_op) || reusable_tirx_call) { ffi::Array args = (call->op.same_as(call_tir_dyn_op) || call->op.same_as(call_tir_packed_op)) @@ -605,7 +611,7 @@ class StorageAllocatorInit : public StorageAllocatorBaseVisitor { } } - void VisitExpr_(const IfExprNode* if_node) final { + void VisitExpr_(const IfExprObj* if_node) final { Tokens cond_tokens = GetTokens(if_node->cond); Tokens then_tokens = GetTokens(if_node->true_branch); Tokens else_tokens = GetTokens(if_node->false_branch); @@ -616,7 +622,7 @@ class StorageAllocatorInit : public StorageAllocatorBaseVisitor { DiscardTokensIn(else_tokens); } - void VisitExpr_(const SeqExprNode* seq) final { + void VisitExpr_(const SeqExprObj* seq) final { for (const BindingBlock& binding_block : seq->blocks) { this->VisitBindingBlock(binding_block); } @@ -634,7 +640,7 @@ class StorageAllocatorInit : public StorageAllocatorBaseVisitor { * \return A boolean indicating if the input op corresponds to a tirx::Function. */ bool IsFunctionGlobalVar(const Expr& op) { - const auto* global_var = op.as(); + const auto* global_var = op.as(); if (global_var == nullptr) { return false; } @@ -642,7 +648,7 @@ class StorageAllocatorInit : public StorageAllocatorBaseVisitor { if (func_it == ctx_mod_->functions.end()) { return false; } - return (*func_it).second->IsInstance(); + return (*func_it).second->IsInstance(); } /*! @@ -650,15 +656,15 @@ class StorageAllocatorInit : public StorageAllocatorBaseVisitor { * \param call The call to be processed. * \return The created token. */ - Tokens CreateToken(const CallNode* call) { + Tokens CreateToken(const CallObj* call) { // Sanity checks about // - the call return value is a Tensor; // - the shape of the tensor is known, in the form of ShapeExpr; // - the tensor has known dtype; // - no storage token was created for this call before. - const auto* ty = call->ty.as(); + const auto* ty = call->ty.as(); TVM_FFI_ICHECK_NOTNULL(ty); - const auto* shape = ty->shape.as(); + const auto* shape = ty->shape.as(); TVM_FFI_ICHECK_NOTNULL(shape); TVM_FFI_ICHECK(!ty->IsUnknownDtype()); TVM_FFI_ICHECK(ty->dtype.value()->dtype == call->args[1].as_or_throw()->value); @@ -673,7 +679,7 @@ class StorageAllocatorInit : public StorageAllocatorBaseVisitor { StringImm storage_scope = call->args[3].as_or_throw(); int64_t vdevice_index = -1; - if (const auto* int_imm = call->args[2].as()) { + if (const auto* int_imm = call->args[2].as()) { vdevice_index = int_imm->value.as().value(); } ffi::Optional vdevice = GetGlobalVDevice(ctx_mod_, vdevice_index); @@ -694,7 +700,7 @@ class StorageAllocatorInit : public StorageAllocatorBaseVisitor { * When we want to discard one token, we use the records to remove the token * from the Expr that are using it. */ - void SetTokens(const ExprNode* expr, Tokens tokens) final { + void SetTokens(const ExprObj* expr, Tokens tokens) final { StorageAllocatorBaseVisitor::SetTokens(expr, tokens); ForEachLeaf(tokens, [this, expr](StorageToken token) { this->token2exprs_[token.get()].push_back(expr); @@ -711,7 +717,7 @@ class StorageAllocatorInit : public StorageAllocatorBaseVisitor { * \param cur_block The pointer to the current block. * \return The tokens used by the input Expr. */ - Tokens GetTokensWithAllocSiteCheck(const Expr& expr, const BindingBlockNode* cur_block) { + Tokens GetTokensWithAllocSiteCheck(const Expr& expr, const BindingBlockObj* cur_block) { Tokens tokens = GetTokens(expr); ForEachLeaf(tokens, [this, cur_block](StorageToken token) { auto it = this->token2block_.find(token.get()); @@ -734,8 +740,8 @@ class StorageAllocatorInit : public StorageAllocatorBaseVisitor { * \param token_to_discard The token to be discarded. */ void DiscardToken(StorageToken token_to_discard) { - const std::vector& exprs = token2exprs_[token_to_discard.get()]; - for (const ExprNode* expr : exprs) { + const std::vector& exprs = token2exprs_[token_to_discard.get()]; + for (const ExprObj* expr : exprs) { token_map_[expr] = MapNestedMsg(token_map_[expr], [token_to_discard](StorageToken token) { return token.same_as(token_to_discard) ? Tokens() : Tokens(token); }); @@ -754,9 +760,9 @@ class StorageAllocatorInit : public StorageAllocatorBaseVisitor { /*! \brief The domain map of dynamic TIR variables for analysis. */ ffi::Map dom_map_; /*! \brief The mapping from each token to the binding block where it is created. */ - std::unordered_map token2block_; + std::unordered_map token2block_; /*! \brief The mapping from each token to the Exprs that are using this token. */ - std::unordered_map> token2exprs_; + std::unordered_map> token2exprs_; }; /*! @@ -775,7 +781,7 @@ class StorageAllocatorInit : public StorageAllocatorBaseVisitor { */ class StorageAllocator : public StorageAllocatorBaseVisitor { public: - explicit StorageAllocator(std::unordered_map token_map, + explicit StorageAllocator(std::unordered_map token_map, sym::AnalyzerObj* analyzer) : allocator_(analyzer) { this->token_map_ = std::move(token_map); @@ -783,7 +789,7 @@ class StorageAllocator : public StorageAllocatorBaseVisitor { void Allocate(const IRModule& mod) { for (auto it : mod->functions) { - const auto* func = it.second.as(); + const auto* func = it.second.as(); if (func == nullptr) { continue; } @@ -797,15 +803,15 @@ class StorageAllocator : public StorageAllocatorBaseVisitor { * \brief The mapping from each `builtin.alloc_tensor` to its corresponding * underlying storage token that it is using. */ - std::unordered_map alloc_tensor2token; + std::unordered_map alloc_tensor2token; /*! \brief The mapping from each binding block to the storage tokens that are create inside. */ - std::unordered_map> block2tokens; + std::unordered_map> block2tokens; private: using ExprVisitor::VisitBinding_; using ExprVisitor::VisitExpr_; - void VisitBindingBlock_(const BindingBlockNode* block) final { + void VisitBindingBlock_(const BindingBlockObj* block) final { StorageAllocatorBaseVisitor::VisitBindingBlock_(block); // Sanity check: each token allocated inside the block should not be // referenced by anyone at the end of the block. @@ -814,7 +820,7 @@ class StorageAllocator : public StorageAllocatorBaseVisitor { } } - void VisitBinding_(const VarBindingNode* binding, const CallNode* call) final { + void VisitBinding_(const VarBindingObj* binding, const CallObj* call) final { static const Op alloc_tensor_op = Op::Get("relax.builtin.alloc_tensor"); if (call->op.same_as(alloc_tensor_op)) { auto it = token_map_.find(call); @@ -910,8 +916,8 @@ class StorageAllocator : public StorageAllocatorBaseVisitor { class StorageAllocationRewriter : public ExprMutator { public: explicit StorageAllocationRewriter( - IRModule mod, std::unordered_map alloc_tensor2token, - std::unordered_map> block2tokens) + IRModule mod, std::unordered_map alloc_tensor2token, + std::unordered_map> block2tokens) : ExprMutator(std::move(mod)), alloc_tensor2token_(std::move(alloc_tensor2token)), block2tokens_(std::move(block2tokens)) {} @@ -919,7 +925,7 @@ class StorageAllocationRewriter : public ExprMutator { IRModule Rewrite() { const IRModule& mod = builder_->GetContextIRModule(); for (const auto& [gv, base_func] : mod->functions) { - const auto* func_ = base_func.as(); + const auto* func_ = base_func.as(); if (func_ == nullptr) { continue; } @@ -942,7 +948,7 @@ class StorageAllocationRewriter : public ExprMutator { private: using ExprMutator::VisitExpr_; - Expr VisitExpr_(const SeqExprNode* seq) final { + Expr VisitExpr_(const SeqExprObj* seq) final { // A storage var is only visible in the scope it is emitted in, such as an if branch. // Forget the vars emitted in this scope on exit, so that a token first used inside a // branch gets a new `alloc_storage` where it is reused in another branch or after the if. @@ -954,7 +960,7 @@ class StorageAllocationRewriter : public ExprMutator { return ret; } - Expr VisitExpr_(const CallNode* call) final { + Expr VisitExpr_(const CallObj* call) final { static const Op alloc_tensor_op = Op::Get("relax.builtin.alloc_tensor"); static const Op mem_alloc_storage = Op::Get("relax.memory.alloc_storage"); static const Op mem_alloc_tensor = Op::Get("relax.memory.alloc_tensor"); @@ -962,9 +968,9 @@ class StorageAllocationRewriter : public ExprMutator { if (it != alloc_tensor2token_.end()) { // Case 1. This `alloc_tensor` is planned for memory reuse. TVM_FFI_ICHECK(call->op.same_as(alloc_tensor_op)); - const auto* ty = call->ty.as(); + const auto* ty = call->ty.as(); TVM_FFI_ICHECK_NOTNULL(ty); - TVM_FFI_ICHECK_NOTNULL(ty->shape.as()); + TVM_FFI_ICHECK_NOTNULL(ty->shape.as()); PrimExpr runtime_device_index = call->args[2].as_or_throw(); // If the token is visited for the first time, create a storage variable using @@ -1001,9 +1007,9 @@ class StorageAllocationRewriter : public ExprMutator { // estimation, we allocate a storage to its upper bound size, and // allocate a tensor out from it with the actual symbolic shape. - const auto* ty = call->ty.as(); + const auto* ty = call->ty.as(); TVM_FFI_ICHECK_NOTNULL(ty); - const auto* shape = ty->shape.as(); + const auto* shape = ty->shape.as(); TVM_FFI_ICHECK_NOTNULL(shape); ffi::Array upper_bounded_shape = GetUpperBoundShape(shape->values, ana_.get(), dom_map_); @@ -1048,9 +1054,9 @@ class StorageAllocationRewriter : public ExprMutator { * \brief The mapping from each memory-reusable `builtin.alloc_tensor` to its corresponding underlying storage token that it is using. */ - std::unordered_map alloc_tensor2token_; + std::unordered_map alloc_tensor2token_; /*! \brief The mapping from each binding block to the storage tokens that are create inside. */ - std::unordered_map> block2tokens_; + std::unordered_map> block2tokens_; /*! \brief The mapping from each token to its storage var visible in the current scope. */ std::unordered_map token2storage_var_; }; @@ -1059,7 +1065,7 @@ IRModule StaticPlanBlockMemory(IRModule mod) { sym::Analyzer ana; // Step 1. Initialize. - std::unordered_map token_map = + std::unordered_map token_map = StorageAllocatorInit::Initialize(mod, ana.get()); // Step 2. Collect the memory allocation info. StorageAllocator allocator(std::move(token_map), ana.get()); @@ -1093,7 +1099,7 @@ PrimExpr GetTextureMemorySizeFromVDevice(ffi::Array pshape, DLDataType struct Shape { const ffi::Array& shape; int64_t operator[](size_t i) const { - const auto* imm = shape[i].as(); + const auto* imm = shape[i].as(); auto value = imm ? imm->value.as() : std::nullopt; TVM_FFI_ICHECK(value.has_value()) << "Dymamic shapes not suported over texture now"; return *value; diff --git a/src/relax/transform/to_mixed_precision.cc b/src/relax/transform/to_mixed_precision.cc index 690da0fd4398..693667bab269 100644 --- a/src/relax/transform/to_mixed_precision.cc +++ b/src/relax/transform/to_mixed_precision.cc @@ -40,8 +40,8 @@ namespace tvm { namespace relax { -int GetMixedPrecisionInfo(const CallNode* call_node) { - const OpNode* op_node = call_node->op.as(); +int GetMixedPrecisionInfo(const CallObj* call_node) { + const OpObj* op_node = call_node->op.as(); if (op_node == nullptr) { return -1; } @@ -151,9 +151,9 @@ class DTypeDecisionCollector : public ExprVisitor { TVM_FFI_ICHECK(args.size() == to.size()) << "Invalid target dtypes"; for (size_t i = 0; i < args.size(); ++i) { auto fvisitleaf = [&](const Expr& expr, NType to) { - if (const auto* var = expr.as()) { + if (const auto* var = expr.as()) { UpdateVarDTypeMap(ffi::GetRef(var), to); - } else if (expr->IsInstance()) { + } else if (expr->IsInstance()) { // GenericConst can be casted anyway, so we don't need to do anything here return; } else { @@ -178,7 +178,7 @@ class DTypeDecisionCollector : public ExprVisitor { RequireArgsToType(std::move(arg_arr), std::move(to_arr)); } - void VisitVars_(const VarNode* op) { + void VisitVars_(const VarObj* op) { Var var = ffi::GetRef(op); if (IsNestedTensor(var)) { // require the var to be fp32 (its original dtype) @@ -188,9 +188,9 @@ class DTypeDecisionCollector : public ExprVisitor { ExprVisitor::VisitExpr_(op); } - void VisitExpr_(const VarNode* op) final { VisitVars_(op); } + void VisitExpr_(const VarObj* op) final { VisitVars_(op); } - void VisitBinding_(const VarBindingNode* binding, const CallNode* call_node) final { + void VisitBinding_(const VarBindingObj* binding, const CallObj* call_node) final { auto policy = GetMixedPrecisionInfo(call_node); if (policy == -1) { ExprVisitor::VisitBinding_(binding, call_node); @@ -207,19 +207,19 @@ class DTypeDecisionCollector : public ExprVisitor { } } - void VisitBinding_(const VarBindingNode* binding, const TupleNode* tuple_node) final { + void VisitBinding_(const VarBindingObj* binding, const TupleObj* tuple_node) final { // require input fields to be the type of the lhs field respectively NType lhs_type = GetTrackedDType(binding->var); RequireArgsToType(tuple_node->fields, lhs_type.NestedArray()); } - void VisitBinding_(const VarBindingNode* binding, - const TupleGetItemNode* tuple_get_item_node) final { + void VisitBinding_(const VarBindingObj* binding, + const TupleGetItemObj* tuple_get_item_node) final { // require the i-th field rhs tuple to be the type of the lhs NType lhs_type = GetTrackedDType(binding->var); std::vector require_rhs; - const TupleTypeNode* ty = tuple_get_item_node->tuple->ty.as(); - TVM_FFI_ICHECK(ty != nullptr) << "TupleGetItemNode must have TupleType"; + const TupleTypeObj* ty = tuple_get_item_node->tuple->ty.as(); + TVM_FFI_ICHECK(ty != nullptr) << "TupleGetItemObj must have TupleType"; for (size_t i = 0; i < ty->fields.size(); ++i) { if (i == static_cast(tuple_get_item_node->index)) { require_rhs.push_back(lhs_type); @@ -231,33 +231,33 @@ class DTypeDecisionCollector : public ExprVisitor { } // override the following methods to visit in backward order - void VisitExpr_(const SeqExprNode* op) final { + void VisitExpr_(const SeqExprObj* op) final { this->VisitLoc(op->loc); this->VisitExpr(op->body); for (auto it = op->blocks.rbegin(); it != op->blocks.rend(); it++) { this->VisitBindingBlock(*it); } - if (auto* ty = op->ty.as()) { + if (auto* ty = op->ty.as()) { this->VisitExprDepTypeField(ffi::GetRef(ty)); } } - void VisitBindingBlock_(const BindingBlockNode* block) { return; } + void VisitBindingBlock_(const BindingBlockObj* block) { return; } - void VisitBindingBlock_(const DataflowBlockNode* block) { + void VisitBindingBlock_(const DataflowBlockObj* block) { for (auto it = block->bindings.rbegin(); it != block->bindings.rend(); it++) { this->VisitBinding(*it); } } - void VisitExpr_(const IfExprNode* op) final { + void VisitExpr_(const IfExprObj* op) final { this->VisitLoc(op->loc); this->VisitExpr(op->true_branch); this->VisitExpr(op->false_branch); this->VisitExpr(op->cond); - if (auto* ty = op->ty.as()) { + if (auto* ty = op->ty.as()) { this->VisitExprDepTypeField(ffi::GetRef(ty)); } } @@ -285,7 +285,7 @@ class ToMixedPrecisionRewriter : public ExprMutator { } else { if (fp16_input_names_.count(var->name)) { auto ty = GetType(var); - if (auto tensor_ty = ty.as()) { + if (auto tensor_ty = ty.as()) { ffi::Optional vdev; if (tensor_ty->vdevice.has_value()) { vdev = tensor_ty->vdevice.value(); @@ -309,7 +309,7 @@ class ToMixedPrecisionRewriter : public ExprMutator { // Note that this function only accepts expr with nested tensor type Expr RewriteExpr(const Expr& expr, const NType& to) { auto fvisitleaf = [&](const Expr& expr, std::array to) -> Expr { - const auto* tensor = GetTypeAs(expr); + const auto* tensor = GetTypeAs(expr); TVM_FFI_ICHECK(tensor != nullptr) << "Only support rewriting tensor expr"; // We only rewrite the expr if the dtype is not the same as the given dtype if (NTypeEqual()(to[0], NTypeFrom(expr))) return expr; @@ -348,14 +348,14 @@ class ToMixedPrecisionRewriter : public ExprMutator { bool AllFP16Castable(const ffi::Array& args) { auto is_fp16 = [](Type ty) { - if (auto tensor_ty = ty.as(); + if (auto tensor_ty = ty.as(); tensor_ty && tensor_ty->dtype == PrimType::Float(16)) { return true; } return false; }; - auto is_in_fp16_range = [this](const GenericConstNode* constant) { + auto is_in_fp16_range = [this](const GenericConstObj* constant) { const auto& data = constant->value.cast(); TVM_FFI_ICHECK(data.defined()); if (data.dtype().lanes > 1) { @@ -394,8 +394,8 @@ class ToMixedPrecisionRewriter : public ExprMutator { for (const Expr& arg : args) { auto ty = GetType(arg); - auto constant = arg.as(); - auto tuple = arg.as(); + auto constant = arg.as(); + auto tuple = arg.as(); if (!IsNestedTensor(arg) || is_fp16(ty) || (constant && is_in_fp16_range(constant)) || (tuple && AllFP16Castable(tuple->fields))) { @@ -441,7 +441,7 @@ class ToMixedPrecisionRewriter : public ExprMutator { return var; } - Expr VisitExpr_(const VarNode* op) final { + Expr VisitExpr_(const VarObj* op) final { if (!builder_->CurrentBlockIsDataFlow()) { return ExprMutator::VisitExpr_(op); } @@ -456,7 +456,7 @@ class ToMixedPrecisionRewriter : public ExprMutator { CastIfFp16Only(binding->var); } - void VisitBinding_(const VarBindingNode* binding, const CallNode* call_node) final { + void VisitBinding_(const VarBindingObj* binding, const CallObj* call_node) final { if (!builder_->CurrentBlockIsDataFlow()) { ExprMutator::VisitBinding_(binding, call_node); return; @@ -468,7 +468,7 @@ class ToMixedPrecisionRewriter : public ExprMutator { return; } // var = Call(op) - const auto* op_node = call_node->op.as(); + const auto* op_node = call_node->op.as(); TVM_FFI_ICHECK(op_node != nullptr); Op op = ffi::GetRef(op_node); if (wrap_param_op.same_as(op)) { @@ -520,10 +520,10 @@ class ToMixedPrecisionRewriter : public ExprMutator { new_value = builder_->Normalize(Call(new_call)); - if (!binding->var->IsInstance()) { + if (!binding->var->IsInstance()) { // Non-Dataflow var: store the tensors to the original dtype new_value = RewriteExpr(new_value, NTypeFrom(binding->var)); - } else if (policy == kAlways && binding->var->IsInstance()) { + } else if (policy == kAlways && binding->var->IsInstance()) { // kAlways: store the tensors to fp16 // But non-dataflow vars will be stored to the original dtype anyway (see above) new_value = RewriteExpr(new_value, NTypeFrom(new_value, new_dtype)); @@ -533,16 +533,16 @@ class ToMixedPrecisionRewriter : public ExprMutator { ReEmitBinding(binding, builder_->Normalize(new_value)); } - void VisitBinding_(const VarBindingNode* binding, const TupleNode* tuple_node) final { + void VisitBinding_(const VarBindingObj* binding, const TupleObj* tuple_node) final { if (!builder_->CurrentBlockIsDataFlow()) { ExprMutator::VisitBinding_(binding, tuple_node); return; } - ffi::ObjectPtr new_tuple = ffi::make_object(*tuple_node); + ffi::ObjectPtr new_tuple = ffi::make_object(*tuple_node); new_tuple->fields = RemapArgs(tuple_node->fields); new_tuple->ty = Type::Missing(); Expr new_value = builder_->Normalize(Tuple(new_tuple)); - if (!binding->var->IsInstance()) { + if (!binding->var->IsInstance()) { // Global var: store the tensors to the original dtype NType to = NTypeFrom(binding->var); new_value = RewriteExpr(new_value, to); @@ -550,19 +550,19 @@ class ToMixedPrecisionRewriter : public ExprMutator { ReEmitBinding(binding, builder_->Normalize(new_value)); } - void VisitBinding_(const VarBindingNode* binding, - const TupleGetItemNode* tuple_get_item_node) final { + void VisitBinding_(const VarBindingObj* binding, + const TupleGetItemObj* tuple_get_item_node) final { if (!builder_->CurrentBlockIsDataFlow()) { // We don't need to rewrite the tuple_get_item in dataflow block ExprMutator::VisitBinding_(binding, tuple_get_item_node); return; } - ffi::ObjectPtr new_tuple_get_item = - ffi::make_object(*tuple_get_item_node); + ffi::ObjectPtr new_tuple_get_item = + ffi::make_object(*tuple_get_item_node); new_tuple_get_item->tuple = RemapArgs({tuple_get_item_node->tuple})[0]; new_tuple_get_item->ty = Type::Missing(); Expr new_value = TupleGetItem(new_tuple_get_item); - if (!binding->var->IsInstance()) { + if (!binding->var->IsInstance()) { // Global var: store the tensors to the original dtype NType to = NTypeFrom(binding->var); new_value = RewriteExpr(new_value, to); @@ -570,7 +570,7 @@ class ToMixedPrecisionRewriter : public ExprMutator { ReEmitBinding(binding, builder_->Normalize(new_value)); } - BindingBlock VisitBindingBlock_(const DataflowBlockNode* block) { + BindingBlock VisitBindingBlock_(const DataflowBlockObj* block) { builder_->BeginDataflowBlock(); // prepare local versions of params here, if they are fp16 expected only for (auto param : params_) { @@ -589,7 +589,7 @@ class ToMixedPrecisionRewriter : public ExprMutator { return builder_->EndBlock(); } - Expr VisitExpr_(const FunctionNode* op) final { + Expr VisitExpr_(const FunctionObj* op) final { params_ = op->params; return ExprMutator::VisitExpr_(op); } diff --git a/src/relax/transform/to_non_dataflow.cc b/src/relax/transform/to_non_dataflow.cc index b6588686167b..fd80979c9dea 100644 --- a/src/relax/transform/to_non_dataflow.cc +++ b/src/relax/transform/to_non_dataflow.cc @@ -34,7 +34,7 @@ using namespace tvm::prim; class ToNonDFMutator : public ExprMutator { public: Var VisitVarDef(const Var& var) final { - if (var.as()) { + if (var.as()) { Var new_var = Var(var->name, GetType(var), var->loc); this->var_remap_.insert_or_assign(var, new_var); return new_var; @@ -42,7 +42,7 @@ class ToNonDFMutator : public ExprMutator { return var; } - BindingBlock VisitBindingBlock_(const DataflowBlockNode* block) final { + BindingBlock VisitBindingBlock_(const DataflowBlockObj* block) final { builder_->BeginBindingBlock(); for (Binding binding : block->bindings) { this->VisitBinding(binding); diff --git a/src/relax/transform/topological_sort.cc b/src/relax/transform/topological_sort.cc index 80b8f554ef8f..858ef177117c 100644 --- a/src/relax/transform/topological_sort.cc +++ b/src/relax/transform/topological_sort.cc @@ -43,8 +43,8 @@ using DataflowNode = std::variant; bool operator==(const DataflowNode& a, const DataflowNode& b) { if (const tvm::Var* var_a = std::get_if(&a)) { if (const tvm::Var* var_b = std::get_if(&b)) { - const tvm::VarNode* ptr_a = var_a->get(); - const tvm::VarNode* ptr_b = var_b->get(); + const tvm::VarObj* ptr_a = var_a->get(); + const tvm::VarObj* ptr_b = var_b->get(); return ptr_a == ptr_b; } } @@ -58,7 +58,7 @@ template <> struct std::hash { std::size_t operator()(const DataflowNode& node) const noexcept { if (const tvm::Var* var = std::get_if(&node)) { - const tvm::VarNode* ptr = var->get(); + const tvm::VarObj* ptr = var->get(); std::hash hasher; return hasher(ptr); } else { @@ -134,7 +134,7 @@ class BindingOrderCollector : ExprVisitor { private: void VisitVarDef(const Var&) override {} - void VisitExpr_(const FunctionNode* op) override { + void VisitExpr_(const FunctionObj* op) override { for (const auto& var : op->params) { graph_vars_.insert(var); dependencies_.downstream_users[InputNode()].push_back(var); @@ -152,7 +152,7 @@ class BindingOrderCollector : ExprVisitor { current_binding_ = cache; } - void VisitExpr_(const VarNode* op) override { + void VisitExpr_(const VarObj* op) override { Var upstream_requirement = ffi::GetRef(op); if (!graph_vars_.count(upstream_requirement)) { return; @@ -173,7 +173,7 @@ class TopologicalSorter : public ExprMutator { TopologicalSorter(TraversalOrder order, StartingLocation starting_location) : order_(order), starting_location_(starting_location) {} - Expr VisitExpr_(const FunctionNode* op) override { + Expr VisitExpr_(const FunctionObj* op) override { auto cached = dependencies_; dependencies_ = BindingOrderCollector::Collect(ffi::GetRef(op)); @@ -191,7 +191,7 @@ class TopologicalSorter : public ExprMutator { return output; } - BindingBlock VisitBindingBlock_(const DataflowBlockNode* op) override { + BindingBlock VisitBindingBlock_(const DataflowBlockObj* op) override { auto block = ffi::GetRef(op); // A map from not-yet-defined variables to the binding that will diff --git a/src/relax/transform/update_param_type.cc b/src/relax/transform/update_param_type.cc index 19e69c002007..0c880afc09f1 100644 --- a/src/relax/transform/update_param_type.cc +++ b/src/relax/transform/update_param_type.cc @@ -46,7 +46,7 @@ class ParamTypeMutator : public ExprMutator { using ExprMutator::VisitExpr_; using ExprMutator::VisitVarDef_; - Expr VisitExpr_(const FunctionNode* op) override { + Expr VisitExpr_(const FunctionObj* op) override { auto func = ffi::GetRef(op); auto params = op->params.Map([this](Var param) { diff --git a/src/relax/transform/update_vdevice.cc b/src/relax/transform/update_vdevice.cc index 2577c0e97d9a..6a57d6eb7b35 100644 --- a/src/relax/transform/update_vdevice.cc +++ b/src/relax/transform/update_vdevice.cc @@ -44,7 +44,7 @@ class VDeviceMutator : public ExprMutator { Expr VisitExpr(const Expr& expr) final { auto visited_expr = ExprMutator::VisitExpr(expr); if (!visited_expr->ty.as().has_value()) { - auto* tinfo = GetTypeAs(visited_expr); + auto* tinfo = GetTypeAs(visited_expr); bool unchanged = true; if (tinfo != nullptr) { if (tinfo->vdevice.has_value()) { @@ -68,7 +68,7 @@ class VDeviceMutator : public ExprMutator { IRModule Run() { for (const auto& [gv, func] : mod_->functions) { - if (func->IsInstance()) { + if (func->IsInstance()) { relax::Function update_func = VisitExpr(func).as_or_throw(); builder_->UpdateFunction(gv, update_func); } diff --git a/src/relax/transform/utils.cc b/src/relax/transform/utils.cc index 869be9775ea3..0f289ce0e0b5 100644 --- a/src/relax/transform/utils.cc +++ b/src/relax/transform/utils.cc @@ -25,14 +25,14 @@ namespace tvm { namespace relax { bool IsScalarTensor(const Type& ty) { - if (!ty->IsInstance()) { + if (!ty->IsInstance()) { return false; } TensorType tensor_ty = ty.as_or_throw(); - if (!tensor_ty->shape.has_value() || !tensor_ty->shape.value()->IsInstance()) { + if (!tensor_ty->shape.has_value() || !tensor_ty->shape.value()->IsInstance()) { return false; } - return tensor_ty->shape.as()->values.size() == 0; + return tensor_ty->shape.as()->values.size() == 0; } bool IsScalarTensor(const Expr& expr) { return IsScalarTensor(GetType(expr)); } @@ -51,7 +51,7 @@ Function ComposeFunctions(Function func_a, Function func_b) { bindings.push_back(VarBinding(func_a_output, func_a->body)); auto func_a_outputs = [&]() -> ffi::Array { - if (auto func_a_output_tuple = func_a->ret_ty.as()) { + if (auto func_a_output_tuple = func_a->ret_ty.as()) { ffi::Array outputs; for (size_t i = 0; i < func_a_output_tuple->fields.size(); i++) { outputs.push_back(TupleGetItem(func_a_output, i)); @@ -62,7 +62,7 @@ Function ComposeFunctions(Function func_a, Function func_b) { } }(); - if (func_b->params.size() == 1 && func_b->params[0]->ty.as()) { + if (func_b->params.size() == 1 && func_b->params[0]->ty.as()) { // Special case where the output of the first function is a tuple // that should be provided as-is to the second function, and // should not be unpacked into individual elements. diff --git a/src/relax/transform/utils.h b/src/relax/transform/utils.h index 4a0246144bc8..7dced90671fd 100644 --- a/src/relax/transform/utils.h +++ b/src/relax/transform/utils.h @@ -86,16 +86,16 @@ class MemoizedExprTranslator : public ExprFunctor { return res; } - virtual OutputType VisitExpr_(const VarNode* vn) { + virtual OutputType VisitExpr_(const VarObj* vn) { TVM_FFI_ICHECK(memo_.count(ffi::GetRef(vn))); return memo_.at(ffi::GetRef(vn)); } - virtual OutputType VisitExpr_(const DataflowVarNode* vn) { - return VisitExpr_(static_cast(vn)); + virtual OutputType VisitExpr_(const DataflowVarObj* vn) { + return VisitExpr_(static_cast(vn)); } - virtual OutputType VisitBinding_(const VarBindingNode* binding) { + virtual OutputType VisitBinding_(const VarBindingObj* binding) { TVM_FFI_ICHECK_EQ(memo_.count(binding->var), 0); auto v = VisitExpr(binding->value); memo_.insert_or_assign(binding->var, v); @@ -178,9 +178,9 @@ bool IsScalarTensor(const Expr& expr); */ template bool IsNestedTensorConditioned(const Type& ty, FType f_condition) { - if (const auto* tensor_ty = ty.as()) { + if (const auto* tensor_ty = ty.as()) { return f_condition(ffi::GetRef(tensor_ty)); - } else if (const auto* tuple_ty = ty.as()) { + } else if (const auto* tuple_ty = ty.as()) { return !std::any_of(tuple_ty->fields.begin(), tuple_ty->fields.end(), [&](const Type& field) { return !IsNestedTensorConditioned(field, f_condition); }); @@ -215,7 +215,7 @@ class VarReplacer : public ExprMutator { } private: - Expr VisitExpr_(const VarNode* op) final { + Expr VisitExpr_(const VarObj* op) final { Var var = ffi::GetRef(op); auto it = var_remap_.find(var); return it == var_remap_.end() ? var : it->second; @@ -244,28 +244,28 @@ class SymbolicVarRenewMutator : public ExprMutator { using relax::ExprMutator::VisitExpr; using relax::ExprMutator::VisitExpr_; - static Var CopyVar(const VarNode* op, Type ty) { + static Var CopyVar(const VarObj* op, Type ty) { ffi::Optional ty_annotation = ty.as().has_value() ? std::nullopt : ffi::Optional(ty); - if (op->IsInstance()) { + if (op->IsInstance()) { return DataflowVar(op->name, std::move(ty_annotation), op->loc); } - return Var(op->name, std::move(ty_annotation), op->loc); + return Var(op->name, std::move(ty), op->loc); } - Type RenewType(const VarNode* op) { + Type RenewType(const VarObj* op) { return op->ty.as().has_value() ? op->ty : this->VisitExprDepTypeField(op->ty); } - Var RenewVarDefinition(const VarNode* op) { + Var RenewVarDefinition(const VarObj* op) { Var old_var = ffi::GetRef(op); if (auto it = var_remap_.find(old_var); it != var_remap_.end()) { return it->second; } Type new_ty = RenewType(op); - bool is_dataflow = op->IsInstance(); - bool renew = renew_all_var_definitions_ || (!is_dataflow && op->ty.as()) || + bool is_dataflow = op->IsInstance(); + bool renew = renew_all_var_definitions_ || (!is_dataflow && op->ty.as()) || !new_ty.same_as(op->ty); if (!renew) { return old_var; @@ -276,12 +276,12 @@ class SymbolicVarRenewMutator : public ExprMutator { return renewed; } - Expr VisitExpr_(const VarNode* op) final { + Expr VisitExpr_(const VarObj* op) final { Var old_var = ffi::GetRef(op); if (auto it = var_remap_.find(old_var); it != var_remap_.end()) { return it->second; } - if (!op->ty.as()) { + if (!op->ty.as()) { return ExprMutator::VisitExpr_(op); } @@ -290,7 +290,7 @@ class SymbolicVarRenewMutator : public ExprMutator { return renewed; } - Expr VisitExpr_(const DataflowVarNode* op) final { + Expr VisitExpr_(const DataflowVarObj* op) final { Var old_var = ffi::GetRef(op); if (auto it = var_remap_.find(old_var); it != var_remap_.end()) { return it->second; @@ -304,11 +304,11 @@ class SymbolicVarRenewMutator : public ExprMutator { return renewed; } - Var VisitVarDef_(const VarNode* op) override { return RenewVarDefinition(op); } + Var VisitVarDef_(const VarObj* op) override { return RenewVarDefinition(op); } - Var VisitVarDef_(const DataflowVarNode* op) override { return RenewVarDefinition(op); } + Var VisitVarDef_(const DataflowVarObj* op) override { return RenewVarDefinition(op); } - Expr VisitExpr_(const FunctionNode* op) final { + Expr VisitExpr_(const FunctionObj* op) final { tvm::ffi::Array params; bool all_params_unchanged = true; for (Var param : op->params) { diff --git a/src/relax/utils.cc b/src/relax/utils.cc index e0cf657da4ca..a4101c8266e2 100644 --- a/src/relax/utils.cc +++ b/src/relax/utils.cc @@ -42,13 +42,13 @@ class ExprBinder : public ExprMutator { private: using ExprMutator::VisitExpr_; - Expr VisitExpr_(const ShapeExprNode* op) final { + Expr VisitExpr_(const ShapeExprObj* op) final { ffi::Array values = op->values.Map([this](const PrimExpr& value) { return BindShapeValue(value); }); return values.same_as(op->values) ? ffi::GetRef(op) : ShapeExpr(values, op->loc); } - Expr VisitExpr_(const FunctionNode* op) final { + Expr VisitExpr_(const FunctionObj* op) final { tvm::ffi::Array params; bool all_params_unchanged = true; for (const Var& param : op->params) { @@ -75,7 +75,7 @@ class ExprBinder : public ExprMutator { } } - Expr VisitExpr_(const VarNode* op) final { + Expr VisitExpr_(const VarObj* op) final { auto id = ffi::GetRef(op); auto it = bindings_.find(id); if (it != bindings_.end()) { @@ -149,11 +149,11 @@ tvm::ffi::Map InferSymbolicVarMap( }; auto bind_from_shape = [&bind_from_prim_expr](const Type& var, const Type& expr) { - auto var_shape = var.as(); + auto var_shape = var.as(); if (!var_shape) return; if (!var_shape->values.has_value()) return; - auto expr_shape = expr.as(); + auto expr_shape = expr.as(); if (!expr_shape) return; if (!expr_shape->values.has_value()) return; @@ -166,11 +166,11 @@ tvm::ffi::Map InferSymbolicVarMap( }; auto bind_from_tensor = [&bind_from_shape](const Type& var, const Type& expr) { - auto var_tensor = var.as(); + auto var_tensor = var.as(); if (!var_tensor) return; if (!var_tensor->shape.has_value()) return; - auto expr_tensor = expr.as(); + auto expr_tensor = expr.as(); if (!expr_tensor) return; if (!expr_tensor->shape.has_value()) return; @@ -179,10 +179,10 @@ tvm::ffi::Map InferSymbolicVarMap( std::function bind_from_ty = nullptr; auto bind_from_tuple = [&bind_from_ty](const Type& var, const Type& expr) { - auto var_tuple = var.as(); + auto var_tuple = var.as(); if (!var_tuple) return; - auto expr_tuple = expr.as(); + auto expr_tuple = expr.as(); if (!expr_tuple) return; if (var_tuple->fields.size() != expr_tuple->fields.size()) return; @@ -211,14 +211,14 @@ bool IsBoolType(const Type& ty, bool permit_unknown_rank, bool permit_unknown_dt DLDataType dtype; int ndim; - if (const auto* tensor = ty.as()) { + if (const auto* tensor = ty.as()) { ndim = tensor->ndim; if (tensor->IsUnknownDtype()) { bool correct_rank = ndim == 0 || (permit_unknown_rank && ndim == -1); return permit_unknown_dtype && correct_rank; } dtype = tensor->dtype.value()->dtype; - } else if (const auto* prim = ty.as()) { + } else if (const auto* prim = ty.as()) { dtype = prim->dtype; ndim = 0; } else { @@ -233,12 +233,12 @@ bool IsBoolType(const Type& ty, bool permit_unknown_rank, bool permit_unknown_dt } bool IsLeafOrTuple(const Expr& expr) { - return !expr.as() && !expr.as() && !expr.as() && - !expr.as() && !expr.as(); + return !expr.as() && !expr.as() && !expr.as() && + !expr.as() && !expr.as(); } bool IsImpureCall(const Call& call) { - if (auto op_ptr = call->op.as()) { + if (auto op_ptr = call->op.as()) { auto op = ffi::GetRef(op_ptr); static auto purity_map = Op::GetAttrMap(tvm::relax::op_attr::kPurity); if (purity_map.count(op)) { @@ -250,18 +250,18 @@ bool IsImpureCall(const Call& call) { auto effect = effect_map[op]; return effect > CallEffectKind::kPure; } - TVM_FFI_CHECK(!call->op->ty.as(), TypeError) + TVM_FFI_CHECK(!call->op->ty.as(), TypeError) << "Ordinary Relax calls cannot invoke a native TIRx function; " << "use R.call_tir or R.call_tir_packed"; // the Type must be FuncType - auto func_ty = GetTypeAs(call->op); + auto func_ty = GetTypeAs(call->op); return !func_ty->purity; } Expr GetBoundValue(const Binding& b) { - if (auto* var_binding = b.as()) { + if (auto* var_binding = b.as()) { return var_binding->value; - } else if (auto* match_binding = b.as()) { + } else if (auto* match_binding = b.as()) { return match_binding->value; } else { TVM_FFI_ICHECK(false) << "Invalid binding (should never happen)"; diff --git a/src/runtime/extra/contrib/tensorrt/tensorrt_ops.h b/src/runtime/extra/contrib/tensorrt/tensorrt_ops.h index 5e4c30ed7f30..c1cfa80f5fc4 100644 --- a/src/runtime/extra/contrib/tensorrt/tensorrt_ops.h +++ b/src/runtime/extra/contrib/tensorrt/tensorrt_ops.h @@ -130,7 +130,7 @@ class TensorRTOpConverter { /*! * \brief Convert to TRT. Implementation should use inputs and attributes - * from the CallNode to add the corresponding TRT layers to network. Outputs + * from the CallObj to add the corresponding TRT layers to network. Outputs * should be pushed to outputs vector. * \param params Parameters for this op. */ diff --git a/src/runtime/rpc/rpc_local_session.cc b/src/runtime/rpc/rpc_local_session.cc index 338bcd1e9b19..ebd23f697b36 100644 --- a/src/runtime/rpc/rpc_local_session.cc +++ b/src/runtime/rpc/rpc_local_session.cc @@ -38,7 +38,7 @@ RPCSession::PackedFuncHandle LocalSession::GetFunction(const std::string& name) if (auto fp = tvm::ffi::Function::GetGlobal(name)) { // return raw handle because the remote need to explicitly manage it. Any ret = *fp; - TVMFFIAny ret_any = ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(std::move(ret)); + TVMFFIAny ret_any = ffi::detail::AnyUnsafe::MoveAnyToTVMFFIAny(std::move(ret)); return ret_any.v_obj; } else { return nullptr; @@ -58,7 +58,7 @@ void LocalSession::EncodeReturn(ffi::Any rv, const FEncodeReturn& encode_return) // We follow a special protocol to return Tensor to client side // The first pack value is the Tensor handle as DLTensor // The second pack value is a customized deleter that deletes the Tensor. - TVMFFIAny ret_any = ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(std::move(rv)); + TVMFFIAny ret_any = ffi::detail::AnyUnsafe::MoveAnyToTVMFFIAny(std::move(rv)); void* opaque_handle = ret_any.v_obj; packed_args[1] = TVMFFITensorGetDLTensorPtr(opaque_handle); packed_args[2] = opaque_handle; @@ -75,7 +75,7 @@ void LocalSession::EncodeReturn(ffi::Any rv, const FEncodeReturn& encode_return) packed_args[1] = (*opt_str).data(); encode_return(ffi::PackedArgs(packed_args, 2)); } else if (rv.as()) { - TVMFFIAny ret_any = ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(std::move(rv)); + TVMFFIAny ret_any = ffi::detail::AnyUnsafe::MoveAnyToTVMFFIAny(std::move(rv)); void* opaque_handle = ret_any.v_obj; packed_args[1] = opaque_handle; encode_return(ffi::PackedArgs(packed_args, 2)); @@ -142,7 +142,7 @@ void LocalSession::CopyFromRemote(DLTensor* from, void* to_bytes, uint64_t nbyte void LocalSession::FreeHandle(void* handle) { // NOTE: the type code is no longer need during free handle. - ffi::details::ObjectUnsafe::DecRefObjectHandle(handle); + ffi::detail::ObjectUnsafe::DecRefObjectHandle(handle); } void LocalSession::SetDevice(Device dev) { DeviceAPI::Get(dev)->SetDevice(dev); } diff --git a/src/runtime/vm/cuda/cuda_graph_builtin.cc b/src/runtime/vm/cuda/cuda_graph_builtin.cc index 0671b93a80c6..bde35b8ce6bd 100644 --- a/src/runtime/vm/cuda/cuda_graph_builtin.cc +++ b/src/runtime/vm/cuda/cuda_graph_builtin.cc @@ -22,7 +22,7 @@ * \brief The CUDA graph related builtin functions for Relax virtual machine. */ -#include +#include #include #include #include @@ -57,7 +57,7 @@ struct CUDAGraphCaptureKeyHash { std::hash hash_fn; size_t hash = hash_fn(key.index); for (const auto& shape : key.shape_expr) { - ffi::details::StableHashCombine(hash, hash_fn(shape)); + ffi::detail::StableHashCombine(hash, hash_fn(shape)); } return hash; } diff --git a/src/runtime/vm/module_utils.h b/src/runtime/vm/module_utils.h index 2897cdf32764..c931b85d12b1 100644 --- a/src/runtime/vm/module_utils.h +++ b/src/runtime/vm/module_utils.h @@ -50,8 +50,8 @@ struct ModuleVTableEntryHelper { using MemFnType = R (T::*)(Args...) const; TVM_FFI_INLINE static void Call(ffi::Any* rv, T* self, MemFnType f, ffi::PackedArgs args) { auto wrapped = [self, f](Args... args) -> R { return (self->*f)(std::forward(args)...); }; - ffi::details::unpack_call(std::make_index_sequence{}, nullptr, wrapped, - args.data(), args.size(), rv); + ffi::detail::unpack_call(std::make_index_sequence{}, nullptr, wrapped, + args.data(), args.size(), rv); } }; @@ -60,8 +60,8 @@ struct ModuleVTableEntryHelper { using MemFnType = R (T::*)(Args...); TVM_FFI_INLINE static void Call(ffi::Any* rv, T* self, MemFnType f, ffi::PackedArgs args) { auto wrapped = [self, f](Args... args) -> R { return (self->*f)(std::forward(args)...); }; - ffi::details::unpack_call(std::make_index_sequence{}, nullptr, wrapped, - args.data(), args.size(), rv); + ffi::detail::unpack_call(std::make_index_sequence{}, nullptr, wrapped, + args.data(), args.size(), rv); } }; @@ -70,8 +70,8 @@ struct ModuleVTableEntryHelper { using MemFnType = void (T::*)(Args...) const; TVM_FFI_INLINE static void Call(ffi::Any* rv, T* self, MemFnType f, ffi::PackedArgs args) { auto wrapped = [self, f](Args... args) -> void { (self->*f)(std::forward(args)...); }; - ffi::details::unpack_call(std::make_index_sequence{}, nullptr, wrapped, - args.data(), args.size(), rv); + ffi::detail::unpack_call(std::make_index_sequence{}, nullptr, wrapped, + args.data(), args.size(), rv); } }; @@ -80,8 +80,8 @@ struct ModuleVTableEntryHelper { using MemFnType = void (T::*)(Args...); TVM_FFI_INLINE static void Call(ffi::Any* rv, T* self, MemFnType f, ffi::PackedArgs args) { auto wrapped = [self, f](Args... args) -> void { (self->*f)(std::forward(args)...); }; - ffi::details::unpack_call(std::make_index_sequence{}, nullptr, wrapped, - args.data(), args.size(), rv); + ffi::detail::unpack_call(std::make_index_sequence{}, nullptr, wrapped, + args.data(), args.size(), rv); } }; diff --git a/src/s_tir/analysis/calculate_allocated_memory.cc b/src/s_tir/analysis/calculate_allocated_memory.cc index 81f9bf716a37..2a8dfdc38808 100644 --- a/src/s_tir/analysis/calculate_allocated_memory.cc +++ b/src/s_tir/analysis/calculate_allocated_memory.cc @@ -45,7 +45,7 @@ namespace s_tir { using namespace tvm::tirx; runtime::MemoryScope GetStorageScope(const Var& var) { - auto* ptr = var->ty.as(); + auto* ptr = var->ty.as(); TVM_FFI_ICHECK(ptr) << "Buffer Var's type annotation must be of PtrType"; return ptr->scope; } @@ -67,15 +67,15 @@ class AllocTensorCalculator : public StmtExprVisitor { } private: - ffi::Optional Visit_(const BindNode* op) final { - if (const auto* call = op->value.as(); + ffi::Optional Visit_(const BindObj* op) final { + if (const auto* call = op->value.as(); call && call->op.same_as(tirx::alloc_tensor_op())) { return DispatchAllocTensor(op, call); } return StmtExprVisitor::Visit_(op); } - ffi::Optional DispatchAllocTensor(const BindNode* op, const CallNode* call) { + ffi::Optional DispatchAllocTensor(const BindObj* op, const CallObj* call) { tvm::Tuple shape = tvm::Tuple(call->ty_args[0].as_or_throw()->shape); DLDataType dtype = call->ty_args[0].as_or_throw()->dtype->dtype; runtime::MemoryScope scope = call->ty_args[0].as_or_throw()->scope; @@ -86,7 +86,7 @@ class AllocTensorCalculator : public StmtExprVisitor { } int64_t size = 1; for (const Expr& e : shape->fields) { - if (auto* imm = e.as()) { + if (auto* imm = e.as()) { size = static_cast(size * imm->value); } else { size = 0; @@ -98,19 +98,19 @@ class AllocTensorCalculator : public StmtExprVisitor { _max_size[scope] = std::max(_current_size[scope], _max_size[scope]); return StmtExprVisitor::Visit_(op); } - ffi::Optional Visit_(const ForNode* op) override { + ffi::Optional Visit_(const ForObj* op) override { auto snapshot = _current_size; TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(op)); _current_size = snapshot; return std::nullopt; } - ffi::Optional Visit_(const IfNode* op) override { + ffi::Optional Visit_(const IfObj* op) override { auto snapshot = _current_size; TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(op)); _current_size = snapshot; return std::nullopt; } - ffi::Optional Visit_(const RegionStmtNode* op) override { + ffi::Optional Visit_(const RegionStmtObj* op) override { auto snapshot = _current_size; TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(op)); _current_size = snapshot; diff --git a/src/s_tir/analysis/check_well_formed.cc b/src/s_tir/analysis/check_well_formed.cc index ea533ba6e3c3..890b66763475 100644 --- a/src/s_tir/analysis/check_well_formed.cc +++ b/src/s_tir/analysis/check_well_formed.cc @@ -64,7 +64,7 @@ class BlockVarAccessVerifier : public StmtExprVisitor { return std::nullopt; } - ffi::Optional Visit_(const VarNode* op) final { + ffi::Optional Visit_(const VarObj* op) final { auto it = loop_vars_.find(op); if (it != loop_vars_.end() && it->second < block_stack_.size()) { has_error_ = true; @@ -88,7 +88,7 @@ class BlockVarAccessVerifier : public StmtExprVisitor { return StmtExprVisitor::Visit_(op); } - ffi::Optional Visit_(const ForNode* op) final { + ffi::Optional Visit_(const ForObj* op) final { TVM_FFI_ICHECK(loop_vars_.find(op->var.get()) == loop_vars_.end()); loop_vars_[op->var.get()] = block_stack_.size(); TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(op)); @@ -96,7 +96,7 @@ class BlockVarAccessVerifier : public StmtExprVisitor { return std::nullopt; } - ffi::Optional Visit_(const SBlockNode* op) final { + ffi::Optional Visit_(const SBlockObj* op) final { // Do not check boundary if it's a opaque block. bool is_non_opaque = op->iter_vars.size(); if (is_non_opaque) { @@ -139,11 +139,11 @@ class BlockVarAccessVerifier : public StmtExprVisitor { private: /*! \brief The map from outside loop vars to its corresponding block level. */ - std::unordered_map loop_vars_; + std::unordered_map loop_vars_; /*! \brief Whether it's in assert mode. */ bool assert_mode_; /*! \brief Current nested block stack level. */ - std::vector block_stack_; + std::vector block_stack_; /*! \brief Whether there is error. */ bool has_error_{false}; }; @@ -206,9 +206,9 @@ class WellFormednessVerifier : public StmtExprVisitor { static void InitVTable(VTable* table) { StmtExprVisitor::InitVTable(table); - StmtExprVisitor::SetDispatch(table); - StmtExprVisitor::SetDispatch(table); - StmtExprVisitor::SetDispatch(table); + StmtExprVisitor::SetDispatch(table); + StmtExprVisitor::SetDispatch(table); + StmtExprVisitor::SetDispatch(table); } Result Visit(ffi::AnyView value) override { @@ -234,8 +234,8 @@ class WellFormednessVerifier : public StmtExprVisitor { TVM_FFI_VISIT_END(object.value()); } if (value.as() && !table_->CanDispatch(value)) { - if (value.as()) return std::nullopt; - if (auto var = value.as()) { + if (value.as()) return std::nullopt; + if (auto var = value.as()) { TVM_FFI_VISIT_BEGIN(); return this->Visit_(var); TVM_FFI_VISIT_END(object.value()); @@ -363,7 +363,7 @@ class WellFormednessVerifier : public StmtExprVisitor { }; ffi::StructuralWalk(tensor.shape()[i], walk); } - if (const auto* tile = tensor.ty()->layout.as(); + if (const auto* tile = tensor.ty()->layout.as(); tile && MatchShapeStridePattern(tensor.ty()).has_value()) { auto site = field->Attr("layout"); for (size_t i = 0; i < tile->shard.size(); ++i) { @@ -381,7 +381,7 @@ class WellFormednessVerifier : public StmtExprVisitor { } } - Result Visit_(const IRModuleNode* op) { + Result Visit_(const IRModuleObj* op) { std::vector globals; for (const auto& [global, function] : op->functions) globals.push_back(global); std::sort(globals.begin(), globals.end(), [&](const GlobalVar& a, const GlobalVar& b) { @@ -395,13 +395,13 @@ class WellFormednessVerifier : public StmtExprVisitor { return std::nullopt; } - Result Visit_(const FunctionNode* op) { + Result Visit_(const FunctionObj* op) { int enclosing_depth = loop_depth_; loop_depth_ = 0; WithScope([&]() -> Result { auto params = AccessPath::Root()->Attr("params"); for (size_t i = 0; i < op->params.size(); ++i) { - if (!op->params[i]->ty.as()) Define(op->params[i], params->ArrayItem(i)); + if (!op->params[i]->ty.as()) Define(op->params[i], params->ArrayItem(i)); } for (size_t i = 0; i < op->params.size(); ++i) { if (auto tensor = op->params[i].as()) { @@ -421,7 +421,7 @@ class WellFormednessVerifier : public StmtExprVisitor { return std::nullopt; } - Result Visit_(const VarNode* op) override { + Result Visit_(const VarObj* op) override { auto var = ffi::GetRef(op); if (check_ == Check::kVariables) { auto verify = Verify(current_.count(var)); @@ -449,13 +449,13 @@ class WellFormednessVerifier : public StmtExprVisitor { } } - Result Visit_(const BindNode* op) override { + Result Visit_(const BindObj* op) override { this->Visit(op->value); Define(op->var, AccessPath::Root()->Attr("var")); return std::nullopt; } - Result Visit_(const RegionStmtNode* op) override { + Result Visit_(const RegionStmtObj* op) override { this->Visit(op->args); this->Visit(op->attrs); for (const auto& param : op->body_params) this->Visit(param->ty); @@ -471,7 +471,7 @@ class WellFormednessVerifier : public StmtExprVisitor { return std::nullopt; } - Result Visit_(const ForNode* op) override { + Result Visit_(const ForObj* op) override { if (check_ == Check::kThreadBinding) { if (auto tag = op->annotations.Get(attr::kThreadBinding.name)) { Verify(op->kind == ForKind::kParallel) @@ -497,7 +497,7 @@ class WellFormednessVerifier : public StmtExprVisitor { }); } - Result Visit_(const WhileNode* op) override { + Result Visit_(const WhileObj* op) override { this->Visit(op->condition); ++loop_depth_; WithScope([&]() { return this->Visit(op->body); }); @@ -505,13 +505,13 @@ class WellFormednessVerifier : public StmtExprVisitor { return std::nullopt; } - Result Visit_(const IfNode* op) override { + Result Visit_(const IfObj* op) override { this->Visit(op->condition); WithScope([&]() { return this->Visit(op->then_case); }); return WithScope([&]() { return this->Visit(op->else_case); }); } - Result Visit_(const prim::LetNode* op) override { + Result Visit_(const prim::LetObj* op) override { this->Visit(op->value); return WithScope([&]() -> Result { Define(op->var, AccessPath::Root()->Attr("var")); @@ -519,7 +519,7 @@ class WellFormednessVerifier : public StmtExprVisitor { }); } - Result Visit_(const LambdaExprNode* op) { + Result Visit_(const LambdaExprObj* op) { this->Visit(op->ty); return WithScope([&]() -> Result { for (size_t i = 0; i < op->vars.size(); ++i) { @@ -529,18 +529,18 @@ class WellFormednessVerifier : public StmtExprVisitor { }); } - Result Visit_(const CallNode* op) override { - if (op->op.as() || op->op.as()) this->Visit(op->op); + Result Visit_(const CallObj* op) override { + if (op->op.as() || op->op.as()) this->Visit(op->op); return this->Visit(op->args); } - Result Visit_(const TensorStoreNode* op) override { + Result Visit_(const TensorStoreObj* op) override { this->Visit(op->value); TensorUse(op->dest.as_or_throw()); return this->Visit(op->indices); } - Result Visit_(const TensorRegionNode* op) override { + Result Visit_(const TensorRegionObj* op) override { if (auto tensor = op->source.as()) { TensorUse(tensor.value()); } else { @@ -549,26 +549,26 @@ class WellFormednessVerifier : public StmtExprVisitor { return this->Visit(op->region); } - Result Visit_(const GenericConstNode* op) override { return std::nullopt; } + Result Visit_(const GenericConstObj* op) override { return std::nullopt; } - Result Visit_(const AssertStmtNode* op) override { + Result Visit_(const AssertStmtObj* op) override { this->Visit(op->condition); this->Visit(op->error_kind); return this->Visit(op->message_parts); } - Result Visit_(const prim::RampNode* op) override { + Result Visit_(const prim::RampObj* op) override { this->Visit(op->base); this->Visit(op->stride); return this->Visit(op->lanes); } - Result Visit_(const prim::BroadcastNode* op) override { + Result Visit_(const prim::BroadcastObj* op) override { this->Visit(op->value); return this->Visit(op->lanes); } - Result Visit_(const BreakNode* op) override { + Result Visit_(const BreakObj* op) override { if (check_ == Check::kLoopControl) { Verify(loop_depth_ > 0) << "ValueError: break at requires an enclosing " "loop in the same function."; @@ -576,7 +576,7 @@ class WellFormednessVerifier : public StmtExprVisitor { return std::nullopt; } - Result Visit_(const ContinueNode* op) override { + Result Visit_(const ContinueObj* op) override { if (check_ == Check::kLoopControl) { Verify(loop_depth_ > 0) << "ValueError: continue at requires an enclosing " "loop in the same function."; @@ -584,9 +584,9 @@ class WellFormednessVerifier : public StmtExprVisitor { return std::nullopt; } - Result Visit_(const TensorLoadNode* op) override; + Result Visit_(const TensorLoadObj* op) override; - Result Visit_(const SBlockNode* op) override { + Result Visit_(const SBlockObj* op) override { return WithScope([&]() -> Result { for (size_t i = 0; i < op->iter_vars.size(); ++i) { const auto& axis = op->iter_vars[i]; @@ -635,7 +635,7 @@ class WellFormednessVerifier : public StmtExprVisitor { std::unordered_map previous_; }; -auto WellFormednessVerifier::Visit_(const TensorLoadNode* op) -> Result { +auto WellFormednessVerifier::Visit_(const TensorLoadObj* op) -> Result { if (check_ != Check::kTensorLoadType) { TensorUse(op->source.as_or_throw()); return this->Visit(op->indices); diff --git a/src/s_tir/analysis/conditional_bounds.cc b/src/s_tir/analysis/conditional_bounds.cc index 1d134c6d3eee..f05fbf848c67 100644 --- a/src/s_tir/analysis/conditional_bounds.cc +++ b/src/s_tir/analysis/conditional_bounds.cc @@ -147,7 +147,13 @@ IntGroupBounds IntGroupBounds::operator+(const Range& r) { ffi::Optional IntGroupBounds::FindBestRange(const ffi::Map& vranges_addl) const { Analyzer analyzer; - analyzer->Bind(vranges_addl); + { + sym::ConstraintContext context; + for (const auto& [var, range] : vranges_addl) { + context->Append(sym::RangeConstraint(var, range)); + } + analyzer->SyncConstraintContext(context); + } ffi::Map var_intsets; for (auto kv : vranges_addl) { @@ -233,22 +239,22 @@ struct ExprLess { */ class NormalizeComparisons : public tvm::ExprMutator { public: - UnchangedOr Mutate_(const prim::EQNode* op, InplaceMode inplace_mode) override { + UnchangedOr Mutate_(const prim::EQObj* op, InplaceMode inplace_mode) override { return Make(op->a, op->b); } - UnchangedOr Mutate_(const prim::NENode* op, InplaceMode inplace_mode) override { + UnchangedOr Mutate_(const prim::NEObj* op, InplaceMode inplace_mode) override { return Make(op->a, op->b); } - UnchangedOr Mutate_(const prim::LTNode* op, InplaceMode inplace_mode) override { + UnchangedOr Mutate_(const prim::LTObj* op, InplaceMode inplace_mode) override { return Make(op->a, op->b); } - UnchangedOr Mutate_(const prim::LENode* op, InplaceMode inplace_mode) override { + UnchangedOr Mutate_(const prim::LEObj* op, InplaceMode inplace_mode) override { return Make(op->a, op->b); } - UnchangedOr Mutate_(const prim::GTNode* op, InplaceMode inplace_mode) override { + UnchangedOr Mutate_(const prim::GTObj* op, InplaceMode inplace_mode) override { return Make(op->b, op->a); } - UnchangedOr Mutate_(const prim::GENode* op, InplaceMode inplace_mode) override { + UnchangedOr Mutate_(const prim::GEObj* op, InplaceMode inplace_mode) override { return Make(op->b, op->a); } @@ -276,9 +282,9 @@ void AddInequality(std::vector* inequality_set, const PrimExpr& new_in // or has already been added return; } - if (const prim::LENode* new_le = new_ineq.as()) { + if (const prim::LEObj* new_le = new_ineq.as()) { for (auto iter = inequality_set->begin(); iter != inequality_set->end();) { - const prim::LENode* le = iter->as(); + const prim::LEObj* le = iter->as(); if (le && analyzer->CanProve(new_le->a - le->a <= 0)) { return; } else if (le && analyzer->CanProve(le->a - new_le->a <= 0)) { @@ -300,9 +306,9 @@ void ClassifyByPolarity(const PrimVar& var, const std::vector& current // Take formulas from current_ineq_set and classify them according to polarity wrt var // and store to coef_pos and coef_neg respectively. for (const PrimExpr& ineq : current_ineq_set) { - if (const prim::LENode* le = ineq.as()) { + if (const prim::LEObj* le = ineq.as()) { ffi::Array coef = tirx::DetectLinearEquation(le->a, {var}); - const auto* imm = !coef.empty() ? coef[0].as() : nullptr; + 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; if (coef0 == 0) { @@ -315,9 +321,9 @@ void ClassifyByPolarity(const PrimVar& var, const std::vector& current } continue; } - } else if (const prim::EQNode* eq = ineq.as()) { + } else if (const prim::EQObj* eq = ineq.as()) { ffi::Array coef = tirx::DetectLinearEquation(eq->a, {var}); - const auto* imm = !coef.empty() ? coef[0].as() : nullptr; + 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; if (coef0 == 0) { @@ -358,7 +364,13 @@ void MoveEquality(std::vector* upper_bounds, std::vector* lo PartialSolvedInequalities SolveLinearInequalities(const IntConstraints& system_to_solve) { sym::Analyzer analyzer; - analyzer->Bind(system_to_solve.ranges); + { + sym::ConstraintContext context; + for (const auto& [var, range] : system_to_solve.ranges) { + context->Append(sym::RangeConstraint(var, range)); + } + analyzer->SyncConstraintContext(context); + } // The algorithm consists in doing the following things for each variable v // - Take formulas from `current_ineq_set_to_solve` and @@ -555,7 +567,13 @@ IntConstraints SolveInequalitiesToRange(const IntConstraints& inequalities) { // This order is needed to compute new ranges. for (auto it = inequalities.variables.rbegin(); it != inequalities.variables.rend(); ++it) { sym::Analyzer analyzer; - analyzer->Bind(vranges); + { + sym::ConstraintContext context; + for (const auto& [var, range] : vranges) { + context->Append(sym::RangeConstraint(var, range)); + } + analyzer->SyncConstraintContext(context); + } const PrimVar& var = *it; TVM_FFI_ICHECK(solved_bounds.count(var)); @@ -591,7 +609,13 @@ IntConstraints SolveInequalitiesToRange(const IntConstraints& inequalities) { // Add the original conditions to the resulting conditions sym::Analyzer analyzer; - analyzer->Bind(vranges); + { + sym::ConstraintContext context; + for (const auto& [var, range] : vranges) { + context->Append(sym::RangeConstraint(var, range)); + } + analyzer->SyncConstraintContext(context); + } for (const PrimExpr& old_cond : AsConditions(inequalities.variables, solved_bounds, solved_other_relations)) { if (!analyzer->CanProve(old_cond)) { @@ -621,24 +645,24 @@ ffi::Optional> ConditionalBoundsContext::TrySolveCondition( ffi::Array equations; ffi::Array vars; std::function fvisit = [&equations, &vars, &fvisit](const PrimExpr& e) { - if (e->IsInstance() || e->IsInstance() || - e->IsInstance() || e->IsInstance() || - e->IsInstance() || e->IsInstance()) { + if (e->IsInstance() || e->IsInstance() || + e->IsInstance() || e->IsInstance() || + e->IsInstance() || e->IsInstance()) { bool is_simple = true; std::vector cand_vars; auto walk_fn = [&cand_vars, &is_simple, &e](const PrimExpr& obj) -> ffi::Expected { if (obj.same_as(e)) { return ffi::WalkResult::Advance(); - } else if (const VarNode* var = obj.as()) { + } else if (const VarObj* var = obj.as()) { PrimType var_ty = var->ty.as_or_throw(); if (var_ty.MatchesCode(DLDataTypeCode::kDLInt, DLDataTypeCode::kDLUInt)) { cand_vars.push_back(ffi::GetRef(var).as_or_throw()); } } else { - is_simple &= obj->IsInstance() || obj->IsInstance() || - obj->IsInstance() || obj->IsInstance() || - obj->IsInstance() || obj->IsInstance(); + is_simple &= obj->IsInstance() || obj->IsInstance() || + obj->IsInstance() || obj->IsInstance() || + obj->IsInstance() || obj->IsInstance(); } return ffi::WalkResult::Advance(); }; @@ -652,11 +676,11 @@ ffi::Optional> ConditionalBoundsContext::TrySolveCondition( } equations.push_back(e.as_or_throw()); } - } else if (e->IsInstance()) { + } else if (e->IsInstance()) { prim::And op = e.as_or_throw(); fvisit(op->a); fvisit(op->b); - } else if (e->IsInstance()) { + } else if (e->IsInstance()) { Call op = e.as_or_throw(); if (op->op.same_as(prim::likely_op())) { fvisit(op->args[0].as_or_throw()); diff --git a/src/s_tir/analysis/domain_touched.cc b/src/s_tir/analysis/domain_touched.cc index 1508032613f8..2acac8be3961 100644 --- a/src/s_tir/analysis/domain_touched.cc +++ b/src/s_tir/analysis/domain_touched.cc @@ -27,7 +27,7 @@ #include #include #include -#include +#include #include #include @@ -62,11 +62,11 @@ using TensorDomainAccess = std::tuple; } // namespace // Find Read region of the tensor in the stmt. -class TensorTouchedDomain final : public s_tir::StmtExprVisitorWithAnalyzer { +class TensorTouchedDomain final : public s_tir::StmtExprVisitorWithConstraintContext { public: - using s_tir::StmtExprVisitorWithAnalyzer::Visit_; + using s_tir::StmtExprVisitorWithConstraintContext::Visit_; - std::unordered_map& GetAccessedTensorRegions() { + std::unordered_map& GetAccessedTensorRegions() { return tensor_access_map_; } @@ -98,9 +98,9 @@ class TensorTouchedDomain final : public s_tir::StmtExprVisitorWithAnalyzer { } private: - using Parent = s_tir::StmtExprVisitorWithAnalyzer; + using Parent = s_tir::StmtExprVisitorWithConstraintContext; - ffi::Optional Visit_(const TensorLoadNode* op) final { + ffi::Optional Visit_(const TensorLoadObj* op) final { TensorVar tensor = op->source.as_or_throw(); // Record load-exclusive tensor access Touch(&std::get(tensor_access_map_[tensor.get()]).set, op->indices); @@ -109,7 +109,7 @@ class TensorTouchedDomain final : public s_tir::StmtExprVisitorWithAnalyzer { return Parent::Visit_(op); } - ffi::Optional Visit_(const TensorStoreNode* op) final { + ffi::Optional Visit_(const TensorStoreObj* op) final { // Record store-exclusive tensor access Touch(&std::get(tensor_access_map_[op->dest.as_or_throw().get()]).set, op->indices); @@ -125,15 +125,15 @@ class TensorTouchedDomain final : public s_tir::StmtExprVisitorWithAnalyzer { bounds->resize(args.size()); } for (size_t i = 0; i < args.size(); ++i) { - if (args[i].as()) { + if (args[i].as()) { (*bounds)[i].emplace_back(IntSet::Vector(args[i])); } else { - (*bounds)[i].emplace_back(analyzer_->int_set(args[i])); + (*bounds)[i].emplace_back(GetSyncedAnalyzer()->int_set(args[i])); } } } - std::unordered_map tensor_access_map_; + std::unordered_map tensor_access_map_; }; ffi::Array> DomainTouched(const Stmt& stmt, const TensorVar& tensor, @@ -149,7 +149,7 @@ ffi::Map> DomainTouchedAccessMap(const Fun auto tensor_access_map = visitor->GetAccessedTensorRegions(); ffi::Map> ret; for (auto& var : func->params) { - if (!var->ty.as()) { + if (!var->ty.as()) { continue; } TensorVar tensor = var.as_or_throw(); diff --git a/src/s_tir/analysis/estimate_flops.cc b/src/s_tir/analysis/estimate_flops.cc index e55ef673d9af..2fc7f7e0fa6a 100644 --- a/src/s_tir/analysis/estimate_flops.cc +++ b/src/s_tir/analysis/estimate_flops.cc @@ -102,24 +102,24 @@ class FlopEstimator : private tvm::ExprFunctor, result.Add(op->ty.as_or_throw()->dtype); \ return result; \ } - TVM_TIR_ESTIMATE_FLOP_VISIT_BINARY(prim::AddNode); - TVM_TIR_ESTIMATE_FLOP_VISIT_BINARY(prim::SubNode); - TVM_TIR_ESTIMATE_FLOP_VISIT_BINARY(prim::MulNode); - TVM_TIR_ESTIMATE_FLOP_VISIT_BINARY(prim::DivNode); - TVM_TIR_ESTIMATE_FLOP_VISIT_BINARY(prim::ModNode); - TVM_TIR_ESTIMATE_FLOP_VISIT_BINARY(prim::FloorDivNode); - TVM_TIR_ESTIMATE_FLOP_VISIT_BINARY(prim::FloorModNode); - TVM_TIR_ESTIMATE_FLOP_VISIT_BINARY(prim::MinNode); - TVM_TIR_ESTIMATE_FLOP_VISIT_BINARY(prim::MaxNode); + TVM_TIR_ESTIMATE_FLOP_VISIT_BINARY(prim::AddObj); + TVM_TIR_ESTIMATE_FLOP_VISIT_BINARY(prim::SubObj); + TVM_TIR_ESTIMATE_FLOP_VISIT_BINARY(prim::MulObj); + TVM_TIR_ESTIMATE_FLOP_VISIT_BINARY(prim::DivObj); + TVM_TIR_ESTIMATE_FLOP_VISIT_BINARY(prim::ModObj); + TVM_TIR_ESTIMATE_FLOP_VISIT_BINARY(prim::FloorDivObj); + TVM_TIR_ESTIMATE_FLOP_VISIT_BINARY(prim::FloorModObj); + TVM_TIR_ESTIMATE_FLOP_VISIT_BINARY(prim::MinObj); + TVM_TIR_ESTIMATE_FLOP_VISIT_BINARY(prim::MaxObj); #undef TVM_TIR_ESTIMATE_FLOP_VISIT_BINARY - TResult Dispatch_(const prim::EQNode* op) override { return TResult(); } - TResult Dispatch_(const prim::NENode* op) override { return TResult(); } - TResult Dispatch_(const prim::LTNode* op) override { return TResult(); } - TResult Dispatch_(const prim::LENode* op) override { return TResult(); } - TResult Dispatch_(const prim::GTNode* op) override { return TResult(); } - TResult Dispatch_(const prim::GENode* op) override { return TResult(); } + TResult Dispatch_(const prim::EQObj* op) override { return TResult(); } + TResult Dispatch_(const prim::NEObj* op) override { return TResult(); } + TResult Dispatch_(const prim::LTObj* op) override { return TResult(); } + TResult Dispatch_(const prim::LEObj* op) override { return TResult(); } + TResult Dispatch_(const prim::GTObj* op) override { return TResult(); } + TResult Dispatch_(const prim::GEObj* op) override { return TResult(); } - int64_t GetLoopExtent(const ForNode* node, const sym::Analyzer& ana) { + int64_t GetLoopExtent(const ForObj* node, const sym::Analyzer& ana) { int64_t bound = ana->const_int_bound(Range::GetExtentFromBeginEnd(node->begin, node->end))->max_value; if (bound == sym::ConstIntBound::kPosInf) { @@ -129,55 +129,53 @@ class FlopEstimator : private tvm::ExprFunctor, } } - TResult Dispatch_(const prim::LShiftNode* op) final { + TResult Dispatch_(const prim::LShiftObj* op) final { TResult result = Dispatch(op->a); result += Dispatch(op->b); return result; } - TResult Dispatch_(const prim::RShiftNode* op) final { + TResult Dispatch_(const prim::RShiftObj* op) final { TResult result = Dispatch(op->a); result += Dispatch(op->b); return result; } - TResult Dispatch_(const prim::BitwiseAndNode* op) final { + TResult Dispatch_(const prim::BitwiseAndObj* op) final { TResult result = Dispatch(op->a); result += Dispatch(op->b); return result; } - TResult Dispatch_(const prim::BitwiseOrNode* op) final { + TResult Dispatch_(const prim::BitwiseOrObj* op) final { TResult result = Dispatch(op->a); result += Dispatch(op->b); return result; } - TResult Dispatch_(const prim::BitwiseXorNode* op) final { + TResult Dispatch_(const prim::BitwiseXorObj* op) final { TResult result = Dispatch(op->a); result += Dispatch(op->b); return result; } - TResult Dispatch_(const prim::BitwiseNotNode* op) final { return Dispatch(op->a); } - TResult Dispatch_(const prim::NotNode* op) override { return Dispatch(op->a); } - TResult Dispatch_(const prim::AndNode* op) final { + TResult Dispatch_(const prim::BitwiseNotObj* op) final { return Dispatch(op->a); } + TResult Dispatch_(const prim::NotObj* op) override { return Dispatch(op->a); } + TResult Dispatch_(const prim::AndObj* op) final { TResult result = Dispatch(op->a); result += Dispatch(op->b); return result; } - TResult Dispatch_(const prim::OrNode* op) final { + TResult Dispatch_(const prim::OrObj* op) final { TResult result = Dispatch(op->a); result += Dispatch(op->b); return result; } - TResult Dispatch_(const TensorLoadNode* op) override { return TResult(); } - TResult Dispatch_(const RegionStmtNode* op) override { + TResult Dispatch_(const TensorLoadObj* op) override { return TResult(); } + TResult Dispatch_(const RegionStmtObj* op) override { TResult result = Dispatch(op->body); for (const Expr& arg : op->args) result += Dispatch(arg); return result; } - TResult Dispatch_(const TensorStoreNode* store) override { return Dispatch(store->value); } - TResult Dispatch_(const SBlockRealizeNode* block) override { - return Dispatch(block->block->body); - } - TResult Dispatch_(const SBlockNode* block) override { + TResult Dispatch_(const TensorStoreObj* store) override { return Dispatch(store->value); } + TResult Dispatch_(const SBlockRealizeObj* block) override { return Dispatch(block->block->body); } + TResult Dispatch_(const SBlockObj* block) override { TResult result; if (block->init.has_value()) { result += Dispatch(block->init.value()); @@ -185,15 +183,19 @@ class FlopEstimator : private tvm::ExprFunctor, result += Dispatch(block->body); return result; } - TResult Dispatch_(const ForNode* loop) override { - ana->Bind(loop->var, Range(loop->begin, loop->end)); + TResult Dispatch_(const ForObj* loop) override { + auto outer_context = ana->GetConstraintContext(); + auto context = outer_context->Snapshot(); + context->Append(sym::RangeConstraint(loop->var, Range(loop->begin, loop->end))); + ana->SyncConstraintContext(context); const auto int_imm = GetLoopExtent(loop, ana); TResult result = Dispatch(loop->body); + ana->SyncConstraintContext(outer_context); result *= int_imm; return result; } - TResult Dispatch_(const IfNode* branch) override { + TResult Dispatch_(const IfObj* branch) override { TResult cond = Dispatch(branch->condition); if (branch->else_case) { cond += Dispatch(branch->then_case).MaxWith(Dispatch(branch->else_case.value())); @@ -203,37 +205,37 @@ class FlopEstimator : private tvm::ExprFunctor, return cond; } - TResult Dispatch_(const WhileNode* op) override { + TResult Dispatch_(const WhileObj* op) override { // TODO(jikechao): Improve while loop FLOP estimation with loop bound analysis TResult result = Dispatch(op->condition); result += Dispatch(op->body); return result; } - TResult Dispatch_(const BindNode* let) override { + TResult Dispatch_(const BindObj* let) override { if (auto value = let->value.as()) return Dispatch(value.value()); return TResult(); } - TResult Dispatch_(const prim::SelectNode* op) override { + TResult Dispatch_(const prim::SelectObj* op) override { TResult cond = Dispatch(op->condition); cond += Dispatch(op->true_value).MaxWith(Dispatch(op->false_value)); return cond; } - TResult Dispatch_(const AssertStmtNode* op) override { + TResult Dispatch_(const AssertStmtObj* op) override { TResult result = Dispatch(op->condition); return result; } - TResult Dispatch_(const VarNode* op) override { return TResult(); } - TResult Dispatch_(const IntImmNode* op) override { return TResult(); } - TResult Dispatch_(const FloatImmNode* op) override { return TResult(); } - TResult Dispatch_(const StringImmNode* op) override { return TResult(); } - TResult Dispatch_(const prim::CastNode* op) override { return Dispatch(op->value); } - TResult Dispatch_(const EvaluateNode* op) override { return TResult(); } + TResult Dispatch_(const VarObj* op) override { return TResult(); } + TResult Dispatch_(const IntImmObj* op) override { return TResult(); } + TResult Dispatch_(const FloatImmObj* op) override { return TResult(); } + TResult Dispatch_(const StringImmObj* op) override { return TResult(); } + TResult Dispatch_(const prim::CastObj* op) override { return Dispatch(op->value); } + TResult Dispatch_(const EvaluateObj* op) override { return TResult(); } - TResult Dispatch_(const SeqStmtNode* seq) override { + TResult Dispatch_(const SeqStmtObj* seq) override { TResult result; for (const Stmt& stmt : seq->seq) { result += Dispatch(stmt); @@ -241,7 +243,7 @@ class FlopEstimator : private tvm::ExprFunctor, return result; } - TResult Dispatch_(const CallNode* op) override { + TResult Dispatch_(const CallObj* op) override { TResult ret; for (const Expr& arg : op->args) { ret += Dispatch(arg); @@ -268,7 +270,7 @@ double EstimateTIRFlops(const IRModule& mod) { TResult result; double cached_result = 0; for (const auto& [_, base_func] : mod->functions) { - if (const auto* func = base_func.as()) { + if (const auto* func = base_func.as()) { if (auto cached = func->attrs.GetAttr(tvm::s_tir::attr::kEstimatedFlops)) { cached_result += cached.value(); } else if (func->body.has_value()) { diff --git a/src/s_tir/analysis/find_anchor_sblock.cc b/src/s_tir/analysis/find_anchor_sblock.cc index c0e75647fcb5..941ea7725cad 100644 --- a/src/s_tir/analysis/find_anchor_sblock.cc +++ b/src/s_tir/analysis/find_anchor_sblock.cc @@ -32,43 +32,43 @@ namespace tvm { namespace tirx { -Stmt GetEnclosingLoop(const s_tir::SBlockNode* block, Stmt func_body) { +Stmt GetEnclosingLoop(const s_tir::SBlockObj* block, Stmt func_body) { struct GetRootSeqStmt : public s_tir::StmtExprVisitor { using s_tir::StmtExprVisitor::Visit_; ffi::Optional Visit(ffi::AnyView value) override { - if (value.as()) return std::nullopt; + if (value.as()) return std::nullopt; return s_tir::StmtExprVisitor::Visit(value); } - ffi::Optional Visit_(const SeqStmtNode* seq) override { + ffi::Optional Visit_(const SeqStmtObj* seq) override { if (seq->size() == 1) { return Visit(seq->seq[0]); } result = seq; return std::nullopt; } - const SeqStmtNode* result = nullptr; + const SeqStmtObj* result = nullptr; }; struct BlockFinder : public s_tir::StmtExprVisitor { using s_tir::StmtExprVisitor::Visit_; ffi::Optional Visit(ffi::AnyView value) override { - if (value.as()) return std::nullopt; + if (value.as()) return std::nullopt; return s_tir::StmtExprVisitor::Visit(value); } - explicit BlockFinder(const s_tir::SBlockNode* tgt) : target(tgt) {} + explicit BlockFinder(const s_tir::SBlockObj* tgt) : target(tgt) {} - ffi::Optional Visit_(const s_tir::SBlockNode* block) override { + ffi::Optional Visit_(const s_tir::SBlockObj* block) override { if (block == target) { found = true; } return std::nullopt; } - const s_tir::SBlockNode* target; + const s_tir::SBlockObj* target; bool found = false; }; @@ -78,7 +78,7 @@ Stmt GetEnclosingLoop(const s_tir::SBlockNode* block, Stmt func_body) { TVM_FFI_ICHECK(seq_finder->result); for (auto stmt : seq_finder->result->seq) { - if (stmt->IsInstance()) { + if (stmt->IsInstance()) { auto finder = ffi::make_object(block); finder->Visit(stmt); if (finder->found) { @@ -92,22 +92,22 @@ Stmt GetEnclosingLoop(const s_tir::SBlockNode* block, Stmt func_body) { TVM_FFI_UNREACHABLE(); } -const s_tir::SBlockNode* FindAnchorBlock(const IRModule& mod) { +const s_tir::SBlockObj* FindAnchorBlock(const IRModule& mod) { struct ReductionSBlockCollector : public s_tir::StmtExprVisitor { using s_tir::StmtExprVisitor::Visit_; ffi::Optional Visit(ffi::AnyView value) override { - if (value.as()) return std::nullopt; + if (value.as()) return std::nullopt; return s_tir::StmtExprVisitor::Visit(value); } - ffi::Optional Visit_(const s_tir::SBlockNode* block) override { + ffi::Optional Visit_(const s_tir::SBlockObj* block) override { if (block->init) { blocks.push_back(block); } return s_tir::StmtExprVisitor::Visit(block->body); } - std::vector blocks; + std::vector blocks; }; if (auto function = FindEntryFunc(mod, nullptr)) { diff --git a/src/s_tir/analysis/identify_memcpy.cc b/src/s_tir/analysis/identify_memcpy.cc index 433463d43a9c..1bd6d6c60378 100644 --- a/src/s_tir/analysis/identify_memcpy.cc +++ b/src/s_tir/analysis/identify_memcpy.cc @@ -27,7 +27,7 @@ #include #include #include -#include +#include #include #include #include @@ -52,7 +52,7 @@ std::variant IdentifyMemCpyImpl(const For& loop, // Walk through the loop nest, stopping at the first loop whose body // is not a loop. Stmt stmt = loop; - while (auto* for_node = stmt.as()) { + while (auto* for_node = stmt.as()) { loop_ranges.Set(for_node->var, Range(for_node->begin, for_node->end)); loop_intervals.Set(for_node->var, sym::IntSet::FromMinExtent( for_node->begin, Range::GetExtentFromBeginEnd( @@ -220,16 +220,16 @@ std::variant IdentifyMemCpyImpl(const For& loop, auto make_comparison_tuple = [](const sym::IterSplitExpr& expr) { auto as_int_or_zero = [](auto& val) -> ffi::BigInt { - if (auto* as_int = val.template as()) { + if (auto* as_int = val.template as()) { return as_int->value; } else { return 0; } }; return std::tuple{ - static_cast(expr->scale.as()), as_int_or_zero(expr->scale), - static_cast(expr->extent.as()), as_int_or_zero(expr->lower_factor), - static_cast(expr->lower_factor.as()), as_int_or_zero(expr->lower_factor), + static_cast(expr->scale.as()), as_int_or_zero(expr->scale), + static_cast(expr->extent.as()), as_int_or_zero(expr->lower_factor), + static_cast(expr->lower_factor.as()), as_int_or_zero(expr->lower_factor), }; }; auto sorting_function = [&make_comparison_tuple](const sym::IterSplitExpr& lhs, @@ -304,17 +304,17 @@ TVM_FFI_STATIC_INIT_BLOCK() { refl::GlobalDef().def("s_tir.analysis._identify_memcpy", [](const Stmt& stmt) { ffi::Array output; - struct Visitor : s_tir::StmtExprVisitorWithAnalyzer { + struct Visitor : s_tir::StmtExprVisitorWithConstraintContext { public: - using s_tir::StmtExprVisitorWithAnalyzer::Visit_; + using s_tir::StmtExprVisitorWithConstraintContext::Visit_; explicit Visitor(ffi::Array* output) : output(output) {} ffi::Array* output; private: - ffi::Optional Visit_(const ForNode* op) override { + ffi::Optional Visit_(const ForObj* op) override { For loop = ffi::GetRef(op); - auto result = IdentifyMemCpyImpl(loop, Visitor::analyzer_.get()); + auto result = IdentifyMemCpyImpl(loop, GetSyncedAnalyzer()); if (auto* ptr = std::get_if(&result)) { output->push_back(ffi::Array{ptr->source, ptr->dest}); } else if (auto* ptr = std::get_if(&result)) { @@ -323,7 +323,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { TVM_FFI_THROW(InternalError) << "Internal error, unhandled std::variant type"; } - return StmtExprVisitorWithAnalyzer::Visit_(op); + return StmtExprVisitorWithConstraintContext::Visit_(op); } }; diff --git a/src/s_tir/analysis/is_pure_function.cc b/src/s_tir/analysis/is_pure_function.cc index 873a7a224c94..b506af6cc253 100644 --- a/src/s_tir/analysis/is_pure_function.cc +++ b/src/s_tir/analysis/is_pure_function.cc @@ -64,15 +64,15 @@ class PurityChecker : public StmtExprVisitor { << "Do not have a default for " << stmt.value()->GetTypeKey(); } if (auto expr = value.as(); expr && !GlobalVTable()->CanDispatch(value)) { - if (value.as() || value.as()) return std::nullopt; + if (value.as() || value.as()) return std::nullopt; TVM_FFI_VISIT_THROW(TypeError, expr.value()) << "Unsupported non-primitive TIR expression " << expr.value()->GetTypeKey(); } return StmtExprVisitor::Visit(value); } - Result Visit_(const BindNode* op) final { - if (const auto* call = op->value.as(); + Result Visit_(const BindObj* op) final { + if (const auto* call = op->value.as(); call && call->op.same_as(tirx::alloc_tensor_op())) { internal_allocations_.insert(op->var); allocation_calls_.insert(call); @@ -81,7 +81,7 @@ class PurityChecker : public StmtExprVisitor { return this->Visit(op->var->ty); } - Result Visit_(const TensorStoreNode* op) override { + Result Visit_(const TensorStoreObj* op) override { this->Visit(op->value); this->Visit(op->indices); @@ -97,8 +97,8 @@ class PurityChecker : public StmtExprVisitor { return std::nullopt; } - Result Visit_(const CallNode* call) override { - if (call->op.as() || call->op.as()) this->Visit(call->op); + Result Visit_(const CallObj* call) override { + if (call->op.as() || call->op.as()) this->Visit(call->op); this->Visit(call->args); if (allocation_calls_.count(call)) return std::nullopt; @@ -124,23 +124,23 @@ class PurityChecker : public StmtExprVisitor { } // Match the analysis' traversal policy without introducing binder state. - Result Visit_(const FunctionNode* op) { + Result Visit_(const FunctionObj* op) { for (const auto& param : op->params) { - if (!param->ty.as()) this->Visit(param->ty); + if (!param->ty.as()) this->Visit(param->ty); } for (const auto& param : op->params) { - if (param->ty.as()) this->Visit(param->ty); + if (param->ty.as()) this->Visit(param->ty); } return this->Visit(op->body); } - Result Visit_(const LambdaExprNode* op) { + Result Visit_(const LambdaExprObj* op) { this->Visit(op->ty); for (const auto& var : op->vars) this->Visit(var->ty); return this->Visit(op->body); } - Result Visit_(const RegionStmtNode* op) override { + Result Visit_(const RegionStmtObj* op) override { this->Visit(op->args); this->Visit(op->attrs); for (const auto& param : op->body_params) this->Visit(param->ty); @@ -149,7 +149,7 @@ class PurityChecker : public StmtExprVisitor { return std::nullopt; } - Result Visit_(const ForNode* op) override { + Result Visit_(const ForObj* op) override { this->Visit(op->begin); this->Visit(op->end); this->Visit(op->step); @@ -157,39 +157,39 @@ class PurityChecker : public StmtExprVisitor { return this->Visit(op->body); } - Result Visit_(const prim::LetNode* op) override { + Result Visit_(const prim::LetObj* op) override { this->Visit(op->value); this->Visit(op->var->ty); return this->Visit(op->body); } - Result Visit_(const TensorLoadNode* op) override { return this->Visit(op->indices); } + Result Visit_(const TensorLoadObj* op) override { return this->Visit(op->indices); } - Result Visit_(const TensorRegionNode* op) override { + Result Visit_(const TensorRegionObj* op) override { if (!op->source.as()) this->Visit(op->source); return this->Visit(op->region); } - Result Visit_(const GenericConstNode* op) override { return std::nullopt; } + Result Visit_(const GenericConstObj* op) override { return std::nullopt; } - Result Visit_(const AssertStmtNode* op) override { + Result Visit_(const AssertStmtObj* op) override { this->Visit(op->condition); this->Visit(op->error_kind); return this->Visit(op->message_parts); } - Result Visit_(const prim::RampNode* op) override { + Result Visit_(const prim::RampObj* op) override { this->Visit(op->base); this->Visit(op->stride); return this->Visit(op->lanes); } - Result Visit_(const prim::BroadcastNode* op) override { + Result Visit_(const prim::BroadcastObj* op) override { this->Visit(op->value); return this->Visit(op->lanes); } - Result Visit_(const SBlockNode* op) override { + Result Visit_(const SBlockObj* op) override { for (const auto& axis : op->iter_vars) { this->Visit(axis->dom); this->Visit(axis->var.as_or_throw()->ty); @@ -209,8 +209,8 @@ class PurityChecker : public StmtExprVisitor { static const VTable table = [] { VTable table; StmtExprVisitor::InitVTable(&table); - SetDispatch(&table); - SetDispatch(&table); + SetDispatch(&table); + SetDispatch(&table); table.Finalize(); return table; }(); @@ -221,7 +221,7 @@ class PurityChecker : public StmtExprVisitor { bool assert_on_error_{false}; bool is_pure_{true}; std::unordered_set internal_allocations_; - std::unordered_set allocation_calls_; + std::unordered_set allocation_calls_; }; } // namespace diff --git a/src/s_tir/analysis/oob_checker.cc b/src/s_tir/analysis/oob_checker.cc index fe24a2c2a23e..cf27e79aa33e 100644 --- a/src/s_tir/analysis/oob_checker.cc +++ b/src/s_tir/analysis/oob_checker.cc @@ -22,7 +22,7 @@ */ #include -#include +#include #include #include "../schedule/error.h" @@ -69,28 +69,28 @@ class OOBError : public s_tir::ScheduleErrorContextObj { IRModule mod_; std::vector locations_; }; -class OOBCheckerVisitor final : public s_tir::StmtExprVisitorWithAnalyzer { +class OOBCheckerVisitor final : public s_tir::StmtExprVisitorWithConstraintContext { public: - using s_tir::StmtExprVisitorWithAnalyzer::Visit_; + using s_tir::StmtExprVisitorWithConstraintContext::Visit_; - ffi::Optional Visit_(const TensorStoreNode* node) final { + ffi::Optional Visit_(const TensorStoreObj* node) final { for (size_t i = 0; i < node->dest.as_or_throw().shape().size(); i++) { CheckBounds(node, node->dest.as_or_throw(), i); } - return StmtExprVisitorWithAnalyzer::Visit_(node); + return StmtExprVisitorWithConstraintContext::Visit_(node); } - ffi::Optional Visit_(const TensorLoadNode* node) final { + ffi::Optional Visit_(const TensorLoadObj* node) final { TensorVar tensor = node->source.as_or_throw(); for (size_t i = 0; i < tensor.shape().size(); i++) { CheckBounds(node, tensor, i); } - return StmtExprVisitorWithAnalyzer::Visit_(node); + return StmtExprVisitorWithConstraintContext::Visit_(node); } template void CheckBounds(const T* node, const TensorVar& tensor, size_t i) { - auto ind_bounds = analyzer_->int_set(node->indices[i]); - auto shape_bounds = analyzer_->int_set(tensor.shape()[i]); + auto ind_bounds = GetSyncedAnalyzer()->int_set(node->indices[i]); + auto shape_bounds = GetSyncedAnalyzer()->int_set(tensor.shape()[i]); // We would expect that // `analyzer_.CanProve(node->indices[i] < 0 || node->indices[i] >= tensor->shape[i])` // would be the way to check if any out of bounds access occurs here, but `CanProve` checks if @@ -102,8 +102,8 @@ class OOBCheckerVisitor final : public s_tir::StmtExprVisitorWithAnalyzer { // has the problem that some valid access patterns maybe be valid but not provably valid. We // prefer that this analysis is conservative and only shows errors that are provable. This leads // us to the following check: are the bounds of the index outside the bounds of the shape. - if (analyzer_->CanProve(ind_bounds.max() >= shape_bounds.min()) || - analyzer_->CanProve(ind_bounds.min() < 0)) { + if (GetSyncedAnalyzer()->CanProve(ind_bounds.max() >= shape_bounds.min()) || + GetSyncedAnalyzer()->CanProve(ind_bounds.min() < 0)) { errors.push_back({tensor, i, node->indices[i], ind_bounds, shape_bounds}); } } diff --git a/src/s_tir/analysis/sblock_access_region_detector.cc b/src/s_tir/analysis/sblock_access_region_detector.cc index 6311da7ce4f8..908c8b9ec48a 100644 --- a/src/s_tir/analysis/sblock_access_region_detector.cc +++ b/src/s_tir/analysis/sblock_access_region_detector.cc @@ -60,10 +60,10 @@ class BlockReadWriteDetector : public s_tir::StmtExprVisitor { /*! \brief Return read regions of the block */ ffi::Array CollectReads( - const std::unordered_set* excluded_tensors = nullptr); + const std::unordered_set* excluded_tensors = nullptr); /*! \brief Return write regions of the block */ ffi::Array CollectWrites( - const std::unordered_set* excluded_tensors = nullptr); + const std::unordered_set* excluded_tensors = nullptr); /*! * \brief Return opaque tensor regions of the block * \note The tensor accessed by load/store or call with tensor.data will @@ -95,9 +95,9 @@ class BlockReadWriteDetector : public s_tir::StmtExprVisitor { /*! \brief The outside tensor data mapping to its tensor */ ffi::Map tensor_var_map_; /*! \brief The target tensor var mapping to its matching */ - std::unordered_map match_tensors_; + std::unordered_map match_tensors_; /*! \brief let bindings inside the block */ - std::unordered_map let_bindings_; + std::unordered_map let_bindings_; /*!\ brief Internal analyzer. */ sym::Analyzer ana_; @@ -115,7 +115,7 @@ class BlockReadWriteDetector : public s_tir::StmtExprVisitor { ffi::Array CollectRegions( const std::vector& tensors, const std::vector>& regions, - const std::unordered_set* excluded_tensors = nullptr); + const std::unordered_set* excluded_tensors = nullptr); /*! \brief Helper function to convert matched access region to source region. */ std::vector ConvertMatchedRegion(const s_tir::MatchTensorRegion& match_tensor, @@ -128,7 +128,7 @@ class BlockReadWriteDetector : public s_tir::StmtExprVisitor { sym::IntSet RelaxAccessIndex(const PrimExpr& index); // Declared regions carry bounds, not opaque runtime accesses. - ffi::Optional Visit_(const TensorRegionNode* op) final { + ffi::Optional Visit_(const TensorRegionObj* op) final { if (!op->source.as()) return StmtExprVisitor::Visit_(op); for (const Range& range : op->region) { TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Visit(range->begin)); @@ -137,18 +137,18 @@ class BlockReadWriteDetector : public s_tir::StmtExprVisitor { return std::nullopt; } - ffi::Optional Visit_(const ForNode* op) override; - ffi::Optional Visit_(const IfNode* op) override; - ffi::Optional Visit_(const s_tir::SBlockRealizeNode* op) override; - ffi::Optional Visit_(const TensorStoreNode* op) override; - ffi::Optional Visit_(const BindNode* op) override; - ffi::Optional Visit_(const TensorLoadNode* op) override; - ffi::Optional Visit_(const VarNode* op) override; - ffi::Optional Visit_(const CallNode* op) override; + ffi::Optional Visit_(const ForObj* op) override; + ffi::Optional Visit_(const IfObj* op) override; + ffi::Optional Visit_(const s_tir::SBlockRealizeObj* op) override; + ffi::Optional Visit_(const TensorStoreObj* op) override; + ffi::Optional Visit_(const BindObj* op) override; + ffi::Optional Visit_(const TensorLoadObj* op) override; + ffi::Optional Visit_(const VarObj* op) override; + ffi::Optional Visit_(const CallObj* op) override; }; void BlockReadWriteDetector::operator()(const Stmt& stmt) { - const auto* block = stmt.as(); + const auto* block = stmt.as(); TVM_FFI_ICHECK(block != nullptr) << "Only visiting Blocks is allowed, but got " << stmt->GetTypeKey(); for (const s_tir::MatchTensorRegion& match_tensor : block->match_tensors) { @@ -163,12 +163,12 @@ void BlockReadWriteDetector::operator()(const Stmt& stmt) { } ffi::Array BlockReadWriteDetector::CollectReads( - const std::unordered_set* excluded_tensors) { + const std::unordered_set* excluded_tensors) { return CollectRegions(read_tensors_, read_regions_, excluded_tensors); } ffi::Array BlockReadWriteDetector::CollectWrites( - const std::unordered_set* excluded_tensors) { + const std::unordered_set* excluded_tensors) { return CollectRegions(writes_tensors_, write_regions_, excluded_tensors); } @@ -176,13 +176,13 @@ ffi::Array BlockReadWriteDetector::CollectOpaques() { return CollectRegions(opaque_tensors_, opaque_regions_); } -ffi::Optional BlockReadWriteDetector::Visit_(const VarNode* op) { +ffi::Optional BlockReadWriteDetector::Visit_(const VarObj* op) { if (def_region_kind() != kTVMFFIDefRegionKindNone) return StmtExprVisitor::Visit_(op); UpdateOpaque(ffi::GetRef(op)); return std::nullopt; } -ffi::Optional BlockReadWriteDetector::Visit_(const TensorLoadNode* op) { +ffi::Optional BlockReadWriteDetector::Visit_(const TensorLoadObj* op) { 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); @@ -208,7 +208,7 @@ ffi::Optional BlockReadWriteDetector::Visit_(const TensorLoadNod return std::nullopt; } -ffi::Optional BlockReadWriteDetector::Visit_(const ForNode* op) { +ffi::Optional BlockReadWriteDetector::Visit_(const ForObj* op) { Range range = Range(op->begin, op->end); dom_map_.Set(op->var, sym::IntSet::FromRange(range)); TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(s_tir::StmtExprVisitor::Visit_(op)); @@ -216,7 +216,7 @@ ffi::Optional BlockReadWriteDetector::Visit_(const ForNode* op) return std::nullopt; } -ffi::Optional BlockReadWriteDetector::Visit_(const IfNode* op) { +ffi::Optional BlockReadWriteDetector::Visit_(const IfObj* op) { TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Visit(op->condition)); { // Visit then branch @@ -233,8 +233,8 @@ ffi::Optional BlockReadWriteDetector::Visit_(const IfNode* op) { return std::nullopt; } -ffi::Optional BlockReadWriteDetector::Visit_(const BindNode* op) { - if (const auto* call = op->value.as(); +ffi::Optional BlockReadWriteDetector::Visit_(const BindObj* op) { + if (const auto* call = op->value.as(); call && call->op.same_as(tirx::decl_tensor_op())) { // A DeclTensor data expression defines the alias source. It is not an // opaque tensor access by the containing block. @@ -246,7 +246,7 @@ ffi::Optional BlockReadWriteDetector::Visit_(const BindNode* op) return s_tir::StmtExprVisitor::Visit_(op); } -ffi::Optional BlockReadWriteDetector::Visit_(const CallNode* op) { +ffi::Optional BlockReadWriteDetector::Visit_(const CallObj* op) { auto update_masked_access = [this](const TensorVar& tensor, const ffi::Array& indices, std::vector* tensors, std::vector>* regions) { @@ -284,7 +284,7 @@ ffi::Optional BlockReadWriteDetector::Visit_(const CallNode* op) return std::nullopt; } if (op->op.same_as(tirx::address_of_op())) { - if (const auto* load = op->args[0].as()) { + if (const auto* load = op->args[0].as()) { // A tensor address is an opaque use, not a load of its pointed-to value. UpdateOpaque(load->source.as_or_throw()); for (const PrimExpr& index : load->indices) { @@ -315,7 +315,7 @@ ffi::Optional BlockReadWriteDetector::Visit_(const CallNode* op) return s_tir::StmtExprVisitor::Visit_(op); } -ffi::Optional BlockReadWriteDetector::Visit_(const TensorStoreNode* op) { +ffi::Optional BlockReadWriteDetector::Visit_(const TensorStoreObj* op) { 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); @@ -341,9 +341,9 @@ ffi::Optional BlockReadWriteDetector::Visit_(const TensorStoreNo return std::nullopt; } -ffi::Optional BlockReadWriteDetector::Visit_(const s_tir::SBlockRealizeNode* op) { +ffi::Optional BlockReadWriteDetector::Visit_(const s_tir::SBlockRealizeObj* op) { /*! \note detector will not visit child block recursively, so it will stop here */ - std::unordered_map vmap; + std::unordered_map vmap; for (size_t i = 0; i < op->block->iter_vars.size(); ++i) { vmap.insert_or_assign(op->block->iter_vars[i]->var.get(), op->iter_values[i]); } @@ -431,7 +431,7 @@ void BlockReadWriteDetector::Update(std::vector* tensors, ffi::Array BlockReadWriteDetector::CollectRegions( const std::vector& tensors, const std::vector>& regions, - const std::unordered_set* excluded_tensors) { + const std::unordered_set* excluded_tensors) { TVM_FFI_ICHECK_EQ(tensors.size(), regions.size()); ffi::Array res; res.reserve(tensors.size()); @@ -476,7 +476,7 @@ ffi::Array> GetSBlockAccessRegion( auto detector = ffi::make_object(tensor_var_map); detector->operator()(block); ffi::Array writes = detector->CollectWrites(); - std::unordered_set excluded_tensors; + std::unordered_set excluded_tensors; // exclude write tensors from read regions for reductions if init block is defined. if (block->init.has_value()) { for (const TensorRegion& write_access : writes) { @@ -493,7 +493,7 @@ ffi::Array> GetSBlockReadWriteRegion( auto detector = ffi::make_object(tensor_var_map); detector->operator()(block); ffi::Array opaques = detector->CollectOpaques(); - std::unordered_set excluded_tensors; + std::unordered_set excluded_tensors; for (const TensorRegion& opaque_access : opaques) { excluded_tensors.insert(opaque_access->source.as_or_throw().get()); } diff --git a/src/s_tir/analysis/sblock_tensor_access_lca_detector.cc b/src/s_tir/analysis/sblock_tensor_access_lca_detector.cc index d7aeb2c45992..874910866b7e 100644 --- a/src/s_tir/analysis/sblock_tensor_access_lca_detector.cc +++ b/src/s_tir/analysis/sblock_tensor_access_lca_detector.cc @@ -91,14 +91,14 @@ class LCADetector : public s_tir::StmtExprVisitor { // The parent scope info const ScopeInfo* parent_scope_info; // The parent scope stmt node - const StmtNode* stmt; + const StmtObj* stmt; // The scope depth in the AST int depth; - ScopeInfo(const ScopeInfo* parent_info, const StmtNode* stmt, int depth) + ScopeInfo(const ScopeInfo* parent_info, const StmtObj* stmt, int depth) : parent_scope_info(parent_info), stmt(stmt), depth(depth) {} }; - ffi::Optional Visit_(const ForNode* op) final { + ffi::Optional Visit_(const ForObj* op) final { int n = ancestor_scopes_.size(); const ScopeInfo* parent_scope = ancestor_scopes_.back(); auto* current_scope = arena_.make(parent_scope, op, n); @@ -118,8 +118,8 @@ class LCADetector : public s_tir::StmtExprVisitor { return std::nullopt; } - ffi::Optional Visit_(const s_tir::SBlockRealizeNode* op) final { - const s_tir::SBlockNode* block = op->block.get(); + ffi::Optional Visit_(const s_tir::SBlockRealizeObj* op) final { + const s_tir::SBlockObj* block = op->block.get(); int n = ancestor_scopes_.size(); for (const TensorVar& tensor : block->alloc_tensors) { tensor_var_map_.emplace(tensor.get(), tensor.get()); @@ -150,9 +150,9 @@ class LCADetector : public s_tir::StmtExprVisitor { return std::nullopt; } - void UpdateDominateScopeOfNonDataParIter(const s_tir::SBlockRealizeNode* block_realize) { + void UpdateDominateScopeOfNonDataParIter(const s_tir::SBlockRealizeObj* block_realize) { // map iter var to the scope which dominate all loop carried dependencies. - std::unordered_map opaque_var_scope; + std::unordered_map opaque_var_scope; // maintain highest scope which dominate all reduce loop iters. null denotes non-reduce block. const ScopeInfo* highest_reduce_scope = nullptr; @@ -163,7 +163,7 @@ class LCADetector : public s_tir::StmtExprVisitor { const ScopeInfo* highest_scope = nullptr; auto walk_fn = [this, &highest_scope](const Var& var) -> ffi::Expected { if (auto prim_var = var.as()) { - const VarNode* loop_var = prim_var.value().get(); + const VarObj* loop_var = prim_var.value().get(); auto it = loop_scope_map_.find(loop_var); if (it == loop_scope_map_.end()) { return ffi::WalkResult::Advance(); @@ -216,7 +216,7 @@ class LCADetector : public s_tir::StmtExprVisitor { auto handle_itervar = [&opaque_var_scope, &scope](const Var& var) -> ffi::Expected { if (auto prim_var = var.as()) { - const VarNode* iter_var = prim_var.value().get(); + const VarObj* iter_var = prim_var.value().get(); auto dom_scope_it = opaque_var_scope.find(iter_var); if (dom_scope_it == opaque_var_scope.end()) { return ffi::WalkResult::Advance(); @@ -259,7 +259,7 @@ class LCADetector : public s_tir::StmtExprVisitor { } } - ffi::Optional Visit_(const RegionStmtNode* op) final { + ffi::Optional Visit_(const RegionStmtObj* op) final { // Operands belong to the enclosing scope, before the local definitions. TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Visit(op->args)); TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Visit(op->attrs)); @@ -279,7 +279,7 @@ class LCADetector : public s_tir::StmtExprVisitor { } // Declared regions carry bounds, not opaque runtime accesses. - ffi::Optional Visit_(const TensorRegionNode* op) final { + ffi::Optional Visit_(const TensorRegionObj* op) final { if (!op->source.as()) return s_tir::StmtExprVisitor::Visit_(op); for (const Range& range : op->region) { TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Visit(range->begin)); @@ -288,7 +288,7 @@ class LCADetector : public s_tir::StmtExprVisitor { return std::nullopt; } - ffi::Optional Visit_(const TensorLoadNode* op) final { + ffi::Optional Visit_(const TensorLoadObj* op) final { UpdateTensorLCA(op->source.as_or_throw().get(), ancestor_scopes_.back()); for (const auto& index : op->indices) { TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Visit(index)); @@ -296,7 +296,7 @@ class LCADetector : public s_tir::StmtExprVisitor { return std::nullopt; } - ffi::Optional Visit_(const TensorStoreNode* op) final { + ffi::Optional Visit_(const TensorStoreObj* op) final { UpdateTensorLCA(op->dest.as_or_throw().get(), ancestor_scopes_.back()); TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Visit(op->value)); for (const auto& index : op->indices) { @@ -306,20 +306,20 @@ class LCADetector : public s_tir::StmtExprVisitor { } // Works for Load/Store and opaque access. - ffi::Optional Visit_(const VarNode* op) final { + ffi::Optional Visit_(const VarObj* op) final { if (def_region_kind() != kTVMFFIDefRegionKindNone) return StmtExprVisitor::Visit_(op); VisitTensorVar(op); return std::nullopt; } - void VisitTensorVar(const VarNode* op) { + void VisitTensorVar(const VarObj* op) { auto it = tensor_var_map_.find(op); if (it != tensor_var_map_.end()) { UpdateTensorLCA(it->second, ancestor_scopes_.back()); } } - void UpdateTensorLCA(const VarNode* tensor, const ScopeInfo* scope) { + void UpdateTensorLCA(const VarObj* tensor, const ScopeInfo* scope) { tensor_var_map_.emplace(tensor, tensor); if (match_tensors_.find(tensor) == match_tensors_.end()) { // Ingore tensor created by block match_tensor @@ -371,16 +371,16 @@ class LCADetector : public s_tir::StmtExprVisitor { * the root scope. */ std::vector ancestor_scopes_ = {}; - /*! \brief The map from TensorVar to its LCA ForNode/BlockNode. */ - std::unordered_map tensor_lca_ = {}; + /*! \brief The map from TensorVar to its LCA ForObj/BlockNode. */ + std::unordered_map tensor_lca_ = {}; /*! \brief The map from TensorVar data to the TensorVar. */ - std::unordered_map tensor_var_map_ = {}; + std::unordered_map tensor_var_map_ = {}; /*! \brief The match tensors inside blocks. */ - std::unordered_set match_tensors_ = {}; + std::unordered_set match_tensors_ = {}; /*! \brief The ForNodes/BlockNodes which contain immediate `blockIdx` launch. */ std::vector blockidx_scopes_ = {}; /*! \brief The map from loop var to the corresponding scope. */ - std::unordered_map loop_scope_map_ = {}; + std::unordered_map loop_scope_map_ = {}; /*! \brief Internal arena. */ support::Arena arena_; }; diff --git a/src/s_tir/analysis/verify_gpu_code.cc b/src/s_tir/analysis/verify_gpu_code.cc index d320ca124c0a..f819b02a481e 100644 --- a/src/s_tir/analysis/verify_gpu_code.cc +++ b/src/s_tir/analysis/verify_gpu_code.cc @@ -70,15 +70,15 @@ class GPUCodeVerifier : public StmtExprVisitor { return errors_; } - ffi::Optional Visit_(const BindNode* op) final { - if (const auto* call = op->value.as(); + ffi::Optional Visit_(const BindObj* op) final { + if (const auto* call = op->value.as(); call && call->op.same_as(tirx::alloc_tensor_op())) { return DispatchAllocTensor(op, call); } return StmtExprVisitor::Visit_(op); } - ffi::Optional DispatchAllocTensor(const BindNode* op, const CallNode* call) { + ffi::Optional DispatchAllocTensor(const BindObj* op, const CallObj* call) { static const auto scope_local = runtime::MemoryScope::Get("local"); static const auto scope_shared = runtime::MemoryScope::Get("shared"); static const auto scope_shared_dyn = runtime::MemoryScope::Get("shared.dyn"); @@ -90,7 +90,7 @@ class GPUCodeVerifier : public StmtExprVisitor { runtime::MemoryScope storage_scope = scope; int64_t const_size = 1; for (const Expr& e : shape->fields) { - if (auto* imm = e.as()) { + if (auto* imm = e.as()) { const_size = static_cast(const_size * imm->value); } else { const_size = 0; @@ -117,7 +117,7 @@ class GPUCodeVerifier : public StmtExprVisitor { return std::nullopt; } - ffi::Optional Visit_(const RegionStmtNode* op) final { + ffi::Optional Visit_(const RegionStmtObj* op) final { if (op->op.same_as(tirx::launch_thread_op())) { if (nest_level_ == 0) { // enter a new kernel, reset statistics @@ -125,7 +125,7 @@ class GPUCodeVerifier : public StmtExprVisitor { kernels_launched_++; } - const auto* extent = op->args[1].as_or_throw().as(); + const auto* extent = op->args[1].as_or_throw().as(); TVM_FFI_ICHECK(extent); std::string name = op->args[0].as_or_throw()->value; @@ -215,10 +215,10 @@ class GPUCodeVerifier : public StmtExprVisitor { return std::nullopt; } - ffi::Optional Visit_(const ForNode* op) { + ffi::Optional Visit_(const ForObj* op) { if (op->var->name == "vthread.s") { PrimExpr extent_expr = Range::GetExtentFromBeginEnd(op->begin, op->end); - const auto* extent = extent_expr.as(); + const auto* extent = extent_expr.as(); TVM_FFI_ICHECK(extent); size_t num_vthread = extent->value.as().value(); @@ -235,7 +235,7 @@ class GPUCodeVerifier : public StmtExprVisitor { void CheckTensorIndicesVectorizable(const ffi::Array indices) { for (const auto index : indices) { - if (const auto* ramp = index.as()) { + if (const auto* ramp = index.as()) { PrimType ramp_ty = ramp->ty.as_or_throw(); if (!IsOne(ramp->stride) && ramp_ty.IsFixedLengthVector() && ElementBytes(ramp_ty) > max_vector_bytes_) { @@ -249,7 +249,7 @@ class GPUCodeVerifier : public StmtExprVisitor { } } - ffi::Optional Visit_(const CastNode* op) { + ffi::Optional Visit_(const CastObj* op) { PrimType op_ty = op->ty.as_or_throw(); if (op_ty.IsFixedLengthVector()) { if (ElementBytes(op_ty) > max_vector_bytes_) { @@ -263,7 +263,7 @@ class GPUCodeVerifier : public StmtExprVisitor { return StmtExprVisitor::Visit_(op); } - ffi::Optional Visit_(const TensorLoadNode* op) { + ffi::Optional Visit_(const TensorLoadObj* op) { PrimType op_ty = op->ty.as_or_throw(); if (op_ty.IsFixedLengthVector()) { if (ElementBytes(op_ty) > max_vector_bytes_) { @@ -278,7 +278,7 @@ class GPUCodeVerifier : public StmtExprVisitor { return StmtExprVisitor::Visit_(op); } - ffi::Optional Visit_(const TensorStoreNode* op) { + ffi::Optional Visit_(const TensorStoreObj* op) { PrimType value_ty = op->value.ty(); if (value_ty.IsFixedLengthVector()) { if (ElementBytes(value_ty) > max_vector_bytes_) { @@ -345,7 +345,7 @@ std::vector VerifyGPUCode_(const Function& func, int64_t max_kernels = INT64_MAX; for (auto iter : constraints) { - const IntImmNode* val = iter.second.as(); + const IntImmObj* val = iter.second.as(); if (iter.first == "max_local_memory_per_block") { max_local_memory_per_block = static_cast(val->value); } else if (iter.first == "max_shared_memory_per_block") { diff --git a/src/s_tir/backend/adreno/inject_texture_alloc.cc b/src/s_tir/backend/adreno/inject_texture_alloc.cc index a09afd3bfec3..22d85be4a1fb 100644 --- a/src/s_tir/backend/adreno/inject_texture_alloc.cc +++ b/src/s_tir/backend/adreno/inject_texture_alloc.cc @@ -27,7 +27,7 @@ #include #include #include -#include +#include #include #include #include @@ -47,49 +47,53 @@ using texture::IsTextureStorage; /*! * \brief Inject Texture Alloc Intrinsic right after tensor allocations are realized. */ -class TextureAllocInjector : public s_tir::StmtExprMutatorWithAnalyzer { +class TextureAllocInjector : public s_tir::StmtExprMutatorWithConstraintContext { public: - using s_tir::StmtExprMutatorWithAnalyzer::Mutate; - using s_tir::StmtExprMutatorWithAnalyzer::Mutate_; + using s_tir::StmtExprMutatorWithConstraintContext::Mutate; + using s_tir::StmtExprMutatorWithConstraintContext::Mutate_; static Function Inject(Function func) { sym::Analyzer ana; auto pass = ffi::make_object(ana); auto writer = func.CopyOnWrite(); - pass->MarkTensorParamShapes(func); - writer->body = pass->Mutate(func->body).ValueOrUnchanged(func->body); + pass->AddParamConstraints(func); + writer->body = pass->Mutate(func->body) + .as_or_throw>>() + .ValueOrUnchanged(func->body); return func; } - explicit TextureAllocInjector(const sym::Analyzer& ana) : StmtExprMutatorWithAnalyzer(ana) {} + explicit TextureAllocInjector(const sym::Analyzer& ana) + : StmtExprMutatorWithConstraintContext(ana->GetConstraintContext()) {} private: - UnchangedOr Mutate_(const BindNode* op, InplaceMode inplace_mode) final { - if (const auto* call = op->value.as(); + UnchangedOr Mutate_(const BindObj* op, InplaceMode inplace_mode) final { + if (const auto* call = op->value.as(); call && call->op.same_as(tirx::alloc_tensor_op())) { return Mutate_AllocTensor(op, call, inplace_mode); } - return StmtExprMutator::Mutate_(op, inplace_mode); + return StmtExprMutatorWithConstraintContext::Mutate_(op, inplace_mode); } - UnchangedOr Mutate_AllocTensor(const BindNode* op, const CallNode* call, + UnchangedOr Mutate_AllocTensor(const BindObj* op, const CallObj* call, InplaceMode inplace_mode) { - Stmt stmt = StmtExprMutator::Mutate_(op, inplace_mode).ValueOrUnchanged(ffi::GetRef(op)); + Stmt stmt = StmtExprMutatorWithConstraintContext::Mutate_(op, inplace_mode) + .ValueOrUnchanged(ffi::GetRef(op)); runtime::MemoryScope scope = call->ty_args[0].as_or_throw()->scope; if (IsTextureStorage(scope)) { - op = stmt.as(); - if (const auto* call = op ? op->value.as() : nullptr; + op = stmt.as(); + if (const auto* call = op ? op->value.as() : nullptr; !call || !call->op.same_as(tirx::alloc_tensor_op())) { TVM_FFI_THROW(InternalError) << "Expected an allocation binding after tensor mutation"; } - const auto* allocation = op->value.as(); + const auto* allocation = op->value.as(); tvm::Tuple shape = tvm::Tuple(allocation->ty_args[0].as_or_throw()->shape); DLDataType dtype = allocation->ty_args[0].as_or_throw()->dtype->dtype; ffi::Array extents = shape->fields.Map([](const Expr& e) { return e.as_or_throw(); }); TVM_FFI_ICHECK(extents.size() >= 3) << "Only 2D Array RGBA texture is currently supported"; const int data_bits = dtype.bits, - vec_length = extents.back().as()->value.as().value(); + vec_length = extents.back().as()->value.as().value(); const int channel_size = data_bits * vec_length; TVM_FFI_ICHECK(channel_size == 128 || channel_size == 64) << "Invalid Channel Size: " << channel_size << " bits"; @@ -114,7 +118,7 @@ class TextureAllocInjector : public s_tir::StmtExprMutatorWithAnalyzer { protected: runtime::MemoryScope GetStorageScope(const Var& buffer_var) { - auto* ptr = buffer_var->ty.as(); + auto* ptr = buffer_var->ty.as(); TVM_FFI_ICHECK(ptr) << "Buffer Var's type annotation must be of PtrType"; return ptr->scope; } diff --git a/src/s_tir/backend/adreno/texture_flatten.cc b/src/s_tir/backend/adreno/texture_flatten.cc index ec774ad78a13..381f4817493e 100644 --- a/src/s_tir/backend/adreno/texture_flatten.cc +++ b/src/s_tir/backend/adreno/texture_flatten.cc @@ -29,7 +29,7 @@ #include #include #include -#include +#include #include #include #include @@ -52,7 +52,7 @@ class TextureLoweringBase : public StmtExprMutator { using StmtExprMutator::Mutate_; explicit TextureLoweringBase(const ffi::Array& params, - StmtExprVisitorWithAnalyzer* bound_analyzer) + StmtExprVisitorWithConstraintContext* bound_analyzer) : bound_analyzer_{bound_analyzer} { for (const Var& param : params) { if (auto tensor = param.as()) { @@ -81,7 +81,7 @@ class TextureLoweringBase : public StmtExprMutator { // Set of all external input and output tensors std::unordered_set extern_tensor_; // Bound analzer - StmtExprVisitorWithAnalyzer* bound_analyzer_; + StmtExprVisitorWithConstraintContext* bound_analyzer_; }; // Lower Nd storage access to 2d texture access using lowering convention @@ -92,12 +92,12 @@ class TextureFlattener : public TextureLoweringBase { using TextureLoweringBase::Mutate_; explicit TextureFlattener(const ffi::Array& params, - StmtExprVisitorWithAnalyzer* bound_analyzer) + StmtExprVisitorWithConstraintContext* bound_analyzer) : TextureLoweringBase(params, bound_analyzer) {} - UnchangedOr Mutate_(const TensorStoreNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const TensorStoreObj* op, InplaceMode inplace_mode) final { Stmt stmt = StmtExprMutator::Mutate_(op, inplace_mode).ValueOrUnchanged(ffi::GetRef(op)); - op = stmt.as(); + op = stmt.as(); runtime::MemoryScope storage_scope = GetStorageScope(op->dest.as_or_throw()); // Lower to two dimensional access if (IsTextureStorage(storage_scope)) { @@ -109,10 +109,10 @@ class TextureFlattener : public TextureLoweringBase { return stmt; } - UnchangedOr Mutate_(const TensorLoadNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const TensorLoadObj* op, InplaceMode inplace_mode) final { PrimExpr expr = StmtExprMutator::Mutate_(op, inplace_mode).ValueOrUnchanged(ffi::GetRef(op)); - op = expr.as(); + op = expr.as(); // Lower to two dimensional access runtime::MemoryScope storage_scope = GetStorageScope(op->source.as_or_throw()); @@ -170,7 +170,7 @@ class TextureFlattener : public TextureLoweringBase { Function TextureFlattenHandler(Function func) { auto fptr = func.CopyOnWrite(); - auto bound_analyzer = ffi::make_object(); + auto bound_analyzer = ffi::make_object(); bound_analyzer->Visit(fptr->body); fptr->body = ffi::make_object(fptr->params, bound_analyzer.get()) ->Mutate(fptr->body, InplaceMode::kAllow) diff --git a/src/s_tir/data_layout.cc b/src/s_tir/data_layout.cc index aea90c011f7e..e9d4234d046d 100644 --- a/src/s_tir/data_layout.cc +++ b/src/s_tir/data_layout.cc @@ -42,7 +42,7 @@ namespace tvm { namespace tirx { using s_tir::IterVar; -using s_tir::IterVarNode; +using s_tir::IterVarObj; using tvm::Var; const SLayoutAxis SLayoutAxis::UPPER_CASE[] = { @@ -80,7 +80,7 @@ const SLayoutAxis& SLayoutAxis::Get(const std::string& name) { } SLayout::SLayout(const ffi::Array& axes) { - auto node = ffi::make_object(); + auto node = ffi::make_object(); node->axes = axes; std::ostringstream repr; @@ -90,7 +90,7 @@ SLayout::SLayout(const ffi::Array& axes) { if (is_grouped) repr << "["; for (const IterVar& axis : unpacked_axes) { - if (const auto* factor = axis->dom.value()->extent.as()) { + if (const auto* factor = axis->dom.value()->extent.as()) { TVM_FFI_ICHECK_GT(factor->value, 0); repr << factor->value; } else { @@ -117,7 +117,7 @@ SLayout::SLayout(const std::string& name, PrimType index_ty) { // NOLINT(*) TVM_FFI_CHECK(!name.empty() && name != "__undef__", ValueError) << "Undefined layout requires Optional"; - auto node = ffi::make_object(); + auto node = ffi::make_object(); node->name = name; // parse layout string @@ -163,7 +163,7 @@ SLayout::SLayout(const std::string& name, PrimType index_ty) { // NOLINT(*) std::stringstream ss; ffi::BigInt extent = 1; for (auto& axis : unpacked_axes) { - TVM_FFI_ICHECK(axis->dom.value()->extent.as()) + TVM_FFI_ICHECK(axis->dom.value()->extent.as()) << "Invalid SLayout " << name << ": can't have variable sized node(" << axis->var->name << ") within a packed axis"; auto axis_name = axis->var->name.operator std::string(); @@ -221,11 +221,11 @@ ffi::Optional SLayout::Create(const std::string& name, PrimType index_t TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - SLayoutNode::RegisterReflection(); - refl::TypeAttrDef().def(ffi::type_attr::kRepr.name, - [](SLayout l, ffi::Function) -> ffi::String { - return "SLayout(" + std::string(l->name) + ")"; - }); + SLayoutObj::RegisterReflection(); + refl::TypeAttrDef().def(ffi::type_attr::kRepr.name, + [](SLayout l, ffi::Function) -> ffi::String { + return "SLayout(" + std::string(l->name) + ")"; + }); refl::GlobalDef().def("s_tir.SLayout", [](std::string name, PrimType dtype) -> ffi::Optional { @@ -299,7 +299,7 @@ int32_t SLayout::FactorOf(const SLayoutAxis& axis) const { for (auto itvar : UnpackIterVar(packed_itvar)) { if (sub == SLayoutAxis::Get(itvar)) { has_sub = true; - int32_t val = itvar->dom.value()->extent.as()->value.as().value(); + int32_t val = itvar->dom.value()->extent.as()->value.as().value(); factor *= val; } } @@ -343,7 +343,7 @@ inline bool GetStoreRule(ffi::Array* index_rule, ffi::Array* std::vector index_divs(src_unpacked_axes.size()); for (size_t j = 0; j < src_unpacked_axes.size(); j++) { index_divs[j] = value; - const auto* extent = src_unpacked_axes[j]->dom.value()->extent.as(); + const auto* extent = src_unpacked_axes[j]->dom.value()->extent.as(); TVM_FFI_ICHECK(extent) << "Expected Integer Extents for Offset Calculation"; index_divs.push_back(value); value *= extent->value; @@ -404,8 +404,8 @@ inline bool GetStoreRule(ffi::Array* index_rule, ffi::Array* for (size_t j = 0; j < dst_unpacked_axes.size(); j++) { const auto& prim_axis = SLayoutAxis::Get(dst_unpacked_axes[j]).ToPrimal(); const auto& sub_axis = SLayoutAxis::Get(dst_unpacked_axes[j]).ToSubordinate(); - const auto* extent = dst_unpacked_axes[j]->dom.value()->extent.as(); - TVM_FFI_ICHECK(extent) << "Expected extent to be IntImmNode"; + const auto* extent = dst_unpacked_axes[j]->dom.value()->extent.as(); + TVM_FFI_ICHECK(extent) << "Expected extent to be IntImmObj"; size_t divfactor = 1; for (size_t k = i; k < dst_layout.ndim(); k++) { @@ -416,7 +416,7 @@ inline bool GetStoreRule(ffi::Array* index_rule, ffi::Array* for (; l < inter_unpacked_axes.size(); l++) { const auto& axis = SLayoutAxis::Get(inter_unpacked_axes[l]); if (sub_axis == axis) { - const auto* sub_extent = inter_unpacked_axes[l]->dom.value()->extent.as(); + const auto* sub_extent = inter_unpacked_axes[l]->dom.value()->extent.as(); TVM_FFI_ICHECK(sub_extent) << "Expected Integer Extents for Offset Calculation"; divfactor = (ffi::BigInt(divfactor) * sub_extent->value).as().value(); } @@ -457,7 +457,7 @@ inline ffi::Array TransformIndex(const ffi::Array& src_index const ffi::Array& transform_rule) { sym::Analyzer ana; ffi::Array result; - std::unordered_map bind_map; + std::unordered_map bind_map; for (size_t i = 0; i < src_index.size(); ++i) { bind_map.insert_or_assign(src_axis[i]->var.get(), src_index[i]); } @@ -476,7 +476,7 @@ inline ffi::Array TransformIndex(const ffi::Array& src_index ffi::Array SBijectiveLayout::ForwardIndex(const ffi::Array& src_index) const { TVM_FFI_ICHECK(defined()) << "Cannot operate on an undefined bijective layout."; - const SBijectiveLayoutNode* self = operator->(); + const SBijectiveLayoutObj* self = operator->(); TVM_FFI_ICHECK_EQ(src_index.size(), self->src_layout->axes.size()) << "Input mismatch with layout " << self->src_layout; return TransformIndex(src_index, self->src_layout->axes, self->index_forward_rule); @@ -484,7 +484,7 @@ ffi::Array SBijectiveLayout::ForwardIndex(const ffi::Array& ffi::Array SBijectiveLayout::BackwardIndex(const ffi::Array& dst_index) const { TVM_FFI_ICHECK(defined()) << "Cannot operate on an undefined bijective layout."; - const SBijectiveLayoutNode* self = operator->(); + const SBijectiveLayoutObj* self = operator->(); TVM_FFI_ICHECK_EQ(dst_index.size(), self->dst_layout->axes.size()) << "Output mismatch with layout " << self->dst_layout; return TransformIndex(dst_index, self->dst_layout->axes, self->index_backward_rule); @@ -502,15 +502,15 @@ inline ffi::Array TransformShape(const ffi::Array& src_shape // for major-axis, bind the corresponding size // for minor-axis, simply bind it as 0, so that we can reuse forward/backward_rule, // e.g., (C * 16 + c) / 32 - std::unordered_map bind_map; + std::unordered_map bind_map; for (size_t i = 0; i < src_shape.size(); ++i) { PrimExpr orig_shape = src_shape[i]; IterVar orig_axis = src_axis[i]; auto layout = SLayout::UnpackIterVar(orig_axis); if (layout.size() != 1 || !SLayoutAxis::Get(layout[0]).IsPrimal()) { if (orig_shape.defined()) { - const auto* orig_shape_const = orig_shape.as(); - const auto* orig_axis_extent = orig_axis->dom.value()->extent.as(); + const auto* orig_shape_const = orig_shape.as(); + const auto* orig_axis_extent = orig_axis->dom.value()->extent.as(); if (orig_shape_const) { TVM_FFI_ICHECK_EQ(orig_shape_const->value, orig_axis_extent->value) << "Input shape mismatch at index " << i << ". Expected " @@ -573,14 +573,14 @@ inline ffi::Array TransformShape(const ffi::Array& src_shape ffi::Array SBijectiveLayout::ForwardShape(const ffi::Array& shape) const { TVM_FFI_ICHECK(defined()) << "Cannot operate on an undefined bijective layout."; - const SBijectiveLayoutNode* self = operator->(); + const SBijectiveLayoutObj* self = operator->(); return TransformShape(shape, self->src_layout->axes, self->dst_layout->axes, self->shape_forward_rule); } ffi::Array SBijectiveLayout::BackwardShape(const ffi::Array& shape) const { TVM_FFI_ICHECK(defined()) << "Cannot operate on an undefined bijective layout."; - const SBijectiveLayoutNode* self = operator->(); + const SBijectiveLayoutObj* self = operator->(); return TransformShape(shape, self->dst_layout->axes, self->src_layout->axes, self->shape_backward_rule); } @@ -592,7 +592,7 @@ SBijectiveLayout::SBijectiveLayout(SLayout src_layout, SLayout dst_layout) { } ffi::Optional SBijectiveLayout::Create(SLayout src_layout, SLayout dst_layout) { - auto n = ffi::make_object(std::move(src_layout), std::move(dst_layout)); + auto n = ffi::make_object(std::move(src_layout), std::move(dst_layout)); if (!GetStoreRule(&n->index_forward_rule, &n->shape_forward_rule, n->src_layout, n->dst_layout)) { return std::nullopt; } @@ -603,8 +603,8 @@ ffi::Optional SBijectiveLayout::Create(SLayout src_layout, SLa TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - SBijectiveLayoutNode::RegisterReflection(); - refl::TypeAttrDef().def( + SBijectiveLayoutObj::RegisterReflection(); + refl::TypeAttrDef().def( ffi::type_attr::kRepr.name, [](SBijectiveLayout bl, ffi::Function) -> ffi::String { return "SBijectiveLayout(" + std::string(bl->src_layout.name()) + "->" + std::string(bl->dst_layout.name()) + ")"; diff --git a/src/s_tir/ir/data_type_rewriter.cc b/src/s_tir/ir/data_type_rewriter.cc index 6334b107d3f9..7bd9da8ca8c4 100644 --- a/src/s_tir/ir/data_type_rewriter.cc +++ b/src/s_tir/ir/data_type_rewriter.cc @@ -33,11 +33,11 @@ namespace s_tir { using namespace tvm::tirx; using namespace tvm::prim; -UnchangedOr DataTypeLegalizer::Mutate_(const ForNode* op, InplaceMode inplace_mode) { +UnchangedOr DataTypeLegalizer::Mutate_(const ForObj* op, InplaceMode inplace_mode) { auto result = StmtExprMutator::Mutate_(op, inplace_mode); if (!result.IsUnchanged()) { - op = ffi::AnyView(result).as(); - TVM_FFI_ICHECK(op != nullptr) << "Expected type to be ForNode, but get " + op = ffi::AnyView(result).as(); + TVM_FFI_ICHECK(op != nullptr) << "Expected type to be ForObj, but get " << ffi::AnyView(result).GetTypeKey(); if (!op->unique()) inplace_mode = InplaceMode::kDisallow; } @@ -45,22 +45,22 @@ UnchangedOr DataTypeLegalizer::Mutate_(const ForNode* op, InplaceMode inpl Var var = e.as_or_throw(); PrimType var_ty = var->ty.as_or_throw(); if (inplace_mode == InplaceMode::kAllow) { - auto* n = const_cast(op); + auto* n = const_cast(op); n->begin = prim::cast(op->begin, var_ty); n->end = prim::cast(op->end, var_ty); n->step = prim::cast(op->step, var_ty); return result; } - auto n = ffi::make_object(*op); + auto n = ffi::make_object(*op); n->begin = prim::cast(op->begin, var_ty); n->end = prim::cast(op->end, var_ty); n->step = prim::cast(op->step, var_ty); return For(n); } -UnchangedOr DataTypeLegalizer::Mutate_(const RegionStmtNode* op, InplaceMode inplace_mode) { +UnchangedOr DataTypeLegalizer::Mutate_(const RegionStmtObj* op, InplaceMode inplace_mode) { Stmt result = StmtExprMutator::Mutate_(op, inplace_mode).ValueOrUnchanged(ffi::GetRef(op)); - const auto* region = result.as(); + const auto* region = result.as(); if (!region || !region->op.same_as(tirx::launch_thread_op())) return result; PrimVar var = region->body_params[0].as_or_throw(); PrimExpr extent = region->args[1].as_or_throw(); @@ -70,8 +70,7 @@ UnchangedOr DataTypeLegalizer::Mutate_(const RegionStmtNode* op, InplaceMo return updated; } -UnchangedOr DataTypeLegalizer::Mutate_(const prim::LetNode* op, - InplaceMode inplace_mode) { +UnchangedOr DataTypeLegalizer::Mutate_(const prim::LetObj* op, InplaceMode inplace_mode) { auto value_result = this->Mutate(op->value, inplace_mode); bool value_unchanged = value_result.UnchangedOrSameAs(op->value); PrimExpr value = std::move(value_result).ValueOrUnchanged(op->value); @@ -96,7 +95,7 @@ UnchangedOr DataTypeLegalizer::Mutate_(const prim::LetNode* op, } } -UnchangedOr DataTypeLegalizer::Mutate_(const BindNode* op, InplaceMode inplace_mode) { +UnchangedOr DataTypeLegalizer::Mutate_(const BindObj* op, InplaceMode inplace_mode) { if (!op->value.as()) { return StmtExprMutator::Mutate_(op, inplace_mode); } @@ -119,7 +118,7 @@ UnchangedOr DataTypeLegalizer::Mutate_(const BindNode* op, InplaceMode inp } } -UnchangedOr DataTypeLegalizer::Mutate_(const prim::SelectNode* op, +UnchangedOr DataTypeLegalizer::Mutate_(const prim::SelectObj* op, InplaceMode inplace_mode) { auto condition_result = this->Mutate(op->condition, inplace_mode); bool condition_unchanged = condition_result.UnchangedOrSameAs(op->condition); @@ -144,7 +143,7 @@ UnchangedOr DataTypeLegalizer::Mutate_(const prim::SelectNode* op, } } -UnchangedOr DataTypeLegalizer::Mutate_(const prim::RampNode* op, +UnchangedOr DataTypeLegalizer::Mutate_(const prim::RampObj* op, InplaceMode inplace_mode) { auto base_result = Mutate(op->base, inplace_mode); bool base_unchanged = base_result.UnchangedOrSameAs(op->base); @@ -167,7 +166,7 @@ UnchangedOr DataTypeLegalizer::Mutate_(const prim::RampNode* op, } } -UnchangedOr DataTypeLegalizer::Mutate_(const prim::BroadcastNode* op, +UnchangedOr DataTypeLegalizer::Mutate_(const prim::BroadcastObj* op, InplaceMode inplace_mode) { auto value = Mutate(op->value, inplace_mode); auto lanes = Mutate(op->lanes, inplace_mode); @@ -179,7 +178,7 @@ UnchangedOr DataTypeLegalizer::Mutate_(const prim::BroadcastNode* op, std::move(lanes).ValueOrUnchanged(op->lanes)); } -UnchangedOr DataTypeLegalizer::Mutate_(const prim::ShuffleNode* op, +UnchangedOr DataTypeLegalizer::Mutate_(const prim::ShuffleObj* op, InplaceMode inplace_mode) { auto vectors = Mutate(op->vectors, inplace_mode).as_or_throw>>(); auto indices = Mutate(op->indices, inplace_mode).as_or_throw>>(); @@ -207,29 +206,29 @@ UnchangedOr DataTypeLegalizer::Mutate_(const prim::ShuffleNode* op, } \ } -TVM_DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(prim::AddNode, operator+); -TVM_DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(prim::SubNode, operator-); -TVM_DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(prim::MulNode, operator*); -TVM_DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(prim::DivNode, div); -TVM_DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(prim::ModNode, truncmod); -TVM_DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(prim::FloorDivNode, floordiv); -TVM_DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(prim::FloorModNode, floormod); -TVM_DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(prim::MinNode, min); -TVM_DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(prim::MaxNode, max); -TVM_DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(prim::EQNode, operator==); -TVM_DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(prim::NENode, operator!=); -TVM_DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(prim::LENode, operator<=); -TVM_DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(prim::LTNode, operator<); // NOLINT(*) -TVM_DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(prim::GTNode, operator>); // NOLINT(*) -TVM_DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(prim::GENode, operator>=); - -TVM_DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(prim::BitwiseAndNode, bitwise_and); -TVM_DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(prim::BitwiseOrNode, bitwise_or); -TVM_DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(prim::BitwiseXorNode, bitwise_xor); +TVM_DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(prim::AddObj, operator+); +TVM_DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(prim::SubObj, operator-); +TVM_DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(prim::MulObj, operator*); +TVM_DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(prim::DivObj, div); +TVM_DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(prim::ModObj, truncmod); +TVM_DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(prim::FloorDivObj, floordiv); +TVM_DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(prim::FloorModObj, floormod); +TVM_DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(prim::MinObj, min); +TVM_DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(prim::MaxObj, max); +TVM_DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(prim::EQObj, operator==); +TVM_DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(prim::NEObj, operator!=); +TVM_DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(prim::LEObj, operator<=); +TVM_DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(prim::LTObj, operator<); // NOLINT(*) +TVM_DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(prim::GTObj, operator>); // NOLINT(*) +TVM_DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(prim::GEObj, operator>=); + +TVM_DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(prim::BitwiseAndObj, bitwise_and); +TVM_DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(prim::BitwiseOrObj, bitwise_or); +TVM_DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(prim::BitwiseXorObj, bitwise_xor); #undef TVM_DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH -UnchangedOr DataTypeLegalizer::Mutate_(const prim::LShiftNode* op, +UnchangedOr DataTypeLegalizer::Mutate_(const prim::LShiftObj* op, InplaceMode inplace_mode) { PrimType before_dtype = op->a.ty(); // Preserve the original operands while computing the narrowed shift. @@ -248,7 +247,7 @@ UnchangedOr DataTypeLegalizer::Mutate_(const prim::LShiftNode* op, return left_shift(lhs, rhs, op->loc); } -UnchangedOr DataTypeLegalizer::Mutate_(const prim::RShiftNode* op, +UnchangedOr DataTypeLegalizer::Mutate_(const prim::RShiftObj* op, InplaceMode inplace_mode) { PrimType before_dtype = op->a.ty(); // Preserve the original operands while computing the narrowed shift. @@ -267,22 +266,22 @@ UnchangedOr DataTypeLegalizer::Mutate_(const prim::RShiftNode* op, return right_shift(lhs, rhs, op->loc); } -UnchangedOr DataTypeLegalizer::Mutate_(const prim::BitwiseNotNode* op, +UnchangedOr DataTypeLegalizer::Mutate_(const prim::BitwiseNotObj* op, InplaceMode inplace_mode) { auto a = Mutate(op->a, inplace_mode); if (a.UnchangedOrSameAs(op->a)) return ffi::Unchanged(); return prim::BitwiseNot(std::move(a).ValueOrUnchanged(op->a), op->loc); } -UnchangedOr DataTypeLegalizer::Mutate_(const CallNode* op, InplaceMode inplace_mode) { +UnchangedOr DataTypeLegalizer::Mutate_(const CallObj* op, InplaceMode inplace_mode) { Call before = ffi::GetRef(op); // Keep the original argument dtype available for clz correction below. Expr e = StmtExprMutator::Mutate_(op, InplaceMode::kDisallow).ValueOrUnchanged(ffi::GetRef(op)); - op = e.as(); - TVM_FFI_ICHECK(op != nullptr) << "Expected type to be CallNode" + op = e.as(); + TVM_FFI_ICHECK(op != nullptr) << "Expected type to be CallObj" << ", but get " << e->GetTypeKey(); - if (!op->ty.as()) { + if (!op->ty.as()) { return e; } PrimExpr prim_e = e.as_or_throw(); @@ -314,7 +313,7 @@ UnchangedOr DataTypeLegalizer::Mutate_(const CallNode* op, InplaceMode inp return prim_e; } -UnchangedOr IndexDataTypeRewriter::Mutate_(const RegionStmtNode* op, +UnchangedOr IndexDataTypeRewriter::Mutate_(const RegionStmtObj* op, InplaceMode inplace_mode) { bool was_enabled = is_enabled_; if (op->op.same_as(tirx::launch_thread_op())) is_enabled_ = true; @@ -325,13 +324,13 @@ UnchangedOr IndexDataTypeRewriter::Mutate_(const RegionStmtNode* op, UnchangedOr IndexDataTypeRewriter::Mutate(ffi::AnyView value, InplaceMode inplace_mode) { bool is_enabled = is_enabled_; - if (value.as()) is_enabled_ = true; + if (value.as()) is_enabled_ = true; auto result = DataTypeLegalizer::Mutate(value, inplace_mode); is_enabled_ = is_enabled; return result; } -UnchangedOr IndexDataTypeRewriter::Mutate_(const TensorStoreNode* op, +UnchangedOr IndexDataTypeRewriter::Mutate_(const TensorStoreObj* op, InplaceMode inplace_mode) { TensorStore store = ffi::GetRef(op); @@ -359,7 +358,7 @@ UnchangedOr IndexDataTypeRewriter::Mutate_(const TensorStoreNode* op, return store; } -UnchangedOr IndexDataTypeRewriter::Mutate_(const TensorLoadNode* op, +UnchangedOr IndexDataTypeRewriter::Mutate_(const TensorLoadObj* op, InplaceMode inplace_mode) { TensorLoad load = ffi::GetRef(op); @@ -386,7 +385,7 @@ ffi::Array IndexDataTypeRewriter::VisitIndices(const ffi::Array IndexDataTypeRewriter::Mutate_(const IfNode* op, InplaceMode inplace_mode) { +UnchangedOr IndexDataTypeRewriter::Mutate_(const IfObj* op, InplaceMode inplace_mode) { bool is_condition = is_condition_; is_condition_ = true; auto cond_result = Mutate(op->condition, inplace_mode); @@ -412,7 +411,7 @@ UnchangedOr IndexDataTypeRewriter::Mutate_(const IfNode* op, InplaceMode i return ffi::Unchanged(); } -UnchangedOr IndexDataTypeRewriter::Mutate_(const ForNode* op, InplaceMode inplace_mode) { +UnchangedOr IndexDataTypeRewriter::Mutate_(const ForObj* op, InplaceMode inplace_mode) { bool is_enabled = is_enabled_; is_enabled_ = true; PrimVar new_loop_var = @@ -447,7 +446,7 @@ UnchangedOr IndexDataTypeRewriter::Mutate_(const ForNode* op, InplaceMode } } -UnchangedOr IndexDataTypeRewriter::Mutate_(const BindNode* op, InplaceMode inplace_mode) { +UnchangedOr IndexDataTypeRewriter::Mutate_(const BindObj* op, InplaceMode inplace_mode) { if (!op->value.as()) { return StmtExprMutator::Mutate_(op, inplace_mode); } @@ -475,14 +474,14 @@ UnchangedOr IndexDataTypeRewriter::Mutate_(const BindNode* op, InplaceMode return result; \ } -TVM_DEFINE_CMPOP_EXPR_MUTATE_WITH_TYPE_MATCH(prim::EQNode, operator==); -TVM_DEFINE_CMPOP_EXPR_MUTATE_WITH_TYPE_MATCH(prim::NENode, operator!=); -TVM_DEFINE_CMPOP_EXPR_MUTATE_WITH_TYPE_MATCH(prim::LENode, operator<=); -TVM_DEFINE_CMPOP_EXPR_MUTATE_WITH_TYPE_MATCH(prim::LTNode, operator<); // NOLINT(*) -TVM_DEFINE_CMPOP_EXPR_MUTATE_WITH_TYPE_MATCH(prim::GTNode, operator>); // NOLINT(*) -TVM_DEFINE_CMPOP_EXPR_MUTATE_WITH_TYPE_MATCH(prim::GENode, operator>=); +TVM_DEFINE_CMPOP_EXPR_MUTATE_WITH_TYPE_MATCH(prim::EQObj, operator==); +TVM_DEFINE_CMPOP_EXPR_MUTATE_WITH_TYPE_MATCH(prim::NEObj, operator!=); +TVM_DEFINE_CMPOP_EXPR_MUTATE_WITH_TYPE_MATCH(prim::LEObj, operator<=); +TVM_DEFINE_CMPOP_EXPR_MUTATE_WITH_TYPE_MATCH(prim::LTObj, operator<); // NOLINT(*) +TVM_DEFINE_CMPOP_EXPR_MUTATE_WITH_TYPE_MATCH(prim::GTObj, operator>); // NOLINT(*) +TVM_DEFINE_CMPOP_EXPR_MUTATE_WITH_TYPE_MATCH(prim::GEObj, operator>=); -UnchangedOr IndexDataTypeRewriter::Mutate_(const CallNode* op, InplaceMode inplace_mode) { +UnchangedOr IndexDataTypeRewriter::Mutate_(const CallObj* op, InplaceMode inplace_mode) { // handle if_then_else condition if (op->op.same_as(prim::if_then_else_op())) { bool is_condition = is_condition_; @@ -503,7 +502,7 @@ UnchangedOr IndexDataTypeRewriter::Mutate_(const CallNode* op, InplaceMode return Parent::Mutate_(op, inplace_mode); } -UnchangedOr IndexDataTypeRewriter::Mutate_(const prim::SelectNode* op, +UnchangedOr IndexDataTypeRewriter::Mutate_(const prim::SelectObj* op, InplaceMode inplace_mode) { bool is_condition = true; std::swap(is_condition_, is_condition); @@ -541,11 +540,11 @@ Function IndexDataTypeNormalizer::Rewrite(Function func) { // are rewritten only after every scalar replacement has been seeded. class IndexVarCollector : public IndexDataTypeNormalizer { public: - explicit IndexVarCollector(std::function collect) + explicit IndexVarCollector(std::function collect) : IndexDataTypeNormalizer(PrimType::Int(64)), collect_(std::move(collect)) {} using IndexDataTypeNormalizer::Mutate; using IndexDataTypeNormalizer::Mutate_; - UnchangedOr Mutate_(const VarNode* op, InplaceMode mode) final { + UnchangedOr Mutate_(const VarObj* op, InplaceMode mode) final { if (def_region_kind() == kTVMFFIDefRegionKindNone && is_enabled_) collect_(op); return IndexDataTypeNormalizer::Mutate_(op, mode); } @@ -554,9 +553,9 @@ Function IndexDataTypeNormalizer::Rewrite(Function func) { bool CanRewriteDType(PrimType dtype) const final { return false; } private: - std::function collect_; + std::function collect_; }; - auto seed = [this](const VarNode* var) { + auto seed = [this](const VarObj* var) { auto dtype = var->ty.as(); if (dtype && CanRewriteDType(dtype.value()) && dtype.value() != target_data_type_ && VarRemapGet(ffi::AnyView(var)) == nullptr) { @@ -578,13 +577,13 @@ Function IndexDataTypeNormalizer::Rewrite(Function func) { return Mutate(param).ValueOrUnchanged(param).as_or_throw(); }); }); - FunctionNode* new_func = func.CopyOnWrite(); + FunctionObj* new_func = func.CopyOnWrite(); new_func->params = std::move(params); new_func->body = Mutate(new_func->body).ValueOrUnchanged(new_func->body); return func; } -UnchangedOr IndexDataTypeNormalizer::Mutate_(const SBlockRealizeNode* op, +UnchangedOr IndexDataTypeNormalizer::Mutate_(const SBlockRealizeObj* op, InplaceMode inplace_mode) { bool is_condition = this->is_condition_; this->is_condition_ = true; @@ -615,7 +614,7 @@ UnchangedOr IndexDataTypeNormalizer::Mutate_(const SBlockRealizeNode* op, } } -UnchangedOr IndexDataTypeNormalizer::Mutate_(const SBlockNode* op, InplaceMode inplace_mode) { +UnchangedOr IndexDataTypeNormalizer::Mutate_(const SBlockObj* op, InplaceMode inplace_mode) { auto new_alloc_tensors = this->WithDefRegionKind(kTVMFFIDefRegionKindSimple, [&] { return this->Mutate(op->alloc_tensors, inplace_mode) .as_or_throw>>() @@ -652,7 +651,7 @@ UnchangedOr IndexDataTypeNormalizer::Mutate_(const SBlockNode* op, Inplace !new_writes.same_as(op->writes) || !new_iter_vars.same_as(op->iter_vars) || !new_annotations.same_as(op->annotations)) { SBlock new_block = ffi::GetRef(op); - SBlockNode* n = new_block.CopyOnWrite(); + SBlockObj* n = new_block.CopyOnWrite(); n->alloc_tensors = std::move(new_alloc_tensors); n->match_tensors = std::move(new_match_tensors); n->reads = std::move(new_reads); @@ -674,7 +673,7 @@ ffi::Map IndexDataTypeNormalizer::VisitBlockAnnotations( if (obj == nullptr) { return obj; } - if (auto var = obj.as(); var && var.value()->ty.as()) { + if (auto var = obj.as(); var && var.value()->ty.as()) { TensorVar tensor = var.value().as_or_throw(); if (TensorVar new_tensor = this->Mutate(tensor, InplaceMode::kDisallow) .as_or_throw>() @@ -712,7 +711,7 @@ IterVar IndexDataTypeNormalizer::VisitIterVar(const IterVar& iter_var) { if (!new_var.same_as(iter_var->var) || !min.same_as(iter_var->dom.value()->begin) || !extent.same_as(iter_var->dom.value()->extent)) { IterVar new_iter_var = iter_var; - IterVarNode* n = new_iter_var.CopyOnWrite(); + IterVarObj* n = new_iter_var.CopyOnWrite(); n->var = std::move(new_var); n->ty = n->var.ty(); n->dom = Range::FromMinExtent(min, extent); @@ -748,7 +747,7 @@ bool IndexDataTypeNormalizer::CanRewriteDType(PrimType dtype) const { return dtype.code() == DLDataTypeCode::kDLInt && dtype.bits() >= 32; } -UnchangedOr IndexDataTypeNormalizer::Mutate_(const IntImmNode* op, +UnchangedOr IndexDataTypeNormalizer::Mutate_(const IntImmObj* op, InplaceMode inplace_mode) { if (is_enabled_ && CanRewriteDType(op->ty.as_or_throw())) { TVM_FFI_ICHECK_LE(op->value, max_value(target_data_type_).as_or_throw()->value); @@ -757,7 +756,7 @@ UnchangedOr IndexDataTypeNormalizer::Mutate_(const IntImmNode* op, return ffi::Unchanged(); } -UnchangedOr IndexDataTypeNormalizer::Mutate_(const prim::CastNode* op, +UnchangedOr IndexDataTypeNormalizer::Mutate_(const prim::CastObj* op, InplaceMode inplace_mode) { // Unwrap the cast only when the dtype of this cast is integer dtype. // When the dtype of this cast is not integer dtype, it means that this cast diff --git a/src/s_tir/ir/data_type_rewriter.h b/src/s_tir/ir/data_type_rewriter.h index 2fc3cd7d6f1f..e5ab41837f5c 100644 --- a/src/s_tir/ir/data_type_rewriter.h +++ b/src/s_tir/ir/data_type_rewriter.h @@ -50,36 +50,36 @@ class DataTypeLegalizer : public tirx::StmtExprMutator { protected: explicit DataTypeLegalizer(const VTable* vtable) : tirx::StmtExprMutator(vtable) {} - UnchangedOr Mutate_(const ForNode* op, InplaceMode inplace_mode) override; - UnchangedOr Mutate_(const RegionStmtNode* op, InplaceMode inplace_mode) override; - UnchangedOr Mutate_(const BindNode* op, InplaceMode inplace_mode) override; - UnchangedOr Mutate_(const CallNode* op, InplaceMode inplace_mode) override; - UnchangedOr Mutate_(const prim::SelectNode* op, InplaceMode inplace_mode) override; - UnchangedOr Mutate_(const prim::RampNode* op, InplaceMode inplace_mode) override; - UnchangedOr Mutate_(const prim::BroadcastNode* op, InplaceMode inplace_mode) override; - UnchangedOr Mutate_(const prim::ShuffleNode* op, InplaceMode inplace_mode) override; - UnchangedOr Mutate_(const prim::LShiftNode* op, InplaceMode inplace_mode) override; - UnchangedOr Mutate_(const prim::RShiftNode* op, InplaceMode inplace_mode) override; - UnchangedOr Mutate_(const prim::BitwiseAndNode* op, InplaceMode inplace_mode) override; - UnchangedOr Mutate_(const prim::BitwiseOrNode* op, InplaceMode inplace_mode) override; - UnchangedOr Mutate_(const prim::BitwiseXorNode* op, InplaceMode inplace_mode) override; - UnchangedOr Mutate_(const prim::BitwiseNotNode* op, InplaceMode inplace_mode) override; - UnchangedOr Mutate_(const prim::AddNode* op, InplaceMode inplace_mode) override; - UnchangedOr Mutate_(const prim::SubNode* op, InplaceMode inplace_mode) override; - UnchangedOr Mutate_(const prim::MulNode* op, InplaceMode inplace_mode) override; - UnchangedOr Mutate_(const prim::DivNode* op, InplaceMode inplace_mode) override; - UnchangedOr Mutate_(const prim::ModNode* op, InplaceMode inplace_mode) override; - UnchangedOr Mutate_(const prim::FloorDivNode* op, InplaceMode inplace_mode) override; - UnchangedOr Mutate_(const prim::FloorModNode* op, InplaceMode inplace_mode) override; - UnchangedOr Mutate_(const prim::MinNode* op, InplaceMode inplace_mode) override; - UnchangedOr Mutate_(const prim::MaxNode* op, InplaceMode inplace_mode) override; - UnchangedOr Mutate_(const prim::EQNode* op, InplaceMode inplace_mode) override; - UnchangedOr Mutate_(const prim::NENode* op, InplaceMode inplace_mode) override; - UnchangedOr Mutate_(const prim::LTNode* op, InplaceMode inplace_mode) override; - UnchangedOr Mutate_(const prim::LENode* op, InplaceMode inplace_mode) override; - UnchangedOr Mutate_(const prim::GTNode* op, InplaceMode inplace_mode) override; - UnchangedOr Mutate_(const prim::GENode* op, InplaceMode inplace_mode) override; - UnchangedOr Mutate_(const prim::LetNode* op, InplaceMode inplace_mode) override; + UnchangedOr Mutate_(const ForObj* op, InplaceMode inplace_mode) override; + UnchangedOr Mutate_(const RegionStmtObj* op, InplaceMode inplace_mode) override; + UnchangedOr Mutate_(const BindObj* op, InplaceMode inplace_mode) override; + UnchangedOr Mutate_(const CallObj* op, InplaceMode inplace_mode) override; + UnchangedOr Mutate_(const prim::SelectObj* op, InplaceMode inplace_mode) override; + UnchangedOr Mutate_(const prim::RampObj* op, InplaceMode inplace_mode) override; + UnchangedOr Mutate_(const prim::BroadcastObj* op, InplaceMode inplace_mode) override; + UnchangedOr Mutate_(const prim::ShuffleObj* op, InplaceMode inplace_mode) override; + UnchangedOr Mutate_(const prim::LShiftObj* op, InplaceMode inplace_mode) override; + UnchangedOr Mutate_(const prim::RShiftObj* op, InplaceMode inplace_mode) override; + UnchangedOr Mutate_(const prim::BitwiseAndObj* op, InplaceMode inplace_mode) override; + UnchangedOr Mutate_(const prim::BitwiseOrObj* op, InplaceMode inplace_mode) override; + UnchangedOr Mutate_(const prim::BitwiseXorObj* op, InplaceMode inplace_mode) override; + UnchangedOr Mutate_(const prim::BitwiseNotObj* op, InplaceMode inplace_mode) override; + UnchangedOr Mutate_(const prim::AddObj* op, InplaceMode inplace_mode) override; + UnchangedOr Mutate_(const prim::SubObj* op, InplaceMode inplace_mode) override; + UnchangedOr Mutate_(const prim::MulObj* op, InplaceMode inplace_mode) override; + UnchangedOr Mutate_(const prim::DivObj* op, InplaceMode inplace_mode) override; + UnchangedOr Mutate_(const prim::ModObj* op, InplaceMode inplace_mode) override; + UnchangedOr Mutate_(const prim::FloorDivObj* op, InplaceMode inplace_mode) override; + UnchangedOr Mutate_(const prim::FloorModObj* op, InplaceMode inplace_mode) override; + UnchangedOr Mutate_(const prim::MinObj* op, InplaceMode inplace_mode) override; + UnchangedOr Mutate_(const prim::MaxObj* op, InplaceMode inplace_mode) override; + UnchangedOr Mutate_(const prim::EQObj* op, InplaceMode inplace_mode) override; + UnchangedOr Mutate_(const prim::NEObj* op, InplaceMode inplace_mode) override; + UnchangedOr Mutate_(const prim::LTObj* op, InplaceMode inplace_mode) override; + UnchangedOr Mutate_(const prim::LEObj* op, InplaceMode inplace_mode) override; + UnchangedOr Mutate_(const prim::GTObj* op, InplaceMode inplace_mode) override; + UnchangedOr Mutate_(const prim::GEObj* op, InplaceMode inplace_mode) override; + UnchangedOr Mutate_(const prim::LetObj* op, InplaceMode inplace_mode) override; /*! \brief Whether to clamp shift amounts after narrowing signed integers. */ virtual bool ShouldClampShiftAmounts() const { return false; } @@ -104,22 +104,22 @@ class IndexDataTypeRewriter : public DataTypeLegalizer { explicit IndexDataTypeRewriter(const VTable* vtable) : DataTypeLegalizer(vtable) {} using Parent = DataTypeLegalizer; UnchangedOr Mutate(ffi::AnyView value, InplaceMode inplace_mode) override; - UnchangedOr Mutate_(const TensorStoreNode* op, InplaceMode inplace_mode) override; - UnchangedOr Mutate_(const RegionStmtNode* op, InplaceMode inplace_mode) override; - UnchangedOr Mutate_(const TensorLoadNode* op, InplaceMode inplace_mode) override; + UnchangedOr Mutate_(const TensorStoreObj* op, InplaceMode inplace_mode) override; + UnchangedOr Mutate_(const RegionStmtObj* op, InplaceMode inplace_mode) override; + UnchangedOr Mutate_(const TensorLoadObj* op, InplaceMode inplace_mode) override; ffi::Array VisitIndices(const ffi::Array& indices, InplaceMode inplace_mode); - UnchangedOr Mutate_(const IfNode* op, InplaceMode inplace_mode) override; - UnchangedOr Mutate_(const BindNode* op, InplaceMode inplace_mode) override; - UnchangedOr Mutate_(const CallNode* op, InplaceMode inplace_mode) override; - UnchangedOr Mutate_(const ForNode* op, InplaceMode inplace_mode) override; - UnchangedOr Mutate_(const prim::EQNode* op, InplaceMode inplace_mode) override; - UnchangedOr Mutate_(const prim::NENode* op, InplaceMode inplace_mode) override; - UnchangedOr Mutate_(const prim::LTNode* op, InplaceMode inplace_mode) override; - UnchangedOr Mutate_(const prim::LENode* op, InplaceMode inplace_mode) override; - UnchangedOr Mutate_(const prim::GTNode* op, InplaceMode inplace_mode) override; - UnchangedOr Mutate_(const prim::GENode* op, InplaceMode inplace_mode) override; - - UnchangedOr Mutate_(const prim::SelectNode* op, InplaceMode inplace_mode) override; + UnchangedOr Mutate_(const IfObj* op, InplaceMode inplace_mode) override; + UnchangedOr Mutate_(const BindObj* op, InplaceMode inplace_mode) override; + UnchangedOr Mutate_(const CallObj* op, InplaceMode inplace_mode) override; + UnchangedOr Mutate_(const ForObj* op, InplaceMode inplace_mode) override; + UnchangedOr Mutate_(const prim::EQObj* op, InplaceMode inplace_mode) override; + UnchangedOr Mutate_(const prim::NEObj* op, InplaceMode inplace_mode) override; + UnchangedOr Mutate_(const prim::LTObj* op, InplaceMode inplace_mode) override; + UnchangedOr Mutate_(const prim::LEObj* op, InplaceMode inplace_mode) override; + UnchangedOr Mutate_(const prim::GTObj* op, InplaceMode inplace_mode) override; + UnchangedOr Mutate_(const prim::GEObj* op, InplaceMode inplace_mode) override; + + UnchangedOr Mutate_(const prim::SelectObj* op, InplaceMode inplace_mode) override; // indicator of index expr to rewrite bool is_enabled_{false}; @@ -138,12 +138,12 @@ class IndexDataTypeNormalizer : public IndexDataTypeRewriter { : Parent(GlobalVTable()), target_data_type_(std::move(target_data_type)) {} tirx::Function Rewrite(tirx::Function func); - virtual UnchangedOr Mutate_(const SBlockNode* op, InplaceMode inplace_mode); - virtual UnchangedOr Mutate_(const SBlockRealizeNode* op, InplaceMode inplace_mode); + virtual UnchangedOr Mutate_(const SBlockObj* op, InplaceMode inplace_mode); + virtual UnchangedOr Mutate_(const SBlockRealizeObj* op, InplaceMode inplace_mode); protected: - UnchangedOr Mutate_(const IntImmNode* op, InplaceMode inplace_mode) override; - UnchangedOr Mutate_(const prim::CastNode* op, InplaceMode inplace_mode) override; + UnchangedOr Mutate_(const IntImmObj* op, InplaceMode inplace_mode) override; + UnchangedOr Mutate_(const prim::CastObj* op, InplaceMode inplace_mode) override; /*! \brief Specifies which data type we can rewrite */ virtual bool CanRewriteDType(PrimType dtype) const; @@ -152,8 +152,8 @@ class IndexDataTypeNormalizer : public IndexDataTypeRewriter { static void InitVTable(VTable* table) { Parent::InitVTable(table); - SetDispatch(table); - SetDispatch(table); + SetDispatch(table); + SetDispatch(table); } static const VTable* GlobalVTable() { static const VTable table = [] { diff --git a/src/s_tir/ir/ir_mutator_with_analyzer.cc b/src/s_tir/ir/ir_mutator_with_analyzer.cc index 84a3a5adac3f..fe7b3314c76c 100644 --- a/src/s_tir/ir/ir_mutator_with_analyzer.cc +++ b/src/s_tir/ir/ir_mutator_with_analyzer.cc @@ -17,20 +17,20 @@ * under the License. */ -#include +#include namespace tvm { namespace s_tir { using namespace tirx; -UnchangedOr StmtExprMutatorWithAnalyzer::Mutate_(const SBlockNode* op, - InplaceMode inplace_mode) { - return constraint_scope_.WithNewScope([&]() -> UnchangedOr { - for (const auto& iter_var : op->iter_vars) { - analyzer_->Bind(iter_var->var, iter_var->dom.value()); - iter_vars_.Set(iter_var->var, iter_var->dom.value()); - } - return s_tir::StmtExprMutator::MutateBlock(this, op, inplace_mode); - }); +UnchangedOr StmtExprMutatorWithConstraintContext::Mutate_(const SBlockObj* op, + InplaceMode inplace_mode) { + return WithConstraintScope([&]() -> ffi::Expected> { + for (const auto& iter_var : op->iter_vars) { + AppendConstraint(sym::RangeConstraint(iter_var->var, iter_var->dom.value())); + } + return s_tir::StmtExprMutator::MutateBlock(this, op, inplace_mode); + }) + .value(); } } // namespace s_tir diff --git a/src/s_tir/ir/ir_visitor_with_analyzer.cc b/src/s_tir/ir/ir_visitor_with_analyzer.cc index 02f0ddb40d40..68c7b025194e 100644 --- a/src/s_tir/ir/ir_visitor_with_analyzer.cc +++ b/src/s_tir/ir/ir_visitor_with_analyzer.cc @@ -17,18 +17,19 @@ * under the License. */ -#include +#include namespace tvm { namespace s_tir { using namespace tirx; -ffi::Optional StmtExprVisitorWithAnalyzer::Visit_(const SBlockNode* op) { - return constraint_scope_.WithNewScope([&]() -> ffi::Optional { - for (const auto& iter_var : op->iter_vars) { - analyzer_->Bind(iter_var->var, iter_var->dom.value()); - } - return s_tir::StmtExprVisitor::VisitBlock(this, op); - }); +ffi::Optional StmtExprVisitorWithConstraintContext::Visit_(const SBlockObj* op) { + return WithConstraintScope([&]() -> ffi::Expected> { + for (const auto& iter_var : op->iter_vars) { + AppendConstraint(sym::RangeConstraint(iter_var->var, iter_var->dom.value())); + } + return s_tir::StmtExprVisitor::VisitBlock(this, op); + }) + .value(); } } // namespace s_tir diff --git a/src/s_tir/iter_var.cc b/src/s_tir/iter_var.cc index c368f8a4a1d5..40e83126db30 100644 --- a/src/s_tir/iter_var.cc +++ b/src/s_tir/iter_var.cc @@ -36,8 +36,8 @@ namespace { TVM_FFI_INLINE ffi::Expected> IterVarVisit( ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept { // skips: iter_type, thread_tag - const IterVarNode* self = - ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value); + const IterVarObj* self = + ffi::detail::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value); TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->dom)); TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->WithDefRegionKind( kTVMFFIDefRegionKindSimple, [&]() { return visitor->VisitExpected(self->var); })); @@ -47,8 +47,8 @@ TVM_FFI_INLINE ffi::Expected> IterVarVisit( TVM_FFI_INLINE ffi::Expected> IterVarMutate( ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept { // skips: iter_type, thread_tag - const IterVarNode* self = - ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value); + const IterVarObj* self = + ffi::detail::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value); TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr>, mapped_dom, mutator->MutateExpected(self->dom)); TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, mapped_var, @@ -58,7 +58,7 @@ TVM_FFI_INLINE ffi::Expected> IterVarMutate( if (mapped_dom.UnchangedOrSameAs(self->dom) && mapped_var.UnchangedOrSameAs(self->var)) { return ffi::Unchanged(); } - ffi::ObjectPtr copy = ffi::make_object(*self); + ffi::ObjectPtr copy = ffi::make_object(*self); copy->dom = std::move(mapped_dom).ValueOrUnchanged(std::move(copy->dom)); copy->var = std::move(mapped_var).ValueOrUnchanged(std::move(copy->var)); copy->ty = copy->var.ty(); @@ -68,8 +68,8 @@ TVM_FFI_INLINE ffi::Expected> IterVarMutate( TVM_FFI_INLINE ffi::Expected> IterVarMaybeInplaceMutate( ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept { // skips: iter_type, thread_tag - IterVarNode* self = const_cast( - ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value)); + IterVarObj* self = const_cast( + ffi::detail::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value)); TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr>, mapped_dom, mutator->MutateExpected(self->dom, ffi::InplaceMode::kAllow)); TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, mapped_var, @@ -89,7 +89,7 @@ TVM_FFI_INLINE ffi::Expected> IterVarMaybeInplaceMuta IterVar::IterVar(ffi::Optional dom, PrimVar var, IterVarType t, ffi::String thread_tag, Location loc) : PrimExpr(ffi::UnsafeInit{}) { - ffi::ObjectPtr n = ffi::make_object(var); + ffi::ObjectPtr n = ffi::make_object(var); if (dom.has_value() && dom.value()->extent.defined()) { PrimType extent_ty = dom.value()->extent.ty(); PrimType var_ty = var.ty(); @@ -110,8 +110,8 @@ IterVar::IterVar(ffi::Optional dom, PrimVar var, IterVarType t, ffi::Stri TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - IterVarNode::RegisterReflection(); - refl::TypeAttrDef() + IterVarObj::RegisterReflection(); + refl::TypeAttrDef() .attr(ffi::type_attr::kStructuralVisit, ffi::FStructuralVisit::FromNative<&IterVarVisit>()) .attr(ffi::type_attr::kStructuralMutate, ffi::FStructuralMutate::FromNative<&IterVarMutate>()) .attr(ffi::type_attr::kStructuralMaybeInplaceMutate, diff --git a/src/s_tir/meta_schedule/arg_info.cc b/src/s_tir/meta_schedule/arg_info.cc index 6c2c99ed322b..fc13f2edc7c3 100644 --- a/src/s_tir/meta_schedule/arg_info.cc +++ b/src/s_tir/meta_schedule/arg_info.cc @@ -39,12 +39,12 @@ namespace meta_schedule { inline tirx::Function FindEntryFunc(const IRModule& mod) { // Priority 1: Function marked as `tirx::attr::kIsEntryFunc.name` int num_function = 0; - const tirx::FunctionNode* main_func = nullptr; - const tirx::FunctionNode* last_func = nullptr; + const tirx::FunctionObj* main_func = nullptr; + const tirx::FunctionObj* last_func = nullptr; for (const auto& kv : mod->functions) { GlobalVar gv = kv.first; BaseFunc base_func = kv.second; - if (const auto* func = base_func.as()) { + if (const auto* func = base_func.as()) { last_func = func; if (func->HasNonzeroAttr(tvm::tirx::attr::kIsEntryFunc.name)) { return ffi::GetRef(func); diff --git a/src/s_tir/meta_schedule/database/database_utils.cc b/src/s_tir/meta_schedule/database/database_utils.cc index 541b7745311f..6620ee48bc19 100644 --- a/src/s_tir/meta_schedule/database/database_utils.cc +++ b/src/s_tir/meta_schedule/database/database_utils.cc @@ -84,7 +84,7 @@ void JSONDumps(Any json_obj, std::ostringstream& os) { JSONDumps(kv.second, os); } os << "}"; - } else if (json_obj.as()) { + } else if (json_obj.as()) { JSONDumps(ffi::String(ffi::json::Stringify( ffi::ToJSONGraph(json_obj, ffi::json::Object{{"tvm_version", TVM_VERSION}}), /*indent=*/2)), 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 e42aaa826d3e..a872f9fb5a78 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 @@ -47,7 +47,7 @@ using MultiIndex = std::vector; /*! \brief Vector of int64_t */ using IntVec = std::vector; /*! \brief Vector of for loops */ -using ForVec = std::vector; +using ForVec = std::vector; /*! * \brief An unordered_map for (for, tensor) => V @@ -55,7 +55,7 @@ using ForVec = std::vector; */ template using ForTensorMap = - std::unordered_map>; /*! \brief Given x, compute log2(|x| + 1) */ @@ -74,7 +74,7 @@ std::vector GetTensorShape(const TensorVar& tensor, sym::AnalyzerObj* a std::vector result; result.reserve(ndim); for (const PrimExpr& i : tensor.shape()) { - if (const IntImmNode* int_imm = i.as()) { + if (const IntImmObj* int_imm = i.as()) { result.push_back(static_cast(int_imm->value)); continue; } @@ -93,7 +93,7 @@ std::vector GetTensorShape(const TensorVar& tensor, sym::AnalyzerObj* a * \param loop The loop to be checked * \return The value of `auto_unroll_max_step` if it exists, or -1 if it does not exist */ -int64_t GetAutoUnrollMaxStep(const ForNode* loop) { +int64_t GetAutoUnrollMaxStep(const ForObj* loop) { if (ffi::Optional auto_unroll = GetAnn(loop, tvm::tirx::attr::kAutoUnrollMaxStep.name)) { return static_cast(auto_unroll.value()->value); @@ -113,7 +113,7 @@ int64_t GetAutoUnrollMaxStep(const ForNode* loop) { int64_t FirstLoopExtent(const ForVec& loops, int64_t default_value) { if (!loops.empty()) { PrimExpr extent_imm_expr = Range::GetExtentFromBeginEnd(loops[0]->begin, loops[0]->end); - const auto* extent_imm = extent_imm_expr.as(); + const auto* extent_imm = extent_imm_expr.as(); if (auto extent = extent_imm ? extent_imm->value.as() : std::nullopt; extent.has_value()) { return *extent; @@ -178,13 +178,13 @@ int64_t GetVarStride(const std::vector& multi_indices, const IntVec& : var(var), stride(0), visited_var(false), visited_add(false), visited_mul(false) {} private: - ffi::Optional Visit_(const MulNode* node) override { + ffi::Optional Visit_(const MulObj* node) override { TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(node)); if (visited_var && !visited_add) { - if (const auto* a = node->a.as()) { + if (const auto* a = node->a.as()) { visited_mul = true; stride = static_cast(a->value); - } else if (const auto* b = node->b.as()) { + } else if (const auto* b = node->b.as()) { visited_mul = true; stride = static_cast(b->value); } @@ -192,7 +192,7 @@ int64_t GetVarStride(const std::vector& multi_indices, const IntVec& return std::nullopt; } - ffi::Optional Visit_(const AddNode* node) override { + ffi::Optional Visit_(const AddObj* node) override { TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(node)); if (visited_var && !visited_mul) { visited_add = true; @@ -201,7 +201,7 @@ int64_t GetVarStride(const std::vector& multi_indices, const IntVec& return std::nullopt; } - ffi::Optional Visit_(const VarNode* node) override { + ffi::Optional Visit_(const VarObj* node) override { if (node == var.get()) { visited_var = true; stride = 2; @@ -291,7 +291,7 @@ Pass SimplifyForFeatureExtraction(bool normalize_thread_bindings = false) { return result.has_value() ? result.value()->value.cast() : false; } - UnchangedOr Mutate_(const SelectNode* node, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const SelectObj* node, InplaceMode inplace_mode) final { if (normalize_thread_bindings_) return StmtExprMutator::Mutate_(node, inplace_mode); if (HasTensorLoad(node->true_value) || HasTensorLoad(node->false_value) || HasTensorLoad(node->condition)) { @@ -300,7 +300,7 @@ Pass SimplifyForFeatureExtraction(bool normalize_thread_bindings = false) { return prim::MakeConst(node->ty.as_or_throw(), 1.0); } - UnchangedOr Mutate_(const ForNode* loop, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const ForObj* loop, InplaceMode inplace_mode) final { auto mutate_retained_loop = [&]() -> UnchangedOr { if (!normalize_thread_bindings_) return StmtExprMutator::Mutate_(loop, inplace_mode); auto annotations = Mutate(loop->annotations, inplace_mode) @@ -326,7 +326,7 @@ Pass SimplifyForFeatureExtraction(bool normalize_thread_bindings = false) { return result; } for (auto it = thread_bindings_.rbegin(); it != thread_bindings_.rend() && *it; ++it) { - const ForNode* outer = *it; + const ForObj* outer = *it; if (tvm::s_tir::GetThreadBinding(outer).value() != tag) continue; // Feature counts describe hardware work, not repeated lexical names // for one axis. Normalize only nested aliases in this analysis input. @@ -367,7 +367,7 @@ Pass SimplifyForFeatureExtraction(bool normalize_thread_bindings = false) { } } - UnchangedOr Mutate_(const SBlockNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const SBlockObj* op, InplaceMode inplace_mode) final { if (!normalize_thread_bindings_) return StmtExprMutator::Mutate_(op, inplace_mode); auto annotations = Mutate(op->annotations, inplace_mode) .as_or_throw>>(); @@ -379,7 +379,7 @@ Pass SimplifyForFeatureExtraction(bool normalize_thread_bindings = false) { return block; } - UnchangedOr Mutate_(const RegionStmtNode* region, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const RegionStmtObj* region, InplaceMode inplace_mode) final { thread_bindings_.push_back(nullptr); auto result = StmtExprMutator::Mutate_(region, inplace_mode); thread_bindings_.pop_back(); @@ -387,11 +387,11 @@ Pass SimplifyForFeatureExtraction(bool normalize_thread_bindings = false) { } bool normalize_thread_bindings_; - std::vector thread_bindings_; + std::vector thread_bindings_; }; auto pass_func = [normalize_thread_bindings](Function f, IRModule m, PassContext ctx) { if (!f->body.has_value()) return f; - FunctionNode* n = f.CopyOnWrite(); + FunctionObj* n = f.CopyOnWrite(); n->body = Simplifier::Run(std::move(n->body).value(), normalize_thread_bindings); return f; }; @@ -443,9 +443,9 @@ struct LoopNest { * \param auto_unroll_attr The auto unroll attribute of the loop * \return A list of for loops that the loop is bound to */ - ForVec* Push(const ForNode* loop, int64_t* auto_unroll_attr) { + ForVec* Push(const ForObj* loop, int64_t* auto_unroll_attr) { PrimExpr extent_imm_expr = Range::GetExtentFromBeginEnd(loop->begin, loop->end); - const auto* extent_imm = extent_imm_expr.as(); + const auto* extent_imm = extent_imm_expr.as(); if (auto extent = extent_imm ? extent_imm->value.as() : std::nullopt; extent.has_value()) { this->prod *= *extent; @@ -493,7 +493,7 @@ struct LoopNest { * \param ref_loops The list of for loops that the loop is bound to * \param auto_unroll_attr The auto unroll attribute of the loop */ - void Pop(const ForNode* loop, ForVec* ref_loops, int auto_unroll_attr) { + void Pop(const ForObj* loop, ForVec* ref_loops, int auto_unroll_attr) { if (ref_loops) { ref_loops->pop_back(); } @@ -501,7 +501,7 @@ struct LoopNest { this->auto_unroll.pop_back(); } PrimExpr extent_imm_expr = Range::GetExtentFromBeginEnd(loop->begin, loop->end); - const auto* extent_imm = extent_imm_expr.as(); + const auto* extent_imm = extent_imm_expr.as(); if (auto extent = extent_imm ? extent_imm->value.as() : std::nullopt; extent.has_value()) { this->prod /= *extent; @@ -541,7 +541,7 @@ struct Feature { static constexpr int64_t kCount = 16; ArithOps() = default; - ArithOps(const TensorStoreNode* store, int64_t prod_loop_extent); + ArithOps(const TensorStoreObj* store, int64_t prod_loop_extent); void Export(std::vector* v) const { double vs[] = { @@ -609,7 +609,7 @@ struct Feature { static constexpr int64_t kCount = ArithOps::kCount + ForKindFeature::kCount * 3 + 8; - explicit Feature(const TensorStoreNode* store, const LoopNest& loop_nest, bool is_gpu) + explicit Feature(const TensorStoreObj* store, const LoopNest& loop_nest, bool is_gpu) : arith_ops(store, loop_nest.prod), vectorize(loop_nest.vectorize), unroll(loop_nest.unroll), @@ -641,7 +641,7 @@ struct Feature { } }; -Feature::ArithOps::ArithOps(const TensorStoreNode* store, int64_t prod_loop_extent) { +Feature::ArithOps::ArithOps(const TensorStoreObj* store, int64_t prod_loop_extent) { class ArithOpCounter : public StmtExprVisitor { public: using StmtExprVisitor::Visit_; @@ -660,35 +660,35 @@ Feature::ArithOps::ArithOps(const TensorStoreNode* store, int64_t prod_loop_exte } \ return StmtExprVisitor::Visit_(op); \ } - TVM_FEATURE_SIMPLE(AndNode, bool_op); - TVM_FEATURE_SIMPLE(OrNode, bool_op); - TVM_FEATURE_SIMPLE(NotNode, bool_op); - TVM_FEATURE_SIMPLE(SelectNode, select_op); - TVM_FEATURE_SIMPLE(prim::LShiftNode, int_math_func); - TVM_FEATURE_SIMPLE(prim::RShiftNode, int_math_func); - TVM_FEATURE_SIMPLE(prim::BitwiseAndNode, int_math_func); - TVM_FEATURE_SIMPLE(prim::BitwiseOrNode, int_math_func); - TVM_FEATURE_SIMPLE(prim::BitwiseXorNode, int_math_func); - TVM_FEATURE_SIMPLE(prim::BitwiseNotNode, int_math_func); - TVM_FEATURE_BINARY(AddNode, float_add_sub, int_add_sub); - TVM_FEATURE_BINARY(SubNode, float_add_sub, int_add_sub); - TVM_FEATURE_BINARY(MulNode, float_mul, int_mul); - TVM_FEATURE_BINARY(DivNode, float_div_mod, int_div_mod); - TVM_FEATURE_BINARY(ModNode, float_div_mod, int_div_mod); - TVM_FEATURE_BINARY(FloorDivNode, float_div_mod, int_div_mod); - TVM_FEATURE_BINARY(FloorModNode, float_div_mod, int_div_mod); - TVM_FEATURE_BINARY(MaxNode, float_cmp, int_cmp); - TVM_FEATURE_BINARY(MinNode, float_cmp, int_cmp); - TVM_FEATURE_BINARY(EQNode, float_cmp, int_cmp); - TVM_FEATURE_BINARY(NENode, float_cmp, int_cmp); - TVM_FEATURE_BINARY(LTNode, float_cmp, int_cmp); - TVM_FEATURE_BINARY(LENode, float_cmp, int_cmp); - TVM_FEATURE_BINARY(GTNode, float_cmp, int_cmp); - TVM_FEATURE_BINARY(GENode, float_cmp, int_cmp); + TVM_FEATURE_SIMPLE(AndObj, bool_op); + TVM_FEATURE_SIMPLE(OrObj, bool_op); + TVM_FEATURE_SIMPLE(NotObj, bool_op); + TVM_FEATURE_SIMPLE(SelectObj, select_op); + TVM_FEATURE_SIMPLE(prim::LShiftObj, int_math_func); + TVM_FEATURE_SIMPLE(prim::RShiftObj, int_math_func); + TVM_FEATURE_SIMPLE(prim::BitwiseAndObj, int_math_func); + TVM_FEATURE_SIMPLE(prim::BitwiseOrObj, int_math_func); + TVM_FEATURE_SIMPLE(prim::BitwiseXorObj, int_math_func); + TVM_FEATURE_SIMPLE(prim::BitwiseNotObj, int_math_func); + TVM_FEATURE_BINARY(AddObj, float_add_sub, int_add_sub); + TVM_FEATURE_BINARY(SubObj, float_add_sub, int_add_sub); + TVM_FEATURE_BINARY(MulObj, float_mul, int_mul); + TVM_FEATURE_BINARY(DivObj, float_div_mod, int_div_mod); + TVM_FEATURE_BINARY(ModObj, float_div_mod, int_div_mod); + TVM_FEATURE_BINARY(FloorDivObj, float_div_mod, int_div_mod); + TVM_FEATURE_BINARY(FloorModObj, float_div_mod, int_div_mod); + TVM_FEATURE_BINARY(MaxObj, float_cmp, int_cmp); + TVM_FEATURE_BINARY(MinObj, float_cmp, int_cmp); + TVM_FEATURE_BINARY(EQObj, float_cmp, int_cmp); + TVM_FEATURE_BINARY(NEObj, float_cmp, int_cmp); + TVM_FEATURE_BINARY(LTObj, float_cmp, int_cmp); + TVM_FEATURE_BINARY(LEObj, float_cmp, int_cmp); + TVM_FEATURE_BINARY(GTObj, float_cmp, int_cmp); + TVM_FEATURE_BINARY(GEObj, float_cmp, int_cmp); #undef TVM_FEATURE_BINARY #undef TVM_FEATURE_SIMPLE - ffi::Optional Visit_(const CallNode* op) final { + ffi::Optional Visit_(const CallObj* op) final { auto result_type = op->ty.as(); if (!result_type) { return StmtExprVisitor::Visit_(op); @@ -731,16 +731,16 @@ Feature::ForKindFeature::ForKindFeature(const ForVec& loops) { } else { PrimExpr last_loop_extent_imm_expr = Range::GetExtentFromBeginEnd(loops.back()->begin, loops.back()->end); - const auto* last_loop_extent_imm = last_loop_extent_imm_expr.as(); + const auto* last_loop_extent_imm = last_loop_extent_imm_expr.as(); auto last_loop_extent = last_loop_extent_imm ? last_loop_extent_imm->value.as() : std::nullopt; this->num = loops.size(); this->len = last_loop_extent.has_value() ? *last_loop_extent : 1; this->pos = ForKindFeature::Pos::kPosMixed; int64_t& prod = this->prod = 1; - for (const ForNode* loop : loops) { + for (const ForObj* loop : loops) { PrimExpr extent_imm_expr = Range::GetExtentFromBeginEnd(loop->begin, loop->end); - const auto* extent_imm = extent_imm_expr.as(); + const auto* extent_imm = extent_imm_expr.as(); if (auto extent = extent_imm ? extent_imm->value.as() : std::nullopt; extent.has_value()) { prod *= *extent; @@ -881,11 +881,11 @@ struct Feature { } } - explicit Feature(const TensorStoreNode* store, const LoopNest& loop_nest, - int64_t cache_line_bytes, IntVec* for_touched_bytes, - ForTensorMap* tensor_touched_under_loop, sym::AnalyzerObj* analyzer); + explicit Feature(const TensorStoreObj* store, const LoopNest& loop_nest, int64_t cache_line_bytes, + IntVec* for_touched_bytes, ForTensorMap* tensor_touched_under_loop, + sym::AnalyzerObj* analyzer); - void Init(const TensorStoreNode* store, int n_loops); + void Init(const TensorStoreObj* store, int n_loops); void SetRegion(const LoopNest& loop_nest, // IntVec* for_touched_bytes, // @@ -895,7 +895,7 @@ struct Feature { std::vector sub_features; }; -void Feature::Init(const TensorStoreNode* store, int n_loops) { +void Feature::Init(const TensorStoreObj* store, int n_loops) { struct Info { AccessType access_type = AccessType::kUnknownRW; std::vector multi_indices; @@ -939,12 +939,25 @@ void Feature::SetRegion(const LoopNest& loop_nest, IntVec* for_touched_bytes, ForTensorMap* tensor_touched_under_loop, sym::AnalyzerObj* analyzer) { int n_loops = loop_nest.loops.size(); - const std::vector& loops = loop_nest.loops; + const std::vector& loops = loop_nest.loops; // Step 1. Initialize and bind all the loop variables to a constant *for_touched_bytes = IntVec(n_loops, 0); for (int i = 0; i < n_loops; ++i) { - const ForNode* loop = loops[i]; - analyzer->Bind(loop->var, loop->begin, /*allow_override=*/true); + const ForObj* loop = loops[i]; + { + sym::ConstraintContext context; + analyzer->GetConstraintContext()->ForEach([&](const sym::Constraint& premise) { + if (auto binding = premise.as(); + binding && binding->var.same_as(loop->var)) + return; + if (auto range = premise.as(); + range && range->var.same_as(loop->var)) + return; + context->Append(premise); + }); + context->Append(sym::BindingConstraint(loop->var, loop->begin)); + analyzer->SyncConstraintContext(context); + } } // Step 2. Corner case: no loops if (n_loops == 0) { @@ -957,9 +970,21 @@ void Feature::SetRegion(const LoopNest& loop_nest, IntVec* for_touched_bytes, // Step 3. Gradually bind the loops from inner to outer, // calculate the area the loops touch on each tensor for (int i = n_loops - 1; i >= 0; --i) { - const ForNode* loop = loops[i]; - analyzer->Bind(loop->var, Range(loop->begin, loop->end), - /*allow_override=*/true); + const ForObj* loop = loops[i]; + { + sym::ConstraintContext context; + analyzer->GetConstraintContext()->ForEach([&](const sym::Constraint& premise) { + if (auto binding = premise.as(); + binding && binding->var.same_as(loop->var)) + return; + if (auto range = premise.as(); + range && range->var.same_as(loop->var)) + return; + context->Append(premise); + }); + context->Append(sym::RangeConstraint(loop->var, Range(loop->begin, loop->end))); + analyzer->SyncConstraintContext(context); + } int64_t& touched_bytes = (*for_touched_bytes)[i] = 0; for (SubFeature& feature : sub_features) { TensorVar tensor = feature.tensor; @@ -977,7 +1002,7 @@ void Feature::SetRegion(const LoopNest& loop_nest, IntVec* for_touched_bytes, 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; + const std::vector& loops = loop_nest.loops; // For each tensor, we find the loop stride on it TensorVar tensor = this->tensor; int ndim = tensor.shape().size(); @@ -1019,7 +1044,7 @@ void Feature::SubFeature::SetStride(const LoopNest& loop_nest, sym::AnalyzerObj* int64_t& prod = this->prod_non_strided_loop_extent = 1; for (int j = n_loops - 1; j > i; --j) { PrimExpr extent_imm_expr = Range::GetExtentFromBeginEnd(loops[j]->begin, loops[j]->end); - const auto* extent_imm = extent_imm_expr.as(); + const auto* extent_imm = extent_imm_expr.as(); if (auto extent = extent_imm ? extent_imm->value.as() : std::nullopt; extent.has_value()) { prod *= *extent; @@ -1031,7 +1056,7 @@ void Feature::SubFeature::SetReuse(const LoopNest& loop_nest, int64_t top_loop_t const ForTensorMap& tensor_touched_under_loop) { TensorVar tensor = this->tensor; // Step 3.1. Collect all `Var`s that appears in the tensor region - std::unordered_set region_vars; + std::unordered_set region_vars; auto walk_fn = [®ion_vars](const Var& var) -> ffi::Expected { if (auto prim_var = var.as()) { region_vars.insert(prim_var.value().get()); @@ -1051,14 +1076,14 @@ void Feature::SubFeature::SetReuse(const LoopNest& loop_nest, int64_t top_loop_t // Step 3.2. Enumerate loops from inner to outer, find the first loop with reuse int n_loops = loop_nest.loops.size(); - const std::vector& loops = loop_nest.loops; + const std::vector& loops = loop_nest.loops; for (int i = n_loops - 1; i >= 0; --i) { - const ForNode* loop = loops[i]; + const ForObj* loop = loops[i]; // Case 1. Find an invariant loop, i.e. reuse with kLoopMultipleRead if (!region_vars.count(loop->var.get())) { reuse_type = ReuseType::kLoopMultipleRead; PrimExpr extent_imm_expr = Range::GetExtentFromBeginEnd(loop->begin, loop->end); - const auto* extent_imm = extent_imm_expr.as(); + const auto* extent_imm = extent_imm_expr.as(); if (auto extent = extent_imm ? extent_imm->value.as() : std::nullopt; extent.has_value()) { reuse_ct = *extent; @@ -1068,7 +1093,7 @@ void Feature::SubFeature::SetReuse(const LoopNest& loop_nest, int64_t top_loop_t reuse_dis_iter = 1; for (int j = n_loops - 1; j > i; --j) { PrimExpr extent_imm_expr = Range::GetExtentFromBeginEnd(loops[j]->begin, loops[j]->end); - const auto* extent_imm = extent_imm_expr.as(); + const auto* extent_imm = extent_imm_expr.as(); if (auto extent = extent_imm ? extent_imm->value.as() : std::nullopt; extent.has_value()) { reuse_dis_iter *= *extent; @@ -1092,7 +1117,7 @@ void Feature::SubFeature::SetReuse(const LoopNest& loop_nest, int64_t top_loop_t if (touched.size() >= 2) { int64_t extent = 1; PrimExpr ext_imm_expr = Range::GetExtentFromBeginEnd(loop->begin, loop->end); - const auto* ext_imm = ext_imm_expr.as(); + const auto* ext_imm = ext_imm_expr.as(); if (auto ext = ext_imm ? ext_imm->value.as() : std::nullopt; ext.has_value()) { extent = *ext; } @@ -1138,7 +1163,7 @@ void Feature::SubFeature::SetFeature(const LoopNest& loop_nest, int64_t cache_li this->unique_lines_d_reuse_ct = this->unique_lines / proxy_reuse_ct; } -Feature::Feature(const TensorStoreNode* store, const LoopNest& loop_nest, int64_t cache_line_bytes, +Feature::Feature(const TensorStoreObj* store, const LoopNest& loop_nest, int64_t cache_line_bytes, IntVec* for_touched_bytes, ForTensorMap* tensor_touched_under_loop, sym::AnalyzerObj* analyzer) { int n_loops = loop_nest.loops.size(); @@ -1209,7 +1234,7 @@ struct Feature { explicit Feature(int n_samples, const LoopNest& loop_nest, const IntVec& for_touched_bytes, const group1::Feature::ArithOps& arith_ops) : arith_intensity_curve(n_samples, 0.0) { - const std::vector& loops = loop_nest.loops; + const std::vector& loops = loop_nest.loops; TVM_FFI_ICHECK_EQ(loops.size(), for_touched_bytes.size()); int n_loops = loops.size(); // Calculate `memory_bytes` @@ -1226,7 +1251,7 @@ struct Feature { total_compute_ops /= loop_nest.prod; for (int i = n_loops - 1; i >= 0; --i) { PrimExpr extent_imm_expr = Range::GetExtentFromBeginEnd(loops[i]->begin, loops[i]->end); - const auto* extent_imm = extent_imm_expr.as(); + const auto* extent_imm = extent_imm_expr.as(); if (auto extent = extent_imm ? extent_imm->value.as() : std::nullopt; extent.has_value()) { total_compute_ops *= *extent; @@ -1339,14 +1364,14 @@ class WorkloadEmbeddingExtractor : public StmtExprVisitor { using StmtExprVisitor::Visit_; ffi::Optional Visit(ffi::AnyView value) override { - if (value.as()) return std::nullopt; + if (value.as()) return std::nullopt; return StmtExprVisitor::Visit(value); } static std::vector Extract(const IRModule& mod) { auto self = ffi::make_object(); for (const auto& kv : mod->functions) { - if (const FunctionNode* func = kv.second.as()) { + if (const FunctionObj* func = kv.second.as()) { self->Visit(func->body); } } @@ -1354,7 +1379,7 @@ class WorkloadEmbeddingExtractor : public StmtExprVisitor { } private: - ffi::Optional Visit_(const SBlockNode* block) final { + ffi::Optional Visit_(const SBlockObj* block) final { TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(block)); std::string name = block->name_hint; std::for_each(name.begin(), name.end(), [](char& c) { c = ::tolower(c); }); @@ -1399,7 +1424,7 @@ struct Feature { /*! \brief The feature extracted */ struct Feature { - const VarNode* tensor = nullptr; + const VarObj* tensor = nullptr; int tensor_order = -1; std::unique_ptr group1 = nullptr; std::unique_ptr group2 = nullptr; @@ -1417,7 +1442,7 @@ class PerStoreFeatureCollector : public StmtExprVisitor { using StmtExprVisitor::Visit_; ffi::Optional Visit(ffi::AnyView value) override { - if (value.as()) return std::nullopt; + if (value.as()) return std::nullopt; return StmtExprVisitor::Visit(value); } @@ -1427,7 +1452,7 @@ class PerStoreFeatureCollector : public StmtExprVisitor { auto collector = ffi::make_object(is_gpu, cache_line_bytes, arith_intensity_curve_num_samples); for (const auto& kv : mod->functions) { - if (const FunctionNode* func = kv.second.as()) { + if (const FunctionObj* func = kv.second.as()) { collector->Visit(func->body); for (const Var& param : func->params) { if (auto tensor = param.as()) { @@ -1456,7 +1481,7 @@ class PerStoreFeatureCollector : public StmtExprVisitor { } private: - ffi::Optional Visit_(const ForNode* loop) final { + ffi::Optional Visit_(const ForObj* loop) final { int64_t auto_unroll; ForVec* for_vec = loop_nest_.Push(loop, &auto_unroll); TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(loop)); @@ -1464,11 +1489,11 @@ class PerStoreFeatureCollector : public StmtExprVisitor { return std::nullopt; } - ffi::Optional Visit_(const TensorStoreNode* store) final { - if (store->value->IsInstance() || store->value->IsInstance()) { + ffi::Optional Visit_(const TensorStoreObj* store) final { + if (store->value->IsInstance() || store->value->IsInstance()) { return std::nullopt; } - const VarNode* tensor = store->dest.as_or_throw().get(); + const VarObj* tensor = store->dest.as_or_throw().get(); Feature& feature = tensor_features_[tensor]; if (feature.tensor == nullptr) { feature.tensor = tensor; @@ -1485,7 +1510,7 @@ class PerStoreFeatureCollector : public StmtExprVisitor { return std::nullopt; } - ffi::Optional Visit_(const SBlockNode* block) final { + ffi::Optional Visit_(const SBlockObj* block) final { TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(block)); for (const TensorVar& tensor : block->alloc_tensors) { HandleTensorAlloc(tensor); @@ -1513,7 +1538,7 @@ class PerStoreFeatureCollector : public StmtExprVisitor { LoopNest loop_nest_ = {}; IntVec for_touched_bytes_ = {}; ForTensorMap tensor_touched_under_loop_ = {}; - std::unordered_map tensor_features_ = {}; + std::unordered_map tensor_features_ = {}; }; } // namespace s_tir diff --git a/src/s_tir/meta_schedule/mutator/mutate_parallel.cc b/src/s_tir/meta_schedule/mutator/mutate_parallel.cc index a4777cb9af7b..2da8f1f7ea82 100644 --- a/src/s_tir/meta_schedule/mutator/mutate_parallel.cc +++ b/src/s_tir/meta_schedule/mutator/mutate_parallel.cc @@ -64,13 +64,13 @@ Instruction ReplaceAnnValue(Instruction inst, int64_t ann_val) { * \param inst The instruction to be checked * \return The output of the instruction Get-Block */ -const SBlockRVNode* GetInstGetSBlockOutput(const Instruction& inst) { +const SBlockRVObj* GetInstGetSBlockOutput(const Instruction& inst) { static const InstructionKind& inst_get_sblock = InstructionKind::Get("GetSBlock"); if (!inst->kind.same_as(inst_get_sblock)) { return nullptr; } TVM_FFI_ICHECK_EQ(inst->outputs.size(), 1); - const SBlockRVNode* block = TVM_TYPE_AS(inst->outputs[0], SBlockRVNode); + const SBlockRVObj* block = TVM_TYPE_AS(inst->outputs[0], SBlockRVObj); return block; } @@ -88,20 +88,20 @@ std::vector> AnalyzeParallel(const ScheduleState& self, ffi::Array block_srefs = GetSBlocks(self, block_name, self->mod->GetGlobalVar(func_name)); TVM_FFI_ICHECK_EQ(block_srefs.size(), 1); - const SBlockNode* block = TVM_SREF_TO_SBLOCK(block_srefs[0]); + const SBlockObj* block = TVM_SREF_TO_SBLOCK(block_srefs[0]); ScopeBlockLoopInfo info = GetScopeBlockLoopInfo(ffi::GetRef(block)); std::vector> results; results.reserve(info.realizes.size()); for (const SBlockRealize& realize : info.realizes) { // Step 1. Extract static loop extents for spatial loops std::vector loop_extents; - const ForNode* loop = nullptr; + const ForObj* loop = nullptr; for (const StmtSRefObj* loop_sref = self->stmt2ref.at(realize->block.get())->parent; - (loop = loop_sref->StmtAs()) != nullptr; // + (loop = loop_sref->StmtAs()) != nullptr; // loop_sref = loop_sref->parent) { int64_t loop_extent = -1; PrimExpr ext_imm_expr = Range::GetExtentFromBeginEnd(loop->begin, loop->end); - const auto* ext_imm = ext_imm_expr.as(); + const auto* ext_imm = ext_imm_expr.as(); if (auto ext = ext_imm ? ext_imm->value.as() : std::nullopt; ext.has_value()) { if (!info.non_spatial_vars.count(loop->var.get())) { loop_extent = *ext; @@ -226,17 +226,17 @@ struct MutateParallelObj::Candidate { */ bool FindParallelDecision(const Trace& trace, TRandState* rand_state, MutateParallelObj::Candidate* candidate) { - using s_tir::InstructionNode; - using s_tir::SBlockRVNode; - std::unordered_map get_sblock_insts; - std::vector ann_insts; + using s_tir::InstructionObj; + using s_tir::SBlockRVObj; + std::unordered_map get_sblock_insts; + std::vector ann_insts; get_sblock_insts.reserve(trace->insts.size()); ann_insts.reserve(trace->insts.size()); for (const Instruction& inst : trace->insts) { if (s_tir::IsAnnotateWithParallel(inst)) { ann_insts.push_back(inst.get()); } - if (const SBlockRVNode* block_rv = s_tir::GetInstGetSBlockOutput(inst)) { + if (const SBlockRVObj* block_rv = s_tir::GetInstGetSBlockOutput(inst)) { get_sblock_insts[block_rv] = inst.get(); } } @@ -244,9 +244,9 @@ bool FindParallelDecision(const Trace& trace, TRandState* rand_state, if (n_ann_insts == 0) { return false; } - const InstructionNode* ann_inst = ann_insts[s_tir::SampleInt(rand_state, 0, n_ann_insts)]; + const InstructionObj* ann_inst = ann_insts[s_tir::SampleInt(rand_state, 0, n_ann_insts)]; TVM_FFI_ICHECK_EQ(ann_inst->inputs.size(), 2); - const InstructionNode* get_sblock_inst = + const InstructionObj* get_sblock_inst = get_sblock_insts.at(ann_inst->inputs[0].as_or_throw().get()); TVM_FFI_ICHECK_EQ(get_sblock_inst->attrs.size(), 2); candidate->inst = ffi::GetRef(ann_inst); diff --git a/src/s_tir/meta_schedule/mutator/mutate_thread_binding.cc b/src/s_tir/meta_schedule/mutator/mutate_thread_binding.cc index 2bfa4f20afe0..31ed871c623f 100644 --- a/src/s_tir/meta_schedule/mutator/mutate_thread_binding.cc +++ b/src/s_tir/meta_schedule/mutator/mutate_thread_binding.cc @@ -86,16 +86,16 @@ class MutateThreadBindingObj : public MutatorObj { */ std::vector MutateThreadBindingObj::FindCandidates( const Trace& trace, TRandState* rand_state) { - using s_tir::InstructionNode; + using s_tir::InstructionObj; static InstructionKind inst_sample_categorical = InstructionKind::Get("SampleCategorical"); static InstructionKind inst_split = InstructionKind::Get("Split"); static InstructionKind inst_bind = InstructionKind::Get("Bind"); std::vector candidates; - std::unordered_map sample_insts; - std::unordered_map sampled_split_insts; - std::vector bind_insts; + std::unordered_map sample_insts; + std::unordered_map sampled_split_insts; + std::vector bind_insts; auto is_split_by_sample = [&sample_insts](const Instruction& inst) -> bool { if (!inst->kind.same_as(inst_split)) { @@ -123,28 +123,28 @@ std::vector MutateThreadBindingObj::FindCandi for (const Instruction& inst : trace->insts) { if (inst->kind.same_as(inst_sample_categorical)) { TVM_FFI_ICHECK_EQ(inst->outputs.size(), 1); - const ExprNode* var_rv = TVM_TYPE_AS(inst->outputs[0], ExprNode); + const ExprObj* var_rv = TVM_TYPE_AS(inst->outputs[0], ExprObj); sample_insts[var_rv] = inst.get(); } else if (is_split_by_sample(inst)) { TVM_FFI_ICHECK_EQ(inst->outputs.size(), 2); // Only consider the inner loop, which can be bound to threadIdx.x - const s_tir::LoopRVNode* var_rv = TVM_TYPE_AS(inst->outputs[1], s_tir::LoopRVNode); + const s_tir::LoopRVObj* var_rv = TVM_TYPE_AS(inst->outputs[1], s_tir::LoopRVObj); sampled_split_insts[var_rv] = inst.get(); } else if (is_thread_binding_by_sample(inst)) { bind_insts.push_back(inst.get()); } } - for (const InstructionNode* bind_inst : bind_insts) { - const auto* loop_rv = TVM_TYPE_AS(bind_inst->inputs[0], s_tir::LoopRVNode); + for (const InstructionObj* bind_inst : bind_insts) { + const auto* loop_rv = TVM_TYPE_AS(bind_inst->inputs[0], s_tir::LoopRVObj); auto split_it = sampled_split_insts.find(loop_rv); TVM_FFI_ICHECK(split_it != sampled_split_insts.end()); - const InstructionNode* split_inst = split_it->second; + const InstructionObj* split_inst = split_it->second; - const auto* expr_rv = TVM_TYPE_AS(split_inst->inputs[2], ExprNode); + const auto* expr_rv = TVM_TYPE_AS(split_inst->inputs[2], ExprObj); auto sample_it = sample_insts.find(expr_rv); TVM_FFI_ICHECK(sample_it != sample_insts.end()); - const InstructionNode* sample_inst = sample_it->second; + const InstructionObj* sample_inst = sample_it->second; // SampleCategorical decision is Optional after the Integer phase-out. int decision = trace->decisions[ffi::GetRef(sample_inst)].cast(); diff --git a/src/s_tir/meta_schedule/mutator/mutate_tile_size.cc b/src/s_tir/meta_schedule/mutator/mutate_tile_size.cc index 9982f083ca05..034b025ad960 100644 --- a/src/s_tir/meta_schedule/mutator/mutate_tile_size.cc +++ b/src/s_tir/meta_schedule/mutator/mutate_tile_size.cc @@ -128,7 +128,7 @@ void FindSampleVectorize(const Trace& trace, std::vector* inst, TVM_FFI_ICHECK_EQ(inst->inputs.size(), 2); if (inst->attrs[0].as_or_throw() == tvm::s_tir::attr::kMetaScheduleCooperativeFetch.name) { - const auto* ann_val = inst->inputs[1].as(); + const auto* ann_val = inst->inputs[1].as(); TVM_FFI_ICHECK(ann_val); annotated.insert(ann_val); } diff --git a/src/s_tir/meta_schedule/mutator/mutate_unroll.cc b/src/s_tir/meta_schedule/mutator/mutate_unroll.cc index fce78165f2ab..ca9233e51477 100644 --- a/src/s_tir/meta_schedule/mutator/mutate_unroll.cc +++ b/src/s_tir/meta_schedule/mutator/mutate_unroll.cc @@ -96,16 +96,16 @@ struct MutateUnrollObj::Candidate { bool FindUnrollDecision(const Trace& trace, TRandState* rand_state, MutateUnrollObj::Candidate* candidate) { using s_tir::InstructionKind; - using s_tir::InstructionNode; + using s_tir::InstructionObj; static const InstructionKind& inst_sample_categorical = InstructionKind::Get("SampleCategorical"); - std::unordered_map sample_insts; - std::vector ann_insts; + std::unordered_map sample_insts; + std::vector ann_insts; sample_insts.reserve(trace->insts.size()); ann_insts.reserve(trace->insts.size()); for (const Instruction& inst : trace->insts) { if (inst->kind.same_as(inst_sample_categorical)) { TVM_FFI_ICHECK_EQ(inst->outputs.size(), 1); - const ExprNode* var_rv = inst->outputs[0].as_or_throw().get(); + const ExprObj* var_rv = inst->outputs[0].as_or_throw().get(); sample_insts[var_rv] = inst.get(); } else if (IsAnnotateWithUnroll(inst)) { ann_insts.push_back(inst.get()); @@ -115,11 +115,11 @@ bool FindUnrollDecision(const Trace& trace, TRandState* rand_state, if (n_ann_insts == 0) { return false; } - const InstructionNode* ann_inst = ann_insts[s_tir::SampleInt(rand_state, 0, n_ann_insts)]; + const InstructionObj* ann_inst = ann_insts[s_tir::SampleInt(rand_state, 0, n_ann_insts)]; TVM_FFI_ICHECK_EQ(ann_inst->inputs.size(), 2); const auto* var_rv = ann_inst->inputs[1].as_or_throw().get(); TVM_FFI_ICHECK(sample_insts.count(var_rv)); - const InstructionNode* sample_inst = sample_insts.at(var_rv); + const InstructionObj* sample_inst = sample_insts.at(var_rv); TVM_FFI_ICHECK_EQ(sample_inst->attrs.size(), 2); candidate->inst = ffi::GetRef(sample_inst); // SampleCategorical decision is Optional after the Integer phase-out. 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 8a5b7f08d647..a8ec8ace9fdd 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 @@ -35,7 +35,7 @@ struct AsyncStridedMemCopyFinder : public StmtExprVisitor { static bool Find(const IRModule& mod) { auto finder = ffi::make_object(); for (const auto& kv : mod->functions) { - if (const auto* function = kv.second.as()) { + if (const auto* function = kv.second.as()) { finder->Visit(function->body); if (finder->found_) { return true; @@ -46,7 +46,7 @@ struct AsyncStridedMemCopyFinder : public StmtExprVisitor { } private: - ffi::Optional Visit_(const ForNode* loop) final { + ffi::Optional Visit_(const ForObj* loop) final { if (!found_) { input_iters.Set(loop->var, Range(loop->begin, loop->end)); TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(loop)); @@ -54,7 +54,7 @@ struct AsyncStridedMemCopyFinder : public StmtExprVisitor { return std::nullopt; } - ffi::Optional Visit_(const RegionStmtNode* op) final { + ffi::Optional Visit_(const RegionStmtObj* op) final { bool previous = in_async_copy_; in_async_copy_ |= op->op.same_as(s_tir::async_copy_scope()); auto result = StmtExprVisitor::Visit_(op); @@ -62,9 +62,9 @@ struct AsyncStridedMemCopyFinder : public StmtExprVisitor { return result; } - ffi::Optional Visit_(const TensorStoreNode* op) final { + ffi::Optional Visit_(const TensorStoreObj* op) final { if (!found_ && in_async_copy_) { - if (const auto* load = op->value.as()) { + if (const auto* load = op->value.as()) { // Inspect each copy, including copies grouped under one commit or predicate. sym::Analyzer analyzer; auto store_map = DetectIterMap(op->indices, input_iters, 1, sym::IterMapLevel::Surjective, @@ -102,7 +102,7 @@ class DisallowAsyncStridedMemCopyObj : public PostprocObj { for (const auto& kv : mod->functions) { const GlobalVar& g_var = kv.first; const BaseFunc& base_func = kv.second; - if (const auto* function = base_func.as()) { + if (const auto* function = base_func.as()) { IRModule lowered{ffi::UnsafeInit()}; try { auto pass_list = ffi::Array(); diff --git a/src/s_tir/meta_schedule/postproc/disallow_dynamic_loop.cc b/src/s_tir/meta_schedule/postproc/disallow_dynamic_loop.cc index 002eb9906cf3..44bc03bb9cea 100644 --- a/src/s_tir/meta_schedule/postproc/disallow_dynamic_loop.cc +++ b/src/s_tir/meta_schedule/postproc/disallow_dynamic_loop.cc @@ -32,7 +32,7 @@ struct DynamicExtentFinder : public StmtExprVisitor { auto finder = ffi::make_object(); for (const auto& kv : mod->functions) { const BaseFunc& func = kv.second; - if (const auto* function = func.as()) { + if (const auto* function = func.as()) { finder->Visit(function->body); if (finder->found_) { return true; @@ -43,8 +43,8 @@ struct DynamicExtentFinder : public StmtExprVisitor { } private: - ffi::Optional Visit_(const ForNode* loop) final { - if (!Range::GetExtentFromBeginEnd(loop->begin, loop->end)->IsInstance()) { + ffi::Optional Visit_(const ForObj* loop) final { + if (!Range::GetExtentFromBeginEnd(loop->begin, loop->end)->IsInstance()) { found_ = true; } else { return StmtExprVisitor::Visit_(loop); @@ -53,7 +53,7 @@ struct DynamicExtentFinder : public StmtExprVisitor { } ffi::Optional Visit(ffi::AnyView stmt) final { - if (stmt.as()) return std::nullopt; + if (stmt.as()) return std::nullopt; if (!found_) { return StmtExprVisitor::Visit(stmt); } diff --git a/src/s_tir/meta_schedule/postproc/rewrite_cooperative_fetch.cc b/src/s_tir/meta_schedule/postproc/rewrite_cooperative_fetch.cc index b907141e6813..32e5cb9dcd60 100644 --- a/src/s_tir/meta_schedule/postproc/rewrite_cooperative_fetch.cc +++ b/src/s_tir/meta_schedule/postproc/rewrite_cooperative_fetch.cc @@ -186,7 +186,7 @@ bool RewriteCooperativeFetchObj::Apply(const s_tir::Schedule& sch) { int64_t fused_extent = -1; PrimExpr extent_imm_expr = Range::GetExtentFromBeginEnd(sch->Get(fused)->begin, sch->Get(fused)->end); - const auto* extent_imm = extent_imm_expr.as(); + const auto* extent_imm = extent_imm_expr.as(); if (auto extent = extent_imm ? extent_imm->value.as() : std::nullopt; extent.has_value()) { fused_extent = *extent; diff --git a/src/s_tir/meta_schedule/postproc/rewrite_layout.cc b/src/s_tir/meta_schedule/postproc/rewrite_layout.cc index 14a5747ca417..ab2992a156ad 100644 --- a/src/s_tir/meta_schedule/postproc/rewrite_layout.cc +++ b/src/s_tir/meta_schedule/postproc/rewrite_layout.cc @@ -47,14 +47,14 @@ class TensorReadPosCollector : public StmtExprVisitor { const ffi::Optional GetTensorIndexMap() const { return tensor_index_map_; } private: - ffi::Optional Visit_(const ForNode* op) final { + ffi::Optional Visit_(const ForObj* op) final { loop_stack_.push_back(ffi::GetRef(op)); TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(op)); loop_stack_.pop_back(); return std::nullopt; } - ffi::Optional Visit_(const SBlockRealizeNode* op) final { + ffi::Optional Visit_(const SBlockRealizeObj* op) final { ffi::Optional outer_block_realize = ffi::GetRef(op); std::swap(outer_block_realize, cur_realize_); TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(op)); @@ -62,7 +62,7 @@ class TensorReadPosCollector : public StmtExprVisitor { return std::nullopt; } - ffi::Optional Visit_(const TensorLoadNode* op) final { + ffi::Optional Visit_(const TensorLoadObj* op) final { TVM_FFI_ICHECK(cur_realize_.has_value()) << "TensorLoad occurred outside of any block"; const TensorVar& tensor = op->source.as_or_throw(); @@ -105,7 +105,7 @@ class TensorReadPosCollector : public StmtExprVisitor { } /*! \brief The tensor of interest. */ - const VarNode* tensor_; + const VarObj* tensor_; /*! \brief The block that consumes the tensor and the corresponding read index. */ std::optional> tensor_loc_; /*! \brief The proposed IndexMap. */ @@ -124,11 +124,11 @@ class LayoutFreeTensorCollector : public StmtExprVisitor { using StmtExprVisitor::Visit_; ffi::Optional Visit(ffi::AnyView value) override { - if (value.as()) return std::nullopt; + if (value.as()) return std::nullopt; return StmtExprVisitor::Visit(value); } - ffi::Optional Visit_(const SBlockNode* block) final { + ffi::Optional Visit_(const SBlockObj* block) final { TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(block)); if (auto ann = block->annotations.Get(tvm::topi::attr::kLayoutFreePlaceholders)) { for (TensorVar tensor : ann.value().as_or_throw>()) { @@ -141,7 +141,7 @@ class LayoutFreeTensorCollector : public StmtExprVisitor { std::unordered_set tensors; }; -ffi::Array CollectLayoutFreeTensors(const FunctionNode* func) { +ffi::Array CollectLayoutFreeTensors(const FunctionObj* func) { // Only rewrite Functions with attr "layout_free_tensors" ffi::Array layout_free_tensor_index = func->GetAttr(s_tir::attr::kLayoutFreeTensors.name, ffi::Array()).value(); @@ -162,8 +162,8 @@ ffi::Array CollectLayoutFreeTensors(const FunctionNode* func) { return layout_free_tensors; } -std::optional> GetSuggestedIndexMap( - TensorVar tensor, const FunctionNode* function) { +std::optional> GetSuggestedIndexMap(TensorVar tensor, + const FunctionObj* function) { auto collector = ffi::make_object(tensor); collector->Visit(function->body); @@ -179,19 +179,19 @@ std::optional> GetSuggestedIndexMap( } /*! \brief Get a chain of cache-read blocks, starting from the one consuming tensor. */ -std::vector GetCacheReadChain(const TensorVar& tensor, const FunctionNode* function) { +std::vector GetCacheReadChain(const TensorVar& tensor, const FunctionObj* function) { class TensorReadChainCollector : public StmtExprVisitor { public: using StmtExprVisitor::Visit_; ffi::Optional Visit(ffi::AnyView value) override { - if (value.as()) return std::nullopt; + if (value.as()) return std::nullopt; return StmtExprVisitor::Visit(value); } explicit TensorReadChainCollector(const TensorVar& tensor) : cur_tensor_(tensor.get()) {} - ffi::Optional Visit_(const SBlockNode* op) final { + ffi::Optional Visit_(const SBlockObj* op) final { // Check if this block is doing cache_read or a similar operation that consumes cur_tensor_. if (!op->init && op->reads.size() == 1 && op->writes.size() == 1 && op->reads[0]->source.as_or_throw().get() == cur_tensor_) { @@ -204,7 +204,7 @@ std::vector GetCacheReadChain(const TensorVar& tensor, const Functi std::vector cache_read_chain; private: - const VarNode* cur_tensor_; + const VarObj* cur_tensor_; }; auto collector = ffi::make_object(tensor); @@ -221,7 +221,7 @@ bool RewriteLayout(const Schedule& sch) { for (const auto& [g_var, base_func] : sch->mod()->functions) { const ffi::String& func_name = g_var->name_hint; - const auto* function = base_func.as(); + const auto* function = base_func.as(); // Only consider Function if (function == nullptr) { continue; 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 68d642aa572d..0cadd3948259 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 @@ -33,7 +33,7 @@ using namespace tvm::tirx; * \param sref The sref of loop * \return Whether the loop has any annotation */ -inline bool HasAnnOrBinding(const ForNode* loop) { +inline bool HasAnnOrBinding(const ForObj* loop) { return tvm::s_tir::GetThreadBinding(loop).has_value() || !loop->annotations.empty(); } @@ -58,14 +58,14 @@ class StrideExtractor : public StmtExprVisitor { explicit StrideExtractor(const Var& var) : var_(var) {} private: - ffi::Optional Visit_(const MulNode* node) final { + ffi::Optional Visit_(const MulObj* node) final { TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(node)); - if (const auto* a = node->a.as()) { + if (const auto* a = node->a.as()) { if (strides_.count(node->b.get())) { strides_[node] = static_cast(strides_[node->b.get()] * a->value); } - } else if (const auto* b = node->b.as()) { + } else if (const auto* b = node->b.as()) { if (strides_.count(node->a.get())) { strides_[node] = static_cast(strides_[node->a.get()] * b->value); } @@ -73,7 +73,7 @@ class StrideExtractor : public StmtExprVisitor { return std::nullopt; } - ffi::Optional Visit_(const AddNode* node) final { + ffi::Optional Visit_(const AddObj* node) final { TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(node)); int64_t stride_a, stride_b; if (strides_.count(node->a.get())) { @@ -92,7 +92,7 @@ class StrideExtractor : public StmtExprVisitor { return std::nullopt; } - ffi::Optional Visit_(const VarNode* node) final { + ffi::Optional Visit_(const VarObj* node) final { if (node == var_.get()) { strides_[node] = 1; } @@ -100,7 +100,7 @@ class StrideExtractor : public StmtExprVisitor { } const Var& var_; - std::unordered_map strides_; + std::unordered_map strides_; }; struct ParsedAnnotation { @@ -163,7 +163,7 @@ int CalculateNumRewritableLoops(const ffi::Array& loop_srefs, TVM_FFI_ICHECK_EQ(loop_srefs.size(), loop_types.size()); for (size_t i = 0; i < loop_srefs.size(); ++i) { const StmtSRef& loop_sref = loop_srefs[i]; - const ForNode* loop = TVM_SREF_TO_FOR(loop_sref); + const ForObj* loop = TVM_SREF_TO_FOR(loop_sref); if (HasAnnOrBinding(loop)) { continue; } @@ -177,7 +177,7 @@ int CalculateNumRewritableLoops(const ffi::Array& loop_srefs, } // Check if the loop extent is valid PrimExpr extent_expr = Range::GetExtentFromBeginEnd(loop->begin, loop->end); - if (const auto* extent = extent_expr.as(); + if (const auto* extent = extent_expr.as(); !extent || !extent->value.as().has_value()) { continue; } @@ -214,7 +214,7 @@ void AdjustParallelVectorize(const Schedule& sch, const SBlockRV& block_rv, ffi::Array tensor_access(realize->block->reads); tensor_access.insert(tensor_access.end(), realize->block->writes.begin(), realize->block->writes.end()); - std::unordered_map binding_map; + std::unordered_map binding_map; for (size_t i = 0; i < realize->iter_values.size(); i++) { binding_map.insert_or_assign(realize->block->iter_vars[i]->var.get(), realize->iter_values[i]); } @@ -234,7 +234,7 @@ void AdjustParallelVectorize(const Schedule& sch, const SBlockRV& block_rv, // get strides for each loop var for (const StmtSRef& loop_sref : loop_srefs) { int64_t stride = 0, tensor_stride = 1; - const auto* var = loop_sref->StmtAs(); + const auto* var = loop_sref->StmtAs(); sym::Analyzer analyzer; for (int i = access->region.size() - 1; i >= 0; i--) { PrimExpr idx = analyzer->Simplify( @@ -246,7 +246,7 @@ void AdjustParallelVectorize(const Schedule& sch, const SBlockRV& block_rv, break; } const auto* shape = - access->source.as_or_throw().shape()[i].as(); + access->source.as_or_throw().shape()[i].as(); if (shape == nullptr) { can_analyze_contiguous_access = false; break; @@ -281,7 +281,7 @@ void AdjustParallelVectorize(const Schedule& sch, const SBlockRV& block_rv, PrimExpr prev_used_iter_extent_imm_expr = Range::GetExtentFromBeginEnd(TVM_SREF_TO_FOR(loop_srefs[prev_used_iter])->begin, TVM_SREF_TO_FOR(loop_srefs[prev_used_iter])->end); - const auto* prev_used_iter_extent_imm = prev_used_iter_extent_imm_expr.as(); + const auto* prev_used_iter_extent_imm = prev_used_iter_extent_imm_expr.as(); auto prev_used_iter_extent = prev_used_iter_extent_imm ? prev_used_iter_extent_imm->value.as() : std::nullopt; @@ -309,13 +309,13 @@ void AdjustParallelVectorize(const Schedule& sch, const SBlockRV& block_rv, int64_t prod_extent = 1; for (int i = 0; i < n_loops && loop_types[i] == IterVarType::kDataPar; ++i) { const StmtSRef& loop_sref = loop_srefs[i]; - const ForNode* loop = TVM_SREF_TO_FOR(loop_sref); + const ForObj* loop = TVM_SREF_TO_FOR(loop_sref); if (HasAnnOrBinding(loop)) { break; } // Check if the loop extent is valid PrimExpr extent_imm_expr = Range::GetExtentFromBeginEnd(loop->begin, loop->end); - const auto* extent_imm = extent_imm_expr.as(); + const auto* extent_imm = extent_imm_expr.as(); auto extent = extent_imm ? extent_imm->value.as() : std::nullopt; if (!extent.has_value()) { break; @@ -340,7 +340,7 @@ void AdjustParallelVectorize(const Schedule& sch, const SBlockRV& block_rv, for (int i = n_loops - 1; i >= 0 && loop_types[i] == IterVarType::kDataPar && num_fusible < max_fusible; --i) { const StmtSRef& loop_sref = loop_srefs[i]; - const ForNode* loop = TVM_SREF_TO_FOR(loop_sref); + const ForObj* loop = TVM_SREF_TO_FOR(loop_sref); if (HasAnnOrBinding(loop)) { break; } @@ -354,7 +354,7 @@ void AdjustParallelVectorize(const Schedule& sch, const SBlockRV& block_rv, } // Check if the loop extent is valid PrimExpr extent_imm_expr = Range::GetExtentFromBeginEnd(loop->begin, loop->end); - const auto* extent_imm = extent_imm_expr.as(); + const auto* extent_imm = extent_imm_expr.as(); auto extent = extent_imm ? extent_imm->value.as() : std::nullopt; if (!extent.has_value()) { break; @@ -389,10 +389,10 @@ bool FindAnnotatedRootBlock(const Schedule& sch, ParsedAnnotation* parsed, SBloc for (const auto& kv : mod->functions) { const GlobalVar& g_var = kv.first; const BaseFunc& base_func = kv.second; - if (const auto* function = base_func.as()) { - const SBlockRealizeNode* block_realize = + if (const auto* function = base_func.as()) { + const SBlockRealizeObj* block_realize = function->body.has_value() && function->body.value()->size() == 1 - ? function->body.value()->seq[0].as() + ? function->body.value()->seq[0].as() : nullptr; if (block_realize != nullptr) { SBlock block = block_realize->block; diff --git a/src/s_tir/meta_schedule/postproc/rewrite_reduction_block.cc b/src/s_tir/meta_schedule/postproc/rewrite_reduction_block.cc index 2f5c0a273e55..85333d9e1869 100644 --- a/src/s_tir/meta_schedule/postproc/rewrite_reduction_block.cc +++ b/src/s_tir/meta_schedule/postproc/rewrite_reduction_block.cc @@ -32,7 +32,7 @@ struct ReductionBlockFinder : public StmtExprVisitor { using StmtExprVisitor::Visit_; ffi::Optional Visit(ffi::AnyView value) override { - if (value.as()) return std::nullopt; + if (value.as()) return std::nullopt; return StmtExprVisitor::Visit(value); } @@ -42,10 +42,10 @@ struct ReductionBlockFinder : public StmtExprVisitor { for (const auto& kv : self->mod->functions) { GlobalVar g_var = kv.first; BaseFunc base_func = kv.second; - if (const auto* function = base_func.as()) { + if (const auto* function = base_func.as()) { auto finder = ffi::make_object(); finder->Visit(function->body); - for (const SBlockNode* block : finder->results_) { + for (const SBlockObj* block : finder->results_) { results.emplace_back(self->stmt2ref.at(block), g_var->name_hint); } } @@ -54,7 +54,7 @@ struct ReductionBlockFinder : public StmtExprVisitor { } private: - ffi::Optional Visit_(const ForNode* loop) final { + ffi::Optional Visit_(const ForObj* loop) final { support::GPUThreadTag thread_scope = GetThreadBindingTag(loop); if (IsThreadIdx(thread_scope) || IsBlockIdx(thread_scope)) { thread_bound_loop_vars_.insert(loop->var.get()); @@ -62,23 +62,23 @@ struct ReductionBlockFinder : public StmtExprVisitor { return StmtExprVisitor::Visit_(loop); } - ffi::Optional Visit_(const SBlockRealizeNode* realize) final { + ffi::Optional Visit_(const SBlockRealizeObj* realize) final { if (realize->block->init.has_value() && AllReductionIterVarAreUnbound(realize)) { results_.push_back(realize->block.get()); } return StmtExprVisitor::Visit_(realize); } - bool AllReductionIterVarAreUnbound(const SBlockRealizeNode* realize) const { + bool AllReductionIterVarAreUnbound(const SBlockRealizeObj* realize) const { if (thread_bound_loop_vars_.empty()) { return true; } - auto f_find = [this](const VarNode* var) -> bool { return thread_bound_loop_vars_.count(var); }; + auto f_find = [this](const VarObj* var) -> bool { return thread_bound_loop_vars_.count(var); }; auto walkfn = [&](const Var& var) -> ffi::Expected { return f_find(var.get()) ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var)) : ffi::WalkResult::Advance(); }; - const SBlockNode* block = realize->block.get(); + const SBlockObj* block = realize->block.get(); TVM_FFI_ICHECK_EQ(block->iter_vars.size(), realize->iter_values.size()); int n = block->iter_vars.size(); for (int i = 0; i < n; ++i) { @@ -94,9 +94,9 @@ struct ReductionBlockFinder : public StmtExprVisitor { } /*! \brief The results of the collection */ - std::vector results_; + std::vector results_; /*! \brief Loop variables that are bound to threads */ - std::unordered_set thread_bound_loop_vars_; + std::unordered_set thread_bound_loop_vars_; }; /*! diff --git a/src/s_tir/meta_schedule/postproc/rewrite_tensorize.cc b/src/s_tir/meta_schedule/postproc/rewrite_tensorize.cc index 71208c56a5c1..8358d7e454c9 100644 --- a/src/s_tir/meta_schedule/postproc/rewrite_tensorize.cc +++ b/src/s_tir/meta_schedule/postproc/rewrite_tensorize.cc @@ -34,12 +34,12 @@ using s_tir::LoopRV; using s_tir::SBlockRV; void CollectTensorizationJobs( - const s_tir::Schedule& sch, const ffi::String& func_name, const tirx::FunctionNode* func, + const s_tir::Schedule& sch, const ffi::String& func_name, const tirx::FunctionObj* func, bool vectorize_init_loop, std::vector>>* jobs) { auto walk_fn = [=, &jobs](const s_tir::SBlock& block) -> ffi::Expected { tirx::StmtSRef block_sref = sch->GetSRef(block.get()); - std::string block_name = block_sref->StmtAs()->name_hint; + std::string block_name = block_sref->StmtAs()->name_hint; if (ffi::Optional intrin_name = s_tir::GetAnn( block_sref, tvm::s_tir::attr::kMetaScheduleAutoTensorize.name)) { if (intrin_name.value() != "") { @@ -93,7 +93,7 @@ bool RewriteTensorizeObj::Apply(const s_tir::Schedule& sch) { for (const auto& kv : sch->mod()->functions) { GlobalVar g_var = kv.first; BaseFunc base_func = kv.second; - if (const tirx::FunctionNode* function = base_func.as()) { + if (const tirx::FunctionObj* function = base_func.as()) { CollectTensorizationJobs(sch, g_var->name_hint, function, vectorize_init_loop, &jobs); } } diff --git a/src/s_tir/meta_schedule/postproc/rewrite_unbound_block.cc b/src/s_tir/meta_schedule/postproc/rewrite_unbound_block.cc index 92e7700fbbc6..10daa88d2851 100644 --- a/src/s_tir/meta_schedule/postproc/rewrite_unbound_block.cc +++ b/src/s_tir/meta_schedule/postproc/rewrite_unbound_block.cc @@ -33,7 +33,7 @@ class UnboundBlockFinder : public StmtExprVisitor { using StmtExprVisitor::Visit_; ffi::Optional Visit(ffi::AnyView value) override { - if (value.as()) return std::nullopt; + if (value.as()) return std::nullopt; return StmtExprVisitor::Visit(value); } @@ -42,7 +42,7 @@ class UnboundBlockFinder : public StmtExprVisitor { for (const auto& kv : self->mod->functions) { GlobalVar g_var = kv.first; BaseFunc base_func = kv.second; - if (const auto* function = base_func.as()) { + if (const auto* function = base_func.as()) { finder->global_var_name_ = g_var->name_hint; TVM_FFI_ICHECK(function->body.has_value()); TVM_FFI_ICHECK_EQ(function->body.value()->size(), 1); @@ -53,7 +53,7 @@ class UnboundBlockFinder : public StmtExprVisitor { } private: - ffi::Optional Visit_(const ForNode* loop) final { + ffi::Optional Visit_(const ForObj* loop) final { support::GPUThreadTag thread_scope = GetThreadBindingTag(loop); if (IsBlockIdx(thread_scope)) { ++n_block_idx_; @@ -71,7 +71,7 @@ class UnboundBlockFinder : public StmtExprVisitor { return std::nullopt; } - ffi::Optional Visit_(const SBlockNode* block) final { + ffi::Optional Visit_(const SBlockObj* block) final { blocks_.emplace_back(self_->stmt2ref.at(block), global_var_name_); return std::nullopt; } diff --git a/src/s_tir/meta_schedule/postproc/verify_gpu_code.cc b/src/s_tir/meta_schedule/postproc/verify_gpu_code.cc index 172842a2651e..d6a31dc26cad 100644 --- a/src/s_tir/meta_schedule/postproc/verify_gpu_code.cc +++ b/src/s_tir/meta_schedule/postproc/verify_gpu_code.cc @@ -36,7 +36,7 @@ class ThreadExtentChecker : public StmtExprVisitor { using StmtExprVisitor::Visit_; ffi::Optional Visit(ffi::AnyView value) override { - if (value.as()) return std::nullopt; + if (value.as()) return std::nullopt; return StmtExprVisitor::Visit(value); } @@ -54,11 +54,11 @@ class ThreadExtentChecker : public StmtExprVisitor { explicit ThreadExtentChecker(int thread_warp_size) : thread_warp_size_(thread_warp_size) {} private: - ffi::Optional Visit_(const ForNode* loop) { + ffi::Optional Visit_(const ForObj* loop) { support::GPUThreadTag thread_scope = GetThreadBindingTag(loop); if (IsThreadIdx(thread_scope)) { PrimExpr p_ext_imm_expr = Range::GetExtentFromBeginEnd(loop->begin, loop->end); - const auto* p_ext_imm = p_ext_imm_expr.as(); + const auto* p_ext_imm = p_ext_imm_expr.as(); if (auto p_ext = p_ext_imm ? p_ext_imm->value.as() : std::nullopt; p_ext.has_value()) { int64_t ext = *p_ext; @@ -85,7 +85,7 @@ class ThreadExtentChecker : public StmtExprVisitor { return StmtExprVisitor::Visit_(loop); } - ffi::Optional Visit_(const SBlockNode* block) { + ffi::Optional Visit_(const SBlockObj* block) { int old_thread_idx_x = thread_idx_x; if (block->annotations.count(tvm::s_tir::attr::kWarpExecution.name)) { thread_idx_x = thread_warp_size_; @@ -166,7 +166,7 @@ class VerifyGPUCodeObj : public PostprocObj { for (const auto& kv : mod->functions) { const GlobalVar& g_var = kv.first; const BaseFunc& base_func = kv.second; - if (const auto* function = base_func.as()) { + if (const auto* function = base_func.as()) { if (function->body.has_value() && !s_tir::ThreadExtentChecker::Check(function->body.value(), thread_warp_size_)) { return false; diff --git a/src/s_tir/meta_schedule/schedule/cuda/thread_bind.cc b/src/s_tir/meta_schedule/schedule/cuda/thread_bind.cc index 0c33cafc5a8b..0e13c279470f 100644 --- a/src/s_tir/meta_schedule/schedule/cuda/thread_bind.cc +++ b/src/s_tir/meta_schedule/schedule/cuda/thread_bind.cc @@ -72,7 +72,7 @@ ffi::Array BindSpatialLoop(Schedule sch, LoopRV loop, int64_t max_thread std::function get_factor) { int64_t extent = -1; PrimExpr e_imm_expr = Range::GetExtentFromBeginEnd(sch->Get(loop)->begin, sch->Get(loop)->end); - const auto* e_imm = e_imm_expr.as(); + const auto* e_imm = e_imm_expr.as(); if (auto e = e_imm ? e_imm->value.as() : std::nullopt; e.has_value()) { extent = *e; } else { @@ -117,7 +117,7 @@ void BindBlockThreadIdx(Schedule sch, SBlockRV block_rv, // int i_spatial_loop = -1; for (int i = 0; i < n; ++i) { const StmtSRef& loop_sref = loops[i]; - const ForNode* loop = TVM_SREF_TO_FOR(loop_sref); + const ForObj* loop = TVM_SREF_TO_FOR(loop_sref); support::GPUThreadTag thread_scope = GetThreadBindingTag(loop); if (IsBlockIdx(thread_scope)) { if (i_block_idx == -1) { diff --git a/src/s_tir/meta_schedule/schedule_rule/auto_inline.cc b/src/s_tir/meta_schedule/schedule_rule/auto_inline.cc index 05b054c75b02..28182702d070 100644 --- a/src/s_tir/meta_schedule/schedule_rule/auto_inline.cc +++ b/src/s_tir/meta_schedule/schedule_rule/auto_inline.cc @@ -54,7 +54,7 @@ bool IsInSpatialFunction(const s_tir::Schedule& sch, const tirx::StmtSRef& block const StmtSRefObj* sref = block_sref.get(); for (; sref->parent != nullptr; sref = sref->parent) { } - TVM_FFI_ICHECK(sref->stmt != nullptr && sref->stmt->IsInstance()); + TVM_FFI_ICHECK(sref->stmt != nullptr && sref->stmt->IsInstance()); return IsSpatialFunction(ffi::GetRef(GetRootFunction(sch->mod(), sref->stmt, nullptr))); } @@ -122,7 +122,7 @@ inline InlineType AutoInlineObj::CheckInline(const s_tir::Schedule& sch, StmtSRef block_sref = sch->GetSRef(block_rv); bool is_pure_sptial = IsInSpatialFunction(sch, block_sref); ScheduleState state = sch->state(); - const SBlockNode* block = TVM_SREF_TO_SBLOCK(block_sref); + const SBlockObj* block = TVM_SREF_TO_SBLOCK(block_sref); SBlockRealize realize = GetSBlockRealize(state, block_sref); // Cond 1. The block has only one write tensor if (block->writes.size() != 1) { @@ -168,10 +168,10 @@ inline InlineType AutoInlineObj::CheckInline(const s_tir::Schedule& sch, // Skip if: // - root block (no parent scope). // - root block is wrapped in a For loop, then parent is non-null, but there is still no enclosing - // SBlockNode. + // SBlockObj. const StmtSRefObj* p = block_sref->parent; for (; p != nullptr; p = p->parent) { - if (p->stmt->IsInstance()) break; + if (p->stmt->IsInstance()) break; } if (p == nullptr) { return InlineType::kNoInline; diff --git a/src/s_tir/meta_schedule/schedule_rule/cross_thread_reduction.cc b/src/s_tir/meta_schedule/schedule_rule/cross_thread_reduction.cc index b5da2d6dbba6..3791eeaa4fc4 100644 --- a/src/s_tir/meta_schedule/schedule_rule/cross_thread_reduction.cc +++ b/src/s_tir/meta_schedule/schedule_rule/cross_thread_reduction.cc @@ -227,7 +227,7 @@ class CrossThreadReductionObj : public ScheduleRuleObj { // - If the lowest common ancestor is a loop, the target block is also the first consumer. const tirx::StmtSRef& lca_sref = s_tir::GetSRefLowestCommonAncestor(s_tir::SBlockRVs2StmtSRefs(sch, consumers)); - if (consumers.size() > 1 && lca_sref->StmtAs() != nullptr) { + if (consumers.size() > 1 && lca_sref->StmtAs() != nullptr) { return std::make_tuple(false, s_tir::LoopRV{ffi::UnsafeInit()}, s_tir::SBlockRV{ffi::UnsafeInit()}, s_tir::LoopRV{ffi::UnsafeInit()}); } diff --git a/src/s_tir/meta_schedule/schedule_rule/multi_level_tiling.cc b/src/s_tir/meta_schedule/schedule_rule/multi_level_tiling.cc index af97da3d2dcf..92eba444c961 100644 --- a/src/s_tir/meta_schedule/schedule_rule/multi_level_tiling.cc +++ b/src/s_tir/meta_schedule/schedule_rule/multi_level_tiling.cc @@ -34,12 +34,12 @@ namespace s_tir { using namespace tvm::tirx; std::vector GetReadTensorNDims(const StmtSRef& block_sref) { - const SBlockNode* block = TVM_SREF_TO_SBLOCK(block_sref); - const VarNode* write_tensor = block->writes[0]->source.as_or_throw().get(); + const SBlockObj* block = TVM_SREF_TO_SBLOCK(block_sref); + const VarObj* write_tensor = block->writes[0]->source.as_or_throw().get(); int n = block->reads.size(); std::vector results(n, -1); for (int i = 0; i < n; ++i) { - const VarNode* read_tensor = block->reads[i]->source.as_or_throw().get(); + const VarObj* read_tensor = block->reads[i]->source.as_or_throw().get(); if (read_tensor != write_tensor) { results[i] = ffi::GetRef(read_tensor).as_or_throw().shape().size(); } @@ -238,7 +238,7 @@ std::vector MultiLevelTilingObj::TileLoopNest(State state, if (spatial_loop_product != -1) { PrimExpr extent_imm_expr = Range::GetExtentFromBeginEnd(sch->Get(loop)->begin, sch->Get(loop)->end); - const auto* extent_imm = extent_imm_expr.as(); + const auto* extent_imm = extent_imm_expr.as(); if (auto extent = extent_imm ? extent_imm->value.as() : std::nullopt; extent.has_value()) { spatial_loop_product *= *extent; @@ -342,13 +342,13 @@ std::vector MultiLevelTilingObj::AddAsyncPipeline(State state) const { // check the reduce loop contains exactly 3 for loops // therefore it matches the notation array size in the following code tirx::StmtSRef r_loop_sref = state->sch->GetSRef(state->tiles[r_indices_[0]].back()); - const ForNode* r_for_loop = TVM_SREF_TO_FOR(r_loop_sref); + const ForObj* r_for_loop = TVM_SREF_TO_FOR(r_loop_sref); ffi::Array seq = r_for_loop->body->seq; if (seq.size() != 3) { return {state}; } for (auto& stmt : seq) { - if (!stmt.as()) { + if (!stmt.as()) { return {state}; } } @@ -372,7 +372,7 @@ std::vector MultiLevelTilingObj::AddAsyncPipeline(State state) const { void MultiLevelTilingObj::AnnotateCooperativeFetching(Schedule* sch, const s_tir::SBlockRV& block) const { // Filter out invalid vector lanes according to the data type. - const s_tir::SBlockNode* block_node = (*sch)->GetSRef(block)->StmtAs(); + const s_tir::SBlockObj* block_node = (*sch)->GetSRef(block)->StmtAs(); TVM_FFI_ICHECK_EQ(block_node->writes.size(), 1); const DLDataType dtype = block_node->writes[0]->source.as_or_throw().dtype()->dtype; diff --git a/src/s_tir/meta_schedule/schedule_rule/multi_level_tiling_wide_vector.cc b/src/s_tir/meta_schedule/schedule_rule/multi_level_tiling_wide_vector.cc index 33a2ac492299..a8c8edb7f3d7 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 @@ -64,9 +64,9 @@ class MultiLevelTilingWideVectorObj : public MultiLevelTilingObj { std::pair>, ffi::Array> MultiLevelTilingWideVectorObj::SplitLoop(const Schedule& sch, SBlockRV block_rv, LoopRV loop_rv, int n_tiles) const { - const ForNode* loop = TVM_SREF_TO_FOR(sch->GetSRef(loop_rv)); + const ForObj* loop = TVM_SREF_TO_FOR(sch->GetSRef(loop_rv)); const tirx::StmtSRef block_sref = sch->GetSRef(block_rv); - const s_tir::SBlockNode* block_node = block_sref->StmtAs(); + const s_tir::SBlockObj* block_node = block_sref->StmtAs(); const s_tir::SBlockRealize block_realize = s_tir::GetSBlockRealize(sch->state(), block_sref); TVM_FFI_ICHECK(block_node && block_node->writes.size() == 1); @@ -100,7 +100,7 @@ MultiLevelTilingWideVectorObj::SplitLoop(const Schedule& sch, SBlockRV block_rv, } else { // We split the innermost spatial loop in a way that always uses the maximum vector length. PrimExpr extent_int_imm_expr = Range::GetExtentFromBeginEnd(loop->begin, loop->end); - const auto* extent_int_imm = extent_int_imm_expr.as(); + const auto* extent_int_imm = extent_int_imm_expr.as(); auto extent_int = extent_int_imm ? extent_int_imm->value.as() : std::nullopt; if (extent_int.has_value() && *extent_int > vec_len) { ffi::Array inner_splits = diff --git a/src/s_tir/meta_schedule/trace_apply.cc b/src/s_tir/meta_schedule/trace_apply.cc index 2bfc43056299..7287598b089f 100644 --- a/src/s_tir/meta_schedule/trace_apply.cc +++ b/src/s_tir/meta_schedule/trace_apply.cc @@ -43,9 +43,9 @@ using s_tir::Instruction; using s_tir::InstructionKind; using s_tir::IsSpatial; using s_tir::LoopRV; -using s_tir::LoopRVNode; +using s_tir::LoopRVObj; using s_tir::SBlockRV; -using s_tir::SBlockRVNode; +using s_tir::SBlockRVObj; using s_tir::Schedule; using s_tir::Trace; using s_tir::TranslateAddOutputRVs; @@ -132,8 +132,8 @@ std::vector ApplyAnchorTrace(Schedule sch, Trace anchor_trace) { auto is_inst_applicable = [&foreign_blocks, &foreign_loops](Instruction inst) { for (auto input : inst->inputs) { if (input == nullptr) continue; - if ((input.as() && foreign_blocks.count(input.as_or_throw())) || - (input.as() && foreign_loops.count(input.as_or_throw()))) { + if ((input.as() && foreign_blocks.count(input.as_or_throw())) || + (input.as() && foreign_loops.count(input.as_or_throw()))) { return false; } } @@ -145,9 +145,9 @@ std::vector ApplyAnchorTrace(Schedule sch, Trace anchor_trace) { // If we find an instruction that is not applicable, its outputs are recorded as "foreign" // to the target schedule. for (auto output : inst->outputs) { - if (output.as()) { + if (output.as()) { foreign_blocks.insert(output.as_or_throw()); - } else if (output.as()) { + } else if (output.as()) { foreign_loops.insert(output.as_or_throw()); } } diff --git a/src/s_tir/meta_schedule/utils.h b/src/s_tir/meta_schedule/utils.h index f089fc16c19b..83a069d7d883 100644 --- a/src/s_tir/meta_schedule/utils.h +++ b/src/s_tir/meta_schedule/utils.h @@ -272,7 +272,7 @@ inline std::string Concat(const ffi::Array& strs, const std::string */ inline s_tir::SBlockRV GetRVFromSRef(const s_tir::Schedule& sch, const tirx::StmtSRef& block_sref, const ffi::String& global_var_name) { - const s_tir::SBlockNode* block = TVM_SREF_TO_SBLOCK(block_sref); + const s_tir::SBlockObj* block = TVM_SREF_TO_SBLOCK(block_sref); return sch->GetSBlock(block->name_hint, global_var_name); } @@ -588,7 +588,7 @@ class SBlockCollector : public s_tir::StmtExprVisitor { using s_tir::StmtExprVisitor::Visit_; ffi::Optional Visit(ffi::AnyView value) override { - if (value.as()) return std::nullopt; + if (value.as()) return std::nullopt; return s_tir::StmtExprVisitor::Visit(value); } @@ -619,7 +619,7 @@ class SBlockCollector : public s_tir::StmtExprVisitor { for (const auto& [gv, base_func] : sch_->mod()->functions) { // `gv->name_hint` is the name of the function // `base_func` can be Function or relax::Function - if (const auto* func = base_func.as()) { + if (const auto* func = base_func.as()) { f_collect(ffi::GetRef(func), gv->name_hint); } } @@ -633,7 +633,7 @@ class SBlockCollector : public s_tir::StmtExprVisitor { private: /*! \brief Override the Stmt visiting behaviour */ - ffi::Optional Visit_(const s_tir::SBlockNode* block) override { + ffi::Optional Visit_(const s_tir::SBlockObj* block) override { TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(s_tir::StmtExprVisitor::Visit_(block)); TVM_FFI_ICHECK(block_names_.count(block->name_hint) == 0) << "Duplicated block name " << block->name_hint << " in function " << func_name_ diff --git a/src/s_tir/sblock_dependence_info.cc b/src/s_tir/sblock_dependence_info.cc index f406d74e2918..ae3e9ebe8387 100644 --- a/src/s_tir/sblock_dependence_info.cc +++ b/src/s_tir/sblock_dependence_info.cc @@ -35,7 +35,7 @@ class SBlockDependenceInfoCollector : public s_tir::StmtExprVisitor { using s_tir::StmtExprVisitor::Visit_; ffi::Optional Visit(ffi::AnyView value) override { - if (value.as()) return std::nullopt; + if (value.as()) return std::nullopt; return s_tir::StmtExprVisitor::Visit(value); } @@ -54,9 +54,9 @@ class SBlockDependenceInfoCollector : public s_tir::StmtExprVisitor { self_->sref2scope[scope] = SBlockScope(child_block_srefs); } - ffi::Optional Visit_(const s_tir::SBlockRealizeNode* realize) final { + ffi::Optional Visit_(const s_tir::SBlockRealizeObj* realize) final { block_frames_.emplace_back(); - const s_tir::SBlockNode* block = realize->block.get(); + const s_tir::SBlockObj* block = realize->block.get(); // Recursive visit TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Visit(block->body)); // `block->init` is not visited // Create SBlockInfo for the block @@ -68,8 +68,8 @@ class SBlockDependenceInfoCollector : public s_tir::StmtExprVisitor { return std::nullopt; } - ffi::Optional Visit_(const SeqStmtNode* seq_stmt) final { - // Set `seq_index` information for SeqStmtNode + ffi::Optional Visit_(const SeqStmtObj* seq_stmt) final { + // Set `seq_index` information for SeqStmtObj TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(s_tir::StmtExprVisitor::Visit_(seq_stmt)); SetSeqIndexInChildren(self_->stmt2ref, seq_stmt, false); return std::nullopt; diff --git a/src/s_tir/sblock_scope.cc b/src/s_tir/sblock_scope.cc index c82f75e044b2..ca567fbc43f1 100644 --- a/src/s_tir/sblock_scope.cc +++ b/src/s_tir/sblock_scope.cc @@ -53,7 +53,7 @@ void AddDependency(SBlockScopeObj* self, const StmtSRef& src, const StmtSRef& ds /******** Constructors ********/ -StmtSRef::StmtSRef(const StmtNode* stmt, StmtSRefObj* parent, int64_t seq_index) { +StmtSRef::StmtSRef(const StmtObj* stmt, StmtSRefObj* parent, int64_t seq_index) { ffi::ObjectPtr n = ffi::make_object(); n->stmt = stmt; n->parent = parent; @@ -84,7 +84,7 @@ SBlockScope::SBlockScope(const ffi::Array& child_block_srefs) { SMap> tensor_readers; SMap>& tensor_writers = n->tensor_writers; for (const StmtSRef& child_block_sref : child_block_srefs) { - const s_tir::SBlockNode* child_block = TVM_SREF_TO_SBLOCK(child_block_sref); + const s_tir::SBlockObj* child_block = TVM_SREF_TO_SBLOCK(child_block_sref); // Step 1. Update `tensor_readers` and `tensor_writers` for each tensor for (const TensorRegion& region : child_block->reads) { tensor_readers[region->source.as_or_throw()].push_back( @@ -149,7 +149,7 @@ ffi::Array SBlockScopeObj::GetDepsByDst(const StmtSRef& block_sref) * \brief Add a new statement to the stack, which becomes the current scope * \param stmt A for-loop statement or a block statement */ -void SRefTreeCreator::PushSRef(const StmtNode* stmt) { +void SRefTreeCreator::PushSRef(const StmtObj* stmt) { if (srefs_.empty()) { srefs_.push_back( StmtSRef(stmt, @@ -170,7 +170,7 @@ void SRefTreeCreator::PopAndRecordSRef() { srefs_.pop_back(); } -ffi::Optional SRefTreeCreator::Visit_(const ForNode* loop) { +ffi::Optional SRefTreeCreator::Visit_(const ForObj* loop) { if (!include_loops_) { TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Visit(loop->body)); } else { @@ -181,16 +181,16 @@ ffi::Optional SRefTreeCreator::Visit_(const ForNode* loop) { return std::nullopt; } -ffi::Optional SRefTreeCreator::Visit_(const s_tir::SBlockRealizeNode* realize) { - const s_tir::SBlockNode* block = realize->block.get(); +ffi::Optional SRefTreeCreator::Visit_(const s_tir::SBlockRealizeObj* realize) { + const s_tir::SBlockObj* block = realize->block.get(); PushSRef(block); TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Visit(block->body)); // `block->init` is not visited PopAndRecordSRef(); return std::nullopt; } -ffi::Optional SRefTreeCreator::Visit_(const SeqStmtNode* seq_stmt) { - // Set `seq_index` information for SeqStmtNode +ffi::Optional SRefTreeCreator::Visit_(const SeqStmtObj* seq_stmt) { + // Set `seq_index` information for SeqStmtObj TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(s_tir::StmtExprVisitor::Visit_(seq_stmt)); SetSeqIndexInChildren(stmt2ref_, seq_stmt, include_loops_); return std::nullopt; diff --git a/src/s_tir/schedule/analysis.h b/src/s_tir/schedule/analysis.h index 60c95d876f2d..7c71168e2dde 100644 --- a/src/s_tir/schedule/analysis.h +++ b/src/s_tir/schedule/analysis.h @@ -65,8 +65,8 @@ void VerifyCachedFlags(const ScheduleState& self); * \return The result Function where the root block belongs to * \note This function returns the pointer instead of ffi::ObjectRef to avoid later copy-on-write */ -const FunctionNode* GetRootFunction(const IRModule& mod, const StmtNode* root_block, - GlobalVar* result_g_var); +const FunctionObj* GetRootFunction(const IRModule& mod, const StmtObj* root_block, + GlobalVar* result_g_var); /*! * \brief Get the root node of the sref tree, which is the root block of the Function. @@ -106,9 +106,9 @@ struct ScopeBlockLoopInfo { /*! \brief A list of the leaf blocks, from left to right */ std::vector realizes; /*! \brief The loop vars bound to spatial block iters */ - std::unordered_set spatial_vars; + std::unordered_set spatial_vars; /*! \brief The loop vars bound to non-spatial block iters */ - std::unordered_set non_spatial_vars; + std::unordered_set non_spatial_vars; }; /*! @@ -288,8 +288,8 @@ ffi::Map GetBindings(const SBlockRealize& realize); * block iter nor a reduction block iter */ bool GetVarsTouchedByBlockIters(const SBlockRealize& block_realize, - std::unordered_set* data_par_vars, - std::unordered_set* reduce_vars); + std::unordered_set* data_par_vars, + std::unordered_set* reduce_vars); /******** Loop properties ********/ /*! @@ -407,7 +407,7 @@ ffi::Array GetConsumers(const StmtSRef& block_sref, const SBlockScope& * \return A list of all blocks that write to some output tensor * block */ -ffi::Array GetOutputBlocks(const ScheduleState& self, const SBlockNode* scope_block); +ffi::Array GetOutputBlocks(const ScheduleState& self, const SBlockObj* scope_block); /*! * \brief A solution to split a ordered list of subtrees into two parts, @@ -438,7 +438,7 @@ struct ProducerConsumerSplit { const ScheduleState& state, const ffi::Array& subtrees, const ffi::Array& producer_block_srefs, const ffi::Array& consumer_block_srefs, - std::unordered_map* block2realize); + std::unordered_map* block2realize); }; /******** Block-tensor relation ********/ diff --git a/src/s_tir/schedule/analysis/analysis.cc b/src/s_tir/schedule/analysis/analysis.cc index cd9c07bf432f..0697befdcd6b 100644 --- a/src/s_tir/schedule/analysis/analysis.cc +++ b/src/s_tir/schedule/analysis/analysis.cc @@ -40,14 +40,14 @@ TVM_FFI_STATIC_INIT_BLOCK() { /******** IR Module ********/ -const FunctionNode* GetRootFunction(const IRModule& mod, const StmtNode* root_block, - GlobalVar* result_g_var) { +const FunctionObj* GetRootFunction(const IRModule& mod, const StmtObj* root_block, + GlobalVar* result_g_var) { for (const auto& kv : mod->functions) { const GlobalVar& g_var = kv.first; const BaseFunc& base_func = kv.second; - if (const auto* func = base_func.as()) { + if (const auto* func = base_func.as()) { if (const auto* realize = func->body.has_value() && func->body.value()->size() == 1 - ? func->body.value()->seq[0].as() + ? func->body.value()->seq[0].as() : nullptr) { if (realize->block.get() == root_block) { if (result_g_var != nullptr) { @@ -111,7 +111,7 @@ Definition of a scope that is a stage pipeline: const StmtSRefObj* p = sref->parent; const StmtSRefObj* subtree = sref.get(); for (; p != nullptr; subtree = p, p = p->parent) { - if (p->stmt->IsInstance()) { + if (p->stmt->IsInstance()) { scope_root_sref = ffi::GetRef(p); scope_root_subtree = ffi::GetRef(subtree); break; @@ -125,7 +125,7 @@ Definition of a scope that is a stage pipeline: if (require_stage_pipeline && self->enable_check) { bool stage_pipeline = self->GetSBlockInfo(scope_root_sref).stage_pipeline; if (stage_pipeline == false) { - const SBlockNode* block = TVM_SREF_TO_SBLOCK(scope_root_sref); + const SBlockObj* block = TVM_SREF_TO_SBLOCK(scope_root_sref); throw MakeScheduleError(self->mod, ffi::GetRef(block)); } } @@ -137,11 +137,11 @@ ScopeBlockLoopInfo GetScopeBlockLoopInfo(const SBlock& scope_block) { using StmtExprVisitor::Visit_; ffi::Optional Visit(ffi::AnyView value) override { - if (value.as()) return std::nullopt; + if (value.as()) return std::nullopt; return StmtExprVisitor::Visit(value); } - ffi::Optional Visit_(const SBlockRealizeNode* realize) final { + ffi::Optional Visit_(const SBlockRealizeObj* realize) final { result.realizes.push_back(ffi::GetRef(realize)); const ffi::Array& iter_vars = realize->block->iter_vars; const ffi::Array& iter_values = realize->iter_values; @@ -150,7 +150,7 @@ ScopeBlockLoopInfo GetScopeBlockLoopInfo(const SBlock& scope_block) { for (int i = 0; i < n; ++i) { const IterVar& iter_var = iter_vars[i]; const PrimExpr& iter_value = iter_values[i]; - std::unordered_set* vars = nullptr; + std::unordered_set* vars = nullptr; if (iter_var->iter_type == IterVarType::kDataPar) { vars = &result.spatial_vars; } else { @@ -202,7 +202,7 @@ bool IsDominantBlock(const ScheduleState& self, const StmtSRef& scope_root_sref, std::unordered_map, ffi::ObjectPtrHash, ffi::ObjectPtrEqual> tensor_writers; CheckSRefHigherOrEqual(scope_root_sref, block_sref); - const SBlockNode* maybe_root_block = scope_root_sref->StmtAs(); + const SBlockObj* maybe_root_block = scope_root_sref->StmtAs(); if (maybe_root_block) { SBlockScope scope = self->GetSBlockScope(scope_root_sref); tensor_writers = scope->tensor_writers; @@ -217,7 +217,7 @@ bool IsDominantBlock(const ScheduleState& self, const StmtSRef& scope_root_sref, } } // Check whether the input block is the only writer of its outputs - const SBlockNode* block = TVM_SREF_TO_SBLOCK(block_sref); + const SBlockObj* block = TVM_SREF_TO_SBLOCK(block_sref); for (const TensorRegion& write_region : block->writes) { if (tensor_writers.count(write_region->source.as_or_throw())) { if (tensor_writers.at(write_region->source.as_or_throw()).size() != 1) { @@ -240,7 +240,7 @@ bool IsDominantBlock(const ScheduleState& self, const StmtSRef& scope_root_sref, int CheckCompleteBlockErrorCode(const ScheduleState& self, const StmtSRef& block_sref, const StmtSRef& scope_root_sref) { // Cond 1. All block vars are data parallel - const SBlockNode* block = TVM_SREF_TO_SBLOCK(block_sref); + const SBlockObj* block = TVM_SREF_TO_SBLOCK(block_sref); for (const IterVar& iter_var : block->iter_vars) { if (iter_var->iter_type != kDataPar) { return 1; @@ -252,7 +252,7 @@ int CheckCompleteBlockErrorCode(const ScheduleState& self, const StmtSRef& block return 2; } // Cond 3. No overlap between the tensors the block reads and writes - std::unordered_set written_tensors; + std::unordered_set written_tensors; written_tensors.reserve(block->writes.size()); for (const TensorRegion& write : block->writes) { written_tensors.insert(write->source.as_or_throw().get()); @@ -316,7 +316,7 @@ void CheckCompleteBlock(const ScheduleState& self, const StmtSRef& block_sref, int error_code = CheckCompleteBlockErrorCode(self, block_sref, scope_root_sref); if (error_code != 0) { - const SBlockNode* block = TVM_SREF_TO_SBLOCK(block_sref); + const SBlockObj* block = TVM_SREF_TO_SBLOCK(block_sref); throw MakeScheduleError(self->mod, ffi::GetRef(block), error_code); } @@ -333,7 +333,7 @@ void CheckCompleteBlock(const ScheduleState& self, const StmtSRef& block_sref, */ int CheckReductionBlockErrorCode(const ScheduleState& self, const StmtSRef& block_sref, const StmtSRef& scope_root_sref) { - const SBlockNode* block = TVM_SREF_TO_SBLOCK(block_sref); + const SBlockObj* block = TVM_SREF_TO_SBLOCK(block_sref); // Cond 1. The block has the `init` statement. if (!block->init.has_value()) { return 1; @@ -391,7 +391,7 @@ void CheckReductionBlock(const ScheduleState& self, const StmtSRef& block_sref, int error_code = CheckReductionBlockErrorCode(self, block_sref, scope_root_sref); if (error_code != 0) { - const SBlockNode* block = TVM_SREF_TO_SBLOCK(block_sref); + const SBlockObj* block = TVM_SREF_TO_SBLOCK(block_sref); throw MakeScheduleError(self->mod, ffi::GetRef(block), error_code); } @@ -439,7 +439,7 @@ void CheckCompleteOrReductionBlock(const ScheduleState& self, const StmtSRef& bl if (reduction_block_error_code == 0) { return; } - const SBlockNode* block = TVM_SREF_TO_SBLOCK(block_sref); + const SBlockObj* block = TVM_SREF_TO_SBLOCK(block_sref); throw MakeScheduleError( self->mod, ffi::GetRef(block), complete_block_error_code, reduction_block_error_code); } @@ -454,8 +454,7 @@ void CheckSubtreeCompactDataflow(const ScheduleState& self, const StmtSRef& subt violate_block_(std::move(violate_block)), local_complete_block_code_(local_complete_block_code), local_reduction_block_code_(local_reduction_block_code) { - TVM_FFI_ICHECK(subtree_root_->IsInstance() || - subtree_root_->IsInstance()); + TVM_FFI_ICHECK(subtree_root_->IsInstance() || subtree_root_->IsInstance()); } ffi::String FastErrorString() const final { return "ScheduleError: The queried subtree root in SRef tree does not have compact dataflow, " @@ -492,7 +491,7 @@ void CheckSubtreeCompactDataflow(const ScheduleState& self, const StmtSRef& subt int local_complete_block_code = CheckCompleteBlockErrorCode(self, block_sref, subtree_root), local_reduction_block_code = CheckReductionBlockErrorCode(self, block_sref, subtree_root); if (local_complete_block_code != 0 && local_reduction_block_code != 0) { - const SBlockNode* block = TVM_SREF_TO_SBLOCK(block_sref); + const SBlockObj* block = TVM_SREF_TO_SBLOCK(block_sref); throw MakeScheduleError( self->mod, ffi::GetRef(subtree_root->stmt), ffi::GetRef(block), local_complete_block_code, local_reduction_block_code); @@ -502,9 +501,9 @@ void CheckSubtreeCompactDataflow(const ScheduleState& self, const StmtSRef& subt bool IsOutputBlock(const ScheduleState& self, const StmtSRef& block_sref, const StmtSRef& scope_root_sref) { - const SBlockNode* scope_root = TVM_SREF_TO_SBLOCK(scope_root_sref); - const SBlockNode* block = TVM_SREF_TO_SBLOCK(block_sref); - std::unordered_set scope_allocated; + const SBlockObj* scope_root = TVM_SREF_TO_SBLOCK(scope_root_sref); + const SBlockObj* block = TVM_SREF_TO_SBLOCK(block_sref); + std::unordered_set scope_allocated; scope_allocated.reserve(scope_root->alloc_tensors.size()); for (const TensorVar& tensor : scope_root->alloc_tensors) { scope_allocated.insert(tensor.get()); @@ -533,12 +532,12 @@ void CheckNotOutputBlock(const ScheduleState& self, const StmtSRef& block_sref, SBlock block_; }; if (IsOutputBlock(self, block_sref, scope_root_sref)) { - const SBlockNode* block = TVM_SREF_TO_SBLOCK(block_sref); + const SBlockObj* block = TVM_SREF_TO_SBLOCK(block_sref); throw MakeScheduleError(self->mod, ffi::GetRef(block)); } } -std::vector GetSBlockVarTypes(const SBlockNode* block) { +std::vector GetSBlockVarTypes(const SBlockObj* block) { std::vector results; results.reserve(block->iter_vars.size()); for (const IterVar& iter_var : block->iter_vars) { @@ -548,12 +547,12 @@ std::vector GetSBlockVarTypes(const SBlockNode* block) { } std::vector GetSBlockVarTypes(const StmtSRef& block_sref) { - const SBlockNode* block = TVM_SREF_TO_SBLOCK(block_sref); + const SBlockObj* block = TVM_SREF_TO_SBLOCK(block_sref); return GetSBlockVarTypes(block); } bool IsWriteCache(const StmtSRef& block_sref) { - const SBlockNode* block = TVM_SREF_TO_SBLOCK(block_sref); + const SBlockObj* block = TVM_SREF_TO_SBLOCK(block_sref); if (block->writes.size() != 1) { return false; } @@ -611,7 +610,7 @@ void CheckPartialAffineBinding(const ScheduleState& self, SBlock block, ffi::Optional high_exclusive) : mod_(std::move(mod)), block_(std::move(block)) { if (high_exclusive.has_value()) { - high_exclusive_loop_ = high_exclusive.value()->StmtAs(); + high_exclusive_loop_ = high_exclusive.value()->StmtAs(); } } ffi::String FastErrorString() const final { @@ -638,7 +637,7 @@ void CheckPartialAffineBinding(const ScheduleState& self, SBlock block, ffi::Array LocationsOfInterest() const final { return {block_}; } IRModule mod_; SBlock block_; - const ForNode* high_exclusive_loop_{nullptr}; + const ForObj* high_exclusive_loop_{nullptr}; }; StmtSRef block_sref = self->stmt2ref.at(block.get()); @@ -688,7 +687,7 @@ void CheckBlockHasTrivialBinding(const ScheduleState& self, const StmtSRef& bloc if (!IsTrivialBinding(self, block_sref)) { throw MakeScheduleError( - self->mod, ffi::GetRef(block_sref->StmtAs())); + self->mod, ffi::GetRef(block_sref->StmtAs())); } } @@ -699,7 +698,7 @@ ffi::Map LoopDomainOfSRefTreePath(const StmtSRef& low_inclusive, const StmtSRefObj* p = low_inclusive.get(); const StmtSRefObj* limit = high_exclusive.has_value() ? high_exclusive.value().get() : nullptr; for (; p != limit; p = p->parent) { - const ForNode* loop = p->StmtAs(); + const ForObj* loop = p->StmtAs(); if (loop == nullptr) { break; } @@ -707,7 +706,7 @@ ffi::Map LoopDomainOfSRefTreePath(const StmtSRef& low_inclusive, } if (!IsGlobalMemoryScope(extra_relax_scope)) { for (; p; p = p->parent) { - if (const ForNode* loop = p->StmtAs()) { + if (const ForObj* loop = p->StmtAs()) { if (tvm::s_tir::GetThreadBinding(loop).has_value()) { const ffi::String thread_tag = tvm::s_tir::GetThreadBinding(loop).value(); if (CanRelaxStorageUnderThread(extra_relax_scope, ParseThreadBindingTag(thread_tag))) { @@ -721,7 +720,7 @@ ffi::Map LoopDomainOfSRefTreePath(const StmtSRef& low_inclusive, } ffi::Map GetBindings(const SBlockRealize& realize) { - const SBlockNode* block = realize->block.get(); + const SBlockObj* block = realize->block.get(); const ffi::Array& all_lhs = block->iter_vars; const ffi::Array& all_rhs = realize->iter_values; TVM_FFI_ICHECK_EQ(all_lhs.size(), all_rhs.size()); @@ -735,8 +734,8 @@ ffi::Map GetBindings(const SBlockRealize& realize) { } bool GetVarsTouchedByBlockIters(const SBlockRealize& block_realize, - std::unordered_set* data_par_vars, - std::unordered_set* reduce_vars) { + std::unordered_set* data_par_vars, + std::unordered_set* reduce_vars) { SBlock block = block_realize->block; TVM_FFI_CHECK(block_realize->block.same_as(block), ValueError) << "The input `block_realize` is required to be the exact BlockRealize of the " @@ -748,7 +747,7 @@ bool GetVarsTouchedByBlockIters(const SBlockRealize& block_realize, for (int i = 0; i < n; ++i) { const IterVar& iter_var = block->iter_vars[i]; const PrimExpr& iter_value = block_realize->iter_values[i]; - std::unordered_set* set = nullptr; + std::unordered_set* set = nullptr; if (iter_var->iter_type == IterVarType::kDataPar) { set = data_par_vars; } else if (iter_var->iter_type == IterVarType::kCommReduce) { @@ -791,7 +790,7 @@ void CheckLoopStartsWithZero(const ScheduleState& self, const StmtSRef& loop_sre IRModule mod_; For loop_; }; - const ForNode* loop = TVM_SREF_TO_FOR(loop_sref); + const ForObj* loop = TVM_SREF_TO_FOR(loop_sref); if (!analyzer->CanProve(loop->begin == 0)) { throw MakeScheduleError(self->mod, ffi::GetRef(loop)); } @@ -816,7 +815,7 @@ ffi::Array GetChildBlockRealizeOnSRefTree(const StmtSRef& parent_ using StmtExprVisitor::Visit_; ffi::Optional Visit(ffi::AnyView value) override { - if (value.as()) return std::nullopt; + if (value.as()) return std::nullopt; return StmtExprVisitor::Visit(value); } @@ -826,7 +825,7 @@ ffi::Array GetChildBlockRealizeOnSRefTree(const StmtSRef& parent_ return std::move(collector->result_); } - ffi::Optional Visit_(const SBlockRealizeNode* block_realize) final { + ffi::Optional Visit_(const SBlockRealizeObj* block_realize) final { result_.push_back(ffi::GetRef(block_realize)); return std::nullopt; } @@ -834,11 +833,11 @@ ffi::Array GetChildBlockRealizeOnSRefTree(const StmtSRef& parent_ ffi::Array result_; }; - if (parent_sref->stmt->IsInstance()) { - const auto* loop = static_cast(parent_sref->stmt); + if (parent_sref->stmt->IsInstance()) { + const auto* loop = static_cast(parent_sref->stmt); return Collector::Collect(loop->body); - } else if (parent_sref->stmt->IsInstance()) { - const auto* block = static_cast(parent_sref->stmt); + } else if (parent_sref->stmt->IsInstance()) { + const auto* block = static_cast(parent_sref->stmt); return Collector::Collect(block->body); } TVM_FFI_ICHECK(false) << "Unreachable"; @@ -851,7 +850,7 @@ SBlockRealize CheckGetSingleChildBlockRealizeOnSRefTree(const ScheduleState& sel public: explicit NonSingleChildBlockError(IRModule mod, const StmtSRef& sref) : mod_(std::move(mod)), stmt_(ffi::GetRef(sref->stmt)) { - sref_type_ = stmt_.as() != nullptr ? "block" : "loop"; + sref_type_ = stmt_.as() != nullptr ? "block" : "loop"; } ffi::String FastErrorString() const final { @@ -885,18 +884,18 @@ SBlockRealize GetSBlockRealize(const ScheduleState& self, const StmtSRef& block_ struct BlockRealizeFinder : public StmtExprVisitor { using StmtExprVisitor::Visit_; - explicit BlockRealizeFinder(const SBlockNode* target_sblock) + explicit BlockRealizeFinder(const SBlockObj* target_sblock) : target_sblock(target_sblock), result(nullptr) {} ffi::Optional Visit(ffi::AnyView stmt) final { - if (stmt.as()) return std::nullopt; + if (stmt.as()) return std::nullopt; if (result != nullptr) { return std::nullopt; } return StmtExprVisitor::Visit(stmt); } - ffi::Optional Visit_(const SBlockRealizeNode* block_realize) final { + ffi::Optional Visit_(const SBlockRealizeObj* block_realize) final { if (block_realize->block.get() == target_sblock) { result = block_realize; } @@ -904,13 +903,13 @@ SBlockRealize GetSBlockRealize(const ScheduleState& self, const StmtSRef& block_ return std::nullopt; } - const SBlockNode* target_sblock; - const SBlockRealizeNode* result; + const SBlockObj* target_sblock; + const SBlockRealizeObj* result; }; - const SBlockNode* block = TVM_SREF_TO_SBLOCK(block_sref); + const SBlockObj* block = TVM_SREF_TO_SBLOCK(block_sref); if (block_sref->parent == nullptr) { - const FunctionNode* func = GetRootFunction(self->mod, block, nullptr); + const FunctionObj* func = GetRootFunction(self->mod, block, nullptr); return func->body.value()->seq[0].as_or_throw(); } else { auto finder = ffi::make_object(block); @@ -929,14 +928,14 @@ TVM_FFI_STATIC_INIT_BLOCK() { } IterVarType GetLoopIterType(const StmtSRef& loop_sref) { - const ForNode* loop = TVM_SREF_TO_FOR(loop_sref); + const ForObj* loop = TVM_SREF_TO_FOR(loop_sref); const Var& loop_var = loop->var; int n_spatial = 0; int n_reduce = 0; int n_other = 0; auto f_visit = [&loop_var, &n_spatial, &n_reduce, &n_other](const SBlockRealize& realize) -> ffi::Expected { - const SBlockNode* block = realize->block.get(); + const SBlockObj* block = realize->block.get(); // Number of block vars and their bindings TVM_FFI_ICHECK_EQ(realize->iter_values.size(), block->iter_vars.size()); size_t n = realize->iter_values.size(); @@ -1025,7 +1024,7 @@ std::pair, std::vector> CollectComputeLocation( // such a loop cannot be found, there is no more candidate and we just return. StmtSRef loop_boundary = consumers.size() > 1 ? GetSRefLowestCommonAncestor(consumers) : ffi::GetRef(consumers[0]->parent); - if (loop_boundary->StmtAs() == nullptr) { + if (loop_boundary->StmtAs() == nullptr) { return std::make_pair(location_srefs, location_indices); } @@ -1114,18 +1113,18 @@ ffi::Array GetConsumers(const StmtSRef& block_sref, const SBlockScope& return results; } -ffi::Array GetOutputBlocks(const ScheduleState& self, const SBlockNode* scope_block) { +ffi::Array GetOutputBlocks(const ScheduleState& self, const SBlockObj* scope_block) { struct OutputSBlockCollector : public StmtExprVisitor { using StmtExprVisitor::Visit_; ffi::Optional Visit(ffi::AnyView value) override { - if (value.as()) return std::nullopt; + if (value.as()) return std::nullopt; return StmtExprVisitor::Visit(value); } explicit OutputSBlockCollector(const ScheduleState& self) : self_(self) {} - ffi::Optional Visit_(const SBlockNode* block) override { + ffi::Optional Visit_(const SBlockObj* block) override { auto it = self_->stmt2ref.find(block); TVM_FFI_ICHECK(it != self_->stmt2ref.end()); auto block_sref = it->second; @@ -1152,7 +1151,7 @@ ProducerConsumerSplit ProducerConsumerSplit::Find( const ScheduleState& self, const ffi::Array& subtrees, const ffi::Array& producer_block_srefs, const ffi::Array& consumer_block_srefs, - std::unordered_map* block2realize) { + std::unordered_map* block2realize) { class InsertionPointNotFoundError : public ScheduleErrorContextObj { public: explicit InsertionPointNotFoundError(IRModule mod, int last_producer_position, @@ -1189,12 +1188,12 @@ ProducerConsumerSplit ProducerConsumerSplit::Find( using StmtExprVisitor::Visit_; ffi::Optional Visit(ffi::AnyView value) override { - if (value.as()) return std::nullopt; + if (value.as()) return std::nullopt; return StmtExprVisitor::Visit(value); } - ffi::Optional Visit_(const SBlockRealizeNode* realize) final { - const SBlockNode* block = realize->block.get(); + ffi::Optional Visit_(const SBlockRealizeObj* realize) final { + const SBlockObj* block = realize->block.get(); if (block2realize_) { block2realize_->emplace(block, realize); } @@ -1207,9 +1206,9 @@ ProducerConsumerSplit ProducerConsumerSplit::Find( return std::nullopt; } - std::unordered_map* block2realize_; - std::unordered_set producer_blocks_; - std::unordered_set consumer_blocks_; + std::unordered_map* block2realize_; + std::unordered_set producer_blocks_; + std::unordered_set consumer_blocks_; int n_producers_visited_ = 0; int n_consumers_visited_ = 0; }; @@ -1324,7 +1323,7 @@ std::pair, bool> GetTensorDefiningSite(const StmtSRef& b // match_tensors. const StmtSRefObj* defining_site_sref = block_sref.get(); while (defining_site_sref != nullptr) { - const auto* block = defining_site_sref->StmtAs(); + const auto* block = defining_site_sref->StmtAs(); // If this sref is not a block sref, skip it. if (block == nullptr) { defining_site_sref = defining_site_sref->parent; @@ -1363,14 +1362,16 @@ void AddShapeVarBounds(const ScheduleState& state, const StmtSRefObj* sref, while (sref->parent != nullptr) { sref = sref->parent; } - const FunctionNode* f = GetRootFunction(state->mod, sref->stmt, nullptr); + const FunctionObj* f = GetRootFunction(state->mod, sref->stmt, nullptr); + auto context = analyzer->GetConstraintContext(); for (const Var& param : f->params) { if (auto tensor = param.as()) { for (const PrimExpr& e : tensor.value().shape()) { - analyzer->MarkGlobalNonNegValue(e); + context->Append(sym::PredicateConstraint(e >= 0, /*assume_always=*/true)); } } } + analyzer->SyncConstraintContext(context); } /******** Misc ********/ @@ -1427,11 +1428,11 @@ AnalyzeReadWritePattern(const TensorRegion& read_region, const TensorRegion& wri std::make_tuple(false, false, false, false, false, false); // Step 1. Extract the write indices int w_dim = write_region->source.as_or_throw().shape().size(); - std::unordered_map var2idx; + std::unordered_map var2idx; var2idx.reserve(w_dim); for (int i = 0; i < w_dim; ++i) { const Range& dom = write_region->region[i]; - if (!dom->extent.as()) { + if (!dom->extent.as()) { return kNotExist; } if (auto var = dom->begin.as()) { @@ -1447,11 +1448,11 @@ AnalyzeReadWritePattern(const TensorRegion& read_region, const TensorRegion& wri std::vector mapped(r_dim, -1); for (int i = 0; i < r_dim; ++i) { const Range& dom = read_region->region[i]; - if (!dom->extent.as()) { + if (!dom->extent.as()) { return kNotExist; } // Case 1. Read index is a constant - if (dom->begin.as()) { + if (dom->begin.as()) { no_const_read = false; continue; } @@ -1459,7 +1460,7 @@ AnalyzeReadWritePattern(const TensorRegion& read_region, const TensorRegion& wri // where `var` is a write index and `const` is an optional constant shift ffi::Optional opt_const = std::nullopt; ffi::Optional opt_var = AnalyzeVarWithShift(dom->begin, &opt_const); - const VarNode* var = opt_var.has_value() ? opt_var.value().get() : nullptr; + const VarObj* var = opt_var.has_value() ? opt_var.value().get() : nullptr; if (var == nullptr || !var2idx.count(var)) { return kNotExist; } @@ -1530,7 +1531,7 @@ void CheckStorageScope(const ScheduleState& self, ffi::String storage_scope) { } bool IsSpatial(const StmtSRef& block_sref) { - const SBlockNode* block = TVM_SREF_TO_SBLOCK(block_sref); + const SBlockObj* block = TVM_SREF_TO_SBLOCK(block_sref); for (const IterVar& iter_var : block->iter_vars) { if (iter_var->iter_type != IterVarType::kDataPar) { return false; @@ -1547,7 +1548,7 @@ bool IsTrivialBinding(const ScheduleState& self, const StmtSRef& block_sref) { return false; } for (int i = 0, n = loops.size(); i < n; ++i) { - const ForNode* loop = TVM_SREF_TO_FOR(loops[i]); + const ForObj* loop = TVM_SREF_TO_FOR(loops[i]); if (binds[i].get() != loop->var.get()) { return false; } @@ -1566,15 +1567,15 @@ bool NeedsMultiLevelTiling(const ScheduleState& self, const StmtSRef& block_sref if (HasBeenMultiLevelTiled(block_sref)) { return false; } - const SBlockNode* block = TVM_SREF_TO_SBLOCK(block_sref); + const SBlockObj* block = TVM_SREF_TO_SBLOCK(block_sref); if (block->writes.size() != 1 || block->reads.empty() || IsSpatial(block_sref) || !IsTrivialBinding(self, block_sref)) { return false; } - const VarNode* write_tensor = block->writes[0]->source.as_or_throw().get(); + const VarObj* write_tensor = block->writes[0]->source.as_or_throw().get(); // Step 1. Sort out spatial block variables. Skip the block iters of domain [0, 1), since such // block iters distracts the following check of the unused block iters. - std::vector spatial_block_vars; + std::vector spatial_block_vars; spatial_block_vars.reserve(block->iter_vars.size()); for (const IterVar& block_var : block->iter_vars) { bool has_trivial_dom = @@ -1586,10 +1587,10 @@ bool NeedsMultiLevelTiling(const ScheduleState& self, const StmtSRef& block_sref // Step 2. Enumerate each read region, check the number of block vars that are not used // to index the read region int total_unused_block_vars = 0; - std::unordered_set read_tensors; + std::unordered_set read_tensors; read_tensors.reserve(block->reads.size()); for (const TensorRegion& tensor_region : block->reads) { - const VarNode* tensor = tensor_region->source.as_or_throw().get(); + const VarObj* tensor = tensor_region->source.as_or_throw().get(); const ffi::Array& regions = tensor_region->region; // Step 2.1. Duplication of read tensors are not allowed if (read_tensors.insert(tensor).second == false) { @@ -1600,9 +1601,9 @@ bool NeedsMultiLevelTiling(const ScheduleState& self, const StmtSRef& block_sref continue; } // Step 2.3. Collect the block vars that are used to index the read region - std::unordered_set vars; + std::unordered_set vars; for (const Range& range : regions) { - if (!range->extent.as()) { + if (!range->extent.as()) { return false; } for (const Var& var : UndefinedVars(range->begin)) { @@ -1611,7 +1612,7 @@ bool NeedsMultiLevelTiling(const ScheduleState& self, const StmtSRef& block_sref } // Step 2.4. Check if the block vars are not used to index the read region int n_unused_block_vars = 0; - for (const VarNode* block_var : spatial_block_vars) { + for (const VarObj* block_var : spatial_block_vars) { if (vars.count(block_var) == 0) { ++n_unused_block_vars; } @@ -1648,7 +1649,7 @@ std::pair GetCumulativeSpaceAndReductionLength(const s_tir::Sc if (type == s_tir::kDataPar) { PrimExpr extent_imm_expr = Range::GetExtentFromBeginEnd(TVM_SREF_TO_FOR(loop_sref)->begin, TVM_SREF_TO_FOR(loop_sref)->end); - const auto* extent_imm = extent_imm_expr.as(); + const auto* extent_imm = extent_imm_expr.as(); auto extent = extent_imm ? extent_imm->value.as() : std::nullopt; if (extent.has_value() && *extent != -1) { cum_space_len *= *extent; @@ -1658,7 +1659,7 @@ std::pair GetCumulativeSpaceAndReductionLength(const s_tir::Sc } else if (type == s_tir::kCommReduce) { PrimExpr extent_imm_expr = Range::GetExtentFromBeginEnd(TVM_SREF_TO_FOR(loop_sref)->begin, TVM_SREF_TO_FOR(loop_sref)->end); - const auto* extent_imm = extent_imm_expr.as(); + const auto* extent_imm = extent_imm_expr.as(); auto extent = extent_imm ? extent_imm->value.as() : std::nullopt; if (extent.has_value() && *extent != -1) { cum_reduce_len *= *extent; @@ -1676,7 +1677,7 @@ bool NeedsRFactorOrCrossThreadReduction(const s_tir::ScheduleState& self, // const tirx::StmtSRef& block_sref, // int64_t max_parallel_extent, // int64_t max_parallel_basic) { - const SBlockNode* block = TVM_SREF_TO_SBLOCK(block_sref); + const SBlockObj* block = TVM_SREF_TO_SBLOCK(block_sref); ffi::Array loops = GetLoops(block_sref); // Cond 1. The block must have at lease one write tensor @@ -1711,15 +1712,15 @@ bool NeedsRFactorOrCrossThreadReduction(const s_tir::ScheduleState& self, // } // Cond 5. - const ForNode* loop_i = TVM_SREF_TO_FOR(loops[i]); + const ForObj* loop_i = TVM_SREF_TO_FOR(loops[i]); if (i < loops.size() - 1) { - const ForNode* loop_i1 = TVM_SREF_TO_FOR(loops[i + 1]); + const ForObj* loop_i1 = TVM_SREF_TO_FOR(loops[i + 1]); if (loop_i->body->size() != 1 || loop_i->body->seq[0].get() != loop_i1) { return false; } } else { const auto* block_realize = - loop_i->body->size() == 1 ? loop_i->body->seq[0].as() : nullptr; + loop_i->body->size() == 1 ? loop_i->body->seq[0].as() : nullptr; if (!block_realize || block_realize->block.get() != block) { return false; } @@ -1747,7 +1748,7 @@ bool NeedsRFactorOrCrossThreadReduction(const s_tir::ScheduleState& self, // PrimExpr SimplifyNonTrivialExpr(const PrimExpr& expr, sym::AnalyzerObj* analyzer) { auto simplified = analyzer->Simplify(expr); - if (simplified->IsInstance()) { + if (simplified->IsInstance()) { return expr; } else { return simplified; @@ -1759,11 +1760,11 @@ struct TensorIntrinDescInfo { /*! \brief The block of the description function, which is the (unique) direct child of the root * block. */ - const SBlockRealizeNode* desc_block = nullptr; + const SBlockRealizeObj* desc_block = nullptr; /*! \brief The loops of the description function, in the order from outer loops to inner ones. */ - std::vector desc_loops; + std::vector desc_loops; /*! \brief The loop variables. */ - std::unordered_set desc_loop_vars; + std::unordered_set desc_loop_vars; }; /*! @@ -1777,7 +1778,7 @@ TensorIntrinDescInfo ExtractTensorIntrinDescInfo(sym::AnalyzerObj* analyzer, TensorIntrinDescInfo info; const auto* desc_scope_realize = desc_func->body.has_value() && desc_func->body.value()->size() == 1 - ? desc_func->body.value()->seq[0].as() + ? desc_func->body.value()->seq[0].as() : nullptr; TVM_FFI_ICHECK(desc_scope_realize); { @@ -1812,11 +1813,11 @@ ffi::Optional GetTensorizeLoopMapping(const s_tir::ScheduleState& // Step 2. Collect loops from block_sref const tirx::StmtSRef& scope_sref = GetScopeRoot(self, block_sref, false); TVM_SREF_TO_SBLOCK(scope_sref); - std::vector block_loops; - std::unordered_set block_loop_vars; + std::vector block_loops; + std::unordered_set block_loop_vars; { for (const tirx::StmtSRefObj* loop_sref = block_sref->parent;; loop_sref = loop_sref->parent) { - const auto* loop = loop_sref->StmtAs(); + const auto* loop = loop_sref->StmtAs(); if (loop == nullptr || loop->body->size() != 1) { break; } @@ -1829,9 +1830,9 @@ ffi::Optional GetTensorizeLoopMapping(const s_tir::ScheduleState& std::reverse(block_loops.begin(), block_loops.end()); } // Step 3. Map from block loops to desc block loops - const std::vector& desc_loops = desc_info.desc_loops; - const std::unordered_set& desc_loop_vars = desc_info.desc_loop_vars; - const SBlockRealizeNode* desc_block = desc_info.desc_block; + const std::vector& desc_loops = desc_info.desc_loops; + const std::unordered_set& desc_loop_vars = desc_info.desc_loop_vars; + const SBlockRealizeObj* desc_block = desc_info.desc_block; ffi::ObjectPtr ret = ffi::make_object(); const int n_block_vars = block->iter_values.size(); const int n_desc_vars = desc_block->iter_values.size(); @@ -1877,7 +1878,7 @@ ffi::Optional GetTensorizeLoopMapping(const s_tir::ScheduleState& for (int i_desc = n_desc_vars - 1; i_desc >= 0; --i_desc) { // Step 3.1. Find the corresponding loop of the i_desc-th block var of desc const PrimExpr& desc_bind = desc_block->iter_values[i_desc]; - const ForNode* desc_loop = nullptr; + const ForObj* desc_loop = nullptr; IterVarType iter_type_desc = iter_types_desc[i_desc]; for (int i = 0, n = desc_loops.size(); i < n; ++i) { // Check if desc_bind = loops[i]->var + stuff-irrelevant-of-loop-vars @@ -1889,13 +1890,12 @@ ffi::Optional GetTensorizeLoopMapping(const s_tir::ScheduleState& } } if (desc_loop == nullptr || - Range::GetExtentFromBeginEnd(desc_loop->begin, desc_loop->end).as() == - nullptr) { + Range::GetExtentFromBeginEnd(desc_loop->begin, desc_loop->end).as() == nullptr) { return std::nullopt; } PrimExpr int_desc_extent_expr = Range::GetExtentFromBeginEnd(desc_loop->begin, desc_loop->end); - const IntImmNode* int_desc_extent = int_desc_extent_expr.as(); + const IntImmObj* int_desc_extent = int_desc_extent_expr.as(); // Step 3.2. Find the corresponding iter_value of the target block with a matching iterator type ffi::Optional block_bind; @@ -1913,7 +1913,7 @@ ffi::Optional GetTensorizeLoopMapping(const s_tir::ScheduleState& // Step 3.3. Find the corresponding loop of the target block for (int i = 0, n = block_loops.size(); i < n; ++i) { // Check if block_bind = block_loops[i]->var + stuff-irrelevant-of-loop-vars - const ForNode* block_loop = block_loops[i]; + const ForObj* block_loop = block_loops[i]; const tirx::StmtSRef& block_loop_sref = self->stmt2ref.at(block_loop); // Skip i-th loop if it has already been mapped if (ret->loop_map.find(block_loop_sref) != ret->loop_map.end()) continue; @@ -1929,7 +1929,7 @@ ffi::Optional GetTensorizeLoopMapping(const s_tir::ScheduleState& PrimExpr int_block_extent_expr = Range::GetExtentFromBeginEnd(block_loops[i]->begin, block_loops[i]->end); - const IntImmNode* int_block_extent = int_block_extent_expr.as(); + const IntImmObj* int_block_extent = int_block_extent_expr.as(); // Check divisibility if (!int_block_extent) { @@ -2042,10 +2042,10 @@ class AutoTensorizeMappingProposer { // Step 2: Compute the tensor mask TVM_FFI_ICHECK_EQ(rhs_tensor_index.size(), lhs_tensor_index.size()); int num_tensors = rhs_tensor_index.size(); - std::unordered_map> rhs_tensor_masks, lhs_tensor_masks; + std::unordered_map> rhs_tensor_masks, lhs_tensor_masks; // helper function to initialize or update the tensor mask - auto update_mask = [&](const VarNode* var, - std::unordered_map>* masks, int i) { + auto update_mask = [&](const VarObj* var, + std::unordered_map>* masks, int i) { if (!masks->count(var)) { (*masks)[var].resize(num_tensors); } @@ -2082,11 +2082,11 @@ class AutoTensorizeMappingProposer { // have the same iter type. std::unordered_map mask_to_rhs_vars; for (const auto& kv : rhs_tensor_masks) { - const VarNode* rhs_var = kv.first; + const VarObj* rhs_var = kv.first; const TensorMask& mask = kv.second; mask_to_rhs_vars[mask].insert(ffi::GetRef(rhs_var)); } - std::unordered_map rhs_var_iter_type; + std::unordered_map rhs_var_iter_type; for (const auto& iter : extractor_->rhs_iters_) { rhs_var_iter_type.emplace(iter->var.get(), iter->iter_type); } diff --git a/src/s_tir/schedule/analysis/layout.cc b/src/s_tir/schedule/analysis/layout.cc index 7799c175c8c2..879347faf5d6 100644 --- a/src/s_tir/schedule/analysis/layout.cc +++ b/src/s_tir/schedule/analysis/layout.cc @@ -82,9 +82,9 @@ class SplitExprCollector { private: void Visit(const sym::IterSplitExpr& expr) { if (auto var = expr->source->source.as()) { - const auto* lower_factor_imm = expr->lower_factor.as(); + const auto* lower_factor_imm = expr->lower_factor.as(); auto lower_factor = lower_factor_imm ? lower_factor_imm->value.as() : std::nullopt; - const auto* extent_imm = expr->extent.as(); + const auto* extent_imm = expr->extent.as(); auto extent = extent_imm ? extent_imm->value.as() : std::nullopt; if (!lower_factor.has_value() || !extent.has_value()) { failed_ = true; @@ -118,7 +118,7 @@ ffi::Optional SuggestIndexMap(const TensorVar& tensor, int n_loops = loops.size(); // Step 1. Collect the domains and indices of loop variables ffi::Map input_iters; - std::unordered_map var2id; + std::unordered_map var2id; var2id.reserve(n_loops); for (int i = 0; i < n_loops; ++i) { const For& loop = loops[i]; @@ -168,7 +168,11 @@ ffi::Optional SuggestIndexMap(const TensorVar& tensor, ](ffi::Array indices) -> ffi::Array { TVM_FFI_ICHECK_EQ(indices.size(), shape.size()); for (int i = 0, n = indices.size(); i < n; ++i) { - analyzer->Bind(indices[i], Range::FromMinExtent(0, shape[i])); + { + auto context = analyzer->GetConstraintContext(); + context->Append(sym::RangeConstraint(indices[i], Range::FromMinExtent(0, shape[i]))); + analyzer->SyncConstraintContext(context); + } } // Step 5.1: Fuse all indices into a flattened one PrimExpr index = @@ -203,7 +207,12 @@ ffi::Optional SuggestIndexMap(const TensorVar& tensor, for (int i = 0, n = indices.size(); i < n; ++i) { const Var& index = indices[inverse_order[i]]; inv_permuted_indices.push_back(index); - analyzer->Bind(index, Range::FromMinExtent(0, IntImm::Int32(split_exprs[i].extent))); + { + auto context = analyzer->GetConstraintContext(); + context->Append(sym::RangeConstraint( + index, Range::FromMinExtent(0, IntImm::Int32(split_exprs[i].extent)))); + analyzer->SyncConstraintContext(context); + } } // Step 6.2: Fuse all the indices. This is the inverse of Step 5.2. diff --git a/src/s_tir/schedule/analysis/reducer.cc b/src/s_tir/schedule/analysis/reducer.cc index 732f038118a6..4e55dc048501 100644 --- a/src/s_tir/schedule/analysis/reducer.cc +++ b/src/s_tir/schedule/analysis/reducer.cc @@ -62,7 +62,7 @@ class PatternMatcher : public StmtExprVisitor { explicit PatternMatcher(ffi::Array pattern) : pattern_(std::move(pattern)) {} - ffi::Optional Visit_(const VarNode* op) final { + ffi::Optional Visit_(const VarObj* op) final { auto it = filled_map_.find(op); if (it == filled_map_.end()) { filled_map_.insert_or_assign(op, expr_to_match_); @@ -76,8 +76,8 @@ class PatternMatcher : public StmtExprVisitor { return std::nullopt; } - ffi::Optional Visit_(const LetNode* op) final { - const auto* ptr = expr_to_match_.as(); + ffi::Optional Visit_(const LetObj* op) final { + const auto* ptr = expr_to_match_.as(); if (ptr == nullptr) { match_success_ = false; } else { @@ -93,8 +93,8 @@ class PatternMatcher : public StmtExprVisitor { return std::nullopt; } - ffi::Optional Visit_(const CallNode* op) final { - const auto* ptr = expr_to_match_.as(); + ffi::Optional Visit_(const CallObj* op) final { + const auto* ptr = expr_to_match_.as(); if (ptr == nullptr) { match_success_ = false; } else { @@ -128,26 +128,26 @@ class PatternMatcher : public StmtExprVisitor { return std::nullopt; \ } - TVM_DECLARE_PATTERN_MATCHER_BIN_OP(AddNode); - TVM_DECLARE_PATTERN_MATCHER_BIN_OP(SubNode); - TVM_DECLARE_PATTERN_MATCHER_BIN_OP(MulNode); - TVM_DECLARE_PATTERN_MATCHER_BIN_OP(DivNode); - TVM_DECLARE_PATTERN_MATCHER_BIN_OP(ModNode); - TVM_DECLARE_PATTERN_MATCHER_BIN_OP(FloorDivNode); - TVM_DECLARE_PATTERN_MATCHER_BIN_OP(FloorModNode); - TVM_DECLARE_PATTERN_MATCHER_BIN_OP(MinNode); - TVM_DECLARE_PATTERN_MATCHER_BIN_OP(MaxNode); - TVM_DECLARE_PATTERN_MATCHER_BIN_OP(EQNode); - TVM_DECLARE_PATTERN_MATCHER_BIN_OP(NENode); - TVM_DECLARE_PATTERN_MATCHER_BIN_OP(LTNode); - TVM_DECLARE_PATTERN_MATCHER_BIN_OP(LENode); - TVM_DECLARE_PATTERN_MATCHER_BIN_OP(GTNode); - TVM_DECLARE_PATTERN_MATCHER_BIN_OP(GENode); - TVM_DECLARE_PATTERN_MATCHER_BIN_OP(AndNode); - TVM_DECLARE_PATTERN_MATCHER_BIN_OP(OrNode); - - ffi::Optional Visit_(const CastNode* op) final { - const auto* ptr = expr_to_match_.as(); + TVM_DECLARE_PATTERN_MATCHER_BIN_OP(AddObj); + TVM_DECLARE_PATTERN_MATCHER_BIN_OP(SubObj); + TVM_DECLARE_PATTERN_MATCHER_BIN_OP(MulObj); + TVM_DECLARE_PATTERN_MATCHER_BIN_OP(DivObj); + TVM_DECLARE_PATTERN_MATCHER_BIN_OP(ModObj); + TVM_DECLARE_PATTERN_MATCHER_BIN_OP(FloorDivObj); + TVM_DECLARE_PATTERN_MATCHER_BIN_OP(FloorModObj); + TVM_DECLARE_PATTERN_MATCHER_BIN_OP(MinObj); + TVM_DECLARE_PATTERN_MATCHER_BIN_OP(MaxObj); + TVM_DECLARE_PATTERN_MATCHER_BIN_OP(EQObj); + TVM_DECLARE_PATTERN_MATCHER_BIN_OP(NEObj); + TVM_DECLARE_PATTERN_MATCHER_BIN_OP(LTObj); + TVM_DECLARE_PATTERN_MATCHER_BIN_OP(LEObj); + TVM_DECLARE_PATTERN_MATCHER_BIN_OP(GTObj); + TVM_DECLARE_PATTERN_MATCHER_BIN_OP(GEObj); + TVM_DECLARE_PATTERN_MATCHER_BIN_OP(AndObj); + TVM_DECLARE_PATTERN_MATCHER_BIN_OP(OrObj); + + ffi::Optional Visit_(const CastObj* op) final { + const auto* ptr = expr_to_match_.as(); if (ptr == nullptr) { match_success_ = false; } else { @@ -163,8 +163,8 @@ class PatternMatcher : public StmtExprVisitor { return std::nullopt; } - ffi::Optional Visit_(const NotNode* op) final { - const auto* ptr = expr_to_match_.as(); + ffi::Optional Visit_(const NotObj* op) final { + const auto* ptr = expr_to_match_.as(); if (ptr == nullptr) { match_success_ = false; } else { @@ -176,8 +176,8 @@ class PatternMatcher : public StmtExprVisitor { return std::nullopt; } - ffi::Optional Visit_(const SelectNode* op) final { - const auto* ptr = expr_to_match_.as(); + ffi::Optional Visit_(const SelectObj* op) final { + const auto* ptr = expr_to_match_.as(); if (ptr == nullptr) { match_success_ = false; } else { @@ -193,8 +193,8 @@ class PatternMatcher : public StmtExprVisitor { return std::nullopt; } - ffi::Optional Visit_(const RampNode* op) final { - const auto* ptr = expr_to_match_.as(); + ffi::Optional Visit_(const RampObj* op) final { + const auto* ptr = expr_to_match_.as(); if (ptr == nullptr) { match_success_ = false; } else { @@ -210,8 +210,8 @@ class PatternMatcher : public StmtExprVisitor { return std::nullopt; } - ffi::Optional Visit_(const BroadcastNode* op) final { - const auto* ptr = expr_to_match_.as(); + ffi::Optional Visit_(const BroadcastObj* op) final { + const auto* ptr = expr_to_match_.as(); if (ptr == nullptr) { match_success_ = false; } else { @@ -225,8 +225,8 @@ class PatternMatcher : public StmtExprVisitor { return std::nullopt; } - ffi::Optional Visit_(const ShuffleNode* op) final { - const auto* ptr = expr_to_match_.as(); + ffi::Optional Visit_(const ShuffleObj* op) final { + const auto* ptr = expr_to_match_.as(); if (ptr == nullptr) { match_success_ = false; } else { @@ -248,26 +248,26 @@ class PatternMatcher : public StmtExprVisitor { return std::nullopt; } - ffi::Optional Visit_(const IntImmNode* op) final { - const auto* ptr = expr_to_match_.as(); + ffi::Optional Visit_(const IntImmObj* op) final { + const auto* ptr = expr_to_match_.as(); match_success_ = ptr != nullptr && op->value == ptr->value; return std::nullopt; } - ffi::Optional Visit_(const FloatImmNode* op) final { - const auto* ptr = expr_to_match_.as(); + ffi::Optional Visit_(const FloatImmObj* op) final { + const auto* ptr = expr_to_match_.as(); match_success_ = ptr != nullptr && op->value == ptr->value; return std::nullopt; } - ffi::Optional Visit_(const StringImmNode* op) final { - const auto* ptr = expr_to_match_.as(); + ffi::Optional Visit_(const StringImmObj* op) final { + const auto* ptr = expr_to_match_.as(); match_success_ = ptr != nullptr && op->value == ptr->value; return std::nullopt; } - ffi::Optional Visit_(const TensorLoadNode* op) final { - const auto* ptr = expr_to_match_.as(); + ffi::Optional Visit_(const TensorLoadObj* op) final { + const auto* ptr = expr_to_match_.as(); if (ptr == nullptr) { match_success_ = false; } else { @@ -312,7 +312,7 @@ class PatternMatcher : public StmtExprVisitor { bool match_success_{true}; ffi::Array pattern_; Expr expr_to_match_{ffi::UnsafeInit{}}; - std::unordered_map filled_map_; + std::unordered_map filled_map_; }; /******** Reduction SBlock Related ********/ @@ -385,8 +385,8 @@ void ErrorRFactorCrossThreadReductionNotApplicable(const ffi::Optional& self, SBlock block, const ffi::Array& stmts, int n_tensors, ffi::Array* updates, - std::unordered_map* tensor2index) { - std::unordered_map var2index; + std::unordered_map* tensor2index) { + std::unordered_map var2index; ffi::Array let_values; let_values.reserve(n_tensors); updates->resize(n_tensors); @@ -399,7 +399,7 @@ void ExtractReductionUpdates(const ffi::Optional& self, SBlock bl if (i >= static_cast(stmts.size())) { ErrorRFactorCrossThreadReductionNotApplicable(self, std::move(block), /*violated_cond=*/3); } - const auto* bind = stmts[i].as(); + const auto* bind = stmts[i].as(); if (bind == nullptr || !bind->value.as()) { ErrorRFactorCrossThreadReductionNotApplicable(self, std::move(block), /*violated_cond=*/3); } @@ -411,7 +411,7 @@ void ExtractReductionUpdates(const ffi::Optional& self, SBlock bl } // There should be no more Bind after the first n_tensors. - if (n_tensors < static_cast(stmts.size()) && stmts[n_tensors]->IsInstance()) { + if (n_tensors < static_cast(stmts.size()) && stmts[n_tensors]->IsInstance()) { ErrorRFactorCrossThreadReductionNotApplicable(self, std::move(block), /*violated_cond=*/3); } @@ -429,7 +429,7 @@ void ExtractReductionUpdates(const ffi::Optional& self, SBlock bl // - Create TensorStores according to the variables being stored. // - Construct the mapping from reduction tensors to the index. for (const Stmt& stmt : seq) { - const auto* tensor_store = stmt.as(); + const auto* tensor_store = stmt.as(); if (tensor_store == nullptr) { ErrorRFactorCrossThreadReductionNotApplicable(self, std::move(block), /*violated_cond=*/5); } @@ -437,7 +437,7 @@ void ExtractReductionUpdates(const ffi::Optional& self, SBlock bl if (!var_ref.has_value()) { ErrorRFactorCrossThreadReductionNotApplicable(self, std::move(block), /*violated_cond=*/7); } - const VarNode* var = var_ref.value().get(); + const VarObj* var = var_ref.value().get(); auto it = var2index.find(var); if (it == var2index.end()) { ErrorRFactorCrossThreadReductionNotApplicable(self, std::move(block), /*violated_cond=*/7); @@ -466,7 +466,7 @@ std::pair, ffi::Array> GetInitValuesAndUpdates // Step 1. Extract the TensorStores serving as block inits. if (block->init.has_value() && !block->init.value()->seq.empty()) { const SeqStmt& seq_init = block->init.value(); - std::unordered_set init_tensors; + std::unordered_set init_tensors; for (const Stmt& stmt : seq_init->seq) { auto init = stmt.as(); if (!init) { @@ -484,9 +484,9 @@ std::pair, ffi::Array> GetInitValuesAndUpdates // Step 2. Extract the block updates, in the form of TensorStores. int n_tensors = inits.size(); - std::unordered_map tensor2index; + std::unordered_map tensor2index; if (const auto* update = - block->body->size() == 1 ? block->body->seq[0].as() : nullptr) { + block->body->size() == 1 ? block->body->seq[0].as() : nullptr) { updates.push_back(ffi::GetRef(update)); tensor2index[update->dest.as_or_throw().get()] = 0; } else { @@ -555,7 +555,7 @@ bool ContainsOnlyDataParAndReductionBlockIter(const ffi::Array& iters) bool ReductionIterNotIndexOutputTensor(const SBlock& block) { // Step 1. Collect the reduction block iters. - std::unordered_set reduction_block_iters; + std::unordered_set reduction_block_iters; reduction_block_iters.reserve(block->iter_vars.size()); for (const IterVar& iter_var : block->iter_vars) { if (iter_var->iter_type == kCommReduce) { @@ -563,13 +563,13 @@ bool ReductionIterNotIndexOutputTensor(const SBlock& block) { } } // Step 2. Check if the reduction block iters are used to index the output tensor. - std::unordered_set tensor_written; + std::unordered_set tensor_written; tensor_written.reserve(block->writes.size()); for (const TensorRegion& write_region : block->writes) { tensor_written.insert(write_region->source.as_or_throw().get()); } - std::unordered_set tensor_allocated; + std::unordered_set tensor_allocated; tensor_allocated.reserve(block->alloc_tensors.size()); for (const TensorVar& tensor : block->alloc_tensors) { tensor_allocated.insert(tensor.get()); @@ -584,7 +584,7 @@ bool ReductionIterNotIndexOutputTensor(const SBlock& block) { return ffi::StructuralWalk(expr, walkfn).has_value(); }; - std::unordered_map match_tensor_sources; + std::unordered_map match_tensor_sources; for (const MatchTensorRegion& region : block->match_tensors) { match_tensor_sources[region->tensor.get()] = region->source->source.as_or_throw().get(); @@ -601,7 +601,7 @@ bool ReductionIterNotIndexOutputTensor(const SBlock& block) { // declare tensor-local scratch storage inside the block body; treat them // the same as block->alloc_tensors entries for the "write-without-signature" // check below. - if (const auto* call = alloc->value.as(); + if (const auto* call = alloc->value.as(); call && call->op.same_as(tirx::alloc_tensor_op())) { tensor_allocated.insert(alloc->var.get()); } diff --git a/src/s_tir/schedule/analysis/verify.cc b/src/s_tir/schedule/analysis/verify.cc index cd7c5122579f..90fadbfb5b96 100644 --- a/src/s_tir/schedule/analysis/verify.cc +++ b/src/s_tir/schedule/analysis/verify.cc @@ -30,7 +30,7 @@ class SRefTreeVerifier : public StmtExprVisitor { using StmtExprVisitor::Visit_; ffi::Optional Visit(ffi::AnyView value) override { - if (value.as()) return std::nullopt; + if (value.as()) return std::nullopt; return StmtExprVisitor::Visit(value); } @@ -44,7 +44,7 @@ class SRefTreeVerifier : public StmtExprVisitor { private: void Verify() { for (const auto& [_, base_func] : self_->mod->functions) { - if (const auto* func = base_func.as()) { + if (const auto* func = base_func.as()) { this->Visit(func->body); } } @@ -64,7 +64,7 @@ class SRefTreeVerifier : public StmtExprVisitor { TVM_FFI_ICHECK_EQ(n_block_sref_visited_, static_cast(self_->block_info.size())); } - ffi::Optional Visit_(const SBlockNode* block) final { + ffi::Optional Visit_(const SBlockObj* block) final { if (init_block_depth_) { TVM_FFI_CHECK(!self_->stmt2ref.count(block), InternalError) << "A block inside init block has its " @@ -98,7 +98,7 @@ class SRefTreeVerifier : public StmtExprVisitor { return std::nullopt; } - ffi::Optional Visit_(const ForNode* loop) final { + ffi::Optional Visit_(const ForObj* loop) final { if (init_block_depth_) { TVM_FFI_CHECK(!self_->stmt2ref.count(loop), InternalError) << "A loop inside init block has its " @@ -106,13 +106,13 @@ class SRefTreeVerifier : public StmtExprVisitor { return StmtExprVisitor::Visit_(loop); } TVM_FFI_CHECK(self_->stmt2ref.count(loop), InternalError) - << "A ForNode should appear in sref map, but it didn't\n" + << "A ForObj should appear in sref map, but it didn't\n" << ffi::GetRef(loop); ++n_sref_visited_; const StmtSRef& sref = self_->stmt2ref.at(loop); ffi::Optional stmt = std::nullopt; TVM_FFI_CHECK(sref->parent == ancestors_.back(), InternalError) - << "Parent information mismatch for ForNode:\n" + << "Parent information mismatch for ForObj:\n" << ffi::GetRef(loop) << "\nIts parent is supposed to be:\n" << ffi::GetRef(ancestors_.back()->stmt) << "\nHowever, its parent is incorrect and is:\n" @@ -124,7 +124,7 @@ class SRefTreeVerifier : public StmtExprVisitor { return std::nullopt; } - ffi::Optional Visit_(const SeqStmtNode* seq_stmt) final { + ffi::Optional Visit_(const SeqStmtObj* seq_stmt) final { // Verify seq_index if (init_block_depth_) { return StmtExprVisitor::Visit_(seq_stmt); @@ -133,11 +133,11 @@ class SRefTreeVerifier : public StmtExprVisitor { for (int i = 0; i < n; ++i) { const Stmt& child = seq_stmt->seq[i]; StmtSRef sref{ffi::UnsafeInit{}}; - if (const auto* realize = child.as()) { + if (const auto* realize = child.as()) { const auto* block = realize->block.get(); TVM_FFI_ICHECK(self_->stmt2ref.count(block)); sref = self_->stmt2ref.at(block); - } else if (child->IsInstance()) { + } else if (child->IsInstance()) { TVM_FFI_ICHECK(self_->stmt2ref.count(child.get())); sref = self_->stmt2ref.at(child.get()); } else { @@ -170,9 +170,9 @@ void VerifyCachedFlags(const ScheduleState& self) { ScheduleState new_state(self->mod); for (const auto& kv : new_state->stmt2ref) { - const StmtNode* stmt = kv.first; + const StmtObj* stmt = kv.first; const StmtSRef& new_sref = kv.second; - if (stmt->IsInstance() || !self->stmt2ref.count(stmt)) { + if (stmt->IsInstance() || !self->stmt2ref.count(stmt)) { continue; } const SBlockInfo& new_block_info = new_state->block_info.at(new_sref); @@ -211,7 +211,7 @@ void VerifyCachedFlags(const ScheduleState& self) { if (has_not_found) { os << "- SBlockInfo not found:"; for (const StmtSRef& block_sref : block_info_not_found) { - const auto* block = block_sref->StmtAs(); + const auto* block = block_sref->StmtAs(); TVM_FFI_ICHECK(block); os << " " << block->name_hint; } @@ -223,7 +223,7 @@ void VerifyCachedFlags(const ScheduleState& self) { const StmtSRef& block_sref = std::get<0>(record); bool expected = std::get<1>(record); bool actual = std::get<2>(record); - const auto* block = block_sref->StmtAs(); + const auto* block = block_sref->StmtAs(); TVM_FFI_ICHECK(block); os << " (" << block->name_hint << ", expected=" << expected << ", actual=" << actual << ")"; } @@ -235,7 +235,7 @@ void VerifyCachedFlags(const ScheduleState& self) { const StmtSRef& block_sref = std::get<0>(record); bool expected = std::get<1>(record); bool actual = std::get<2>(record); - const auto* block = block_sref->StmtAs(); + const auto* block = block_sref->StmtAs(); TVM_FFI_ICHECK(block); os << " (" << block->name_hint << ", expected=" << expected << ", actual=" << actual << ")"; } @@ -247,7 +247,7 @@ void VerifyCachedFlags(const ScheduleState& self) { const StmtSRef& block_sref = std::get<0>(record); bool expected = std::get<1>(record); bool actual = std::get<2>(record); - const auto* block = block_sref->StmtAs(); + const auto* block = block_sref->StmtAs(); TVM_FFI_ICHECK(block); os << " (" << block->name_hint << ", expected=" << expected << ", actual=" << actual << ")"; } diff --git a/src/s_tir/schedule/concrete_schedule.cc b/src/s_tir/schedule/concrete_schedule.cc index 0ef220e37275..2be2363d38a6 100644 --- a/src/s_tir/schedule/concrete_schedule.cc +++ b/src/s_tir/schedule/concrete_schedule.cc @@ -162,11 +162,11 @@ class ScheduleCopier { } /*! \brief Copy the stmt2ref */ - UMap Copy(const UMap& stmt2ref) { - UMap result; + UMap Copy(const UMap& stmt2ref) { + UMap result; result.reserve(stmt2ref.size()); for (const auto& kv : stmt2ref) { - const StmtNode* stmt = kv.first; + const StmtObj* stmt = kv.first; const StmtSRef& sref = kv.second; result.emplace(stmt, Copy(sref)); } @@ -213,7 +213,7 @@ Schedule ConcreteScheduleObj::Copy() { #define TVM_TIR_SCHEDULE_BEGIN() try { /*! \brief Keep the payload on rendered errors for outer scheduling error handlers. */ #define TVM_TIR_SCHEDULE_THROW(error) \ - ffi::details::ErrorBuilder( \ + ffi::detail::ErrorBuilder( \ "ScheduleError", TVMFFIBacktrace(__FILE__, __LINE__, TVM_FFI_FUNC_SIG, 0), \ TVM_FFI_ALWAYS_LOG_BEFORE_THROW, std::nullopt, (error).extra_context()) \ .stream() @@ -298,7 +298,7 @@ SBlockRV ConcreteScheduleObj::GetSBlock(const ffi::String& name, : name_(name), mod_(mod), blocks_{} { blocks_.reserve(blocks.size()); for (const StmtSRef& block_sref : blocks) { - const SBlockNode* block = TVM_SREF_TO_SBLOCK(block_sref); + const SBlockObj* block = TVM_SREF_TO_SBLOCK(block_sref); blocks_.push_back(ffi::GetRef(block)); } } @@ -490,7 +490,7 @@ ffi::Array ConcreteScheduleObj::Split(const LoopRV& loop_rv, bool preserve_unit_iters, bool disable_predication) { // Prepare for the splitting StmtSRef loop_sref = this->GetSRef(loop_rv); - const ForNode* loop = TVM_SREF_TO_FOR(loop_sref); + const ForObj* loop = TVM_SREF_TO_FOR(loop_sref); ffi::Array factors; factors.reserve(factor_rvs.size()); int infer_index = -1; @@ -508,7 +508,7 @@ ffi::Array ConcreteScheduleObj::Split(const LoopRV& loop_rv, } else { PrimExpr factor = this->Get(factor_rvs[i].value()); if (IsConstInt(factor) && !IsPositiveConst(factor)) { - throw MakeScheduleError(state_->mod, factor.as()->value, + throw MakeScheduleError(state_->mod, factor.as()->value, i); } if (factor.ty().bits() > Range::GetExtentFromBeginEnd(loop->begin, loop->end).ty().bits()) { @@ -558,7 +558,7 @@ ffi::Array ConcreteScheduleObj::LoopPartition( // Prepare for the loop_partitioning StmtSRef loop_sref = this->GetSRef(loop_rv); - const ForNode* loop = TVM_SREF_TO_FOR(loop_sref); + const ForObj* loop = TVM_SREF_TO_FOR(loop_sref); ffi::Array factors; factors.reserve(factor_rvs.size()); int infer_index = -1; @@ -580,7 +580,7 @@ ffi::Array ConcreteScheduleObj::LoopPartition( } else { PrimExpr factor = this->Get(factor_rvs[i].value()); if (IsConstInt(factor) && !IsPositiveConst(factor)) { - throw MakeScheduleError(state_->mod, factor.as()->value, + throw MakeScheduleError(state_->mod, factor.as()->value, i); } if (factor.ty().bits() > Range::GetExtentFromBeginEnd(loop->begin, loop->end).ty().bits()) { @@ -966,7 +966,7 @@ Any ConcreteScheduleObj::CheckAndGetAnnotationValue(const ffi::Any& ann_val) { } if (auto expr = ann_val.as()) { - TVM_FFI_CHECK(!expr.value().as(), TypeError) + TVM_FFI_CHECK(!expr.value().as(), TypeError) << "ffi::String is expected, but gets StringImm"; auto res_expr = this->Get(expr.value()); // prefer to return int/float literals for annotations @@ -991,7 +991,7 @@ Any ConcreteScheduleObj::CheckAndGetAnnotationValue(const ffi::Any& ann_val) { for (auto it = dict->begin(); it != dict->end(); ++it) { const auto& key = it->first; auto value = CheckAndGetAnnotationValue(it->second); - if (const StringImmNode* imm = key.as()) { + if (const StringImmObj* imm = key.as()) { result.Set(imm->value, value); } else if (auto opt_str = key.try_cast()) { result.Set(opt_str.value(), value); diff --git a/src/s_tir/schedule/concrete_schedule.h b/src/s_tir/schedule/concrete_schedule.h index de322aa9f357..b506bcf76041 100644 --- a/src/s_tir/schedule/concrete_schedule.h +++ b/src/s_tir/schedule/concrete_schedule.h @@ -257,7 +257,7 @@ inline SBlock ConcreteScheduleObj::Get(const SBlockRV& block_rv) const { inline For ConcreteScheduleObj::Get(const LoopRV& loop_rv) const { StmtSRef sref = this->GetSRef(loop_rv); - const ForNode* loop = TVM_SREF_TO_FOR(sref); + const ForObj* loop = TVM_SREF_TO_FOR(sref); return ffi::GetRef(loop); } @@ -268,7 +268,7 @@ inline PrimExpr ConcreteScheduleObj::Get(const ExprRV& expr_rv) const { TVM_FFI_THROW(IndexError) << "Cannot find corresponding ExprRV: " << var; } const ffi::ObjectRef& obj = (*it).second; - const auto* int_imm = TVM_TYPE_AS(obj, IntImmNode); + const auto* int_imm = TVM_TYPE_AS(obj, IntImmObj); return ffi::Any(IntImm::Int32(int_imm->value)); }; PrimExpr transformed = diff --git a/src/s_tir/schedule/instruction.cc b/src/s_tir/schedule/instruction.cc index c00354a980e7..f9326b79aa56 100644 --- a/src/s_tir/schedule/instruction.cc +++ b/src/s_tir/schedule/instruction.cc @@ -30,7 +30,7 @@ using namespace tvm::tirx; namespace { -ffi::String InstructionAsPythonRepr(const InstructionNode* self) { +ffi::String InstructionAsPythonRepr(const InstructionObj* self) { ffi::Array inputs; inputs.reserve(self->inputs.size()); for (const Any& obj : self->inputs) { @@ -38,11 +38,11 @@ ffi::String InstructionAsPythonRepr(const InstructionNode* self) { inputs.push_back(ffi::String("None")); } else if (auto opt_str = obj.as()) { inputs.push_back(ffi::String('"' + (*opt_str).operator std::string() + '"')); - } else if (obj.as() || obj.as()) { + } else if (obj.as() || obj.as()) { inputs.push_back(ffi::String("_")); } else if (obj.type_index() < ffi::TypeIndex::kTVMFFISmallStr) { inputs.push_back(obj); - } else if (obj.as() || obj.as()) { + } else if (obj.as() || obj.as()) { inputs.push_back(obj); } else if (auto expr = obj.as()) { auto f_substitute = [](const Var& var) -> ffi::Expected> { @@ -53,7 +53,7 @@ ffi::String InstructionAsPythonRepr(const InstructionNode* self) { std::ostringstream os; os << new_expr; inputs.push_back(ffi::String(os.str())); - } else if (obj.as()) { + } else if (obj.as()) { inputs.push_back(obj); } else { TVM_FFI_THROW(TypeError) << "Stringifying is not supported for type: " << obj.GetTypeKey(); @@ -78,7 +78,7 @@ bool InstructionKindObj::IsPostproc() const { Instruction::Instruction(InstructionKind kind, ffi::Array inputs, ffi::Array attrs, ffi::Array outputs) { - ffi::ObjectPtr n = ffi::make_object(std::move(kind)); + ffi::ObjectPtr n = ffi::make_object(std::move(kind)); n->inputs = std::move(inputs); n->attrs = std::move(attrs); n->outputs = std::move(outputs); @@ -87,11 +87,11 @@ Instruction::Instruction(InstructionKind kind, ffi::Array inputs, ffi::Arra TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - InstructionNode::RegisterReflection(); - refl::TypeAttrDef().def(ffi::type_attr::kRepr.name, - [](Instruction inst, ffi::Function) -> ffi::String { - return InstructionAsPythonRepr(inst.get()); - }); + InstructionObj::RegisterReflection(); + refl::TypeAttrDef().def(ffi::type_attr::kRepr.name, + [](Instruction inst, ffi::Function) -> ffi::String { + return InstructionAsPythonRepr(inst.get()); + }); refl::GlobalDef().def("s_tir.schedule.Instruction", [](InstructionKind kind, ffi::Array inputs, ffi::Array attrs, diff --git a/src/s_tir/schedule/instruction_traits.h b/src/s_tir/schedule/instruction_traits.h index cd771b361c03..4d0fcbc56afa 100644 --- a/src/s_tir/schedule/instruction_traits.h +++ b/src/s_tir/schedule/instruction_traits.h @@ -339,9 +339,9 @@ ffi::Array UnpackedInstTraits::ApplyToSchedule(const Schedule& sch ffi::Function pf([](const ffi::PackedArgs& args, ffi::Any* rv) -> void { constexpr size_t kNumArgs = details::NumArgs; TVM_FFI_ICHECK_EQ(args.size(), kNumArgs); - ffi::details::unpack_call(std::make_index_sequence{}, nullptr, - TTraits::UnpackedApplyToSchedule, args.data(), - args.size(), rv); + ffi::detail::unpack_call(std::make_index_sequence{}, nullptr, + TTraits::UnpackedApplyToSchedule, args.data(), + args.size(), rv); }); ffi::Any rv; pf.CallPacked(ffi::PackedArgs(packed_args, kNumArgs), &rv); @@ -374,8 +374,8 @@ ffi::String UnpackedInstTraits::AsPython(const ffi::Array& inputs, ffi::Function pf([](const ffi::PackedArgs& args, ffi::Any* rv) -> void { constexpr size_t kNumArgs = details::NumArgs; TVM_FFI_ICHECK_EQ(args.size(), kNumArgs); - ffi::details::unpack_call(std::make_index_sequence{}, nullptr, - TTraits::UnpackedAsPython, args.data(), args.size(), rv); + ffi::detail::unpack_call(std::make_index_sequence{}, nullptr, + TTraits::UnpackedAsPython, args.data(), args.size(), rv); }); ffi::Any rv; pf.CallPacked(ffi::PackedArgs(packed_args, kNumArgs), &rv); diff --git a/src/s_tir/schedule/ir_comparator.cc b/src/s_tir/schedule/ir_comparator.cc index 47c4c334fdf6..062dfe4a7bd1 100644 --- a/src/s_tir/schedule/ir_comparator.cc +++ b/src/s_tir/schedule/ir_comparator.cc @@ -28,7 +28,7 @@ namespace tvm { namespace { // File-local helper: true if `expr` is a call to prim::vscale_op(). bool IsVScaleCall(const PrimExpr& expr) { - if (const auto* call = expr.as()) { + if (const auto* call = expr.as()) { return call->op.same_as(prim::vscale_op()); } return false; @@ -62,7 +62,7 @@ class TensorIntrinMismatchError : public ScheduleErrorContextObj { lhs_stmt_(std::move(lhs_stmt)), rhs_stmt_(std::move(rhs_stmt)), error_messages_(std::move(error_messages)) { - TVM_FFI_ICHECK(lhs_stmt_->IsInstance() || lhs_stmt_->IsInstance()); + TVM_FFI_ICHECK(lhs_stmt_->IsInstance() || lhs_stmt_->IsInstance()); } ffi::String FastErrorString() const final { @@ -95,7 +95,7 @@ class TensorIntrinMismatchError : public ScheduleErrorContextObj { bool TensorizeComparator::Dispatch(const Stmt& n, const Stmt& other) { bool equal = n.same_as(other) || ((n->type_index() == other->type_index()) && StmtComparator::Dispatch(n, other)); - if (!equal && assert_mode_ && (n->IsInstance() || n->IsInstance())) { + if (!equal && assert_mode_ && (n->IsInstance() || n->IsInstance())) { throw MakeScheduleError(lhs_mod_, n, other, std::move(error_messages_)); } @@ -142,13 +142,13 @@ bool TensorizeComparator::CompareExpr(const Expr& lhs, const Expr& rhs) { return ffi::StructuralEqual()(lhs, rhs); } -bool TensorizeComparator::Dispatch_(const CallNode* op, const PrimExpr& other) { - const auto* rhs = other.as(); +bool TensorizeComparator::Dispatch_(const CallObj* op, const PrimExpr& other) { + const auto* rhs = other.as(); if (!rhs->op.same_as(op->op)) return false; if (op->ty.as_or_throw().code() != rhs->ty.as_or_throw().code()) { if (assert_mode_) { std::ostringstream os; - os << "CallNode data type codes do not match: op->dtype.code()=" + os << "CallObj data type codes do not match: op->dtype.code()=" << op->ty.as_or_throw().code() << " vs rhs->dtype.code()=" << rhs->ty.as_or_throw().code(); EmitError(os.str()); @@ -158,7 +158,7 @@ bool TensorizeComparator::Dispatch_(const CallNode* op, const PrimExpr& other) { if (op->args.size() != rhs->args.size()) { if (assert_mode_) { std::ostringstream os; - os << "CallNode arg size mismatch: op->args.size()=" << op->args.size() + os << "CallObj arg size mismatch: op->args.size()=" << op->args.size() << " vs rhs->args.size()=" << rhs->args.size(); EmitError(os.str()); } @@ -168,7 +168,7 @@ bool TensorizeComparator::Dispatch_(const CallNode* op, const PrimExpr& other) { if (!CompareExpr(op->args[i], rhs->args[i])) { if (assert_mode_) { std::ostringstream os; - os << "CallNode args do not match at index " << i << ": op->args[i]=" << op->args[i] + os << "CallObj args do not match at index " << i << ": op->args[i]=" << op->args[i] << " vs rhs->args[i]=" << rhs->args[i]; EmitError(os.str()); } @@ -178,12 +178,12 @@ bool TensorizeComparator::Dispatch_(const CallNode* op, const PrimExpr& other) { return true; } -bool TensorizeComparator::Dispatch_(const ForNode* op, const Stmt& other) { - const auto* rhs = other.as(); +bool TensorizeComparator::Dispatch_(const ForObj* op, const Stmt& other) { + const auto* rhs = other.as(); if (!DefEqual(op->var, rhs->var)) { if (assert_mode_) { std::ostringstream os; - os << "ForNode loop vars do not match: op->var=" << op->var << " vs rhs->var=" << rhs->var; + os << "ForObj loop vars do not match: op->var=" << op->var << " vs rhs->var=" << rhs->var; EmitError(os.str()); } return false; @@ -191,7 +191,7 @@ bool TensorizeComparator::Dispatch_(const ForNode* op, const Stmt& other) { if (!Dispatch(op->begin, rhs->begin)) { if (assert_mode_) { std::ostringstream os; - os << "ForNode begin values do not match: op->begin=" << op->begin + os << "ForObj begin values do not match: op->begin=" << op->begin << " vs rhs->begin=" << rhs->begin; EmitError(os.str()); } @@ -200,21 +200,21 @@ bool TensorizeComparator::Dispatch_(const ForNode* op, const Stmt& other) { if (!Dispatch(op->end, rhs->end)) { if (assert_mode_) { std::ostringstream os; - os << "ForNode end values do not match: op->end=" << op->end << " vs rhs->end=" << rhs->end; + os << "ForObj end values do not match: op->end=" << op->end << " vs rhs->end=" << rhs->end; EmitError(os.str()); } return false; } if (!Dispatch(op->step, rhs->step)) { if (assert_mode_) { - EmitError("ForNode step values do not match"); + EmitError("ForObj step values do not match"); } return false; } if (op->kind != rhs->kind) { if (assert_mode_) { std::ostringstream os; - os << "ForNode kinds do not match: op->kind=" << op->kind << " vs rhs->kind=" << rhs->kind; + os << "ForObj kinds do not match: op->kind=" << op->kind << " vs rhs->kind=" << rhs->kind; EmitError(os.str()); } return false; @@ -222,7 +222,7 @@ bool TensorizeComparator::Dispatch_(const ForNode* op, const Stmt& other) { if (!CompareAnnotationMap(op->annotations, rhs->annotations)) { if (assert_mode_) { std::ostringstream os; - os << "ForNode annotation maps do not match: op->annotations=" << op->annotations + os << "ForObj annotation maps do not match: op->annotations=" << op->annotations << " vs rhs->annotations=" << rhs->annotations; EmitError(os.str()); } @@ -231,20 +231,20 @@ bool TensorizeComparator::Dispatch_(const ForNode* op, const Stmt& other) { return Dispatch(op->body, rhs->body); } -bool TensorizeComparator::Dispatch_(const SeqStmtNode* op, const Stmt& other) { - const auto* rhs = other.as(); +bool TensorizeComparator::Dispatch_(const SeqStmtObj* op, const Stmt& other) { + const auto* rhs = other.as(); return CompareArray(op->seq, rhs->seq, static_cast( &TensorizeComparator::Dispatch)); } -bool TensorizeComparator::Dispatch_(const TensorStoreNode* op, const Stmt& other) { - const auto* rhs = other.as(); +bool TensorizeComparator::Dispatch_(const TensorStoreObj* op, const Stmt& other) { + const auto* rhs = other.as(); return CompareTensorAccess(op, rhs) && Dispatch(op->value, rhs->value); } -bool TensorizeComparator::Dispatch_(const SBlockRealizeNode* op, const Stmt& other) { - const auto* rhs = other.as(); +bool TensorizeComparator::Dispatch_(const SBlockRealizeObj* op, const Stmt& other) { + const auto* rhs = other.as(); if (!is_scope_block) { if (!CompareArray(op->iter_values, rhs->iter_values, static_cast( @@ -261,10 +261,14 @@ bool TensorizeComparator::Dispatch_(const SBlockRealizeNode* op, const Stmt& oth return Dispatch(op->predicate, rhs->predicate) && Dispatch(op->block, rhs->block); } -bool TensorizeComparator::Dispatch_(const SBlockNode* op, const Stmt& other) { - const auto* rhs = other.as(); +bool TensorizeComparator::Dispatch_(const SBlockObj* op, const Stmt& other) { + const auto* rhs = other.as(); for (const IterVar& iter : op->iter_vars) { - lhs_analyzer_->Bind(iter->var, iter->dom.value()); + { + auto context = lhs_analyzer_->GetConstraintContext(); + context->Append(sym::RangeConstraint(iter->var, iter->dom.value())); + lhs_analyzer_->SyncConstraintContext(context); + } } // Check block equality. // All iter vars and tensor regions including the order should match. @@ -318,41 +322,41 @@ bool TensorizeComparator::Dispatch_(const SBlockNode* op, const Stmt& other) { return Dispatch(op->a, rhs->a) && Dispatch(op->b, rhs->b); \ } -TVM_DECLARE_TENSORIZE_COMPARATOR_BINOP(AddNode); -TVM_DECLARE_TENSORIZE_COMPARATOR_BINOP(SubNode); -TVM_DECLARE_TENSORIZE_COMPARATOR_BINOP(MulNode); -TVM_DECLARE_TENSORIZE_COMPARATOR_BINOP(DivNode); -TVM_DECLARE_TENSORIZE_COMPARATOR_BINOP(ModNode); -TVM_DECLARE_TENSORIZE_COMPARATOR_BINOP(EQNode); -TVM_DECLARE_TENSORIZE_COMPARATOR_BINOP(NENode); -TVM_DECLARE_TENSORIZE_COMPARATOR_BINOP(LTNode); -TVM_DECLARE_TENSORIZE_COMPARATOR_BINOP(LENode); -TVM_DECLARE_TENSORIZE_COMPARATOR_BINOP(GTNode); -TVM_DECLARE_TENSORIZE_COMPARATOR_BINOP(GENode); -TVM_DECLARE_TENSORIZE_COMPARATOR_BINOP(AndNode); -TVM_DECLARE_TENSORIZE_COMPARATOR_BINOP(OrNode); -TVM_DECLARE_TENSORIZE_COMPARATOR_BINOP(MinNode); -TVM_DECLARE_TENSORIZE_COMPARATOR_BINOP(MaxNode); -TVM_DECLARE_TENSORIZE_COMPARATOR_BINOP(FloorDivNode); -TVM_DECLARE_TENSORIZE_COMPARATOR_BINOP(FloorModNode); - -TVM_DECLARE_TENSORIZE_COMPARATOR_BINOP(prim::LShiftNode); -TVM_DECLARE_TENSORIZE_COMPARATOR_BINOP(prim::RShiftNode); -TVM_DECLARE_TENSORIZE_COMPARATOR_BINOP(prim::BitwiseAndNode); -TVM_DECLARE_TENSORIZE_COMPARATOR_BINOP(prim::BitwiseOrNode); -TVM_DECLARE_TENSORIZE_COMPARATOR_BINOP(prim::BitwiseXorNode); - -bool TensorizeComparator::Dispatch_(const prim::BitwiseNotNode* op, const PrimExpr& other) { - const auto* rhs = other.as(); +TVM_DECLARE_TENSORIZE_COMPARATOR_BINOP(AddObj); +TVM_DECLARE_TENSORIZE_COMPARATOR_BINOP(SubObj); +TVM_DECLARE_TENSORIZE_COMPARATOR_BINOP(MulObj); +TVM_DECLARE_TENSORIZE_COMPARATOR_BINOP(DivObj); +TVM_DECLARE_TENSORIZE_COMPARATOR_BINOP(ModObj); +TVM_DECLARE_TENSORIZE_COMPARATOR_BINOP(EQObj); +TVM_DECLARE_TENSORIZE_COMPARATOR_BINOP(NEObj); +TVM_DECLARE_TENSORIZE_COMPARATOR_BINOP(LTObj); +TVM_DECLARE_TENSORIZE_COMPARATOR_BINOP(LEObj); +TVM_DECLARE_TENSORIZE_COMPARATOR_BINOP(GTObj); +TVM_DECLARE_TENSORIZE_COMPARATOR_BINOP(GEObj); +TVM_DECLARE_TENSORIZE_COMPARATOR_BINOP(AndObj); +TVM_DECLARE_TENSORIZE_COMPARATOR_BINOP(OrObj); +TVM_DECLARE_TENSORIZE_COMPARATOR_BINOP(MinObj); +TVM_DECLARE_TENSORIZE_COMPARATOR_BINOP(MaxObj); +TVM_DECLARE_TENSORIZE_COMPARATOR_BINOP(FloorDivObj); +TVM_DECLARE_TENSORIZE_COMPARATOR_BINOP(FloorModObj); + +TVM_DECLARE_TENSORIZE_COMPARATOR_BINOP(prim::LShiftObj); +TVM_DECLARE_TENSORIZE_COMPARATOR_BINOP(prim::RShiftObj); +TVM_DECLARE_TENSORIZE_COMPARATOR_BINOP(prim::BitwiseAndObj); +TVM_DECLARE_TENSORIZE_COMPARATOR_BINOP(prim::BitwiseOrObj); +TVM_DECLARE_TENSORIZE_COMPARATOR_BINOP(prim::BitwiseXorObj); + +bool TensorizeComparator::Dispatch_(const prim::BitwiseNotObj* op, const PrimExpr& other) { + const auto* rhs = other.as(); return Dispatch(op->a, rhs->a); } -bool TensorizeComparator::Dispatch_(const IntImmNode* op, const PrimExpr& other) { - const auto* rhs = other.as(); +bool TensorizeComparator::Dispatch_(const IntImmObj* op, const PrimExpr& other) { + const auto* rhs = other.as(); if (op->value != rhs->value) { if (assert_mode_) { std::ostringstream os; - os << "IntImmNode values do not match: op->value=" << op->value + os << "IntImmObj values do not match: op->value=" << op->value << " vs rhs->value=" << rhs->value; EmitError(os.str()); } @@ -361,12 +365,12 @@ bool TensorizeComparator::Dispatch_(const IntImmNode* op, const PrimExpr& other) return true; } -bool TensorizeComparator::Dispatch_(const FloatImmNode* op, const PrimExpr& other) { - const auto* rhs = other.as(); +bool TensorizeComparator::Dispatch_(const FloatImmObj* op, const PrimExpr& other) { + const auto* rhs = other.as(); if (op->value != rhs->value) { if (assert_mode_) { std::ostringstream os; - os << "FloatImmNode values do not match: op->value=" << op->value + os << "FloatImmObj values do not match: op->value=" << op->value << " vs rhs->value=" << rhs->value; EmitError(os.str()); } @@ -375,12 +379,12 @@ bool TensorizeComparator::Dispatch_(const FloatImmNode* op, const PrimExpr& othe return true; } -bool TensorizeComparator::Dispatch_(const CastNode* op, const PrimExpr& other) { - const auto* rhs = other.as(); +bool TensorizeComparator::Dispatch_(const CastObj* op, const PrimExpr& other) { + const auto* rhs = other.as(); return Dispatch(op->value, rhs->value); } -bool TensorizeComparator::Dispatch_(const VarNode* op, const PrimExpr& other) { +bool TensorizeComparator::Dispatch_(const VarObj* op, const PrimExpr& other) { auto rhs_ref = other.as(); if (!rhs_ref.has_value()) return false; const auto* rhs = rhs_ref.value().get(); @@ -391,7 +395,7 @@ bool TensorizeComparator::Dispatch_(const VarNode* op, const PrimExpr& other) { if (lhs_ty.code() != rhs_ty.code()) { if (assert_mode_) { std::ostringstream os; - os << "VarNode data type codes do not match: op->dtype.code()=" << lhs_ty.code() + os << "VarObj data type codes do not match: op->dtype.code()=" << lhs_ty.code() << " vs rhs->dtype.code()=" << rhs_ty.code(); EmitError(os.str()); } @@ -401,13 +405,13 @@ bool TensorizeComparator::Dispatch_(const VarNode* op, const PrimExpr& other) { return it != equal_map_.end() && it->second.same_as(other); } -bool TensorizeComparator::Dispatch_(const TensorLoadNode* op, const PrimExpr& other) { - const auto* rhs = other.as(); +bool TensorizeComparator::Dispatch_(const TensorLoadObj* op, const PrimExpr& other) { + const auto* rhs = other.as(); return CompareTensorAccess(op, rhs); } -bool TensorizeComparator::Dispatch_(const SelectNode* op, const PrimExpr& other) { - const auto* rhs = other.as(); +bool TensorizeComparator::Dispatch_(const SelectObj* op, const PrimExpr& other) { + const auto* rhs = other.as(); return Dispatch(op->condition, rhs->condition) && Dispatch(op->true_value, rhs->true_value) && Dispatch(op->false_value, rhs->false_value); } @@ -428,7 +432,12 @@ bool TensorizeComparator::DefEqual(const Var& lhs, const Var& rhs) { equal_map_.insert_or_assign(lhs, rhs); // Cast if necessary. This allows the workload and the tensor intrin to have different dtypes in // the indices. - analyzer_->Bind(lhs, cast(rhs.as_or_throw(), lhs_prim_type.value())); + { + auto context = analyzer_->GetConstraintContext(); + context->Append( + sym::BindingConstraint(lhs, cast(rhs.as_or_throw(), lhs_prim_type.value()))); + analyzer_->SyncConstraintContext(context); + } return true; } @@ -644,11 +653,11 @@ bool TensorizeComparator::CompareTensorRegion(const TensorRegion& lhs, const Ten return true; } -// Comparator for TensorStoreNode and TensorLoadNode -inline TensorVar GetTensorAccessTensor(const TensorStoreNode* op) { +// Comparator for TensorStoreObj and TensorLoadObj +inline TensorVar GetTensorAccessTensor(const TensorStoreObj* op) { return op->dest.as_or_throw(); } -inline TensorVar GetTensorAccessTensor(const TensorLoadNode* op) { +inline TensorVar GetTensorAccessTensor(const TensorLoadObj* op) { return op->source.as_or_throw(); } @@ -728,8 +737,8 @@ bool AutoTensorizeComparator::DispatchDefault_(const ffi::Object* op, const Stmt return false; } -bool AutoTensorizeComparator::Dispatch_(const SBlockNode* op, const Stmt& other) { - const auto* rhs = other.as(); +bool AutoTensorizeComparator::Dispatch_(const SBlockObj* op, const Stmt& other) { + const auto* rhs = other.as(); // Check block equality. // All iter vars and tensor regions including the order should match. // When checking iter vars, DefEqual is used to remap variables. @@ -748,9 +757,13 @@ bool AutoTensorizeComparator::Dispatch_(const SBlockNode* op, const Stmt& other) inner_iter_dom_map_.Set(block_iter->var, sym::IntSet::FromRange(block_iter->dom.value())); } } else { - auto collect_iter = [&](const SBlockNode* op, std::vector& iters) -> bool { + auto collect_iter = [&](const SBlockObj* op, std::vector& iters) -> bool { for (const auto& iter : op->iter_vars) { - analyzer_->Bind(iter->var, iter->dom.value()); + { + auto context = analyzer_->GetConstraintContext(); + context->Append(sym::RangeConstraint(iter->var, iter->dom.value())); + analyzer_->SyncConstraintContext(context); + } if (iter->iter_type == IterVarType::kDataPar || iter->iter_type == IterVarType::kCommReduce) { iters.push_back(iter); @@ -798,13 +811,13 @@ bool AutoTensorizeComparator::CompareTensor(const TensorVar& lhs, const TensorVa return equal; } -bool AutoTensorizeComparator::Dispatch_(const TensorStoreNode* op, const Stmt& other) { - const auto* rhs = other.as(); +bool AutoTensorizeComparator::Dispatch_(const TensorStoreObj* op, const Stmt& other) { + const auto* rhs = other.as(); return CompareTensorAccess(op, rhs) && Dispatch(op->value, rhs->value); } -bool AutoTensorizeComparator::Dispatch_(const TensorLoadNode* op, const PrimExpr& other) { - const auto* rhs = other.as(); +bool AutoTensorizeComparator::Dispatch_(const TensorLoadObj* op, const PrimExpr& other) { + const auto* rhs = other.as(); return CompareTensorAccess(op, rhs); } @@ -826,7 +839,7 @@ bool AutoTensorizeComparator::CompareTensorAccess(const T* lhs, const T* rhs) { auto is_scalar_access = [](const ffi::Array& indices, PrimExpr index) { // Check if the indexing is of the form C[0] if (indices.size() > 1) return false; - auto int_imm = index.template as(); + auto int_imm = index.template as(); if (int_imm && int_imm->value == 0) return true; return false; }; diff --git a/src/s_tir/schedule/ir_comparator.h b/src/s_tir/schedule/ir_comparator.h index 1be7c6f59032..d19e1e160883 100644 --- a/src/s_tir/schedule/ir_comparator.h +++ b/src/s_tir/schedule/ir_comparator.h @@ -50,42 +50,42 @@ class TensorizeComparator : public ExprComparator, public StmtComparator { bool Dispatch(const Expr& n, const PrimExpr& other) override; bool Dispatch(const Stmt& n, const Stmt& other) override; - bool Dispatch_(const CallNode* op, const PrimExpr& other) override; - bool Dispatch_(const ForNode* op, const Stmt& other) override; - bool Dispatch_(const SeqStmtNode* op, const Stmt& other) override; - bool Dispatch_(const TensorStoreNode* op, const Stmt& other) override; - bool Dispatch_(const SBlockRealizeNode* op, const Stmt& other) override; - bool Dispatch_(const SBlockNode* op, const Stmt& other) override; - - bool Dispatch_(const prim::LShiftNode* op, const PrimExpr& other) override; - bool Dispatch_(const prim::RShiftNode* op, const PrimExpr& other) override; - bool Dispatch_(const prim::BitwiseAndNode* op, const PrimExpr& other) override; - bool Dispatch_(const prim::BitwiseOrNode* op, const PrimExpr& other) override; - bool Dispatch_(const prim::BitwiseXorNode* op, const PrimExpr& other) override; - bool Dispatch_(const prim::BitwiseNotNode* op, const PrimExpr& other) override; - bool Dispatch_(const AddNode* op, const PrimExpr& other) override; - bool Dispatch_(const SubNode* op, const PrimExpr& other) override; - bool Dispatch_(const MulNode* op, const PrimExpr& other) override; - bool Dispatch_(const DivNode* op, const PrimExpr& other) override; - bool Dispatch_(const ModNode* op, const PrimExpr& other) override; - bool Dispatch_(const EQNode* op, const PrimExpr& other) override; - bool Dispatch_(const NENode* op, const PrimExpr& other) override; - bool Dispatch_(const LTNode* op, const PrimExpr& other) override; - bool Dispatch_(const LENode* op, const PrimExpr& other) override; - bool Dispatch_(const GTNode* op, const PrimExpr& other) override; - bool Dispatch_(const GENode* op, const PrimExpr& other) override; - bool Dispatch_(const AndNode* op, const PrimExpr& other) override; - bool Dispatch_(const OrNode* op, const PrimExpr& other) override; - bool Dispatch_(const MinNode* op, const PrimExpr& other) override; - bool Dispatch_(const MaxNode* op, const PrimExpr& other) override; - bool Dispatch_(const FloorDivNode* op, const PrimExpr& other) override; - bool Dispatch_(const FloorModNode* op, const PrimExpr& other) override; - bool Dispatch_(const IntImmNode* op, const PrimExpr& other) override; - bool Dispatch_(const FloatImmNode* op, const PrimExpr& other) override; - bool Dispatch_(const CastNode* op, const PrimExpr& other) override; - bool Dispatch_(const VarNode* op, const PrimExpr& other) override; - bool Dispatch_(const TensorLoadNode* op, const PrimExpr& other) override; - bool Dispatch_(const SelectNode* op, const PrimExpr& other) override; + bool Dispatch_(const CallObj* op, const PrimExpr& other) override; + bool Dispatch_(const ForObj* op, const Stmt& other) override; + bool Dispatch_(const SeqStmtObj* op, const Stmt& other) override; + bool Dispatch_(const TensorStoreObj* op, const Stmt& other) override; + bool Dispatch_(const SBlockRealizeObj* op, const Stmt& other) override; + bool Dispatch_(const SBlockObj* op, const Stmt& other) override; + + bool Dispatch_(const prim::LShiftObj* op, const PrimExpr& other) override; + bool Dispatch_(const prim::RShiftObj* op, const PrimExpr& other) override; + bool Dispatch_(const prim::BitwiseAndObj* op, const PrimExpr& other) override; + bool Dispatch_(const prim::BitwiseOrObj* op, const PrimExpr& other) override; + bool Dispatch_(const prim::BitwiseXorObj* op, const PrimExpr& other) override; + bool Dispatch_(const prim::BitwiseNotObj* op, const PrimExpr& other) override; + bool Dispatch_(const AddObj* op, const PrimExpr& other) override; + bool Dispatch_(const SubObj* op, const PrimExpr& other) override; + bool Dispatch_(const MulObj* op, const PrimExpr& other) override; + bool Dispatch_(const DivObj* op, const PrimExpr& other) override; + bool Dispatch_(const ModObj* op, const PrimExpr& other) override; + bool Dispatch_(const EQObj* op, const PrimExpr& other) override; + bool Dispatch_(const NEObj* op, const PrimExpr& other) override; + bool Dispatch_(const LTObj* op, const PrimExpr& other) override; + bool Dispatch_(const LEObj* op, const PrimExpr& other) override; + bool Dispatch_(const GTObj* op, const PrimExpr& other) override; + bool Dispatch_(const GEObj* op, const PrimExpr& other) override; + bool Dispatch_(const AndObj* op, const PrimExpr& other) override; + bool Dispatch_(const OrObj* op, const PrimExpr& other) override; + bool Dispatch_(const MinObj* op, const PrimExpr& other) override; + bool Dispatch_(const MaxObj* op, const PrimExpr& other) override; + bool Dispatch_(const FloorDivObj* op, const PrimExpr& other) override; + bool Dispatch_(const FloorModObj* op, const PrimExpr& other) override; + bool Dispatch_(const IntImmObj* op, const PrimExpr& other) override; + bool Dispatch_(const FloatImmObj* op, const PrimExpr& other) override; + bool Dispatch_(const CastObj* op, const PrimExpr& other) override; + bool Dispatch_(const VarObj* op, const PrimExpr& other) override; + bool Dispatch_(const TensorLoadObj* op, const PrimExpr& other) override; + bool Dispatch_(const SelectObj* op, const PrimExpr& other) override; /*! \brief Map from RHS tensor to LHS tensor */ std::unordered_map rhs_tensor_map_; @@ -148,10 +148,10 @@ class AutoTensorizeComparator : public TensorizeComparator { bool DispatchDefault_(const ffi::Object* op, const PrimExpr& other) override; bool DispatchDefault_(const ffi::Object* op, const Stmt& other) override; - bool Dispatch_(const SBlockNode* op, const Stmt& other) override; - bool Dispatch_(const TensorStoreNode* op, const Stmt& other) override; + bool Dispatch_(const SBlockObj* op, const Stmt& other) override; + bool Dispatch_(const TensorStoreObj* op, const Stmt& other) override; - bool Dispatch_(const TensorLoadNode* op, const PrimExpr& other) override; + bool Dispatch_(const TensorLoadObj* op, const PrimExpr& other) override; bool CompareTensor(const TensorVar& lhs, const TensorVar& rhs) override; template diff --git a/src/s_tir/schedule/primitive/annotate.cc b/src/s_tir/schedule/primitive/annotate.cc index 221b4463b1a8..2b8ea989fa56 100644 --- a/src/s_tir/schedule/primitive/annotate.cc +++ b/src/s_tir/schedule/primitive/annotate.cc @@ -30,11 +30,11 @@ void Annotate(ScheduleState self, const StmtSRef& sref, const ffi::String& ann_k const Any& ann_val) { // Extract annotation const ffi::Map* annotations = nullptr; - if (const auto* loop = sref->StmtAs()) { + if (const auto* loop = sref->StmtAs()) { TVM_FFI_CHECK(ann_key != tvm::s_tir::attr::kThreadBinding.name, ValueError) << "Use loop scheduling primitives to change thread binding"; annotations = &loop->annotations; - } else if (const auto* block = sref->StmtAs()) { + } else if (const auto* block = sref->StmtAs()) { annotations = &block->annotations; } else { TVM_FFI_THROW(TypeError) << "Unknown type of sref: " << sref->stmt->GetTypeKey(); @@ -47,12 +47,12 @@ void Annotate(ScheduleState self, const StmtSRef& sref, const ffi::String& ann_k ffi::Map new_ann(*annotations); new_ann.Set(ann_key, ann_val); // Create the new stmt - if (const auto* loop = sref->StmtAs()) { - ffi::ObjectPtr n = ffi::make_object(*loop); + if (const auto* loop = sref->StmtAs()) { + ffi::ObjectPtr n = ffi::make_object(*loop); n->annotations = std::move(new_ann); self->Replace(sref, For(n), {}); - } else if (const auto* block = sref->StmtAs()) { - ffi::ObjectPtr n = ffi::make_object(*block); + } else if (const auto* block = sref->StmtAs()) { + ffi::ObjectPtr n = ffi::make_object(*block); n->annotations = std::move(new_ann); SBlock p(n); self->Replace(sref, p, {{ffi::GetRef(block), p}}); @@ -65,11 +65,11 @@ void Annotate(ScheduleState self, const StmtSRef& sref, const ffi::String& ann_k void Unannotate(ScheduleState self, const StmtSRef& sref, const ffi::String& ann_key) { // Extract annotation const ffi::Map* annotations = nullptr; - if (const auto* loop = sref->StmtAs()) { + if (const auto* loop = sref->StmtAs()) { TVM_FFI_CHECK(ann_key != tvm::s_tir::attr::kThreadBinding.name, ValueError) << "Use loop scheduling primitives to change thread binding"; annotations = &loop->annotations; - } else if (const auto* block = sref->StmtAs()) { + } else if (const auto* block = sref->StmtAs()) { annotations = &block->annotations; } else { TVM_FFI_THROW(TypeError) << "Unknown type of sref: " << sref->stmt->GetTypeKey(); @@ -80,12 +80,12 @@ void Unannotate(ScheduleState self, const StmtSRef& sref, const ffi::String& ann ffi::Map new_ann(*annotations); new_ann.erase(ann_key); // Create the new stmt - if (const auto* loop = sref->StmtAs()) { - ffi::ObjectPtr n = ffi::make_object(*loop); + if (const auto* loop = sref->StmtAs()) { + ffi::ObjectPtr n = ffi::make_object(*loop); n->annotations = std::move(new_ann); self->Replace(sref, For(n), {}); - } else if (const auto* block = sref->StmtAs()) { - ffi::ObjectPtr n = ffi::make_object(*block); + } else if (const auto* block = sref->StmtAs()) { + ffi::ObjectPtr n = ffi::make_object(*block); n->annotations = std::move(new_ann); SBlock p(n); self->Replace(sref, p, {{ffi::GetRef(block), p}}); diff --git a/src/s_tir/schedule/primitive/annotate_tensor_access.cc b/src/s_tir/schedule/primitive/annotate_tensor_access.cc index 663456379303..d3171ffb828e 100644 --- a/src/s_tir/schedule/primitive/annotate_tensor_access.cc +++ b/src/s_tir/schedule/primitive/annotate_tensor_access.cc @@ -37,7 +37,7 @@ class AnnotateRegionRewriter : public StmtExprMutator { new_region_(new_region), tensor_index_type_(tensor_index_type) {} - UnchangedOr Mutate_(const SBlockNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const SBlockObj* op, InplaceMode inplace_mode) final { SBlock block = StmtExprMutator::Mutate_(op, inplace_mode) .ValueOrUnchanged(ffi::GetRef(op)) .as_or_throw(); @@ -49,7 +49,7 @@ class AnnotateRegionRewriter : public StmtExprMutator { << "Tensor index out of range"; regions.Set(tensor_index_, new_region_); - SBlockNode* n = block.CopyOnWrite(); + SBlockObj* n = block.CopyOnWrite(); if (tensor_index_type_ == TensorIndexType::kWrite) { n->writes = std::move(regions); } else { @@ -92,7 +92,7 @@ class AnnotateRegionRewriter : public StmtExprMutator { void AnnotateTensorAccess(ScheduleState self, const StmtSRef& block_sref, int tensor_index, TensorIndexType tensor_index_type, const IndexMap& index_map) { - const SBlockNode* block = TVM_SREF_TO_SBLOCK(block_sref); + const SBlockObj* block = TVM_SREF_TO_SBLOCK(block_sref); TensorVar tensor = GetNthAccessTensor(self, ffi::GetRef(block), tensor_index, tensor_index_type); diff --git a/src/s_tir/schedule/primitive/block_annotate.cc b/src/s_tir/schedule/primitive/block_annotate.cc index 0caaea5065eb..99a34383382d 100644 --- a/src/s_tir/schedule/primitive/block_annotate.cc +++ b/src/s_tir/schedule/primitive/block_annotate.cc @@ -232,7 +232,7 @@ class StorageScopeMutator : public ReplaceTensorMutator { void StorageAlign(ScheduleState self, const StmtSRef& block_sref, int tensor_index, int axis, int factor, int offset) { - const SBlockNode* block_ptr = TVM_SREF_TO_SBLOCK(block_sref); + const SBlockObj* block_ptr = TVM_SREF_TO_SBLOCK(block_sref); TensorVar tensor = GetNthAccessTensor(self, ffi::GetRef(block_ptr), tensor_index, TensorIndexType::kWrite); StorageAlignInvalidFactorError::Check(self->mod, factor); @@ -267,7 +267,7 @@ void StorageAlign(ScheduleState self, const StmtSRef& block_sref, int tensor_ind void SetScope(ScheduleState self, const StmtSRef& block_sref, int tensor_index, const ffi::String& storage_scope) { - const SBlockNode* block = TVM_SREF_TO_SBLOCK(block_sref); + const SBlockObj* block = TVM_SREF_TO_SBLOCK(block_sref); TensorVar tensor = GetNthAccessTensor(self, ffi::GetRef(block), tensor_index, TensorIndexType::kWrite); @@ -282,7 +282,7 @@ void SetScope(ScheduleState self, const StmtSRef& block_sref, int tensor_index, // Step 3. Get the allocation site of the target tensor. StmtSRef alloc_site_sref = NonAllocatedTensorError::CheckAndGetTensorAllocationSite(self->mod, block_sref, tensor); - const SBlockNode* alloc_site = TVM_SREF_TO_SBLOCK(alloc_site_sref); + const SBlockObj* alloc_site = TVM_SREF_TO_SBLOCK(alloc_site_sref); // Step 4. Recursively replace the old tensor to a new tensor, where the new tensor has the given // storage scope. In the meanwhile, collect the block sref reuse information. @@ -337,7 +337,7 @@ class DTypeMutator : public ReplaceTensorMutator { } } - UnchangedOr Mutate_(const TensorStoreNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const TensorStoreObj* op, InplaceMode inplace_mode) final { TensorVar original_tensor = op->dest.as_or_throw(); TensorStore node = StmtExprMutator::Mutate_(op, inplace_mode) .ValueOrUnchanged(ffi::GetRef(op)) @@ -349,7 +349,7 @@ class DTypeMutator : public ReplaceTensorMutator { return node; } - UnchangedOr Mutate_(const TensorLoadNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const TensorLoadObj* op, InplaceMode inplace_mode) final { Expr original_tensor = op->source; TensorLoad node = StmtExprMutator::Mutate_(op, inplace_mode) .ValueOrUnchanged(ffi::GetRef(op)) @@ -365,7 +365,7 @@ class DTypeMutator : public ReplaceTensorMutator { void UnsafeSetDType(ScheduleState self, const StmtSRef& block_sref, int tensor_index, const ffi::String& dtype) { - const SBlockNode* block = TVM_SREF_TO_SBLOCK(block_sref); + const SBlockObj* block = TVM_SREF_TO_SBLOCK(block_sref); TensorVar tensor = GetNthAccessTensor(self, ffi::GetRef(block), tensor_index, TensorIndexType::kWrite); PrimType target_dtype(ffi::StringToDLDataType(dtype)); @@ -378,7 +378,7 @@ void UnsafeSetDType(ScheduleState self, const StmtSRef& block_sref, int tensor_i // Step 2. Get the allocation site of the target tensor. StmtSRef alloc_site_sref = NonAllocatedTensorError::CheckAndGetTensorAllocationSite(self->mod, block_sref, tensor); - const SBlockNode* alloc_site = TVM_SREF_TO_SBLOCK(alloc_site_sref); + const SBlockObj* alloc_site = TVM_SREF_TO_SBLOCK(alloc_site_sref); // Step 3. Recursively replace old tensor to a new tensor, where the new tensor has the given // dtype, and insert data type conversions. diff --git a/src/s_tir/schedule/primitive/blockize_tensorize.cc b/src/s_tir/schedule/primitive/blockize_tensorize.cc index 2c8d3cd9f298..6f7a9c2b909e 100644 --- a/src/s_tir/schedule/primitive/blockize_tensorize.cc +++ b/src/s_tir/schedule/primitive/blockize_tensorize.cc @@ -99,12 +99,12 @@ ffi::Array> TrivialSubspaceDivision(const ffi::Array& inner_iters) { if (!IsOne(predicate)) return {}; ffi::Array> res; - std::unordered_set outer_loop_vars; - std::unordered_set inner_loop_vars; + std::unordered_set outer_loop_vars; + std::unordered_set inner_loop_vars; auto make_uses_var = [](const ffi::Array& vars) -> std::function { - std::unordered_set var_set; + std::unordered_set var_set; var_set.reserve(vars.size()); for (const Var& var : vars) { var_set.insert(var.get()); @@ -167,7 +167,7 @@ ffi::Array> TrivialSubspaceDivision(const ffi::Array> SubspaceDivide(const SBlockRealize& realize, const StmtSRef& block_sref, // const StmtSRef& loop_sref, // - std::vector* loops, + std::vector* loops, sym::AnalyzerObj* analyzer, bool preserve_unit_iters, bool loop_sref_as_outer = false) { @@ -176,10 +176,10 @@ ffi::Array> SubspaceDivide(const SBlockRealize& realiz ffi::Array primitive_inner_vars; ffi::Map primitive_loop_var_domain; bool inner = true; - for (StmtSRefObj* sref = block_sref->parent; // - sref && sref->stmt->IsInstance(); // + for (StmtSRefObj* sref = block_sref->parent; // + sref && sref->stmt->IsInstance(); // sref = sref->parent) { - const ForNode* loop = static_cast(sref->stmt); + const ForObj* loop = static_cast(sref->stmt); if (inner) { loops->push_back(loop); inner_vars.push_back(loop->var); @@ -225,7 +225,7 @@ ffi::Map DeriveBlockBinding( ffi::Array* inner_bindings, // bool preserve_unit_iters, bool reuse_outer = false) { using sym::IterMapExpr; - using sym::IterMapExprNode; + using sym::IterMapExprObj; using sym::NormalizeIterMapToExpr; ffi::Map block_var_subst; TVM_FFI_ICHECK_EQ(iter_vars.size() + 1, division.size()); @@ -301,7 +301,7 @@ SBlockRealize GenerateInner(bool is_write_reduction, const ffi::Array& iter_values, // const PrimExpr& predicate, // SBlock block) { - SBlockNode* n = block.CopyOnWrite(); + SBlockObj* n = block.CopyOnWrite(); n->iter_vars = iter_vars; n->init = std::nullopt; if (is_write_reduction) { @@ -323,7 +323,7 @@ SBlockRealize GenerateInner(bool is_write_reduction, * \return The subtree of the init block and its outer loops. */ Stmt GenerateOuterInit(const Stmt& block_init, const SBlockRealize& inner_realize, - const std::vector& loops, ffi::String block_name) { + const std::vector& loops, ffi::String block_name) { const SBlock& inner_block = inner_realize->block; ffi::Map subst_map; // Step 1: Create new block vars for the block inside the init stmt of outer block @@ -346,7 +346,7 @@ Stmt GenerateOuterInit(const Stmt& block_init, const SBlockRealize& inner_realiz }; if (old_iter_var->iter_type == IterVarType::kDataPar && ffi::StructuralWalk(block_init, walkfn).has_value()) { - ffi::ObjectPtr new_iter_var = ffi::make_object(*old_iter_var.get()); + ffi::ObjectPtr new_iter_var = ffi::make_object(*old_iter_var.get()); new_iter_var->var = new_iter_var->var.CopyWithSuffix("_init"); subst_map.Set(old_iter_var->var, new_iter_var->var); iter_vars.push_back(IterVar(new_iter_var)); @@ -365,7 +365,7 @@ Stmt GenerateOuterInit(const Stmt& block_init, const SBlockRealize& inner_realiz /*body=*/block_init, /*init=*/std::nullopt)); // Step 3. Create the loop nest on top of the block - for (const ForNode* loop : loops) { + for (const ForObj* loop : loops) { bool is_init_loop = false; auto walkfn = [target = loop->var.get()](const Var& var) -> ffi::Expected { return var.get() == target ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var)) @@ -378,7 +378,7 @@ Stmt GenerateOuterInit(const Stmt& block_init, const SBlockRealize& inner_realiz } } if (is_init_loop) { - ffi::ObjectPtr new_loop = ffi::make_object(*loop); + ffi::ObjectPtr new_loop = ffi::make_object(*loop); new_loop->var = loop->var.CopyWithSuffix(""); new_loop->body = std::move(stmt); subst_map.Set(loop->var, new_loop->var); @@ -419,7 +419,7 @@ Stmt ReplaceAndSimplify(const Stmt& stmt, const ffi::Map& sub, UnchangedOr Mutate(ffi::AnyView value, InplaceMode inplace_mode) final { auto result = StmtExprMutator::Mutate(value, inplace_mode); - if (!value.as() || result.UnchangedOrSameAs(value)) return result; + if (!value.as() || result.UnchangedOrSameAs(value)) return result; auto replacement = std::move(result).ValueOrUnchanged(value); if (auto prim_result = replacement.as()) { return ffi::Any(analyzer_->Simplify(prim_result.value())); @@ -427,7 +427,7 @@ Stmt ReplaceAndSimplify(const Stmt& stmt, const ffi::Map& sub, return replacement; } - UnchangedOr Mutate_(const SBlockNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const SBlockObj* op, InplaceMode inplace_mode) final { SBlock src = ffi::GetRef(op); SBlock tgt = StmtExprMutator::Mutate_(op, inplace_mode) .ValueOrUnchanged(ffi::GetRef(op)) @@ -509,9 +509,9 @@ ffi::Array UnionRegions(const ffi::Array& regions) { * \param loops The loop nests * \return The wrapped stmt. */ -Stmt MakeLoopNest(Stmt stmt, const std::vector& loops) { - for (const ForNode* loop : loops) { - ffi::ObjectPtr new_loop = ffi::make_object(*loop); +Stmt MakeLoopNest(Stmt stmt, const std::vector& loops) { + for (const ForObj* loop : loops) { + ffi::ObjectPtr new_loop = ffi::make_object(*loop); new_loop->body = std::move(stmt); stmt = For(new_loop); } @@ -527,7 +527,7 @@ SBlockRealize BlockizeImpl(const ScheduleState& self, const StmtSRef& loop_sref, SBlock block = block_realize->block; StmtSRef block_sref = self->stmt2ref.at(block.get()); // Step 2: Derive subspace division - std::vector loops; + std::vector loops; ffi::Array> division = SubspaceDivide(block_realize, block_sref, loop_sref, &loops, analyzer, preserve_unit_iters); if (division.empty()) { @@ -550,7 +550,11 @@ SBlockRealize BlockizeImpl(const ScheduleState& self, const StmtSRef& loop_sref, ffi::Map inner_iter_dom; for (const IterVar& iter : inner_iter_vars) { inner_iter_dom.Set(iter->var, sym::IntSet::FromRange(iter->dom.value())); - analyzer->Bind(iter->var, iter->dom.value()); + { + auto context = analyzer->GetConstraintContext(); + context->Append(sym::RangeConstraint(iter->var, iter->dom.value())); + analyzer->SyncConstraintContext(context); + } } SBlock block_subst = ReplaceAndSimplify(block, block_var_subst, block_sref_reuse, analyzer).as_or_throw(); @@ -620,7 +624,7 @@ SBlockRealize BlockizeBlocks(const ScheduleState& self, const ffi::Arrayblock; // Step 1: Derive subspace division - std::vector loops; + std::vector loops; ffi::Array> division = SubspaceDivide( block_realize, block_sref, lca, &loops, analyzer.get(), preserve_unit_iters, true); if (division.empty()) { @@ -658,7 +662,11 @@ SBlockRealize BlockizeBlocks(const ScheduleState& self, const ffi::Array(); Range dom = Range::FromMinExtent(min, extent); inner_iter_dom.Set(iter->var, sym::IntSet::FromRange(dom)); - analyzer->Bind(iter->var, dom); + { + auto context = analyzer->GetConstraintContext(); + context->Append(sym::RangeConstraint(iter->var, dom)); + analyzer->SyncConstraintContext(context); + } } SBlock block_subst = ReplaceAndSimplify(block, block_var_subst, block_sref_reuse, analyzer.get()) @@ -687,8 +695,8 @@ SBlockRealize BlockizeBlocks(const ScheduleState& self, const ffi::ArraySet(block, inner_realize->block); Stmt stmt = inner_realize; - for (const ForNode* loop : loops) { - ffi::ObjectPtr new_loop = ffi::make_object(*loop); + for (const ForObj* loop : loops) { + ffi::ObjectPtr new_loop = ffi::make_object(*loop); new_loop->body = std::move(stmt); new_loop->end = new_loop->begin + @@ -717,7 +725,7 @@ class BlockizeRewriter : public StmtExprMutator { using StmtExprMutator::Mutate; using StmtExprMutator::Mutate_; UnchangedOr Mutate(ffi::AnyView value, InplaceMode inplace_mode) override { - if (value.as()) return ffi::Unchanged(); + if (value.as()) return ffi::Unchanged(); return StmtExprMutator::Mutate(value, inplace_mode); } @@ -734,7 +742,7 @@ class BlockizeRewriter : public StmtExprMutator { private: Stmt RewriteSeq(const Stmt& stmt) { - const SeqStmtNode* seq = stmt.as(); + const SeqStmtObj* seq = stmt.as(); TVM_FFI_ICHECK(seq) << "Target blocks must not be nested with each other!"; int idx_start = -1; int last_found_idx = -1; @@ -761,7 +769,7 @@ class BlockizeRewriter : public StmtExprMutator { return SeqStmt(new_seq, seq->loc); } - UnchangedOr Mutate_(const ForNode* loop, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const ForObj* loop, InplaceMode inplace_mode) final { if (loop == lca_->stmt) { return For(loop->var, loop->begin, loop->end, loop->step, loop->kind, RewriteSeq(loop->body), loop->annotations, loop->loc); @@ -769,7 +777,7 @@ class BlockizeRewriter : public StmtExprMutator { return StmtExprMutator::Mutate_(loop, inplace_mode); } - UnchangedOr Mutate_(const SBlockNode* block, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const SBlockObj* block, InplaceMode inplace_mode) final { if (block == lca_->stmt) { return SBlock(block->iter_vars, block->reads, block->writes, block->name_hint, RewriteSeq(block->body), block->init, block->alloc_tensors, @@ -809,11 +817,11 @@ void Tensorize(ScheduleState self, const StmtSRef& sref, const TensorIntrin& int // Step 1: Blockize the subtree rooted at the given loop if needed ffi::Optional old_block = std::nullopt; SBlockRealize block_realize = [&]() -> SBlockRealize { - if (sref->stmt->IsInstance()) { + if (sref->stmt->IsInstance()) { SBlockRealize result = GetSBlockRealize(self, sref); old_block = result->block; return result; - } else if (sref->stmt->IsInstance()) { + } else if (sref->stmt->IsInstance()) { sym::Analyzer analyzer; ffi::Map block_sref_reuse; return BlockizeImpl(self, sref, &block_sref_reuse, analyzer.get(), preserve_unit_iters); @@ -905,7 +913,7 @@ void Tensorize(ScheduleState self, const StmtSRef& sref, const TensorIntrin& int } // Step 5: Replace the subtree in the original IR with the tensor intrin impl. { - SBlockNode* block = block_realize.CopyOnWrite()->block.CopyOnWrite(); + SBlockObj* block = block_realize.CopyOnWrite()->block.CopyOnWrite(); block->body = impl_block->body; block->match_tensors = std::move(match_tensor_regions); for (const auto& [key, val] : impl_block->annotations) { @@ -926,7 +934,7 @@ void Tensorize(ScheduleState self, const StmtSRef& sref, const TensorIntrin& int // Step 6: Update the cached flags. StmtSRef result = self->stmt2ref.at(block_realize->block.get()); StmtSRef scope_root = GetScopeRoot(self, result, /*require_stage_pipeline=*/false); - self->UpdateScopeSBlockInfo(scope_root->StmtAs()->body); + self->UpdateScopeSBlockInfo(scope_root->StmtAs()->body); } /******** InstructionKind Registration ********/ diff --git a/src/s_tir/schedule/primitive/cache_index.cc b/src/s_tir/schedule/primitive/cache_index.cc index d2bc494d6db3..7ed8b9d1409e 100644 --- a/src/s_tir/schedule/primitive/cache_index.cc +++ b/src/s_tir/schedule/primitive/cache_index.cc @@ -107,7 +107,7 @@ class IndexInfoCollector : public StmtExprVisitor { : self_(self), block_sref_(block_sref), scope_sref_(scope_sref), cse_thresh_(cse_thresh) {} private: - ffi::Optional Visit_(const SeqStmtNode* seq_stmt) final { + ffi::Optional Visit_(const SeqStmtObj* seq_stmt) final { for (size_t i = 0; i < seq_stmt->size(); ++i) { if (loc_pos_ != -1) { break; @@ -123,7 +123,7 @@ class IndexInfoCollector : public StmtExprVisitor { return std::nullopt; } - ffi::Optional Visit_(const SBlockNode* block) final { + ffi::Optional Visit_(const SBlockObj* block) final { visiting_target_sblock = static_cast(block_sref_->stmt == block); TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(block)); visiting_target_sblock = false; @@ -143,7 +143,7 @@ class IndexInfoCollector : public StmtExprVisitor { return std::nullopt; } - ffi::Optional Visit_(const ForNode* loop) final { + ffi::Optional Visit_(const ForObj* loop) final { range_map_.Set(loop->var, Range(loop->begin, loop->end)); TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(loop)); // Update seq pos only at top scope @@ -153,12 +153,12 @@ class IndexInfoCollector : public StmtExprVisitor { return std::nullopt; } - ffi::Optional Visit_(const TensorStoreNode* store) final { + ffi::Optional Visit_(const TensorStoreObj* store) final { // Only analyze the cache candidate for stores in target block if (visiting_target_sblock) { auto IsEligibleComputation = [](const PrimExpr& expr) { return (SideEffect(expr) <= CallEffectKind::kPure && CalculateExprComplexity(expr) > 1 && - (expr.as() == nullptr) && (expr.as() == nullptr)); + (expr.as() == nullptr) && (expr.as() == nullptr)); }; // Analyze sub expr candidates @@ -287,7 +287,7 @@ ffi::Array MakeIndexCacheStage(IndexInfo* info, const ffi::String& stora info->cache_tensor.push_back( TensorVar(index_tensor_name, TensorType(tensor_shape, data_ty, runtime::MemoryScope::Get(storage_scope), 0, - tvm::tirx::TileLayoutNode::DefaultLayout(tensor_shape)))); + tvm::tirx::TileLayout::DefaultLayout(tensor_shape)))); // Create loop vars and block vars' binding_value std::vector loop_vars; @@ -379,7 +379,7 @@ ffi::Array MakeIndexCacheStage(IndexInfo* info, const ffi::String& stora * \return A SeqStmt, the result after insertion */ Stmt InsertIndexStage(const Stmt& stmt, int pos, const Stmt& stage) { - if (const auto* seq_stmt = stmt.as()) { + if (const auto* seq_stmt = stmt.as()) { ffi::Array result = seq_stmt->seq; result.insert(result.begin() + pos, stage); return SeqStmt(result, seq_stmt->loc); @@ -421,7 +421,7 @@ class CacheIndexRewriter : public StmtExprMutator { } private: - UnchangedOr Mutate_(const SBlockNode* block, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const SBlockObj* block, InplaceMode inplace_mode) final { SBlock old_stmt = ffi::GetRef(block); // Mutate the body visiting_target_sblock = static_cast(block == info_->target_sblock->stmt); @@ -433,7 +433,7 @@ class CacheIndexRewriter : public StmtExprMutator { // Check if it is the block corresponding to the parent scope if (block == scope_sref_->stmt) { // If so, put tensor allocation and insert cache stages on the parent scope - ffi::ObjectPtr n = ffi::make_object(*stmt.as()); + ffi::ObjectPtr n = ffi::make_object(*stmt.as()); n->body = InsertIndexStage(n->body, info_->loc_pos, info_->cache_stage); for (const TensorVar& it : info_->cache_tensor) { n->alloc_tensors.push_back(it); @@ -444,7 +444,7 @@ class CacheIndexRewriter : public StmtExprMutator { return stmt; } - UnchangedOr Mutate_(const TensorStoreNode* store, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const TensorStoreObj* store, InplaceMode inplace_mode) final { Stmt ret_stmt = StmtExprMutator::Mutate_(store, inplace_mode).ValueOrUnchanged(ffi::GetRef(store)); // Replace common sub expr for target block, with cached tensor load diff --git a/src/s_tir/schedule/primitive/cache_index_helpers.cc b/src/s_tir/schedule/primitive/cache_index_helpers.cc index cf0e6736b66e..bdbae8bcd913 100644 --- a/src/s_tir/schedule/primitive/cache_index_helpers.cc +++ b/src/s_tir/schedule/primitive/cache_index_helpers.cc @@ -150,8 +150,8 @@ ComputationTable BuildTableForThreeChildrenNode(const ComputationTable& table_ch ComputationTable ComputationsDoneBy::GetComputationsDoneBy( const PrimExpr& expr, std::function is_eligible_computation, std::function can_contain_computations) { - if (expr.as() != nullptr || expr.as() != nullptr || - expr.as() != nullptr || expr.as()) { + if (expr.as() != nullptr || expr.as() != nullptr || + expr.as() != nullptr || expr.as()) { return {}; } @@ -218,8 +218,8 @@ ffi::Optional ComputationsDoneBy::Visit(ffi::AnyView expr_value) return s_tir::StmtExprVisitor::Visit(expr_value); } PrimExpr expr = opt_expr.value(); - if (expr.as() != nullptr || expr.as() != nullptr || - expr.as() != nullptr || expr.as()) { + if (expr.as() != nullptr || expr.as() != nullptr || + expr.as() != nullptr || expr.as()) { return std::nullopt; } @@ -244,9 +244,9 @@ ffi::Optional ComputationsDoneBy::Visit(ffi::AnyView expr_value) } /*! - * \brief The method which overrides the specific treatment for an IfNode + * \brief The method which overrides the specific treatment for an IfObj */ -ffi::Optional ComputationsDoneBy::Visit_(const IfNode* op) { +ffi::Optional ComputationsDoneBy::Visit_(const IfObj* op) { TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Visit(op->condition)); ComputationTable computations_done_by_cond = table_of_computations_; table_of_computations_.clear(); @@ -271,9 +271,9 @@ ffi::Optional ComputationsDoneBy::Visit_(const IfNode* op) { } /*! - * \brief The method which overrides the specific treatment for a ForNode + * \brief The method which overrides the specific treatment for a ForObj */ -ffi::Optional ComputationsDoneBy::Visit_(const ForNode* op) { +ffi::Optional ComputationsDoneBy::Visit_(const ForObj* op) { TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Visit(op->begin)); ComputationTable computations_done_by_min = table_of_computations_; table_of_computations_.clear(); @@ -296,9 +296,9 @@ ffi::Optional ComputationsDoneBy::Visit_(const ForNode* op) { } /*! - * \brief The method which overrides the specific treatment for a WhileNode + * \brief The method which overrides the specific treatment for a WhileObj */ -ffi::Optional ComputationsDoneBy::Visit_(const WhileNode* op) { +ffi::Optional ComputationsDoneBy::Visit_(const WhileObj* op) { TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Visit(op->condition)); ComputationTable computations_done_by_condition = table_of_computations_; table_of_computations_.clear(); diff --git a/src/s_tir/schedule/primitive/cache_index_helpers.h b/src/s_tir/schedule/primitive/cache_index_helpers.h index 34382e0f0b96..83b2a87cdad2 100644 --- a/src/s_tir/schedule/primitive/cache_index_helpers.h +++ b/src/s_tir/schedule/primitive/cache_index_helpers.h @@ -95,9 +95,9 @@ class ComputationsDoneBy : public s_tir::StmtExprVisitor { protected: ffi::Optional Visit(ffi::AnyView expr) override; - ffi::Optional Visit_(const IfNode* op) override; - ffi::Optional Visit_(const ForNode* op) override; - ffi::Optional Visit_(const WhileNode* op) override; + ffi::Optional Visit_(const IfObj* op) override; + ffi::Optional Visit_(const ForObj* op) override; + ffi::Optional Visit_(const WhileObj* op) override; private: static ComputationTable ComputationsDoneByChildrenOf( diff --git a/src/s_tir/schedule/primitive/cache_read_write.cc b/src/s_tir/schedule/primitive/cache_read_write.cc index 26b6e63dff56..2b1cc62f2053 100644 --- a/src/s_tir/schedule/primitive/cache_read_write.cc +++ b/src/s_tir/schedule/primitive/cache_read_write.cc @@ -48,7 +48,7 @@ class NotSingleWriteBlock : public ScheduleErrorContextObj { TVM_FFI_ICHECK_GT(write_blocks.size(), 1); write_blocks_.reserve(write_blocks.size()); for (const StmtSRef& block_sref : write_blocks) { - const SBlockNode* block = TVM_SREF_TO_SBLOCK(block_sref); + const SBlockObj* block = TVM_SREF_TO_SBLOCK(block_sref); write_blocks_.push_back(ffi::GetRef(block)); } } @@ -484,7 +484,7 @@ bool CalculateAffineFlag(const ScheduleState& self, const StmtSRef& block_sref) * \return A SeqStmt, the result after insertion */ Stmt InsertCacheStage(const Stmt& stmt, int pos, const Stmt& stage) { - if (const auto* seq_stmt = stmt.as()) { + if (const auto* seq_stmt = stmt.as()) { ffi::Array seq = seq_stmt->seq; TVM_FFI_ICHECK_LE(pos, seq.size()) << "Cannot insert at position " << pos << " into sequence of length " << seq.size(); @@ -545,9 +545,9 @@ ffi::Optional GetOnlyWriteBlock(ScheduleState self, const StmtSRef& sc bool AllConsumersUnderStmt(ScheduleState self, TensorVar tensor, StmtSRef scope_sref, StmtSRef stmt_sref) { // Collect all children blocks of the target stmt. - std::unordered_set blocks_under_target; + std::unordered_set blocks_under_target; for (const StmtSRef& block_sref : GetChildBlocks(self, stmt_sref)) { - const auto* block = block_sref->StmtAs(); + const auto* block = block_sref->StmtAs(); TVM_FFI_ICHECK(block != nullptr); blocks_under_target.insert(block); } @@ -556,7 +556,7 @@ bool AllConsumersUnderStmt(ScheduleState self, TensorVar tensor, StmtSRef scope_ // input tensor, check if it is also a child block of the // target stmt. for (const StmtSRef& block_sref : GetChildBlocks(self, scope_sref)) { - const auto* block = block_sref->StmtAs(); + const auto* block = block_sref->StmtAs(); TVM_FFI_ICHECK(block != nullptr); if (GetTensorRegionFromTensor(block->reads, tensor).has_value()) { if (blocks_under_target.find(block) == blocks_under_target.end()) { @@ -587,15 +587,15 @@ static PrimExpr CollectNestedBlockPredicates(const Stmt& body, const TensorVar& using StmtExprVisitor::Visit_; ffi::Optional Visit(ffi::AnyView value) override { - if (value.as()) return std::nullopt; + if (value.as()) return std::nullopt; return StmtExprVisitor::Visit(value); } Collector(const TensorVar& tensor, TensorIndexType idx_type) : tensor_(tensor), index_type_(idx_type), result_(IntImm::Bool(false)), found_(false) {} - ffi::Optional Visit_(const SBlockRealizeNode* realize) final { - const SBlockNode* block = realize->block.get(); + ffi::Optional Visit_(const SBlockRealizeObj* realize) final { + const SBlockObj* block = realize->block.get(); const auto& regions = (index_type_ == TensorIndexType::kRead) ? block->reads : block->writes; bool accesses_tensor = false; for (const TensorRegion& region : regions) { @@ -700,7 +700,7 @@ class CacheLocDetector : public StmtExprVisitor { using StmtExprVisitor::Visit_; ffi::Optional Visit(ffi::AnyView value) override { - if (value.as()) return std::nullopt; + if (value.as()) return std::nullopt; return StmtExprVisitor::Visit(value); } @@ -750,7 +750,7 @@ class CacheLocDetector : public StmtExprVisitor { } else { info->loc_sref = scope_sref; - info->loc_pos = scope_sref->StmtAs()->body->size(); + info->loc_pos = scope_sref->StmtAs()->body->size(); } } @@ -770,7 +770,7 @@ class CacheLocDetector : public StmtExprVisitor { related_blocks_(related_blocks) {} private: - ffi::Optional Visit_(const SeqStmtNode* seq_stmt) final { + ffi::Optional Visit_(const SeqStmtObj* seq_stmt) final { bool previous_visited_block = visited_block_; visited_block_ = false; @@ -793,7 +793,7 @@ class CacheLocDetector : public StmtExprVisitor { return std::nullopt; } - ffi::Optional Visit_(const SBlockNode* block) final { + ffi::Optional Visit_(const SBlockObj* block) final { // Only visit the current scope under tensor writer's parent block if (block == scope_sref_->stmt) { // The block visited is the current parent scope @@ -823,7 +823,7 @@ class CacheLocDetector : public StmtExprVisitor { return std::nullopt; } - ffi::Optional Visit_(const ForNode* loop) final { + ffi::Optional Visit_(const ForObj* loop) final { TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(loop)); if (visited_block_ && visited_related_ && !loc_sref_.has_value() && loc_pos_ != -1) { loc_sref_ = self_->stmt2ref.at(loop); @@ -855,7 +855,7 @@ class CacheInplaceLocDetector : public StmtExprVisitor { using StmtExprVisitor::Visit_; ffi::Optional Visit(ffi::AnyView value) override { - if (value.as()) return std::nullopt; + if (value.as()) return std::nullopt; return StmtExprVisitor::Visit(value); } @@ -886,7 +886,7 @@ class CacheInplaceLocDetector : public StmtExprVisitor { : self_(self), block_sref_(block_sref), scope_sref_(scope_sref) {} private: - ffi::Optional Visit_(const SeqStmtNode* seq_stmt) final { + ffi::Optional Visit_(const SeqStmtObj* seq_stmt) final { for (size_t i = 0; i < seq_stmt->size(); ++i) { if (loc_pos_ != -1) { break; @@ -902,7 +902,7 @@ class CacheInplaceLocDetector : public StmtExprVisitor { return std::nullopt; } - ffi::Optional Visit_(const SBlockNode* block) final { + ffi::Optional Visit_(const SBlockObj* block) final { // Only visit the current scope under tensor writer's parent block if (block == scope_sref_->stmt) { // The block visited is the current parent scope @@ -921,7 +921,7 @@ class CacheInplaceLocDetector : public StmtExprVisitor { return std::nullopt; } - ffi::Optional Visit_(const ForNode* loop) final { + ffi::Optional Visit_(const ForObj* loop) final { TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(loop)); if (visited_block_ && !loc_sref_.has_value()) { loc_sref_ = self_->stmt2ref.at(loop); @@ -1024,27 +1024,27 @@ class CacheReadRewriter : public StmtExprMutator { } private: - UnchangedOr Mutate_(const ForNode* loop, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const ForObj* loop, InplaceMode inplace_mode) final { Stmt stmt = StmtExprMutator::Mutate_(loop, inplace_mode).ValueOrUnchanged(ffi::GetRef(loop)); // Check the insertion point if (loop == info_->loc_sref.value()->stmt) { // Insert cache stage into the loop if it is the right place - ffi::ObjectPtr n = ffi::make_object(*stmt.as()); + ffi::ObjectPtr n = ffi::make_object(*stmt.as()); n->body = InsertCacheStage(n->body, info_->loc_pos, info_->cache_stage); stmt = Stmt(n); } return stmt; } - UnchangedOr Mutate_(const SBlockNode* block, InplaceMode inplace_mode) override { + UnchangedOr Mutate_(const SBlockObj* block, InplaceMode inplace_mode) override { SBlock old_stmt = ffi::GetRef(block); // Check if this block is one of the specified consumers. // If no consumer blocks are specified, all blocks should be considered consumers. bool is_consumer = info_->consumer_blocks.empty(); // Otherwise check if this is one of the specified blocks. for (StmtSRef consumer_sref : info_->consumer_blocks) { - const SBlockNode* consumer_node = TVM_SREF_TO_SBLOCK(consumer_sref); + const SBlockObj* consumer_node = TVM_SREF_TO_SBLOCK(consumer_sref); SBlock consumer_block = ffi::GetRef(consumer_node); if (old_stmt.same_as(consumer_block)) { is_consumer = true; @@ -1059,7 +1059,7 @@ class CacheReadRewriter : public StmtExprMutator { } // Mutate the body // Cache accesses change storage; the original allocation still owns its source. - auto input_node = ffi::make_object(*block); + auto input_node = ffi::make_object(*block); input_node->alloc_tensors.clear(); SBlock input(std::move(input_node)); SBlock stmt = StmtExprMutator::Mutate_(input.get(), InplaceMode::kDisallow) @@ -1069,14 +1069,14 @@ class CacheReadRewriter : public StmtExprMutator { // Check the insertion point if (block == info_->loc_sref.value()->stmt) { // Insert cache stage into the block if it is the right place - ffi::ObjectPtr n = ffi::make_object(*stmt.as()); + ffi::ObjectPtr n = ffi::make_object(*stmt.as()); n->body = InsertCacheStage(n->body, info_->loc_pos, info_->cache_stage); stmt = SBlock(n); } // Check if it is the block corresponding to the parent scope if (block == scope_sref_->stmt) { // If so, put tensor allocation on the parent scope - ffi::ObjectPtr n = ffi::make_object(*stmt.as()); + ffi::ObjectPtr n = ffi::make_object(*stmt.as()); // In cache_inplace case, alloc_tensor may be already exits. if (info_->alloc.has_value()) { n->alloc_tensors.push_back(info_->alloc.value()); @@ -1090,7 +1090,7 @@ class CacheReadRewriter : public StmtExprMutator { ffi::Array reads = update_access_regions(stmt->reads); ffi::Array match_tensors = update_match_tensors(stmt->match_tensors); if (!reads.same_as(stmt->reads) || !match_tensors.same_as(stmt->match_tensors)) { - ffi::ObjectPtr n = ffi::make_object(*stmt.as()); + ffi::ObjectPtr n = ffi::make_object(*stmt.as()); n->reads = std::move(reads); n->match_tensors = std::move(match_tensors); stmt = SBlock(n); @@ -1109,10 +1109,10 @@ class CacheReadRewriter : public StmtExprMutator { return ret; } - UnchangedOr Mutate_(const CallNode* op, InplaceMode inplace_mode) override { + UnchangedOr Mutate_(const CallObj* op, InplaceMode inplace_mode) override { auto result = StmtExprMutator::Mutate_(op, inplace_mode); if (!result.IsUnchanged()) { - op = ffi::AnyView(result).as(); + op = ffi::AnyView(result).as(); if (!op->unique()) inplace_mode = InplaceMode::kDisallow; } // Cache remapping can change pointer storage scope; the base Call hook preserves its type. @@ -1120,15 +1120,15 @@ class CacheReadRewriter : public StmtExprMutator { PtrType type = op->args[0].as_or_throw().ty()->DataPtrType(); if (ffi::StructuralEqual()(op->ty, type)) return result; if (inplace_mode == InplaceMode::kAllow) { - const_cast(op)->ty = std::move(type); + const_cast(op)->ty = std::move(type); return result; } - auto copy = ffi::make_object(*op); + auto copy = ffi::make_object(*op); copy->ty = std::move(type); return Expr(std::move(copy)); } - UnchangedOr Mutate_(const TensorLoadNode* load, InplaceMode inplace_mode) override { + UnchangedOr Mutate_(const TensorLoadObj* load, InplaceMode inplace_mode) override { if (load->source.as_or_throw().same_as(info_->read_tensor) && current_block_consumes) { ffi::Array indices = load->indices; @@ -1149,7 +1149,7 @@ class CacheReadRewriter : public StmtExprMutator { return node; } - UnchangedOr Mutate_(const TensorRegionNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const TensorRegionObj* op, InplaceMode inplace_mode) final { if (!op->source.as()) { return StmtExprMutator::Mutate_(op, inplace_mode); } @@ -1238,7 +1238,7 @@ class ReindexCacheReadRewriter : public CacheReadRewriter { } private: - UnchangedOr Mutate_(const TensorLoadNode* load, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const TensorLoadObj* load, InplaceMode inplace_mode) final { if (load->source.as_or_throw().same_as(info_->read_tensor) && current_block_consumes) { TensorLoad node = ffi::GetRef(load); @@ -1337,26 +1337,26 @@ class CacheWriteRewriter : public StmtExprMutator { } private: - UnchangedOr Mutate_(const ForNode* loop, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const ForObj* loop, InplaceMode inplace_mode) final { Stmt stmt = StmtExprMutator::Mutate_(loop, inplace_mode).ValueOrUnchanged(ffi::GetRef(loop)); // Check the insertion point if (loop == info_->loc_sref.value()->stmt) { // Insert cache stage into the loop if it is the right place - ffi::ObjectPtr n = ffi::make_object(*stmt.as()); + ffi::ObjectPtr n = ffi::make_object(*stmt.as()); n->body = InsertCacheStage(n->body, info_->loc_pos, info_->cache_stage); stmt = Stmt(n); } return stmt; } - UnchangedOr Mutate_(const SBlockNode* block, InplaceMode inplace_mode) override { + UnchangedOr Mutate_(const SBlockObj* block, InplaceMode inplace_mode) override { SBlock old_stmt = ffi::GetRef(block); // Check if this block is one of the specified cache consumers. // update the read tensor to the cache. for (StmtSRef consumer_sref : info_->consumer_blocks) { - const SBlockNode* consumer_node = TVM_SREF_TO_SBLOCK(consumer_sref); + const SBlockObj* consumer_node = TVM_SREF_TO_SBLOCK(consumer_sref); SBlock consumer_block = ffi::GetRef(consumer_node); if (old_stmt.same_as(consumer_block)) { ffi::Array writes = update_access_regions(block->writes); @@ -1365,7 +1365,7 @@ class CacheWriteRewriter : public StmtExprMutator { if (!writes.same_as(block->writes) || !reads.same_as(block->reads) || !match_tensors.same_as(block->match_tensors)) { SBlock new_consumer = old_stmt; - SBlockNode* n = new_consumer.CopyOnWrite(); + SBlockObj* n = new_consumer.CopyOnWrite(); n->writes = std::move(writes); n->reads = std::move(reads); n->match_tensors = std::move(match_tensors); @@ -1386,7 +1386,7 @@ class CacheWriteRewriter : public StmtExprMutator { bool under_scope = under_writer_block_ || block == writer_block_sref_->stmt; std::swap(under_scope, under_writer_block_); // Cache accesses change storage; the original allocation still owns its source. - auto input_node = ffi::make_object(*block); + auto input_node = ffi::make_object(*block); input_node->alloc_tensors.clear(); SBlock input(std::move(input_node)); SBlock stmt = StmtExprMutator::Mutate_(input.get(), InplaceMode::kDisallow) @@ -1397,13 +1397,13 @@ class CacheWriteRewriter : public StmtExprMutator { // Find the insertion point if (block == info_->loc_sref.value()->stmt) { - ffi::ObjectPtr n = ffi::make_object(*stmt.as()); + ffi::ObjectPtr n = ffi::make_object(*stmt.as()); n->body = InsertCacheStage(n->body, info_->loc_pos, info_->cache_stage); stmt = SBlock(n); } // Put tensor allocation on the parent scope if (block == scope_sref_->stmt) { - ffi::ObjectPtr n = ffi::make_object(*stmt.as()); + ffi::ObjectPtr n = ffi::make_object(*stmt.as()); // In cache_inplace case, alloc_tensor may be already exits. if (info_->alloc.has_value()) { n->alloc_tensors.push_back(info_->alloc.value()); @@ -1416,7 +1416,7 @@ class CacheWriteRewriter : public StmtExprMutator { auto match_tensors = update_match_tensors(block->match_tensors); if (!writes.same_as(block->writes) || !reads.same_as(block->reads) || !match_tensors.same_as(block->match_tensors)) { - ffi::ObjectPtr n = ffi::make_object(*stmt.as()); + ffi::ObjectPtr n = ffi::make_object(*stmt.as()); n->writes = std::move(writes); n->reads = std::move(reads); n->match_tensors = std::move(match_tensors); @@ -1435,7 +1435,7 @@ class CacheWriteRewriter : public StmtExprMutator { return ret; } - UnchangedOr Mutate_(const TensorStoreNode* store, InplaceMode inplace_mode) override { + UnchangedOr Mutate_(const TensorStoreObj* store, InplaceMode inplace_mode) override { bool rewrite_tensor = store->dest.as_or_throw().same_as(info_->write_tensor); auto value = Mutate(store->value); auto indices = Mutate(store->indices).as_or_throw>>(); @@ -1446,7 +1446,7 @@ class CacheWriteRewriter : public StmtExprMutator { n->indices = std::move(indices).ValueOrUnchanged(store->indices); } if (rewrite_tensor) { - TensorStoreNode* n = stmt.CopyOnWrite(); + TensorStoreObj* n = stmt.CopyOnWrite(); n->dest = info_->read_tensor; if (!cache_full_region_) { n->indices = RewriteIndices(n->indices); @@ -1457,10 +1457,10 @@ class CacheWriteRewriter : public StmtExprMutator { } } - UnchangedOr Mutate_(const CallNode* op, InplaceMode inplace_mode) override { + UnchangedOr Mutate_(const CallObj* op, InplaceMode inplace_mode) override { auto result = StmtExprMutator::Mutate_(op, inplace_mode); if (!result.IsUnchanged()) { - op = ffi::AnyView(result).as(); + op = ffi::AnyView(result).as(); if (!op->unique()) inplace_mode = InplaceMode::kDisallow; } // Cache remapping can change pointer storage scope; the base Call hook preserves its type. @@ -1468,15 +1468,15 @@ class CacheWriteRewriter : public StmtExprMutator { PtrType type = op->args[0].as_or_throw().ty()->DataPtrType(); if (ffi::StructuralEqual()(op->ty, type)) return result; if (inplace_mode == InplaceMode::kAllow) { - const_cast(op)->ty = std::move(type); + const_cast(op)->ty = std::move(type); return result; } - auto copy = ffi::make_object(*op); + auto copy = ffi::make_object(*op); copy->ty = std::move(type); return Expr(std::move(copy)); } - UnchangedOr Mutate_(const TensorLoadNode* load, InplaceMode inplace_mode) override { + UnchangedOr Mutate_(const TensorLoadObj* load, InplaceMode inplace_mode) override { if (load->source.as_or_throw().same_as(info_->write_tensor)) { ffi::Array indices = load->indices; if (!cache_full_region_) { @@ -1496,7 +1496,7 @@ class CacheWriteRewriter : public StmtExprMutator { return node; } - UnchangedOr Mutate_(const TensorRegionNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const TensorRegionObj* op, InplaceMode inplace_mode) final { if (!op->source.as()) { return StmtExprMutator::Mutate_(op, inplace_mode); } @@ -1592,7 +1592,7 @@ class ReindexCacheWriteRewriter : public CacheWriteRewriter { } private: - UnchangedOr Mutate_(const TensorStoreNode* store, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const TensorStoreObj* store, InplaceMode inplace_mode) final { bool rewrite_tensor = store->dest.as_or_throw().same_as(info_->write_tensor); auto value = Mutate(store->value); auto indices = Mutate(store->indices).as_or_throw>>(); @@ -1603,7 +1603,7 @@ class ReindexCacheWriteRewriter : public CacheWriteRewriter { n->indices = std::move(indices).ValueOrUnchanged(store->indices); } if (rewrite_tensor) { - TensorStoreNode* n = stmt.CopyOnWrite(); + TensorStoreObj* n = stmt.CopyOnWrite(); n->dest = info_->read_tensor; n->indices = new_indices_; return stmt; @@ -1612,7 +1612,7 @@ class ReindexCacheWriteRewriter : public CacheWriteRewriter { } } - UnchangedOr Mutate_(const TensorLoadNode* load, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const TensorLoadObj* load, InplaceMode inplace_mode) final { if (load->source.as_or_throw().same_as(info_->write_tensor)) { TensorLoad node = ffi::GetRef(load); auto* n = node.CopyOnWrite(); @@ -1641,7 +1641,7 @@ class ReindexCacheWriteRewriter : public CacheWriteRewriter { */ TensorVar CreateReindexTensor(const TensorVar& tensor, const ffi::Array& block_iters, const std::unordered_set& covered) { - ffi::ObjectPtr new_tensor = ffi::make_object(*tensor.ty().get()); + ffi::ObjectPtr new_tensor = ffi::make_object(*tensor.ty().get()); std::vector new_shape; for (const auto& iter : block_iters) { if (covered.count(iter->var)) { @@ -1652,7 +1652,7 @@ TensorVar CreateReindexTensor(const TensorVar& tensor, const ffi::Array if (new_tensor->layout.has_value()) { auto tile = new_tensor->layout.value().as_or_throw(); new_tensor->layout = - tirx::TileLayout(tirx::TileLayoutNode::DefaultLayout(new_shape)->shard, {}, tile->offset); + tirx::TileLayout(tirx::TileLayout::DefaultLayout(new_shape)->shard, {}, tile->offset); } return TensorVar(tensor->name + "_reindex", TensorType(std::move(new_tensor)), tensor->loc); } @@ -1737,7 +1737,7 @@ class ReIndexCollector : public StmtExprVisitor { : mod_(mod), tensor_(tensor), block_(block) {} private: - ffi::Optional Visit_(const TensorLoadNode* load) final { + ffi::Optional Visit_(const TensorLoadObj* load) final { for (const PrimExpr& index : load->indices) { TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Visit(index)); } @@ -1747,12 +1747,12 @@ class ReIndexCollector : public StmtExprVisitor { return std::nullopt; } - ffi::Optional Visit_(const SBlockNode* block) final { + ffi::Optional Visit_(const SBlockObj* block) final { // no sub-blocks under this block throw MakeScheduleError(mod_, block_); } - ffi::Optional Visit_(const TensorStoreNode* store) final { + ffi::Optional Visit_(const TensorStoreObj* store) final { TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Visit(store->value)); for (const PrimExpr& index : store->indices) { TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Visit(index)); @@ -1775,7 +1775,7 @@ class ReIndexCollector : public StmtExprVisitor { } } - ffi::Optional Visit_(const VarNode* var) final { + ffi::Optional Visit_(const VarObj* var) final { if (def_region_kind() != kTVMFFIDefRegionKindNone) return StmtExprVisitor::Visit_(var); if (var == tensor_.get()) { throw MakeScheduleError( @@ -1816,7 +1816,7 @@ class ReIndexRewriter : public StmtExprMutator { new_tensor_(info->alloc.value()) {} private: - UnchangedOr Mutate_(const SBlockNode* block, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const SBlockObj* block, InplaceMode inplace_mode) final { SBlock old_stmt = ffi::GetRef(block); if (is_scope_) { is_scope_ = false; @@ -1824,7 +1824,7 @@ class ReIndexRewriter : public StmtExprMutator { .ValueOrUnchanged(ffi::GetRef(block)) .as_or_throw(); // Insert cache stage into the loop - ffi::ObjectPtr n = ffi::make_object(*stmt.as()); + ffi::ObjectPtr n = ffi::make_object(*stmt.as()); n->body = InsertCacheStage(n->body, info_->loc_pos, info_->cache_stage); n->alloc_tensors.push_back(info_->alloc.value()); stmt = SBlock(n); @@ -1854,7 +1854,7 @@ class ReIndexRewriter : public StmtExprMutator { TensorRegion(new_tensor_, region_, TensorRegionType())); if (!writes.same_as(block->writes) || !reads.same_as(block->reads) || !match_tensors.same_as(block->match_tensors)) { - ffi::ObjectPtr n = ffi::make_object(*stmt.as()); + ffi::ObjectPtr n = ffi::make_object(*stmt.as()); n->writes = std::move(writes); n->reads = std::move(reads); n->match_tensors = std::move(match_tensors); @@ -1882,7 +1882,7 @@ class ReIndexRewriter : public StmtExprMutator { } return node; } - UnchangedOr Mutate_(const TensorStoreNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const TensorStoreObj* op, InplaceMode inplace_mode) final { auto value = Mutate(op->value); auto indices = Mutate(op->indices).as_or_throw>>(); TensorStore tensor_store = ffi::GetRef(op); @@ -1894,7 +1894,7 @@ class ReIndexRewriter : public StmtExprMutator { return VisitTensorAccess(std::move(tensor_store)); } - UnchangedOr Mutate_(const TensorLoadNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const TensorLoadObj* op, InplaceMode inplace_mode) final { auto indices = Mutate(op->indices).as_or_throw>>(); TensorLoad tensor_load = ffi::GetRef(op); if (!indices.UnchangedOrSameAs(op->indices)) { @@ -1940,11 +1940,11 @@ The region cover property require to hold for every of its child blocks }; for (const auto& child_block_sref : GetChildBlocks(self, scope_root)) { - const SBlockNode* child_block = TVM_SREF_TO_SBLOCK(child_block_sref); + const SBlockObj* child_block = TVM_SREF_TO_SBLOCK(child_block_sref); for (const TensorRegion& region : child_block->reads) { if (region->source.as_or_throw().same_as(read_tensor)) { if (!self->block_info.at(child_block_sref).region_cover) { - const SBlockNode* block = TVM_SREF_TO_SBLOCK(scope_root); + const SBlockObj* block = TVM_SREF_TO_SBLOCK(scope_root); throw MakeScheduleError(self->mod, ffi::GetRef(block)); } } @@ -1971,13 +1971,13 @@ StmtSRef CacheRead(ScheduleState self, const StmtSRef& block_sref, int read_tens CheckStorageScope(self, storage_scope); // Step 1. Check index, getting the target tensor and the parent scope - const SBlockNode* block = TVM_SREF_TO_SBLOCK(block_sref); + const SBlockObj* block = TVM_SREF_TO_SBLOCK(block_sref); TensorVar read_tensor = GetNthAccessTensor(self, ffi::GetRef(block), read_tensor_index, TensorIndexType::kRead); StmtSRef scope_sref = GetScopeRoot(self, block_sref, /*require_stage_pipeline=*/false); // Check required region cover for cache_read CheckRegionCover(self, scope_sref, read_tensor); - const SBlockNode* scope_block = TVM_SREF_TO_SBLOCK(scope_sref); + const SBlockObj* scope_block = TVM_SREF_TO_SBLOCK(scope_sref); // Step 2. Create CacheStageInfo CacheStageInfo info; @@ -1997,7 +1997,7 @@ StmtSRef CacheRead(ScheduleState self, const StmtSRef& block_sref, int read_tens GetOnlyWriteBlock(self, scope_sref, read_tensor)) { // Case 1. The tensor is written inside the block. StmtSRef write_block_sref = _write_block_sref.value(); - const SBlockNode* write_block = TVM_SREF_TO_SBLOCK(write_block_sref); + const SBlockObj* write_block = TVM_SREF_TO_SBLOCK(write_block_sref); // Find the producing region TensorRegion region = GetTensorRegionFromTensor(write_block->writes, read_tensor).value(); StmtSRef parent_sref = ffi::GetRef(write_block_sref->parent); @@ -2035,12 +2035,12 @@ StmtSRef CacheRead(ScheduleState self, const StmtSRef& block_sref, int read_tens // Step 4. Making new cache stage block and rewrite readers. bool cache_full_region = - info.loc_sref.value()->StmtAs() == nullptr || + info.loc_sref.value()->StmtAs() == nullptr || !AllConsumersUnderStmt(self, read_tensor, scope_sref, info.loc_sref.value()); info.cache_region = cache_region; info.write_tensor = WithScope(read_tensor, runtime::MemoryScope::Get(storage_scope)); if (!cache_full_region) { - auto write_tensor = ffi::make_object(*info.write_tensor.ty().get()); + auto write_tensor = ffi::make_object(*info.write_tensor.ty().get()); std::vector shape; for (auto cache_range : info.cache_region->region) { shape.push_back(cache_range->extent); @@ -2084,7 +2084,7 @@ StmtSRef CacheWrite(ScheduleState self, const StmtSRef& block_sref, int write_te CheckStorageScope(self, storage_scope); // Step 1. Checking index, getting the target tensor and the parent scope - const SBlockNode* block = TVM_SREF_TO_SBLOCK(block_sref); + const SBlockObj* block = TVM_SREF_TO_SBLOCK(block_sref); TensorVar write_tensor = GetNthAccessTensor(self, ffi::GetRef(block), write_tensor_index, TensorIndexType::kWrite); StmtSRef scope_sref = GetScopeRoot(self, block_sref, /*require_stage_pipeline=*/false); @@ -2127,12 +2127,12 @@ StmtSRef CacheWrite(ScheduleState self, const StmtSRef& block_sref, int write_te } bool cache_full_region = - info.loc_sref.value()->StmtAs() == nullptr || + info.loc_sref.value()->StmtAs() == nullptr || !AllConsumersUnderStmt(self, write_tensor, scope_sref, info.loc_sref.value()); info.cache_region = cache_region; info.read_tensor = WithScope(write_tensor, runtime::MemoryScope::Get(storage_scope)); if (!cache_full_region) { - auto read_tensor_type = ffi::make_object(*info.read_tensor.ty().get()); + auto read_tensor_type = ffi::make_object(*info.read_tensor.ty().get()); std::vector shape; for (auto cache_range : info.cache_region->region) { shape.push_back(cache_range->extent); @@ -2164,7 +2164,7 @@ StmtSRef CacheWrite(ScheduleState self, const StmtSRef& block_sref, int write_te ffi::Array GetLoopsUnderScope(const StmtSRef& block_sref, const ffi::Optional& top_sref) { std::vector result; - for (StmtSRefObj* parent = block_sref->parent; parent && parent->stmt->IsInstance(); + for (StmtSRefObj* parent = block_sref->parent; parent && parent->stmt->IsInstance(); parent = parent->parent) { if (top_sref.has_value() && parent == top_sref.value().get()) break; result.push_back(ffi::GetRef(parent)); @@ -2270,7 +2270,7 @@ void CollectReindexCacheStageInfoAndCreateTensor( auto iter_value_vars = collect_vars(info->block_iter_values); for (const StmtSRef& loop_sref : GetLoopsUnderScope(block_sref, info->loc_sref)) { - const ForNode* loop = TVM_SREF_TO_FOR(loop_sref); + const ForObj* loop = TVM_SREF_TO_FOR(loop_sref); if (iter_value_vars.count(loop->var)) { info->loop_vars.push_back(loop->var); info->loop_ranges.push_back(Range(loop->begin, loop->end)); @@ -2278,8 +2278,8 @@ void CollectReindexCacheStageInfoAndCreateTensor( } // Create new tensor - ffi::ObjectPtr new_tensor = - ffi::make_object(*old_tensor.ty().get()); + ffi::ObjectPtr new_tensor = + ffi::make_object(*old_tensor.ty().get()); new_tensor->scope = runtime::MemoryScope::Get(storage_scope); new_tensor->shape = new_shape; TensorVar rebuilt = TensorVar(old_tensor->name + "_" + storage_scope, @@ -2300,7 +2300,7 @@ void CheckSinglePoint(ScheduleState self, const SBlock& block, const TensorRegio bool single_point = true; for (const Range& range : cache_region->region) { PrimExpr ext_int_expr = range->extent; - const auto* ext_int = ext_int_expr.as(); + const auto* ext_int = ext_int_expr.as(); if (!ext_int || ext_int->value != 1) { single_point = false; } @@ -2479,7 +2479,7 @@ ffi::Array CacheInplace(ScheduleState self, const StmtSRef& block_sref CheckStorageScope(self, storage_scope); // Check 1. Check index, get the target tensor and the parent scope - const SBlockNode* block = TVM_SREF_TO_SBLOCK(block_sref); + const SBlockObj* block = TVM_SREF_TO_SBLOCK(block_sref); TensorVar tensor = GetNthAccessTensor(self, ffi::GetRef(block), read_tensor_index, TensorIndexType::kRead); StmtSRef scope_sref = GetScopeRoot(self, block_sref, /*require_stage_pipeline=*/false); @@ -2488,7 +2488,7 @@ ffi::Array CacheInplace(ScheduleState self, const StmtSRef& block_sref CheckRegionCover(self, scope_sref, tensor); // Check 4. Check if target block both read & write target tensor. - const SBlockNode* rw_block = TVM_SREF_TO_SBLOCK(block_sref); + const SBlockObj* rw_block = TVM_SREF_TO_SBLOCK(block_sref); ffi::Optional read_region = GetTensorRegionFromTensor(rw_block->reads, tensor); ffi::Optional write_region = GetTensorRegionFromTensor(rw_block->writes, tensor); if (!read_region.has_value() || !write_region.has_value()) { @@ -2560,7 +2560,7 @@ ffi::Array CacheInplace(ScheduleState self, const StmtSRef& block_sref StmtSRef ReIndex(ScheduleState self, const StmtSRef& block_sref, int tensor_index, TensorIndexType tensor_index_type) { - const SBlockNode* block_ptr = TVM_SREF_TO_SBLOCK(block_sref); + const SBlockObj* block_ptr = TVM_SREF_TO_SBLOCK(block_sref); SBlock block = ffi::GetRef(block_ptr); TensorVar tensor = GetNthAccessTensor(self, block, tensor_index, tensor_index_type); StmtSRef scope_sref = GetScopeRoot(self, block_sref, /*require_stage_pipeline=*/true); @@ -2571,7 +2571,11 @@ StmtSRef ReIndex(ScheduleState self, const StmtSRef& block_sref, int tensor_inde ffi::Array original_indices = ReIndexCollector::Collect(self->mod, tensor, block); // Simplify the indices if possible for (const IterVar& iter : block->iter_vars) { - analyzer->Bind(iter->var, iter->dom.value()); + { + auto context = analyzer->GetConstraintContext(); + context->Append(sym::RangeConstraint(iter->var, iter->dom.value())); + analyzer->SyncConstraintContext(context); + } } original_indices.MutateByApply( [&analyzer](const PrimExpr& expr) { return SimplifyNonTrivialExpr(expr, analyzer.get()); }); @@ -2603,8 +2607,8 @@ StmtSRef ReIndex(ScheduleState self, const StmtSRef& block_sref, int tensor_inde // and get the insert location const StmtSRefObj* loop; for (loop = block_sref->parent; loop->parent != scope_sref.get();) { - const ForNode* outer = loop->parent->StmtAs(); - const ForNode* inner = loop->StmtAs(); + const ForObj* outer = loop->parent->StmtAs(); + const ForObj* inner = loop->StmtAs(); TVM_FFI_ICHECK(outer != nullptr && inner != nullptr); TVM_FFI_ICHECK(outer->body->size() == 1 && outer->body->seq[0].get() == inner); loop = loop->parent; diff --git a/src/s_tir/schedule/primitive/compute_at.cc b/src/s_tir/schedule/primitive/compute_at.cc index 0f801bee21c2..e50900ef55cc 100644 --- a/src/s_tir/schedule/primitive/compute_at.cc +++ b/src/s_tir/schedule/primitive/compute_at.cc @@ -44,7 +44,7 @@ class NotAllRequiredBlocksAreVisitedError : public ScheduleErrorContextObj { : mod_(mod), num_not_visited_(num_not_visited) { required_.reserve(required.size()); for (const StmtSRef& block_sref : required) { - const SBlockNode* block = TVM_SREF_TO_SBLOCK(block_sref); + const SBlockObj* block = TVM_SREF_TO_SBLOCK(block_sref); required_.push_back(ffi::GetRef(block)); } } @@ -88,8 +88,12 @@ class NotInSameScopeError : public ScheduleErrorContextObj { const StmtSRef& loop_sref, const StmtSRef& scope_root_sref, sym::AnalyzerObj* analyzer) { for (const StmtSRefObj* p = loop_sref.get();; p = p->parent) { - if (const ForNode* loop = p->StmtAs()) { - analyzer->Bind(loop->var, Range(loop->begin, loop->end)); + if (const ForObj* loop = p->StmtAs()) { + { + auto context = analyzer->GetConstraintContext(); + context->Append(sym::RangeConstraint(loop->var, Range(loop->begin, loop->end))); + analyzer->SyncConstraintContext(context); + } } else if (p != scope_root_sref.get()) { throw MakeScheduleError(self->mod, block_sref, loop_sref); } else { @@ -116,8 +120,8 @@ class NotInSameScopeError : public ScheduleErrorContextObj { explicit NotInSameScopeError(IRModule mod, const StmtSRef& block_sref, const StmtSRef& loop_sref) : mod_(mod), - block_(ffi::GetRef(block_sref->StmtAs())), - loop_(ffi::GetRef(loop_sref->StmtAs())) {} + block_(ffi::GetRef(block_sref->StmtAs())), + loop_(ffi::GetRef(loop_sref->StmtAs())) {} private: IRModule mod_; @@ -145,10 +149,11 @@ class NotInSameScopeError : public ScheduleErrorContextObj { * \throws ScheduleError if there is no such insertion point found */ template -int FindInsertionPoint( - const ScheduleState& self, const ffi::Array& subtrees, - const ffi::Array& producer_srefs, const ffi::Array& consumer_srefs, - std::unordered_map* block2realize, int index) { +int FindInsertionPoint(const ScheduleState& self, const ffi::Array& subtrees, + const ffi::Array& producer_srefs, + const ffi::Array& consumer_srefs, + std::unordered_map* block2realize, + int index) { ProducerConsumerSplit split = ProducerConsumerSplit::Find(self, subtrees, producer_srefs, consumer_srefs, block2realize); // Step 1. Check if all the producers are visited in the subtrees, if required to @@ -248,7 +253,7 @@ class ScopeReconstructor : public StmtExprMutator { using StmtExprMutator::Mutate; using StmtExprMutator::Mutate_; UnchangedOr Mutate(ffi::AnyView value, InplaceMode inplace_mode) override { - if (value.as()) return ffi::Unchanged(); + if (value.as()) return ffi::Unchanged(); return StmtExprMutator::Mutate(value, inplace_mode); } @@ -281,8 +286,13 @@ class ScopeReconstructor : public StmtExprMutator { loop_vars.push_back(var); loop_extents.push_back(analyzer->Simplify(iter_dom->extent)); iter_values.push_back(iter_dom->begin + var.as_or_throw()); - analyzer->Bind(var, Range::FromMinExtent(IntImm(var->ty.as_or_throw(), 0), - iter_dom->extent)); + { + auto context = analyzer->GetConstraintContext(); + context->Append(sym::RangeConstraint( + var, + Range::FromMinExtent(IntImm(var->ty.as_or_throw(), 0), iter_dom->extent))); + analyzer->SyncConstraintContext(context); + } } else { iter_values.push_back(iter_dom->begin); } @@ -315,26 +325,26 @@ class ScopeReconstructor : public StmtExprMutator { } ffi::Array subtrees = AsArray(loop_->body); subtrees.insert(subtrees.begin() + insert_position, std::move(new_subtree)); - ffi::ObjectPtr new_loop = ffi::make_object(*loop_.get()); + ffi::ObjectPtr new_loop = ffi::make_object(*loop_.get()); new_loop->body = SeqStmt(std::move(subtrees)); this->new_loop_ = For(std::move(new_loop)); } private: - UnchangedOr Mutate_(const SBlockNode* block, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const SBlockObj* block, InplaceMode inplace_mode) final { if (block != scope_root_.get()) { return ffi::GetRef(block); } if (block == rm_src_stmt_.get()) { - block = TVM_TYPE_AS(rm_tgt_stmt_, SBlockNode); + block = TVM_TYPE_AS(rm_tgt_stmt_, SBlockObj); } return StmtExprMutator::Mutate_(block, block->unique() ? inplace_mode : InplaceMode::kDisallow) .ValueOrUnchanged(ffi::GetRef(block)); } - UnchangedOr Mutate_(const ForNode* loop, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const ForObj* loop, InplaceMode inplace_mode) final { if (loop == rm_src_stmt_.get()) { - loop = TVM_TYPE_AS(rm_tgt_stmt_, ForNode); + loop = TVM_TYPE_AS(rm_tgt_stmt_, ForObj); } if (loop == loop_.get()) { return new_loop_.value(); @@ -376,7 +386,7 @@ void RelaxTensorRegions(const ffi::Map& binding, const ffi::Array& tensor_regions, const StmtSRef& relax_path_low_inclusive, const StmtSRef& relax_path_high_exclusive, - std::unordered_map>* relaxed) { + std::unordered_map>* relaxed) { runtime::MemoryScope global_scope = runtime::MemoryScope::Global(); // We cache the variable domains runtime::MemoryScope previous_scope = global_scope; @@ -436,7 +446,7 @@ std::pair SolveBlockVarDomain(const sym::IntSet& provid sym::IntSet var_dom{ffi::UnsafeInit{}}, var_bound{ffi::UnsafeInit{}}; ffi::Optional var; ffi::Optional factor; - if (const auto* mul = provided_min.as()) { + if (const auto* mul = provided_min.as()) { if (auto lhs = mul->a.as()) { var = lhs; factor = mul->b; @@ -455,7 +465,7 @@ std::pair SolveBlockVarDomain(const sym::IntSet& provid var_dom = sym::IntSet::Interval(required_min, required_max); var_bound = sym::IntSet::Interval(0, dim_max); } else { - if (const auto* div = provided_min.as()) { + if (const auto* div = provided_min.as()) { PrimExpr var_expr = div->a; PrimExpr fac = div->b; if (analyzer->CanProveGreaterEqual(fac, 1)) { @@ -472,7 +482,7 @@ std::pair SolveBlockVarDomain(const sym::IntSet& provid return SolveBlockVarDomain(new_provided, new_required, dim_max, analyzer); } } - } else if (const auto* mod = provided_min.as()) { + } else if (const auto* mod = provided_min.as()) { PrimExpr var_expr = mod->a; if (var_expr.as()) { // generally domain of (x % fac) enforce no constraints to domain of x @@ -508,9 +518,9 @@ std::pair SolveBlockVarDomain(const sym::IntSet& provid * \param analyzer The arithmetic analyzer * \param iter_doms The result iteration domains to be updated */ -void UpdateBlockVarDomainDimwise( - const VarNode* tensor, const NDIntSet& provided_region, const NDIntSet& required_region, - sym::AnalyzerObj* analyzer, std::unordered_map* iter_doms) { +void UpdateBlockVarDomainDimwise(const VarObj* tensor, const NDIntSet& provided_region, + const NDIntSet& required_region, sym::AnalyzerObj* analyzer, + std::unordered_map* iter_doms) { size_t ndim = ffi::GetRef(tensor).as_or_throw().shape().size(); for (size_t i = 0; i < ndim; ++i) { sym::IntSet provided = provided_region[i]; @@ -573,10 +583,10 @@ ffi::Map InverseAffineIterMap(const ffi::Array& iter_vars, +bool UpdateBlockVarDomainAffine(const VarObj* tensor, const ffi::Array& iter_vars, const NDIntSet& provided_region, const NDIntSet& required_region, sym::AnalyzerObj* analyzer, - std::unordered_map* iter_doms) { + std::unordered_map* iter_doms) { // we only support single point provided region now, which could cover most cases sym::Analyzer analyzer_ref = ffi::GetRef(analyzer); for (const auto& intset : provided_region) { @@ -628,19 +638,19 @@ bool UpdateBlockVarDomainAffine(const VarNode* tensor, const ffi::Array */ std::vector CalculateBlockVarDomain( const ffi::Array& iter_vars, - std::unordered_map> provided_regions, - std::unordered_map> required_regions, + std::unordered_map> provided_regions, + std::unordered_map> required_regions, 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; + std::unordered_map iter_doms; iter_doms.reserve(n_iters); for (const IterVar& iter_var : iter_vars) { iter_doms[iter_var->var.get()] = BlockVarDomainInfo(); } // Step 2. For each tensor, update the domain according to the provided and required regions for (const auto& kv : provided_regions) { - const VarNode* tensor = kv.first; + const VarObj* tensor = kv.first; const std::vector& many_provided_regions = kv.second; // Calculate `provided_region` and `required_region` auto it = required_regions.find(tensor); @@ -691,25 +701,24 @@ std::vector CalculateBlockVarDomain( */ template void CalculateProvidedRequiredRegions( - const SBlockNode* block, const StmtSRef& loop_sref, - std::unordered_map block2realize, + const SBlockObj* block, const StmtSRef& loop_sref, + std::unordered_map block2realize, ffi::Array producer_srefs, ffi::Array consumer_srefs, - std::unordered_map>* provided_regions, - std::unordered_map>* required_regions) { + std::unordered_map>* provided_regions, + std::unordered_map>* required_regions) { // Step 1. Calculate the region provided by a single execution instance of `block` const ffi::Array& provided_tensors = is_compute_at ? block->writes : block->reads; provided_regions->reserve(provided_tensors.size()); required_regions->reserve(provided_tensors.size()); for (const TensorRegion& provided_tensor_region : provided_tensors) { - const VarNode* tensor = - provided_tensor_region->source.as_or_throw().get(); + const VarObj* tensor = provided_tensor_region->source.as_or_throw().get(); const ffi::Array& region = provided_tensor_region->region; (*provided_regions)[tensor].push_back(support::NDIntSetFromRegion(region)); (*required_regions)[tensor].clear(); } // Step 2. Calculate the region required by dependent blocks under `loop` for (const StmtSRef& required_block_sref : is_compute_at ? consumer_srefs : producer_srefs) { - const SBlockNode* required_block = TVM_SREF_TO_SBLOCK(required_block_sref); + const SBlockObj* required_block = TVM_SREF_TO_SBLOCK(required_block_sref); TVM_FFI_ICHECK(block2realize.count(required_block)); RelaxTensorRegions( /*binding=*/GetBindings(ffi::GetRef(block2realize.at(required_block))), @@ -726,13 +735,13 @@ void ComputeAtOrReverseComputeAtImpl(ScheduleState self, const StmtSRef& block_s const StmtSRef& loop_sref, bool preserve_unit_loops, 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); + const SBlockObj* block = TVM_SREF_TO_SBLOCK(block_sref); + const ForObj* loop = TVM_SREF_TO_FOR(loop_sref); // Step 1. Bunch of checks // Check condition 1) : scope stage pipeline StmtSRef scope_root_sref = GetScopeRoot(self, block_sref, /*require_stage_pipeline=*/true); - SBlock scope_root = ffi::GetRef(scope_root_sref->StmtAs()); + SBlock scope_root = ffi::GetRef(scope_root_sref->StmtAs()); AddShapeVarBounds(self, scope_root_sref.get(), analyzer); SBlockScope scope = self->GetSBlockScope(scope_root_sref); ffi::Array producer_srefs = GetProducers(block_sref, scope); @@ -754,7 +763,7 @@ void ComputeAtOrReverseComputeAtImpl(ScheduleState self, const StmtSRef& block_s &reconstructor->rm_tgt_stmt_); // Step 3. Find the insertion point under `loop` // Check condition 5): all the required block are under the given loop - std::unordered_map block2realize; + std::unordered_map block2realize; block2realize.reserve(self->block_info.size()); int insert_position = FindInsertionPoint( /*self=*/self, @@ -767,8 +776,8 @@ void ComputeAtOrReverseComputeAtImpl(ScheduleState self, const StmtSRef& block_s // Here is the definition of `provide` and `require`: // - In compute-at, `provide` means `produce`, and `require` means `consume` // - In reverse-compute-at, `provide` means `consume`, and `require` means `produce` - std::unordered_map> provided_regions; - std::unordered_map> required_regions; + std::unordered_map> provided_regions; + std::unordered_map> required_regions; CalculateProvidedRequiredRegions( /*block=*/block, /*loop_sref=*/loop_sref, /*block2realize=*/std::move(block2realize), /*producer_srefs=*/std::move(producer_srefs), diff --git a/src/s_tir/schedule/primitive/compute_inline.cc b/src/s_tir/schedule/primitive/compute_inline.cc index c7f13ff23414..de64b9b79069 100644 --- a/src/s_tir/schedule/primitive/compute_inline.cc +++ b/src/s_tir/schedule/primitive/compute_inline.cc @@ -100,9 +100,9 @@ class NotSingleReadWriteTensor : public ScheduleErrorContextObj { const std::unordered_map, ffi::ObjectPtrHash, ffi::ObjectPtrEqual>& tensor_writers = self->block_info.at(scope_root_sref).scope->tensor_writers; - const VarNode* read_tensor = nullptr; + const VarObj* read_tensor = nullptr; for (const TensorRegion& read_region : block->reads) { - const VarNode* tensor = read_region->source.as_or_throw().get(); + const VarObj* tensor = read_region->source.as_or_throw().get(); if (tensor == read_tensor) { continue; } @@ -182,8 +182,8 @@ class NonSingleProducerError : public ScheduleErrorContextObj { */ static StmtSRef Check(const ScheduleState& self, const StmtSRef& consumer_block_sref, const StmtSRef& scope_root_sref) { - const SBlockNode* scope_block = TVM_SREF_TO_SBLOCK(scope_root_sref); - const SBlockNode* consumer_block = TVM_SREF_TO_SBLOCK(consumer_block_sref); + const SBlockObj* scope_block = TVM_SREF_TO_SBLOCK(scope_root_sref); + const SBlockObj* consumer_block = TVM_SREF_TO_SBLOCK(consumer_block_sref); TensorVar consumer_tensor = NotSingleReadWriteTensor::GetSingleRead( self, ffi::GetRef(consumer_block), scope_root_sref); class ProducerFinder : public StmtExprVisitor { @@ -191,7 +191,7 @@ class NonSingleProducerError : public ScheduleErrorContextObj { using StmtExprVisitor::Visit_; ffi::Optional Visit(ffi::AnyView value) override { - if (value.as()) return std::nullopt; + if (value.as()) return std::nullopt; return StmtExprVisitor::Visit(value); } @@ -210,7 +210,7 @@ class NonSingleProducerError : public ScheduleErrorContextObj { } private: - ffi::Optional Visit_(const SBlockNode* node) final { + ffi::Optional Visit_(const SBlockObj* node) final { producer_across_scope_.push_back({}); TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(node)); // not a leaf block @@ -257,7 +257,7 @@ class OpaqueAccessError : public ScheduleErrorContextObj { public: explicit OpaqueAccessError(IRModule mod, StmtSRef scope_root_sref) : mod_(mod), scope_root_(ffi::GetRef(TVM_SREF_TO_SBLOCK(scope_root_sref))) { - const SBlockNode* scope_root = TVM_SREF_TO_SBLOCK(scope_root_sref); + const SBlockObj* scope_root = TVM_SREF_TO_SBLOCK(scope_root_sref); this->scope_root_ = ffi::GetRef(scope_root); } @@ -319,14 +319,14 @@ class BaseInliner : public StmtExprMutator { const StmtSRef& scope_root_sref) : inlined_tensor_(inlined_tensor), inlined_store_(inlined_block->body->size() == 1 - ? inlined_block->body->seq[0].as() + ? inlined_block->body->seq[0].as() : nullptr), scope_root_sref_(scope_root_sref) { AddTensorsInBlockSignature(inlined_block.get()); } protected: - UnchangedOr Mutate_(const TensorRegionNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const TensorRegionObj* op, InplaceMode inplace_mode) final { if (!op->source.as()) { return StmtExprMutator::Mutate_(op, inplace_mode); } @@ -337,7 +337,7 @@ class BaseInliner : public StmtExprMutator { return node; } - UnchangedOr Mutate_(const TensorLoadNode* op, InplaceMode inplace_mode) override { + UnchangedOr Mutate_(const TensorLoadObj* op, InplaceMode inplace_mode) override { auto indices = Mutate(op->indices).as_or_throw>>(); TensorLoad node = ffi::GetRef(op); if (!indices.UnchangedOrSameAs(op->indices)) { @@ -346,7 +346,7 @@ class BaseInliner : public StmtExprMutator { return node; } - UnchangedOr Mutate_(const TensorStoreNode* op, InplaceMode inplace_mode) override { + UnchangedOr Mutate_(const TensorStoreObj* op, InplaceMode inplace_mode) override { auto value = Mutate(op->value); auto indices = Mutate(op->indices).as_or_throw>>(); TensorStore node = ffi::GetRef(op); @@ -358,16 +358,16 @@ class BaseInliner : public StmtExprMutator { return node; } - UnchangedOr Mutate_(const VarNode* var, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const VarObj* var, InplaceMode inplace_mode) final { if (def_region_kind() == kTVMFFIDefRegionKindNone) { CheckOpaqueAccess(var); } return StmtExprMutator::Mutate_(var, inplace_mode); } - UnchangedOr Mutate_(const CallNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const CallObj* op, InplaceMode inplace_mode) final { if (op->op.same_as(tirx::address_of_op())) { - if (const auto* load = op->args[0].as()) { + if (const auto* load = op->args[0].as()) { CheckOpaqueAccess(load->source.as_or_throw().get()); // An address operand must remain a TensorLoad, not the inlined value. auto operand = BaseInliner::Mutate_(load, inplace_mode); @@ -380,21 +380,21 @@ class BaseInliner : public StmtExprMutator { return StmtExprMutator::Mutate_(op, inplace_mode); } - UnchangedOr Mutate_(const ForNode* loop, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const ForObj* loop, InplaceMode inplace_mode) final { if (src_stmt.get() == loop) { - loop = tgt_stmt.as(); + loop = tgt_stmt.as(); TVM_FFI_ICHECK(loop != nullptr); } return StmtExprMutator::Mutate_(loop, loop->unique() ? inplace_mode : InplaceMode::kDisallow) .ValueOrUnchanged(ffi::GetRef(loop)); } - UnchangedOr Mutate_(const SBlockNode* block, InplaceMode inplace_mode) { + UnchangedOr Mutate_(const SBlockObj* block, InplaceMode inplace_mode) { CheckMatchTensorRegion(block); AddTensorsInBlockSignature(block); SBlock src_block = ffi::GetRef(block); if (src_block.same_as(src_stmt)) { - block = tgt_stmt.as(); + block = tgt_stmt.as(); TVM_FFI_ICHECK(block != nullptr); } SBlock tgt_block = @@ -413,7 +413,7 @@ class BaseInliner : public StmtExprMutator { * which is used for auto-completion of a block's read/write region * \param block The block whose signature to be added */ - void AddTensorsInBlockSignature(const SBlockNode* block) { + void AddTensorsInBlockSignature(const SBlockObj* block) { for (const TensorRegion& tensor_region : block->reads) { const TensorVar& tensor = tensor_region->source.as_or_throw(); tensor_var_map_.Set(tensor, tensor); @@ -463,7 +463,7 @@ class BaseInliner : public StmtExprMutator { writes = inspected[1]; } // Step 3. Assemble the result - SBlockNode* n = block.CopyOnWrite(); + SBlockObj* n = block.CopyOnWrite(); n->reads = std::move(reads); n->writes = std::move(writes); n->alloc_tensors = std::move(alloc_tensors); @@ -475,7 +475,7 @@ class BaseInliner : public StmtExprMutator { * This method checks if a tensor var belongs to the tensor * \param tensor_var The tensor var to be checked */ - void CheckOpaqueAccess(const VarNode* tensor_var) { + void CheckOpaqueAccess(const VarObj* tensor_var) { if (inlined_tensor_.get() == tensor_var) { this->has_opaque_access = true; } @@ -486,7 +486,7 @@ class BaseInliner : public StmtExprMutator { * This method checks if a block has the disallowed behavior of tensor region match. * \param block The block to be checked */ - void CheckMatchTensorRegion(const SBlockNode* block) { + void CheckMatchTensorRegion(const SBlockObj* block) { for (const MatchTensorRegion& match_tensor_region : block->match_tensors) { const TensorVar& matched = match_tensor_region->source->source.as_or_throw(); @@ -500,7 +500,7 @@ class BaseInliner : public StmtExprMutator { /*! \brief The tensor to be inlined */ TensorVar inlined_tensor_; /*! \brief The body of the block to be inlined */ - const TensorStoreNode* inlined_store_{nullptr}; + const TensorStoreObj* inlined_store_{nullptr}; /*! \brief The scope root */ StmtSRef scope_root_sref_; /*! \brief Maps a tensor's data field to itself */ @@ -614,7 +614,7 @@ class ComputeInliner : public BaseInliner { } private: - UnchangedOr Mutate_(const TensorLoadNode* _load, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const TensorLoadObj* _load, InplaceMode inplace_mode) final { TensorLoad load = BaseInliner::Mutate_(_load, inplace_mode) .ValueOrUnchanged(ffi::GetRef(_load)) .as_or_throw(); @@ -666,7 +666,7 @@ class ReverseComputeInliner : public BaseInliner { } private: - UnchangedOr Mutate_(const TensorLoadNode* _load, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const TensorLoadObj* _load, InplaceMode inplace_mode) final { TensorLoad load = StmtExprMutator::Mutate_(_load, inplace_mode) .ValueOrUnchanged(ffi::GetRef(_load)) .as_or_throw(); @@ -691,7 +691,7 @@ class ReverseComputeInliner : public BaseInliner { } private: - UnchangedOr Mutate_(const TensorLoadNode* _load, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const TensorLoadObj* _load, InplaceMode inplace_mode) final { TensorLoad load = StmtExprMutator::Mutate_(_load, inplace_mode) .ValueOrUnchanged(ffi::GetRef(_load)) .as_or_throw(); @@ -706,7 +706,7 @@ class ReverseComputeInliner : public BaseInliner { ReverseComputeInliner* self_; }; - explicit ReverseComputeInliner(const TensorVar& inlined_tensor, const SBlockNode* producer_block, + explicit ReverseComputeInliner(const TensorVar& inlined_tensor, const SBlockObj* producer_block, const SBlockRealize& consumer_block_realize, const StmtSRef& scope_root_sref, const IRModule& mod) : BaseInliner(inlined_tensor, consumer_block_realize->block, scope_root_sref), @@ -733,7 +733,7 @@ class ReverseComputeInliner : public BaseInliner { // Failure: block body is not TensorStore return false; } - std::vector loads = ExtractTensorLoad(inlined_tensor_, inlined_store_); + std::vector loads = ExtractTensorLoad(inlined_tensor_, inlined_store_); if (loads.size() == 0) { // Failure: no TensorLoad from the `inlined_tensor_` return false; @@ -750,7 +750,7 @@ class ReverseComputeInliner : public BaseInliner { } } - for (const TensorLoadNode* load : loads) { + for (const TensorLoadObj* load : loads) { if (!UpdateAndCheckIndexExprs(load->indices)) { return false; } @@ -772,27 +772,27 @@ class ReverseComputeInliner : public BaseInliner { if (producer_block_->body->size() != 1) { return false; } - const TensorStoreNode* producer_store = nullptr; + const TensorStoreObj* producer_store = nullptr; const Stmt& producer_body = producer_block_->body->seq[0]; - if (const auto* producer_if = producer_body.as()) { + if (const auto* producer_if = producer_body.as()) { if (producer_if->else_case.has_value()) { return false; } producer_store = producer_if->then_case->size() == 1 - ? producer_if->then_case->seq[0].as() + ? producer_if->then_case->seq[0].as() : nullptr; } else { - producer_store = producer_body.as(); + producer_store = producer_body.as(); if (producer_block_->annotations.count(tvm::s_tir::attr::kAutoCopy.name) != 0) { - const ForNode* producer_inner_loop = producer_body.as(); + const ForObj* producer_inner_loop = producer_body.as(); if (producer_inner_loop == nullptr) return false; for (;;) { if (producer_inner_loop->body->size() != 1) return false; const Stmt& child = producer_inner_loop->body->seq[0]; - if (const auto* inner_loop = child.as()) { + if (const auto* inner_loop = child.as()) { producer_inner_loop = inner_loop; } else { - producer_store = child.as(); + producer_store = child.as(); break; } } @@ -821,21 +821,29 @@ class ReverseComputeInliner : public BaseInliner { const IterVar& iter = producer_block->iter_vars[i]; const PrimExpr& binding = producer_block_realize->iter_values[i]; subst_map.Set(iter->var, binding); - analyzer_->Bind(iter->var, iter->dom.value()); + { + auto context = analyzer_->GetConstraintContext(); + context->Append(sym::RangeConstraint(iter->var, iter->dom.value())); + analyzer_->SyncConstraintContext(context); + } } if (producer_block->annotations.count(tvm::s_tir::attr::kAutoCopy.name) != 0) { - auto bind = [&](const ForNode* loop) { - analyzer_->Bind(loop->var, - Range::FromMinExtent( - IntImm(Range::GetExtentFromBeginEnd(loop->begin, loop->end).ty(), 0), - Range::GetExtentFromBeginEnd(loop->begin, loop->end))); + auto bind = [&](const ForObj* loop) { + { + auto context = analyzer_->GetConstraintContext(); + context->Append(sym::RangeConstraint( + loop->var, Range::FromMinExtent( + IntImm(Range::GetExtentFromBeginEnd(loop->begin, loop->end).ty(), 0), + Range::GetExtentFromBeginEnd(loop->begin, loop->end)))); + analyzer_->SyncConstraintContext(context); + } }; TVM_FFI_ICHECK_EQ(producer_block->body->size(), 1); - const ForNode* producer_inner_loop = producer_block->body->seq[0].as(); + const ForObj* producer_inner_loop = producer_block->body->seq[0].as(); while (producer_inner_loop != nullptr) { bind(producer_inner_loop); TVM_FFI_ICHECK_EQ(producer_inner_loop->body->size(), 1); - producer_inner_loop = producer_inner_loop->body->seq[0].as(); + producer_inner_loop = producer_inner_loop->body->seq[0].as(); } } // Substitute the consumer block iters with the corresponding iters in the producer blocks @@ -848,7 +856,7 @@ class ReverseComputeInliner : public BaseInliner { return producer_block_realize; } if (const auto* if_ = producer_block->body->size() == 1 - ? producer_block->body->seq[0].as() + ? producer_block->body->seq[0].as() : nullptr) { if (!if_->else_case.has_value()) { PrimExpr if_predicate = analyzer_->Simplify(if_->condition); @@ -871,7 +879,7 @@ class ReverseComputeInliner : public BaseInliner { return ffi::GetRef(n); } - UnchangedOr Mutate_(const SBlockRealizeNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const SBlockRealizeObj* op, InplaceMode inplace_mode) final { SBlock src_block = op->block; SBlockRealize tgt_block_realize = StmtExprMutator::Mutate_(op, inplace_mode) .ValueOrUnchanged(ffi::GetRef(op)) @@ -883,7 +891,7 @@ class ReverseComputeInliner : public BaseInliner { return tgt_block_realize; } - UnchangedOr Mutate_(const TensorStoreNode* _store, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const TensorStoreObj* _store, InplaceMode inplace_mode) final { TensorStore store = BaseInliner::Mutate_(_store, inplace_mode) .ValueOrUnchanged(ffi::GetRef(_store)) .as_or_throw(); @@ -948,19 +956,19 @@ class ReverseComputeInliner : public BaseInliner { * \param from The TensorStore statement to be extracted from * \return A list of `TensorLoad` expressions */ - static std::vector ExtractTensorLoad(const TensorVar& tensor, - const TensorStoreNode* from) { + static std::vector ExtractTensorLoad(const TensorVar& tensor, + const TensorStoreObj* from) { struct Extractor : public StmtExprVisitor { using StmtExprVisitor::Visit_; - ffi::Optional Visit_(const TensorLoadNode* load) final { + ffi::Optional Visit_(const TensorLoadObj* load) final { if (load->source.as_or_throw().get() == tensor) { result.push_back(load); } return StmtExprVisitor::Visit_(load); } - const VarNode* tensor; - std::vector result; + const VarObj* tensor; + std::vector result; }; auto extractor = ffi::make_object(); extractor->tensor = tensor.get(); @@ -996,9 +1004,9 @@ class ReverseComputeInliner : public BaseInliner { /*! \brief The IterMap representing the indices of the consumer's TensorLoad */ ffi::Array tensor_load_iter_map_; /*! \brief The producer block */ - const SBlockNode* producer_block_{nullptr}; + const SBlockObj* producer_block_{nullptr}; /* \brief The consumer block */ - const SBlockNode* consumer_block_{nullptr}; + const SBlockObj* consumer_block_{nullptr}; /*! \brief The predicate to ensure the consumer block iters are in-bound. It will be inserted * as the predicate of the producer block after inlining. */ @@ -1009,7 +1017,7 @@ class ReverseComputeInliner : public BaseInliner { void ComputeInlineImpl(ScheduleState self, const StmtSRef& producer_block_sref, bool check_only = false) { - const SBlockNode* _producer_block = TVM_SREF_TO_SBLOCK(producer_block_sref); + const SBlockObj* _producer_block = TVM_SREF_TO_SBLOCK(producer_block_sref); SBlock producer_block = ffi::GetRef(_producer_block); HasInitBlock::Check(self->mod, producer_block); TensorVar inlined_tensor = NotSingleReadWriteTensor::GetSingleWrite(self, producer_block); @@ -1055,7 +1063,7 @@ bool CanComputeInline(const ScheduleState& self, const StmtSRef& producer_block_ void ReverseComputeInlineImpl(ScheduleState self, const StmtSRef& consumer_block_sref, bool check_only = false) { - const SBlockNode* _consumer_block = TVM_SREF_TO_SBLOCK(consumer_block_sref); + const SBlockObj* _consumer_block = TVM_SREF_TO_SBLOCK(consumer_block_sref); SBlock consumer_block = ffi::GetRef(_consumer_block); SBlockRealize consumer_block_realize = GetSBlockRealize(self, consumer_block_sref); HasInitBlock::Check(self->mod, consumer_block); @@ -1073,7 +1081,7 @@ void ReverseComputeInlineImpl(ScheduleState self, const StmtSRef& consumer_block CheckNotOutputBlock(self, producer_block_sref, scope_root_sref); // Step 4. Analyze the block body auto inliner = ffi::make_object( - inlined_tensor, producer_block_sref->StmtAs(), consumer_block_realize, + inlined_tensor, producer_block_sref->StmtAs(), consumer_block_realize, scope_root_sref, self->mod); if (!inliner->BodyPatternAllowInline(consumer_block_realize)) { throw MakeScheduleError(true, self->mod, consumer_block); @@ -1125,7 +1133,7 @@ class ReductionEpilogueFuser : public BaseInliner { using BaseInliner::Mutate_; explicit ReductionEpilogueFuser(const TensorVar& reduction_tensor, - const SBlockNode* reduction_block, + const SBlockObj* reduction_block, const SBlockRealize& epilogue_block_realize, const StmtSRef& scope_root_sref) : BaseInliner(reduction_tensor, epilogue_block_realize->block, scope_root_sref), @@ -1137,7 +1145,7 @@ class ReductionEpilogueFuser : public BaseInliner { } // Override CheckOpaqueAccess to allow multiple tensor reads - void CheckOpaqueAccess(const VarNode* tensor_var) { + void CheckOpaqueAccess(const VarObj* tensor_var) { // For epilogue fusion, we allow multiple tensor reads (temp + bias) // So we don't check for opaque access // BaseInliner::CheckOpaqueAccess(tensor_var); // Don't call base class @@ -1146,19 +1154,19 @@ class ReductionEpilogueFuser : public BaseInliner { bool BodyPatternAllowFusion(const SBlockRealize& epilogue_block_realize); // Step 2: Create single fused reduction block - SBlock CreateFusedReductionBlock(const SBlockNode* reduction_block, - const SBlockRealizeNode* reduction_realize); + SBlock CreateFusedReductionBlock(const SBlockObj* reduction_block, + const SBlockRealizeObj* reduction_realize); private: - bool IsReductionBlock(const SBlockNode* block); + bool IsReductionBlock(const SBlockObj* block); void ExtractEpilogueInfo(); // Helper function to extract TensorLoad nodes from TensorStore - static std::vector ExtractTensorLoad(const TensorVar& tensor, - const TensorStoreNode* from) { + static std::vector ExtractTensorLoad(const TensorVar& tensor, + const TensorStoreObj* from) { struct Extractor : public StmtExprVisitor { using StmtExprVisitor::Visit_; - ffi::Optional Visit_(const TensorLoadNode* load) final { + ffi::Optional Visit_(const TensorLoadObj* load) final { if (load->source.as_or_throw().same_as(tensor)) { result.push_back(load); } @@ -1166,7 +1174,7 @@ class ReductionEpilogueFuser : public BaseInliner { return StmtExprVisitor::Visit_(load); } TensorVar tensor{ffi::UnsafeInit{}}; - std::vector result; + std::vector result; }; auto extractor = ffi::make_object(); extractor->tensor = tensor; @@ -1179,12 +1187,12 @@ class ReductionEpilogueFuser : public BaseInliner { return std::move(extractor->result); } - const SBlockNode* reduction_block_; - const SBlockNode* epilogue_block_; + const SBlockObj* reduction_block_; + const SBlockObj* epilogue_block_; // Generalized approach: store the entire epilogue expression PrimExpr epilogue_expression_{ ffi::UnsafeInit{}}; // The entire epilogue expression (e.g., temp + C, max(temp + C, 0)) - const TensorLoadNode* reduction_tensor_load_{ + const TensorLoadObj* reduction_tensor_load_{ nullptr}; // The reduction tensor load in epilogue expression TensorVar epilogue_output_tensor_{ffi::UnsafeInit{}}; // Output tensor D ffi::Array epilogue_output_indices_; // Indices of D[vi, vj] @@ -1207,7 +1215,7 @@ bool ReductionEpilogueFuser::BodyPatternAllowFusion(const SBlockRealize& epilogu } // 3. Check if epilogue reads from reduction tensor - std::vector loads = ExtractTensorLoad(inlined_tensor_, inlined_store_); + std::vector loads = ExtractTensorLoad(inlined_tensor_, inlined_store_); if (loads.size() == 0) { // Failure: no TensorLoad from the reduction tensor return false; @@ -1255,7 +1263,7 @@ bool ReductionEpilogueFuser::BodyPatternAllowFusion(const SBlockRealize& epilogu } private: - ffi::Optional Visit_(const TensorLoadNode* op) final { + ffi::Optional Visit_(const TensorLoadObj* op) final { if (op->source.as_or_throw().same_as(tensor_)) { found_ = true; return std::nullopt; @@ -1270,7 +1278,7 @@ bool ReductionEpilogueFuser::BodyPatternAllowFusion(const SBlockRealize& epilogu // Helper to check if a subtree contains a load from the reduction tensor bool ContainsTarget(const PrimExpr& expr) { return finder_->Find(expr); } - ffi::Optional Visit_(const MulNode* op) final { + ffi::Optional Visit_(const MulObj* op) final { if (has_scaling_) return std::nullopt; // If either operand subtree contains the reduction tensor load, // we treat this as invalid scaling of the reduction result. @@ -1281,7 +1289,7 @@ bool ReductionEpilogueFuser::BodyPatternAllowFusion(const SBlockRealize& epilogu return StmtExprVisitor::Visit_(op); } - ffi::Optional Visit_(const DivNode* op) final { + ffi::Optional Visit_(const DivObj* op) final { if (has_scaling_) return std::nullopt; if (ContainsTarget(op->a) || ContainsTarget(op->b)) { has_scaling_ = true; @@ -1290,7 +1298,7 @@ bool ReductionEpilogueFuser::BodyPatternAllowFusion(const SBlockRealize& epilogu return StmtExprVisitor::Visit_(op); } - ffi::Optional Visit_(const ModNode* op) final { + ffi::Optional Visit_(const ModObj* op) final { if (has_scaling_) return std::nullopt; if (ContainsTarget(op->a) || ContainsTarget(op->b)) { has_scaling_ = true; @@ -1323,7 +1331,7 @@ bool ReductionEpilogueFuser::BodyPatternAllowFusion(const SBlockRealize& epilogu return true; } -bool ReductionEpilogueFuser::IsReductionBlock(const SBlockNode* block) { +bool ReductionEpilogueFuser::IsReductionBlock(const SBlockObj* block) { // Check if block has reduction iter vars for (const IterVar& iter : block->iter_vars) { if (iter->iter_type == kCommReduce) { @@ -1351,14 +1359,14 @@ void ReductionEpilogueFuser::ExtractEpilogueInfo() { struct TensorExtractor : public StmtExprVisitor { using StmtExprVisitor::Visit_; - ffi::Optional Visit_(const TensorLoadNode* load) final { + ffi::Optional Visit_(const TensorLoadObj* load) final { if (!load->source.as_or_throw().same_as(reduction_tensor)) { other_tensors.insert(load->source.as_or_throw().get()); } return StmtExprVisitor::Visit_(load); } TensorVar reduction_tensor{ffi::UnsafeInit{}}; - std::unordered_set other_tensors; + std::unordered_set other_tensors; }; auto extractor = ffi::make_object(); extractor->reduction_tensor = inlined_tensor_; @@ -1367,7 +1375,7 @@ void ReductionEpilogueFuser::ExtractEpilogueInfo() { // Extract the first non-reduction tensor and its region // In most cases, there's one additional tensor (e.g., bias tensor) if (!extractor->other_tensors.empty()) { - const VarNode* first_tensor = *extractor->other_tensors.begin(); + const VarObj* first_tensor = *extractor->other_tensors.begin(); epilogue_addend_tensor_ = ffi::GetRef(first_tensor).as_or_throw(); // Find the read region from epilogue block reads for (const TensorRegion& read : epilogue_block_->reads) { @@ -1380,8 +1388,8 @@ void ReductionEpilogueFuser::ExtractEpilogueInfo() { } SBlock ReductionEpilogueFuser::CreateFusedReductionBlock( - const SBlockNode* reduction_block, const SBlockRealizeNode* reduction_realize) { - ffi::ObjectPtr new_block = ffi::make_object(*reduction_block); + const SBlockObj* reduction_block, const SBlockRealizeObj* reduction_realize) { + ffi::ObjectPtr new_block = ffi::make_object(*reduction_block); // 1. Map epilogue block vars to reduction block vars std::vector reduction_data_vars; @@ -1422,7 +1430,7 @@ SBlock ReductionEpilogueFuser::CreateFusedReductionBlock( InitSubstituter(const TensorVar& target_tensor, PrimExpr identity_elem) : target_tensor_(target_tensor), identity_elem_(identity_elem) {} - UnchangedOr Mutate_(const TensorLoadNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const TensorLoadObj* op, InplaceMode inplace_mode) final { TensorLoad load = StmtExprMutator::Mutate_(op, inplace_mode) .ValueOrUnchanged(ffi::GetRef(op)) .as_or_throw(); @@ -1478,7 +1486,7 @@ SBlock ReductionEpilogueFuser::CreateFusedReductionBlock( epilogue_expression_(epilogue_expr), var_map_(var_map) {} - UnchangedOr Mutate_(const TensorStoreNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const TensorStoreObj* op, InplaceMode inplace_mode) final { TensorStore store = StmtExprMutator::Mutate_(op, inplace_mode) .ValueOrUnchanged(ffi::GetRef(op)) .as_or_throw(); @@ -1493,7 +1501,7 @@ SBlock ReductionEpilogueFuser::CreateFusedReductionBlock( ReductionUpdateReplacer(const TensorVar& old_tensor, const TensorVar& new_tensor) : old_tensor_(old_tensor), new_tensor_(new_tensor) {} - UnchangedOr Mutate_(const TensorLoadNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const TensorLoadObj* op, InplaceMode inplace_mode) final { TensorLoad load = StmtExprMutator::Mutate_(op, inplace_mode) .ValueOrUnchanged(ffi::GetRef(op)) .as_or_throw(); @@ -1529,19 +1537,19 @@ SBlock ReductionEpilogueFuser::CreateFusedReductionBlock( replacement_(replacement), found_target_load_(false) {} - UnchangedOr Mutate_(const TensorLoadNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const TensorLoadObj* op, InplaceMode inplace_mode) final { TensorLoad load = StmtExprMutator::Mutate_(op, inplace_mode) .ValueOrUnchanged(ffi::GetRef(op)) .as_or_throw(); if (load->source.as_or_throw().same_as(target_tensor_)) { found_target_load_ = true; - // Check if parent is Add (will be checked in Dispatch_(const AddNode*)) + // Check if parent is Add (will be checked in Dispatch_(const AddObj*)) return replacement_; } return load; } - UnchangedOr Mutate_(const AddNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const AddObj* op, InplaceMode inplace_mode) final { // Visit children first to see if we find the target tensor load bool found_before = found_target_load_; found_target_load_ = false; @@ -1560,7 +1568,7 @@ SBlock ReductionEpilogueFuser::CreateFusedReductionBlock( bool other_is_reduction = false; if (found_in_a) { // Check if b is from reduction tensor - if (const auto* load_b = b.as()) { + if (const auto* load_b = b.as()) { other_is_reduction = load_b->source.as_or_throw().same_as(reduction_tensor_); } @@ -1570,7 +1578,7 @@ SBlock ReductionEpilogueFuser::CreateFusedReductionBlock( } } else { // found_in_b // Check if a is from reduction tensor - if (const auto* load_a = a.as()) { + if (const auto* load_a = a.as()) { other_is_reduction = load_a->source.as_or_throw().same_as(reduction_tensor_); } @@ -1615,7 +1623,7 @@ SBlock ReductionEpilogueFuser::CreateFusedReductionBlock( return store; } - UnchangedOr Mutate_(const TensorLoadNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const TensorLoadObj* op, InplaceMode inplace_mode) final { TensorLoad load = StmtExprMutator::Mutate_(op, inplace_mode) .ValueOrUnchanged(ffi::GetRef(op)) .as_or_throw(); @@ -1664,7 +1672,7 @@ SBlock ReductionEpilogueFuser::CreateFusedReductionBlock( // 5. Update read regions: add all tensors from epilogue expression (except reduction tensor) ffi::Array new_reads; - std::unordered_set read_tensors; + std::unordered_set read_tensors; // Add all non-reduction tensors from epilogue expression for (const TensorRegion& read : epilogue_block_->reads) { @@ -1708,7 +1716,7 @@ static bool CheckTensorStillUsed(const SBlock& scope_root, const TensorVar& tens using StmtExprVisitor::Visit_; ffi::Optional Visit(ffi::AnyView value) override { - if (value.as()) return std::nullopt; + if (value.as()) return std::nullopt; return StmtExprVisitor::Visit(value); } @@ -1721,14 +1729,14 @@ static bool CheckTensorStillUsed(const SBlock& scope_root, const TensorVar& tens } private: - ffi::Optional Visit_(const SBlockRealizeNode* op) final { + ffi::Optional Visit_(const SBlockRealizeObj* op) final { if (found_usage_) return std::nullopt; if (!op || !op->block.defined()) { return StmtExprVisitor::Visit_(op); } - const SBlockNode* block = op->block.get(); + const SBlockObj* block = op->block.get(); if (!block) { return StmtExprVisitor::Visit_(op); } @@ -1753,7 +1761,7 @@ static bool CheckTensorStillUsed(const SBlock& scope_root, const TensorVar& tens return StmtExprVisitor::Visit_(op); } - ffi::Optional Visit_(const SBlockNode* op) final { + ffi::Optional Visit_(const SBlockObj* op) final { if (found_usage_) return std::nullopt; if (!op) return std::nullopt; @@ -1788,7 +1796,7 @@ class SingleBlockFusionReplacer : public StmtExprMutator { using StmtExprMutator::Mutate; using StmtExprMutator::Mutate_; UnchangedOr Mutate(ffi::AnyView value, InplaceMode inplace_mode) override { - if (value.as()) return ffi::Unchanged(); + if (value.as()) return ffi::Unchanged(); return StmtExprMutator::Mutate(value, inplace_mode); } @@ -1806,7 +1814,7 @@ class SingleBlockFusionReplacer : public StmtExprMutator { // Remove intermediate temp tensor only if it's not used by other blocks if (!tensor_still_used) { - SBlockNode* p = result.CopyOnWrite(); + SBlockObj* p = result.CopyOnWrite(); ffi::Array new_alloc_tensors; for (const TensorVar& tensor : p->alloc_tensors) { if (!tensor.same_as(reduction_tensor)) { @@ -1827,7 +1835,7 @@ class SingleBlockFusionReplacer : public StmtExprMutator { reduction_tensor_(std::move(reduction_tensor)) {} private: - UnchangedOr Mutate_(const ForNode* loop, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const ForObj* loop, InplaceMode inplace_mode) final { SeqStmt mutated_body = Mutate(loop->body, inplace_mode).ValueOrUnchanged(loop->body); if (mutated_body->seq.empty()) { return mutated_body; @@ -1837,10 +1845,10 @@ class SingleBlockFusionReplacer : public StmtExprMutator { loop->annotations, loop->loc); } - UnchangedOr Mutate_(const SBlockRealizeNode* realize, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const SBlockRealizeObj* realize, InplaceMode inplace_mode) final { if (realize->block.same_as(old_reduction_block_)) { // Replace reduction block with new fused block - ffi::ObjectPtr new_realize = ffi::make_object(*realize); + ffi::ObjectPtr new_realize = ffi::make_object(*realize); new_realize->block = new_fused_block_; return SBlockRealize(new_realize); } else if (realize->block.same_as(old_epilogue_block_)) { @@ -1850,7 +1858,7 @@ class SingleBlockFusionReplacer : public StmtExprMutator { return StmtExprMutator::Mutate_(realize, inplace_mode); } - UnchangedOr Mutate_(const SeqStmtNode* seq, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const SeqStmtObj* seq, InplaceMode inplace_mode) final { ffi::Array new_stmts; for (size_t i = 0; i < seq->seq.size(); ++i) { Stmt stmt = seq->seq[i]; @@ -1868,8 +1876,8 @@ class SingleBlockFusionReplacer : public StmtExprMutator { void FuseReductionEpilogueImpl(ScheduleState self, const StmtSRef& reduction_block_sref, const StmtSRef& epilogue_block_sref, bool check_only = false) { - const SBlockNode* _reduction_block = TVM_SREF_TO_SBLOCK(reduction_block_sref); - const SBlockNode* _epilogue_block = TVM_SREF_TO_SBLOCK(epilogue_block_sref); + const SBlockObj* _reduction_block = TVM_SREF_TO_SBLOCK(reduction_block_sref); + const SBlockObj* _epilogue_block = TVM_SREF_TO_SBLOCK(epilogue_block_sref); SBlock reduction_block = ffi::GetRef(_reduction_block); SBlock epilogue_block = ffi::GetRef(_epilogue_block); @@ -1903,7 +1911,7 @@ void FuseReductionEpilogueImpl(ScheduleState self, const StmtSRef& reduction_blo SBlock fused_block = fuser->CreateFusedReductionBlock(_reduction_block, reduction_realize.get()); // Step 6. Transform and replace IR - const SBlockNode* old_scope_root = TVM_SREF_TO_SBLOCK(scope_root_sref); + const SBlockObj* old_scope_root = TVM_SREF_TO_SBLOCK(scope_root_sref); SBlock new_scope_root = SingleBlockFusionReplacer::Replace(ffi::GetRef(old_scope_root), fused_block, diff --git a/src/s_tir/schedule/primitive/decompose_padding.cc b/src/s_tir/schedule/primitive/decompose_padding.cc index 0305648bccaa..3d1bc316c7f8 100644 --- a/src/s_tir/schedule/primitive/decompose_padding.cc +++ b/src/s_tir/schedule/primitive/decompose_padding.cc @@ -71,7 +71,7 @@ class PaddingPatternMatchError : public ScheduleErrorContextObj { */ class PaddingInfoAnalyzer { public: - static PaddingSBlockInfo CheckAndGetPaddingInfo(IRModule mod, const SBlockRealizeNode* realize, + static PaddingSBlockInfo CheckAndGetPaddingInfo(IRModule mod, const SBlockRealizeObj* realize, const ffi::Map& dom_map, sym::AnalyzerObj* analyzer) { PaddingInfoAnalyzer padding_analyzer(analyzer); @@ -86,22 +86,22 @@ class PaddingInfoAnalyzer { explicit PaddingInfoAnalyzer(sym::AnalyzerObj* analyzer) : analyzer_(analyzer) {} /*! \brief Detect padding pattern and update result. */ - bool MatchPadding(const SBlockRealizeNode* realize, const ffi::Map& dom_map) { + bool MatchPadding(const SBlockRealizeObj* realize, const ffi::Map& dom_map) { // Step 1. Check match padding computation pattern. // A[...] = T.if_then_else(predicate, B[...], imm) SBlock block = realize->block; - std::unordered_map iter_values; + std::unordered_map iter_values; for (size_t i = 0; i < realize->iter_values.size(); ++i) { Var block_var = block->iter_vars[i]->var; iter_values.insert_or_assign(block_var.get(), realize->iter_values[i]); } - const TensorStoreNode* store = - block->body->size() == 1 ? block->body->seq[0].as() : nullptr; + const TensorStoreObj* store = + block->body->size() == 1 ? block->body->seq[0].as() : nullptr; if (!store) { SetError("Block body expect a TensorStore to the write tensor"); return false; } - const CallNode* if_then_else = store->value.as(); + const CallObj* if_then_else = store->value.as(); if (!if_then_else || !if_then_else->op.same_as(prim::if_then_else_op())) { SetError("Value of TensorStore expect to be constrained by a padding predicate"); return false; @@ -152,11 +152,11 @@ class PaddingInfoAnalyzer { PrimExpr RewritePredicate(const PrimExpr& predicate) { PrimExpr res = IntImm::Bool(true); std::function update = [&res, &update](PrimExpr e) { - if (const auto* and_node = e.as()) { + if (const auto* and_node = e.as()) { update(and_node->a); update(and_node->b); } else { - if (const CallNode* call = e.as()) { + if (const CallObj* call = e.as()) { if (call->op.same_as(prim::likely_op())) { e = call->args[0].as_or_throw(); } @@ -207,7 +207,7 @@ class PaddingInfoAnalyzer { }; /*! \brief Create block to fill constant pad values into full region */ -static std::pair CreateConstBlock(const SBlockRealizeNode* realize, +static std::pair CreateConstBlock(const SBlockRealizeObj* realize, const PaddingSBlockInfo& info, const ffi::Array& loops, const Stmt& highest_pos_inclusive, @@ -285,7 +285,7 @@ static std::pair CreateConstBlock(const SBlockRealizeNode* } /*! \brief Create block to fill in-bound region values. */ -static std::pair CreateInBoundBlock(const SBlockRealizeNode* realize, +static std::pair CreateInBoundBlock(const SBlockRealizeObj* realize, const PaddingSBlockInfo& info, const ffi::Array& loops, @@ -323,7 +323,20 @@ static std::pair CreateInBoundBlock(const SBlockRealizeNode new_loop_ranges.Set(loop_var, new_range); new_iter_binding.push_back(realize->iter_values[i]); repl_dict.Set(loop_var, loop_var.as_or_throw() + info.in_bound_region[i]->begin); - analyzer->Bind(loop_var, new_range, /*allow_override=*/true); + { + sym::ConstraintContext context; + analyzer->GetConstraintContext()->ForEach([&](const sym::Constraint& premise) { + if (auto binding = premise.as(); + binding && binding->var.same_as(loop_var)) + return; + if (auto range = premise.as(); + range && range->var.same_as(loop_var)) + return; + context->Append(premise); + }); + context->Append(sym::RangeConstraint(loop_var, new_range)); + analyzer->SyncConstraintContext(context); + } } else { new_iter_binding.push_back( analyzer->Simplify(realize->iter_values[i] - info.in_bound_region[i]->begin)); @@ -393,7 +406,7 @@ class DecomposePaddingBlockReplacer : public StmtExprMutator { using StmtExprMutator::Mutate; using StmtExprMutator::Mutate_; UnchangedOr Mutate(ffi::AnyView value, InplaceMode inplace_mode) override { - if (value.as()) return ffi::Unchanged(); + if (value.as()) return ffi::Unchanged(); return StmtExprMutator::Mutate(value, inplace_mode); } @@ -423,7 +436,7 @@ class DecomposePaddingBlockReplacer : public StmtExprMutator { explicit DecomposePaddingBlockReplacer(const ReplaceDesc& desc) : desc_(desc) {} private: - UnchangedOr Mutate_(const ForNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const ForObj* op, InplaceMode inplace_mode) final { Stmt new_loop{ffi::UnsafeInit{}}; if (op == desc_.in_bound_filling_pos.get()) { // position to rewrite inbound filling code @@ -453,8 +466,8 @@ StmtSRef DecomposePaddingImpl(ScheduleState self, const StmtSRef& block_sref, * - trim original block to write non-padding part only */ // Condition Checks and Information Collection - const SBlockNode* block = TVM_SREF_TO_SBLOCK(block_sref); - const SBlockRealizeNode* realize = GetSBlockRealize(self, block_sref).get(); + const SBlockObj* block = TVM_SREF_TO_SBLOCK(block_sref); + const SBlockRealizeObj* realize = GetSBlockRealize(self, block_sref).get(); ffi::Map dom_map; sym::Analyzer analyzer; @@ -469,13 +482,17 @@ StmtSRef DecomposePaddingImpl(ScheduleState self, const StmtSRef& block_sref, ffi::Array loops; bool found_const_filling_pos = false; bool found_in_bound_filling_pos = false; - For const_filling_pos = ffi::GetRef(loop_sref->StmtAs()); + For const_filling_pos = ffi::GetRef(loop_sref->StmtAs()); ffi::Optional in_bound_filling_pos; for (auto it = loop_srefs.rbegin(); it != loop_srefs.rend(); ++it) { - For cur_loop = ffi::GetRef((*it)->StmtAs()); + For cur_loop = ffi::GetRef((*it)->StmtAs()); Range range = Range(cur_loop->begin, cur_loop->end); dom_map.Set(cur_loop->var, range); - analyzer->Bind(cur_loop->var, range); + { + auto context = analyzer->GetConstraintContext(); + context->Append(sym::RangeConstraint(cur_loop->var, range)); + analyzer->SyncConstraintContext(context); + } loops.push_back(cur_loop); if (cur_loop.same_as(const_filling_pos)) { @@ -487,8 +504,8 @@ StmtSRef DecomposePaddingImpl(ScheduleState self, const StmtSRef& block_sref, } } else if (!found_in_bound_filling_pos) { if (cur_loop->body->size() != 1 || - (!cur_loop->body->seq[0]->IsInstance() && - !cur_loop->body->seq[0]->IsInstance())) { + (!cur_loop->body->seq[0]->IsInstance() && + !cur_loop->body->seq[0]->IsInstance())) { found_in_bound_filling_pos = true; } else { in_bound_filling_pos = cur_loop; @@ -516,7 +533,7 @@ StmtSRef DecomposePaddingImpl(ScheduleState self, const StmtSRef& block_sref, in_bound_block}; // Step 2. Execute IR replacement. - SBlock old_scope_root_block = ffi::GetRef(scope_root_sref->StmtAs()); + SBlock old_scope_root_block = ffi::GetRef(scope_root_sref->StmtAs()); SBlock new_scope_root = DecomposePaddingBlockReplacer::Replace(old_scope_root_block, replace_desc); if (check_only) { diff --git a/src/s_tir/schedule/primitive/for_kind.cc b/src/s_tir/schedule/primitive/for_kind.cc index a55daa499a87..0bd1c85bfa1e 100644 --- a/src/s_tir/schedule/primitive/for_kind.cc +++ b/src/s_tir/schedule/primitive/for_kind.cc @@ -157,7 +157,7 @@ void CheckParallelizability(const ScheduleState& self, const For& loop, ForKind */ void ParallelizeComputation(const ScheduleState& self, const StmtSRef& loop_sref, ForKind for_kind, ffi::Optional thread_axis) { - const ForNode* loop = TVM_SREF_TO_FOR(loop_sref); + const ForObj* loop = TVM_SREF_TO_FOR(loop_sref); /* * Check: @@ -180,7 +180,7 @@ void ParallelizeComputation(const ScheduleState& self, const StmtSRef& loop_sref : support::GPUThreadTag{-1, -1}); // Step 3. Loop update and IR replacement - ffi::ObjectPtr new_loop = ffi::make_object(*loop); + ffi::ObjectPtr new_loop = ffi::make_object(*loop); new_loop->kind = for_kind; if (thread_axis.has_value()) { new_loop->annotations.Set(tvm::s_tir::attr::kThreadBinding.name, thread_axis.value()); @@ -203,8 +203,8 @@ void Bind(ScheduleState self, const StmtSRef& loop_sref, const ffi::String& thre } void Unroll(ScheduleState self, const StmtSRef& loop_sref) { - const ForNode* loop = TVM_SREF_TO_FOR(loop_sref); - ffi::ObjectPtr new_loop = ffi::make_object(*loop); + const ForObj* loop = TVM_SREF_TO_FOR(loop_sref); + ffi::ObjectPtr new_loop = ffi::make_object(*loop); new_loop->kind = ForKind::kUnrolled; new_loop->annotations.erase(tvm::s_tir::attr::kThreadBinding.name); self->Replace(loop_sref, For(new_loop), {}); diff --git a/src/s_tir/schedule/primitive/get_block_loop.cc b/src/s_tir/schedule/primitive/get_block_loop.cc index b0372ed0e75c..2b2a21c68b38 100644 --- a/src/s_tir/schedule/primitive/get_block_loop.cc +++ b/src/s_tir/schedule/primitive/get_block_loop.cc @@ -32,14 +32,14 @@ ffi::Array GetSBlocks(const ScheduleState& self, const ffi::String& na using StmtExprVisitor::Visit_; ffi::Optional Visit(ffi::AnyView value) override { - if (value.as()) return std::nullopt; + if (value.as()) return std::nullopt; return StmtExprVisitor::Visit(value); } explicit Finder(const ScheduleState& self, const ffi::String& name) : self_(self), name_(name) {} - ffi::Optional Visit_(const SBlockNode* block) override { + ffi::Optional Visit_(const SBlockObj* block) override { if (block->name_hint == name_) { auto it = self_->stmt2ref.find(block); TVM_FFI_ICHECK(it != self_->stmt2ref.end()); @@ -54,7 +54,7 @@ ffi::Array GetSBlocks(const ScheduleState& self, const ffi::String& na }; BaseFunc func = self->mod->Lookup(gv); - const auto* function = TVM_TYPE_AS(func, FunctionNode); + const auto* function = TVM_TYPE_AS(func, FunctionObj); auto finder = ffi::make_object(self, name); finder->Visit(function->body); return std::move(finder->results_); @@ -62,7 +62,7 @@ ffi::Array GetSBlocks(const ScheduleState& self, const ffi::String& na ffi::Array GetLoops(const StmtSRef& block_sref) { std::vector result; - for (StmtSRefObj* parent = block_sref->parent; parent && parent->stmt->IsInstance(); + for (StmtSRefObj* parent = block_sref->parent; parent && parent->stmt->IsInstance(); parent = parent->parent) { result.push_back(ffi::GetRef(parent)); } @@ -74,12 +74,12 @@ ffi::Array GetChildBlocks(const ScheduleState& self, const StmtSRef& p using StmtExprVisitor::Visit_; ffi::Optional Visit(ffi::AnyView value) override { - if (value.as()) return std::nullopt; + if (value.as()) return std::nullopt; return StmtExprVisitor::Visit(value); } private: - ffi::Optional Visit_(const SBlockNode* block) final { + ffi::Optional Visit_(const SBlockObj* block) final { result.push_back(self->stmt2ref.at(block)); return std::nullopt; } @@ -91,11 +91,11 @@ ffi::Array GetChildBlocks(const ScheduleState& self, const StmtSRef& p ffi::Array result; }; auto collector = ffi::make_object(self); - if (parent_sref->stmt->IsInstance()) { - const auto* loop = static_cast(parent_sref->stmt); + if (parent_sref->stmt->IsInstance()) { + const auto* loop = static_cast(parent_sref->stmt); collector->Visit(loop->body); - } else if (parent_sref->stmt->IsInstance()) { - const auto* block = static_cast(parent_sref->stmt); + } else if (parent_sref->stmt->IsInstance()) { + const auto* block = static_cast(parent_sref->stmt); collector->Visit(block->body); } return std::move(collector->result); diff --git a/src/s_tir/schedule/primitive/hide_tensor_access.cc b/src/s_tir/schedule/primitive/hide_tensor_access.cc index 847e1f1edd1b..44479dc4378c 100644 --- a/src/s_tir/schedule/primitive/hide_tensor_access.cc +++ b/src/s_tir/schedule/primitive/hide_tensor_access.cc @@ -91,7 +91,7 @@ void UnsafeHideTensorAccess(ScheduleState self, const StmtSRef& block_sref, * - validity of tensor_index_array * - validity of tensor_type */ - const SBlockNode* block = TVM_SREF_TO_SBLOCK(block_sref); + const SBlockObj* block = TVM_SREF_TO_SBLOCK(block_sref); int num_access_regions = 0; if (tensor_type == "read") { num_access_regions = block->reads.size(); @@ -135,7 +135,7 @@ void UnsafeHideTensorAccess(ScheduleState self, const StmtSRef& block_sref, /* Step 1: Replace old block with the new block */ - auto n = ffi::make_object(*block); + auto n = ffi::make_object(*block); n->reads = reads; n->writes = writes; SBlock new_block = SBlock(n); diff --git a/src/s_tir/schedule/primitive/layout_transformation.cc b/src/s_tir/schedule/primitive/layout_transformation.cc index ff5b1e6ee02b..4c50bf2fd38e 100644 --- a/src/s_tir/schedule/primitive/layout_transformation.cc +++ b/src/s_tir/schedule/primitive/layout_transformation.cc @@ -24,7 +24,7 @@ #include #include #include -#include +#include #include #include @@ -139,22 +139,22 @@ class TransformLayoutPlanner : public StmtExprVisitor { explicit TransformLayoutPlanner(TensorVar old_tensor) : old_tensor_(old_tensor) {} private: - ffi::Optional Visit_(const ForNode* op) override { + ffi::Optional Visit_(const ForObj* op) override { BindLoopVar context(this, ffi::GetRef(op)); return StmtExprVisitor::Visit_(op); } - ffi::Optional Visit_(const BindNode* op) override { + ffi::Optional Visit_(const BindObj* op) override { BindVariableDefinition context(this, op->var, op->value); return StmtExprVisitor::Visit_(op); } - ffi::Optional Visit_(const SBlockRealizeNode* op) override { + ffi::Optional Visit_(const SBlockRealizeObj* op) override { BindBlockRealize context(this, ffi::GetRef(op)); return StmtExprVisitor::Visit_(op); } - ffi::Optional Visit_(const TensorStoreNode* op) override { + ffi::Optional Visit_(const TensorStoreObj* op) override { if (!op->dest.as_or_throw().same_as(old_tensor_)) { return std::nullopt; } @@ -371,7 +371,7 @@ class TransformLayoutPlanner : public StmtExprVisitor { this->new_iter_values = new_iter_values; } - UnchangedOr Mutate_(const TensorStoreNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const TensorStoreObj* op, InplaceMode inplace_mode) final { bool can_replace = [&]() -> bool { if (!op->dest.as_or_throw().same_as(info.store->dest)) { return false; @@ -405,7 +405,7 @@ class TransformLayoutPlanner : public StmtExprVisitor { .ValueOrUnchanged(store); } - UnchangedOr Mutate_(const SBlockRealizeNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const SBlockRealizeObj* op, InplaceMode inplace_mode) final { SBlockRealize realize = StmtExprMutator::Mutate_(op, inplace_mode) .ValueOrUnchanged(ffi::GetRef(op)) .as_or_throw(); @@ -429,7 +429,7 @@ class TransformLayoutPlanner : public StmtExprVisitor { return realize; } - UnchangedOr Mutate_(const SBlockNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const SBlockObj* op, InplaceMode inplace_mode) final { SBlock orig = ffi::GetRef(op); SBlock input = orig; if (info.innermost_block_realize.has_value() && @@ -676,7 +676,7 @@ class TransformLayoutPlanner : public StmtExprVisitor { // Under SSA, each variable is bound exactly once, so the lookup maps // grow monotonically and cleanup is unnecessary. Omitting cleanup also - // ensures correctness for flat BindNode (which has no body): the + // ensures correctness for flat BindObj (which has no body): the // binding must remain visible to subsequent sibling statements. struct BindVariableDefinition { BindVariableDefinition() {} @@ -742,7 +742,7 @@ class TransformLayoutPlanner : public StmtExprVisitor { * * Used to fill the `WriteInfo::dependent_loopnest` field. */ - std::unordered_map> loop_depth_lookup_; + std::unordered_map> loop_depth_lookup_; /*! \brief The variable mappings that are currently in-scope * @@ -750,7 +750,7 @@ class TransformLayoutPlanner : public StmtExprVisitor { * row-major traversal, even if they are rebound in let/block * mappings. */ - std::unordered_map active_var_bindings_; + std::unordered_map active_var_bindings_; /*! \brief The innermost BlockRealize surrounding the current node * @@ -771,7 +771,7 @@ class ReuseBlocksCollector : public s_tir::StmtExprVisitor { using s_tir::StmtExprVisitor::Visit_; ffi::Optional Visit(ffi::AnyView value) override { - if (value.as()) return std::nullopt; + if (value.as()) return std::nullopt; return s_tir::StmtExprVisitor::Visit(value); } @@ -793,7 +793,7 @@ class ReuseBlocksCollector : public s_tir::StmtExprVisitor { private: /*! \brief Override the Stmt visiting behaviour */ - ffi::Optional Visit_(const s_tir::SBlockNode* block) override { + ffi::Optional Visit_(const s_tir::SBlockObj* block) override { SBlock block_ref = ffi::GetRef(block); auto it = new_block_to_old_.find(block_ref); if (it != new_block_to_old_.end()) { @@ -808,10 +808,10 @@ class ReuseBlocksCollector : public s_tir::StmtExprVisitor { ffi::Map new_block_to_old_; }; -class TransformLayoutRewriter : public s_tir::StmtExprMutatorWithAnalyzer { +class TransformLayoutRewriter : public s_tir::StmtExprMutatorWithConstraintContext { public: - using s_tir::StmtExprMutatorWithAnalyzer::Mutate; - using s_tir::StmtExprMutatorWithAnalyzer::Mutate_; + using s_tir::StmtExprMutatorWithConstraintContext::Mutate; + using s_tir::StmtExprMutatorWithConstraintContext::Mutate_; /*! * \brief Rewrite the access to the tensor after the transformation @@ -851,7 +851,7 @@ class TransformLayoutRewriter : public s_tir::StmtExprMutatorWithAnalyzer { const IndexMap& index_map, const TransformLayoutPlanner::TransformPlan& plan, const sym::Analyzer& analyzer) - : StmtExprMutatorWithAnalyzer(analyzer), + : StmtExprMutatorWithConstraintContext(analyzer->GetConstraintContext()), old_tensor_(old_tensor), new_tensor_(new_tensor), index_map_(index_map), @@ -869,10 +869,10 @@ class TransformLayoutRewriter : public s_tir::StmtExprMutatorWithAnalyzer { *indices = this->IterMapSimplifyWithContext(*indices, true); } - using Parent = s_tir::StmtExprMutatorWithAnalyzer; + using Parent = s_tir::StmtExprMutatorWithConstraintContext; UnchangedOr Mutate(ffi::AnyView value, InplaceMode inplace_mode) final { - const auto* stmt = value.as(); + const auto* stmt = value.as(); auto plan_ptr = std::get_if(&plan_); bool insert_after = stmt && plan_ptr && plan_ptr->insert_after.get() == stmt; auto result = Parent::Mutate(value, inplace_mode); @@ -881,7 +881,7 @@ class TransformLayoutRewriter : public s_tir::StmtExprMutatorWithAnalyzer { return ffi::Any(SeqStmt({output, plan_ptr->new_block})); } - UnchangedOr Mutate_(const ForNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const ForObj* op, InplaceMode inplace_mode) final { // Some replacements may include the original string, such as // replacing `loop` with `{loop, post_proc}`. In this case, avoid // infinite recursion. @@ -896,7 +896,7 @@ class TransformLayoutRewriter : public s_tir::StmtExprMutatorWithAnalyzer { return Parent::Mutate_(op, inplace_mode); } - UnchangedOr Mutate_(const TensorLoadNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const TensorLoadObj* op, InplaceMode inplace_mode) final { TensorLoad tensor_load = Parent::Mutate_(op, inplace_mode) .ValueOrUnchanged(ffi::GetRef(op)) .as_or_throw(); @@ -909,7 +909,7 @@ class TransformLayoutRewriter : public s_tir::StmtExprMutatorWithAnalyzer { return tensor_load; } - UnchangedOr Mutate_(const TensorStoreNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const TensorStoreObj* op, InplaceMode inplace_mode) final { TensorStore tensor_store = Parent::Mutate_(op, inplace_mode) .ValueOrUnchanged(ffi::GetRef(op)) .as_or_throw(); @@ -941,7 +941,7 @@ class TransformLayoutRewriter : public s_tir::StmtExprMutatorWithAnalyzer { (*old_access_regions).MutateByApply(fmutate); } - UnchangedOr Mutate_(const SBlockNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const SBlockObj* op, InplaceMode inplace_mode) final { SBlock orig = [&]() { SBlock block = ffi::GetRef(op); while (true) { @@ -958,32 +958,39 @@ class TransformLayoutRewriter : public s_tir::StmtExprMutatorWithAnalyzer { .ValueOrUnchanged(ffi::GetRef(op)) .as_or_throw(); - auto infered_access_regions = GetSBlockReadWriteRegion(block, tensor_data_to_tensor_); - auto* n = block.CopyOnWrite(); - RewriteAccessRegion(&n->reads, infered_access_regions[0]); - RewriteAccessRegion(&n->writes, infered_access_regions[1]); - // Update match_tensors whose source references old_tensor_ - n->match_tensors.MutateByApply([this](const MatchTensorRegion& match_tensor) { - if (match_tensor->source->source.as_or_throw().same_as(old_tensor_)) { - auto new_source = match_tensor->source; - auto* source_n = new_source.CopyOnWrite(); - source_n->source = new_tensor_; - auto new_match = match_tensor; - new_match.CopyOnWrite()->source = new_source; - return new_match; - } - return match_tensor; - }); - n->alloc_tensors.MutateByApply([this](const TensorVar& tensor) { - if (tensor.same_as(old_tensor_)) { - return new_tensor_; - } else { - return tensor; - } - }); - - RecordReplacement(orig, block); - return block; + return WithConstraintScope([&]() -> ffi::Expected> { + for (const auto& iter_var : block->iter_vars) { + AppendConstraint(sym::RangeConstraint(iter_var->var, iter_var->dom.value())); + } + auto infered_access_regions = GetSBlockReadWriteRegion(block, tensor_data_to_tensor_); + auto* n = block.CopyOnWrite(); + RewriteAccessRegion(&n->reads, infered_access_regions[0]); + RewriteAccessRegion(&n->writes, infered_access_regions[1]); + // Update match_tensors whose source references old_tensor_ + n->match_tensors.MutateByApply([this](const MatchTensorRegion& match_tensor) { + if (match_tensor->source->source.as_or_throw().same_as( + old_tensor_)) { + auto new_source = match_tensor->source; + auto* source_n = new_source.CopyOnWrite(); + source_n->source = new_tensor_; + auto new_match = match_tensor; + new_match.CopyOnWrite()->source = new_source; + return new_match; + } + return match_tensor; + }); + n->alloc_tensors.MutateByApply([this](const TensorVar& tensor) { + if (tensor.same_as(old_tensor_)) { + return new_tensor_; + } else { + return tensor; + } + }); + + RecordReplacement(orig, block); + return block; + }) + .value(); } void RecordReplacement(SBlock before, SBlock after) { @@ -1116,7 +1123,7 @@ class TransformationPaddingExpressionError : public ScheduleErrorContextObj { explicit Visitor(const TensorVar& tensor) : tensor_(tensor) {} - ffi::Optional Visit_(const TensorLoadNode* op) final { + ffi::Optional Visit_(const TensorLoadObj* op) final { if (!op->source.as_or_throw().same_as(tensor_)) { illegal_load = ffi::GetRef(op); } @@ -1224,12 +1231,12 @@ IndexMap LegalizeIndexMapDType(const IndexMap& index_map, const ffi::Arrayfinal_indices.Map([&](PrimExpr index) { - if (auto* ptr = index.as()) { + if (auto* ptr = index.as()) { TVM_FFI_ICHECK(index_dtype.has_value()); return tvm::prim::MakeConst(PrimType(*index_dtype), ptr->value); } else { - return SubstituteWithDataTypeLegalization(index, - [&](const Var& var) { return var_map.Get(var); }); + return SubstituteWithDataTypeLegalization( + index, [&](const Var& var) -> ffi::Optional { return var_map.Get(var); }); } }); ffi::Optional opt_inverse_index_map = std::nullopt; @@ -1251,7 +1258,7 @@ void TransformLayout(ScheduleState self, const StmtSRef& block_sref, int tensor_ 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); + const SBlockObj* block_ptr = TVM_SREF_TO_SBLOCK(block_sref); TensorVar old_tensor = GetNthAccessTensor(self, ffi::GetRef(block_ptr), tensor_index, tensor_index_type); @@ -1276,18 +1283,20 @@ void TransformLayout(ScheduleState self, const StmtSRef& block_sref, int tensor_ StmtSRef scope_sref = defining_site_sref.has_value() ? defining_site_sref.value() : GetScopeRoot(self, block_sref, /*require_stage_pipeline=*/false); - const SBlockNode* scope_block = TVM_SREF_TO_SBLOCK(scope_sref); + const SBlockObj* scope_block = TVM_SREF_TO_SBLOCK(scope_sref); ffi::Optional opt_inverse = std::nullopt; PrimExpr padding_predicate = IntImm::Bool(false); if (!assume_injective_transform) { - std::tie(opt_inverse, padding_predicate) = [&]() { + auto inverse_and_padding = [&]() { ffi::Array region; for (const auto& dim : old_tensor.shape()) { region.push_back(Range::FromMinExtent(IntImm(dim.ty(), 0), dim)); } return index_map.NonSurjectiveInverse(region, analyzer); }(); + opt_inverse = inverse_and_padding.get<0>(); + padding_predicate = inverse_and_padding.get<1>(); } bool has_padding = !IsZero(padding_predicate); @@ -1297,7 +1306,7 @@ void TransformLayout(ScheduleState self, const StmtSRef& block_sref, int tensor_ } // Step 2: Infer the shape of the new tensor - auto new_tensor_type = ffi::make_object(*old_tensor.ty().get()); + auto new_tensor_type = ffi::make_object(*old_tensor.ty().get()); new_tensor_type->shape = index_map->MapShape(old_tensor.shape(), analyzer); TensorVar new_tensor = TensorVar(old_tensor->name, TensorType(std::move(new_tensor_type)), old_tensor->loc); @@ -1313,7 +1322,7 @@ void TransformLayout(ScheduleState self, const StmtSRef& block_sref, int tensor_ if (!defining_site_sref.has_value()) { GlobalVar g_var{ffi::UnsafeInit{}}; const auto* old_func = GetRootFunction(self->mod, scope_block, &g_var); - IRModuleNode* new_mod = self->mod.CopyOnWrite(); + IRModuleObj* new_mod = self->mod.CopyOnWrite(); ffi::MapObj* new_map = new_mod->functions.CopyOnWrite(); ffi::Array new_params = old_func->params.Map([&](Var param) -> Var { @@ -1345,7 +1354,7 @@ void TransformLayout(ScheduleState self, const StmtSRef& block_sref, int tensor_ */ IterVarType DetectNewBlockIterType( const PrimExpr& expr, - const std::unordered_map& block_iter_type_map) { + const std::unordered_map& block_iter_type_map) { IterVarType result{kOpaque}; bool found = false; auto walk_fn = [&](const Var& var) -> ffi::Expected { @@ -1452,17 +1461,17 @@ class OpaqueNewIterTypeError : public ScheduleErrorContextObj { void TransformBlockLayout(ScheduleState self, const StmtSRef& block_sref, const IndexMap& index_map) { - const SBlockNode* block_ptr = TVM_SREF_TO_SBLOCK(block_sref); + const SBlockObj* block_ptr = TVM_SREF_TO_SBLOCK(block_sref); const SBlock& block = ffi::GetRef(block_ptr); sym::Analyzer analyzer; AddShapeVarBounds(self, block_sref.get(), analyzer.get()); // Step 1: Collect outer loops and loop vars ffi::Array loops = GetLoops(block_sref); // outer loops of the block - std::unordered_set loop_vars; // loop vars of the outer loops + std::unordered_set loop_vars; // loop vars of the outer loops for (const StmtSRef& loop_sref : loops) { CheckLoopStartsWithZero(self, loop_sref, analyzer.get()); - loop_vars.emplace(loop_sref->StmtAs()->var.get()); + loop_vars.emplace(loop_sref->StmtAs()->var.get()); } // Step 2: Check the all outer loops have a single child and the block bindings are trivial (all @@ -1481,7 +1490,7 @@ void TransformBlockLayout(ScheduleState self, const StmtSRef& block_sref, // Step 3: Collect information of block iter vars ffi::Array block_vars; // iter_var->var of each block iter ffi::Map block_iter_dom; // domain of block iter - std::unordered_map block_iter_type; // iter type of block iter + std::unordered_map block_iter_type; // iter type of block iter ffi::Array block_iter_range_array; // array of block iter extents in the same order as block iters @@ -1570,7 +1579,7 @@ void TransformBlockLayout(ScheduleState self, const StmtSRef& block_sref, } // Make new block realize - SBlockRealizeNode* new_block_realize = block_realize.CopyOnWrite(); + SBlockRealizeObj* new_block_realize = block_realize.CopyOnWrite(); new_block_realize->iter_values = new_loop_vars; new_block_realize->block = new_block; @@ -1583,7 +1592,7 @@ void TransformBlockLayout(ScheduleState self, const StmtSRef& block_sref, } // Step 6: Do the actual replacement - if (scope_sref->StmtAs()) { + if (scope_sref->StmtAs()) { TVM_FFI_ICHECK(new_loop_vars.empty()) << "Invalid block to loop replacement due to layout transform " << index_map; } diff --git a/src/s_tir/schedule/primitive/loop_transformation.cc b/src/s_tir/schedule/primitive/loop_transformation.cc index 665509e3b407..15f6f2616142 100644 --- a/src/s_tir/schedule/primitive/loop_transformation.cc +++ b/src/s_tir/schedule/primitive/loop_transformation.cc @@ -34,7 +34,7 @@ class BlockPredicateAppender : public StmtExprMutator { using StmtExprMutator::Mutate; using StmtExprMutator::Mutate_; UnchangedOr Mutate(ffi::AnyView value, InplaceMode inplace_mode) override { - if (value.as()) return ffi::Unchanged(); + if (value.as()) return ffi::Unchanged(); return StmtExprMutator::Mutate(value, inplace_mode); } @@ -46,14 +46,14 @@ class BlockPredicateAppender : public StmtExprMutator { private: // For each direct child of type BlockRealizeNode, append the predicate - UnchangedOr Mutate_(const SBlockRealizeNode* realize, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const SBlockRealizeObj* realize, InplaceMode inplace_mode) final { // We do not recursively do this PrimExpr predicate = realize->predicate && to_append_; if (inplace_mode == InplaceMode::kAllow) { - const_cast(realize)->predicate = std::move(predicate); + const_cast(realize)->predicate = std::move(predicate); return ffi::Unchanged(); } else { - auto copy = ffi::make_object(*realize); + auto copy = ffi::make_object(*realize); copy->predicate = std::move(predicate); return SBlockRealize(std::move(copy)); } @@ -78,7 +78,7 @@ class SubstituteVarAndCollectOpaqueBlock : public StmtExprMutator { } private: - UnchangedOr Mutate_(const SBlockRealizeNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const SBlockRealizeObj* op, InplaceMode inplace_mode) final { SBlockRealize realize = StmtExprMutator::Mutate_(op, inplace_mode) .ValueOrUnchanged(ffi::GetRef(op)) .as_or_throw(); @@ -109,7 +109,7 @@ class IterMapSimplifyBlockBinding : public StmtExprMutator { ffi::MapObj* opaque_blocks, bool preserve_unit_iters) { ffi::Map loop_var2extent; for (const StmtSRef& sref : loop_srefs) { - const ForNode* loop = TVM_SREF_TO_FOR(sref); + const ForObj* loop = TVM_SREF_TO_FOR(sref); loop_var2extent.Set(loop->var, Range(loop->begin, loop->end)); } return ffi::make_object(opaque_blocks, std::move(loop_var2extent), @@ -120,14 +120,14 @@ class IterMapSimplifyBlockBinding : public StmtExprMutator { } private: - UnchangedOr Mutate_(const ForNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const ForObj* op, InplaceMode inplace_mode) final { loop_var2extent_.Set(op->var, Range(op->begin, op->end)); Stmt res = StmtExprMutator::Mutate_(op, inplace_mode).ValueOrUnchanged(ffi::GetRef(op)); loop_var2extent_.erase(op->var); return res; } - UnchangedOr Mutate_(const SBlockRealizeNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const SBlockRealizeObj* op, InplaceMode inplace_mode) final { // skip opaque block and update mapping if (op->iter_values.empty()) { SBlock block = op->block; @@ -153,10 +153,10 @@ class IterMapSimplifyBlockBinding : public StmtExprMutator { return ffi::Unchanged(); } else { if (inplace_mode == InplaceMode::kAllow) { - const_cast(op)->iter_values = std::move(v); + const_cast(op)->iter_values = std::move(v); return ffi::Unchanged(); } else { - auto copy = ffi::make_object(*op); + auto copy = ffi::make_object(*op); copy->iter_values = std::move(v); return SBlockRealize(std::move(copy)); } @@ -188,7 +188,7 @@ class BlockPropertyError : public ScheduleErrorContextObj { using StmtExprVisitor::Visit_; ffi::Optional Visit(ffi::AnyView value) override { - if (value.as()) return std::nullopt; + if (value.as()) return std::nullopt; return StmtExprVisitor::Visit(value); } @@ -197,7 +197,7 @@ class BlockPropertyError : public ScheduleErrorContextObj { : state_(state), top_(top) {} private: - ffi::Optional Visit_(const SBlockNode* op) final { + ffi::Optional Visit_(const SBlockObj* op) final { for (const IterVar& iter_var : op->iter_vars) { if (iter_var->iter_type != kDataPar && iter_var->iter_type != kCommReduce) { throw MakeScheduleError(state_->mod, ffi::GetRef(op)); @@ -436,7 +436,7 @@ ffi::Array Split(ScheduleState self, const StmtSRef& loop_sref, // - The execution order has not changed. (The block executes with the same args and the same // order with before. // Step 1. Check correctness - const ForNode* loop = TVM_SREF_TO_FOR(loop_sref); + const ForObj* loop = TVM_SREF_TO_FOR(loop_sref); if (!loop->annotations.empty()) { throw MakeScheduleError(self->mod, ffi::GetRef(loop)); } @@ -461,7 +461,12 @@ ffi::Array Split(ScheduleState self, const StmtSRef& loop_sref, const PrimExpr& factor = factors[i]; Var var = loop->var.CopyWithSuffix("_" + std::to_string(i)).CopyWithDType(dtype); substitute_value = substitute_value * factor + var.as_or_throw(); - analyzer->Bind(var, Range::FromMinExtent(IntImm(dtype, 0), tvm::prim::cast(factor, dtype))); + { + auto context = analyzer->GetConstraintContext(); + context->Append(sym::RangeConstraint( + var, Range::FromMinExtent(IntImm(dtype, 0), tvm::prim::cast(factor, dtype)))); + analyzer->SyncConstraintContext(context); + } new_loop_vars.emplace_back(std::move(var)); } ffi::Map opaque_block_reuse; @@ -492,7 +497,7 @@ ffi::Array Split(ScheduleState self, const StmtSRef& loop_sref, for (int i = 0; i < n; i++) { result_srefs.push_back(self->stmt2ref.at(new_stmt.get())); if (i + 1 < n) { - const ForNode* outer_loop = TVM_TYPE_AS(new_stmt, ForNode); + const ForObj* outer_loop = TVM_TYPE_AS(new_stmt, ForObj); TVM_FFI_ICHECK_EQ(outer_loop->body->size(), 1); new_stmt = outer_loop->body->seq[0]; } @@ -513,7 +518,7 @@ class TensorIndicesMapExtractor : public StmtExprVisitor { } private: - ffi::Optional Visit_(const TensorStoreNode* store) final { + ffi::Optional Visit_(const TensorStoreObj* store) final { ffi::Array indices; bool check_ = false; for (size_t i = 0; i < store->indices.size(); i++) { @@ -530,7 +535,7 @@ class TensorIndicesMapExtractor : public StmtExprVisitor { return StmtExprVisitor::Visit_(store); } - ffi::Optional Visit_(const TensorLoadNode* load) final { + ffi::Optional Visit_(const TensorLoadObj* load) final { ffi::Array indices; bool check_ = false; for (size_t i = 0; i < load->indices.size(); i++) { @@ -582,7 +587,7 @@ class BlockMutator : public StmtExprMutator { : new_loop_var_(new_loop_var), min_(min), extent_(extent) {} private: - UnchangedOr Mutate_(const SBlockNode* _op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const SBlockObj* _op, InplaceMode inplace_mode) final { SBlock new_block = StmtExprMutator::Mutate_(_op, inplace_mode) .ValueOrUnchanged(ffi::GetRef(_op)) .as_or_throw(); @@ -662,7 +667,7 @@ class BlockMutator : public StmtExprMutator { return block_stmt; } - UnchangedOr Mutate_(const SBlockRealizeNode* realize, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const SBlockRealizeObj* realize, InplaceMode inplace_mode) final { ffi::Array iter_values = realize->iter_values; for (size_t i = 0; i < iter_values.size(); i++) { if (new_loop_var_.same_as(iter_values[i])) { @@ -677,7 +682,7 @@ class BlockMutator : public StmtExprMutator { return stmt; } - UnchangedOr Mutate_(const ForNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const ForObj* op, InplaceMode inplace_mode) final { For res = StmtExprMutator::Mutate_(op, inplace_mode) .ValueOrUnchanged(ffi::GetRef(op)) .as_or_throw(); @@ -707,7 +712,7 @@ class BlockMutator : public StmtExprMutator { const ffi::String get_sblock_name(SeqStmt loop_body) { TVM_FFI_ICHECK_EQ(loop_body->size(), 1); const Stmt& child = loop_body->seq[0]; - if (const auto* block_realize = child.as()) { + if (const auto* block_realize = child.as()) { return block_realize->block->name_hint; } return get_sblock_name(child.as_or_throw()->body); @@ -715,7 +720,7 @@ const ffi::String get_sblock_name(SeqStmt loop_body) { ffi::Array LoopPartition(ScheduleState self, const StmtSRef& loop_sref, const ffi::Array& factors, bool preserve_unit_iters) { - const ForNode* loop = TVM_SREF_TO_FOR(loop_sref); + const ForObj* loop = TVM_SREF_TO_FOR(loop_sref); if (!loop->annotations.empty()) { throw MakeScheduleError(self->mod, ffi::GetRef(loop)); } @@ -788,7 +793,7 @@ ffi::Array LoopPartition(ScheduleState self, const StmtSRef& loop_sref ffi::Array partition_srefs; partition_srefs.reserve(n); for (int i = 0; i < n; i++) { - const SeqStmt& body = block_partitions[i].as()->block->body; + const SeqStmt& body = block_partitions[i].as()->block->body; TVM_FFI_ICHECK_EQ(body->size(), 1); StmtSRef partition_loop_sref = self->stmt2ref.at(body->seq[0].get()); partition_srefs.push_back(partition_loop_sref); @@ -801,7 +806,7 @@ class LoopReconstructor : public StmtExprMutator { using StmtExprMutator::Mutate; using StmtExprMutator::Mutate_; UnchangedOr Mutate(ffi::AnyView value, InplaceMode inplace_mode) override { - if (value.as()) return ffi::Unchanged(); + if (value.as()) return ffi::Unchanged(); return StmtExprMutator::Mutate(value, inplace_mode); } @@ -854,14 +859,14 @@ class LoopReconstructor : public StmtExprMutator { } private: - UnchangedOr Mutate_(const SBlockNode* block, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const SBlockObj* block, InplaceMode inplace_mode) final { if (block != scope_root_.get()) { return ffi::Unchanged(); } return StmtExprMutator::Mutate_(block, inplace_mode); } - UnchangedOr Mutate_(const ForNode* loop, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const ForObj* loop, InplaceMode inplace_mode) final { if (ffi::GetRef(loop) == need_remove_loop_.back()) { return new_outer_loop_.value(); } else if (std::count(need_remove_loop_.begin(), need_remove_loop_.end(), @@ -902,7 +907,7 @@ StmtSRef Merge(ScheduleState self, const ffi::Array& loop_srefs) { std::vector nest_loop_i_extents; std::vector nest_loop_i_loops; for (auto p = sref.get(); p != lca.get(); p = p->parent) { - if (auto loop = p->StmtAs()) { + if (auto loop = p->StmtAs()) { if (!loop->annotations.empty()) { throw MakeScheduleError(self->mod, ffi::GetRef(loop)); @@ -913,7 +918,7 @@ StmtSRef Merge(ScheduleState self, const ffi::Array& loop_srefs) { } } lca_nest_loops.push_back(nest_loop_i_loops); - const ForNode* outer_loop = nullptr; + const ForObj* outer_loop = nullptr; for (auto iter = nest_loop_i_loops.rbegin(); iter != nest_loop_i_loops.rend(); ++iter) { if (outer_loop && (outer_loop->body->size() != 1 || !outer_loop->body->seq[0].same_as(*iter))) { @@ -946,7 +951,7 @@ StmtSRef Merge(ScheduleState self, const ffi::Array& loop_srefs) { } } // Step 2. Create merged loops and replace the original loops - SBlock scope_root = ffi::GetRef(scope_root_sref.value()->StmtAs()); + SBlock scope_root = ffi::GetRef(scope_root_sref.value()->StmtAs()); auto reconstructor = ffi::make_object(scope_root, lca_nest_loops); reconstructor->MakeNewLoop(); SBlock new_scope_root = @@ -962,15 +967,15 @@ StmtSRef Fuse(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.) - std::vector loops; + std::vector loops; loops.reserve(loop_srefs.size()); ffi::Optional outer_loop_sref; - const ForNode* outer_loop = nullptr; + const ForObj* outer_loop = nullptr; sym::Analyzer analyzer; - std::unordered_set outer_loop_vars; + std::unordered_set outer_loop_vars; // Step 1. check correctness for (const StmtSRef& sref : loop_srefs) { - const ForNode* loop = TVM_SREF_TO_FOR(sref); + const ForObj* loop = TVM_SREF_TO_FOR(sref); if (!loop->annotations.empty()) { throw MakeScheduleError(self->mod, ffi::GetRef(loop)); } @@ -1065,7 +1070,7 @@ std::unordered_set CollectLoopsIntoSet( for (const StmtSRef& loop_sref : ordered_loop_srefs) { auto inserted = loop_srefs.insert(loop_sref.get()); if (!inserted.second) { - const ForNode* loop = TVM_SREF_TO_FOR(loop_sref); + const ForObj* loop = TVM_SREF_TO_FOR(loop_sref); throw MakeScheduleError(self->mod, ffi::GetRef(loop)); } } @@ -1092,7 +1097,7 @@ std::pair GetBoundaryOfReorderRange( } for (const StmtSRefObj* v = loop_sref;; v = v->parent) { // Case 1. If `v` corresponds to a block, stop traversal. - if (v->stmt->IsInstance()) { + if (v->stmt->IsInstance()) { if (scope_block_visited) { throw MakeScheduleError( self->mod, std::nullopt, LoopsNotAChainError::ProblemKind::kNotUnderAScope); @@ -1138,8 +1143,8 @@ std::vector GetLoopsInReorderRange(const ScheduleState& self std::vector chain; for (const StmtSRefObj* loop_sref = bottom; loop_sref != top;) { const StmtSRefObj* parent_loop_sref = loop_sref->parent; - const ForNode* outer = parent_loop_sref->StmtAs(); - const ForNode* inner = loop_sref->StmtAs(); + const ForObj* outer = parent_loop_sref->StmtAs(); + const ForObj* inner = loop_sref->StmtAs(); TVM_FFI_ICHECK(outer != nullptr && inner != nullptr); if (outer->body->size() != 1 || outer->body->seq[0].get() != inner) { throw MakeScheduleError( @@ -1166,24 +1171,24 @@ std::vector GetLoopsInReorderRange(const ScheduleState& self For ConstructNewLoopChain(const ScheduleState& self, std::vector chain, const ffi::Array& ordered_loop_srefs, const std::unordered_set& loop_srefs) { - std::unordered_set inner_vars; + std::unordered_set inner_vars; inner_vars.reserve(chain.size()); ffi::Optional new_loop; int index = static_cast(ordered_loop_srefs.size()) - 1; for (const StmtSRefObj* loop_sref : chain) { - const ForNode* copy = nullptr; + const ForObj* copy = nullptr; if (loop_srefs.count(loop_sref)) { - copy = ordered_loop_srefs[index]->StmtAs(); + copy = ordered_loop_srefs[index]->StmtAs(); --index; } else { - copy = loop_sref->StmtAs(); + copy = loop_sref->StmtAs(); } TVM_FFI_ICHECK(copy != nullptr); - ffi::ObjectPtr n = ffi::make_object(*copy); + ffi::ObjectPtr n = ffi::make_object(*copy); if (new_loop.has_value()) { n->body = new_loop.value(); } else { - n->body = loop_sref->StmtAs()->body; + n->body = loop_sref->StmtAs()->body; } auto walkfn = [&inner_vars](const Var& var) -> ffi::Expected { return inner_vars.count(var.get()) ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var)) @@ -1231,7 +1236,7 @@ void Reorder(ScheduleState self, const ffi::Array& ordered_loop_srefs) } StmtSRef AddUnitLoop(ScheduleState self, StmtSRef sref) { - if (sref->stmt->IsInstance()) { + if (sref->stmt->IsInstance()) { For new_loop = For(PrimVar("u", PrimType::Int(32)), 0, 1, IntImm(PrimType::Int(32), 1), ForKind::kDefault, ffi::GetRef(sref->stmt)); self->Replace(sref, new_loop, {}); @@ -1242,13 +1247,13 @@ StmtSRef AddUnitLoop(ScheduleState self, StmtSRef sref) { using StmtExprMutator::Mutate; using StmtExprMutator::Mutate_; UnchangedOr Mutate(ffi::AnyView value, InplaceMode inplace_mode) override { - if (value.as()) return ffi::Unchanged(); + if (value.as()) return ffi::Unchanged(); return StmtExprMutator::Mutate(value, inplace_mode); } - explicit NewLoopCreator(const StmtNode* src_block) : src_block_(src_block) {} + explicit NewLoopCreator(const StmtObj* src_block) : src_block_(src_block) {} - UnchangedOr Mutate_(const SBlockRealizeNode* realize, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const SBlockRealizeObj* realize, InplaceMode inplace_mode) final { if (realize->block.get() == src_block_) { new_loop_ = For(PrimVar("u", PrimType::Int(32)), 0, 1, IntImm(PrimType::Int(32), 1), ForKind::kDefault, ffi::GetRef(realize)); @@ -1257,7 +1262,7 @@ StmtSRef AddUnitLoop(ScheduleState self, StmtSRef sref) { return StmtExprMutator::Mutate_(realize, inplace_mode); } - const StmtNode* src_block_; + const StmtObj* src_block_; ffi::Optional new_loop_; }; @@ -1266,10 +1271,10 @@ StmtSRef AddUnitLoop(ScheduleState self, StmtSRef sref) { auto creator = ffi::make_object(sref->stmt); Stmt new_stmt = creator->Mutate(ffi::GetRef(parent_sref->stmt)) .ValueOrUnchanged(ffi::GetRef(parent_sref->stmt)); - if (new_stmt->IsInstance()) { + if (new_stmt->IsInstance()) { self->Replace(parent_sref, std::move(new_stmt), {}); } else { - SBlock old_parent_block = ffi::GetRef(parent_sref->StmtAs()); + SBlock old_parent_block = ffi::GetRef(parent_sref->StmtAs()); SBlock new_parent_block = new_stmt.as_or_throw(); self->Replace(parent_sref, new_stmt, {{old_parent_block, new_parent_block}}); } diff --git a/src/s_tir/schedule/primitive/pad_einsum.cc b/src/s_tir/schedule/primitive/pad_einsum.cc index a9a719e9cfc5..b2acb42d5f26 100644 --- a/src/s_tir/schedule/primitive/pad_einsum.cc +++ b/src/s_tir/schedule/primitive/pad_einsum.cc @@ -37,7 +37,7 @@ ffi::Optional> CheckTrivialTensorIndices( const ffi::Array& tensor_access) { ffi::Array indices; for (const PrimExpr& index : tensor_access) { - if (index->IsInstance()) { + if (index->IsInstance()) { continue; } auto var = index.as(); @@ -56,7 +56,7 @@ ffi::Optional> CheckTrivialTensorAccess(const TensorRegion& tens if (!tvm::prim::IsOne(range->extent)) { return std::nullopt; } - if (range->begin->IsInstance()) { + if (range->begin->IsInstance()) { continue; } if (auto var = range->begin.as()) { @@ -146,8 +146,8 @@ struct TensorPadding { int ndim = tensor_region->region.size(); for (int i = 0; i < ndim; ++i) { PrimExpr pos = tensor_region->region[i]->begin; - TVM_FFI_ICHECK(pos->IsInstance() || pos.as()); - if (pos->IsInstance()) { + TVM_FFI_ICHECK(pos->IsInstance() || pos.as()); + if (pos->IsInstance()) { shape.push_back(IntImm(pos.ty(), 1)); } else if (ffi::Optional extent = iter_extents.Get(pos.as_or_throw())) { shape.push_back(extent.value()); @@ -215,7 +215,7 @@ struct TensorPadding { Einsum ExtractEinsum(const ScheduleState& self, const SBlock& block) { Einsum result; - std::unordered_set tensor_used; + std::unordered_set tensor_used; int n_reads = block->reads.size(); for (int i = 0; i < n_reads; ++i) { const TensorVar& tensor = block->reads[i]->source.as_or_throw(); @@ -302,7 +302,7 @@ class PadEinsumTensorReplacer : public StmtExprMutator { using StmtExprMutator::Mutate; using StmtExprMutator::Mutate_; - UnchangedOr Mutate_(const SBlockNode* old_block_ptr, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const SBlockObj* old_block_ptr, InplaceMode inplace_mode) final { SBlock old_block = ffi::GetRef(old_block_ptr); SBlock block = StmtExprMutator::Mutate_(old_block_ptr, inplace_mode) .ValueOrUnchanged(ffi::GetRef(old_block_ptr)) @@ -311,7 +311,7 @@ class PadEinsumTensorReplacer : public StmtExprMutator { iter_vars.reserve(block->iter_vars.size()); for (const IterVar& iter_var : block->iter_vars) { if (ffi::Optional new_dom = iter2padded_extents.Get(iter_var->var)) { - ffi::ObjectPtr new_iter_var = ffi::make_object(*iter_var.get()); + ffi::ObjectPtr new_iter_var = ffi::make_object(*iter_var.get()); new_iter_var->dom = Range::FromMinExtent(iter_var->dom.value()->begin, new_dom.value()); iter_vars.push_back(IterVar(new_iter_var)); } else { @@ -345,20 +345,20 @@ class PadEinsumTensorReplacer : public StmtExprMutator { return new_block; } - UnchangedOr Mutate_(const ForNode* old_for_ptr, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const ForObj* old_for_ptr, InplaceMode inplace_mode) final { For old_for = ffi::GetRef(old_for_ptr); For new_for = StmtExprMutator::Mutate_(old_for_ptr, inplace_mode) .ValueOrUnchanged(ffi::GetRef(old_for_ptr)) .as_or_throw(); if (ffi::Optional new_extent = loop_var2padded_extent.Get(new_for->var)) { - ffi::ObjectPtr new_for_ptr = ffi::make_object(*new_for.get()); + ffi::ObjectPtr new_for_ptr = ffi::make_object(*new_for.get()); new_for_ptr->end = new_for_ptr->begin + new_extent.value(); new_for = For(new_for_ptr); } return new_for; } - UnchangedOr Mutate_(const TensorStoreNode* old_store_ptr, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const TensorStoreObj* old_store_ptr, InplaceMode inplace_mode) final { TensorStore store = StmtExprMutator::Mutate_(old_store_ptr, inplace_mode) .ValueOrUnchanged(ffi::GetRef(old_store_ptr)) .as_or_throw(); @@ -370,8 +370,7 @@ class PadEinsumTensorReplacer : public StmtExprMutator { } } - UnchangedOr Mutate_(const TensorLoadNode* old_load_ptr, - InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const TensorLoadObj* old_load_ptr, InplaceMode inplace_mode) final { TensorLoad load = StmtExprMutator::Mutate_(old_load_ptr, inplace_mode) .ValueOrUnchanged(ffi::GetRef(old_load_ptr)) .as_or_throw(); @@ -390,10 +389,10 @@ class PadEinsumTensorReplacer : public StmtExprMutator { void PadEinsum(ScheduleState self, const StmtSRef& block_sref, const ffi::Array& padding) { sym::Analyzer analyzer; // Step 1: Input checking and error handling - const SBlockNode* block = TVM_SREF_TO_SBLOCK(block_sref); + const SBlockObj* block = TVM_SREF_TO_SBLOCK(block_sref); SBlockRealize realize = GetSBlockRealize(self, block_sref); StmtSRef scope_sref = GetScopeRoot(self, block_sref, /*require_stage_pipeline=*/true); - const SBlockNode* scope_block = TVM_SREF_TO_SBLOCK(scope_sref); + const SBlockObj* scope_block = TVM_SREF_TO_SBLOCK(scope_sref); InvalidPaddingError::Check(self, ffi::GetRef(block), padding); // Step 2. Extract the Einsum pattern ExtractEinsum(self, ffi::GetRef(block)); @@ -476,7 +475,7 @@ void PadEinsum(ScheduleState self, const StmtSRef& block_sref, const ffi::Array< } // Step 7. Create new scope SBlock new_scope_block = [&]() { - ffi::ObjectPtr n = ffi::make_object(*scope_block); + ffi::ObjectPtr n = ffi::make_object(*scope_block); n->body = SeqStmt(new_scope_body, scope_block->body->loc); n->alloc_tensors.insert(n->alloc_tensors.end(), alloc_tensors.begin(), alloc_tensors.end()); return SBlock(n); diff --git a/src/s_tir/schedule/primitive/read_write_at.cc b/src/s_tir/schedule/primitive/read_write_at.cc index 392e56602d75..5e199bb667dc 100644 --- a/src/s_tir/schedule/primitive/read_write_at.cc +++ b/src/s_tir/schedule/primitive/read_write_at.cc @@ -73,9 +73,9 @@ class ScopeReplacer : public StmtExprMutator { using StmtExprMutator::Mutate; using StmtExprMutator::Mutate_; - static SBlock Replace(const SBlockNode* scope_block, const TensorVar& dst, - const ForNode* old_loop, const ForNode* new_loop) { - ffi::ObjectPtr new_scope_block = ffi::make_object(*scope_block); + static SBlock Replace(const SBlockObj* scope_block, const TensorVar& dst, const ForObj* old_loop, + const ForObj* new_loop) { + ffi::ObjectPtr new_scope_block = ffi::make_object(*scope_block); new_scope_block->body = ffi::make_object(old_loop, new_loop) ->Mutate(new_scope_block->body, InplaceMode::kAllow) .ValueOrUnchanged(std::move(new_scope_block->body)); @@ -83,18 +83,18 @@ class ScopeReplacer : public StmtExprMutator { return SBlock(new_scope_block); } - explicit ScopeReplacer(const ForNode* old_loop, const ForNode* new_loop) + explicit ScopeReplacer(const ForObj* old_loop, const ForObj* new_loop) : old_loop_(old_loop), new_loop_(new_loop), found_(false) {} private: UnchangedOr Mutate(ffi::AnyView value, InplaceMode inplace_mode) final { - if (value.as() || found_) return ffi::Unchanged(); + if (value.as() || found_) return ffi::Unchanged(); return StmtExprMutator::Mutate(value, inplace_mode); } - UnchangedOr Mutate_(const SBlockNode* block, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const SBlockObj* block, InplaceMode inplace_mode) final { return ffi::Unchanged(); } - UnchangedOr Mutate_(const ForNode* loop, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const ForObj* loop, InplaceMode inplace_mode) final { if (loop == old_loop_) { found_ = true; return ffi::GetRef(new_loop_); @@ -102,8 +102,8 @@ class ScopeReplacer : public StmtExprMutator { return StmtExprMutator::Mutate_(loop, inplace_mode); } - const ForNode* old_loop_; - const ForNode* new_loop_; + const ForObj* old_loop_; + const ForObj* new_loop_; bool found_; }; @@ -117,19 +117,19 @@ class ReadWriteAtTensorReplacer : public StmtExprMutator { : src_(src), dst_(dst), block_sref_reuse_(block_sref_reuse) {} private: - UnchangedOr Mutate_(const TensorStoreNode* _store, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const TensorStoreObj* _store, InplaceMode inplace_mode) final { TensorStore store = StmtExprMutator::Mutate_(_store, inplace_mode) .ValueOrUnchanged(ffi::GetRef(_store)) .as_or_throw(); if (store->dest.as_or_throw().same_as(src_)) { - ffi::ObjectPtr new_store = ffi::make_object(*store.get()); + ffi::ObjectPtr new_store = ffi::make_object(*store.get()); new_store->dest = dst_; return TensorStore(new_store); } return store; } - UnchangedOr Mutate_(const TensorLoadNode* _load, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const TensorLoadObj* _load, InplaceMode inplace_mode) final { TensorLoad load = StmtExprMutator::Mutate_(_load, inplace_mode) .ValueOrUnchanged(ffi::GetRef(_load)) .as_or_throw(); @@ -139,12 +139,12 @@ class ReadWriteAtTensorReplacer : public StmtExprMutator { return load; } - UnchangedOr Mutate_(const SBlockNode* _block, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const SBlockObj* _block, InplaceMode inplace_mode) final { SBlock old_block = ffi::GetRef(_block); SBlock block = StmtExprMutator::Mutate_(_block, InplaceMode::kDisallow) .ValueOrUnchanged(ffi::GetRef(_block)) .as_or_throw(); - ffi::ObjectPtr new_block = ffi::make_object(*block.get()); + ffi::ObjectPtr new_block = ffi::make_object(*block.get()); new_block->reads = ReplaceTensor(new_block->reads, src_, dst_); new_block->writes = ReplaceTensor(new_block->writes, src_, dst_); block_sref_reuse_->Set(old_block, SBlock(new_block)); @@ -161,7 +161,7 @@ struct ReadWriteAtImpl { static StmtSRef Main(ScheduleState self, const StmtSRef& loop_sref, const StmtSRef& block_sref, int tensor_index, const ffi::String& storage_scope, ffi::Map annotations) { - const SBlockNode* block = TVM_SREF_TO_SBLOCK(block_sref); + const SBlockObj* block = TVM_SREF_TO_SBLOCK(block_sref); TensorVar src = GetNthAccessTensor(self, ffi::GetRef(block), tensor_index, is_read ? TensorIndexType::kRead : TensorIndexType::kWrite); TensorVar dst = WithScope(src, runtime::MemoryScope::Get(storage_scope)); @@ -177,17 +177,17 @@ struct ReadWriteAtImpl { private: static ffi::Map GetLoopDomain(const StmtSRefObj* loop_sref) { ffi::Map result; - for (const ForNode* loop; (loop = loop_sref->StmtAs()) != nullptr; + for (const ForObj* loop; (loop = loop_sref->StmtAs()) != nullptr; loop_sref = loop_sref->parent) { result.Set(loop->var, Range(loop->begin, loop->end)); } return result; } - StmtSRef ReplaceScopeBlock(const ForNode* new_loop, const SBlockNode* new_block) { + StmtSRef ReplaceScopeBlock(const ForObj* new_loop, const SBlockObj* new_block) { StmtSRef scope_root_sref = GetScopeRoot(self_, loop_sref_, /*require_stage_pipeline=*/true); - const SBlockNode* scope_block = TVM_SREF_TO_SBLOCK(scope_root_sref); + const SBlockObj* scope_block = TVM_SREF_TO_SBLOCK(scope_root_sref); SBlock new_scope_block = ScopeReplacer::Replace(scope_block, dst_, loop_, new_loop); block_sref_reuse_.Set(ffi::GetRef(scope_block), new_scope_block); self_->Replace(scope_root_sref, new_scope_block, block_sref_reuse_); @@ -218,7 +218,7 @@ struct ReadWriteAtImpl { bool w_visited = false; auto f_visit = [this, &relaxed_regions, &r_visited, &w_visited, &scope](const SBlockRealize& realize) -> ffi::Expected { - const SBlockNode* block = realize->block.get(); + const SBlockObj* block = realize->block.get(); bool has_r = HasTensor(block->reads, src_); bool has_w = HasTensor(block->writes, src_); r_visited = r_visited || has_r; @@ -288,7 +288,7 @@ struct ReadWriteAtImpl { ? MakeSBlock(src_, dst_, new_block_name_hint, GetLoopDomain(loop_sref_.get()), domain) : MakeSBlock(dst_, src_, new_block_name_hint, GetLoopDomain(loop_sref_.get()), domain); subtrees.insert(subtrees.begin() + insert_pos, realize); - ffi::ObjectPtr new_loop = ffi::make_object(*loop_); + ffi::ObjectPtr new_loop = ffi::make_object(*loop_); new_loop->body = SeqStmt(std::move(subtrees)); return {For(new_loop), realize}; } @@ -323,7 +323,7 @@ struct ReadWriteAtImpl { iter_vars.push_back(IterVar(range, v.as_or_throw(), IterVarType::kDataPar)); return ffi::Any(v.as_or_throw()); }; - ffi::ObjectPtr dom = ffi::make_object(*domain[i].get()); + ffi::ObjectPtr dom = ffi::make_object(*domain[i].get()); dom->begin = ffi::StructuralMap(std::move(dom->begin), f_substitute) .as_or_throw(); @@ -369,7 +369,7 @@ struct ReadWriteAtImpl { ScheduleState self_; const StmtSRef& loop_sref_; - const ForNode* loop_; + const ForObj* loop_; const TensorVar& src_; const TensorVar& dst_; ffi::Map annotations_; diff --git a/src/s_tir/schedule/primitive/reduction.cc b/src/s_tir/schedule/primitive/reduction.cc index 754ff501a697..5b34ff370e1c 100644 --- a/src/s_tir/schedule/primitive/reduction.cc +++ b/src/s_tir/schedule/primitive/reduction.cc @@ -40,7 +40,7 @@ class DecomposeReductionBlockReplacer : public StmtExprMutator { using StmtExprMutator::Mutate; using StmtExprMutator::Mutate_; UnchangedOr Mutate(ffi::AnyView value, InplaceMode inplace_mode) override { - if (value.as()) return ffi::Unchanged(); + if (value.as()) return ffi::Unchanged(); return StmtExprMutator::Mutate(value, inplace_mode); } @@ -69,7 +69,7 @@ class DecomposeReductionBlockReplacer : public StmtExprMutator { old_reduction_block_(std::move(old_reduction_block)) {} private: - UnchangedOr Mutate_(const ForNode* loop, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const ForObj* loop, InplaceMode inplace_mode) final { Stmt mutated_stmt = StmtExprMutator::Mutate_(loop, inplace_mode).ValueOrUnchanged(ffi::GetRef(loop)); if (loop == target_loop_.get()) { @@ -79,14 +79,14 @@ class DecomposeReductionBlockReplacer : public StmtExprMutator { } } - UnchangedOr Mutate_(const SBlockNode* block, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const SBlockObj* block, InplaceMode inplace_mode) final { if (block == old_reduction_block_.get()) { - auto p_new_block = ffi::make_object(*block); + auto p_new_block = ffi::make_object(*block); p_new_block->name_hint = p_new_block->name_hint + "_update"; p_new_block->init = std::nullopt; // Add write regions back to read regions in update block. ffi::Array new_reads; - std::unordered_set read_tensors; + std::unordered_set read_tensors; for (const TensorRegion& read_access : block->reads) { read_tensors.insert(read_access->source.as_or_throw().get()); } @@ -115,8 +115,8 @@ class DecomposeReductionBlockReplacer : public StmtExprMutator { class LoopHeightError : public ScheduleErrorContextObj { public: - static void CheckLoopHigherThanReduceLoops(const IRModule& mod, const SBlockNode* block, - const SBlockRealizeNode* realize, + static void CheckLoopHigherThanReduceLoops(const IRModule& mod, const SBlockObj* block, + const SBlockRealizeObj* realize, const ffi::Array& loops, const StmtSRef& loop_sref) { for (int i = 0, n = block->iter_vars.size(); i < n; ++i) { @@ -132,13 +132,13 @@ class LoopHeightError : public ScheduleErrorContextObj { break; } // loop_var of a higher loop shouldn't contain loop var - const Var& loop_var = higher_loop->StmtAs()->var; + const Var& loop_var = higher_loop->StmtAs()->var; auto walkfn = [v = loop_var.get()](const Var& var) -> ffi::Expected { return var.get() == v ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var)) : ffi::WalkResult::Advance(); }; if (ffi::StructuralWalk(binding, walkfn).has_value()) { - const ForNode* loop = TVM_SREF_TO_FOR(loop_sref); + const ForObj* loop = TVM_SREF_TO_FOR(loop_sref); throw MakeScheduleError(mod, ffi::GetRef(loop), ffi::GetRef(block)); } @@ -170,13 +170,13 @@ class LoopHeightError : public ScheduleErrorContextObj { }; PrimExpr RewriteInitPredicate(PrimExpr pred, - const std::unordered_set& discarded_loops) { + const std::unordered_set& discarded_loops) { if (IsOne(pred)) return IntImm::Bool(true); - if (const auto* and_node = pred.as()) { + if (const auto* and_node = pred.as()) { return RewriteInitPredicate(and_node->a, discarded_loops) && RewriteInitPredicate(and_node->b, discarded_loops); } - auto uses_discarded_loop = [&discarded_loops](const VarNode* var) { + auto uses_discarded_loop = [&discarded_loops](const VarObj* var) { return discarded_loops.count(var); }; auto walkfn = [&](const Var& var) -> ffi::Expected { @@ -199,11 +199,11 @@ StmtSRef DecomposeReduction(ScheduleState self, const StmtSRef& block_sref, * - generate corresponding init block and update block */ // Condition Checks and Information Collection - const SBlockNode* block = TVM_SREF_TO_SBLOCK(block_sref); - const ForNode* loop = TVM_SREF_TO_FOR(loop_sref); + const SBlockObj* block = TVM_SREF_TO_SBLOCK(block_sref); + const ForObj* loop = TVM_SREF_TO_FOR(loop_sref); // Get the outer loops from high to low ffi::Array loops = GetLoops(block_sref); - const SBlockRealizeNode* realize = GetSBlockRealize(self, block_sref).get(); + const SBlockRealizeObj* realize = GetSBlockRealize(self, block_sref).get(); StmtSRef scope_root_sref = GetScopeRoot(self, block_sref, /*require_stage_pipeline=*/false); if (self->enable_check) { @@ -252,7 +252,7 @@ StmtSRef DecomposeReduction(ScheduleState self, const StmtSRef& block_sref, } return ffi::Unchanged(); }; - auto init_block = ffi::make_object( + auto init_block = ffi::make_object( ffi::StructuralMap(block->init.value(), map_block_var) .as_or_throw()); init_block->name_hint = block->name_hint + "_init"; @@ -272,10 +272,10 @@ StmtSRef DecomposeReduction(ScheduleState self, const StmtSRef& block_sref, // Step 3. Scan loops not higher than the specified loop above the reduction block. // If the loop is used in the init block binding, then it is chosen. // Otherwise, it is discarded. - std::unordered_set discarded_loops; + std::unordered_set discarded_loops; std::vector chosen_loops; for (int i = static_cast(loops.size()) - 1; i >= 0; --i) { - const VarNode* loop_var = loops[i]->StmtAs()->var.get(); + const VarObj* loop_var = loops[i]->StmtAs()->var.get(); bool discarded = true; auto walkfn = [v = loop_var](const Var& var) -> ffi::Expected { return var.get() == v ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var)) @@ -300,7 +300,7 @@ StmtSRef DecomposeReduction(ScheduleState self, const StmtSRef& block_sref, } // Step 4. Derive the predicate for the init block realize. Omit conjunction clauses that // depend on discarded loops. - auto init_realize = ffi::make_object( + auto init_realize = ffi::make_object( RewriteInitPredicate(realize->predicate, discarded_loops), SBlock(init_block)); init_realize->iter_values = std::move(init_iter_values); // Step 5. Create new loops above init block @@ -328,7 +328,7 @@ StmtSRef DecomposeReduction(ScheduleState self, const StmtSRef& block_sref, }; body = ffi::StructuralMap(body, f_substitute).as_or_throw(); // Step 6. Mutate IR - const SBlockNode* old_scope_root = TVM_SREF_TO_SBLOCK(scope_root_sref); + const SBlockObj* old_scope_root = TVM_SREF_TO_SBLOCK(scope_root_sref); auto [new_scope_root, new_reduction_block] = DecomposeReductionBlockReplacer::Replace( ffi::GetRef(old_scope_root), ffi::GetRef(loop), body, ffi::GetRef(block)); @@ -670,9 +670,9 @@ class LoopPropertyError : public ScheduleErrorContextObj { ffi::Array LocationsOfInterest() const final { return {loop_}; } static void CheckLoopProperty(const ScheduleState& self, const ffi::Array& loops, - const ForNode* rf_loop, const SBlock& block, - const std::unordered_set& data_par_loop_vars, - const std::unordered_set& reduce_loop_vars) { + const ForObj* rf_loop, const SBlock& block, + const std::unordered_set& data_par_loop_vars, + const std::unordered_set& reduce_loop_vars) { ffi::Array children_of_outermost_loop = GetChildBlockRealizeOnSRefTree(self->stmt2ref.at(loops[0].get())); if (!children_of_outermost_loop[0]->block.same_as(block)) { @@ -717,8 +717,8 @@ class LoopPropertyError : public ScheduleErrorContextObj { * \param loops The loops to be analyzed * \return A mapping from loops to their corresponding loop vars */ -std::unordered_map GetLoopVar2LoopMap(const ffi::Array& loops) { - std::unordered_map loop_vars2loop; +std::unordered_map GetLoopVar2LoopMap(const ffi::Array& loops) { + std::unordered_map loop_vars2loop; loop_vars2loop.reserve(loops.size()); for (const For& loop : loops) { loop_vars2loop.insert_or_assign(loop->var.get(), loop); @@ -736,7 +736,7 @@ std::unordered_map GetLoopVar2LoopMap(const ffi::Array * \return The new created intermediate rfactor tensor */ ffi::Array CreateRFactorTensors(const ffi::Array& tensor_stores, - int factor_axis, const ForNode* rf_loop) { + int factor_axis, const ForObj* rf_loop) { ffi::Array rf_tensors; rf_tensors.reserve(tensor_stores.size()); for (const TensorStore& tensor_store : tensor_stores) { @@ -745,7 +745,7 @@ ffi::Array CreateRFactorTensors(const ffi::Array& tensor rf_shape.insert(rf_shape.begin() + factor_axis, Range::GetExtentFromBeginEnd(rf_loop->begin, rf_loop->end)); - ffi::ObjectPtr n = ffi::make_object(*tensor.ty().get()); + ffi::ObjectPtr n = ffi::make_object(*tensor.ty().get()); n->shape = rf_shape; rf_tensors.push_back(TensorVar(tensor->name + ".rf", TensorType(std::move(n)), tensor->loc)); } @@ -965,7 +965,7 @@ class RFactorBlockCreator : public BaseBlockCreator { explicit RFactorBlockCreator(SBlockRealize old_block_realize, For rf_loop, ffi::Array old_reduction_updates, te::CommReducer reducer, ffi::Array rf_tensors, - std::unordered_map loop_vars2loop, + std::unordered_map loop_vars2loop, int factor_axis, ffi::Array combiner_rhs) : BaseBlockCreator(std::move(old_block_realize), std::move(rf_loop), std::move(old_reduction_updates), std::move(reducer), @@ -1099,7 +1099,7 @@ class RFactorBlockCreator : public BaseBlockCreator { * \brief A mapping which maps a loop var to its corresponding For loop for all the reduction * block's outer loops */ - std::unordered_map loop_vars2loop_; + std::unordered_map loop_vars2loop_; /*! \brief The factor_axis specified for rfactor */ int factor_axis_; /*! \brief The RHS values of the reduction in the old block */ @@ -1109,7 +1109,7 @@ class RFactorBlockCreator : public BaseBlockCreator { * substitute the loop vars which appear in the bindings of some old block iters with the new * created block iters */ - std::unordered_map loop_var2block_binding_; + std::unordered_map loop_var2block_binding_; }; /*! @@ -1178,7 +1178,7 @@ class WriteBackBlockCreator : public BaseBlockCreator { void CreateRegion(const ffi::Array& tensor_loads, bool is_read) { ffi::Array& tensor_regions = is_read ? read_regions_ : write_regions_; for (const PrimExpr& expr : tensor_loads) { - const auto* tensor_load = expr.as(); + const auto* tensor_load = expr.as(); TVM_FFI_ICHECK(tensor_load != nullptr); ffi::Array region; region.reserve(tensor_load->indices.size()); @@ -1209,7 +1209,7 @@ Stmt CreateLoopOutsideRfactorBlock(SBlockRealize rf_block_realize, const ffi::Ar // Step 1. Create new loop vars. ffi::Array new_loops; - std::unordered_map new_loop_var_map; + std::unordered_map new_loop_var_map; new_loops.reserve(n_loops); new_loop_var_map.reserve(n_loops); for (const For& old_loop : loops) { @@ -1232,7 +1232,7 @@ Stmt CreateLoopOutsideRfactorBlock(SBlockRealize rf_block_realize, const ffi::Ar .as_or_throw()); } { - SBlockRealizeNode* p_rf_block_realize = rf_block_realize.CopyOnWrite(); + SBlockRealizeObj* p_rf_block_realize = rf_block_realize.CopyOnWrite(); p_rf_block_realize->iter_values = new_bindings; p_rf_block_realize->predicate = ffi::StructuralMap(rf_block_realize->predicate, map_loop_var) @@ -1242,7 +1242,7 @@ Stmt CreateLoopOutsideRfactorBlock(SBlockRealize rf_block_realize, const ffi::Ar // Step 3. Wrap `rf_block_realize` with outer loops. Stmt rf_body = rf_block_realize; for (int i = n_loops - 1; i >= 0; --i) { - ffi::ObjectPtr p_loop = ffi::make_object(*loops[i].get()); + ffi::ObjectPtr p_loop = ffi::make_object(*loops[i].get()); p_loop->var = new_loop_var_map.at(loops[i]->var.get()).as_or_throw(); p_loop->body = rf_body; rf_body = For(std::move(p_loop)); @@ -1256,7 +1256,7 @@ class BlockReplacer : public StmtExprMutator { using StmtExprMutator::Mutate; using StmtExprMutator::Mutate_; UnchangedOr Mutate(ffi::AnyView value, InplaceMode inplace_mode) override { - if (value.as()) return ffi::Unchanged(); + if (value.as()) return ffi::Unchanged(); return StmtExprMutator::Mutate(value, inplace_mode); } @@ -1284,8 +1284,8 @@ class BlockReplacer : public StmtExprMutator { */ static SBlock Replace(SBlock scope_root_block, Stmt rf_body, For outermost_loop, SBlockRealize wb_block_realize, SBlockRealize old_block_realize, - For rf_loop, std::unordered_set reduce_loop_vars, - std::unordered_map loop_vars2loop, + For rf_loop, std::unordered_set reduce_loop_vars, + std::unordered_map loop_vars2loop, const ffi::Array& rf_tensors) { auto replacer = ffi::make_object( std::move(rf_body), std::move(outermost_loop), std::move(wb_block_realize), @@ -1294,7 +1294,7 @@ class BlockReplacer : public StmtExprMutator { SBlock new_scope_root = replacer->Mutate(scope_root_block, InplaceMode::kAllow) .ValueOrUnchanged(std::move(scope_root_block)) .as_or_throw(); - SBlockNode* p = new_scope_root.CopyOnWrite(); + SBlockObj* p = new_scope_root.CopyOnWrite(); for (const TensorVar& rf_tensor : rf_tensors) { p->alloc_tensors.push_back(rf_tensor); } @@ -1303,8 +1303,8 @@ class BlockReplacer : public StmtExprMutator { explicit BlockReplacer(Stmt rf_body, For outermost_loop, SBlockRealize wb_block_realize, SBlockRealize old_block_realize, For rf_loop, - std::unordered_set reduce_loop_vars, - std::unordered_map loop_vars2loop) + std::unordered_set reduce_loop_vars, + std::unordered_map loop_vars2loop) : rf_body_(std::move(rf_body)), outermost_loop_(std::move(outermost_loop)), wb_block_realize_(std::move(wb_block_realize)), @@ -1314,7 +1314,7 @@ class BlockReplacer : public StmtExprMutator { loop_vars2loop_(std::move(loop_vars2loop)) {} private: - UnchangedOr Mutate_(const ForNode* loop, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const ForObj* loop, InplaceMode inplace_mode) final { // Step 1. Check whether this loop is outside the reduction block. Given that we've made sure // that the scope root block has stage-pipeline property, if this loop is not outside the // reduction block, there's no need to recursively mutate. @@ -1331,10 +1331,10 @@ class BlockReplacer : public StmtExprMutator { // should be kept outside the write-back block. Otherwise it shouldn't. if (loop == rf_loop_.get() || !reduce_loop_vars_.count(loop->var.get())) { if (inplace_mode == InplaceMode::kAllow) { - const_cast(loop)->body = std::move(body); + const_cast(loop)->body = std::move(body); body = ffi::GetRef(loop); } else { - auto copy = ffi::make_object(*loop); + auto copy = ffi::make_object(*loop); copy->body = std::move(body); body = For(std::move(copy)); } @@ -1345,9 +1345,8 @@ class BlockReplacer : public StmtExprMutator { return loop == outermost_loop_.get() ? SeqStmt({rf_body_, body}) : body; } - UnchangedOr Mutate_(const SBlockRealizeNode* block_realize, - InplaceMode inplace_mode) final { - // Due to the visitor's behavior on ForNode, this block-realize must be the reduction block's + UnchangedOr Mutate_(const SBlockRealizeObj* block_realize, InplaceMode inplace_mode) final { + // Due to the visitor's behavior on ForObj, this block-realize must be the reduction block's // block-realize. And we directly return the new `wb_block_realize`. TVM_FFI_ICHECK_EQ(block_realize, old_block_realize_.get()); return wb_block_realize_; @@ -1358,8 +1357,8 @@ class BlockReplacer : public StmtExprMutator { SBlockRealize wb_block_realize_; SBlockRealize old_block_realize_; For rf_loop_; - std::unordered_set reduce_loop_vars_; - std::unordered_map loop_vars2loop_; + std::unordered_set reduce_loop_vars_; + std::unordered_map loop_vars2loop_; }; StmtSRef RFactor(ScheduleState self, const StmtSRef& rf_loop_sref, int factor_axis) { @@ -1376,21 +1375,21 @@ StmtSRef RFactor(ScheduleState self, const StmtSRef& rf_loop_sref, int factor_ax if (self->enable_check) { CheckReductionBlock(self, block_sref, scope_root); } - const ForNode* rf_loop = TVM_SREF_TO_FOR(rf_loop_sref); + const ForObj* rf_loop = TVM_SREF_TO_FOR(rf_loop_sref); if (rf_loop->kind != ForKind::kDefault) { throw MakeScheduleError(self->mod, ffi::GetRef(rf_loop)); } // Step 2. Collect loop vars that are touched by data parallel block iters and reduction block // iters, respectively. - std::unordered_set data_par_loop_vars; - std::unordered_set reduce_loop_vars; + std::unordered_set data_par_loop_vars; + std::unordered_set reduce_loop_vars; GetVarsTouchedByBlockIters(block_realize, &data_par_loop_vars, &reduce_loop_vars); // Step 3. Collect the loops of the reduction block. Construct a mapping from loops to // corresponding loop vars. ffi::Array loops = LoopSRefs2Loops(GetLoops(block_sref)); - std::unordered_map loop_vars2loop = GetLoopVar2LoopMap(loops); + std::unordered_map loop_vars2loop = GetLoopVar2LoopMap(loops); // Step 4. Check four properties that the loops should have: // - the rfactor loop cannot be touched by any data parallel block iter; @@ -1452,7 +1451,7 @@ StmtSRef RFactor(ScheduleState self, const StmtSRef& rf_loop_sref, int factor_ax // ***************************************************** // Step 1. Substitute the old scope root block with the new scope root block. - SBlock old_scope_root_block = ffi::GetRef(scope_root->StmtAs()); + SBlock old_scope_root_block = ffi::GetRef(scope_root->StmtAs()); SBlock new_scope_root_block = BlockReplacer::Replace( old_scope_root_block, rf_body, loops[0], wb_block_creator.new_block_realize_, block_realize, ffi::GetRef(rf_loop), reduce_loop_vars, loop_vars2loop, rf_tensors); diff --git a/src/s_tir/schedule/primitive/reorder_block_iter_var.cc b/src/s_tir/schedule/primitive/reorder_block_iter_var.cc index e670d560aeda..ea9e1d7141ae 100644 --- a/src/s_tir/schedule/primitive/reorder_block_iter_var.cc +++ b/src/s_tir/schedule/primitive/reorder_block_iter_var.cc @@ -58,18 +58,18 @@ class BlockIterVarRewriter : public StmtExprMutator { using StmtExprMutator::Mutate; using StmtExprMutator::Mutate_; UnchangedOr Mutate(ffi::AnyView value, InplaceMode inplace_mode) override { - if (value.as()) return ffi::Unchanged(); + if (value.as()) return ffi::Unchanged(); return StmtExprMutator::Mutate(value, inplace_mode); } ffi::Map block_map; - explicit BlockIterVarRewriter(const SBlockNode* block_n, std::vector order) + explicit BlockIterVarRewriter(const SBlockObj* block_n, std::vector order) : order_(std::move(order)), block_to_rewrite(block_n) {} private: std::vector order_; - const SBlockNode* block_to_rewrite; - UnchangedOr Mutate_(const SBlockRealizeNode* op, InplaceMode inplace_mode) final { + const SBlockObj* block_to_rewrite; + UnchangedOr Mutate_(const SBlockRealizeObj* op, InplaceMode inplace_mode) final { if (op->block.get() == block_to_rewrite) { bool inplace_block = inplace_mode == InplaceMode::kAllow && op->block.unique(); SBlock block = op->block; @@ -81,20 +81,20 @@ class BlockIterVarRewriter : public StmtExprMutator { } SBlock new_block = block; if (inplace_block) { - const_cast(block.get())->iter_vars = std::move(new_iter_vars); + const_cast(block.get())->iter_vars = std::move(new_iter_vars); } else { - auto copy = ffi::make_object(*block.get()); + auto copy = ffi::make_object(*block.get()); copy->iter_vars = std::move(new_iter_vars); new_block = SBlock(std::move(copy)); } block_map.Set(block, new_block); if (inplace_mode == InplaceMode::kAllow) { - auto* writable = const_cast(op); + auto* writable = const_cast(op); writable->block = std::move(new_block); writable->iter_values = std::move(new_iter_values); return ffi::Unchanged(); } else { - auto copy = ffi::make_object(*op); + auto copy = ffi::make_object(*op); copy->block = std::move(new_block); copy->iter_values = std::move(new_iter_values); return SBlockRealize(std::move(copy)); @@ -107,7 +107,7 @@ class BlockIterVarRewriter : public StmtExprMutator { void ReorderBlockIterVar(ScheduleState self, const StmtSRef& block_sref, const ffi::Array& new_order) { - const SBlockNode* block_n = TVM_SREF_TO_SBLOCK(block_sref); + const SBlockObj* block_n = TVM_SREF_TO_SBLOCK(block_sref); std::vector new_order_vec; for (int64_t x : new_order) { new_order_vec.push_back(static_cast(x)); @@ -126,10 +126,10 @@ void ReorderBlockIterVar(ScheduleState self, const StmtSRef& block_sref, } // find parent block - const SBlockNode* parent_block_n = nullptr; + const SBlockObj* parent_block_n = nullptr; const StmtSRefObj* p = block_sref.get()->parent; while (p != nullptr) { - if (p->stmt->IsInstance()) { + if (p->stmt->IsInstance()) { parent_block_n = TVM_SREF_TO_SBLOCK(ffi::GetRef(p)); break; } diff --git a/src/s_tir/schedule/primitive/rolling_buffer.cc b/src/s_tir/schedule/primitive/rolling_buffer.cc index e72d8eac6453..fa76be1b7db3 100644 --- a/src/s_tir/schedule/primitive/rolling_buffer.cc +++ b/src/s_tir/schedule/primitive/rolling_buffer.cc @@ -99,14 +99,14 @@ class RollingBufferDependencyError : public ScheduleErrorContextObj { SBlockScope scope = self->GetSBlockScope(scope_root_sref); for (const Dependency& producers : scope->GetDepsByDst(block_sref)) { if (!(producers->kind == DepKind::kRAW)) { - const SBlockNode* block = TVM_SREF_TO_SBLOCK(block_sref); + const SBlockObj* block = TVM_SREF_TO_SBLOCK(block_sref); throw MakeScheduleError(self->mod, ffi::GetRef(block)); } } for (const Dependency& consumers : scope->GetDepsBySrc(block_sref)) { if (!(consumers->kind == DepKind::kRAW)) { - const SBlockNode* block = TVM_SREF_TO_SBLOCK(block_sref); + const SBlockObj* block = TVM_SREF_TO_SBLOCK(block_sref); throw MakeScheduleError(self->mod, ffi::GetRef(block)); } @@ -176,7 +176,7 @@ class RollingBufferInfoCollector { const TensorRegion& tensor_region) { RollingBufferInfoCollector collector; if (!collector.MatchRollingBuffer(block_sref, tensor_region)) { - const SBlockNode* block = TVM_SREF_TO_SBLOCK(block_sref); + const SBlockObj* block = TVM_SREF_TO_SBLOCK(block_sref); throw MakeScheduleError(mod, ffi::GetRef(block), tensor_region); } @@ -196,8 +196,8 @@ class RollingBufferInfoCollector { auto divisor = 1; ffi::Optional iter_var; - if (const auto* div = bound->begin.as()) { - const auto* mul = div->a.as(); + if (const auto* div = bound->begin.as()) { + const auto* mul = div->a.as(); if (!mul) return false; iter_var = mul->a.as(); auto stride_imm = mul->b.as(); @@ -213,7 +213,7 @@ class RollingBufferInfoCollector { if (!divisor_value.has_value() || !stride_value.has_value()) return false; divisor = *divisor_value; stride = std::ceil(static_cast(*stride_value) / divisor); - } else if (const auto* mul = bound->begin.as()) { + } else if (const auto* mul = bound->begin.as()) { // The bound is the iter var multiplied by the stride iter_var = mul->a.as(); auto stride_imm = mul->b.as(); @@ -253,7 +253,7 @@ class RollingBufferInfoCollector { ffi::Optional roll_iter_var; int roll_axis = 0; for (const tirx::StmtSRef& loop_sref : loop_srefs) { - auto loop_var = loop_sref->StmtAs()->var; + auto loop_var = loop_sref->StmtAs()->var; auto it{std::find_if( bound_iter_vars.begin(), bound_iter_vars.end(), @@ -271,7 +271,7 @@ class RollingBufferInfoCollector { } ffi::Array new_shape = tensor.shape(); new_shape.Set(roll_axis, region[roll_axis]->extent); - auto new_tensor_type = ffi::make_object(*tensor.ty().get()); + auto new_tensor_type = ffi::make_object(*tensor.ty().get()); new_tensor_type->shape = new_shape; TensorVar new_tensor = TensorVar(tensor->name, TensorType(std::move(new_tensor_type)), tensor->loc); @@ -329,12 +329,12 @@ class RollingBufferRewriter : public StmtExprMutator { *indices = std::move(new_indices); } - UnchangedOr Mutate_(const SBlockNode* block, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const SBlockObj* block, InplaceMode inplace_mode) final { SBlock old_stmt = ffi::GetRef(block); SBlock stmt = StmtExprMutator::Mutate_(block, inplace_mode) .ValueOrUnchanged(ffi::GetRef(block)) .as_or_throw(); - SBlockNode* n = stmt.CopyOnWrite(); + SBlockObj* n = stmt.CopyOnWrite(); if (block == scope_sref_->stmt) { ffi::Array new_alloc_tensors; for (const TensorVar& tensor : stmt->alloc_tensors) { @@ -374,7 +374,7 @@ class RollingBufferRewriter : public StmtExprMutator { return stmt; } - UnchangedOr Mutate_(const SBlockRealizeNode* realize, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const SBlockRealizeObj* realize, InplaceMode inplace_mode) final { SBlockRealize stmt = StmtExprMutator::Mutate_(realize, inplace_mode) .ValueOrUnchanged(ffi::GetRef(realize)) .as_or_throw(); @@ -395,18 +395,18 @@ class RollingBufferRewriter : public StmtExprMutator { GE(term_2, info_->axis_overlaps[i])))); } } - SBlockRealizeNode* n = stmt.CopyOnWrite(); + SBlockRealizeObj* n = stmt.CopyOnWrite(); n->predicate = condition; } return stmt; } - UnchangedOr Mutate_(const TensorStoreNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const TensorStoreObj* op, InplaceMode inplace_mode) final { TensorStore stmt = StmtExprMutator::Mutate_(op, inplace_mode) .ValueOrUnchanged(ffi::GetRef(op)) .as_or_throw(); if (stmt->dest.as_or_throw().same_as(info_->old_tensor)) { - TensorStoreNode* n = stmt.CopyOnWrite(); + TensorStoreObj* n = stmt.CopyOnWrite(); TensorVar tensor = n->dest.as_or_throw(); RewriteTensorAccess(&tensor, &n->indices); n->dest = std::move(tensor); @@ -416,7 +416,7 @@ class RollingBufferRewriter : public StmtExprMutator { return stmt; } - UnchangedOr Mutate_(const TensorLoadNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const TensorLoadObj* op, InplaceMode inplace_mode) final { TensorLoad stmt = StmtExprMutator::Mutate_(op, inplace_mode) .ValueOrUnchanged(ffi::GetRef(op)) .as_or_throw(); @@ -466,7 +466,7 @@ void RollingBuffer(ScheduleState self, const StmtSRef& block_sref, int write_ten ffi::Array consumers_sref = GetConsumers(self, block_sref); consumers_sref.push_back(block_sref); StmtSRef lca = GetSRefLowestCommonAncestor(consumers_sref); - if (!lca->StmtAs()) { + if (!lca->StmtAs()) { throw MakeScheduleError( self->mod, tensor_region->source.as_or_throw(), block); } @@ -477,7 +477,7 @@ void RollingBuffer(ScheduleState self, const StmtSRef& block_sref, int write_ten if (stmt == lca) { break; } - For cur_loop = ffi::GetRef(stmt->StmtAs()); + For cur_loop = ffi::GetRef(stmt->StmtAs()); Range range = Range(cur_loop->begin, cur_loop->end); dom_map.Set(cur_loop->var, sym::IntSet::FromRange(range)); } diff --git a/src/s_tir/schedule/primitive/sampling.cc b/src/s_tir/schedule/primitive/sampling.cc index 3a20bc95e8cd..d4976d2bc60f 100644 --- a/src/s_tir/schedule/primitive/sampling.cc +++ b/src/s_tir/schedule/primitive/sampling.cc @@ -313,9 +313,9 @@ std::vector SamplePerfectTile(LinearCongruentialEngine::TRandState* ran const tirx::StmtSRef& loop_sref, int32_t n_splits, int32_t max_innermost_factor, ffi::Optional>* decision) { - const ForNode* loop = TVM_SREF_TO_FOR(loop_sref); + const ForObj* loop = TVM_SREF_TO_FOR(loop_sref); PrimExpr extent_imm_expr = Range::GetExtentFromBeginEnd(loop->begin, loop->end); - const auto* extent_imm = extent_imm_expr.as(); + const auto* extent_imm = extent_imm_expr.as(); auto extent = extent_imm ? extent_imm->value.as() : std::nullopt; std::vector result; if (!extent.has_value()) { @@ -376,9 +376,9 @@ std::vector SamplePartitionedTile(LinearCongruentialEngine::TRandState* const tirx::StmtSRef& loop_sref, int32_t n_splits, int32_t partition_pos, int32_t innerpart_factor, ffi::Optional>* decision) { - const ForNode* loop = TVM_SREF_TO_FOR(loop_sref); + const ForObj* loop = TVM_SREF_TO_FOR(loop_sref); PrimExpr extent_imm_expr = Range::GetExtentFromBeginEnd(loop->begin, loop->end); - const auto* extent_imm = extent_imm_expr.as(); + const auto* extent_imm = extent_imm_expr.as(); auto extent = extent_imm ? extent_imm->value.as() : std::nullopt; std::vector result; if (!extent.has_value() || *extent % innerpart_factor != 0) { diff --git a/src/s_tir/schedule/schedule.cc b/src/s_tir/schedule/schedule.cc index 6924c6326895..95722ebf62ae 100644 --- a/src/s_tir/schedule/schedule.cc +++ b/src/s_tir/schedule/schedule.cc @@ -27,19 +27,19 @@ using namespace tvm::tirx; TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; refl::ObjectDef(); - SBlockRVNode::RegisterReflection(); - LoopRVNode::RegisterReflection(); + SBlockRVObj::RegisterReflection(); + LoopRVObj::RegisterReflection(); } /**************** Constructor ****************/ -SBlockRV::SBlockRV() { this->data_ = ffi::make_object(); } +SBlockRV::SBlockRV() { this->data_ = ffi::make_object(); } -LoopRV::LoopRV() { this->data_ = ffi::make_object(); } +LoopRV::LoopRV() { this->data_ = ffi::make_object(); } /**************** GetSRef ****************/ -StmtSRef ScheduleObj::GetSRef(const StmtNode* stmt) const { +StmtSRef ScheduleObj::GetSRef(const StmtObj* stmt) const { ScheduleState state = this->state(); auto it = state->stmt2ref.find(stmt); if (it == state->stmt2ref.end()) { diff --git a/src/s_tir/schedule/state.cc b/src/s_tir/schedule/state.cc index 5ee0d2c23b27..d6aec5ff3670 100644 --- a/src/s_tir/schedule/state.cc +++ b/src/s_tir/schedule/state.cc @@ -149,9 +149,9 @@ bool ProducerCoversConsumer(const ffi::Array& tensor_shape, * \param sref The sref to be updated * \param new_stmt The statement that replaces the statement inside the sref */ -void UpdateSRef(ScheduleStateObj* self, StmtSRefObj* sref, const StmtNode* new_stmt) { - TVM_FFI_ICHECK(new_stmt->IsInstance() || new_stmt->IsInstance()); - const StmtNode* old_stmt = sref->stmt; +void UpdateSRef(ScheduleStateObj* self, StmtSRefObj* sref, const StmtObj* new_stmt) { + TVM_FFI_ICHECK(new_stmt->IsInstance() || new_stmt->IsInstance()); + const StmtObj* old_stmt = sref->stmt; TVM_FFI_ICHECK_NE(new_stmt, old_stmt); self->stmt2ref.insert_or_assign(new_stmt, ffi::GetRef(sref)); self->stmt2ref.erase(sref->stmt); @@ -165,7 +165,7 @@ class SBlockInfoCollector : public StmtExprVisitor { using StmtExprVisitor::Visit_; ffi::Optional Visit(ffi::AnyView value) override { - if (value.as()) return std::nullopt; + if (value.as()) return std::nullopt; return StmtExprVisitor::Visit(value); } @@ -185,7 +185,7 @@ class SBlockInfoCollector : public StmtExprVisitor { * \param stmt A for-loop statement or a block statement * \return A sref to the stmt */ - void PushSRef(const StmtNode* stmt) { srefs_.push_back(self_->stmt2ref.at(stmt)); } + void PushSRef(const StmtObj* stmt) { srefs_.push_back(self_->stmt2ref.at(stmt)); } /*! \brief Pop the top of the scope */ StmtSRef PopSRef() { @@ -203,7 +203,7 @@ class SBlockInfoCollector : public StmtExprVisitor { if (is_root_block) { // If the block doesn't have outer loops and BlockRealize, // then we set the affine binding flag as true only if the block has no block vars - const SBlockNode* block = TVM_SREF_TO_SBLOCK(scope_root); + const SBlockObj* block = TVM_SREF_TO_SBLOCK(scope_root); if (block->iter_vars.empty()) info.affine_binding = true; } else { info.affine_binding = @@ -227,7 +227,7 @@ class SBlockInfoCollector : public StmtExprVisitor { block_reads_unbound.reserve(child_block_srefs.size()); block_writes_unbound.reserve(child_block_srefs.size()); for (const StmtSRef& block_sref : child_block_srefs) { - const SBlockNode* block = TVM_SREF_TO_SBLOCK(block_sref); + const SBlockObj* block = TVM_SREF_TO_SBLOCK(block_sref); ffi::Map binding = GetBindings(block2realize_.at(block)); auto f_substitute = [&binding](const Var& var) -> ffi::Expected> { if (auto repl = binding.Get(var)) return ffi::Any(*std::move(repl)); @@ -270,7 +270,7 @@ class SBlockInfoCollector : public StmtExprVisitor { for (const auto& kv : info.scope->dst2deps) { const StmtSRef& consumer_block_sref = kv.first; const ffi::Array& deps = kv.second; - const SBlockNode* consumer_block = TVM_SREF_TO_SBLOCK(consumer_block_sref); + const SBlockObj* consumer_block = TVM_SREF_TO_SBLOCK(consumer_block_sref); const SBlockRealize& consumer_realize = block2realize_.at(consumer_block); bool& region_cover = self_->block_info.at(consumer_block_sref).region_cover = true; // Step 2.1. Extract the path to the scope root @@ -372,17 +372,21 @@ class SBlockInfoCollector : public StmtExprVisitor { return stage_pipeline; } - ffi::Optional Visit_(const ForNode* loop) final { - analyzer_->Bind(loop->var, Range(loop->begin, loop->end)); + ffi::Optional Visit_(const ForObj* loop) final { + { + auto context = analyzer_->GetConstraintContext(); + context->Append(sym::RangeConstraint(loop->var, Range(loop->begin, loop->end))); + analyzer_->SyncConstraintContext(context); + } PushSRef(loop); TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Visit(loop->body)); PopSRef(); return std::nullopt; } - ffi::Optional Visit_(const SBlockRealizeNode* realize) final { + ffi::Optional Visit_(const SBlockRealizeObj* realize) final { block_frames_.emplace_back(); - const SBlockNode* block = realize->block.get(); + const SBlockObj* block = realize->block.get(); block2realize_.emplace(block, ffi::GetRef(realize)); // Recursive visit PushSRef(block); @@ -396,8 +400,8 @@ class SBlockInfoCollector : public StmtExprVisitor { return std::nullopt; } - ffi::Optional Visit_(const SeqStmtNode* seq_stmt) final { - // Set `seq_index` information for SeqStmtNode + ffi::Optional Visit_(const SeqStmtObj* seq_stmt) final { + // Set `seq_index` information for SeqStmtObj TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(seq_stmt)); SetSeqIndexInChildren(self_->stmt2ref, seq_stmt); return std::nullopt; @@ -408,7 +412,7 @@ class SBlockInfoCollector : public StmtExprVisitor { /*! \brief The stack frame used to indicate the current scope */ std::vector srefs_; /*! \brief The BlockRealize corresponding to blocks */ - std::unordered_map block2realize_; + std::unordered_map block2realize_; /*! \brief The stack frames of blocks in the DFS visit. */ std::vector> block_frames_; /*! \brief The auxiliary analyzer */ @@ -428,7 +432,7 @@ ScheduleState::ScheduleState(IRModule mod, int debug_mask, bool enable_check) { n->debug_mask = debug_mask; // Set `n->enable_check` n->enable_check = enable_check; - std::unordered_map stmt2ref; + std::unordered_map stmt2ref; // Set `n->stmt2ref` self->stmt2ref = SRefTreeCreator::Create(self->mod); // Set `n->block_info` @@ -483,7 +487,7 @@ struct ReuseInfo { * \brief Kind 1. Intact reuse. If a stmt is in `intact`, it means its corresponding * sref is reused and it is intact reuse. */ - std::unordered_set intact; + std::unordered_set intact; /*! * \brief Kind 2.1. Loop sref reuse * If the loop var of a loop is in `loop_sref_possible_reuse`, @@ -492,13 +496,13 @@ struct ReuseInfo { * \note For each loop var in `loop_sref_possible_reuse`, it is possible that `src_stmt` doesn't * contain a loop that uses this loop var, and that is the reason why it is named "possible". */ - std::unordered_set loop_sref_possible_reuse; + std::unordered_set loop_sref_possible_reuse; /*! * \brief Kind 2.2. SBlock sref reuse. * Maps an old SBlock in `src_stmt` to a new block in `tgt_stmt`, * indicating the sref to the old block should be reused in the sref to the new block. */ - std::unordered_map block_sref_reuse; + std::unordered_map block_sref_reuse; }; /*! @@ -518,7 +522,7 @@ class ReuseCollector : public StmtExprVisitor { using StmtExprVisitor::Visit_; ffi::Optional Visit(ffi::AnyView value) override { - if (value.as()) return std::nullopt; + if (value.as()) return std::nullopt; return StmtExprVisitor::Visit(value); } @@ -536,7 +540,7 @@ class ReuseCollector : public StmtExprVisitor { explicit ReuseCollector(const ScheduleStateObj* self) : self_(self) {} private: - ffi::Optional Visit_(const ForNode* op) final { + ffi::Optional Visit_(const ForObj* op) final { if (self_->stmt2ref.count(op)) { intact_.push_back(op); } else { @@ -547,7 +551,7 @@ class ReuseCollector : public StmtExprVisitor { return std::nullopt; } - ffi::Optional Visit_(const SBlockNode* op) final { + ffi::Optional Visit_(const SBlockObj* op) final { if (self_->stmt2ref.count(op)) { intact_.push_back(op); } else { @@ -559,9 +563,9 @@ class ReuseCollector : public StmtExprVisitor { /*! \brief The schedule state to be worked on */ const ScheduleStateObj* self_; /*! \brief The intact statements we have collected along the way of visiting */ - std::vector intact_; + std::vector intact_; /*! \brief The loop variable we collected in the tgt_stmt */ - std::vector loop_vars_; + std::vector loop_vars_; }; /*! @@ -577,7 +581,7 @@ class SRefTreePruner : public StmtExprVisitor { using StmtExprVisitor::Visit_; ffi::Optional Visit(ffi::AnyView value) override { - if (value.as()) return std::nullopt; + if (value.as()) return std::nullopt; return StmtExprVisitor::Visit(value); } @@ -604,7 +608,7 @@ class SRefTreePruner : public StmtExprVisitor { : self_(self), reuse_info_(reuse_info) {} private: - ffi::Optional Visit_(const ForNode* op) final { + ffi::Optional Visit_(const ForObj* op) final { if (reuse_info_.intact.count(op)) { return std::nullopt; } @@ -614,7 +618,7 @@ class SRefTreePruner : public StmtExprVisitor { << ffi::GetRef(op); StmtSRef& sref = it->second; // Detect reuse - const VarNode* loop_var = op->var.get(); + const VarObj* loop_var = op->var.get(); if (reuse_info_.loop_sref_possible_reuse.count(loop_var)) { // sref can be reused reused_srefs_.emplace(loop_var, std::move(sref)); @@ -627,7 +631,7 @@ class SRefTreePruner : public StmtExprVisitor { return Visit(op->body); } - ffi::Optional Visit_(const SBlockNode* op) final { + ffi::Optional Visit_(const SBlockObj* op) final { if (reuse_info_.intact.count(op)) { return std::nullopt; } @@ -639,7 +643,7 @@ class SRefTreePruner : public StmtExprVisitor { // Detect reuse const auto& sref_reuse = reuse_info_.block_sref_reuse; if (auto reuse_it = sref_reuse.find(op); reuse_it != sref_reuse.end()) { - const SBlockNode* to_reuse = reuse_it->second; + const SBlockObj* to_reuse = reuse_it->second; // sref can be reused reused_srefs_.emplace(to_reuse, std::move(sref)); } else { @@ -678,7 +682,7 @@ class SRefUpdater : public StmtExprVisitor { using StmtExprVisitor::Visit_; ffi::Optional Visit(ffi::AnyView value) override { - if (value.as()) return std::nullopt; + if (value.as()) return std::nullopt; return StmtExprVisitor::Visit(value); } @@ -695,7 +699,7 @@ class SRefUpdater : public StmtExprVisitor { reused_srefs_(reused_srefs) {} private: - ffi::Optional Visit_(const ForNode* op) final { + ffi::Optional Visit_(const ForObj* op) final { auto existing = self_->stmt2ref.find(op); // Detect intact reuse if (existing != self_->stmt2ref.end()) { @@ -726,7 +730,7 @@ class SRefUpdater : public StmtExprVisitor { return std::nullopt; } - ffi::Optional Visit_(const SBlockNode* op) final { + ffi::Optional Visit_(const SBlockObj* op) final { auto existing = self_->stmt2ref.find(op); // Detect intact if (existing != self_->stmt2ref.end()) { @@ -759,7 +763,7 @@ class SRefUpdater : public StmtExprVisitor { return std::nullopt; } - ffi::Optional Visit_(const SeqStmtNode* seq_stmt) final { + ffi::Optional Visit_(const SeqStmtObj* seq_stmt) final { TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(seq_stmt)); SetSeqIndexInChildren(self_->stmt2ref, seq_stmt); return std::nullopt; @@ -805,14 +809,14 @@ class ChildReplacer : public StmtExprMutator { using StmtExprMutator::Mutate; using StmtExprMutator::Mutate_; - static Stmt Replace(const StmtNode* parent_stmt, const StmtNode* child_src_stmt, + static Stmt Replace(const StmtObj* parent_stmt, const StmtObj* child_src_stmt, const Stmt& child_tgt_stmt, int seq_index, bool allow_copy_on_write) { // Check the invariant - TVM_FFI_ICHECK(child_src_stmt->IsInstance() || // - child_src_stmt->IsInstance()); - TVM_FFI_ICHECK(child_tgt_stmt->IsInstance() || // - child_tgt_stmt->IsInstance() || // - child_tgt_stmt->IsInstance()); + TVM_FFI_ICHECK(child_src_stmt->IsInstance() || // + child_src_stmt->IsInstance()); + TVM_FFI_ICHECK(child_tgt_stmt->IsInstance() || // + child_tgt_stmt->IsInstance() || // + child_tgt_stmt->IsInstance()); auto replacer = ffi::make_object(child_src_stmt, child_tgt_stmt, seq_index); // ScheduleState has proved the complete ancestor chain. Check the borrowed // parent before creating an owning result handle. @@ -821,31 +825,31 @@ class ChildReplacer : public StmtExprMutator { return replacer->CopyOnWriteAndMutate(parent_stmt, mode); } - explicit ChildReplacer(const StmtNode* src_stmt, const Stmt& tgt_stmt, int seq_index) + explicit ChildReplacer(const StmtObj* src_stmt, const Stmt& tgt_stmt, int seq_index) : src_stmt_(src_stmt), tgt_stmt_(tgt_stmt), seq_index_(seq_index) {} private: UnchangedOr Mutate(ffi::AnyView value, InplaceMode inplace_mode) final { - if (value.as()) return ffi::Unchanged(); - if (value.as() == src_stmt_) return ffi::Any(tgt_stmt_); + if (value.as()) return ffi::Unchanged(); + if (value.as() == src_stmt_) return ffi::Any(tgt_stmt_); return StmtExprMutator::Mutate(value, inplace_mode); } // Skipping sibling blocks and loops other than `src_stmt_` - UnchangedOr Mutate_(const SBlockNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const SBlockObj* op, InplaceMode inplace_mode) final { return ffi::Unchanged(); } - UnchangedOr Mutate_(const ForNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const ForObj* op, InplaceMode inplace_mode) final { return ffi::Unchanged(); } - UnchangedOr Mutate_(const SeqStmtNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const SeqStmtObj* op, InplaceMode inplace_mode) final { int i = this->seq_index_; int n = static_cast(op->seq.size()); if (0 <= i && i < n) { const Stmt& stmt = op->seq[i]; ffi::Optional new_stmt = std::nullopt; - const StmtNode* src_stmt = this->src_stmt_; + const StmtObj* src_stmt = this->src_stmt_; // `stmt` can be For or BlockRealize // `src_stmt` can be For or Block // so the match from `stmt` to `src_stmt` can be @@ -854,12 +858,12 @@ class ChildReplacer : public StmtExprMutator { if (stmt.get() == src_stmt) { // Case 1. src_stmt is For, stmt is For new_stmt = tgt_stmt_; - } else if (const auto* realize = stmt.as()) { + } else if (const auto* realize = stmt.as()) { // Case 2. stmt is BlockRealize, src_stmt is Block if (realize->block.get() == src_stmt) { - const auto* tgt_block = TVM_TYPE_AS(tgt_stmt_, SBlockNode); - ffi::ObjectPtr new_realize = - ffi::make_object(*realize); + const auto* tgt_block = TVM_TYPE_AS(tgt_stmt_, SBlockObj); + ffi::ObjectPtr new_realize = + ffi::make_object(*realize); new_realize->block = ffi::GetRef(tgt_block); new_stmt = SBlockRealize(std::move(new_realize)); } @@ -867,10 +871,10 @@ class ChildReplacer : public StmtExprMutator { // Move new_stmt to position i if (new_stmt.has_value()) { if (inplace_mode == InplaceMode::kAllow) { - const_cast(op)->seq.Set(i, new_stmt.value()); + const_cast(op)->seq.Set(i, new_stmt.value()); return ffi::Unchanged(); } else { - auto copy = ffi::make_object(*op); + auto copy = ffi::make_object(*op); copy->seq.Set(i, new_stmt.value()); return SeqStmt(std::move(copy)); } @@ -879,30 +883,30 @@ class ChildReplacer : public StmtExprMutator { return StmtExprMutator::Mutate_(op, inplace_mode); } - Stmt CopyOnWriteAndMutate(const StmtNode* parent_stmt, InplaceMode inplace_mode) { + Stmt CopyOnWriteAndMutate(const StmtObj* parent_stmt, InplaceMode inplace_mode) { // Step 1. Copy-on-write the `parent_stmt` and extract its `body`, // where `body` means the body of either a block or a loop // Step 2. Mutate the `block/loop->body`, searching for `child_old_stmt` // and replace it with `child_tgt_stmt` - if (parent_stmt->IsInstance()) { - auto* block = static_cast(parent_stmt); + if (parent_stmt->IsInstance()) { + auto* block = static_cast(parent_stmt); if (inplace_mode == InplaceMode::kAllow) { - auto* writable = const_cast(block); + auto* writable = const_cast(block); writable->body = this->Mutate(block->body, inplace_mode).ValueOrUnchanged(block->body); return ffi::GetRef(block); } else { - auto copy = ffi::make_object(*block); + auto copy = ffi::make_object(*block); copy->body = this->Mutate(copy->body, inplace_mode).ValueOrUnchanged(copy->body); return SBlock(std::move(copy)); } - } else if (parent_stmt->IsInstance()) { - auto* loop = static_cast(parent_stmt); + } else if (parent_stmt->IsInstance()) { + auto* loop = static_cast(parent_stmt); if (inplace_mode == InplaceMode::kAllow) { - auto* writable = const_cast(loop); + auto* writable = const_cast(loop); writable->body = this->Mutate(loop->body, inplace_mode).ValueOrUnchanged(loop->body); return ffi::GetRef(loop); } else { - auto copy = ffi::make_object(*loop); + auto copy = ffi::make_object(*loop); copy->body = this->Mutate(copy->body, inplace_mode).ValueOrUnchanged(copy->body); return For(std::move(copy)); } @@ -912,7 +916,7 @@ class ChildReplacer : public StmtExprMutator { } /*! \brief The `src_stmt` to be replaced */ - const StmtNode* src_stmt_; + const StmtObj* src_stmt_; /*! \brief The `tgt_stmt` to be replaced in */ const Stmt& tgt_stmt_; /*! @@ -925,11 +929,11 @@ class ChildReplacer : public StmtExprMutator { void ScheduleStateObj::Replace(const tirx::StmtSRef& _src_sref, const Stmt& tgt_stmt, const ffi::Map& _block_sref_reuse) { if (this->debug_mask != 0) { - const StmtNode* src_stmt = _src_sref->stmt; + const StmtObj* src_stmt = _src_sref->stmt; bool input_correct = - (src_stmt->IsInstance() && tgt_stmt->IsInstance()) || - (src_stmt->IsInstance() && tgt_stmt->IsInstance()) || - (src_stmt->IsInstance() && tgt_stmt->IsInstance()); + (src_stmt->IsInstance() && tgt_stmt->IsInstance()) || + (src_stmt->IsInstance() && tgt_stmt->IsInstance()) || + (src_stmt->IsInstance() && tgt_stmt->IsInstance()); if (!input_correct) { TVM_FFI_THROW(TypeError) << "src_stmt has type: " << src_stmt->GetTypeKey() << ". tgt_stmt has type: " << tgt_stmt->GetTypeKey() @@ -952,7 +956,7 @@ void ScheduleStateObj::Replace(const tirx::StmtSRef& _src_sref, const Stmt& tgt_ // 3) all `stmt`s are correct, except for the root { // Step 0. Setup block_sref_reuse - std::unordered_map block_sref_reuse; + std::unordered_map block_sref_reuse; block_sref_reuse.reserve(_block_sref_reuse.size() + 1); for (const auto& kv : _block_sref_reuse) { block_sref_reuse.emplace(kv.first.get(), kv.second.get()); @@ -982,7 +986,7 @@ void ScheduleStateObj::Replace(const tirx::StmtSRef& _src_sref, const Stmt& tgt_ // 3) `g_var` and `g_func`. Indicate which GlobalVar and Function the sref corresponds to int num_copy_steps = -1; bool need_module_copy = false; - const FunctionNode* g_func = nullptr; + const FunctionObj* g_func = nullptr; GlobalVar g_var{ffi::UnsafeInit{}}; { int i = 0; @@ -992,13 +996,13 @@ void ScheduleStateObj::Replace(const tirx::StmtSRef& _src_sref, const Stmt& tgt_ if (child != nullptr) { // The sref tree omits body sequences and BlockRealize wrappers. A shared // container can own a unique child, so include it in the ancestor proof. - const auto* loop = p->StmtAs(); + const auto* loop = p->StmtAs(); const SeqStmt& body = loop ? loop->body : TVM_SREF_TO_SBLOCK(p)->body; bool unique_direct_child = false; if (0 <= child->seq_index && child->seq_index < static_cast(body->size())) { ffi::AnyView child_value = body->seq.GetArrayObj()->at(child->seq_index); - const auto* child_stmt = child_value.as(); - const auto* realize = child_value.as(); + const auto* child_stmt = child_value.as(); + const auto* realize = child_value.as(); bool is_direct_child = child_stmt == child->stmt || (realize != nullptr && realize->block.get() == child->stmt); unique_direct_child = is_direct_child && child_stmt->unique(); @@ -1021,7 +1025,7 @@ void ScheduleStateObj::Replace(const tirx::StmtSRef& _src_sref, const Stmt& tgt_ } // Find `g_func` and `g_var` where the `src_sref` is in g_func = GetRootFunction(this->mod, p->stmt, &g_var); - const auto* function_body = g_func->body.as(); + const auto* function_body = g_func->body.as(); TVM_FFI_ICHECK(function_body != nullptr && function_body->size() == 1); need_module_copy = num_copy_steps == i || // !this->mod.unique() || // @@ -1029,7 +1033,7 @@ void ScheduleStateObj::Replace(const tirx::StmtSRef& _src_sref, const Stmt& tgt_ !g_func->unique() || // !function_body->unique() || // !function_body->seq.unique() || // - !function_body->seq.GetArrayObj()->at(0).as()->unique(); + !function_body->seq.GetArrayObj()->at(0).as()->unique(); } // Loop invariant: // @@ -1049,8 +1053,8 @@ void ScheduleStateObj::Replace(const tirx::StmtSRef& _src_sref, const Stmt& tgt_ for (int i = 0; (need_module_copy || i <= num_copy_steps) && child_sref->parent != nullptr; ++i) { bool can_directly_mutate_parent = !need_module_copy && i == num_copy_steps; // Replace `child_sref->stmt` to `child_tgt_stmt`. - const StmtNode* parent_stmt = child_sref->parent->stmt; - const StmtNode* child_src_stmt = child_sref->stmt; + const StmtObj* parent_stmt = child_sref->parent->stmt; + const StmtObj* child_src_stmt = child_sref->stmt; // Step 2.1. Link `child_sref` to `child_tgt_stmt` if (i == 0) { // As the invariance of SRefUpdater, @@ -1084,12 +1088,12 @@ void ScheduleStateObj::Replace(const tirx::StmtSRef& _src_sref, const Stmt& tgt_ UpdateSRef(this, child_sref, child_tgt_stmt.get()); } // Ensure the uniqueness of `this->mod` and `this->mod->functions` - IRModuleNode* new_mod = this->mod.CopyOnWrite(); + IRModuleObj* new_mod = this->mod.CopyOnWrite(); ffi::MapObj* new_map = new_mod->functions.CopyOnWrite(); // Move out the Function where the sref belong while ensuring uniqueness Function ref_new_func = std::move(new_map->at(g_var)).as_or_throw(); TVM_FFI_ICHECK(ref_new_func.get() == g_func); - FunctionNode* new_func = ref_new_func.CopyOnWrite(); + FunctionObj* new_func = ref_new_func.CopyOnWrite(); // If `g_func` was not unique, after the 3 lines above: // `ref_new_func` points to a unique Function // `g_func` points to the previous Function if it is not unique @@ -1098,13 +1102,13 @@ void ScheduleStateObj::Replace(const tirx::StmtSRef& _src_sref, const Stmt& tgt_ // Update the body of the function the sref belongs to Assign TVM_FFI_ICHECK(g_func->body.has_value()); TVM_FFI_ICHECK_EQ(g_func->body.value()->size(), 1); - const auto* realize = TVM_TYPE_AS(g_func->body.value()->seq[0], SBlockRealizeNode); + const auto* realize = TVM_TYPE_AS(g_func->body.value()->seq[0], SBlockRealizeObj); // Make `child_tgt_stmt` the root block - const auto* child_block = TVM_TYPE_AS(child_tgt_stmt, SBlockNode); - ffi::ObjectPtr new_realize = ffi::make_object(*realize); + const auto* child_block = TVM_TYPE_AS(child_tgt_stmt, SBlockObj); + ffi::ObjectPtr new_realize = ffi::make_object(*realize); new_realize->block = ffi::GetRef(child_block); new_func->body = - SeqStmt({SBlockRealize(std::move(new_realize))}, g_func->body.as()->loc); + SeqStmt({SBlockRealize(std::move(new_realize))}, g_func->body.as()->loc); SetSeqIndexInChildren(this->stmt2ref, new_func->body.value().get()); // Finally, move the `ref_new_func` back and update `this->mod` new_map->at(g_var) = std::move(ref_new_func); diff --git a/src/s_tir/schedule/trace.cc b/src/s_tir/schedule/trace.cc index a3fec1b9a60d..ea0d14597d82 100644 --- a/src/s_tir/schedule/trace.cc +++ b/src/s_tir/schedule/trace.cc @@ -32,7 +32,7 @@ using namespace tvm::tirx; namespace { -ffi::String TraceAsPythonRepr(const TraceNode* self) { +ffi::String TraceAsPythonRepr(const TraceObj* self) { std::ostringstream os; os << "# from tvm import s_tir\n"; os << "def apply_trace(sch: s_tir.Schedule) -> None:\n"; @@ -56,10 +56,10 @@ ffi::String TraceAsPythonRepr(const TraceNode* self) { /**************** Constructors ****************/ -Trace::Trace() { data_ = ffi::make_object(); } +Trace::Trace() { data_ = ffi::make_object(); } Trace::Trace(ffi::Array insts, ffi::Map decisions) { - ffi::ObjectPtr n = ffi::make_object(); + ffi::ObjectPtr n = ffi::make_object(); n->insts = std::move(insts); n->decisions = std::move(decisions); data_ = std::move(n); @@ -67,9 +67,9 @@ Trace::Trace(ffi::Array insts, ffi::Map decisions TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - TraceNode::RegisterReflection(); + TraceObj::RegisterReflection(); // Register __ffi_repr__ so str(trace) returns the Python script format - refl::TypeAttrDef().def( + refl::TypeAttrDef().def( ffi::type_attr::kRepr.name, [](Trace trace, ffi::Function) -> ffi::String { return TraceAsPythonRepr(trace.get()); }); @@ -110,9 +110,9 @@ ffi::Array TranslateInputRVs( return std::nullopt; } const ffi::Object* dst = it->second; - TVM_FFI_CHECK(dst->IsInstance(), TypeError) + TVM_FFI_CHECK(dst->IsInstance(), TypeError) << "Expect 'tirx.Var', but gets: " << dst->GetTypeKey(); - return ffi::GetRef(static_cast(dst)); + return ffi::GetRef(static_cast(dst)); }; auto f_subst_with_rv_map_prim = [&](const Var& var) -> ffi::Optional { if (auto replacement = f_subst_with_rv_map(var)) { @@ -139,9 +139,9 @@ ffi::Array TranslateInputRVs( result.push_back(input); } else if (auto expr = input.as()) { result.push_back(expr.value()); - } else if (input.as() || // RV: block - input.as() || // RV: loop - input.as()) { // RV: var + } else if (input.as() || // RV: block + input.as() || // RV: loop + input.as()) { // RV: var auto it = rv_map.find(input.as()); TVM_FFI_CHECK(it != rv_map.end(), IndexError) << "Random variable doesn't exist: " << input; result.push_back(ffi::GetRef(it->second)); @@ -183,9 +183,9 @@ ffi::Array TranslateInputRVs( } else if (input.type_index() < ffi::TypeIndex::kTVMFFISmallStr) { // directly put back POD type and not string results.push_back(input); - } else if (input.as() || // RV: block - input.as() || // RV: loop - input.as()) { // RV: var + } else if (input.as() || // RV: block + input.as() || // RV: loop + input.as()) { // RV: var auto it = rv_names.find(input.cast()); if (it != rv_names.end()) { // Case 1. SBlockRV, LoopRV, VarRV @@ -194,7 +194,7 @@ ffi::Array TranslateInputRVs( TVM_FFI_THROW(IndexError) << "Random variable is not defined " << input; throw; } - } else if (input.as() || input.as()) { + } else if (input.as() || input.as()) { // Case 3. integer or floating-point number results.push_back(input); } else if (input.as()) { @@ -203,7 +203,7 @@ ffi::Array TranslateInputRVs( } else if (input.as()) { // Case 5: dict results.push_back(input); - } else if (input.as()) { + } else if (input.as()) { // // Case 6: IndexMap IndexMap index_map = input.as_or_throw(); index_map = @@ -244,7 +244,7 @@ ffi::Array TranslateInputRVs( continue; } // Case 3. integer or floating-point immediate - if (input.as() || input.as()) { + if (input.as() || input.as()) { results.push_back(input); continue; } @@ -267,7 +267,7 @@ ffi::Array TranslateInputRVs( if (name[0] == '{' && name[size - 1] == '}') { Any obj = ffi::FromJSONGraph(ffi::json::Parse(name)); // Case 6. IndexMap - if (obj.as()) { + if (obj.as()) { IndexMap index_map = obj.as_or_throw(); index_map = ffi::StructuralMap(index_map, f_substitute) .as_or_throw(); @@ -327,11 +327,11 @@ ffi::Array TranslateAddOutputRVs( ffi::String result; if (output == nullptr) { result = "_"; - } else if (output.as()) { + } else if (output.as()) { result = "b" + std::to_string(i); - } else if (output.as()) { + } else if (output.as()) { result = "l" + std::to_string(i); - } else if (output.as()) { + } else if (output.as()) { result = "v" + std::to_string(i); } else { TVM_FFI_THROW(TypeError) << "Cannot recognize the type of the random variable: " @@ -356,18 +356,18 @@ void TranslateAddOutputRVs(const ffi::Array& old_outputs, /**************** Add/Remove/Get ****************/ -Any TraceNode::GetDecision(const Instruction& inst) const { +Any TraceObj::GetDecision(const Instruction& inst) const { return this->decisions.Get(inst).value_or(Any{nullptr}); } -void TraceNode::Append(Instruction inst) { insts.push_back(std::move(inst)); } +void TraceObj::Append(Instruction inst) { insts.push_back(std::move(inst)); } -void TraceNode::Append(Instruction inst, Any decision) { +void TraceObj::Append(Instruction inst, Any decision) { decisions.Set(inst, std::move(decision)); insts.push_back(std::move(inst)); } -ffi::Optional TraceNode::Pop() { +ffi::Optional TraceObj::Pop() { if (insts.empty()) { return std::nullopt; } @@ -381,7 +381,7 @@ ffi::Optional TraceNode::Pop() { /**************** Interfacing with InstructionKind ****************/ -void TraceNode::ApplyToSchedule( +void TraceObj::ApplyToSchedule( Schedule sch, bool remove_postproc, ffi::TypedFunction& inputs, // const ffi::Array& attrs, // @@ -403,7 +403,7 @@ void TraceNode::ApplyToSchedule( } } -ffi::ObjectRef TraceNode::AsJSON(bool remove_postproc) const { +ffi::ObjectRef TraceObj::AsJSON(bool remove_postproc) const { std::unordered_map rv_names; ffi::Array json_insts; ffi::Array json_decisions; @@ -441,7 +441,7 @@ ffi::ObjectRef TraceNode::AsJSON(bool remove_postproc) const { }; } -ffi::Array TraceNode::AsPython(bool remove_postproc) const { +ffi::Array TraceObj::AsPython(bool remove_postproc) const { std::unordered_map rv_names; ffi::Array py_trace; py_trace.reserve(this->insts.size()); @@ -545,7 +545,7 @@ void Trace::ApplyJSONToSchedule(ffi::ObjectRef json, Schedule sch) { /**************** Creation ****************/ -Trace TraceNode::WithDecision(Instruction inst, Any decision, bool remove_postproc) const { +Trace TraceObj::WithDecision(Instruction inst, Any decision, bool remove_postproc) const { int n_insts = GetNumValidInstructions(this->insts, remove_postproc); ffi::Array new_insts = ffi::Array{this->insts.begin(), this->insts.begin() + n_insts}; @@ -554,7 +554,7 @@ Trace TraceNode::WithDecision(Instruction inst, Any decision, bool remove_postpr return Trace(new_insts, new_decisions); } -Trace TraceNode::Simplified(bool remove_postproc) const { +Trace TraceObj::Simplified(bool remove_postproc) const { int n_insts = GetNumValidInstructions(this->insts, remove_postproc); std::unordered_set used_rvs; std::vector new_insts; @@ -590,7 +590,7 @@ Trace TraceNode::Simplified(bool remove_postproc) const { for (const Any& obj : inst->inputs) { if (obj == nullptr) { continue; - } else if (obj.as() || obj.as() || obj.as()) { + } else if (obj.as() || obj.as() || obj.as()) { used_rvs.insert(obj.as()); continue; } else if (auto prim_expr = obj.as()) { @@ -637,7 +637,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { RegisterInstructionKind(); } TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; refl::GlobalDef() - .def_method("s_tir.schedule.TraceGetDecision", &TraceNode::GetDecision) + .def_method("s_tir.schedule.TraceGetDecision", &TraceObj::GetDecision) .def("s_tir.schedule.TraceAppend", [](Trace self, Instruction inst, ffi::Optional decision) { if (decision.has_value()) { @@ -646,12 +646,12 @@ TVM_FFI_STATIC_INIT_BLOCK() { return self->Append(inst); } }) - .def_method("s_tir.schedule.TracePop", &TraceNode::Pop) - .def_method("s_tir.schedule.TraceApplyToSchedule", &TraceNode::ApplyToSchedule) - .def_method("s_tir.schedule.TraceAsJSON", &TraceNode::AsJSON) - .def_method("s_tir.schedule.TraceAsPython", &TraceNode::AsPython) - .def_method("s_tir.schedule.TraceWithDecision", &TraceNode::WithDecision) - .def_method("s_tir.schedule.TraceSimplified", &TraceNode::Simplified) + .def_method("s_tir.schedule.TracePop", &TraceObj::Pop) + .def_method("s_tir.schedule.TraceApplyToSchedule", &TraceObj::ApplyToSchedule) + .def_method("s_tir.schedule.TraceAsJSON", &TraceObj::AsJSON) + .def_method("s_tir.schedule.TraceAsPython", &TraceObj::AsPython) + .def_method("s_tir.schedule.TraceWithDecision", &TraceObj::WithDecision) + .def_method("s_tir.schedule.TraceSimplified", &TraceObj::Simplified) .def("s_tir.schedule.TraceApplyJSONToSchedule", Trace::ApplyJSONToSchedule); } diff --git a/src/s_tir/schedule/transform.cc b/src/s_tir/schedule/transform.cc index f3dd164307cb..93259c089ee3 100644 --- a/src/s_tir/schedule/transform.cc +++ b/src/s_tir/schedule/transform.cc @@ -36,24 +36,24 @@ using namespace tvm::tirx; /******** Annotation ********/ -SBlock WithAnnotation(const SBlockNode* block, const ffi::String& attr_key, +SBlock WithAnnotation(const SBlockObj* block, const ffi::String& attr_key, const ffi::ObjectRef& attr_value) { ffi::Map annotations = block->annotations; annotations.Set(attr_key, attr_value); - ffi::ObjectPtr new_block = ffi::make_object(*block); + ffi::ObjectPtr new_block = ffi::make_object(*block); new_block->annotations = std::move(annotations); return SBlock(new_block); } /******** Tensor Related ********/ TensorVar WithScope(const TensorVar& tensor, runtime::MemoryScope scope) { - auto type = ffi::make_object(*tensor.ty().get()); + auto type = ffi::make_object(*tensor.ty().get()); type->scope = scope; return TensorVar(tensor->name + "_" + scope.name(), TensorType(std::move(type)), tensor->loc); } TensorVar WithDType(const TensorVar& tensor, PrimType dtype) { - auto type = ffi::make_object(*tensor.ty().get()); + auto type = ffi::make_object(*tensor.ty().get()); type->dtype = dtype; return TensorVar(tensor->name, TensorType(std::move(type)), tensor->loc); } @@ -62,7 +62,7 @@ ffi::Array ReplaceTensor(ffi::Array regions, const T const TensorVar& target) { regions.MutateByApply([&source, &target](TensorRegion region) -> TensorRegion { if (region->source.as_or_throw().same_as(source)) { - ffi::ObjectPtr n = ffi::make_object(*region.get()); + ffi::ObjectPtr n = ffi::make_object(*region.get()); n->source = target; return TensorRegion(n); } @@ -75,7 +75,7 @@ ffi::Array ReplaceTensor(ffi::Array regions, const ffi::Map& tensor_map) { regions.MutateByApply([&tensor_map](TensorRegion region) -> TensorRegion { if (tensor_map.count(region->source.as_or_throw())) { - ffi::ObjectPtr n = ffi::make_object(*region.get()); + ffi::ObjectPtr n = ffi::make_object(*region.get()); n->source = tensor_map[region->source.as_or_throw()]; return TensorRegion(n); } @@ -89,8 +89,8 @@ ffi::Array ReplaceTensor(ffi::Array match_ match_tensors.MutateByApply( [&source, &target](MatchTensorRegion match_tensor) -> MatchTensorRegion { if (match_tensor->source->source.as_or_throw().same_as(source)) { - ffi::ObjectPtr n = - ffi::make_object(*match_tensor.get()); + ffi::ObjectPtr n = + ffi::make_object(*match_tensor.get()); n->source = TensorRegion(target, n->source->region, TensorRegionType()); return MatchTensorRegion(ffi::UnsafeInit{}, n); } @@ -117,8 +117,8 @@ ffi::Array ReplaceTensorRegion(ffi::Array match_tensors.MutateByApply([&source_tensor, &target]( const MatchTensorRegion& match_tensor) -> MatchTensorRegion { if (match_tensor->source->source.as_or_throw().same_as(source_tensor)) { - ffi::ObjectPtr n = - ffi::make_object(*match_tensor.get()); + ffi::ObjectPtr n = + ffi::make_object(*match_tensor.get()); n->source = target; return MatchTensorRegion(ffi::UnsafeInit{}, n); } @@ -142,20 +142,20 @@ ReplaceTensorMutator::ReplaceTensorMutator(const ffi::Map& } } -UnchangedOr ReplaceTensorMutator::Mutate_(const CallNode* op, InplaceMode inplace_mode) { +UnchangedOr ReplaceTensorMutator::Mutate_(const CallObj* op, InplaceMode inplace_mode) { auto result = StmtExprMutator::Mutate_(op, inplace_mode); if (!result.IsUnchanged()) { - op = ffi::AnyView(result).as(); + op = ffi::AnyView(result).as(); if (!op->unique()) inplace_mode = InplaceMode::kDisallow; } if (!op->op.same_as(tirx::tensor_data_ptr_op()) || op->args.size() != 1) return result; PtrType type = op->args[0].as_or_throw().ty()->DataPtrType(); if (ffi::StructuralEqual()(op->ty, type)) return result; if (inplace_mode == InplaceMode::kAllow) { - const_cast(op)->ty = std::move(type); + const_cast(op)->ty = std::move(type); return result; } - auto copy = ffi::make_object(*op); + auto copy = ffi::make_object(*op); copy->ty = std::move(type); return Expr(std::move(copy)); } @@ -173,7 +173,7 @@ MatchTensorRegion ReplaceTensorMutator::VisitMatchTensorRegion( } } -UnchangedOr ReplaceTensorMutator::Mutate_(const SBlockNode* block, InplaceMode inplace_mode) { +UnchangedOr ReplaceTensorMutator::Mutate_(const SBlockObj* block, InplaceMode inplace_mode) { // To reduce the number of blocks in block sref reuse map, we check whether the block is really // mutated (i.e., the old tensor appears in the block). If so, we return the block after // mutation. Otherwise we just return the original block. @@ -229,7 +229,7 @@ UnchangedOr ReplaceTensorMutator::Mutate_(const SBlockNode* block, Inplace match_tensors.same_as(mutated_block->match_tensors)) { return ffi::Unchanged(); } else { - SBlockNode* n = mutated_block.CopyOnWrite(); + SBlockObj* n = mutated_block.CopyOnWrite(); n->reads = std::move(reads); n->writes = std::move(writes); n->alloc_tensors = std::move(alloc_tensors); @@ -272,10 +272,10 @@ void LeafBlockRemovalPlan(const ScheduleState& self, const StmtSRef& leaf_block_ }; // Go upwards until find an ancestor with more than one child - const StmtNode* last_stmt = leaf_block_sref->stmt; + const StmtObj* last_stmt = leaf_block_sref->stmt; StmtSRefObj* sref = leaf_block_sref->parent; for (;; last_stmt = sref->stmt, sref = sref->parent) { - if (const auto* loop = sref->StmtAs()) { + if (const auto* loop = sref->StmtAs()) { if (loop->body->size() > 1) { break; } @@ -285,9 +285,9 @@ void LeafBlockRemovalPlan(const ScheduleState& self, const StmtSRef& leaf_block_ break; } } - if (const auto* block = sref->StmtAs()) { + if (const auto* block = sref->StmtAs()) { if (block->body->size() > 1) { - ffi::ObjectPtr n = ffi::make_object(*block); + ffi::ObjectPtr n = ffi::make_object(*block); auto new_seq = RemoveFromSeqStmt(block->body, ffi::GetRef(last_stmt)); n->body = new_seq; *src_stmt = ffi::GetRef(block); @@ -295,9 +295,9 @@ void LeafBlockRemovalPlan(const ScheduleState& self, const StmtSRef& leaf_block_ return; } } - if (const auto* loop = sref->StmtAs()) { + if (const auto* loop = sref->StmtAs()) { if (loop->body->size() > 1) { - ffi::ObjectPtr n = ffi::make_object(*loop); + ffi::ObjectPtr n = ffi::make_object(*loop); n->body = RemoveFromSeqStmt(loop->body, ffi::GetRef(last_stmt)); *src_stmt = ffi::GetRef(loop); *tgt_stmt = Stmt(std::move(n)); @@ -394,16 +394,16 @@ ffi::Optional TileWithTensorIntrin(const s_tir::Schedule& sch, for (const auto& kv : info->loop_map) { // Extract mapping (block_loop => desc_loop) const tirx::StmtSRef& block_loop_sref = kv.first; - const ForNode* block_loop = block_loop_sref->StmtAs(); - const ForNode* desc_loop = kv.second.get(); + const ForObj* block_loop = block_loop_sref->StmtAs(); + const ForObj* desc_loop = kv.second.get(); TVM_FFI_ICHECK(block_loop != nullptr && desc_loop != nullptr); // Extract the loop extent PrimExpr block_extent = analyzer->Simplify(Range::GetExtentFromBeginEnd(block_loop->begin, block_loop->end)); PrimExpr desc_extent = analyzer->Simplify(Range::GetExtentFromBeginEnd(desc_loop->begin, desc_loop->end)); - const auto* int_block_extent = block_extent.as(); - const auto* int_desc_extent = desc_extent.as(); + const auto* int_block_extent = block_extent.as(); + const auto* int_desc_extent = desc_extent.as(); TVM_FFI_ICHECK(int_block_extent != nullptr && int_desc_extent != nullptr); // Check divisibility const ffi::BigInt& total = int_block_extent->value; @@ -454,7 +454,7 @@ void BlockTensorAccessSimplifier::SimplifyAccessRegion( int n = tensor_region->region.size(); for (int i = 0; i < n; ++i) { PrimExpr min = simplified_min[i]; - PrimExpr extent = analyzer_->Simplify(tensor_region->region[i]->extent); + PrimExpr extent = GetSyncedAnalyzer()->Simplify(tensor_region->region[i]->extent); new_tensor_region.push_back(Range::FromMinExtent(min, extent)); } return TensorRegion(tensor_region->source.as_or_throw(), @@ -467,29 +467,35 @@ void BlockTensorAccessSimplifier::SimplifyTensorIndices(ffi::Array* in *indices = this->IterMapSimplifyWithContext(*indices, true); } -UnchangedOr BlockTensorAccessSimplifier::Mutate_(const SBlockNode* op, +UnchangedOr BlockTensorAccessSimplifier::Mutate_(const SBlockObj* op, InplaceMode inplace_mode) { - SBlock block = s_tir::StmtExprMutatorWithAnalyzer::Mutate_(op, inplace_mode) + SBlock block = s_tir::StmtExprMutatorWithConstraintContext::Mutate_(op, inplace_mode) .ValueOrUnchanged(ffi::GetRef(op)) .as_or_throw(); - auto* n = block.CopyOnWrite(); - SimplifyAccessRegion(&n->reads); - SimplifyAccessRegion(&n->writes); - return block; + return WithConstraintScope([&]() -> ffi::Expected> { + for (const auto& iter_var : block->iter_vars) { + AppendConstraint(sym::RangeConstraint(iter_var->var, iter_var->dom.value())); + } + auto* n = block.CopyOnWrite(); + SimplifyAccessRegion(&n->reads); + SimplifyAccessRegion(&n->writes); + return block; + }) + .value(); } -UnchangedOr BlockTensorAccessSimplifier::Mutate_(const TensorStoreNode* op, +UnchangedOr BlockTensorAccessSimplifier::Mutate_(const TensorStoreObj* op, InplaceMode inplace_mode) { - TensorStore node = s_tir::StmtExprMutatorWithAnalyzer::Mutate_(op, inplace_mode) + TensorStore node = s_tir::StmtExprMutatorWithConstraintContext::Mutate_(op, inplace_mode) .ValueOrUnchanged(ffi::GetRef(op)) .as_or_throw(); SimplifyTensorIndices(&node.CopyOnWrite()->indices); return node; } -UnchangedOr BlockTensorAccessSimplifier::Mutate_(const TensorLoadNode* op, +UnchangedOr BlockTensorAccessSimplifier::Mutate_(const TensorLoadObj* op, InplaceMode inplace_mode) { - TensorLoad node = s_tir::StmtExprMutatorWithAnalyzer::Mutate_(op, inplace_mode) + TensorLoad node = s_tir::StmtExprMutatorWithConstraintContext::Mutate_(op, inplace_mode) .ValueOrUnchanged(ffi::GetRef(op)) .as_or_throw(); SimplifyTensorIndices(&node.CopyOnWrite()->indices); @@ -522,7 +528,7 @@ ffi::Optional NormalizeFunction(Schedule sch) { return std::nullopt; } for (int i = 0, n = loops.size(); i < n; ++i) { - const ForNode* loop = TVM_SREF_TO_FOR(loops[i]); + const ForObj* loop = TVM_SREF_TO_FOR(loops[i]); if (binds[i].get() != loop->var.get()) { return std::nullopt; } diff --git a/src/s_tir/schedule/transform.h b/src/s_tir/schedule/transform.h index fc9ff52c1cb3..2c4e1cab0c6d 100644 --- a/src/s_tir/schedule/transform.h +++ b/src/s_tir/schedule/transform.h @@ -25,7 +25,7 @@ #include #include #include -#include +#include #include #include @@ -43,7 +43,7 @@ using namespace tvm::tirx; * \param attr_value The annotation value to be added * \return A new block with the given annotation as its last annotation */ -SBlock WithAnnotation(const SBlockNode* block, const ffi::String& attr_key, +SBlock WithAnnotation(const SBlockObj* block, const ffi::String& attr_key, const ffi::ObjectRef& attr_value); /******** Tensor Related ********/ @@ -142,11 +142,11 @@ class ReplaceTensorMutator : public StmtExprMutator { ffi::Map* block_sref_reuse); protected: - UnchangedOr Mutate_(const CallNode* op, InplaceMode inplace_mode) override; + UnchangedOr Mutate_(const CallObj* op, InplaceMode inplace_mode) override; virtual MatchTensorRegion VisitMatchTensorRegion(const MatchTensorRegion& match_tensor); - UnchangedOr Mutate_(const SBlockNode* block, InplaceMode inplace_mode) override; + UnchangedOr Mutate_(const SBlockObj* block, InplaceMode inplace_mode) override; /*! \brief The block sref reuse map for the following replacement */ ffi::Map* block_sref_reuse_; @@ -211,10 +211,10 @@ ffi::Optional TileWithTensorIntrin(const s_tir::Schedule& sch, /*! * \brief Simplifier for indices of tensor access and block tensor access regions. */ -class BlockTensorAccessSimplifier : public s_tir::StmtExprMutatorWithAnalyzer { +class BlockTensorAccessSimplifier : public s_tir::StmtExprMutatorWithConstraintContext { public: - using s_tir::StmtExprMutatorWithAnalyzer::Mutate; - using s_tir::StmtExprMutatorWithAnalyzer::Mutate_; + using s_tir::StmtExprMutatorWithConstraintContext::Mutate; + using s_tir::StmtExprMutatorWithConstraintContext::Mutate_; /*! * \brief Simplify indices of tensor access and block tensor access regions in the statement @@ -228,15 +228,15 @@ class BlockTensorAccessSimplifier : public s_tir::StmtExprMutatorWithAnalyzer { } explicit BlockTensorAccessSimplifier(const sym::Analyzer& analyzer) - : StmtExprMutatorWithAnalyzer(analyzer) {} + : StmtExprMutatorWithConstraintContext(analyzer->GetConstraintContext()) {} private: void SimplifyAccessRegion(ffi::Array* old_access_regions); void SimplifyTensorIndices(ffi::Array* indices); - UnchangedOr Mutate_(const SBlockNode* op, InplaceMode inplace_mode) final; - UnchangedOr Mutate_(const TensorStoreNode* op, InplaceMode inplace_mode) final; - UnchangedOr Mutate_(const TensorLoadNode* op, InplaceMode inplace_mode) final; + UnchangedOr Mutate_(const SBlockObj* op, InplaceMode inplace_mode) final; + UnchangedOr Mutate_(const TensorStoreObj* op, InplaceMode inplace_mode) final; + UnchangedOr Mutate_(const TensorLoadObj* op, InplaceMode inplace_mode) final; }; } // namespace s_tir diff --git a/src/s_tir/schedule/utils.h b/src/s_tir/schedule/utils.h index 911d943c104d..8bc048f6a592 100644 --- a/src/s_tir/schedule/utils.h +++ b/src/s_tir/schedule/utils.h @@ -67,7 +67,7 @@ inline ffi::Array LoopSRefs2Loops(const ffi::Array& loop_srefs) { ffi::Array loops; loops.reserve(loop_srefs.size()); for (StmtSRef loop_sref : loop_srefs) { - const ForNode* loop = TVM_SREF_TO_FOR(loop_sref); + const ForObj* loop = TVM_SREF_TO_FOR(loop_sref); loops.push_back(ffi::GetRef(loop)); } return loops; @@ -132,7 +132,7 @@ inline SeqStmt RemoveFromSeqStmt(const SeqStmt& seq, const Stmt& to_remove) { if (to_remove.same_as(stmt)) { continue; } - if (const auto* realize = stmt.as()) { + if (const auto* realize = stmt.as()) { if (to_remove.same_as(realize->block)) { continue; } @@ -183,7 +183,7 @@ inline support::GPUThreadTag ParseThreadBindingTag(const ffi::String& tag) { * \param loop The loop to be inspected * \return The thread scope bound to the loop */ -inline support::GPUThreadTag GetThreadBindingTag(const ForNode* loop) { +inline support::GPUThreadTag GetThreadBindingTag(const ForObj* loop) { if (tvm::s_tir::GetThreadBinding(loop).has_value()) { return ParseThreadBindingTag(tvm::s_tir::GetThreadBinding(loop).value()); } @@ -222,7 +222,7 @@ inline ffi::Optional AnalyzeVarWithShift(const PrimExpr& expr, return static_cast(var.value()); } // match: "var + shift" - if (const auto* add = expr.as()) { + if (const auto* add = expr.as()) { if (auto var = add->a.as()) { if (auto shift = add->b.as()) { *constant = shift; @@ -237,7 +237,7 @@ inline ffi::Optional AnalyzeVarWithShift(const PrimExpr& expr, } } // match: "var - shift" - if (const auto* sub = expr.as()) { + if (const auto* sub = expr.as()) { if (auto var = sub->a.as()) { if (auto shift = sub->b.as()) { *constant = IntImm(shift.value().ty(), -shift.value()->value); @@ -277,10 +277,10 @@ inline ffi::Optional GetAnn(const TStmtNode* stmt, const ffi::String */ template inline ffi::Optional GetAnn(const StmtSRef& sref, const ffi::String& ann_key) { - if (const auto* loop = sref->StmtAs()) { - return GetAnn(loop, ann_key); - } else if (const auto* block = sref->StmtAs()) { - return GetAnn(block, ann_key); + if (const auto* loop = sref->StmtAs()) { + return GetAnn(loop, ann_key); + } else if (const auto* block = sref->StmtAs()) { + return GetAnn(block, ann_key); } else { TVM_FFI_THROW(TypeError) << "Unknown type of sref: " << sref->stmt->GetTypeKey(); throw; @@ -391,11 +391,11 @@ inline std::unordered_set GetSBlockNames(const IRModule& mod) { using s_tir::StmtExprVisitor::Visit_; ffi::Optional Visit(ffi::AnyView value) override { - if (value.as()) return std::nullopt; + if (value.as()) return std::nullopt; return s_tir::StmtExprVisitor::Visit(value); } - ffi::Optional Visit_(const s_tir::SBlockNode* block) override { + ffi::Optional Visit_(const s_tir::SBlockObj* block) override { block_names.insert(block->name_hint); return StmtExprVisitor::Visit(block->body); } diff --git a/src/s_tir/script/ir_builder/ir.cc b/src/s_tir/script/ir_builder/ir.cc index 01a74ab4ef71..5aaf77150b0a 100644 --- a/src/s_tir/script/ir_builder/ir.cc +++ b/src/s_tir/script/ir_builder/ir.cc @@ -316,11 +316,11 @@ ffi::Array Remap(ffi::String kinds, ffi::Array bindings, PrimType } // namespace axis TVM_FFI_STATIC_INIT_BLOCK() { - tvm::script::ir_builder::details::Namer::vtable().SetDispatch( + tvm::script::ir_builder::detail::Namer::vtable().SetDispatch( [](const ffi::ObjectRef& node, ffi::String name) -> void { using namespace tvm::s_tir; - IterVarNode* var = const_cast(node.as()); - tvm::script::ir_builder::details::Namer::Name(var->var, name); + IterVarObj* var = const_cast(node.as()); + tvm::script::ir_builder::detail::Namer::Name(var->var, name); }); } diff --git a/src/s_tir/script/ir_builder/script_complete.cc b/src/s_tir/script/ir_builder/script_complete.cc index b5ac212cefd2..9ec8ef55b462 100644 --- a/src/s_tir/script/ir_builder/script_complete.cc +++ b/src/s_tir/script/ir_builder/script_complete.cc @@ -44,7 +44,7 @@ class ScriptCompleter : public s_tir::StmtExprMutator { public: using s_tir::StmtExprMutator::Mutate; UnchangedOr Mutate(ffi::AnyView value, InplaceMode inplace_mode) final { - if (value.as()) return ffi::Unchanged(); + if (value.as()) return ffi::Unchanged(); return s_tir::StmtExprMutator::Mutate(value, inplace_mode); } @@ -54,12 +54,12 @@ class ScriptCompleter : public s_tir::StmtExprMutator { private: ffi::Map* tensor_var_map_; - UnchangedOr Mutate_(const ForNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const ForObj* op, InplaceMode inplace_mode) final { GetThreadBinding(op); return s_tir::StmtExprMutator::Mutate_(op, inplace_mode); } - UnchangedOr Mutate_(const s_tir::SBlockRealizeNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const s_tir::SBlockRealizeObj* op, InplaceMode inplace_mode) final { for (const PrimExpr& value : op->iter_values) { PrimType value_ty = value.ty(); TVM_FFI_ICHECK(value_ty.code() == DLDataTypeCode::kDLInt) @@ -68,7 +68,7 @@ class ScriptCompleter : public s_tir::StmtExprMutator { return s_tir::StmtExprMutator::Mutate_(op, inplace_mode); } - UnchangedOr Mutate_(const s_tir::SBlockNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const s_tir::SBlockObj* op, InplaceMode inplace_mode) final { // Tensors allocated in the block can be accessed by its body. for (const auto& alloc_tensor : op->alloc_tensors) { tensor_var_map_->Set(alloc_tensor, alloc_tensor); @@ -122,10 +122,10 @@ class ScriptCompleter : public s_tir::StmtExprMutator { } } - UnchangedOr Mutate_(const BindNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const BindObj* op, InplaceMode inplace_mode) final { Stmt stmt = s_tir::StmtExprMutator::Mutate_(op, inplace_mode).ValueOrUnchanged(ffi::GetRef(op)); - const Var& var = stmt.as()->var; + const Var& var = stmt.as()->var; // A flat definition registers its tensor for access detection in subsequent siblings. if (auto tensor = var.as(); tensor && !tensor_var_map_->count(var)) { tensor_var_map_->Set(var, tensor.value()); @@ -160,12 +160,12 @@ Function ScriptComplete(Function func, const ffi::Array& root_allocat return true; } auto* block_realize = func->body.value()->seq.size() == 1 - ? func->body.value()->seq[0].as() + ? func->body.value()->seq[0].as() : nullptr; if (block_realize && block_realize->block->iter_vars.size()) { return true; } - if (!block_realize && ContainsNode(func->body.value())) { + if (!block_realize && ContainsNode(func->body.value())) { return true; } return false; diff --git a/src/s_tir/script/printer/expr.cc b/src/s_tir/script/printer/expr.cc index 81ea8f3727ae..06381a3fc265 100644 --- a/src/s_tir/script/printer/expr.cc +++ b/src/s_tir/script/printer/expr.cc @@ -27,13 +27,13 @@ namespace script { namespace printer { namespace details { -TVM_FFI_STATIC_INIT_BLOCK() { RegisterScriptRepr(); } +TVM_FFI_STATIC_INIT_BLOCK() { RegisterScriptRepr(); } namespace { ffi::Optional IterVarDocTranslate(DocTranslatorObj* d, ffi::AnyView input, const ffi::Object*) { const auto* iter = - ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(input); + ffi::detail::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(input); ExprDoc domain = d->Translate(iter->dom).value(); if (iter->dom.has_value()) { domain.as_or_throw()->callee = NamespaceDoc("tirx")->Attr("Range"); @@ -46,7 +46,7 @@ ffi::Optional IterVarDocTranslate(DocTranslatorObj* d, ffi::AnyView inp } TVM_FFI_STATIC_INIT_BLOCK() { - ffi::reflection::TypeAttrDef().attr( + ffi::reflection::TypeAttrDef().attr( tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&IterVarDocTranslate>()); } diff --git a/src/s_tir/script/printer/stmt.cc b/src/s_tir/script/printer/stmt.cc index ac06b31a0ee8..b7f463fa8a17 100644 --- a/src/s_tir/script/printer/stmt.cc +++ b/src/s_tir/script/printer/stmt.cc @@ -42,7 +42,7 @@ CallDoc TensorHelperCall(DocTranslatorObj* d, const Type& type, const ffi::Strin return rhs; } -void PrintMatchTensor(DocTranslatorObj* d, const s_tir::MatchTensorRegionNode* match, +void PrintMatchTensor(DocTranslatorObj* d, const s_tir::MatchTensorRegionObj* match, const ffi::Array& allocated_addr = {}) { ExprDoc source = d->Translate(match->source).value(); CallDoc rhs = TensorHelperCall(d, match->tensor->ty, "match_tensor"); @@ -55,8 +55,8 @@ void PrintMatchTensor(DocTranslatorObj* d, const s_tir::MatchTensorRegionNode* m d->Emit(AssignDoc(lhs, rhs, std::nullopt), ffi::GetRef(match)); } -ffi::Array SBlockBody(DocTranslatorObj* d, const s_tir::SBlockNode* block, - const s_tir::SBlockRealizeNode* realize) { +ffi::Array SBlockBody(DocTranslatorObj* d, const s_tir::SBlockObj* block, + const s_tir::SBlockRealizeObj* realize) { auto docs = d->WithDocScope([&]() { static const auto scope_shared = runtime::MemoryScope::Get("shared"); static const auto scope_shared_dyn = runtime::MemoryScope::Get("shared.dyn"); @@ -134,7 +134,7 @@ ffi::Array SBlockBody(DocTranslatorObj* d, const s_tir::SBlockNode* blo } for (const tirx::TensorVar& tensor : block->alloc_tensors) { CallDoc rhs = TensorHelperCall(d, tensor->ty, "sblock_alloc_tensor"); - const auto* tensor_type = tensor->ty.as(); + const auto* tensor_type = tensor->ty.as(); auto allocated_addr = addresses.Get(tensor).value_or(ffi::Array{}); TVM_FFI_CHECK( allocated_addr.empty() || @@ -168,11 +168,11 @@ ffi::Array SBlockBody(DocTranslatorObj* d, const s_tir::SBlockNode* blo ffi::Optional SBlockRealizeDocTranslate(DocTranslatorObj* d, ffi::AnyView input, const ffi::Object* destination) { const auto* realize = - ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck( + ffi::detail::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck( input); TVM_FFI_CHECK(destination == nullptr, TypeError) << "printer statement-only node cannot fulfill a destination"; - const s_tir::SBlockNode* block = realize->block.get(); + const s_tir::SBlockObj* block = realize->block.get(); d->Emit(ScopeDoc(std::nullopt, NamespaceDoc("s_tir")->Attr("sblock")->Call( {LiteralDoc::Str(block->name_hint, std::nullopt)}), @@ -182,7 +182,7 @@ ffi::Optional SBlockRealizeDocTranslate(DocTranslatorObj* d, ffi::AnyVi } TVM_FFI_STATIC_INIT_BLOCK() { - ffi::reflection::TypeAttrDef().attr( + ffi::reflection::TypeAttrDef().attr( tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&SBlockRealizeDocTranslate>()); } @@ -190,7 +190,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { ffi::Optional SBlockDocTranslate(DocTranslatorObj* d, ffi::AnyView input, const ffi::Object* destination) { const auto* block = - ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(input); + ffi::detail::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(input); TVM_FFI_CHECK(destination == nullptr, TypeError) << "printer statement-only node cannot fulfill a destination"; d->Emit(ScopeDoc(std::nullopt, @@ -203,15 +203,16 @@ ffi::Optional SBlockDocTranslate(DocTranslatorObj* d, ffi::AnyView inpu } TVM_FFI_STATIC_INIT_BLOCK() { - ffi::reflection::TypeAttrDef().attr( + ffi::reflection::TypeAttrDef().attr( tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&SBlockDocTranslate>()); } ffi::Optional MatchTensorRegionDocTranslate(DocTranslatorObj* d, ffi::AnyView input, const ffi::Object* destination) { - const auto* match = ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck< - const s_tir::MatchTensorRegionNode>(input); + const auto* match = + ffi::detail::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck( + input); TVM_FFI_CHECK(destination == nullptr, TypeError) << "printer statement-only node cannot fulfill a destination"; PrintMatchTensor(d, match); @@ -219,15 +220,15 @@ ffi::Optional MatchTensorRegionDocTranslate(DocTranslatorObj* d, ffi::A } TVM_FFI_STATIC_INIT_BLOCK() { - ffi::reflection::TypeAttrDef().attr( + ffi::reflection::TypeAttrDef().attr( tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&MatchTensorRegionDocTranslate>()); } TVM_FFI_STATIC_INIT_BLOCK() { - RegisterScriptRepr(); - RegisterScriptRepr(); - RegisterScriptRepr(); + RegisterScriptRepr(); + RegisterScriptRepr(); + RegisterScriptRepr(); } } // namespace diff --git a/src/s_tir/stmt.cc b/src/s_tir/stmt.cc index 979adb445631..8b4c826dbc6d 100644 --- a/src/s_tir/stmt.cc +++ b/src/s_tir/stmt.cc @@ -57,7 +57,7 @@ const Op& manual_sync() { return op; } -static ffi::Array RegionNoBodyParams(const CallNode*) { return {}; } +static ffi::Array RegionNoBodyParams(const CallObj*) { return {}; } TVM_FFI_STATIC_INIT_BLOCK() { OpDef("s_tir.async_copy_scope", "Mark eligible copies for asynchronous lowering.") @@ -80,9 +80,8 @@ namespace { TVM_FFI_INLINE ffi::Expected> MatchTensorRegionVisit( ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept { - const MatchTensorRegionNode* self = - ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck( - value); + const MatchTensorRegionObj* self = + ffi::detail::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value); TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->WithDefRegionKind( kTVMFFIDefRegionKindPattern, [&]() { return visitor->VisitExpected(self->tensor); })); TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->source)); @@ -91,9 +90,8 @@ TVM_FFI_INLINE ffi::Expected> MatchTensorRegi TVM_FFI_INLINE ffi::Expected> MatchTensorRegionMutate( ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept { - const MatchTensorRegionNode* self = - ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck( - value); + const MatchTensorRegionObj* self = + ffi::detail::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value); TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, mapped_tensor, mutator->WithDefRegionKind(kTVMFFIDefRegionKindPattern, [&]() { return mutator->MutateExpected(self->tensor); @@ -104,7 +102,7 @@ TVM_FFI_INLINE ffi::Expected> MatchTensorRegionMutate mapped_source.UnchangedOrSameAs(self->source)) { return ffi::Unchanged(); } - ffi::ObjectPtr copy = ffi::make_object(*self); + ffi::ObjectPtr copy = ffi::make_object(*self); copy->tensor = std::move(mapped_tensor).ValueOrUnchanged(std::move(copy->tensor)); copy->source = std::move(mapped_source).ValueOrUnchanged(std::move(copy->source)); return ffi::Any(std::move(copy)); @@ -112,9 +110,8 @@ TVM_FFI_INLINE ffi::Expected> MatchTensorRegionMutate TVM_FFI_INLINE ffi::Expected> MatchTensorRegionMaybeInplaceMutate( ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept { - MatchTensorRegionNode* self = const_cast( - ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck( - value)); + MatchTensorRegionObj* self = const_cast( + ffi::detail::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value)); TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, mapped_tensor, mutator->WithDefRegionKind(kTVMFFIDefRegionKindPattern, [&]() { return mutator->MutateExpected(self->tensor, @@ -137,8 +134,8 @@ TVM_FFI_INLINE ffi::Expected> SBlockVisit( // Establish allocation and match-tensor definitions before their region uses. // Whole iterators and annotations remain part of structural traversal. // skips: name_hint - const SBlockNode* self = - ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value); + const SBlockObj* self = + ffi::detail::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value); TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->iter_vars)); TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->WithDefRegionKind( kTVMFFIDefRegionKindSimple, [&]() { return visitor->VisitExpected(self->alloc_tensors); })); @@ -156,8 +153,8 @@ TVM_FFI_INLINE ffi::Expected> SBlockMutate( // Establish allocation and match-tensor definitions before their region uses. // Whole iterators and annotations remain part of structural traversal. // skips: name_hint - const SBlockNode* self = - ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value); + const SBlockObj* self = + ffi::detail::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value); TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr>, mapped_iter_vars, mutator->MutateExpected(self->iter_vars)); TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr>, mapped_alloc_tensors, @@ -188,7 +185,7 @@ TVM_FFI_INLINE ffi::Expected> SBlockMutate( mapped_body.UnchangedOrSameAs(self->body)) { return ffi::Unchanged(); } - ffi::ObjectPtr copy = ffi::make_object(*self); + ffi::ObjectPtr copy = ffi::make_object(*self); copy->iter_vars = std::move(mapped_iter_vars).ValueOrUnchanged(std::move(copy->iter_vars)); copy->reads = std::move(mapped_reads).ValueOrUnchanged(std::move(copy->reads)); copy->writes = std::move(mapped_writes).ValueOrUnchanged(std::move(copy->writes)); @@ -212,8 +209,8 @@ TVM_FFI_INLINE ffi::Expected> SBlockMaybeInplaceMutat // Establish allocation and match-tensor definitions before their region uses. // Whole iterators and annotations remain part of structural traversal. // skips: name_hint - SBlockNode* self = const_cast( - ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value)); + SBlockObj* self = const_cast( + ffi::detail::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value)); TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN( ffi::UnchangedOr>, mapped_iter_vars, mutator->MutateExpected(self->iter_vars, ffi::InplaceMode::kAllow)); @@ -272,8 +269,8 @@ TVM_FFI_INLINE ffi::Expected> SBlockMaybeInplaceMutat TVM_FFI_INLINE ffi::Expected> SBlockRealizeVisit( ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept { - const SBlockRealizeNode* self = - ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value); + const SBlockRealizeObj* self = + ffi::detail::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value); TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->iter_values)); TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->predicate)); TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->block)); @@ -282,8 +279,8 @@ TVM_FFI_INLINE ffi::Expected> SBlockRealizeVi TVM_FFI_INLINE ffi::Expected> SBlockRealizeMutate( ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept { - const SBlockRealizeNode* self = - ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value); + const SBlockRealizeObj* self = + ffi::detail::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value); TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr>, mapped_iter_values, mutator->MutateExpected(self->iter_values)); TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, mapped_predicate, @@ -295,7 +292,7 @@ TVM_FFI_INLINE ffi::Expected> SBlockRealizeMutate( mapped_block.UnchangedOrSameAs(self->block)) { return ffi::Unchanged(); } - ffi::ObjectPtr copy = ffi::make_object(*self); + ffi::ObjectPtr copy = ffi::make_object(*self); copy->iter_values = std::move(mapped_iter_values).ValueOrUnchanged(std::move(copy->iter_values)); copy->predicate = std::move(mapped_predicate).ValueOrUnchanged(std::move(copy->predicate)); copy->block = std::move(mapped_block).ValueOrUnchanged(std::move(copy->block)); @@ -304,8 +301,8 @@ TVM_FFI_INLINE ffi::Expected> SBlockRealizeMutate( TVM_FFI_INLINE ffi::Expected> SBlockRealizeMaybeInplaceMutate( ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept { - SBlockRealizeNode* self = const_cast( - ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value)); + SBlockRealizeObj* self = const_cast( + ffi::detail::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value)); TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN( ffi::UnchangedOr>, mapped_iter_values, mutator->MutateExpected(self->iter_values, ffi::InplaceMode::kAllow)); @@ -370,15 +367,15 @@ MatchTensorRegion::MatchTensorRegion(TensorVar tensor, TensorRegion source) { } } // Note that we do not check elem_offset and strides in this function - ffi::ObjectPtr node = - ffi::make_object(std::move(tensor), std::move(source)); + ffi::ObjectPtr node = + ffi::make_object(std::move(tensor), std::move(source)); data_ = std::move(node); } TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - MatchTensorRegionNode::RegisterReflection(); - refl::TypeAttrDef() + MatchTensorRegionObj::RegisterReflection(); + refl::TypeAttrDef() .attr(ffi::type_attr::kStructuralVisit, ffi::FStructuralVisit::FromNative<&MatchTensorRegionVisit>()) .attr(ffi::type_attr::kStructuralMutate, @@ -405,7 +402,7 @@ SBlock::SBlock(ffi::Array iter_vars, ffi::Array reads, << "SBlock region must match its tensor rank"; } } - ffi::ObjectPtr node = ffi::make_object(std::move(body)); + ffi::ObjectPtr node = ffi::make_object(std::move(body)); node->iter_vars = std::move(iter_vars); node->reads = std::move(reads); node->writes = std::move(writes); @@ -421,7 +418,7 @@ SBlock::SBlock(ffi::Array iter_vars, ffi::Array reads, SBlock::SBlock(ffi::String name_hint, SeqStmt body, ffi::Array alloc_tensors, Location loc) : Stmt(ffi::UnsafeInit{}) { - ffi::ObjectPtr node = ffi::make_object(std::move(body)); + ffi::ObjectPtr node = ffi::make_object(std::move(body)); node->iter_vars = {}; node->reads = {}; node->writes = {}; @@ -436,8 +433,8 @@ SBlock::SBlock(ffi::String name_hint, SeqStmt body, ffi::Array alloc_ TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - SBlockNode::RegisterReflection(); - refl::TypeAttrDef() + SBlockObj::RegisterReflection(); + refl::TypeAttrDef() .attr(ffi::type_attr::kStructuralVisit, ffi::FStructuralVisit::FromNative<&SBlockVisit>()) .attr(ffi::type_attr::kStructuralMutate, ffi::FStructuralMutate::FromNative<&SBlockMutate>()) .attr(ffi::type_attr::kStructuralMaybeInplaceMutate, @@ -463,8 +460,8 @@ SBlockRealize::SBlockRealize(ffi::Array values, PrimExpr predicate, SB PrimType predicate_ty = predicate.ty(); TVM_FFI_CHECK(predicate_ty.MatchesCode(DLDataTypeCode::kDLBool), TypeError) << "Expect Block.predicate to be a bool expression"; - ffi::ObjectPtr node = - ffi::make_object(std::move(predicate), std::move(block)); + ffi::ObjectPtr node = + ffi::make_object(std::move(predicate), std::move(block)); node->iter_values = std::move(values); node->loc = loc; data_ = std::move(node); @@ -472,8 +469,8 @@ SBlockRealize::SBlockRealize(ffi::Array values, PrimExpr predicate, SB TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - SBlockRealizeNode::RegisterReflection(); - refl::TypeAttrDef() + SBlockRealizeObj::RegisterReflection(); + refl::TypeAttrDef() .attr(ffi::type_attr::kStructuralVisit, ffi::FStructuralVisit::FromNative<&SBlockRealizeVisit>()) .attr(ffi::type_attr::kStructuralMutate, diff --git a/src/s_tir/stmt_functor.cc b/src/s_tir/stmt_functor.cc index b319bcb06665..5cca253edf2e 100644 --- a/src/s_tir/stmt_functor.cc +++ b/src/s_tir/stmt_functor.cc @@ -32,89 +32,94 @@ using namespace tirx; void StmtExprVisitor::InitVTable(VTable* vtable) { tirx::StmtExprVisitor::InitVTable(vtable); - SetDispatch(vtable); - SetDispatch(vtable); + SetDispatch(vtable); + SetDispatch(vtable); } void StmtExprMutator::InitVTable(VTable* vtable) { tirx::StmtExprMutator::InitVTable(vtable); - SetDispatch(vtable); - SetDispatch(vtable); + SetDispatch(vtable); + SetDispatch(vtable); } -ffi::Optional StmtExprVisitor::Visit_(const SBlockNode* op) { +ffi::Optional StmtExprVisitor::Visit_(const SBlockObj* op) { return VisitBlock(this, op); } -ffi::Optional StmtExprVisitor::VisitBlock(tirx::StmtExprVisitor* visitor, - const SBlockNode* op) { +ffi::Optional StmtExprVisitor::VisitBlock(ffi::StructuralVisitorObj* visitor, + const SBlockObj* op) { for (const IterVar& iter_var : op->iter_vars) { - TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->Visit(iter_var->dom.value()->begin)); - TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->Visit(iter_var->dom.value()->extent)); + TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN( + visitor->VisitExpected(iter_var->dom.value()->begin).value()); + TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN( + visitor->VisitExpected(iter_var->dom.value()->extent).value()); } for (const TensorVar& tensor : op->alloc_tensors) { TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->WithDefRegionKind( - kTVMFFIDefRegionKindSimple, [&]() { return visitor->Visit(tensor); })); + kTVMFFIDefRegionKindSimple, [&]() { return visitor->VisitExpected(tensor).value(); })); } // Define match-tensor targets before visiting reads/writes that may use them. // This differs from the old TIRX native order (reads/writes before matches) // and agrees with structural traversal's definition-before-use contract. for (const MatchTensorRegion& match_tensor_region : op->match_tensors) { TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->WithDefRegionKind( - kTVMFFIDefRegionKindSimple, [&]() { return visitor->Visit(match_tensor_region->tensor); })); - TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->Visit(match_tensor_region->source)); + kTVMFFIDefRegionKindSimple, + [&]() { return visitor->VisitExpected(match_tensor_region->tensor).value(); })); + TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(match_tensor_region->source).value()); } if (auto addresses = op->annotations.Get(attr::kTensorAllocatedAddr.name)) { // The owner reference identifies the allocation; it is not a data access. for (const auto& entry : addresses.value().cast()) { - TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->Visit(entry.get<1>())); + TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(entry.get<1>()).value()); } } for (const TensorRegion& region : op->reads) { - TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->Visit(region)); + TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(region).value()); } for (const TensorRegion& region : op->writes) { - TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->Visit(region)); + TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(region).value()); } if (op->init.has_value()) { - TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->Visit(op->init.value())); + TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(op->init.value()).value()); } - return visitor->Visit(op->body); + return visitor->VisitExpected(op->body).value(); } -ffi::Optional StmtExprVisitor::Visit_(const SBlockRealizeNode* op) { +ffi::Optional StmtExprVisitor::Visit_(const SBlockRealizeObj* op) { return VisitBlockRealize(this, op); } -ffi::Optional StmtExprVisitor::VisitBlockRealize(tirx::StmtExprVisitor* visitor, - const SBlockRealizeNode* op) { +ffi::Optional StmtExprVisitor::VisitBlockRealize(ffi::StructuralVisitorObj* visitor, + const SBlockRealizeObj* op) { for (const auto& child : op->iter_values) { - TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->Visit(child)); + TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(child).value()); } - TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->Visit(op->predicate)); - return visitor->Visit(op->block); + TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(op->predicate).value()); + return visitor->VisitExpected(op->block).value(); } -UnchangedOr StmtExprMutator::Mutate_(const SBlockNode* op, InplaceMode inplace_mode) { +UnchangedOr StmtExprMutator::Mutate_(const SBlockObj* op, InplaceMode inplace_mode) { return MutateBlock(this, op, inplace_mode); } -UnchangedOr StmtExprMutator::MutateBlock(tirx::StmtExprMutator* mutator, const SBlockNode* op, - InplaceMode inplace_mode) { +UnchangedOr StmtExprMutator::MutateBlock(ffi::StructuralMutatorObj* mutator, + const SBlockObj* op, InplaceMode inplace_mode) { // SBlock iteration variables keep their binders; only their domains are expressions here. const auto* iters = op->iter_vars.GetArrayObj(); InplaceMode iter_mode = iters->unique() ? inplace_mode : InplaceMode::kDisallow; std::vector> replacements; for (size_t i = 0; i < iters->size(); ++i) { - const auto* iter = (*iters)[i].as(); + const auto* iter = (*iters)[i].as(); if (!iter->dom.has_value()) continue; InplaceMode domain_mode = iter->unique() ? iter_mode : InplaceMode::kDisallow; - auto domain = mutator->Mutate(iter->dom.value(), domain_mode).as_or_throw>(); + auto domain = mutator->MutateExpected(iter->dom.value(), domain_mode) + .value() + .as_or_throw>(); if (domain.UnchangedOrSameAs(iter->dom.value())) continue; if (domain_mode == InplaceMode::kAllow) { - const_cast(iter)->dom = std::move(domain).ValueUnchecked(); + const_cast(iter)->dom = std::move(domain).ValueUnchecked(); } else { - auto updated = ffi::make_object(*iter); + auto updated = ffi::make_object(*iter); updated->dom = std::move(domain).ValueUnchecked(); replacements.emplace_back(i, IterVar(std::move(updated))); } @@ -133,24 +138,31 @@ UnchangedOr StmtExprMutator::MutateBlock(tirx::StmtExprMutator* mutator, c } auto alloc_tensors = mutator - ->WithDefRegionKind(kTVMFFIDefRegionKindSimple, - [&] { return mutator->Mutate(op->alloc_tensors, inplace_mode); }) + ->WithDefRegionKind( + kTVMFFIDefRegionKindSimple, + [&] { return mutator->MutateExpected(op->alloc_tensors, inplace_mode).value(); }) .as_or_throw>>(); - auto match_tensors = mutator->Mutate(op->match_tensors, inplace_mode) + auto match_tensors = mutator->MutateExpected(op->match_tensors, inplace_mode) + .value() .as_or_throw>>(); auto annotations = op->annotations; if (auto addresses = annotations.Get(attr::kTensorAllocatedAddr.name)) { - auto updated = mutator->Mutate(addresses.value(), InplaceMode::kDisallow); + auto updated = mutator->MutateExpected(addresses.value(), InplaceMode::kDisallow).value(); if (!updated.IsUnchanged()) { annotations.Set(attr::kTensorAllocatedAddr.name, std::move(updated).ValueUnchecked()); } } - auto reads = - mutator->Mutate(op->reads, inplace_mode).as_or_throw>>(); - auto writes = mutator->Mutate(op->writes, inplace_mode) + auto reads = mutator->MutateExpected(op->reads, inplace_mode) + .value() + .as_or_throw>>(); + auto writes = mutator->MutateExpected(op->writes, inplace_mode) + .value() .as_or_throw>>(); - auto init = mutator->Mutate(op->init, inplace_mode); - auto body = mutator->Mutate(op->body, inplace_mode); + auto init = mutator->MutateExpected(op->init, inplace_mode) + .value() + .as_or_throw>>(); + auto body = + mutator->MutateExpected(op->body, inplace_mode).value().as_or_throw>(); if (iter_vars.UnchangedOrSameAs(op->iter_vars) && alloc_tensors.UnchangedOrSameAs(op->alloc_tensors) && reads.UnchangedOrSameAs(op->reads) && writes.UnchangedOrSameAs(op->writes) && match_tensors.UnchangedOrSameAs(op->match_tensors) && @@ -158,7 +170,7 @@ UnchangedOr StmtExprMutator::MutateBlock(tirx::StmtExprMutator* mutator, c annotations.same_as(op->annotations)) return ffi::Unchanged(); if (inplace_mode == InplaceMode::kAllow) { - auto* writable = const_cast(op); + auto* writable = const_cast(op); writable->annotations = std::move(annotations); if (!iter_vars.IsUnchanged()) writable->iter_vars = std::move(iter_vars).ValueUnchecked(); if (!alloc_tensors.IsUnchanged()) @@ -171,7 +183,7 @@ UnchangedOr StmtExprMutator::MutateBlock(tirx::StmtExprMutator* mutator, c if (!body.IsUnchanged()) writable->body = std::move(body).ValueUnchecked(); return ffi::Unchanged(); } - auto copy = ffi::make_object(*op); + auto copy = ffi::make_object(*op); copy->annotations = std::move(annotations); if (!iter_vars.IsUnchanged()) copy->iter_vars = std::move(iter_vars).ValueUnchecked(); if (!alloc_tensors.IsUnchanged()) copy->alloc_tensors = std::move(alloc_tensors).ValueUnchecked(); @@ -183,28 +195,32 @@ UnchangedOr StmtExprMutator::MutateBlock(tirx::StmtExprMutator* mutator, c return Stmt(std::move(copy)); } -UnchangedOr StmtExprMutator::Mutate_(const SBlockRealizeNode* op, InplaceMode inplace_mode) { +UnchangedOr StmtExprMutator::Mutate_(const SBlockRealizeObj* op, InplaceMode inplace_mode) { return MutateBlockRealize(this, op, inplace_mode); } -UnchangedOr StmtExprMutator::MutateBlockRealize(tirx::StmtExprMutator* mutator, - const SBlockRealizeNode* op, +UnchangedOr StmtExprMutator::MutateBlockRealize(ffi::StructuralMutatorObj* mutator, + const SBlockRealizeObj* op, InplaceMode inplace_mode) { - auto iter_values = mutator->Mutate(op->iter_values, inplace_mode) + auto iter_values = mutator->MutateExpected(op->iter_values, inplace_mode) + .value() .as_or_throw>>(); - auto predicate = mutator->Mutate(op->predicate, inplace_mode); - auto block = mutator->Mutate(op->block, inplace_mode).as_or_throw>(); + auto predicate = mutator->MutateExpected(op->predicate, inplace_mode) + .value() + .as_or_throw>(); + auto block = + mutator->MutateExpected(op->block, inplace_mode).value().as_or_throw>(); if (iter_values.UnchangedOrSameAs(op->iter_values) && predicate.UnchangedOrSameAs(op->predicate) && block.UnchangedOrSameAs(op->block)) return ffi::Unchanged(); if (inplace_mode == InplaceMode::kAllow) { - auto* writable = const_cast(op); + auto* writable = const_cast(op); if (!iter_values.IsUnchanged()) writable->iter_values = std::move(iter_values).ValueUnchecked(); if (!predicate.IsUnchanged()) writable->predicate = std::move(predicate).ValueUnchecked(); if (!block.IsUnchanged()) writable->block = std::move(block).ValueUnchecked(); return ffi::Unchanged(); } - auto copy = ffi::make_object(*op); + auto copy = ffi::make_object(*op); if (!iter_values.IsUnchanged()) copy->iter_values = std::move(iter_values).ValueUnchecked(); if (!predicate.IsUnchanged()) copy->predicate = std::move(predicate).ValueUnchecked(); if (!block.IsUnchanged()) copy->block = std::move(block).ValueUnchecked(); diff --git a/src/s_tir/tensor_intrin.cc b/src/s_tir/tensor_intrin.cc index 427f987e12e7..71077778ee60 100644 --- a/src/s_tir/tensor_intrin.cc +++ b/src/s_tir/tensor_intrin.cc @@ -27,9 +27,9 @@ namespace tvm { namespace s_tir { using tirx::Function; -using tirx::TensorTypeNode; +using tirx::TensorTypeObj; -TVM_FFI_STATIC_INIT_BLOCK() { TensorIntrinNode::RegisterReflection(); } +TVM_FFI_STATIC_INIT_BLOCK() { TensorIntrinObj::RegisterReflection(); } class TensorIntrinManager { public: @@ -47,7 +47,7 @@ TensorIntrin::TensorIntrin(Function desc, Function impl) { << "The number of parameters of the description and the implementation of the " "tensor intrinsic doesn't match."; auto is_handle = [](const Var& param) { - return param->ty.as() != nullptr || param->ty.as() != nullptr; + return param->ty.as() != nullptr || param->ty.as() != nullptr; }; for (size_t i = 0; i < desc->params.size(); i++) { TVM_FFI_CHECK(is_handle(desc->params[i]), ValueError) @@ -57,8 +57,8 @@ TensorIntrin::TensorIntrin(Function desc, Function impl) { << "Parameters of the implementation of " "the tensor intrinsic should be handle only."; } - ffi::ObjectPtr n = - ffi::make_object(std::move(desc), std::move(impl)); + ffi::ObjectPtr n = + ffi::make_object(std::move(desc), std::move(impl)); data_ = std::move(n); } diff --git a/src/s_tir/tensor_region.cc b/src/s_tir/tensor_region.cc index 2837f06b40eb..76fdb488fd52 100644 --- a/src/s_tir/tensor_region.cc +++ b/src/s_tir/tensor_region.cc @@ -42,7 +42,7 @@ TensorRegion TensorRegionFromPoint(tirx::TensorVar tensor, ffi::Array << "Tensor rank and region dimension mismatch"; ffi::Array region; for (const PrimExpr& index : indices) { - if (const prim::RampNode* ramp_index = index.as()) { + if (const prim::RampObj* ramp_index = index.as()) { region.push_back( Range::FromMinExtent(ramp_index->base, ramp_index->stride * ramp_index->lanes)); } else { diff --git a/src/s_tir/transform/annotate_irregular_loop.cc b/src/s_tir/transform/annotate_irregular_loop.cc index 069c62e56db9..0edc27ac23e5 100644 --- a/src/s_tir/transform/annotate_irregular_loop.cc +++ b/src/s_tir/transform/annotate_irregular_loop.cc @@ -37,7 +37,7 @@ class IrregularLoopAnnotator : public StmtExprMutator { using StmtExprMutator::Mutate; using StmtExprMutator::Mutate_; UnchangedOr Mutate(ffi::AnyView value, InplaceMode inplace_mode) override { - if (value.as()) return ffi::Unchanged(); + if (value.as()) return ffi::Unchanged(); return StmtExprMutator::Mutate(value, inplace_mode); } @@ -48,7 +48,7 @@ class IrregularLoopAnnotator : public StmtExprMutator { IrregularLoopAnnotator() = default; private: - UnchangedOr Mutate_(const ForNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const ForObj* op, InplaceMode inplace_mode) final { bool cur_has_jump = has_jump_; has_jump_ = false; For res = StmtExprMutator::Mutate_(op, inplace_mode) @@ -70,7 +70,7 @@ class IrregularLoopAnnotator : public StmtExprMutator { return res; } - UnchangedOr Mutate_(const WhileNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const WhileObj* op, InplaceMode inplace_mode) final { bool cur_has_jump = has_jump_; has_jump_ = false; Stmt res = StmtExprMutator::Mutate_(op, inplace_mode).ValueOrUnchanged(ffi::GetRef(op)); @@ -78,12 +78,12 @@ class IrregularLoopAnnotator : public StmtExprMutator { return res; } - UnchangedOr Mutate_(const BreakNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const BreakObj* op, InplaceMode inplace_mode) final { has_jump_ = true; return ffi::Unchanged(); } - UnchangedOr Mutate_(const ContinueNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const ContinueObj* op, InplaceMode inplace_mode) final { has_jump_ = true; return ffi::Unchanged(); } diff --git a/src/s_tir/transform/canonicalize_loop.cc b/src/s_tir/transform/canonicalize_loop.cc index 1ba9701744ba..5f339b35cead 100644 --- a/src/s_tir/transform/canonicalize_loop.cc +++ b/src/s_tir/transform/canonicalize_loop.cc @@ -46,7 +46,7 @@ class LoopCanonicalizer : public StmtExprMutator { LoopCanonicalizer() = default; private: - UnchangedOr Mutate_(const ForNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const ForObj* op, InplaceMode inplace_mode) final { if (IsZero(op->begin) && op->HasTrivialStep()) { return StmtExprMutator::Mutate_(op, inplace_mode); } @@ -67,14 +67,14 @@ class LoopCanonicalizer : public StmtExprMutator { PrimExpr extent = analyzer_->Simplify(ceildiv(Range::GetExtentFromBeginEnd(op->begin, op->end), step)); if (inplace_mode == InplaceMode::kAllow) { - auto* writable = const_cast(op); + auto* writable = const_cast(op); writable->body = std::move(body); writable->begin = std::move(min); writable->end = std::move(extent); writable->step = IntImm(loop_var_ty, 1); return ffi::Unchanged(); } else { - auto copy = ffi::make_object(*op); + auto copy = ffi::make_object(*op); copy->body = std::move(body); copy->begin = std::move(min); copy->end = std::move(extent); @@ -90,7 +90,7 @@ namespace transform { Pass CanonicalizeLoop() { auto pass_func = [=](Function func, IRModule m, PassContext ctx) { - FunctionNode* fptr = func.CopyOnWrite(); + FunctionObj* fptr = func.CopyOnWrite(); fptr->body = ffi::make_object() ->Mutate(fptr->body, InplaceMode::kAllow) .ValueOrUnchanged(std::move(fptr->body)); diff --git a/src/s_tir/transform/compact_tensor_region.cc b/src/s_tir/transform/compact_tensor_region.cc index 5846dfa347f1..4805d0ce2cf8 100644 --- a/src/s_tir/transform/compact_tensor_region.cc +++ b/src/s_tir/transform/compact_tensor_region.cc @@ -81,18 +81,18 @@ class Var2TensorCollector : public StmtExprVisitor { var2tensor_; private: - ffi::Optional Visit_(const TensorStoreNode* op) final { + ffi::Optional Visit_(const TensorStoreObj* op) final { var2tensor_[op->dest.as_or_throw()].insert(op->dest.as_or_throw()); return StmtExprVisitor::Visit_(op); } - ffi::Optional Visit_(const TensorLoadNode* op) final { + ffi::Optional Visit_(const TensorLoadObj* op) final { TensorVar tensor = op->source.as_or_throw(); var2tensor_[tensor].insert(tensor); return StmtExprVisitor::Visit_(op); } - ffi::Optional Visit_(const SBlockNode* op) final { + ffi::Optional Visit_(const SBlockObj* op) final { for (const TensorVar& tensor : op->alloc_tensors) { var2tensor_[tensor].insert(tensor); } @@ -104,7 +104,7 @@ class Var2TensorCollector : public StmtExprVisitor { return StmtExprVisitor::Visit_(op); } - ffi::Optional Visit_(const BindNode* op) final { + ffi::Optional Visit_(const BindObj* op) final { if (auto tensor = op->var.as()) { var2tensor_[op->var].insert(tensor.value()); } @@ -158,7 +158,7 @@ class TensorAccessRegionCollector : public StmtExprVisitor { /**************** Visitor overload ****************/ // Declared regions carry bounds, not opaque runtime accesses. - ffi::Optional Visit_(const TensorRegionNode* op) final { + ffi::Optional Visit_(const TensorRegionObj* op) final { if (!op->source.as()) return StmtExprVisitor::Visit_(op); for (const Range& range : op->region) { TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Visit(range->begin)); @@ -167,12 +167,12 @@ class TensorAccessRegionCollector : public StmtExprVisitor { return std::nullopt; } - ffi::Optional Visit_(const TensorStoreNode* op) final { + ffi::Optional Visit_(const TensorStoreObj* op) final { VisitTensorAccess(s_tir::TensorRegionFromPoint(op->dest.as_or_throw(), op->indices)); return Visit(op->value); } - ffi::Optional Visit_(const TensorLoadNode* op) final { + ffi::Optional Visit_(const TensorLoadObj* op) final { TensorVar tensor = op->source.as_or_throw(); auto explicit_it = explicit_access_annotations_.find(tensor); if (explicit_it != explicit_access_annotations_.end()) { @@ -186,60 +186,82 @@ class TensorAccessRegionCollector : public StmtExprVisitor { return std::nullopt; } - ffi::Optional Visit_(const VarNode* op) final { + ffi::Optional Visit_(const VarObj* op) final { if (def_region_kind() != kTVMFFIDefRegionKindNone) return StmtExprVisitor::Visit_(op); VisitTensorVar(ffi::GetRef(op)); return std::nullopt; } - ffi::Optional Visit_(const ForNode* op) final { + ffi::Optional Visit_(const ForObj* op) final { + auto outer_context = dom_analyzer_->GetConstraintContext(); Range loop_range = Range(op->begin, op->end); IterVar iter = tvm::s_tir::GetThreadBinding(op).has_value() ? IterVar(std::nullopt, op->var, IterVarType::kThreadIndex, tvm::s_tir::GetThreadBinding(op).value()) : IterVar(std::nullopt, op->var, IterVarType::kDataPar); ancestor_iters_.push_back(iter); - dom_analyzer_->Bind(op->var, loop_range); + { + auto context = dom_analyzer_->GetConstraintContext(); + context->Append(sym::RangeConstraint(op->var, loop_range)); + dom_analyzer_->SyncConstraintContext(context); + } dom_map_.Set(op->var, sym::IntSet::FromRange(loop_range)); size_t n_pending_before = pending_flat_alloc_tensors_.size(); - TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(op)); + if (auto interrupt = StmtExprVisitor::Visit_(op)) { + dom_analyzer_->SyncConstraintContext(outer_context); + return interrupt; + } // Compact flat AllocTensors defined inside this For scope CompactPendingFlatAllocTensors(n_pending_before); + dom_analyzer_->SyncConstraintContext(outer_context); dom_map_.erase(op->var); ancestor_iters_.pop_back(); return std::nullopt; } - ffi::Optional Visit_(const BindNode* op) final { - if (const auto* call = op->value.as(); + ffi::Optional Visit_(const BindObj* op) final { + if (const auto* call = op->value.as(); call && call->op.same_as(tirx::alloc_tensor_op())) { return DispatchAllocTensor(op); } - if (const auto* call = op->value.as(); + if (const auto* call = op->value.as(); call && call->op.same_as(tirx::decl_tensor_op())) return StmtExprVisitor::Visit_(op); TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit(op->value)); if (auto value = op->value.as(); value && IsIndexTypedExpr(value.value())) { - dom_analyzer_->Bind(op->var, value.value()); + { + auto context = dom_analyzer_->GetConstraintContext(); + context->Append(sym::BindingConstraint(op->var, value.value())); + dom_analyzer_->SyncConstraintContext(context); + } dom_map_.Set(op->var, sym::IntSet::SinglePoint(value.value())); } return std::nullopt; } - ffi::Optional Visit_(const LetNode* op) final { + ffi::Optional Visit_(const LetObj* op) final { TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit(op->value)); + auto outer_context = dom_analyzer_->GetConstraintContext(); if (IsIndexTypedExpr(op->value)) { - dom_analyzer_->Bind(op->var, op->value); + { + auto context = dom_analyzer_->GetConstraintContext(); + context->Append(sym::BindingConstraint(op->var, op->value)); + dom_analyzer_->SyncConstraintContext(context); + } dom_map_.Set(op->var, sym::IntSet::SinglePoint(op->value)); } - TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit(op->body)); + if (auto interrupt = StmtExprVisitor::Visit(op->body)) { + dom_analyzer_->SyncConstraintContext(outer_context); + return interrupt; + } + dom_analyzer_->SyncConstraintContext(outer_context); if (IsIndexTypedExpr(op->value)) { dom_map_.erase(op->var); } return std::nullopt; } - ffi::Optional Visit_(const IfNode* op) final { + ffi::Optional Visit_(const IfObj* op) final { // Visit condition TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit(op->condition)); { @@ -257,7 +279,7 @@ class TensorAccessRegionCollector : public StmtExprVisitor { return std::nullopt; } - ffi::Optional Visit_(const CallNode* op) final { + ffi::Optional Visit_(const CallObj* op) final { if (op->op.same_as(prim::if_then_else_op())) { PrimExpr condition = op->args[0].as_or_throw(); PrimExpr then_value = op->args[1].as_or_throw(); @@ -279,7 +301,7 @@ class TensorAccessRegionCollector : public StmtExprVisitor { return StmtExprVisitor::Visit_(op); } - ffi::Optional Visit_(const SBlockNode* op) final { + ffi::Optional Visit_(const SBlockObj* op) final { // Step 0. Check there is no init part and block is opaque TVM_FFI_ICHECK(!op->init.has_value()); TVM_FFI_ICHECK_EQ(op->iter_vars.size(), 0) << "CompactTensorRegion only works on opaque blocks"; @@ -349,30 +371,39 @@ class TensorAccessRegionCollector : public StmtExprVisitor { return std::nullopt; } - ffi::Optional Visit_(const SBlockRealizeNode* op) final { + ffi::Optional Visit_(const SBlockRealizeObj* op) final { With ctx(op->predicate, &dom_map_, &hint_map_, &pending_conditions_); return StmtExprVisitor::Visit_(op); } - ffi::Optional DispatchAllocTensor(const BindNode* op) { + ffi::Optional DispatchAllocTensor(const BindObj* op) { // AllocTensor is flat: register the tensor def and track for post-scope compaction. RecordTensorDefinition(op->var.as_or_throw()); pending_flat_alloc_tensors_.push_back(op->var.as_or_throw()); return StmtExprVisitor::Visit_(op); } - ffi::Optional Visit_(const RegionStmtNode* op) final { + ffi::Optional Visit_(const RegionStmtObj* op) final { if (op->op.same_as(tirx::launch_thread_op())) { + auto outer_context = dom_analyzer_->GetConstraintContext(); PrimExpr extent = op->args[1].as_or_throw(); Range dom = Range::FromMinExtent(IntImm(extent.ty(), 0), extent); IterVar iter(dom, op->body_params[0].as_or_throw(), IterVarType::kThreadIndex, op->args[0].as_or_throw()->value); ancestor_iters_.push_back(iter); - dom_analyzer_->Bind(iter->var, dom); + { + auto context = dom_analyzer_->GetConstraintContext(); + context->Append(sym::RangeConstraint(iter->var, dom)); + dom_analyzer_->SyncConstraintContext(context); + } dom_map_.Set(iter->var, sym::IntSet::FromRange(dom)); size_t n_pending_before = pending_flat_alloc_tensors_.size(); - TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(op)); + if (auto interrupt = StmtExprVisitor::Visit_(op)) { + dom_analyzer_->SyncConstraintContext(outer_context); + return interrupt; + } CompactPendingFlatAllocTensors(n_pending_before); + dom_analyzer_->SyncConstraintContext(outer_context); dom_map_.erase(iter->var); ancestor_iters_.pop_back(); return std::nullopt; @@ -511,7 +542,7 @@ class TensorAccessRegionCollector : public StmtExprVisitor { result_region.Set(i, original); continue; } - auto is_loop_var = [this](const VarNode* v) { + auto is_loop_var = [this](const VarObj* v) { return std::any_of(ancestor_iters_.begin(), ancestor_iters_.end(), [v](const IterVar& n) { return n->var.get() == v; }); }; @@ -626,19 +657,19 @@ class TensorCompactor : public StmtExprMutator { explicit TensorCompactor(std::unordered_map tensor_info) : tensor_info_(std::move(tensor_info)) {} - UnchangedOr Mutate_(const TensorStoreNode* _op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const TensorStoreObj* _op, InplaceMode inplace_mode) final { TensorVar original_tensor = _op->dest.as_or_throw(); TensorStore store = StmtExprMutator::Mutate_(_op, inplace_mode) .ValueOrUnchanged(ffi::GetRef(_op)) .as_or_throw(); - TensorStoreNode* op = store.CopyOnWrite(); + TensorStoreObj* op = store.CopyOnWrite(); TensorVar tensor = op->dest.as_or_throw(); RewriteTensorAccess(original_tensor, &tensor, &op->indices); op->dest = std::move(tensor); return store; } - UnchangedOr Mutate_(const TensorLoadNode* _op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const TensorLoadObj* _op, InplaceMode inplace_mode) final { TensorVar original_tensor = _op->source.as_or_throw(); TensorLoad load = StmtExprMutator::Mutate_(_op, inplace_mode) .ValueOrUnchanged(ffi::GetRef(_op)) @@ -649,12 +680,12 @@ class TensorCompactor : public StmtExprMutator { return TensorLoad(tensor, indices, load->loc); } - UnchangedOr Mutate_(const SBlockNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const SBlockObj* op, InplaceMode inplace_mode) final { // Step 0. Check there is no Init part. TVM_FFI_ICHECK(!op->init.has_value()); // Rewrite the signature while its tensor identities still match tensor_info_. SBlock block = ffi::GetRef(op); - SBlockNode* n = block.CopyOnWrite(); + SBlockObj* n = block.CopyOnWrite(); RewriteTensorRegions(&n->reads); RewriteTensorRegions(&n->writes); RewriteMatchTensors(&n->match_tensors); @@ -666,8 +697,8 @@ class TensorCompactor : public StmtExprMutator { .ValueOrUnchanged(block); } - UnchangedOr Mutate_(const BindNode* op, InplaceMode inplace_mode) final { - const auto* call = op->value.as(); + UnchangedOr Mutate_(const BindObj* op, InplaceMode inplace_mode) final { + const auto* call = op->value.as(); if (!call || (!call->op.same_as(tirx::alloc_tensor_op()) && !call->op.same_as(tirx::decl_tensor_op()))) { return StmtExprMutator::Mutate_(op, inplace_mode); @@ -682,7 +713,7 @@ class TensorCompactor : public StmtExprMutator { // Update the producer before generic Bind mutation propagates its result type. ffi::Array args = call->args; - auto rewritten = ffi::make_object(*op); + auto rewritten = ffi::make_object(*op); rewritten->value = Call(new_tensor.ty(), call->op, args, call->attrs, {new_tensor.ty()}, call->loc); tvm::Bind binding(std::move(rewritten)); @@ -743,7 +774,7 @@ class TensorCompactor : public StmtExprMutator { new_regions.reserve(regions->size()); for (const auto& region : *regions) { TensorRegion tensor_region = region; - TensorRegionNode* p = tensor_region.CopyOnWrite(); + TensorRegionObj* p = tensor_region.CopyOnWrite(); TensorVar source = p->source.as_or_throw(); RewriteTensorRegion(&source, &p->region); p->source = source; @@ -757,7 +788,7 @@ class TensorCompactor : public StmtExprMutator { result.reserve(match_tensors->size()); for (const auto& match_tensor : *match_tensors) { const TensorRegion& tensor_region = match_tensor->source; - auto p = ffi::make_object(*tensor_region.get()); + auto p = ffi::make_object(*tensor_region.get()); TensorVar source = p->source.as_or_throw(); RewriteTensorRegion(&source, &p->region); p->source = source; @@ -822,13 +853,13 @@ Stmt TensorCompactorCompact( // prepare new tensor ffi::Array shape = region.Map([](const Range& range) { return range->extent; }); ffi::Array strides = CalcStrides(alloc_info, shape); - ffi::ObjectPtr n = ffi::make_object(*tensor.ty().get()); + ffi::ObjectPtr n = ffi::make_object(*tensor.ty().get()); n->shape = std::move(shape); if (strides.empty()) { if (n->layout.has_value()) { auto tile = n->layout.value().as_or_throw(); - n->layout = tvm::tirx::TileLayout(tvm::tirx::TileLayoutNode::DefaultLayout(n->shape)->shard, - {}, tile->offset); + n->layout = tvm::tirx::TileLayout(tvm::tirx::TileLayout::DefaultLayout(n->shape)->shard, {}, + tile->offset); } } else { static const auto memory_axis = tvm::tirx::Axis::Get("m"); @@ -855,7 +886,7 @@ namespace transform { Pass CompactTensorAllocation(bool is_strict) { auto pass_func = [=](Function f, IRModule m, PassContext ctx) { if (!f->body.has_value()) return f; - FunctionNode* fptr = f.CopyOnWrite(); + FunctionObj* fptr = f.CopyOnWrite(); auto region = TensorAccessRegionCollector::Collect(f, /*collect_inbound=*/is_strict); auto storage_align = CollectStorageAlignAnnotation(f->body.value()); fptr->body = TensorCompactorCompact(f, region, storage_align); diff --git a/src/s_tir/transform/convert_blocks_to_opaque.cc b/src/s_tir/transform/convert_blocks_to_opaque.cc index f62d9a69b853..b66702a4c41b 100644 --- a/src/s_tir/transform/convert_blocks_to_opaque.cc +++ b/src/s_tir/transform/convert_blocks_to_opaque.cc @@ -49,7 +49,7 @@ class OpaqueBlockConverter : public StmtExprMutator { OpaqueBlockConverter() = default; private: - UnchangedOr Mutate_(const VarNode* var, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const VarObj* var, InplaceMode inplace_mode) final { if (def_region_kind() == kTVMFFIDefRegionKindNone) { TVM_FFI_ICHECK(!forbidden_iter_vars_.count(var)) << "Variable " << var->name << " occurs in the predicate or iter_values of a block, " @@ -58,16 +58,16 @@ class OpaqueBlockConverter : public StmtExprMutator { return StmtExprMutator::Mutate_(var, inplace_mode); } - UnchangedOr Mutate_(const SBlockNode* block, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const SBlockObj* block, InplaceMode inplace_mode) final { TVM_FFI_ICHECK(!block->init.has_value()) << "Block Init part is not allowed in pass ConvertBlocksToOpaque"; - auto node = ffi::make_object(*block); + auto node = ffi::make_object(*block); node->iter_vars.clear(); return StmtExprMutator::Mutate_(node.get(), inplace_mode) .ValueOrUnchanged(SBlock(std::move(node))); } - UnchangedOr Mutate_(const SBlockRealizeNode* realize, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const SBlockRealizeObj* realize, InplaceMode inplace_mode) final { const auto* block_op = realize->block.get(); TVM_FFI_ICHECK(!block_op->init.has_value()); @@ -105,14 +105,14 @@ class OpaqueBlockConverter : public StmtExprMutator { } /* \brief Variables that may not occur in the current context */ - std::unordered_set forbidden_iter_vars_; + std::unordered_set forbidden_iter_vars_; }; namespace transform { Pass ConvertBlocksToOpaque() { auto pass_func = [=](Function f, IRModule m, PassContext ctx) { - FunctionNode* fptr = f.CopyOnWrite(); + FunctionObj* fptr = f.CopyOnWrite(); fptr->body = OpaqueBlockConverter::Convert(f); return f; }; diff --git a/src/s_tir/transform/default_gpu_schedule.cc b/src/s_tir/transform/default_gpu_schedule.cc index 8059a92f8cd7..ccd310895787 100644 --- a/src/s_tir/transform/default_gpu_schedule.cc +++ b/src/s_tir/transform/default_gpu_schedule.cc @@ -69,10 +69,10 @@ void ThreadBind(s_tir::Schedule sch, const s_tir::SBlockRV& block, int64_t max_t s_tir::LoopRV fused = sch->Fuse(data_parallel_loops, /*preserve_unit_iters=*/false); int64_t product = std::numeric_limits::max(); if (Range::GetExtentFromBeginEnd(sch->Get(fused)->begin, sch->Get(fused)->end) - ->IsInstance()) { + ->IsInstance()) { product = static_cast( Range::GetExtentFromBeginEnd(sch->Get(fused)->begin, sch->Get(fused)->end) - .as() + .as() ->value); } // schedule the fused loop @@ -96,7 +96,7 @@ IRModule MarkScheduled(const IRModule& mod) { ffi::Map result; for (const auto& [gv, base_func] : mod->functions) { - if (const auto* function_node = base_func.as()) { + if (const auto* function_node = base_func.as()) { tirx::Function function = ffi::GetRef(function_node); tirx::Function new_function = WithAttr(std::move(function), tvm::s_tir::attr::kIsScheduled.name, true); @@ -127,7 +127,7 @@ IRModule MarkScheduled(const IRModule& mod) { */ tirx::Function WrapBareSBlockBody(const tirx::Function& func) { const auto* realize = func->body.has_value() && func->body.value()->size() == 1 - ? func->body.value()->seq[0].as() + ? func->body.value()->seq[0].as() : nullptr; if (realize == nullptr || !realize->block->iter_vars.empty()) { return func; @@ -137,8 +137,8 @@ tirx::Function WrapBareSBlockBody(const tirx::Function& func) { // whose block body is a For loop (or a nested SBlockRealize) — that case // already has somewhere to put thread bindings, so leave it alone. const SeqStmt& inner = realize->block->body; - if (inner->size() == 1 && (inner->seq[0]->IsInstance() || - inner->seq[0]->IsInstance())) { + if (inner->size() == 1 && (inner->seq[0]->IsInstance() || + inner->seq[0]->IsInstance())) { return func; } tvm::IntImm zero(tvm::PrimType::Int(32), 0); @@ -193,7 +193,7 @@ Pass DefaultGPUSchedule() { ffi::Map wrapped; bool any_wrapped = false; for (const auto& [gv, base_func] : m->functions) { - if (const auto* function_node = base_func.as(); + if (const auto* function_node = base_func.as(); function_node != nullptr && IsScheduledOnGPU(base_func) && !base_func->HasNonzeroAttr(tvm::s_tir::attr::kIsScheduled.name)) { tirx::Function func = ffi::GetRef(function_node); @@ -212,7 +212,7 @@ Pass DefaultGPUSchedule() { s_tir::Schedule sch = s_tir::Schedule::Traced(m, /*seed=*/-1, /*debug_mask=*/0, s_tir::ScheduleErrorRenderLevel::kDetail); for (const auto& [gv, func] : m->functions) { - if (func->IsInstance() && + if (func->IsInstance() && !func->HasNonzeroAttr(tvm::s_tir::attr::kIsScheduled.name) && IsScheduledOnGPU(func)) { // get the target from context. diff --git a/src/s_tir/transform/inject_permuted_layout.cc b/src/s_tir/transform/inject_permuted_layout.cc index 98661dccf63c..d31a78bdcc93 100644 --- a/src/s_tir/transform/inject_permuted_layout.cc +++ b/src/s_tir/transform/inject_permuted_layout.cc @@ -27,7 +27,7 @@ #include #include #include -#include +#include #include #include #include @@ -39,16 +39,17 @@ using namespace tvm::tirx; using namespace sym; using namespace runtime; -class PermutedLayoutInjector : public StmtExprMutatorWithAnalyzer { +class PermutedLayoutInjector : public StmtExprMutatorWithConstraintContext { public: - using StmtExprMutatorWithAnalyzer::Mutate; - using StmtExprMutatorWithAnalyzer::Mutate_; + using StmtExprMutatorWithConstraintContext::Mutate; + using StmtExprMutatorWithConstraintContext::Mutate_; static Function Transform(Function func) { Analyzer analyzer; auto new_body = ffi::make_object(analyzer) ->Mutate(func->body) + .as_or_throw>>() .ValueOrUnchanged(func->body); auto func_node = func.CopyOnWrite(); func_node->body = new_body; @@ -56,7 +57,7 @@ class PermutedLayoutInjector : public StmtExprMutatorWithAnalyzer { } explicit PermutedLayoutInjector(const Analyzer& analyzer) - : StmtExprMutatorWithAnalyzer(analyzer) {} + : StmtExprMutatorWithConstraintContext(analyzer->GetConstraintContext()) {} private: ffi::Array PermuteIndices(PrimExpr row_idx, PrimExpr col_idx, int row_size) { @@ -100,14 +101,14 @@ class PermutedLayoutInjector : public StmtExprMutatorWithAnalyzer { auto row_idx_sub = floormod(row_idx, 8); new_col_idx_outer = col_idx_outer ^ floordiv(row_idx_sub, 2); } - return {row_idx, analyzer_->Simplify(new_col_idx_outer * 8 + col_idx_inner)}; + return {row_idx, GetSyncedAnalyzer()->Simplify(new_col_idx_outer * 8 + col_idx_inner)}; } static bool CheckAnnotation(const Any& annotation) { if (auto opt_str = annotation.as()) { // Support string annotation for backward compatibility return *opt_str != ""; - } else if (auto* node = annotation.as()) { + } else if (auto* node = annotation.as()) { return node->value != 0; } else if (auto opt_val = annotation.try_cast()) { return *opt_val != 0; @@ -116,16 +117,16 @@ class PermutedLayoutInjector : public StmtExprMutatorWithAnalyzer { } } - UnchangedOr Mutate_(const SBlockNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const SBlockObj* op, InplaceMode inplace_mode) final { if (op->annotations.count(tvm::s_tir::attr::kPermutedLayout) == 0 || !CheckAnnotation(op->annotations.at(tvm::s_tir::attr::kPermutedLayout))) { - return StmtExprMutatorWithAnalyzer::Mutate_(op, inplace_mode); + return StmtExprMutatorWithConstraintContext::Mutate_(op, inplace_mode); } auto prev_permute = permute_; permute_ = true; - SBlock block = StmtExprMutatorWithAnalyzer::Mutate_(op, inplace_mode) + SBlock block = StmtExprMutatorWithConstraintContext::Mutate_(op, inplace_mode) .ValueOrUnchanged(ffi::GetRef(op)) .as_or_throw(); @@ -143,8 +144,8 @@ class PermutedLayoutInjector : public StmtExprMutatorWithAnalyzer { << " should be at least 2"; auto dim = tensor.shape().size(); - auto tensor_row_size = tensor.shape()[dim - 1].as()->value; - auto tensor_col_size = tensor.shape()[dim - 2].as()->value; + auto tensor_row_size = tensor.shape()[dim - 1].as()->value; + auto tensor_col_size = tensor.shape()[dim - 2].as()->value; if (tensor_row_size % 64 != 0) { TVM_FFI_ICHECK(tensor_row_size % 32 == 0) @@ -174,14 +175,14 @@ class PermutedLayoutInjector : public StmtExprMutatorWithAnalyzer { return indices; } - UnchangedOr Mutate_(const TensorStoreNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const TensorStoreObj* op, InplaceMode inplace_mode) final { static const auto scope_shared = runtime::MemoryScope::Get("shared"); static const auto scope_shared_dyn = runtime::MemoryScope::Get("shared.dyn"); // Rewrite write from global to shared.dyn or shared // We assume the shape of the shared memory is [..., row_size, col_size], // where row_size is divisible by 64, or divisible by 32 and col_size is divisible by 2. - auto store = StmtExprMutatorWithAnalyzer::Mutate_(op, inplace_mode) + auto store = StmtExprMutatorWithConstraintContext::Mutate_(op, inplace_mode) .ValueOrUnchanged(ffi::GetRef(op)) .as_or_throw(); @@ -200,12 +201,12 @@ class PermutedLayoutInjector : public StmtExprMutatorWithAnalyzer { return store; } - UnchangedOr Mutate_(const TensorLoadNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const TensorLoadObj* op, InplaceMode inplace_mode) final { static const auto scope_shared = runtime::MemoryScope::Get("shared"); static const auto scope_shared_dyn = runtime::MemoryScope::Get("shared.dyn"); // Rewrite load from shared or shared.dyn to global - auto load = StmtExprMutatorWithAnalyzer::Mutate_(op, inplace_mode) + auto load = StmtExprMutatorWithConstraintContext::Mutate_(op, inplace_mode) .ValueOrUnchanged(ffi::GetRef(op)) .as_or_throw(); diff --git a/src/s_tir/transform/inject_ptx_async_copy.cc b/src/s_tir/transform/inject_ptx_async_copy.cc index d74452569b3f..38b2e0acef22 100644 --- a/src/s_tir/transform/inject_ptx_async_copy.cc +++ b/src/s_tir/transform/inject_ptx_async_copy.cc @@ -29,9 +29,14 @@ #include #include #include +#include #include #include +#include +#include +#include + #include "storage_access.h" #include "tvm/s_tir/stmt.h" #include "tvm/tirx/stmt.h" @@ -45,11 +50,11 @@ class PTXAsyncCopyInjector : public StmtExprMutator { using StmtExprMutator::Mutate; using StmtExprMutator::Mutate_; UnchangedOr Mutate(ffi::AnyView value, InplaceMode inplace_mode) override { - if (value.as()) return ffi::Unchanged(); + if (value.as()) return ffi::Unchanged(); return StmtExprMutator::Mutate(value, inplace_mode); } - UnchangedOr Mutate_(const RegionStmtNode* region, InplaceMode inplace_mode) { + UnchangedOr Mutate_(const RegionStmtObj* region, InplaceMode inplace_mode) { if (region->op.same_as(s_tir::async_copy_scope())) { TVM_FFI_ICHECK(in_async == false) << "Nested async scopes not supported"; in_async = true; @@ -60,7 +65,41 @@ class PTXAsyncCopyInjector : public StmtExprMutator { return StmtExprMutator::Mutate_(region, inplace_mode); } - Stmt InjectPTX(const TensorLoadNode* load, const TensorStoreNode* store, bool predicated = false, + UnchangedOr Mutate_(const BindObj* op, InplaceMode inplace_mode) final { + const auto* call = op->value.as(); + if (call) { + const auto* attrs = call->attrs.as(); + // Track only ordinary fresh allocations and exact, unannotated views. + if (!call->attrs.has_value() || (attrs && attrs->dict.empty())) { + if (call->op.same_as(alloc_tensor_op()) && call->args.empty()) { + allocations_.emplace(op->var.get(), op->var.as_or_throw()); + } else if (call->op.same_as(decl_tensor_op()) && call->args.size() == 1) { + const auto* data = call->args[0].as(); + if (data && data->op.same_as(tensor_data_ptr_op()) && data->args.size() == 1) { + const auto* source = data->args[0].as(); + auto it = allocations_.find(source); + if (it != allocations_.end()) { + TensorVar view = op->var.as_or_throw(); + if (view.dtype() == it->second.dtype() && view.scope() == it->second.scope() && + !view.ty()->layout.has_value()) { + allocations_.emplace(op->var.get(), it->second); + } + } + } + } + } + } + return StmtExprMutator::Mutate_(op, inplace_mode); + } + + UnchangedOr Mutate_(const ForObj* op, InplaceMode inplace_mode) final { + if (in_async) { + if (auto copy = CoalesceHalfCopy(op)) return *copy; + } + return StmtExprMutator::Mutate_(op, inplace_mode); + } + + Stmt InjectPTX(const TensorLoadObj* load, const TensorStoreObj* store, bool predicated = false, ffi::Optional predicate_value = std::nullopt) { static const auto scope_shared_dyn = runtime::MemoryScope::Get("shared.dyn"); @@ -112,24 +151,24 @@ class PTXAsyncCopyInjector : public StmtExprMutator { if (!predicated) { // Only some vectorized indexing patterns are supported for now. auto src_offset = [=]() -> ffi::Optional { - if (load->indices[0]->IsInstance()) { - return load->indices[0].as()->base; + if (load->indices[0]->IsInstance()) { + return load->indices[0].as()->base; } return std::nullopt; }(); auto dst_offset = [=]() -> ffi::Optional { - if (store->indices[0].as()) { - return store->indices[0].as()->base; - } else if (store->indices[0].as()) { + if (store->indices[0].as()) { + return store->indices[0].as()->base; + } else if (store->indices[0].as()) { // The case where the dst tensor is a byte tensor generated by merging dynamic // shared memory. // A_shared.dyn[(ramp(...), 1, 8) + x8(17408))] = A_global[ramp(...),1, 8)] - auto* add = store->indices[0].as(); - if (!add->a->IsInstance()) return std::nullopt; - if (!add->b->IsInstance()) return std::nullopt; - return prim::Add(add->a.as()->base, - add->b.as()->value); + auto* add = store->indices[0].as(); + if (!add->a->IsInstance()) return std::nullopt; + if (!add->b->IsInstance()) return std::nullopt; + return prim::Add(add->a.as()->base, + add->b.as()->value); } return std::nullopt; }(); @@ -145,24 +184,24 @@ class PTXAsyncCopyInjector : public StmtExprMutator { } else { // Only some vectorized indexing patterns are supported for now. auto src_offset = [=]() -> ffi::Optional { - if (load->indices[0]->IsInstance()) { - return load->indices[0].as()->base; + if (load->indices[0]->IsInstance()) { + return load->indices[0].as()->base; } return std::nullopt; }(); auto dst_offset = [=]() -> ffi::Optional { - if (store->indices[0].as()) { - return store->indices[0].as()->base; - } else if (store->indices[0].as()) { + if (store->indices[0].as()) { + return store->indices[0].as()->base; + } else if (store->indices[0].as()) { // The case where the dst tensor is a byte tensor generated by merging dynamic // shared memory. // A_shared.dyn[(ramp(...), 1, 8) + x8(17408))] = A_global[ramp(...),1, 8)] - auto* add = store->indices[0].as(); - if (!add->a->IsInstance()) return std::nullopt; - if (!add->b->IsInstance()) return std::nullopt; - return prim::Add(add->a.as()->base, - add->b.as()->value); + auto* add = store->indices[0].as(); + if (!add->a->IsInstance()) return std::nullopt; + if (!add->b->IsInstance()) return std::nullopt; + return prim::Add(add->a.as()->base, + add->b.as()->value); } return std::nullopt; }(); @@ -183,28 +222,28 @@ class PTXAsyncCopyInjector : public StmtExprMutator { .ValueOrUnchanged(ffi::GetRef(store)); } - UnchangedOr Mutate_(const TensorStoreNode* store, InplaceMode inplace_mode) { + UnchangedOr Mutate_(const TensorStoreObj* store, InplaceMode inplace_mode) { static const auto scope_shared = runtime::MemoryScope::Get("shared"); static const auto scope_shared_dyn = runtime::MemoryScope::Get("shared.dyn"); if (in_async && (store->dest.as_or_throw().scope() == scope_shared || store->dest.as_or_throw().scope() == scope_shared_dyn)) { - if (auto* load = store->value.as()) { + if (auto* load = store->value.as()) { return InjectPTX(load, store); - } else if (auto* call = store->value.as()) { + } else if (auto* call = store->value.as()) { // tirx.if_then_else is a call to prim::if_then_else_op() if (call->op.same_as(prim::if_then_else_op()) && call->args.size() == 3) { - if (auto* load = call->args[1].as()) { + if (auto* load = call->args[1].as()) { // Only default value of 0 is supported since 0 is the default value used by cp.async // ptx. @see section 9.7.8.22.3. of // https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#asynchronous-memory-operations bool else_value_is_zero = false; - if (auto* b = call->args[2].as()) { - if (auto* f = b->value.as()) { + if (auto* b = call->args[2].as()) { + if (auto* f = b->value.as()) { else_value_is_zero = f->value == 0.0f; } } - if (auto* f = call->args[2].as()) { + if (auto* f = call->args[2].as()) { else_value_is_zero = f->value == 0.0f; } if (else_value_is_zero) { @@ -218,6 +257,76 @@ class PTXAsyncCopyInjector : public StmtExprMutator { } private: + ffi::Optional CoalesceHalfCopy(const ForObj* loop) { + // A short, contiguous half copy can be one hardware transaction even when + // the general loop vectorizer leaves it scalar. Do not widen arbitrary + // computation, predicates, or loops with other effects. + const auto* extent = loop->end.as(); + if (loop->kind != ForKind::kDefault || !prim::IsZero(loop->begin) || !loop->HasTrivialStep() || + !loop->annotations.empty() || !extent || + (extent->value != 2 && extent->value != 4 && extent->value != 8) || + loop->body->seq.empty()) { + return std::nullopt; + } + const int bytes = extent->value.as().value() * 2; + static const auto shared = runtime::MemoryScope::Get("shared"); + static const auto shared_dyn = runtime::MemoryScope::Get("shared.dyn"); + static const Op copy_op = Op::Get("tirx.s_tir.cp_async_raw"); + std::vector> destinations; + ffi::Array copies; + for (const Stmt& stmt : loop->body->seq) { + const auto* store = stmt.as(); + if (!store) return std::nullopt; + const auto* load = store->value.as(); + if (!load || load->indices.size() != 1 || store->indices.size() != 1) { + return std::nullopt; + } + TensorVar src = load->source.as_or_throw(); + TensorVar dst = store->dest.as_or_throw(); + auto allocation = allocations_.find(dst.get()); + if (src.scope() != runtime::MemoryScope::Global() || + (dst.scope() != shared && dst.scope() != shared_dyn) || + src.dtype() != PrimType::Float(16) || store->value.ty() != src.dtype() || + dst.dtype() != src.dtype() || src.ty()->layout.has_value() || + dst.ty()->layout.has_value() || !load->indices[0].ty().IsScalar() || + !store->indices[0].ty().IsScalar() || allocation == allocations_.end() || + allocation->second.ty()->layout.has_value() || + allocation->second.ty()->data_alignment < bytes || + allocation->second.ty()->data_alignment % bytes != 0 || + src.ty()->data_alignment < bytes || src.ty()->data_alignment % bytes != 0 || + dst.ty()->data_alignment < bytes || dst.ty()->data_alignment % bytes != 0) { + return std::nullopt; + } + auto source = DetectLinearEquation(load->indices[0], {loop->var}); + auto destination = DetectLinearEquation(store->indices[0], {loop->var}); + if (source.size() != 2 || destination.size() != 2 || !prim::IsOne(source[0]) || + !prim::IsOne(destination[0]) || SideEffect(source[1]) != CallEffectKind::kPure || + SideEffect(destination[1]) != CallEffectKind::kPure || + !analyzer_->CanProveEqual(floormod(source[1], PrimExpr(bytes / 2)), 0) || + !analyzer_->CanProveEqual(floormod(destination[1], PrimExpr(bytes / 2)), 0)) { + return std::nullopt; + } + // Storage rewriting can merge destinations into views of one allocation. + // Prove their copied intervals disjoint before reordering their stores. + for (const auto& [root, offset] : destinations) { + if (root == allocation->second.get() && + !analyzer_->CanProve(destination[1] + bytes / 2 <= offset) && + !analyzer_->CanProve(offset + bytes / 2 <= destination[1])) { + return std::nullopt; + } + } + destinations.emplace_back(allocation->second.get(), destination[1]); + // Reads are global and destinations are disjoint fresh shared storage. + copies.push_back( + Evaluate(Call(dst.dtype(), copy_op, + {dst.data(), destination[1], src.data(), source[1], PrimExpr(bytes)}) + .as_or_throw())); + } + return SeqStmt(copies); + } + + sym::Analyzer analyzer_; + std::unordered_map allocations_; bool in_async{false}; }; diff --git a/src/s_tir/transform/inject_software_pipeline.cc b/src/s_tir/transform/inject_software_pipeline.cc index bfeaf6536ba6..f019ece2d412 100644 --- a/src/s_tir/transform/inject_software_pipeline.cc +++ b/src/s_tir/transform/inject_software_pipeline.cc @@ -55,7 +55,7 @@ ffi::Optional GetBufferDataVar(const ffi::Any& data) { if (auto var = data.as()) { return var; } - if (const auto* call = data.as(); + if (const auto* call = data.as(); call && call->op.same_as(tirx::tensor_data_ptr_op()) && call->args.size() == 1) { return call->args[0].as(); } @@ -76,8 +76,8 @@ ffi::Optional GetBufferDataVar(const ffi::Any& data) { * \return The result block. */ SBlock MakeSBlock(const SeqStmt& body, const ffi::Map& tensor_data_to_tensor) { - if (const SBlockRealizeNode* block_realize = - body->size() == 1 ? body->seq[0].as() : nullptr) { + if (const SBlockRealizeObj* block_realize = + body->size() == 1 ? body->seq[0].as() : nullptr) { if (IsOne(block_realize->predicate)) { // no need to create a new block return block_realize->block; @@ -86,7 +86,7 @@ SBlock MakeSBlock(const SeqStmt& body, const ffi::Map& tensor_da SBlock block(/*iter_vars=*/{}, /*reads=*/{}, /*writes=*/{}, /*name_hint=*/"", /*body*/ body); ffi::Array> access = GetSBlockReadWriteRegion(block, tensor_data_to_tensor); - SBlockNode* n = block.CopyOnWrite(); + SBlockObj* n = block.CopyOnWrite(); n->reads = access[0]; n->writes = access[1]; return block; @@ -117,10 +117,10 @@ class PipelineOpaqueAccessRewriter { * \param pipeline_loop The original loop to be software pipelined. * \param fragment_info Information about tensor core fragment */ - PipelineOpaqueAccessRewriter( - const ffi::Map& tensor_data_to_tensor, - const ffi::Map& tensor_remap, const For& pipeline_loop, - const std::unordered_map& fragment_info) + PipelineOpaqueAccessRewriter(const ffi::Map& tensor_data_to_tensor, + const ffi::Map& tensor_remap, + const For& pipeline_loop, + const std::unordered_map& fragment_info) : tensor_data_to_tensor_(tensor_data_to_tensor), tensor_remap_(tensor_remap), pipeline_loop_(pipeline_loop), @@ -187,7 +187,7 @@ class PipelineOpaqueAccessRewriter { const ffi::Map& tensor_data_to_tensor_; const ffi::Map& tensor_remap_; const For& pipeline_loop_; - const std::unordered_map& fragment_info_; + const std::unordered_map& fragment_info_; }; /*! @@ -213,7 +213,7 @@ class PipelineBodyRewriter : public StmtExprMutator { PipelineBodyRewriter(const ffi::Map& tensor_data_to_tensor, const ffi::Map& tensor_remap, For pipeline_loop, bool access_all_versions, - const std::unordered_map& fragment_info) + const std::unordered_map& fragment_info) : tensor_data_to_tensor_(tensor_data_to_tensor), pipeline_loop_(pipeline_loop), access_all_versions_(access_all_versions), @@ -245,7 +245,7 @@ class PipelineBodyRewriter : public StmtExprMutator { return tensor_region; } - UnchangedOr Mutate_(const SBlockNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const SBlockObj* op, InplaceMode inplace_mode) final { for (const TensorVar& alloc_tensor : op->alloc_tensors) { tensor_data_to_tensor_.Set(alloc_tensor, alloc_tensor); } @@ -258,7 +258,7 @@ class PipelineBodyRewriter : public StmtExprMutator { return block; } - UnchangedOr Mutate_(const TensorRegionNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const TensorRegionObj* op, InplaceMode inplace_mode) final { if (!op->source.as()) { return StmtExprMutator::Mutate_(op, inplace_mode); } @@ -266,7 +266,7 @@ class PipelineBodyRewriter : public StmtExprMutator { return StmtExprMutator::Mutate_(region.get(), InplaceMode::kDisallow).ValueOrUnchanged(region); } - UnchangedOr Mutate_(const TensorStoreNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const TensorStoreObj* op, InplaceMode inplace_mode) final { auto replacement = VarRemapGet(op->dest.as_or_throw()).as(); TensorStore store = StmtExprMutator::Mutate_(op, inplace_mode) .ValueOrUnchanged(ffi::GetRef(op)) @@ -283,7 +283,7 @@ class PipelineBodyRewriter : public StmtExprMutator { return store; } - UnchangedOr Mutate_(const TensorLoadNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const TensorLoadObj* op, InplaceMode inplace_mode) final { auto replacement = VarRemapGet(op->source).as(); if (!replacement) return StmtExprMutator::Mutate_(op, inplace_mode); auto indices = Mutate(op->indices, inplace_mode) @@ -296,7 +296,7 @@ class PipelineBodyRewriter : public StmtExprMutator { return TensorLoad(new_tensor, indices, op->loc); } - UnchangedOr Mutate_(const CallNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const CallObj* op, InplaceMode inplace_mode) final { const auto scope_mma_matrixA = GetRegisteredMemoryScope("mma.matrixA"); const auto scope_mma_matrixB = GetRegisteredMemoryScope("mma.matrixB"); @@ -342,7 +342,7 @@ class PipelineRewriter : public StmtExprMutator { const std::unordered_set& double_buffers, const ffi::Array pipeline_allocs, const For& pipeline_loop, const PipelineInfo& pipeline_info, - const std::unordered_map& fragment_info, + const std::unordered_map& fragment_info, const ffi::Map preserved_annotations) { auto rewriter = ffi::make_object( tensor_data_to_tensor, double_buffers, pipeline_allocs, pipeline_loop, pipeline_info, @@ -355,7 +355,7 @@ class PipelineRewriter : public StmtExprMutator { const std::unordered_set& double_buffers, const ffi::Array& pipeline_allocs, const For& pipeline_loop, const PipelineInfo& pipeline_info, - const std::unordered_map& fragment_info, + const std::unordered_map& fragment_info, const ffi::Map preserved_annotations) : tensor_data_to_tensor_(std::move(tensor_data_to_tensor)), @@ -561,8 +561,7 @@ class PipelineRewriter : public StmtExprMutator { * \return The resized tensor. */ TensorVar RewriteAllocTensor(const TensorVar& tensor, int num_versions) { - ffi::ObjectPtr new_tensor = - ffi::make_object(*tensor.ty().get()); + ffi::ObjectPtr new_tensor = ffi::make_object(*tensor.ty().get()); auto strides = tvm::tirx::MatchShapeStridePattern(tensor.ty()).value().get<1>(); new_tensor->shape.insert(new_tensor->shape.begin(), PrimExpr(num_versions)); if (new_tensor->layout.has_value()) { @@ -585,7 +584,7 @@ class PipelineRewriter : public StmtExprMutator { // Per-stage states that need to be tracked across pipeline prologue, body, and epilogue. struct AsyncStateGlobal { // Tensors that this stage asynchronously writes. - std::unordered_set dst_tensors; + std::unordered_set dst_tensors; // An imaginary index that the latest async operation associated with this stage has written // into. Only valid if all associated predicates are true, so that we can count the number of // async invocations exactly. When it is valid, it is the "sum of extents of loops that have @@ -619,7 +618,7 @@ class PipelineRewriter : public StmtExprMutator { // until a point where we encounter a consumer of async result tensors. This is used to decide // if the producer_head of each tensor points to a copy written in the current or previous // iteration. - std::unordered_set seen; + std::unordered_set seen; // A symbolic expression representing the index the latest async operation associated with this // stage has written into, at the "current" iteration. @@ -647,7 +646,7 @@ class PipelineRewriter : public StmtExprMutator { // Determine where to insert async_wait and the corresponding wait count. void PopulateWaitCounts(const std::vector& new_blocks, sym::AnalyzerObj* ana_normalized, - const std::unordered_map& tensor_to_commit_group, + const std::unordered_map& tensor_to_commit_group, std::map* async_states_local) { for (size_t i = 0; i < new_blocks.size(); ++i) { if (new_blocks[i].is_async) { @@ -799,7 +798,7 @@ class PipelineRewriter : public StmtExprMutator { if (state.pending_wait.valid()) { auto attach_wait = [&new_blocks](int i, int stage_id, PrimExpr wait_count) { auto& block = new_blocks[i].block; - SBlockNode* n = block.CopyOnWrite(); + SBlockObj* n = block.CopyOnWrite(); n->body = SeqStmt({Evaluate(Call(PrimType::Void(), s_tir::async_wait(), {IntImm::Int32(stage_id), wait_count})), n->body}); @@ -873,21 +872,30 @@ class PipelineRewriter : public StmtExprMutator { new_loop_var = start; // use constants as the loop var for unit loops } else { new_loop_var = pipeline_loop_->var.CopyWithSuffix(""); - analyzer_->Bind(new_loop_var.as_or_throw(), Range(start, end)); + { + auto context = analyzer_->GetConstraintContext(); + context->Append(sym::RangeConstraint(new_loop_var.as_or_throw(), Range(start, end))); + analyzer_->SyncConstraintContext(context); + } } // In contrast to analyzer_ which is bound to [start, end), this one is bound to // the "normalized" range, [pipeline_loop_->begin, extent). sym::Analyzer ana_normalized; if (!is_unit_loop) { - ana_normalized->Bind(new_loop_var.as_or_throw(), Range(pipeline_loop_->begin, extent)); + { + auto context = ana_normalized->GetConstraintContext(); + context->Append(sym::RangeConstraint(new_loop_var.as_or_throw(), + Range(pipeline_loop_->begin, extent))); + ana_normalized->SyncConstraintContext(context); + } } std::vector new_blocks; // Async related std::map async_states_local; - std::unordered_map tensor_to_commit_group; + std::unordered_map tensor_to_commit_group; for (const SBlock& block : ordered_stmts_) { int stage = pipeline_info_.at(block).stage; @@ -976,7 +984,7 @@ class PipelineRewriter : public StmtExprMutator { local_state.predicate = ana_normalized->Simplify(local_state.predicate.value() & inbound); } - SBlockNode* n = new_block.CopyOnWrite(); + SBlockObj* n = new_block.CopyOnWrite(); n->body = RegionStmt(s_tir::async_copy_scope(), {}, {}, DictAttrs(), n->body); } @@ -1039,7 +1047,7 @@ class PipelineRewriter : public StmtExprMutator { ffi::Array pipeline_allocs_; For pipeline_loop_; PipelineInfo pipeline_info_; - const std::unordered_map& fragment_info_; + const std::unordered_map& fragment_info_; int max_stage_ = -1; ffi::Map tensor_remap_; ffi::Array ordered_stmts_; @@ -1148,7 +1156,7 @@ class PipelineInjector : public StmtExprMutator { } } - UnchangedOr Mutate_(const ForNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const ForObj* op, InplaceMode inplace_mode) final { // Step 1: Recursively rewrite the children first. For for_node = StmtExprMutator::Mutate_(op, inplace_mode) .ValueOrUnchanged(ffi::GetRef(op)) @@ -1161,9 +1169,8 @@ class PipelineInjector : public StmtExprMutator { // child of the block. SeqStmt pipeline_body = for_node->body; ffi::Array pipeline_allocs; - if (const auto* realize = for_node->body->size() == 1 - ? for_node->body->seq[0].as() - : nullptr) { + if (const auto* realize = + for_node->body->size() == 1 ? for_node->body->seq[0].as() : nullptr) { const auto& block = realize->block; for (const auto& tensor : block->alloc_tensors) { tensor_data_to_tensor_.Set(tensor, tensor); @@ -1172,7 +1179,7 @@ class PipelineInjector : public StmtExprMutator { pipeline_allocs = block->alloc_tensors; } - const SeqStmtNode* pipeline_body_seq = pipeline_body.get(); + const SeqStmtObj* pipeline_body_seq = pipeline_body.get(); TVM_FFI_CHECK_GT(pipeline_body_seq->size(), 1, ValueError) << "The body of the software pipeline should contain multiple statements"; @@ -1185,7 +1192,7 @@ class PipelineInjector : public StmtExprMutator { original_order.push_back(MakeSBlock(child, tensor_data_to_tensor_)); }; for (size_t i = 0; i < pipeline_body_seq->seq.size(); i++) { - const auto* nested_block_realize = pipeline_body_seq->seq[i].as(); + const auto* nested_block_realize = pipeline_body_seq->seq[i].as(); if (nested_block_realize && IsOne(nested_block_realize->predicate) && nested_block_realize->block->body->size() > 1) { const SBlock& nested_pipeline_block = nested_block_realize->block; @@ -1250,7 +1257,7 @@ class PipelineInjector : public StmtExprMutator { fragment_info_, preserved_annotations); if (const auto* realize = - op->body->size() == 1 ? op->body->seq[0].as() : nullptr) { + op->body->size() == 1 ? op->body->seq[0].as() : nullptr) { const auto& block = realize->block; for (const auto& tensor : block->alloc_tensors) { tensor_data_to_tensor_.erase(tensor); @@ -1264,7 +1271,7 @@ class PipelineInjector : public StmtExprMutator { * \param n The block pointer to which the tensor allocations are added. * \param alloc_tensors The tensor allocations to be added. */ - void AddAllocTensors(SBlockNode* n, const ffi::Array alloc_tensors) { + void AddAllocTensors(SBlockObj* n, const ffi::Array alloc_tensors) { for (const TensorVar& alloc_tensor : alloc_tensors) { n->alloc_tensors.push_back(alloc_tensor); ffi::Array region; @@ -1276,7 +1283,7 @@ class PipelineInjector : public StmtExprMutator { } } - UnchangedOr Mutate_(const SBlockNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const SBlockObj* op, InplaceMode inplace_mode) final { for (const auto& tensor : op->alloc_tensors) { tensor_data_to_tensor_.Set(tensor, tensor); } @@ -1300,7 +1307,7 @@ class PipelineInjector : public StmtExprMutator { return block; } - bool HasPipelineAnnotation(const ForNode* op) const { + bool HasPipelineAnnotation(const ForObj* op) const { auto it1 = op->annotations.find(tvm::s_tir::attr::kSoftwarePipelineStage.name); auto it2 = op->annotations.find(tvm::s_tir::attr::kSoftwarePipelineOrder.name); bool has_stage = it1 != op->annotations.end(); @@ -1318,7 +1325,7 @@ class PipelineInjector : public StmtExprMutator { } ffi::Map tensor_data_to_tensor_; - std::unordered_map fragment_info_; + std::unordered_map fragment_info_; std::unordered_set double_buffers; ffi::Optional global_symbol_; }; diff --git a/src/s_tir/transform/inject_virtual_thread.cc b/src/s_tir/transform/inject_virtual_thread.cc index 4673f74809c9..f6ddcd390ce5 100644 --- a/src/s_tir/transform/inject_virtual_thread.cc +++ b/src/s_tir/transform/inject_virtual_thread.cc @@ -29,7 +29,7 @@ #include #include #include -#include +#include #include #include #include @@ -49,7 +49,7 @@ ffi::Optional GetBufferDataVar(const ffi::Any& data) { if (auto var = data.as()) { return var; } - if (const auto* call = data.as(); + if (const auto* call = data.as(); call && call->op.same_as(tirx::tensor_data_ptr_op()) && call->args.size() == 1) { return call->args[0].as(); } @@ -62,7 +62,7 @@ ffi::Optional GetBufferDataVar(const ffi::Any& data) { class ExprTouched final : public StmtExprVisitor { public: using StmtExprVisitor::Visit_; - explicit ExprTouched(const std::unordered_set& touched) : touched_var_(touched) {} + explicit ExprTouched(const std::unordered_set& touched) : touched_var_(touched) {} void Reset(bool check_write) { expr_touched_ = false; @@ -77,18 +77,18 @@ class ExprTouched final : public StmtExprVisitor { return StmtExprVisitor::Visit(n); } - ffi::Optional Visit_(const TensorLoadNode* op) final { + ffi::Optional Visit_(const TensorLoadObj* op) final { HandleUseVar(op->source.as_or_throw().get()); return StmtExprVisitor::Visit_(op); } - ffi::Optional Visit_(const VarNode* op) final { + ffi::Optional Visit_(const VarObj* op) final { HandleUseVar(op); return std::nullopt; } - ffi::Optional Visit_(const CallNode* op) final { + ffi::Optional Visit_(const CallObj* op) final { if (op->op.same_as(tirx::masked_load_op()) || op->op.same_as(tirx::masked_store_op())) { bool is_load = op->op.same_as(tirx::masked_load_op()); - const VarNode* tensor = op->args[0].as_or_throw().get(); + const VarObj* tensor = op->args[0].as_or_throw().get(); if (is_load) { HandleUseVar(tensor); } else { @@ -98,15 +98,15 @@ class ExprTouched final : public StmtExprVisitor { TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->Visit(op->args[i])); } } else if (op->op.same_as(tirx::tensor_data_ptr_op())) { - const VarNode* tensor = op->args[0].as_or_throw().get(); + const VarObj* tensor = op->args[0].as_or_throw().get(); HandleUseVar(tensor); // An escaping address carries no direction information. Conservatively // connect it to the other operands of an opaque statement so virtual // threads never share storage that the statement may write. if (check_write_) HandleWriteVar(tensor); - } else if (op->op.same_as(tirx::address_of_op()) && op->args[0].as()) { - const auto* load = op->args[0].as(); - const VarNode* tensor = load->source.as_or_throw().get(); + } else if (op->op.same_as(tirx::address_of_op()) && op->args[0].as()) { + const auto* load = op->args[0].as(); + const VarObj* tensor = load->source.as_or_throw().get(); HandleUseVar(tensor); if (check_write_) HandleWriteVar(tensor); for (const PrimExpr& index : load->indices) { @@ -117,7 +117,7 @@ class ExprTouched final : public StmtExprVisitor { } return std::nullopt; } - void HandleUseVar(const VarNode* var) { + void HandleUseVar(const VarObj* var) { auto it = touched_var_.find(var); if (it != touched_var_.end()) { expr_touched_ = true; @@ -128,12 +128,12 @@ class ExprTouched final : public StmtExprVisitor { used_vars_.push_back(var); } } - void HandleWriteVar(const VarNode* var) { write_vars_.push_back(var); } + void HandleWriteVar(const VarObj* var) { write_vars_.push_back(var); } // the fields. bool expr_touched_{false}; - std::vector used_vars_; - std::vector write_vars_; - const std::unordered_set& touched_var_; + std::vector used_vars_; + std::vector write_vars_; + const std::unordered_set& touched_var_; bool check_write_{false}; }; @@ -142,15 +142,15 @@ class VarTouchedAnalysis : public StmtExprVisitor { public: using StmtExprVisitor::Visit_; ffi::Optional Visit(ffi::AnyView value) override { - if (value.as()) return std::nullopt; + if (value.as()) return std::nullopt; return StmtExprVisitor::Visit(value); } - ffi::Optional Visit_(const BindNode* op) final { - if (const auto* call = op->value.as(); + ffi::Optional Visit_(const BindObj* op) final { + if (const auto* call = op->value.as(); call && call->op.same_as(tirx::alloc_tensor_op())) { return DispatchAllocTensor(op, call); } - if (const auto* call = op->value.as(); + if (const auto* call = op->value.as(); call && call->op.same_as(tirx::decl_tensor_op())) return StmtExprVisitor::Visit_(op); expr_touched_->Reset(false); @@ -159,7 +159,7 @@ class VarTouchedAnalysis : public StmtExprVisitor { return std::nullopt; } - ffi::Optional Visit_(const TensorStoreNode* op) final { + ffi::Optional Visit_(const TensorStoreObj* op) final { expr_touched_->Reset(false); expr_touched_->Visit(op->value); for (const auto& index : op->indices) { @@ -168,7 +168,7 @@ class VarTouchedAnalysis : public StmtExprVisitor { Record(op->dest.as_or_throw().get(), *expr_touched_); return std::nullopt; } - ffi::Optional Visit_(const ForNode* op) final { + ffi::Optional Visit_(const ForObj* op) final { expr_touched_->Reset(false); expr_touched_->Visit(op->begin); expr_touched_->Visit(op->end); @@ -176,7 +176,7 @@ class VarTouchedAnalysis : public StmtExprVisitor { Record(op->var.get(), *expr_touched_); return this->Visit(op->body); } - ffi::Optional Visit_(const RegionStmtNode* op) final { + ffi::Optional Visit_(const RegionStmtObj* op) final { if (!op->op.same_as(tirx::launch_thread_op())) has_opaque_region_ = true; expr_touched_->Reset(false); for (const Expr& arg : op->args) expr_touched_->Visit(arg); @@ -187,15 +187,15 @@ class VarTouchedAnalysis : public StmtExprVisitor { bool has_opaque_region_{false}; // external function call - ffi::Optional Visit_(const EvaluateNode* op) final { + ffi::Optional Visit_(const EvaluateObj* op) final { expr_touched_->Reset(true); expr_touched_->Visit(op->value); - for (const VarNode* var : expr_touched_->write_vars_) { + for (const VarObj* var : expr_touched_->write_vars_) { Record(var, *expr_touched_); } return std::nullopt; } - ffi::Optional DispatchAllocTensor(const BindNode* op, const CallNode* call) { + ffi::Optional DispatchAllocTensor(const BindObj* op, const CallObj* call) { expr_touched_->Reset(false); tvm::Tuple shape = tvm::Tuple(call->ty_args[0].as_or_throw()->shape); for (const Expr& extent : shape->fields) { @@ -204,12 +204,12 @@ class VarTouchedAnalysis : public StmtExprVisitor { Record(op->var.as_or_throw().get(), *expr_touched_); return StmtExprVisitor::Visit_(op); } - void Record(const VarNode* var, const ExprTouched& tc) { + void Record(const VarObj* var, const ExprTouched& tc) { if (touched_var_.count(var)) return; if (tc.expr_touched_) { touched_var_.insert(var); } else { - for (const VarNode* r : tc.used_vars_) { + for (const VarObj* r : tc.used_vars_) { if (r != var) { affect_[r].push_back(var); } @@ -217,15 +217,15 @@ class VarTouchedAnalysis : public StmtExprVisitor { } } - std::unordered_set TouchedVar(const Stmt& stmt, const VarNode* var) { + std::unordered_set TouchedVar(const Stmt& stmt, const VarObj* var) { touched_var_.insert(var); this->Visit(stmt); // do a DFS to push affect around dependency. - std::vector pending(touched_var_.begin(), touched_var_.end()); + std::vector pending(touched_var_.begin(), touched_var_.end()); while (!pending.empty()) { - const VarNode* v = pending.back(); + const VarObj* v = pending.back(); pending.pop_back(); - for (const VarNode* r : affect_[v]) { + for (const VarObj* r : affect_[v]) { if (!touched_var_.count(r)) { touched_var_.insert(r); pending.push_back(r); @@ -237,23 +237,23 @@ class VarTouchedAnalysis : public StmtExprVisitor { private: // Whether variable is touched by the thread variable. - std::unordered_set touched_var_; + std::unordered_set touched_var_; ffi::ObjectPtr expr_touched_ = ffi::make_object(touched_var_); // x -> all the buffers x read from - std::unordered_map> affect_; + std::unordered_map> affect_; }; // Inject virtual thread loop // rewrite the buffer access pattern when necessary. -class VTInjector : public s_tir::StmtExprMutatorWithAnalyzer { +class VTInjector : public s_tir::StmtExprMutatorWithConstraintContext { public: - using s_tir::StmtExprMutatorWithAnalyzer::Mutate; - using s_tir::StmtExprMutatorWithAnalyzer::Mutate_; + using s_tir::StmtExprMutatorWithConstraintContext::Mutate; + using s_tir::StmtExprMutatorWithConstraintContext::Mutate_; // constructor VTInjector(sym::AnalyzerObj* analyzer, Var var, int num_threads, - const std::unordered_set& touched_var, bool allow_share) - : StmtExprMutatorWithAnalyzer(analyzer), + const std::unordered_set& touched_var, bool allow_share) + : StmtExprMutatorWithConstraintContext(analyzer->GetConstraintContext()), var_(var), num_threads_(num_threads), touched_var_(touched_var), @@ -261,9 +261,10 @@ class VTInjector : public s_tir::StmtExprMutatorWithAnalyzer { // Inject VTLoop when needed. UnchangedOr Mutate(ffi::AnyView value, InplaceMode inplace_mode = InplaceMode::kDisallow) final { - if (!value.as()) return StmtExprMutator::Mutate(value, inplace_mode); + if (!value.as()) + return StmtExprMutatorWithConstraintContext::Mutate(value, inplace_mode); TVM_FFI_ICHECK(!visit_touched_var_); - auto result = StmtExprMutator::Mutate(value, inplace_mode); + auto result = StmtExprMutatorWithConstraintContext::Mutate(value, inplace_mode); if (visit_touched_var_ || trigger_base_inject_) { if (!vt_loop_injected_) { Stmt stmt = std::move(result).ValueOrUnchanged(value).as_or_throw(); @@ -275,9 +276,9 @@ class VTInjector : public s_tir::StmtExprMutatorWithAnalyzer { return result; } // Variable - UnchangedOr Mutate_(const TensorRegionNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const TensorRegionObj* op, InplaceMode inplace_mode) final { if (!op->source.as()) { - return StmtExprMutator::Mutate_(op, inplace_mode); + return StmtExprMutatorWithConstraintContext::Mutate_(op, inplace_mode); } auto region = Mutate(op->region).as_or_throw>>(); if (region.UnchangedOrSameAs(op->region)) return ffi::Unchanged(); @@ -286,7 +287,7 @@ class VTInjector : public s_tir::StmtExprMutatorWithAnalyzer { return node; } - UnchangedOr Mutate_(const VarNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const VarObj* op, InplaceMode inplace_mode) final { if (def_region_kind() == kTVMFFIDefRegionKindNone) { TVM_FFI_ICHECK(!alloc_remap_.count(op)) << "TensorVar address may get rewritten in virtual thread"; @@ -294,13 +295,13 @@ class VTInjector : public s_tir::StmtExprMutatorWithAnalyzer { visit_touched_var_ = true; } } - return StmtExprMutator::Mutate_(op, inplace_mode); + return StmtExprMutatorWithConstraintContext::Mutate_(op, inplace_mode); } - PrimExpr RewriteIndex(PrimExpr index, PrimExpr alloc_extent) const { - return analyzer_->Simplify(index + var_.as_or_throw() * alloc_extent); + PrimExpr RewriteIndex(PrimExpr index, PrimExpr alloc_extent) { + return GetSyncedAnalyzer()->Simplify(index + var_.as_or_throw() * alloc_extent); } // Expression. - UnchangedOr Mutate_(const CallNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const CallObj* op, InplaceMode inplace_mode) final { if (op->op.same_as(tirx::masked_load_op()) || op->op.same_as(tirx::masked_store_op())) { bool is_load = op->op.same_as(tirx::masked_load_op()); TensorVar tensor = op->args[0].as_or_throw(); @@ -331,7 +332,7 @@ class VTInjector : public s_tir::StmtExprMutatorWithAnalyzer { auto tensor_var = GetBufferDataVar(ffi::GetRef(op)).value(); auto it = alloc_remap_.find(tensor_var.get()); if (it == alloc_remap_.end()) { - return StmtExprMutator::Mutate_(op, inplace_mode); + return StmtExprMutatorWithConstraintContext::Mutate_(op, inplace_mode); } visit_touched_var_ = true; TensorVar tensor = tensor_var.as_or_throw(); @@ -342,15 +343,15 @@ class VTInjector : public s_tir::StmtExprMutatorWithAnalyzer { RewriteIndex(PrimExpr(0), it->second) * bytes}, {}, {}, op->loc); } else { - return StmtExprMutator::Mutate_(op, inplace_mode); + return StmtExprMutatorWithConstraintContext::Mutate_(op, inplace_mode); } } - UnchangedOr Mutate_(const EvaluateNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const EvaluateObj* op, InplaceMode inplace_mode) final { trigger_base_inject_ = !allow_share_; - return StmtExprMutator::Mutate_(op, inplace_mode); + return StmtExprMutatorWithConstraintContext::Mutate_(op, inplace_mode); } // TensorLoad - UnchangedOr Mutate_(const TensorLoadNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const TensorLoadObj* op, InplaceMode inplace_mode) final { auto indices = Mutate(op->indices).as_or_throw>>(); TensorLoad node = ffi::GetRef(op); if (!indices.UnchangedOrSameAs(op->indices)) { @@ -359,7 +360,7 @@ class VTInjector : public s_tir::StmtExprMutatorWithAnalyzer { return VisitTensorAccess(std::move(node)); } // TensorStore - UnchangedOr Mutate_(const TensorStoreNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const TensorStoreObj* op, InplaceMode inplace_mode) final { auto value = Mutate(op->value); auto indices = Mutate(op->indices).as_or_throw>>(); TensorStore node = ffi::GetRef(op); @@ -406,20 +407,21 @@ class VTInjector : public s_tir::StmtExprMutatorWithAnalyzer { } TensorVar GetRemappedTensor(TensorVar tensor, PrimExpr alloc_extent) { - if (auto replacement = VarRemapGet(tensor).as()) return replacement.value(); + if (auto replacement = VarRemapGetExpected(tensor).value().as()) + return replacement.value(); TensorVar original = tensor; TVM_FFI_ICHECK_EQ(tensor.shape().size(), 1) << "Expected tensors being rewritten to already be flattened."; - auto writer = ffi::make_object(*tensor.ty().get()); + auto writer = ffi::make_object(*tensor.ty().get()); writer->shape = {tensor.shape()[0] * alloc_extent}; tensor = TensorVar(tensor->name, TensorType(std::move(writer)), tensor->loc); - VarRemapSet(original, tensor); + VarRemapSetExpected(original, tensor).value(); return tensor; } - UnchangedOr Mutate_(const RegionStmtNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const RegionStmtObj* op, InplaceMode inplace_mode) final { // A launch extent is evaluated before entering its body. If it depends on // the virtual index, put the virtual loop outside the entire launch. auto args = @@ -438,14 +440,14 @@ class VTInjector : public s_tir::StmtExprMutatorWithAnalyzer { } // Bind - UnchangedOr Mutate_(const BindNode* op, InplaceMode inplace_mode) final { - if (const auto* call = op->value.as(); + UnchangedOr Mutate_(const BindObj* op, InplaceMode inplace_mode) final { + if (const auto* call = op->value.as(); call && call->op.same_as(tirx::alloc_tensor_op())) { return Mutate_AllocTensor(op, call, inplace_mode); } - if (const auto* call = op->value.as(); + if (const auto* call = op->value.as(); call && call->op.same_as(tirx::decl_tensor_op())) - return StmtExprMutator::Mutate_(op, inplace_mode); + return StmtExprMutatorWithConstraintContext::Mutate_(op, inplace_mode); auto value_result = this->Mutate(op->value, inplace_mode); bool value_unchanged = value_result.UnchangedOrSameAs(op->value); Expr value = std::move(value_result).ValueOrUnchanged(op->value); @@ -460,7 +462,7 @@ class VTInjector : public s_tir::StmtExprMutatorWithAnalyzer { } } // For - UnchangedOr Mutate_(const ForNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const ForObj* op, InplaceMode inplace_mode) final { TVM_FFI_ICHECK(IsZero(op->begin)); auto extent_result = this->Mutate(op->end, inplace_mode); bool extent_unchanged = extent_result.UnchangedOrSameAs(op->end); @@ -479,12 +481,12 @@ class VTInjector : public s_tir::StmtExprMutatorWithAnalyzer { return ffi::Unchanged(); } else { if (inplace_mode == InplaceMode::kAllow) { - auto* writable = const_cast(op); + auto* writable = const_cast(op); writable->end = std::move(extent); writable->body = std::move(body); return ffi::Unchanged(); } else { - auto copy = ffi::make_object(*op); + auto copy = ffi::make_object(*op); copy->end = std::move(extent); copy->body = std::move(body); return For(std::move(copy)); @@ -492,7 +494,7 @@ class VTInjector : public s_tir::StmtExprMutatorWithAnalyzer { } } // If - UnchangedOr Mutate_(const IfNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const IfObj* op, InplaceMode inplace_mode) final { auto condition_result = this->Mutate(op->condition, inplace_mode); bool condition_unchanged = condition_result.UnchangedOrSameAs(op->condition); PrimExpr condition = std::move(condition_result).ValueOrUnchanged(op->condition); @@ -520,9 +522,9 @@ class VTInjector : public s_tir::StmtExprMutatorWithAnalyzer { } // While - UnchangedOr Mutate_(const WhileNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const WhileObj* op, InplaceMode inplace_mode) final { // TODO(masahi): What should we do for While nodes? - TVM_FFI_THROW(InternalError) << "WhileNode in InjectVirtualThread not supported yet"; + TVM_FFI_THROW(InternalError) << "WhileObj in InjectVirtualThread not supported yet"; TVM_FFI_UNREACHABLE(); } @@ -533,7 +535,7 @@ class VTInjector : public s_tir::StmtExprMutatorWithAnalyzer { // With flat Bind (no body), a Bind whose value touches vt_var must be // grouped with all remaining siblings and wrapped in a VT loop together. // This preserves the scoping that was implicit when Bind carried a body. - UnchangedOr Mutate_(const SeqStmtNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const SeqStmtObj* op, InplaceMode inplace_mode) final { TVM_FFI_ICHECK_EQ(max_loop_depth_, 0); ffi::Array new_seq; bool changed = false; @@ -543,7 +545,7 @@ class VTInjector : public s_tir::StmtExprMutatorWithAnalyzer { // For Bind children, pre-check if the value touches vt_var before // visiting. If so, group the Bind with all remaining siblings and // wrap the group with InjectVTLoop. - if (const auto* bind = op->seq[i].as(); bind && !vt_loop_injected_) { + if (const auto* bind = op->seq[i].as(); bind && !vt_loop_injected_) { // Visit just the value expression to probe for vt_var dependency. TVM_FFI_ICHECK(!visit_touched_var_); // This is a probe: the original Bind is revisited when forming the group. @@ -582,7 +584,7 @@ class VTInjector : public s_tir::StmtExprMutatorWithAnalyzer { } // Allocate // AllocTensor - UnchangedOr Mutate_AllocTensor(const BindNode* op, const CallNode* call, + UnchangedOr Mutate_AllocTensor(const BindObj* op, const CallObj* call, InplaceMode inplace_mode) { tvm::Tuple shape = tvm::Tuple(call->ty_args[0].as_or_throw()->shape); ffi::Array original_shape = shape->fields.as_or_throw>(); @@ -612,12 +614,12 @@ class VTInjector : public s_tir::StmtExprMutatorWithAnalyzer { if (new_shape.same_as(original_shape) && args.same_as(call->args)) { return ffi::Unchanged(); } else { - auto type = ffi::make_object(*op->var.as_or_throw().ty().get()); + auto type = ffi::make_object(*op->var.as_or_throw().ty().get()); type->shape = new_shape; TensorVar new_tensor = TensorVar(op->var.as_or_throw()->name, TensorType(std::move(type)), op->var.as_or_throw()->loc); - VarRemapSet(op->var.as_or_throw(), new_tensor); + VarRemapSetExpected(op->var.as_or_throw(), new_tensor).value(); return Bind(new_tensor, Call(new_tensor.ty(), tirx::alloc_tensor_op(), args, call->attrs, @@ -684,7 +686,7 @@ class VTInjector : public s_tir::StmtExprMutatorWithAnalyzer { // the counter of loops in after mutation. int max_loop_depth_{0}; // The variables that get touched. - const std::unordered_set& touched_var_; + const std::unordered_set& touched_var_; // Whether allow shareding. bool allow_share_; /* \brief The allocations that get touched -> extent @@ -693,7 +695,7 @@ class VTInjector : public s_tir::StmtExprMutatorWithAnalyzer { * extent of the allocation. Used when rewriting the indices of * TensorLoad/TensorStore. */ - std::unordered_map alloc_remap_; + std::unordered_map alloc_remap_; /*! \brief Map of tensors that are modified. * * Tensors allocated or written to within the virtual thread loop @@ -703,20 +705,21 @@ class VTInjector : public s_tir::StmtExprMutatorWithAnalyzer { */ }; -class VirtualThreadInjector : public s_tir::StmtExprMutatorWithAnalyzer { +class VirtualThreadInjector : public s_tir::StmtExprMutatorWithConstraintContext { public: - using s_tir::StmtExprMutatorWithAnalyzer::Mutate; - using s_tir::StmtExprMutatorWithAnalyzer::Mutate_; + using s_tir::StmtExprMutatorWithConstraintContext::Mutate; + using s_tir::StmtExprMutatorWithConstraintContext::Mutate_; - using StmtExprMutatorWithAnalyzer::StmtExprMutatorWithAnalyzer; + using StmtExprMutatorWithConstraintContext::StmtExprMutatorWithConstraintContext; - UnchangedOr Mutate_(const RegionStmtNode* op, InplaceMode inplace_mode) final { - Stmt stmt = StmtExprMutator::Mutate_(op, inplace_mode).ValueOrUnchanged(ffi::GetRef(op)); - op = stmt.as(); + UnchangedOr Mutate_(const RegionStmtObj* op, InplaceMode inplace_mode) final { + Stmt stmt = StmtExprMutatorWithConstraintContext::Mutate_(op, inplace_mode) + .ValueOrUnchanged(ffi::GetRef(op)); + op = stmt.as(); if (op && op->op.same_as(tirx::launch_thread_op()) && std::string(op->args[0].as_or_throw()->value).rfind("vthread", 0) == 0) { PrimVar var = op->body_params[0].as_or_throw(); - const auto* extent = op->args[1].as_or_throw().as(); + const auto* extent = op->args[1].as_or_throw().as(); TVM_FFI_ICHECK(extent) << "Virtual thread extent must be a constant integer"; int nthread = extent->value.as().value(); TVM_FFI_ICHECK_GT(nthread, 0) << "Virtual thread extent must be positive"; @@ -728,8 +731,8 @@ class VirtualThreadInjector : public s_tir::StmtExprMutatorWithAnalyzer { return For(var, IntImm(var.ty(), 0), op->args[1].as_or_throw(), IntImm(var.ty(), 1), ForKind::kDefault, op->body); } - auto injector = - ffi::make_object(analyzer_, var, nthread, touched, /*allow_share=*/true); + auto injector = ffi::make_object(GetSyncedAnalyzer(), var, nthread, touched, + /*allow_share=*/true); return injector->Mutate(op->body).ValueOrUnchanged(op->body); } else { return stmt; @@ -746,7 +749,7 @@ Pass InjectVirtualThread() { sym::Analyzer analyzer; - n->body = ffi::make_object(analyzer) + n->body = ffi::make_object(analyzer->GetConstraintContext()) ->Mutate(n->body.value(), InplaceMode::kAllow) .ValueOrUnchanged(std::move(n->body).value()); n->body = s_tir::ConvertSSA(std::move(n->body).value()); diff --git a/src/s_tir/transform/ir_utils.cc b/src/s_tir/transform/ir_utils.cc index 7c2c8983e8a6..2d5326e5e8c8 100644 --- a/src/s_tir/transform/ir_utils.cc +++ b/src/s_tir/transform/ir_utils.cc @@ -49,7 +49,7 @@ class SIRConvertSSA final : public StmtExprMutator { Function VisitFunction(Function func); IRModule VisitIRModule(IRModule mod); - UnchangedOr Mutate_(const SBlockNode* op, InplaceMode mode) { + UnchangedOr Mutate_(const SBlockObj* op, InplaceMode mode) { SBlock block = ffi::GetRef(op); return scope_.WithNewScope([&]() -> Stmt { auto iter_vars = op->iter_vars.Map([&](IterVar iter) { @@ -64,31 +64,31 @@ class SIRConvertSSA final : public StmtExprMutator { }); } - UnchangedOr Mutate_(const SBlockRealizeNode* op, InplaceMode mode) { + UnchangedOr Mutate_(const SBlockRealizeObj* op, InplaceMode mode) { return s_tir::StmtExprMutator::MutateBlockRealize(this, op, mode); } private: - UnchangedOr Mutate_(const VarNode* op, InplaceMode inplace_mode) final; + UnchangedOr Mutate_(const VarObj* op, InplaceMode inplace_mode) final; Var DefineVar(Var var); ffi::Optional GetVarRemap(const Var& var); - UnchangedOr Mutate_(const BindNode* op, InplaceMode inplace_mode) final; - UnchangedOr Mutate_(const IfNode* op, InplaceMode inplace_mode) final; - UnchangedOr Mutate_(const ForNode* op, InplaceMode inplace_mode) final; - UnchangedOr Mutate_(const WhileNode* op, InplaceMode inplace_mode) final; - UnchangedOr Mutate_(const RegionStmtNode* op, InplaceMode inplace_mode) final; - UnchangedOr Mutate_(const prim::LetNode* op, InplaceMode inplace_mode) final; + UnchangedOr Mutate_(const BindObj* op, InplaceMode inplace_mode) final; + UnchangedOr Mutate_(const IfObj* op, InplaceMode inplace_mode) final; + UnchangedOr Mutate_(const ForObj* op, InplaceMode inplace_mode) final; + UnchangedOr Mutate_(const WhileObj* op, InplaceMode inplace_mode) final; + UnchangedOr Mutate_(const RegionStmtObj* op, InplaceMode inplace_mode) final; + UnchangedOr Mutate_(const prim::LetObj* op, InplaceMode inplace_mode) final; static Var MakeNewVar(const Var& old_var); void SetVarRemap(const Var& old_var, const Var& new_var); // Lexical visibility ends with each scope, but definitions stay unique across the module. - ScopeStack> scope_; - std::unordered_set defined_; + ScopeStack> scope_; + std::unordered_set defined_; }; Function SIRConvertSSA::VisitFunction(Function func) { return scope_.WithNewScope([&]() -> Function { - std::unordered_set parameter_symbols; + std::unordered_set parameter_symbols; // Define explicit parameters before the symbolic values in their types. for (const Var& param : func->params) { DefineVar(param); @@ -130,7 +130,7 @@ Function SIRConvertSSA::VisitFunction(Function func) { }); } -UnchangedOr SIRConvertSSA::Mutate_(const VarNode* op, InplaceMode inplace_mode) { +UnchangedOr SIRConvertSSA::Mutate_(const VarObj* op, InplaceMode inplace_mode) { Var var = ffi::GetRef(op); if (def_region_kind() != kTVMFFIDefRegionKindNone) { if (def_region_kind() == kTVMFFIDefRegionKindPattern) { @@ -143,7 +143,7 @@ UnchangedOr SIRConvertSSA::Mutate_(const VarNode* op, InplaceMode inplace_ return ffi::Unchanged(); } -UnchangedOr SIRConvertSSA::Mutate_(const prim::LetNode* op, InplaceMode inplace_mode) { +UnchangedOr SIRConvertSSA::Mutate_(const prim::LetObj* op, InplaceMode inplace_mode) { ffi::Any previous_remap = VarRemapGet(op->var); PrimExpr result = scope_.WithNewScope([&] { return StmtExprMutator::Mutate_(op, inplace_mode).ValueOrUnchanged(ffi::GetRef(op)); @@ -172,18 +172,18 @@ ffi::Optional SIRConvertSSA::GetVarRemap(const Var& var) { return std::nullopt; } -UnchangedOr SIRConvertSSA::Mutate_(const BindNode* op, InplaceMode inplace_mode) { +UnchangedOr SIRConvertSSA::Mutate_(const BindObj* op, InplaceMode inplace_mode) { Var var = op->var; ffi::Any previous_remap = VarRemapGet(var); // The ordinary Bind path rewrites the RHS before defining the Var and propagates its type. Stmt stmt = StmtExprMutator::Mutate_(op, inplace_mode).ValueOrUnchanged(ffi::GetRef(op)); - Var bound_var = stmt.as()->var; + Var bound_var = stmt.as()->var; if (!bound_var.same_as(GetVarRemap(var).value_or(var))) SetVarRemap(var, bound_var); VarRemapSet(var, previous_remap); return stmt; } -UnchangedOr SIRConvertSSA::Mutate_(const IfNode* op, InplaceMode inplace_mode) { +UnchangedOr SIRConvertSSA::Mutate_(const IfObj* op, InplaceMode inplace_mode) { // Each branch gets its own scope so Bind remaps in one branch // do not leak into the other. auto condition_result = Mutate(op->condition, inplace_mode); @@ -204,7 +204,7 @@ UnchangedOr SIRConvertSSA::Mutate_(const IfNode* op, InplaceMode inplace_m return If(condition, then_case, else_case); } -UnchangedOr SIRConvertSSA::Mutate_(const ForNode* op, InplaceMode inplace_mode) { +UnchangedOr SIRConvertSSA::Mutate_(const ForObj* op, InplaceMode inplace_mode) { const Var& v = op->var; if (defined_.count(v.get())) { return scope_.WithNewScope([&]() -> Stmt { @@ -212,7 +212,7 @@ UnchangedOr SIRConvertSSA::Mutate_(const ForNode* op, InplaceMode inplace_ SetVarRemap(v, new_var); Stmt stmt = StmtExprMutator::Mutate_(op, inplace_mode).ValueOrUnchanged(ffi::GetRef(op)); - auto n = ffi::make_object(*stmt.as()); + auto n = ffi::make_object(*stmt.as()); n->var = new_var.as_or_throw(); return For(n); }); @@ -225,13 +225,13 @@ UnchangedOr SIRConvertSSA::Mutate_(const ForNode* op, InplaceMode inplace_ } } -UnchangedOr SIRConvertSSA::Mutate_(const WhileNode* op, InplaceMode inplace_mode) { +UnchangedOr SIRConvertSSA::Mutate_(const WhileObj* op, InplaceMode inplace_mode) { return scope_.WithNewScope([&]() -> Stmt { return StmtExprMutator::Mutate_(op, inplace_mode).ValueOrUnchanged(ffi::GetRef(op)); }); } -UnchangedOr SIRConvertSSA::Mutate_(const RegionStmtNode* op, InplaceMode inplace_mode) { +UnchangedOr SIRConvertSSA::Mutate_(const RegionStmtObj* op, InplaceMode inplace_mode) { auto attrs = Mutate(op->attrs, InplaceMode::kDisallow) .as_or_throw>() .ValueOrUnchanged(op->attrs); @@ -261,7 +261,7 @@ IRModule SIRConvertSSA::VisitIRModule(IRModule mod) { ffi::Map functions; bool made_change = false; for (auto [gvar, base_func] : mod->functions) { - if (auto* ptr = base_func.as()) { + if (auto* ptr = base_func.as()) { auto updated = VisitFunction(ffi::GetRef(ptr)); if (!updated.same_as(base_func)) { made_change = true; @@ -349,7 +349,7 @@ ffi::Array ConvertRegion(const MatchTensorRegion& match_tensor, class StorageAlignCollector : public StmtExprVisitor { public: ffi::Optional Visit(ffi::AnyView value) override { - if (value.as()) return std::nullopt; + if (value.as()) return std::nullopt; return StmtExprVisitor::Visit(value); } @@ -358,7 +358,7 @@ class StorageAlignCollector : public StmtExprVisitor { const Stmt& body); /*! \brief SBlock: resolve each annotation's tensor index through the write regions. */ - ffi::Optional Visit_(const SBlockNode* op) final { + ffi::Optional Visit_(const SBlockObj* op) final { auto it = op->annotations.find(attr::kTensorDimAlign.name); if (it != op->annotations.end()) { auto annotation = (*it).second.as_or_throw(); @@ -370,15 +370,15 @@ class StorageAlignCollector : public StmtExprVisitor { } /*! \brief AllocTensor: check for tensor_dim_align annotations. */ - ffi::Optional Visit_(const BindNode* op) final { - if (const auto* call = op->value.as(); + ffi::Optional Visit_(const BindObj* op) final { + if (const auto* call = op->value.as(); call && call->op.same_as(tirx::alloc_tensor_op())) { return DispatchAllocTensor(op, call); } return StmtExprVisitor::Visit_(op); } - ffi::Optional DispatchAllocTensor(const BindNode* op, const CallNode* call) { + ffi::Optional DispatchAllocTensor(const BindObj* op, const CallObj* call) { DictAttrs annotations = call->attrs.as_or_throw(); auto it = annotations->dict.find(attr::kTensorDimAlign.name); if (it != annotations->dict.end()) { diff --git a/src/s_tir/transform/ir_utils.h b/src/s_tir/transform/ir_utils.h index 6bcc93785dff..a36823da4367 100644 --- a/src/s_tir/transform/ir_utils.h +++ b/src/s_tir/transform/ir_utils.h @@ -86,7 +86,7 @@ struct FragmentInfo { * \param stmt The stmt to visit. * \return Map from tensor variables to the fragment info. */ -std::unordered_map GetTensorCoreFragmentInfo(const Stmt& stmt); +std::unordered_map GetTensorCoreFragmentInfo(const Stmt& stmt); /*! \brief The quad used by StorageAlign for (tensor_idx, axis, factor, offset) */ using StorageAlignTuple = ffi::Tuple; diff --git a/src/s_tir/transform/lift_thread_binding.cc b/src/s_tir/transform/lift_thread_binding.cc index fdf23c7cf9c4..41e5f2709f3b 100644 --- a/src/s_tir/transform/lift_thread_binding.cc +++ b/src/s_tir/transform/lift_thread_binding.cc @@ -40,10 +40,10 @@ FindLoopLCA(const Stmt& root) { public: using StmtExprVisitor::Visit_; ffi::Optional Visit(ffi::AnyView value) override { - if (value.as()) return std::nullopt; + if (value.as()) return std::nullopt; return StmtExprVisitor::Visit(value); } - ffi::Optional Visit_(const ForNode* op) final { + ffi::Optional Visit_(const ForObj* op) final { stack.push_back(ffi::GetRef(op)); TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(op)); if (tvm::s_tir::GetThreadBinding(op).has_value()) { @@ -53,7 +53,7 @@ FindLoopLCA(const Stmt& root) { return std::nullopt; } - void UpdateLCA(const ForNode* loop) { + void UpdateLCA(const ForObj* loop) { std::string thread_tag = tvm::s_tir::GetThreadBinding(loop).value(); { ffi::Map* tgt = &annotations[thread_tag]; @@ -121,7 +121,7 @@ class ThreadBindingLifter : public StmtExprMutator { using StmtExprMutator::Mutate; using StmtExprMutator::Mutate_; - UnchangedOr Mutate_(const ForNode* _op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const ForObj* _op, InplaceMode inplace_mode) final { For op = ffi::GetRef(_op); bool is_kernel_root = false; if (tvm::s_tir::GetThreadBinding(op).has_value()) { @@ -155,7 +155,7 @@ class ThreadBindingLifter : public StmtExprMutator { } } - void SetKernelRoot(const ForNode* op) { + void SetKernelRoot(const ForObj* op) { auto result = FindLoopLCA(ffi::GetRef(op)); this->iter_lca = std::move(result.first); for (const auto& [var, replacement] : result.second) VarRemapSet(var, replacement); @@ -170,7 +170,7 @@ namespace transform { Pass LiftThreadBinding() { auto pass_func = [=](Function f, IRModule m, PassContext ctx) { - FunctionNode* fptr = f.CopyOnWrite(); + FunctionObj* fptr = f.CopyOnWrite(); fptr->body = ffi::make_object() ->Mutate(fptr->body, InplaceMode::kAllow) .ValueOrUnchanged(std::move(fptr->body)); diff --git a/src/s_tir/transform/loop_partition.cc b/src/s_tir/transform/loop_partition.cc index ca3133a0df93..df5161525f67 100644 --- a/src/s_tir/transform/loop_partition.cc +++ b/src/s_tir/transform/loop_partition.cc @@ -111,13 +111,13 @@ using Partition = std::unordered_map; // Virtual threads are lowered separately and are not hardware partition scopes. -static bool IsVirtualThread(const ForNode* op) { +static bool IsVirtualThread(const ForObj* op) { if (!tvm::s_tir::GetThreadBinding(op).has_value()) return false; const auto tag = tvm::s_tir::GetThreadBinding(op).value(); return tag.starts_with("vthread"); } -static bool HasPartitionHint(const ForNode* op) { +static bool HasPartitionHint(const ForObj* op) { auto it = op->annotations.find(tvm::s_tir::attr::kLoopPartitionHint.name); return it != op->annotations.end() && IsOne((*it).second.as_or_throw()); } @@ -129,7 +129,7 @@ class NormalizePartitionHints : public StmtExprMutator { using StmtExprMutator::Mutate; using StmtExprMutator::Mutate_; - UnchangedOr Mutate_(const ForNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const ForObj* op, InplaceMode inplace_mode) final { For loop = StmtExprMutator::Mutate_(op, inplace_mode) .ValueOrUnchanged(ffi::GetRef(op)) .as_or_throw(); @@ -162,9 +162,9 @@ class CandidateSelector final : public StmtExprVisitor { explicit CandidateSelector(bool partition_const_loop) : partition_const_loop_(partition_const_loop) {} - ffi::Optional Visit_(const ForNode* op) final { + ffi::Optional Visit_(const ForObj* op) final { if (IsVirtualThread(op)) return StmtExprVisitor::Visit_(op); - const VarNode* var = op->var.get(); + const VarObj* var = op->var.get(); if (HasPartitionHint(op)) { candidates.insert(ffi::GetRef(op)); return StmtExprVisitor::Visit_(op); @@ -188,7 +188,7 @@ class CandidateSelector final : public StmtExprVisitor { return std::nullopt; } - ffi::Optional Visit_(const RegionStmtNode* op) final { + ffi::Optional Visit_(const RegionStmtObj* op) final { if (op->op.same_as(tirx::launch_thread_op()) && !op->args[0].as_or_throw()->value.starts_with("vthread")) { Var var = op->body_params[0].as_or_throw(); @@ -208,7 +208,7 @@ class CandidateSelector final : public StmtExprVisitor { return StmtExprVisitor::Visit_(op); } - ffi::Optional Visit_(const SeqStmtNode* op) final { + ffi::Optional Visit_(const SeqStmtObj* op) final { bool init_no_split = no_split_; for (Stmt stmt : op->seq) { // erase the no split state of before visiting the next one. @@ -221,7 +221,7 @@ class CandidateSelector final : public StmtExprVisitor { return std::nullopt; } - ffi::Optional Visit_(const CallNode* op) final { + ffi::Optional Visit_(const CallObj* op) final { if (op->op.same_as(prim::likely_op())) { in_likely_ = true; TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(op)); @@ -236,7 +236,7 @@ class CandidateSelector final : public StmtExprVisitor { return std::nullopt; } - ffi::Optional Visit_(const VarNode* op) final { + ffi::Optional Visit_(const VarObj* op) final { if (in_likely_ && record_.count(op)) { record_.at(op) = true; } @@ -249,7 +249,7 @@ class CandidateSelector final : public StmtExprVisitor { bool in_likely_{false}; bool no_split_{false}; bool partition_const_loop_{false}; - std::unordered_map record_; + std::unordered_map record_; }; // Finder try best to find partitions for hinted vars @@ -281,9 +281,9 @@ class PartitionFinder : public StmtExprVisitor { } } - ffi::Optional Visit_(const ForNode* op) final { + ffi::Optional Visit_(const ForObj* op) final { if (IsVirtualThread(op)) return StmtExprVisitor::Visit_(op); - auto f_vset_contains = [this](const VarNode* var) { return out_vars_.count(var); }; + auto f_vset_contains = [this](const VarObj* var) { return out_vars_.count(var); }; auto walkfn = [&](const Var& var) -> ffi::Expected { return f_vset_contains(var.get()) ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var)) : ffi::WalkResult::Advance(); @@ -304,7 +304,7 @@ class PartitionFinder : public StmtExprVisitor { return result; } - ffi::Optional Visit_(const RegionStmtNode* op) final { + ffi::Optional Visit_(const RegionStmtObj* op) final { // handle thread_axis if (op->op.same_as(tirx::launch_thread_op()) && !op->args[0].as_or_throw()->value.starts_with("vthread")) { @@ -323,7 +323,7 @@ class PartitionFinder : public StmtExprVisitor { return std::nullopt; } - ffi::Optional Visit_(const CallNode* op) final { + ffi::Optional Visit_(const CallObj* op) final { if (op->op.same_as(prim::likely_op())) { DeduceCondition(op->args[0].as_or_throw()); } else if (op->op.same_as(tirx::ignore_loop_partition_op())) { @@ -334,12 +334,12 @@ class PartitionFinder : public StmtExprVisitor { return std::nullopt; } - DEFINE_PARTITION_FINDER_VISIT_CMP_OP(GENode); - DEFINE_PARTITION_FINDER_VISIT_CMP_OP(GTNode); - DEFINE_PARTITION_FINDER_VISIT_CMP_OP(LENode); - DEFINE_PARTITION_FINDER_VISIT_CMP_OP(LTNode); - DEFINE_PARTITION_FINDER_VISIT_CMP_OP(EQNode); - DEFINE_PARTITION_FINDER_VISIT_CMP_OP(NENode); + DEFINE_PARTITION_FINDER_VISIT_CMP_OP(GEObj); + DEFINE_PARTITION_FINDER_VISIT_CMP_OP(GTObj); + DEFINE_PARTITION_FINDER_VISIT_CMP_OP(LEObj); + DEFINE_PARTITION_FINDER_VISIT_CMP_OP(LTObj); + DEFINE_PARTITION_FINDER_VISIT_CMP_OP(EQObj); + DEFINE_PARTITION_FINDER_VISIT_CMP_OP(NEObj); Partition partitions; @@ -363,7 +363,7 @@ class PartitionFinder : public StmtExprVisitor { if (interval.IsNothing()) { // `DeduceBound` do not support NE now, thus when // deduce l==r failed, just only try (l<=r && l>=r) - if (const EQNode* op = cond.as()) { + if (const EQObj* op = cond.as()) { IntSet part1 = DeduceBound(current_var_.as_or_throw(), GE(op->a, op->b), hint_map_, relax_map_); IntSet part2 = DeduceBound(current_var_.as_or_throw(), LE(op->a, op->b), @@ -391,23 +391,23 @@ class PartitionFinder : public StmtExprVisitor { ffi::Optional InverseCond(const PrimExpr& cond) { ffi::Optional inverse_cond; - if (const LTNode* op = cond.as()) { + if (const LTObj* op = cond.as()) { // a < b -> a >= b inverse_cond = GE(op->a, op->b); - } else if (const GTNode* op = cond.as()) { + } else if (const GTObj* op = cond.as()) { // a > b -> a <= b inverse_cond = LE(op->a, op->b); - } else if (const LENode* op = cond.as()) { + } else if (const LEObj* op = cond.as()) { // a <= b -> a > b inverse_cond = GT(op->a, op->b); - } else if (const GENode* op = cond.as()) { + } else if (const GEObj* op = cond.as()) { // a >= b -> a < b inverse_cond = LT(op->a, op->b); - } else if (const EQNode* op = cond.as()) { + } else if (const EQObj* op = cond.as()) { // a == b -> a != b inverse_cond = NE(op->a, op->b); // a != b -> a == b - } else if (const NENode* op = cond.as()) { + } else if (const NEObj* op = cond.as()) { inverse_cond = EQ(op->a, op->b); } return inverse_cond; @@ -415,7 +415,7 @@ class PartitionFinder : public StmtExprVisitor { Var current_var_; bool has_partition_hint_; - std::unordered_set out_vars_; + std::unordered_set out_vars_; // Borrow the caller's maps; each visit restores its scoped bindings. ffi::Map& hint_map_; ffi::Map& relax_map_; @@ -450,14 +450,14 @@ class ThreadPartitionInserter : public StmtExprMutator { using StmtExprMutator::Mutate; using StmtExprMutator::Mutate_; UnchangedOr Mutate(ffi::AnyView value, InplaceMode inplace_mode) override { - if (value.as()) return ffi::Unchanged(); + if (value.as()) return ffi::Unchanged(); return StmtExprMutator::Mutate(value, inplace_mode); } explicit ThreadPartitionInserter(const ExpressionSet& ps, PrimExpr cond) : ps_(ps), cond_(cond), innermost_thread_scope_(false) {} - UnchangedOr Mutate_(const ForNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const ForObj* op, InplaceMode inplace_mode) final { if (!tvm::s_tir::GetThreadBinding(op).has_value() || IsVirtualThread(op)) { return StmtExprMutator::Mutate_(op, inplace_mode); } @@ -476,7 +476,7 @@ class ThreadPartitionInserter : public StmtExprMutator { return loop; } - UnchangedOr Mutate_(const RegionStmtNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const RegionStmtObj* op, InplaceMode inplace_mode) final { if (op->op.same_as(tirx::launch_thread_op()) && !op->args[0].as_or_throw()->value.starts_with("vthread")) { innermost_thread_scope_ = true; @@ -511,7 +511,7 @@ class LoopPartitioner : public StmtExprMutator { using StmtExprMutator::Mutate; using StmtExprMutator::Mutate_; UnchangedOr Mutate(ffi::AnyView value, InplaceMode inplace_mode) override { - if (value.as()) return ffi::Unchanged(); + if (value.as()) return ffi::Unchanged(); return StmtExprMutator::Mutate(value, inplace_mode); } @@ -526,10 +526,23 @@ class LoopPartitioner : public StmtExprMutator { return Mutate(stmt, InplaceMode::kAllow).ValueOrUnchanged(std::move(stmt)); } - UnchangedOr Mutate_(const ForNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const ForObj* op, InplaceMode inplace_mode) final { if (IsVirtualThread(op)) return StmtExprMutator::Mutate_(op, inplace_mode); bool thread_scope = tvm::s_tir::GetThreadBinding(op).has_value(); - analyzer_->Bind(op->var, Range(op->begin, op->end), true); + { + sym::ConstraintContext context; + analyzer_->GetConstraintContext()->ForEach([&](const sym::Constraint& premise) { + if (auto binding = premise.as(); + binding && binding->var.same_as(op->var)) + return; + if (auto range = premise.as(); + range && range->var.same_as(op->var)) + return; + context->Append(premise); + }); + context->Append(sym::RangeConstraint(op->var, Range(op->begin, op->end))); + analyzer_->SyncConstraintContext(context); + } auto fs = ffi::GetRef(op); if (selector->candidates.count(fs)) { auto s = TryPartition(fs, op->var, op->begin, op->end - 1, op->body, thread_scope); @@ -547,7 +560,7 @@ class LoopPartitioner : public StmtExprMutator { return res; } - UnchangedOr Mutate_(const RegionStmtNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const RegionStmtObj* op, InplaceMode inplace_mode) final { if (!op->op.same_as(tirx::launch_thread_op()) || op->args[0].as_or_throw()->value.starts_with("vthread")) { return StmtExprMutator::Mutate_(op, inplace_mode); @@ -702,7 +715,7 @@ ffi::Optional LoopPartitioner::TryPartition(const Stmt& stmt, Var var, Pri // include hint of var. hint_map_.Set(var, IntSet::Interval(min, max)); - const auto* loop = stmt.as(); + const auto* loop = stmt.as(); bool has_partition_hint_ = loop && HasPartitionHint(loop); auto finder = ffi::make_object(var, hint_map_, relax_map_, has_partition_hint_); finder->Visit(body); @@ -778,7 +791,7 @@ ffi::Optional LoopPartitioner::TryPartition(const Stmt& stmt, Var var, Pri if (!partition_thread_scope && has_partition_hint_ && unroll_loop_with_partition_hint_no_interval_ && analyzer_->CanProve(max - min > 0)) { auto new_body = VisitAndMutate(body); - auto unrolled = ffi::make_object(*loop); + auto unrolled = ffi::make_object(*loop); unrolled->kind = ForKind::kUnrolled; unrolled->body = std::move(new_body); return For(unrolled); @@ -914,7 +927,7 @@ ffi::Optional LoopPartitioner::TryPartition(const Stmt& stmt, Var var, Pri } inline Stmt LoopPartitioner::MakeFor(const ffi::Object* node, PrimExpr extent, Stmt body) { - const ForNode* for_node = static_cast(node); + const ForObj* for_node = static_cast(node); TVM_FFI_ICHECK(for_node); if (analyzer_->CanProve(extent == IntImm::Int32(1)) && !no_unroll_loop_with_extent_one_ && @@ -929,7 +942,7 @@ inline Stmt LoopPartitioner::MakeFor(const ffi::Object* node, PrimExpr extent, S return ffi::StructuralMap(body, f_substitute).as_or_throw(); } else { TVM_FFI_ICHECK(!tvm::s_tir::GetThreadBinding(for_node).has_value()); - auto new_loop = ffi::make_object(*for_node); + auto new_loop = ffi::make_object(*for_node); new_loop->begin = IntImm(for_node->begin.ty(), 0); new_loop->end = extent; new_loop->body = body; @@ -942,7 +955,7 @@ class RemoveLikelyTagsAndHints : public StmtExprMutator { using StmtExprMutator::Mutate; using StmtExprMutator::Mutate_; - UnchangedOr Mutate_(const CallNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const CallObj* op, InplaceMode inplace_mode) final { if (op->op.same_as(prim::likely_op())) { TVM_FFI_ICHECK_EQ(op->args.size(), 1); return StmtExprMutator::Mutate(op->args[0]).ValueOrUnchanged(op->args[0]).as_or_throw(); @@ -954,7 +967,7 @@ class RemoveLikelyTagsAndHints : public StmtExprMutator { } } - UnchangedOr Mutate_(const ForNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const ForObj* op, InplaceMode inplace_mode) final { For loop = StmtExprMutator::Mutate_(op, inplace_mode) .ValueOrUnchanged(ffi::GetRef(op)) .as_or_throw(); diff --git a/src/s_tir/transform/lower_cross_thread_reduction.cc b/src/s_tir/transform/lower_cross_thread_reduction.cc index e78387842925..f287e07ede22 100644 --- a/src/s_tir/transform/lower_cross_thread_reduction.cc +++ b/src/s_tir/transform/lower_cross_thread_reduction.cc @@ -68,7 +68,7 @@ struct GPUThreadTagEqual { * \brief loop The loop to be checked * \return True if the loop is bound to threadIdx.x/y/z */ -bool IsBoundToThreadIdx(const ForNode* loop) { +bool IsBoundToThreadIdx(const ForObj* loop) { if (!tvm::s_tir::GetThreadBinding(loop).has_value()) { return false; } @@ -85,7 +85,7 @@ bool IsBoundToThreadIdx(const ForNode* loop) { */ bool IsDominantBlock(const SBlock& scope_block, const SBlock& block) { // Step 1. Count the number of writers for each tensor written by the scope block. - std::unordered_map tensor_writer_cnt; + std::unordered_map tensor_writer_cnt; auto walk_fn = [&tensor_writer_cnt](const SBlock& block) -> ffi::Expected { for (const TensorRegion& tensor_region : block->writes) { ++tensor_writer_cnt[tensor_region->source.as_or_throw().get()]; @@ -117,7 +117,7 @@ bool IsDominantBlock(const SBlock& scope_block, const SBlock& block) { */ bool IsReductionBlock(const SBlockRealize& realize, const ffi::Map& loop_range_map, const SBlock& scope_block, sym::AnalyzerObj* analyzer) { - const auto* block = realize->block.as(); + const auto* block = realize->block.as(); // Cond 1. The block has the `init` statement. if (!block->init.has_value()) { return false; @@ -158,7 +158,7 @@ ffi::Array MakeScratchpads(const ffi::Array& reduction_ten new_tensors.push_back( TensorVar(name, TensorType(/*shape=*/{IntImm::Int32(1)}, /*dtype=*/tensor.dtype(), scope_local, /*data_alignment=*/0, - tvm::tirx::TileLayoutNode::DefaultLayout({IntImm::Int32(1)})))); + tvm::tirx::TileLayout::DefaultLayout({IntImm::Int32(1)})))); } return new_tensors; } @@ -189,12 +189,12 @@ class TensorReplacer : public StmtExprMutator { } private: - UnchangedOr Mutate_(const TensorLoadNode* load, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const TensorLoadObj* load, InplaceMode inplace_mode) final { auto replacement = VarRemapGet(load->source).as(); return replacement ? TensorLoad(replacement.value(), {0}) : ffi::GetRef(load); } - UnchangedOr Mutate_(const TensorStoreNode* store, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const TensorStoreObj* store, InplaceMode inplace_mode) final { if (auto replacement = VarRemapGet(store->dest.as_or_throw()).as()) { PrimExpr value = StmtExprMutator::Mutate(ffi::AnyView(store->value), inplace_mode) .ValueOrUnchanged(store->value) @@ -215,7 +215,7 @@ class InThreadReducerMaker : public StmtExprMutator { using StmtExprMutator::Mutate; using StmtExprMutator::Mutate_; UnchangedOr Mutate(ffi::AnyView value, InplaceMode inplace_mode) override { - if (value.as()) return ffi::Unchanged(); + if (value.as()) return ffi::Unchanged(); return StmtExprMutator::Mutate(value, inplace_mode); } @@ -226,7 +226,7 @@ class InThreadReducerMaker : public StmtExprMutator { public: using StmtExprVisitor::Visit_; ffi::Optional Visit(ffi::AnyView value) override { - if (value.as()) return std::nullopt; + if (value.as()) return std::nullopt; return StmtExprVisitor::Visit(value); } /*! @@ -241,7 +241,7 @@ class InThreadReducerMaker : public StmtExprMutator { } private: - ffi::Optional Visit_(const SBlockNode* block) final { + ffi::Optional Visit_(const SBlockObj* block) final { ffi::Array iter_vars = block->iter_vars; for (const IterVar& iter_var : block->iter_vars) { if (iter_var->iter_type == kCommReduce) { @@ -255,13 +255,13 @@ class InThreadReducerMaker : public StmtExprMutator { ffi::Array reduction_block_vars_; }; - static ffi::Optional Make(const SBlockRealizeNode* src_realize, + static ffi::Optional Make(const SBlockRealizeObj* src_realize, ffi::Optional tgt_realize, Stmt stmt) { return ffi::make_object(src_realize, std::move(tgt_realize)) ->Rewrite(std::move(stmt)); } - explicit InThreadReducerMaker(const SBlockRealizeNode* src_realize, + explicit InThreadReducerMaker(const SBlockRealizeObj* src_realize, ffi::Optional tgt_realize) : src_realize_(src_realize), tgt_realize_(tgt_realize) {} @@ -277,7 +277,7 @@ class InThreadReducerMaker : public StmtExprMutator { return result; } - UnchangedOr Mutate_(const SBlockRealizeNode* realize, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const SBlockRealizeObj* realize, InplaceMode inplace_mode) final { if (realize == src_realize_) { if (tgt_realize_.has_value()) { return tgt_realize_.value(); @@ -287,7 +287,7 @@ class InThreadReducerMaker : public StmtExprMutator { return ffi::Unchanged(); } - UnchangedOr Mutate_(const ForNode* loop, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const ForObj* loop, InplaceMode inplace_mode) final { auto body = Rewrite(loop->body); if (!body.has_value()) { removed_ = true; @@ -304,7 +304,7 @@ class InThreadReducerMaker : public StmtExprMutator { return res; } - UnchangedOr Mutate_(const SeqStmtNode* seq, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const SeqStmtObj* seq, InplaceMode inplace_mode) final { ffi::Array stmts; stmts.reserve(seq->size()); for (const Stmt& stmt : seq->seq) { @@ -321,7 +321,7 @@ class InThreadReducerMaker : public StmtExprMutator { bool removed_ = false; - const SBlockRealizeNode* src_realize_; + const SBlockRealizeObj* src_realize_; ffi::Optional tgt_realize_; }; @@ -337,18 +337,18 @@ class InThreadReducerMaker : public StmtExprMutator { * \param combiner_rhs The RHS values of the combiner * \param reduction_loops The reduction loops */ -Stmt TransformReductionBlock(const SBlockRealizeNode* realize, // +Stmt TransformReductionBlock(const SBlockRealizeObj* realize, // const ffi::Optional>& it_tensors, // const ffi::Array& ct_tensors, // const ffi::Array& wb_tensors, // const ffi::Array& old_wb_indices, // const te::CommReducer& reducer, // const ffi::Array& combiner_rhs, // - const std::vector& reduction_loops) { + const std::vector& reduction_loops) { static const auto scope_local = runtime::MemoryScope::Get("local"); int n_tensors = wb_tensors.size(); - const SBlockNode* block = realize->block.get(); + const SBlockObj* block = realize->block.get(); auto f_create_tensor_regions = [](ffi::Array tensors) { ffi::Array regions; @@ -395,13 +395,13 @@ Stmt TransformReductionBlock(const SBlockRealizeNode* realize, // then replace `wb_tensors` with `it_tensors` accordingly in given BlockRealize // otherwise, directly remove given BlockRealize if (it_tensors.has_value()) { - ffi::ObjectPtr new_block = ffi::make_object(*block); + ffi::ObjectPtr new_block = ffi::make_object(*block); new_block->writes = it_tensor_regions.value(); new_block->name_hint = new_block->name_hint + "_in_thread"; new_block->body = TensorReplacer::Run(wb_tensors, it_tensors.value(), std::move(new_block->body)); new_block->init = std::nullopt; - ffi::ObjectPtr n = ffi::make_object(*realize); + ffi::ObjectPtr n = ffi::make_object(*realize); n->block = SBlock(new_block); new_realize = SBlockRealize(n); } @@ -431,7 +431,7 @@ Stmt TransformReductionBlock(const SBlockRealizeNode* realize, destinations.push_back(TensorLoad(ct_tensors[i], {0})); } ffi::Array thread_axes; - for (const ForNode* reduction_loop : reduction_loops) { + for (const ForObj* reduction_loop : reduction_loops) { if (tvm::s_tir::GetThreadBinding(reduction_loop).has_value()) { thread_axes.push_back(reduction_loop->var); } @@ -484,7 +484,7 @@ Stmt TransformReductionBlock(const SBlockRealizeNode* realize, if (iter_var->iter_type != kCommReduce) { IterVar new_iter_var{ffi::UnsafeInit{}}; { - ffi::ObjectPtr n = ffi::make_object(*iter_var.get()); + ffi::ObjectPtr n = ffi::make_object(*iter_var.get()); Var v(iter_var->var->name, iter_var->var->ty, iter_var->var->loc); n->var = v.as_or_throw(); new_iter_var = IterVar(n); @@ -528,9 +528,9 @@ Stmt TransformReductionBlock(const SBlockRealizeNode* realize, // - each predicate clause of the original block which contains spatial loop var, and // - `t == 0` for each reduction thread dim when the write-back tensor is not local. PrimExpr wb_predicate = IntImm::Bool(true); - std::unordered_set reduction_loop_vars; + std::unordered_set reduction_loop_vars; reduction_loop_vars.reserve(reduction_loops.size()); - for (const ForNode* reduction_loop : reduction_loops) { + for (const ForObj* reduction_loop : reduction_loops) { reduction_loop_vars.insert(reduction_loop->var.get()); } std::unordered_set visited_predicate_nodes; @@ -541,7 +541,7 @@ Stmt TransformReductionBlock(const SBlockRealizeNode* realize, } ffi::Array sub_exprs = {and_expr->a, and_expr->b}; for (PrimExpr sub_expr : sub_exprs) { - if (sub_expr->IsInstance()) { + if (sub_expr->IsInstance()) { continue; } bool is_reduction = [sub_expr, &reduction_loop_vars]() { @@ -561,7 +561,7 @@ Stmt TransformReductionBlock(const SBlockRealizeNode* realize, }; ffi::StructuralWalk(realize->predicate, walk_fn); if (wb_tensors[0].scope() != scope_local) { - for (const ForNode* loop : reduction_loops) { + for (const ForObj* loop : reduction_loops) { if (tvm::s_tir::GetThreadBinding(loop).has_value()) { wb_predicate = wb_predicate && (static_cast(loop->var) == IntImm(loop->var.ty(), 0)); @@ -582,9 +582,9 @@ Stmt TransformReductionBlock(const SBlockRealizeNode* realize, // Final step: Wrap all the above four statements with the reduction loops bound to threadIdx Stmt new_stmt = SeqStmt(std::move(stmts)); for (auto rit = reduction_loops.rbegin(); rit != reduction_loops.rend(); ++rit) { - const ForNode* loop = *rit; + const ForObj* loop = *rit; if (tvm::s_tir::GetThreadBinding(loop).has_value()) { - ffi::ObjectPtr n = ffi::make_object(*loop); + ffi::ObjectPtr n = ffi::make_object(*loop); n->body = std::move(new_stmt); new_stmt = For(n); } @@ -602,7 +602,7 @@ class CrossThreadReductionTransformer : public StmtExprMutator { private: // Check if the input block needs cross-thread reduction. - std::vector NeedCrossThreadReduction(const SBlockRealizeNode* realize) { + std::vector NeedCrossThreadReduction(const SBlockRealizeObj* realize) { // Step 0. If the block is the root block, just return. if (block_stack_.empty()) { return {}; @@ -615,7 +615,7 @@ class CrossThreadReductionTransformer : public StmtExprMutator { } // Step 2. Collect all the vars that appear in the bindings of reduction block iters. - std::unordered_set reduction_vars; + std::unordered_set reduction_vars; GetVarsTouchedByBlockIters(ffi::GetRef(realize), nullptr, &reduction_vars); // Step 3. Collect the loops whose loop vars appear in the bindings of reduction block iters. @@ -624,8 +624,8 @@ class CrossThreadReductionTransformer : public StmtExprMutator { // so, cross-thread reduction is needed. If none of the reduction-related loops is bound to // thread axis, cross-thread reduction is not needed for the input block. bool need = false; - std::vector reduction_loops; - for (const ForNode* loop : loop_stack_) { + std::vector reduction_loops; + for (const ForObj* loop : loop_stack_) { if (reduction_vars.count(loop->var.get())) { // Step 3. Collect the loop. reduction_loops.push_back(loop); @@ -635,7 +635,7 @@ class CrossThreadReductionTransformer : public StmtExprMutator { } } } - return need ? reduction_loops : std::vector{}; + return need ? reduction_loops : std::vector{}; } // Check if the input block needs thread broadcast rewrite. @@ -646,7 +646,7 @@ class CrossThreadReductionTransformer : public StmtExprMutator { // 3. at least one of the reduction thread vars of the cross-thread reduction // is free to this block (i.e., not bound to the block). std::vector> NeedCrossThreadBroadcast( - const SBlockRealizeNode* realize) { + const SBlockRealizeObj* realize) { static const auto scope_local = runtime::MemoryScope::Get("local"); SBlock block = realize->block; @@ -673,7 +673,7 @@ class CrossThreadReductionTransformer : public StmtExprMutator { } // Erase those threads which are not free to this block. - for (const ForNode* loop : loop_stack_) { + for (const ForObj* loop : loop_stack_) { if (tvm::s_tir::GetThreadBinding(loop).has_value()) { GPUThreadTag scope = ParseThreadBindingTag(tvm::s_tir::GetThreadBinding(loop).value()); thread2range.erase(scope); @@ -701,21 +701,21 @@ class CrossThreadReductionTransformer : public StmtExprMutator { */ std::tuple, ffi::Array, ffi::Array> - CheckCanApplyCrossThreadReduction(const SBlockNode* block, - const std::vector& reduction_loops) const { + CheckCanApplyCrossThreadReduction(const SBlockObj* block, + const std::vector& reduction_loops) const { static const auto scope_local = runtime::MemoryScope::Get("local"); // Condition 1. All the reduction-related loops should be the deepest among all statements // outside the block (ignoring SeqStmt here). int n_deepest_reduction_loops = 0; for (auto rit = statement_stack_.rbegin() + 1; rit != statement_stack_.rend(); ++rit) { - const StmtNode* stmt = *rit; - if ((*rit)->IsInstance()) { + const StmtObj* stmt = *rit; + if ((*rit)->IsInstance()) { // Skip SeqStmt. continue; } if (std::find(reduction_loops.begin(), reduction_loops.end(), - reinterpret_cast(stmt)) == reduction_loops.end()) { + reinterpret_cast(stmt)) == reduction_loops.end()) { break; } ++n_deepest_reduction_loops; @@ -730,7 +730,7 @@ class CrossThreadReductionTransformer : public StmtExprMutator { // Condition 2. All the reduction-related loops that are bound to thread axes should only be // bound to `threadIdx.x/y/z`. int n_bound_reduction_loops = 0; - for (const ForNode* reduction_loop : reduction_loops) { + for (const ForObj* reduction_loop : reduction_loops) { if (tvm::s_tir::GetThreadBinding(reduction_loop).has_value()) { ++n_bound_reduction_loops; TVM_FFI_CHECK(IsBoundToThreadIdx(reduction_loop), ValueError) @@ -794,8 +794,8 @@ class CrossThreadReductionTransformer : public StmtExprMutator { } UnchangedOr Mutate(ffi::AnyView value, InplaceMode inplace_mode) final { - if (value.as()) return ffi::Unchanged(); - const auto* stmt = value.as(); + if (value.as()) return ffi::Unchanged(); + const auto* stmt = value.as(); if (!stmt) return StmtExprMutator::Mutate(value, inplace_mode); statement_stack_.push_back(stmt); auto result = StmtExprMutator::Mutate(value, inplace_mode); @@ -803,7 +803,7 @@ class CrossThreadReductionTransformer : public StmtExprMutator { return result; } - UnchangedOr Mutate_(const ForNode* loop, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const ForObj* loop, InplaceMode inplace_mode) final { loop_stack_.push_back(loop); loop_range_map_.Set(loop->var, Range(loop->begin, loop->end)); @@ -848,7 +848,7 @@ class CrossThreadReductionTransformer : public StmtExprMutator { } } - UnchangedOr Mutate_(const SBlockNode* block, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const SBlockObj* block, InplaceMode inplace_mode) final { ffi::Map old_loop_range_map; block_stack_.push_back(block); std::swap(old_loop_range_map, loop_range_map_); @@ -861,7 +861,7 @@ class CrossThreadReductionTransformer : public StmtExprMutator { // Insert the new allocated tensors into the block's `alloc_tensors` field. auto it = block2new_tensors_.find(block); if (it != block2new_tensors_.end()) { - SBlockNode* p_new_block = new_block.CopyOnWrite(); + SBlockObj* p_new_block = new_block.CopyOnWrite(); for (const TensorVar& new_tensor : it->second) { if (new_tensor.defined()) { p_new_block->alloc_tensors.push_back(new_tensor); @@ -871,9 +871,9 @@ class CrossThreadReductionTransformer : public StmtExprMutator { return new_block; } - void MakeCrossThreadReduction(const SBlockRealizeNode* realize, - const std::vector reduction_loops) { - const SBlockNode* block = realize->block.get(); + void MakeCrossThreadReduction(const SBlockRealizeObj* realize, + const std::vector reduction_loops) { + const SBlockObj* block = realize->block.get(); // Step 1. Check whether cross-thread reduction can be applied. If no, throw an exception on // which condition the block violates. @@ -911,7 +911,7 @@ class CrossThreadReductionTransformer : public StmtExprMutator { // The information is used for consumer block broadcasting detection. std::vector> reduction_threads; reduction_threads.reserve(reduction_loops.size()); - for (const ForNode* loop : reduction_loops) { + for (const ForObj* loop : reduction_loops) { if (tvm::s_tir::GetThreadBinding(loop).has_value()) { reduction_threads.emplace_back( ParseThreadBindingTag(tvm::s_tir::GetThreadBinding(loop).value()), @@ -924,7 +924,7 @@ class CrossThreadReductionTransformer : public StmtExprMutator { } Stmt MakeCrossThreadBroadcast( - const SBlockRealizeNode* realize, + const SBlockRealizeObj* realize, const std::vector>& unbound_thread2range) { // Step 1. Generate loop var for each unbound thread. // Update the block predicate with clauses of `thread_var == min`. @@ -939,7 +939,7 @@ class CrossThreadReductionTransformer : public StmtExprMutator { } // Step 2. Update the BlockRealize with the new predicate. - ffi::ObjectPtr p_realize = ffi::make_object(*realize); + ffi::ObjectPtr p_realize = ffi::make_object(*realize); p_realize->predicate = std::move(predicate); // Step 3. Wrap the updated BlockRealize with the new loops. @@ -954,9 +954,9 @@ class CrossThreadReductionTransformer : public StmtExprMutator { return body; } - UnchangedOr Mutate_(const SBlockRealizeNode* realize, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const SBlockRealizeObj* realize, InplaceMode inplace_mode) final { // Part 1. Check if the block needs cross-thread reduction rewrite. - std::vector reduction_loops = NeedCrossThreadReduction(realize); + std::vector reduction_loops = NeedCrossThreadReduction(realize); if (!reduction_loops.empty()) { // Keep the original statement during traversal. The complete transformed // subtree replaces it when returning to the first reduction-related loop. @@ -981,17 +981,17 @@ class CrossThreadReductionTransformer : public StmtExprMutator { } bool has_cross_thread_reduction_ = false; - std::vector statement_stack_; - std::vector loop_stack_; - std::vector block_stack_; - std::unordered_map> block2new_tensors_; - std::unordered_map loop2new_stmt_; + std::vector statement_stack_; + std::vector loop_stack_; + std::vector block_stack_; + std::unordered_map> block2new_tensors_; + std::unordered_map loop2new_stmt_; ffi::Map loop_range_map_; sym::Analyzer analyzer_; int block_idx_depth = 0; int thread_idx_depth = 0; - std::unordered_map>> + std::unordered_map>> crt_tensor2threads_; }; @@ -999,7 +999,7 @@ namespace transform { Pass LowerCrossThreadReduction() { auto pass_func = [=](Function f, IRModule m, PassContext ctx) { - FunctionNode* fptr = f.CopyOnWrite(); + FunctionObj* fptr = f.CopyOnWrite(); fptr->body = ffi::make_object() ->Mutate(fptr->body) .ValueOrUnchanged(fptr->body); diff --git a/src/s_tir/transform/lower_init_block.cc b/src/s_tir/transform/lower_init_block.cc index d3f9e614be1b..b3d4ee44f2be 100644 --- a/src/s_tir/transform/lower_init_block.cc +++ b/src/s_tir/transform/lower_init_block.cc @@ -36,12 +36,12 @@ class InitBlockLower : public StmtExprMutator { using StmtExprMutator::Mutate; using StmtExprMutator::Mutate_; UnchangedOr Mutate(ffi::AnyView value, InplaceMode inplace_mode) override { - if (value.as()) return ffi::Unchanged(); + if (value.as()) return ffi::Unchanged(); return StmtExprMutator::Mutate(value, inplace_mode); } private: - UnchangedOr Mutate_(const SBlockNode* block, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const SBlockObj* block, InplaceMode inplace_mode) final { if (!block->init.has_value()) { return StmtExprMutator::Mutate_(block, inplace_mode); } @@ -49,12 +49,12 @@ class InitBlockLower : public StmtExprMutator { Stmt body = Mutate(block->body, inplace_mode).ValueOrUnchanged(block->body); body = SeqStmt({init, body}); if (inplace_mode == InplaceMode::kAllow) { - auto* writable = const_cast(block); + auto* writable = const_cast(block); writable->init = std::nullopt; writable->body = std::move(body); return ffi::Unchanged(); } else { - auto copy = ffi::make_object(*block); + auto copy = ffi::make_object(*block); copy->init = std::nullopt; copy->body = std::move(body); return SBlock(std::move(copy)); diff --git a/src/s_tir/transform/lower_match_tensor.cc b/src/s_tir/transform/lower_match_tensor.cc index c6514c0fea67..fcfe3f5c331e 100644 --- a/src/s_tir/transform/lower_match_tensor.cc +++ b/src/s_tir/transform/lower_match_tensor.cc @@ -59,7 +59,7 @@ class MatchTensorLower : public StmtExprMutator { } private: - UnchangedOr Mutate_(const SBlockNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const SBlockObj* op, InplaceMode inplace_mode) final { for (const MatchTensorRegion& match_tensor : op->match_tensors) { CheckAndUpdateVarMap(match_tensor); } @@ -80,7 +80,7 @@ class MatchTensorLower : public StmtExprMutator { } } } - op = stmt.as(); + op = stmt.as(); TVM_FFI_ICHECK(op != nullptr); ffi::Array reads = op->reads.Map(std::bind(&MatchTensorLower::VisitTensorRegion, this, std::placeholders::_1)); @@ -115,15 +115,19 @@ class MatchTensorLower : public StmtExprMutator { } } - UnchangedOr Mutate_(const ForNode* op, InplaceMode inplace_mode) final { - analyzer_->Bind(op->var, Range(op->begin, op->end)); + UnchangedOr Mutate_(const ForObj* op, InplaceMode inplace_mode) final { + { + auto context = analyzer_->GetConstraintContext(); + context->Append(sym::RangeConstraint(op->var, Range(op->begin, op->end))); + analyzer_->SyncConstraintContext(context); + } return StmtExprMutator::Mutate_(op, inplace_mode); } - UnchangedOr Mutate_(const CallNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const CallObj* op, InplaceMode inplace_mode) final { if ((op->op.same_as(tirx::masked_load_op()) || op->op.same_as(tirx::masked_store_op())) && !op->args.empty()) { - if (auto var = op->args[0].as(); var && var.value()->ty.as()) { + if (auto var = op->args[0].as(); var && var.value()->ty.as()) { TensorVar tensor = var.value().as_or_throw(); TVM_FFI_ICHECK(!match_tensors_.count(tensor)) << "Predicated tensor access is not currently supported in lower match tensor pass."; @@ -131,7 +135,7 @@ class MatchTensorLower : public StmtExprMutator { } if (op->op.same_as(tirx::tensor_data_ptr_op()) && op->args.size() == 1) { if (auto var = op->args[0].as(); - var.has_value() && var.value()->ty.as()) { + var.has_value() && var.value()->ty.as()) { auto it = match_tensors_.find(var.value().as_or_throw()); if (it != match_tensors_.end()) { return (*it).second->source.as_or_throw().data(); @@ -141,13 +145,13 @@ class MatchTensorLower : public StmtExprMutator { return StmtExprMutator::Mutate_(op, inplace_mode); } - UnchangedOr Mutate_(const TensorStoreNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const TensorStoreObj* op, InplaceMode inplace_mode) final { // Save the original tensor before base class mutation may remap it TensorVar orig_tensor = op->dest.as_or_throw(); TensorStore stmt = StmtExprMutator::Mutate_(op, inplace_mode) .ValueOrUnchanged(ffi::GetRef(op)) .as_or_throw(); - op = stmt.as(); + op = stmt.as(); TVM_FFI_ICHECK(op != nullptr); // Look up using original tensor (before the inherited Var environment may have remapped it) @@ -165,12 +169,12 @@ class MatchTensorLower : public StmtExprMutator { } } - UnchangedOr Mutate_(const TensorLoadNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const TensorLoadObj* op, InplaceMode inplace_mode) final { // Save the original tensor before base class mutation may remap it TensorVar orig_tensor = op->source.as_or_throw(); PrimExpr expr = StmtExprMutator::Mutate_(op, inplace_mode).ValueOrUnchanged(ffi::GetRef(op)); - op = expr.as(); + op = expr.as(); TVM_FFI_ICHECK(op != nullptr); auto it = match_tensors_.find(orig_tensor); @@ -298,13 +302,17 @@ class MatchTensorLower : public StmtExprMutator { }; value = ffi::StructuralMap(std::move(value), f_substitute) .as_or_throw(); - if (arg->IsInstance()) { + if (arg->IsInstance()) { Var v = arg.as_or_throw(); auto replacement = VarRemapGet(v); if (replacement == nullptr) { VarRemapSet(v, value); if (auto prim_value = value.as()) { - analyzer_->Bind(v, prim_value.value()); + { + auto context = analyzer_->GetConstraintContext(); + context->Append(sym::BindingConstraint(v, prim_value.value())); + analyzer_->SyncConstraintContext(context); + } } } else { AssertBinding(replacement.as_or_throw(), value, arg_name); @@ -317,7 +325,7 @@ class MatchTensorLower : public StmtExprMutator { } if (unknown.size() == 1 && arg_prim) { auto coefficients = tirx::DetectLinearEquation(arg_prim.value(), unknown); - if (coefficients.size() == 2 && coefficients[0].as() && + if (coefficients.size() == 2 && coefficients[0].as() && !IsZero(coefficients[0])) { PrimExpr delta = analyzer_->Simplify(value.as_or_throw() - coefficients[1]); TVM_FFI_ICHECK(analyzer_->CanProve(truncmod(delta, coefficients[0]) == 0)) diff --git a/src/s_tir/transform/lower_opaque_block.cc b/src/s_tir/transform/lower_opaque_block.cc index f859ec89d820..0b9d23de14b5 100644 --- a/src/s_tir/transform/lower_opaque_block.cc +++ b/src/s_tir/transform/lower_opaque_block.cc @@ -60,7 +60,7 @@ class OpaqueBlockLower : public StmtExprMutator { std::optional unroll_explicit; }; - UnchangedOr Mutate_(const SBlockRealizeNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const SBlockRealizeObj* op, InplaceMode inplace_mode) final { // We have convert blocks into opaque blocks in previous passes. TVM_FFI_ICHECK(op->iter_values.empty()) << "Non-opaque blocks are not allowed in FlattenTensor. Please " @@ -108,7 +108,7 @@ class OpaqueBlockLower : public StmtExprMutator { return body; } - UnchangedOr Mutate_(const ForNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const ForObj* op, InplaceMode inplace_mode) final { // Step 1. Update unit loop info. PrimExpr min = this->Mutate(op->begin, inplace_mode).ValueOrUnchanged(op->begin); PrimExpr end = this->Mutate(op->end, inplace_mode).ValueOrUnchanged(op->end); diff --git a/src/s_tir/transform/lower_thread_binding.cc b/src/s_tir/transform/lower_thread_binding.cc index da9bd20a05af..41ce24fdfd06 100644 --- a/src/s_tir/transform/lower_thread_binding.cc +++ b/src/s_tir/transform/lower_thread_binding.cc @@ -36,7 +36,7 @@ class ThreadBindingLowerer : public StmtExprMutator { using StmtExprMutator::Mutate_; private: - UnchangedOr Mutate_(const SBlockNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const SBlockObj* op, InplaceMode inplace_mode) final { auto annotations = Mutate(op->annotations, inplace_mode) .as_or_throw>>(); auto result = StmtExprMutator::Mutate_(op, inplace_mode); @@ -46,7 +46,7 @@ class ThreadBindingLowerer : public StmtExprMutator { return block; } - UnchangedOr Mutate_(const ForNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const ForObj* op, InplaceMode inplace_mode) final { if (!tvm::s_tir::GetThreadBinding(op).has_value()) { auto annotations = Mutate(op->annotations, inplace_mode) .as_or_throw>>(); diff --git a/src/s_tir/transform/lower_vtcm_alloc.cc b/src/s_tir/transform/lower_vtcm_alloc.cc index a490cbd1c176..40f6df927930 100644 --- a/src/s_tir/transform/lower_vtcm_alloc.cc +++ b/src/s_tir/transform/lower_vtcm_alloc.cc @@ -43,15 +43,15 @@ class VtcmAllocator : public StmtExprMutator { VtcmAllocator() {} - UnchangedOr Mutate_(const BindNode* op, InplaceMode inplace_mode) final { - if (const auto* call = op->value.as(); + UnchangedOr Mutate_(const BindObj* op, InplaceMode inplace_mode) final { + if (const auto* call = op->value.as(); call && call->op.same_as(tirx::alloc_tensor_op())) { return Mutate_AllocTensor(op, call, inplace_mode); } return StmtExprMutator::Mutate_(op, inplace_mode); } - UnchangedOr Mutate_AllocTensor(const BindNode* op, const CallNode* call, + UnchangedOr Mutate_AllocTensor(const BindObj* op, const CallObj* call, InplaceMode inplace_mode) { runtime::MemoryScope scope = call->ty_args[0].as_or_throw()->scope; if (IsVtcmStorage(scope)) { @@ -73,7 +73,7 @@ class VtcmAllocator : public StmtExprMutator { protected: runtime::MemoryScope GetStorageScope(const Var& var) { - auto* ptr = var->ty.as(); + auto* ptr = var->ty.as(); TVM_FFI_ICHECK(ptr) << "Buffer Var's type annotation must be of PtrType"; return ptr->scope; } 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 12340f2f930b..636cb9f539a2 100644 --- a/src/s_tir/transform/manifest_shared_memory_local_stage.cc +++ b/src/s_tir/transform/manifest_shared_memory_local_stage.cc @@ -55,12 +55,12 @@ class IntermediateStageRewriter { explicit IntermediateStageRewriter(const std::vector& ancestor_loop_or_blocks) : ancestor_loop_or_blocks_(ancestor_loop_or_blocks) {} - std::tuple Rewrite(const SBlockNode* block) { + std::tuple Rewrite(const SBlockObj* block) { static const auto scope_shared = runtime::MemoryScope::Get("shared"); static const auto scope_shared_dyn = runtime::MemoryScope::Get("shared.dyn"); - const TensorStoreNode* store = - block->body->size() == 1 ? block->body->seq[0].as() : nullptr; + const TensorStoreObj* store = + block->body->size() == 1 ? block->body->seq[0].as() : nullptr; TVM_FFI_CHECK( store != nullptr && (store->dest.as_or_throw().scope() == scope_shared || store->dest.as_or_throw().scope() == scope_shared_dyn), @@ -70,7 +70,7 @@ class IntermediateStageRewriter { const TensorVar& target_tensor = store->dest.as_or_throw(); // Step 0: Collect relaxed loops - std::vector relaxed_loops = CollectRelaxedOuterLoops(block, target_tensor); + std::vector relaxed_loops = CollectRelaxedOuterLoops(block, target_tensor); // Step 1: Create tensor for the local stage auto [new_tensor, tensor_indices] = CreateIntermediateTensor(relaxed_loops, target_tensor); @@ -90,17 +90,17 @@ class IntermediateStageRewriter { private: /*! \brief Collect relaxed outer loops from innermost to outermost */ - std::vector CollectRelaxedOuterLoops(const SBlockNode* block, - const TensorVar& target_tensor) { - std::vector relaxed_loops; + std::vector CollectRelaxedOuterLoops(const SBlockObj* block, + const TensorVar& target_tensor) { + std::vector relaxed_loops; for (int n = static_cast(ancestor_loop_or_blocks_.size()) - 1, i = n - 1; i >= 0; --i) { const Stmt& ancestor = ancestor_loop_or_blocks_[i]; - if (const ForNode* ancestor_loop = ancestor.as()) { + if (const ForObj* ancestor_loop = ancestor.as()) { TVM_FFI_CHECK( ancestor_loop->kind == ForKind::kDefault || ancestor_loop->kind == ForKind::kVectorized, ValueError) << "Expect the ancestor loops to be serial or vectorized, got " << ancestor_loop->kind; - relaxed_loops.push_back(ancestor.as()); + relaxed_loops.push_back(ancestor.as()); if (i < n - 1) { TVM_FFI_CHECK(ancestor_loop->body->size() == 1 && @@ -108,17 +108,17 @@ class IntermediateStageRewriter { ValueError) << "Expect the ancestor loops to have a single child."; } else { - const SBlockRealizeNode* block_realize = - ancestor_loop->body->size() == 1 ? ancestor_loop->body->seq[0].as() + const SBlockRealizeObj* block_realize = + ancestor_loop->body->size() == 1 ? ancestor_loop->body->seq[0].as() : nullptr; TVM_FFI_ICHECK(block_realize != nullptr); TVM_FFI_CHECK(block_realize != nullptr && block_realize->block.get() == block, ValueError) << "Expect the ancestor loops to have a single child."; } } else { - const SBlockRealizeNode* ancestor_block_realize = ancestor.as(); + const SBlockRealizeObj* ancestor_block_realize = ancestor.as(); TVM_FFI_ICHECK(ancestor_block_realize != nullptr); - const SBlockNode* ancestor_block = ancestor_block_realize->block.get(); + const SBlockObj* ancestor_block = ancestor_block_realize->block.get(); auto it = std::find_if( ancestor_block->alloc_tensors.begin(), ancestor_block->alloc_tensors.end(), [&target_tensor](const TensorVar& tensor) { return tensor.same_as(target_tensor); }); @@ -131,9 +131,9 @@ class IntermediateStageRewriter { } /*! \brief Create the intermediate stage. */ - Stmt MakeLocalStage(const SBlockNode* block, const TensorVar& new_tensor, + Stmt MakeLocalStage(const SBlockObj* block, const TensorVar& new_tensor, ffi::Array local_stage_indices, - std::vector relaxed_loops, const TensorStoreNode* store) { + std::vector relaxed_loops, const TensorStoreObj* store) { // Step 0: Create the body of the local stage, which is TensorStore to the intermediate tensor. Stmt local_stage = TensorStore(new_tensor, local_stage_indices, store->value); @@ -145,13 +145,13 @@ class IntermediateStageRewriter { /*body=*/std::move(local_stage)); local_stage = SBlockRealize( /*iter_values=*/{}, - /*predicate=*/ancestor_loop_or_blocks_.back().as()->predicate, + /*predicate=*/ancestor_loop_or_blocks_.back().as()->predicate, local_stage.as_or_throw()); // Step 2: Add outer loops ffi::Map subst_map; - for (const ForNode* relaxed_loop : relaxed_loops) { - ffi::ObjectPtr for_node = ffi::make_object(*relaxed_loop); + for (const ForObj* relaxed_loop : relaxed_loops) { + ffi::ObjectPtr for_node = ffi::make_object(*relaxed_loop); for_node->var = for_node->var.CopyWithSuffix(""); for_node->body = std::move(local_stage); local_stage = For(for_node); @@ -168,7 +168,7 @@ class IntermediateStageRewriter { /*! \brief Create the intermediate tensor with the extents of the relaxed outer loops. */ std::pair> CreateIntermediateTensor( - const std::vector relaxed_loops, const TensorVar& tensor) const { + const std::vector relaxed_loops, const TensorVar& tensor) const { ffi::Array tensor_indices; ffi::Array new_tensor_shape; @@ -176,14 +176,14 @@ class IntermediateStageRewriter { // extents of the relaxed outer loops. for (auto it = relaxed_loops.rbegin(); it != relaxed_loops.rend(); ++it) { - const ForNode* relaxed_loop = *it; + const ForObj* relaxed_loop = *it; tensor_indices.push_back(relaxed_loop->begin + relaxed_loop->var); new_tensor_shape.push_back( Range::GetExtentFromBeginEnd(relaxed_loop->begin, relaxed_loop->end)); } static const auto scope_local = runtime::MemoryScope::Get("local"); TensorVar new_tensor = WithScope(tensor, scope_local); - ffi::ObjectPtr type = ffi::make_object(*new_tensor.ty().get()); + ffi::ObjectPtr type = ffi::make_object(*new_tensor.ty().get()); type->shape = new_tensor_shape; new_tensor = TensorVar(new_tensor->name, TensorType(std::move(type)), new_tensor->loc); return {new_tensor, tensor_indices}; @@ -197,11 +197,11 @@ class SharedMemoryLocalStageInserter : public StmtExprMutator { using StmtExprMutator::Mutate; using StmtExprMutator::Mutate_; UnchangedOr Mutate(ffi::AnyView value, InplaceMode inplace_mode) override { - if (value.as()) return ffi::Unchanged(); + if (value.as()) return ffi::Unchanged(); return StmtExprMutator::Mutate(value, inplace_mode); } - UnchangedOr Mutate_(const ForNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const ForObj* op, InplaceMode inplace_mode) final { ancestor_loop_or_blocks_.push_back(ffi::GetRef(op)); Stmt new_stmt = StmtExprMutator::Mutate_(op, InplaceMode::kDisallow) .ValueOrUnchanged(ffi::GetRef(op)); @@ -209,7 +209,7 @@ class SharedMemoryLocalStageInserter : public StmtExprMutator { return new_stmt; } - UnchangedOr Mutate_(const SBlockRealizeNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const SBlockRealizeObj* op, InplaceMode inplace_mode) final { ancestor_loop_or_blocks_.push_back(ffi::GetRef(op)); Stmt new_stmt = StmtExprMutator::Mutate_(op, InplaceMode::kDisallow) .ValueOrUnchanged(ffi::GetRef(op)); @@ -217,7 +217,7 @@ class SharedMemoryLocalStageInserter : public StmtExprMutator { return new_stmt; } - UnchangedOr Mutate_(const SBlockNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const SBlockObj* op, InplaceMode inplace_mode) final { if (op->annotations.count(s_tir::attr::kManifestSharedMemoryLocalStage.name)) { // Rewrite the shared memory access to load from the intermediate tensor. // The annotated block must be a leaf block (will be checked during rewriting). No need to @@ -271,7 +271,7 @@ class SharedMemoryLocalStageInserter : public StmtExprMutator { } SBlock new_block = ffi::GetRef(op); - SBlockNode* new_block_node = new_block.CopyOnWrite(); + SBlockObj* new_block_node = new_block.CopyOnWrite(); // Add new tensor allocations if any. if (new_alloc_tensors.size() > 0) { new_block_node->alloc_tensors = Concat(new_block_node->alloc_tensors, new_alloc_tensors); diff --git a/src/s_tir/transform/memhammer_coalesce.cc b/src/s_tir/transform/memhammer_coalesce.cc index 0c7133753f28..f6dfc02307de 100644 --- a/src/s_tir/transform/memhammer_coalesce.cc +++ b/src/s_tir/transform/memhammer_coalesce.cc @@ -33,8 +33,8 @@ using namespace tvm::tirx; * \return the fused loop */ Stmt FuseNestLoops(Stmt body) { - std::vector loops; - while (const ForNode* loop = body.as()) { + std::vector loops; + while (const ForObj* loop = body.as()) { loops.push_back(loop); body = loop->body->size() == 1 ? loop->body->seq[0] : loop->body; } @@ -74,7 +74,7 @@ Stmt FuseNestLoops(Stmt body) { * \return The stmt after transformation */ Stmt SplitBindVectorize(const Stmt& stmt, const ConstraintSet& constraints) { - const ForNode* loop = TVM_TYPE_AS(stmt, ForNode); + const ForObj* loop = TVM_TYPE_AS(stmt, ForObj); int loop_extent = Range::GetExtentFromBeginEnd(loop->begin, loop->end) .as_or_throw() ->value.as() @@ -116,7 +116,11 @@ Stmt SplitBindVectorize(const Stmt& stmt, const ConstraintSet& constraints) { for (int i = 0; i < n; i++) { const PrimExpr& factor = factors[i]; PrimVar var = loop->var.CopyWithSuffix("_" + std::to_string(i)); - analyzer->Bind(var, Range::FromMinExtent(0, factor)); + { + auto context = analyzer->GetConstraintContext(); + context->Append(sym::RangeConstraint(var, Range::FromMinExtent(0, factor))); + analyzer->SyncConstraintContext(context); + } new_loop_vars.push_back(var); } // substitute fused loop var with new loop vars @@ -172,10 +176,10 @@ Stmt CoalescedAccess::Rewrite(const Stmt& stmt, const ConstraintSet& constraints */ ffi::Array GetMapping(const Stmt& stmt, const ConstraintSet& constraints) { Stmt body = stmt; - while (const ForNode* loop = body.as()) { + while (const ForObj* loop = body.as()) { body = loop->body->size() == 1 ? loop->body->seq[0] : loop->body; } - const TensorStoreNode* tensor_store = TVM_TYPE_AS(body, TensorStoreNode); + const TensorStoreObj* tensor_store = TVM_TYPE_AS(body, TensorStoreObj); TensorRegion write_region = constraints.write_region; const ffi::Array& write_index = tensor_store->indices; TVM_FFI_ICHECK(write_region->region.size() == write_index.size() && @@ -199,7 +203,7 @@ Stmt InverseMapping::Rewrite(const Stmt& stmt, const ConstraintSet& constraints, ffi::Array loop_vars; // Step 1. Get index mapping ffi::Array mapping_pattern = GetMapping(stmt, constraints); - while (const ForNode* loop = body.as()) { + while (const ForObj* loop = body.as()) { var_range.Set(loop->var, Range(loop->begin, loop->end)); loop_vars.push_back(loop->var); body = loop->body->size() == 1 ? loop->body->seq[0] : loop->body; diff --git a/src/s_tir/transform/memhammer_intermediate_stage.cc b/src/s_tir/transform/memhammer_intermediate_stage.cc index 8c7f0188bcb4..fe5166c7e00f 100644 --- a/src/s_tir/transform/memhammer_intermediate_stage.cc +++ b/src/s_tir/transform/memhammer_intermediate_stage.cc @@ -28,11 +28,11 @@ namespace tvm { namespace s_tir { using namespace tvm::tirx; -Stmt CopyLoopChain(const std::vector loops, const Stmt& inner_body, int ith = -1, +Stmt CopyLoopChain(const std::vector loops, const Stmt& inner_body, int ith = -1, ffi::Optional* ith_loop = nullptr) { Stmt ret = inner_body; for (int i = static_cast(loops.size() - 1); i >= 0; i--) { - ffi::ObjectPtr new_loop = ffi::make_object(*loops[i]); + ffi::ObjectPtr new_loop = ffi::make_object(*loops[i]); new_loop->body = ret; ret = For(new_loop); if (ith == i) { @@ -49,10 +49,10 @@ Stmt CopyLoopChain(const std::vector loops, const Stmt& inner_bo * The second is the lowest thread binding loop. */ std::pair> LiftThreadBindingLoops(Stmt stmt) { - std::vector normal_loops; - std::vector thread_binding_loops; + std::vector normal_loops; + std::vector thread_binding_loops; Stmt body = stmt; - while (const ForNode* loop = body.as()) { + while (const ForObj* loop = body.as()) { if (tvm::s_tir::GetThreadBinding(loop).has_value()) { thread_binding_loops.push_back(loop); } else { @@ -123,13 +123,13 @@ class IndexPatternFinder : public StmtExprVisitor { } private: - ffi::Optional Visit_(const VarNode* op) final { + ffi::Optional Visit_(const VarObj* op) final { if (!success_) { return std::nullopt; } if (ffi::Optional range = var_range_.Get(ffi::GetRef(op))) { PrimExpr index = ffi::GetRef(op).as_or_throw(); - int64_t max = static_cast(range.value()->extent.as()->value); + int64_t max = static_cast(range.value()->extent.as()->value); int64_t extent = max; for (int i = static_cast(operator_stack.size()) - 1; i >= 0; i--) { Operator o = operator_stack[i]; @@ -178,24 +178,24 @@ class IndexPatternFinder : public StmtExprVisitor { return std::nullopt; } - ffi::Optional Visit_(const FloorDivNode* op) final { - int64_t b = static_cast(op->b.as()->value); + ffi::Optional Visit_(const FloorDivObj* op) final { + int64_t b = static_cast(op->b.as()->value); operator_stack.push_back(Operator{Operator::OpKind::FloorDiv, b}); TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(op)); operator_stack.pop_back(); return std::nullopt; } - ffi::Optional Visit_(const FloorModNode* op) final { - int64_t b = static_cast(op->b.as()->value); + ffi::Optional Visit_(const FloorModObj* op) final { + int64_t b = static_cast(op->b.as()->value); operator_stack.push_back(Operator{Operator::OpKind::FloorMod, b}); TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(op)); operator_stack.pop_back(); return std::nullopt; } - ffi::Optional Visit_(const MulNode* op) final { - int64_t b = static_cast(op->b.as()->value); + ffi::Optional Visit_(const MulObj* op) final { + int64_t b = static_cast(op->b.as()->value); operator_stack.push_back(Operator{Operator::OpKind::Mul, b}); TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(op)); operator_stack.pop_back(); @@ -217,7 +217,7 @@ class TensorLoadReplacer : public StmtExprMutator { TensorLoadReplacer(const TensorVar& tgt_tensor, const TensorLoad& new_tensor_load) : tgt_tensor_(tgt_tensor), new_tensor_load_(new_tensor_load) {} - UnchangedOr Mutate_(const TensorLoadNode* op, InplaceMode inplace_mode) { + UnchangedOr Mutate_(const TensorLoadObj* op, InplaceMode inplace_mode) { if (op->source.as_or_throw().same_as(tgt_tensor_)) { return new_tensor_load_; } @@ -248,16 +248,16 @@ std::pair InsertCacheStage(Stmt stmt, bool is_write_cache, ffi::S static const auto scope_mma_matrixC = runtime::MemoryScope::Get("mma.matrixC"); Stmt body = stmt; - std::vector loops; - std::vector loops_under_compute_location; - std::vector relaxed_thread_loops; + std::vector loops; + std::vector loops_under_compute_location; + std::vector relaxed_thread_loops; bool need_relax = !compute_location.has_value(); ffi::Map var_range; PrimExpr vector_bytes = -1; // Step 1. Perform rank promotion on the tensor access, turning a strided-changing dimension into // several contiguous-changing dimensions // Step 1.1 collect loop var range for rank promotion - while (const ForNode* loop = body.as()) { + while (const ForObj* loop = body.as()) { if (need_relax) { var_range.Set(loop->var, Range(loop->begin, loop->end)); loops_under_compute_location.push_back(loop); @@ -273,7 +273,7 @@ std::pair InsertCacheStage(Stmt stmt, bool is_write_cache, ffi::S body = loop->body->size() == 1 ? loop->body->seq[0] : loop->body; } ffi::Optional predicate; - if (const auto* op = body.as()) { + if (const auto* op = body.as()) { // the predicate is generated by coalescing predicate = op->condition; TVM_FFI_ICHECK_EQ(op->then_case->size(), 1); @@ -296,7 +296,7 @@ std::pair InsertCacheStage(Stmt stmt, bool is_write_cache, ffi::S } sym::Analyzer analyzer; - const TensorLoadNode* target_tensor_load = nullptr; + const TensorLoadObj* target_tensor_load = nullptr; if (is_write_cache) { auto walk_fn = [&](const TensorLoad& tensor_load) -> ffi::Expected { if (tensor_load->source.as_or_throw().scope() == @@ -321,13 +321,13 @@ std::pair InsertCacheStage(Stmt stmt, bool is_write_cache, ffi::S TVM_FFI_ICHECK(target_tensor_load); } - const TensorStoreNode* tensor_store = TVM_TYPE_AS(body, TensorStoreNode); + const TensorStoreObj* tensor_store = TVM_TYPE_AS(body, TensorStoreObj); ffi::Array cache_indices; ffi::Array new_shape; bool use_rank_promotion = false; - if (!is_write_cache && tensor_store->value.as()) { + if (!is_write_cache && tensor_store->value.as()) { ffi::Array indices = - is_write_cache ? tensor_store->indices : tensor_store->value.as()->indices; + is_write_cache ? tensor_store->indices : tensor_store->value.as()->indices; new_shape = IndexPatternFinder::getRankPromotedShape(indices, var_range, &cache_indices); // write cache disabled for now // rank promotion for write cache cannot guarantee the shape fits wmma.accumulator @@ -339,16 +339,16 @@ std::pair InsertCacheStage(Stmt stmt, bool is_write_cache, ffi::S ffi::Map subst_map; if (!use_rank_promotion) { cache_indices.clear(); - for (const ForNode* loop : relaxed_thread_loops) { + for (const ForObj* loop : relaxed_thread_loops) { new_shape.push_back(Range::GetExtentFromBeginEnd(loop->begin, loop->end)); } - for (const ForNode* loop : loops_under_compute_location) { + for (const ForObj* loop : loops_under_compute_location) { new_shape.push_back(Range::GetExtentFromBeginEnd(loop->begin, loop->end)); } } for (int i = 0; i < static_cast(relaxed_thread_loops.size()); i++) { - const ForNode* loop = relaxed_thread_loops[i]; + const ForObj* loop = relaxed_thread_loops[i]; PrimVar new_loop_var = loop->var.CopyWithSuffix("_cache"); new_loop_vars.push_back(new_loop_var); subst_map.Set(loop->var, new_loop_var); @@ -357,7 +357,7 @@ std::pair InsertCacheStage(Stmt stmt, bool is_write_cache, ffi::S } } for (int i = 0; i < static_cast(loops_under_compute_location.size()); i++) { - const ForNode* loop = loops_under_compute_location[i]; + const ForObj* loop = loops_under_compute_location[i]; PrimVar new_loop_var = loop->var.CopyWithSuffix("_cache"); new_loop_vars.push_back(new_loop_var); subst_map.Set(loop->var, new_loop_var); @@ -392,8 +392,8 @@ std::pair InsertCacheStage(Stmt stmt, bool is_write_cache, ffi::S new_tensor = WithScope(tensor_store->dest.as_or_throw(), runtime::MemoryScope::Get(storage_scope)); } - ffi::ObjectPtr tensor_type = - ffi::make_object(*new_tensor.ty().get()); + ffi::ObjectPtr tensor_type = + ffi::make_object(*new_tensor.ty().get()); tensor_type->shape = new_shape; new_tensor = TensorVar(new_tensor->name, TensorType(std::move(tensor_type)), new_tensor->loc); *alloc_tensor = new_tensor; @@ -427,15 +427,15 @@ std::pair InsertCacheStage(Stmt stmt, bool is_write_cache, ffi::S } for (int i = static_cast(loops_under_compute_location.size()) - 1; i >= 0; i--) { - const ForNode* orig_loop = loops_under_compute_location[i]; - ffi::ObjectPtr new_loop = ffi::make_object(*orig_loop); + const ForObj* orig_loop = loops_under_compute_location[i]; + ffi::ObjectPtr new_loop = ffi::make_object(*orig_loop); new_loop->var = new_loop_vars[i + relaxed_thread_loops.size()]; new_loop->body = generate_body; generate_body = For(new_loop); } for (int i = static_cast(relaxed_thread_loops.size()) - 1; i >= 0; i--) { - const ForNode* orig_loop = relaxed_thread_loops[i]; - ffi::ObjectPtr new_loop = ffi::make_object(*orig_loop); + const ForObj* orig_loop = relaxed_thread_loops[i]; + ffi::ObjectPtr new_loop = ffi::make_object(*orig_loop); new_loop->var = new_loop_vars[i]; new_loop->body = generate_body; new_loop->kind = ForKind::kDefault; @@ -455,8 +455,8 @@ std::pair InsertCacheStage(Stmt stmt, bool is_write_cache, ffi::S rewrite_body = If(predicate.value(), rewrite_body); } for (int i = static_cast(loops_under_compute_location.size()) - 1; i >= 0; i--) { - const ForNode* orig_loop = loops_under_compute_location[i]; - ffi::ObjectPtr new_loop = ffi::make_object(*orig_loop); + const ForObj* orig_loop = loops_under_compute_location[i]; + ffi::ObjectPtr new_loop = ffi::make_object(*orig_loop); new_loop->body = rewrite_body; rewrite_body = For(new_loop); } diff --git a/src/s_tir/transform/memhammer_lower_auto_copy.cc b/src/s_tir/transform/memhammer_lower_auto_copy.cc index 82f5a5a302d2..af496a183819 100644 --- a/src/s_tir/transform/memhammer_lower_auto_copy.cc +++ b/src/s_tir/transform/memhammer_lower_auto_copy.cc @@ -127,12 +127,12 @@ class AutoPadder { for (int k = n - 2; k >= 0; k--) { // dims int max_pad_size = static_cast(std::min( max_pad_factor_ * - static_cast((stride * tensor.shape()[k + 1]).as()->value), + static_cast((stride * tensor.shape()[k + 1]).as()->value), static_cast(32 * 32 / data_bits))); int min_conflict = INT32_MAX; int min_conflict_pad = -1; for (int pad = 0; pad <= max_pad_size; pad += pad_min) { // select padding - int padded_stride = (((stride * tensor.shape()[k + 1]).as()->value + pad) % + int padded_stride = (((stride * tensor.shape()[k + 1]).as()->value + pad) % (32 * 32 / data_bits)) .as() .value(); @@ -160,7 +160,7 @@ class AutoPadder { auto iter_space = iter_spaces[i][k]; if (!iter_space.empty()) { int padded_stride = - (((stride * tensor.shape()[k + 1]).as()->value + min_conflict_pad) % + (((stride * tensor.shape()[k + 1]).as()->value + min_conflict_pad) % (32 * 32 / data_bits)) .as() .value(); @@ -177,7 +177,7 @@ class AutoPadder { reverse_strides.push_back(stride); } // Step 3. create the new padded tensor - ffi::ObjectPtr b = ffi::make_object(*tensor.ty().get()); + ffi::ObjectPtr b = ffi::make_object(*tensor.ty().get()); ffi::Array strides; for (int i = static_cast(reverse_strides.size()) - 1; i >= 0; i--) { strides.push_back(reverse_strides[i]); @@ -218,7 +218,7 @@ class AutoPadder { } private: - UnchangedOr Mutate_(const TensorLoadNode* _op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const TensorLoadObj* _op, InplaceMode inplace_mode) final { TensorLoad load = StmtExprMutator::Mutate_(_op, inplace_mode) .ValueOrUnchanged(ffi::GetRef(_op)) .as_or_throw(); @@ -229,18 +229,18 @@ class AutoPadder { return load; } - UnchangedOr Mutate_(const TensorStoreNode* _op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const TensorStoreObj* _op, InplaceMode inplace_mode) final { TensorStore store = StmtExprMutator::Mutate_(_op, inplace_mode) .ValueOrUnchanged(ffi::GetRef(_op)) .as_or_throw(); - TensorStoreNode* op = store.CopyOnWrite(); + TensorStoreObj* op = store.CopyOnWrite(); if (auto replacement = VarRemapGet(op->dest.as_or_throw()).as()) { op->dest = replacement.value(); } return store; } - UnchangedOr Mutate_(const SBlockNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const SBlockObj* op, InplaceMode inplace_mode) final { // To reduce the number of blocks in block sref reuse map, we check whether the block is // really mutated (i.e., the old tensor appears in the block). If so, we return the block // after mutation. Otherwise we just return the original block. @@ -292,7 +292,7 @@ class AutoPadder { if (changed) { SBlock block = std::move(res).as_or_throw(); - SBlockNode* n = block.CopyOnWrite(); + SBlockObj* n = block.CopyOnWrite(); n->reads = std::move(reads); n->writes = std::move(writes); n->match_tensors = std::move(match_tensors); @@ -322,12 +322,12 @@ class AutoPadder { using StmtExprVisitor::Visit_; private: - ffi::Optional Visit_(const VarNode* op) final { + ffi::Optional Visit_(const VarObj* op) final { if (!success_) { return std::nullopt; } int extent = - var_range_[ffi::GetRef(op)]->extent.as()->value.as().value(); + var_range_[ffi::GetRef(op)]->extent.as()->value.as().value(); if (extent > 1) { stack_.push({{extent, 1}}); } else { @@ -336,7 +336,7 @@ class AutoPadder { return std::nullopt; } - ffi::Optional Visit_(const AddNode* op) final { + ffi::Optional Visit_(const AddObj* op) final { TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(op)); if (!success_) { return std::nullopt; @@ -370,14 +370,14 @@ class AutoPadder { return std::nullopt; } - ffi::Optional Visit_(const FloorDivNode* op) final { + ffi::Optional Visit_(const FloorDivObj* op) final { TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(op)); if (!success_) { return std::nullopt; } std::vector inner = stack_.top(); stack_.pop(); - int lower_factor = op->b.as()->value.as().value(); + int lower_factor = op->b.as()->value.as().value(); std::vector ret; for (const Pattern& pattern : inner) { if (pattern.scale >= lower_factor) { @@ -398,14 +398,14 @@ class AutoPadder { return std::nullopt; } - ffi::Optional Visit_(const FloorModNode* op) final { + ffi::Optional Visit_(const FloorModObj* op) final { TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(op)); if (!success_) { return std::nullopt; } std::vector inner = stack_.top(); stack_.pop(); - int extent = op->b.as()->value.as().value(); + int extent = op->b.as()->value.as().value(); std::vector ret; for (const Pattern& pattern : inner) { if (pattern.scale < extent) { @@ -424,14 +424,14 @@ class AutoPadder { return std::nullopt; } - ffi::Optional Visit_(const MulNode* op) final { + ffi::Optional Visit_(const MulObj* op) final { TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(op)); if (!success_) { return std::nullopt; } std::vector inner = stack_.top(); stack_.pop(); - int scale = op->b.as()->value.as().value(); + int scale = op->b.as()->value.as().value(); std::vector ret; for (const Pattern& pattern : inner) { ret.push_back({pattern.extent, pattern.scale * scale}); @@ -527,7 +527,7 @@ class AutoPadder { return !analyzer->CanProve(delta != 1); } - ffi::Optional Visit_(const ForNode* op) final { + ffi::Optional Visit_(const ForObj* op) final { if (!tvm::s_tir::GetThreadBinding(op).has_value()) { substitute_map_.Set(op->var, op->begin); } else { @@ -538,7 +538,7 @@ class AutoPadder { if (op->kind == ForKind::kVectorized) { vector_var = op->var; vector_length_ = Range::GetExtentFromBeginEnd(op->begin, op->end) - .as() + .as() ->value.as() .value(); } @@ -558,7 +558,7 @@ class AutoPadder { * The iteration space would be {{0, 1}, {0, 4, ..., 60}}. * \param op the tensor store */ - ffi::Optional Visit_(const TensorStoreNode* op) final { + ffi::Optional Visit_(const TensorStoreObj* op) final { static const auto scope_shared = runtime::MemoryScope::Get("shared"); static const auto scope_shared_dyn = runtime::MemoryScope::Get("shared.dyn"); @@ -596,7 +596,7 @@ class AutoPadder { * The iteration space would be {{0, 1}, {0, 4, ..., 60}}. * \param op the tensor load */ - ffi::Optional Visit_(const TensorLoadNode* op) final { + ffi::Optional Visit_(const TensorLoadObj* op) final { static const auto scope_shared = runtime::MemoryScope::Get("shared"); static const auto scope_shared_dyn = runtime::MemoryScope::Get("shared.dyn"); @@ -635,13 +635,12 @@ class AutoPadder { * threadIdx. The iteration space would be {{0, 1, ..., 15}, {0, 1, ..., 15}}. * \param op the call node */ - ffi::Optional Visit_(const SBlockNode* op) final { + ffi::Optional Visit_(const SBlockObj* op) final { static const auto scope_shared = runtime::MemoryScope::Get("shared"); static const auto scope_shared_dyn = runtime::MemoryScope::Get("shared.dyn"); - if (const auto* eval = - op->body->size() == 1 ? op->body->seq[0].as() : nullptr) { - if (const auto* call = eval->value.as()) { + if (const auto* eval = op->body->size() == 1 ? op->body->seq[0].as() : nullptr) { + if (const auto* call = eval->value.as()) { static const Op gpu_load_matrix_sync_op = Op::Get("tirx.gpu_load_matrix_sync"); static const Op gpu_store_matrix_sync_op = Op::Get("tirx.gpu_store_matrix_sync"); if (call->op.same_as(gpu_load_matrix_sync_op) || @@ -729,7 +728,7 @@ class AutoPadder { /*! \brief A map from the old tensors to the new padded tensors */ ffi::Map padded_tensor_map_; /*! \brief A map from each tensor to the iteration spaces of the accesses*/ - std::unordered_map>>> iter_spaces_; + std::unordered_map>>> iter_spaces_; /*! \brief A map from each tensor to their minimal padding size */ ffi::Map padding_min_; /*! \brief max padding size in relative to the original shape*/ @@ -753,7 +752,7 @@ class AutoCopyMutator : public StmtExprMutator { Stmt RewritePaddingBody(const Stmt& stmt) { return padder.RewriteTensorAccess(stmt); } private: - UnchangedOr Mutate_(const SBlockNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const SBlockObj* op, InplaceMode inplace_mode) final { static const auto scope_wmma_accumulator = runtime::MemoryScope::Get("wmma.accumulator"); SBlock block = StmtExprMutator::Mutate_(op, inplace_mode) @@ -761,7 +760,7 @@ class AutoCopyMutator : public StmtExprMutator { .as_or_throw(); // only rewrite the block annotated with "auto_copy" if (!GetAnn(op, tvm::s_tir::attr::kAutoCopy.name).value_or(false)) { - SBlockNode* n = block.CopyOnWrite(); + SBlockObj* n = block.CopyOnWrite(); n->alloc_tensors = padder.PadSharedMemory(std::move(n->alloc_tensors)); return block; } @@ -788,7 +787,7 @@ class AutoCopyMutator : public StmtExprMutator { block->writes[0], // data_bits, // block->annotations); - SBlockNode* n = block.CopyOnWrite(); + SBlockObj* n = block.CopyOnWrite(); OutputSet outputs; TVM_FFI_ICHECK_EQ(n->body->size(), 1); Stmt rewritten = n->body->seq[0]; @@ -808,7 +807,7 @@ class AutoCopyMutator : public StmtExprMutator { return block; } - UnchangedOr Mutate_(const ForNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const ForObj* op, InplaceMode inplace_mode) final { outer_loops_.push_back(ffi::GetRef(op)); Stmt stmt = StmtExprMutator::Mutate_(op, InplaceMode::kDisallow) .ValueOrUnchanged(ffi::GetRef(op)); @@ -840,7 +839,7 @@ class ThreadExtentCollector : public StmtExprVisitor { public: using StmtExprVisitor::Visit_; ffi::Optional Visit(ffi::AnyView value) override { - if (value.as()) return std::nullopt; + if (value.as()) return std::nullopt; return StmtExprVisitor::Visit(value); } static ffi::Map CollectThreadExtent(const Stmt& stmt) { @@ -850,7 +849,7 @@ class ThreadExtentCollector : public StmtExprVisitor { } private: - ffi::Optional Visit_(const SBlockNode* op) final { + ffi::Optional Visit_(const SBlockObj* op) final { if (ffi::Optional warp_execution = GetAnn(op, tvm::s_tir::attr::kWarpExecution.name)) { if (warp_execution.value() != 0) { @@ -859,10 +858,10 @@ class ThreadExtentCollector : public StmtExprVisitor { } return StmtExprVisitor::Visit_(op); } - ffi::Optional Visit_(const ForNode* op) final { + ffi::Optional Visit_(const ForObj* op) final { if (tvm::s_tir::GetThreadBinding(op).has_value()) { PrimExpr extent_expr = Range::GetExtentFromBeginEnd(op->begin, op->end); - if (const auto* extent = extent_expr.as()) { + if (const auto* extent = extent_expr.as()) { thread_extent_.Set(tvm::s_tir::GetThreadBinding(op).value(), static_cast(extent->value)); } diff --git a/src/s_tir/transform/memhammer_tensorcore_rewrite.cc b/src/s_tir/transform/memhammer_tensorcore_rewrite.cc index 93020e386e8b..0eeab1a8d70c 100644 --- a/src/s_tir/transform/memhammer_tensorcore_rewrite.cc +++ b/src/s_tir/transform/memhammer_tensorcore_rewrite.cc @@ -42,8 +42,8 @@ using namespace tvm::tirx; */ std::pair> TileWmmaBlock(Stmt stmt) { Stmt body = stmt; - std::vector loops; - while (const ForNode* loop = body.as()) { + std::vector loops; + while (const ForObj* loop = body.as()) { loops.push_back(loop); body = loop->body->size() == 1 ? loop->body->seq[0] : loop->body; } @@ -135,8 +135,8 @@ Stmt RewriteWmmaLoad(Stmt stmt) { const PrimType int32_ty = PrimType::Int(32); Stmt body = stmt; - std::vector loops; - while (const ForNode* loop = body.as()) { + std::vector loops; + while (const ForObj* loop = body.as()) { loops.push_back(loop); body = loop->body->size() == 1 ? loop->body->seq[0] : loop->body; } @@ -151,8 +151,8 @@ Stmt RewriteWmmaLoad(Stmt stmt) { Range::GetExtentFromBeginEnd(loops[n - 2]->begin, loops[n - 2]->end))}, }; // TODO(tian): the assumption that the RHS of TensorStore is TensorLoad may not be accurate - const TensorStoreNode* tensor_store = TVM_TYPE_AS(body, TensorStoreNode); - const TensorLoadNode* tensor_load = TVM_TYPE_AS(tensor_store->value, TensorLoadNode); + const TensorStoreObj* tensor_store = TVM_TYPE_AS(body, TensorStoreObj); + const TensorLoadObj* tensor_load = TVM_TYPE_AS(tensor_store->value, TensorLoadObj); TensorVar src_tensor = tensor_load->source.as_or_throw(); TensorVar tgt_tensor = tensor_store->dest.as_or_throw(); @@ -244,8 +244,8 @@ Stmt RewriteWmmaStore(Stmt stmt) { const PrimType int32_ty = PrimType::Int(32); Stmt body = stmt; - std::vector loops; - while (const ForNode* loop = body.as()) { + std::vector loops; + while (const ForObj* loop = body.as()) { loops.push_back(loop); body = loop->body->size() == 1 ? loop->body->seq[0] : loop->body; } @@ -260,8 +260,8 @@ Stmt RewriteWmmaStore(Stmt stmt) { Range::GetExtentFromBeginEnd(loops[n - 2]->begin, loops[n - 2]->end))}, }; // TODO(tian): the assumption that the RHS of TensorStore is TensorLoad may not be accurate - const TensorStoreNode* tensor_store = TVM_TYPE_AS(body, TensorStoreNode); - const TensorLoadNode* tensor_load = nullptr; + const TensorStoreObj* tensor_store = TVM_TYPE_AS(body, TensorStoreObj); + const TensorLoadObj* tensor_load = nullptr; auto walk_fn = [&](const TensorLoad& load) -> ffi::Expected { if (load->source.as_or_throw().scope() == scope_wmma_accumulator) { TVM_FFI_ICHECK(tensor_load == nullptr || @@ -364,11 +364,11 @@ class WmmaToGlobalRewriter : public StmtExprMutator { using StmtExprMutator::Mutate; using StmtExprMutator::Mutate_; - WmmaToGlobalRewriter(const SeqStmtNode* tgt_stmt, const ConstraintSet& constraints) + WmmaToGlobalRewriter(const SeqStmtObj* tgt_stmt, const ConstraintSet& constraints) : tgt_stmt_(tgt_stmt), constraints_(constraints) {} private: - UnchangedOr Mutate_(const SeqStmtNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const SeqStmtObj* op, InplaceMode inplace_mode) final { if (op == tgt_stmt_) { TVM_FFI_ICHECK_EQ(op->seq.size(), 2); Stmt wmma_to_shared = RewriteWmmaStore(op->seq[0]); @@ -379,7 +379,7 @@ class WmmaToGlobalRewriter : public StmtExprMutator { } } - const SeqStmtNode* tgt_stmt_; + const SeqStmtObj* tgt_stmt_; const ConstraintSet& constraints_; }; @@ -404,8 +404,8 @@ std::pair> TileMmaToGlobalBlock(Stmt stmt) { // j_0, j_1 = sch.split(j, factors=[None, 8]) // sch.reorder(i_0, j_0, i_1, j_1) Stmt body = stmt; - std::vector loops; - while (const ForNode* loop = body.as()) { + std::vector loops; + while (const ForObj* loop = body.as()) { loops.push_back(loop); body = loop->body->size() == 1 ? loop->body->seq[0] : loop->body; } @@ -480,8 +480,8 @@ Stmt RewriteMmaStore(Stmt stmt) { // Step 1. Get inner loop body Stmt body = stmt; - std::vector loops; - while (const ForNode* loop = body.as()) { + std::vector loops; + while (const ForObj* loop = body.as()) { loops.push_back(loop); body = loop->body->size() == 1 ? loop->body->seq[0] : loop->body; } @@ -497,8 +497,8 @@ Stmt RewriteMmaStore(Stmt stmt) { }; // Step 2. Find matrixC tensor - const TensorStoreNode* tensor_store = TVM_TYPE_AS(body, TensorStoreNode); - const TensorLoadNode* tensor_load = nullptr; + const TensorStoreObj* tensor_store = TVM_TYPE_AS(body, TensorStoreObj); + const TensorLoadObj* tensor_load = nullptr; auto walk_fn = [&](const TensorLoad& load) -> ffi::Expected { if (load->source.as_or_throw().scope() == scope_mma_matrixC) { TVM_FFI_ICHECK(tensor_load == nullptr || @@ -588,11 +588,11 @@ class MmaToGlobalRewriter : public StmtExprMutator { using StmtExprMutator::Mutate; using StmtExprMutator::Mutate_; - MmaToGlobalRewriter(const SeqStmtNode* tgt_stmt, const ConstraintSet& constraints) + MmaToGlobalRewriter(const SeqStmtObj* tgt_stmt, const ConstraintSet& constraints) : tgt_stmt_(tgt_stmt), constraints_(constraints) {} private: - UnchangedOr Mutate_(const SeqStmtNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const SeqStmtObj* op, InplaceMode inplace_mode) final { if (op == tgt_stmt_) { TVM_FFI_ICHECK_EQ(op->seq.size(), 2); // Rewrite for local to shared.dyn @@ -606,7 +606,7 @@ class MmaToGlobalRewriter : public StmtExprMutator { } } - const SeqStmtNode* tgt_stmt_; + const SeqStmtObj* tgt_stmt_; const ConstraintSet& constraints_; }; diff --git a/src/s_tir/transform/merge_shared_memory_allocations.cc b/src/s_tir/transform/merge_shared_memory_allocations.cc index 6101d57d71e4..2fcdb3b72983 100644 --- a/src/s_tir/transform/merge_shared_memory_allocations.cc +++ b/src/s_tir/transform/merge_shared_memory_allocations.cc @@ -78,7 +78,7 @@ ffi::Optional GetBufferDataVar(const ffi::Any& data) { if (auto var = data.as()) { return var; } - if (const auto* call = data.as(); + if (const auto* call = data.as(); call && call->op.same_as(tirx::tensor_data_ptr_op()) && call->args.size() == 1) { return call->args[0].as(); } @@ -91,8 +91,8 @@ bool IsDynamicSharedMemory(Var buffer_var) { static const auto scope_shared_dyn = runtime::MemoryScope::Get("shared.dyn"); MemoryScope storage_scope = - (buffer_var->ty.as() ? buffer_var->ty.as()->scope - : buffer_var->ty.as_or_throw()->scope); + (buffer_var->ty.as() ? buffer_var->ty.as()->scope + : buffer_var->ty.as_or_throw()->scope); return storage_scope == scope_shared_dyn; } @@ -100,8 +100,8 @@ bool IsStaticSharedMemory(Var buffer_var) { static const auto scope_shared = runtime::MemoryScope::Get("shared"); MemoryScope storage_scope = - (buffer_var->ty.as() ? buffer_var->ty.as()->scope - : buffer_var->ty.as_or_throw()->scope); + (buffer_var->ty.as() ? buffer_var->ty.as()->scope + : buffer_var->ty.as_or_throw()->scope); return storage_scope == scope_shared; } @@ -126,7 +126,7 @@ bool IsStaticSharedMemory(const TensorVar& tensor) { static int64_t ConstantAllocationSize(const ffi::Array& extents) { int64_t result = 1; for (size_t i = 0; i < extents.size(); ++i) { - if (const IntImmNode* int_size = extents[i].as()) { + if (const IntImmObj* int_size = extents[i].as()) { auto product = (result * int_size->value).as(); if (!product.has_value()) return 0; result = *product; @@ -145,15 +145,15 @@ class AllocateCollector : public StmtExprVisitor { using StmtExprVisitor::Visit_; explicit AllocateCollector(bool is_dynamic) : is_dynamic_(is_dynamic) {} - ffi::Optional Visit_(const BindNode* op) final { - if (const auto* call = op->value.as(); + ffi::Optional Visit_(const BindObj* op) final { + if (const auto* call = op->value.as(); call && call->op.same_as(tirx::alloc_tensor_op())) { return DispatchAllocTensor(op, call); } return StmtExprVisitor::Visit_(op); } - ffi::Optional DispatchAllocTensor(const BindNode* op, const CallNode* call) { + ffi::Optional DispatchAllocTensor(const BindObj* op, const CallObj* call) { static const auto scope_shared_dyn = runtime::MemoryScope::Get("shared.dyn"); static const auto scope_shared = runtime::MemoryScope::Get("shared"); @@ -169,14 +169,14 @@ class AllocateCollector : public StmtExprVisitor { return StmtExprVisitor::Visit_(op); } - ffi::Optional Visit_(const RegionStmtNode* op) final { + ffi::Optional Visit_(const RegionStmtObj* op) final { if (!op->op.same_as(tirx::launch_thread_op())) has_opaque_region_ = true; return StmtExprVisitor::Visit_(op); } bool has_opaque_region_{false}; // The mapping from the original buffer var to its TensorVar - std::unordered_map shmem_allocs_; + std::unordered_map shmem_allocs_; private: bool is_dynamic_; @@ -209,7 +209,7 @@ class SharedMemLinearAccessPatternFinder final : public StmtExprVisitor { // if offset < 0, means this is the end, the begin entry is current_index + offset int64_t scope_pair_offset{0}; // The buffer variables this statement touched. - std::vector touched; + std::vector touched; }; // The scope of each allocation struct AllocEntry { @@ -219,14 +219,14 @@ class SharedMemLinearAccessPatternFinder final : public StmtExprVisitor { TensorVar tensor; }; - ffi::Optional DispatchAllocTensor(const BindNode* op) { + ffi::Optional DispatchAllocTensor(const BindObj* op) { size_t level = scope_.size(); - const VarNode* buf = op->var.as_or_throw().get(); + const VarObj* buf = op->var.as_or_throw().get(); alloc_info_.insert_or_assign(buf, AllocEntry{level, op->var.as_or_throw()}); return StmtExprVisitor::Visit_(op); } - ffi::Optional Visit_(const TensorStoreNode* op) final { + ffi::Optional Visit_(const TensorStoreObj* op) final { scope_.push_back(StmtEntry()); // visit subexpr TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Visit(op->value)); @@ -234,7 +234,7 @@ class SharedMemLinearAccessPatternFinder final : public StmtExprVisitor { TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Visit(index)); } // Add write access. - const VarNode* buf = ResolveAlias(op->dest.as_or_throw().get()); + const VarObj* buf = ResolveAlias(op->dest.as_or_throw().get()); auto it = alloc_info_.find(buf); if (it != alloc_info_.end() && it->second.tensor.defined()) { TVM_FFI_ICHECK_LT(it->second.level, scope_.size()); @@ -251,7 +251,7 @@ class SharedMemLinearAccessPatternFinder final : public StmtExprVisitor { return std::nullopt; } - ffi::Optional Visit_(const EvaluateNode* op) final { + ffi::Optional Visit_(const EvaluateObj* op) final { scope_.push_back(StmtEntry()); // visit subexpr TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(op)); @@ -264,21 +264,21 @@ class SharedMemLinearAccessPatternFinder final : public StmtExprVisitor { return std::nullopt; } - ffi::Optional Visit_(const BindNode* op) final { - if (const auto* call = op->value.as(); + ffi::Optional Visit_(const BindObj* op) final { + if (const auto* call = op->value.as(); call && call->op.same_as(tirx::alloc_tensor_op())) { return DispatchAllocTensor(op); } - if (const auto* call = op->value.as(); + if (const auto* call = op->value.as(); call && call->op.same_as(tirx::decl_tensor_op())) { return DispatchDeclTensor(op, call); } return StmtExprVisitor::Visit_(op); } - ffi::Optional DispatchDeclTensor(const BindNode* op, const CallNode* call) { + ffi::Optional DispatchDeclTensor(const BindObj* op, const CallObj* call) { if (auto source = GetBufferDataVar(call->args[0])) { - const VarNode* allocation = ResolveAlias(source.value().get()); + const VarObj* allocation = ResolveAlias(source.value().get()); if (alloc_info_.count(allocation)) { buffer_alias_sources_[op->var.as_or_throw().get()] = allocation; return std::nullopt; @@ -287,12 +287,12 @@ class SharedMemLinearAccessPatternFinder final : public StmtExprVisitor { return StmtExprVisitor::Visit_(op); } - ffi::Optional Visit_(const TensorLoadNode* op) final { + ffi::Optional Visit_(const TensorLoadObj* op) final { // Add read access. for (const auto& index : op->indices) { TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Visit(index)); } - const VarNode* buf = ResolveAlias(op->source.as_or_throw().get()); + const VarObj* buf = ResolveAlias(op->source.as_or_throw().get()); auto it = alloc_info_.find(buf); if (it != alloc_info_.end() && it->second.tensor.defined()) { TVM_FFI_ICHECK_LT(it->second.level, scope_.size()) @@ -304,9 +304,9 @@ class SharedMemLinearAccessPatternFinder final : public StmtExprVisitor { return std::nullopt; } - ffi::Optional Visit_(const CallNode* op) final { + ffi::Optional Visit_(const CallObj* op) final { if (op->op.same_as(tirx::address_of_op())) { - if (const auto* load = op->args[0].as()) { + if (const auto* load = op->args[0].as()) { TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->Visit(load->source)); for (const auto& index : load->indices) { TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->Visit(index)); @@ -320,7 +320,7 @@ class SharedMemLinearAccessPatternFinder final : public StmtExprVisitor { return std::nullopt; } - ffi::Optional Visit_(const VarNode* buf) final { + ffi::Optional Visit_(const VarObj* buf) final { if (def_region_kind() != kTVMFFIDefRegionKindNone) return StmtExprVisitor::Visit_(buf); // Directly reference to the variable count as a read. buf = ResolveAlias(buf); @@ -356,7 +356,7 @@ class SharedMemLinearAccessPatternFinder final : public StmtExprVisitor { return std::nullopt; } - ffi::Optional Visit_(const RegionStmtNode* op) final { + ffi::Optional Visit_(const RegionStmtObj* op) final { // Only record the outer most thread extent. if (op->op.same_as(tirx::launch_thread_op()) && std::string(op->args[0].as_or_throw()->value).rfind("vthread", 0) != 0 && @@ -373,23 +373,23 @@ class SharedMemLinearAccessPatternFinder final : public StmtExprVisitor { return std::nullopt; } - ffi::Optional Visit_(const IfNode* op) final { return VisitNewScope(op); } + ffi::Optional Visit_(const IfObj* op) final { return VisitNewScope(op); } - ffi::Optional Visit_(const ForNode* op) final { return VisitNewScope(op); } + ffi::Optional Visit_(const ForObj* op) final { return VisitNewScope(op); } - ffi::Optional Visit_(const WhileNode* op) final { return VisitNewScope(op); } + ffi::Optional Visit_(const WhileObj* op) final { return VisitNewScope(op); } - ffi::Optional Visit_(const AssertStmtNode* op) final { return VisitNewScope(op); } + ffi::Optional Visit_(const AssertStmtObj* op) final { return VisitNewScope(op); } // linearized access sequence. std::vector linear_seq_; // The storage scope of each buffer - std::unordered_map alloc_info_; + std::unordered_map alloc_info_; // Typed views of an allocated shared-memory buffer. - std::unordered_map buffer_alias_sources_; + std::unordered_map buffer_alias_sources_; private: - const VarNode* ResolveAlias(const VarNode* buffer) const { + const VarObj* ResolveAlias(const VarObj* buffer) const { while (true) { auto it = buffer_alias_sources_.find(buffer); if (it == buffer_alias_sources_.end()) break; @@ -436,15 +436,15 @@ class SharedMemoryRewriter : public StmtExprMutator { // | A: 4096 bytes | B: 4096 bytes | // | C: 8192 bytes | // Then the allocs = {{A, B}, {C}} - std::vector> allocs; + std::vector> allocs; }; // Event entry in liveness analysis struct EventEntry { // variables we generate - std::vector gen; + std::vector gen; // variables we kill - std::vector kill; + std::vector kill; }; /*! @@ -456,13 +456,13 @@ class SharedMemoryRewriter : public StmtExprMutator { // Total byte size of THIS kernel's merged buffer. PrimExpr merged_alloc_size{0}; // Allocations from THIS kernel's subtree. - std::unordered_map shmem_allocs; + std::unordered_map shmem_allocs; // Per-buffer byte offset into merged_buf_var. - std::unordered_map buffer_byte_offsets; + std::unordered_map buffer_byte_offsets; // TensorVar-object remap: original TensorVar -> merged-data-var TensorVar. - std::unordered_map tensor_remap; + std::unordered_map tensor_remap; // Typed views whose physical source is one of shmem_allocs. - std::unordered_map buffer_alias_sources; + std::unordered_map buffer_alias_sources; // Remapped buffers in first-use order, for deterministic alias emission. std::vector tensor_remap_order; // Has any original alloc in this scope been marked volatile? @@ -471,7 +471,7 @@ class SharedMemoryRewriter : public StmtExprMutator { std::unordered_map event_map; std::multimap const_free_map; std::list sym_free_list; - std::unordered_map alloc_map; + std::unordered_map alloc_map; }; /*! @@ -486,7 +486,7 @@ class SharedMemoryRewriter : public StmtExprMutator { is_dynamic_ ? scope_shared_dyn : scope_shared)); } - UnchangedOr Mutate_(const RegionStmtNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const RegionStmtObj* op, InplaceMode inplace_mode) final { if (op->op.same_as(tirx::launch_thread_op()) && std::string(op->args[0].as_or_throw()->value).rfind("vthread", 0) != 0 && !in_thread_env_) { @@ -561,7 +561,7 @@ class SharedMemoryRewriter : public StmtExprMutator { return StmtExprMutator::Mutate_(op, inplace_mode); } - UnchangedOr Mutate_AllocTensor(const BindNode* op, const CallNode* call, + UnchangedOr Mutate_AllocTensor(const BindObj* op, const CallObj* call, InplaceMode inplace_mode) { static const auto scope_shared_dyn = runtime::MemoryScope::Get("shared.dyn"); static const auto scope_shared = runtime::MemoryScope::Get("shared"); @@ -572,7 +572,7 @@ class SharedMemoryRewriter : public StmtExprMutator { if (!scope_stack_.empty()) { KernelScope& scope = scope_stack_.back(); if (scope.shmem_allocs.count(op->var.as_or_throw().get())) { - if (call->attrs.as()->dict.count(tvm::tirx::attr::kVolatile.name)) { + if (call->attrs.as()->dict.count(tvm::tirx::attr::kVolatile.name)) { scope.has_volatile_alloc = true; } return Evaluate(0); @@ -584,24 +584,24 @@ class SharedMemoryRewriter : public StmtExprMutator { return StmtExprMutator::Mutate_(op, inplace_mode); } - UnchangedOr Mutate_(const BindNode* op, InplaceMode inplace_mode) final { - if (const auto* call = op->value.as(); + UnchangedOr Mutate_(const BindObj* op, InplaceMode inplace_mode) final { + if (const auto* call = op->value.as(); call && call->op.same_as(tirx::alloc_tensor_op())) { return Mutate_AllocTensor(op, call, inplace_mode); } - if (const auto* call = op->value.as(); + if (const auto* call = op->value.as(); call && call->op.same_as(tirx::decl_tensor_op())) { return Mutate_DeclTensor(op, call, inplace_mode); } return StmtExprMutator::Mutate_(op, inplace_mode); } - UnchangedOr Mutate_DeclTensor(const BindNode* op, const CallNode* call, + UnchangedOr Mutate_DeclTensor(const BindObj* op, const CallObj* call, InplaceMode inplace_mode) { if (!scope_stack_.empty()) { if (auto source = GetBufferDataVar(call->args[0])) { KernelScope& scope = scope_stack_.back(); - if (const VarNode* allocation = ResolveAllocation(source.value().get(), scope)) { + if (const VarObj* allocation = ResolveAllocation(source.value().get(), scope)) { scope.buffer_alias_sources[op->var.as_or_throw().get()] = allocation; return Evaluate(0); } @@ -613,7 +613,7 @@ class SharedMemoryRewriter : public StmtExprMutator { .as_or_throw(); if (auto new_buf = GetUpdatedTensor(node->var.as_or_throw()); !new_buf.same_as(node->var)) { - const auto* new_call = node->value.as(); + const auto* new_call = node->value.as(); return Bind(new_buf, Call(new_buf.ty(), tirx::decl_tensor_op(), {new_call->args[0]}, new_call->attrs, {new_buf.ty()}, new_call->loc), @@ -622,14 +622,14 @@ class SharedMemoryRewriter : public StmtExprMutator { return node; } - UnchangedOr Mutate_(const TensorLoadNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const TensorLoadObj* op, InplaceMode inplace_mode) final { auto node = StmtExprMutator::Mutate_(op, inplace_mode) .ValueOrUnchanged(ffi::GetRef(op)) .as_or_throw(); return VisitTensorAccess(std::move(node)); } - UnchangedOr Mutate_(const TensorStoreNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const TensorStoreObj* op, InplaceMode inplace_mode) final { auto node = StmtExprMutator::Mutate_(op, inplace_mode) .ValueOrUnchanged(ffi::GetRef(op)) .as_or_throw(); @@ -687,7 +687,7 @@ class SharedMemoryRewriter : public StmtExprMutator { << "and is to be run after " << "FlattenTensor"; tensor = - TensorVar(tensor->name, TensorType(ffi::make_object(*tensor.ty().get())), + TensorVar(tensor->name, TensorType(ffi::make_object(*tensor.ty().get())), tensor->loc); } @@ -696,7 +696,7 @@ class SharedMemoryRewriter : public StmtExprMutator { return tensor; } - UnchangedOr Mutate_(const CallNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const CallObj* op, InplaceMode inplace_mode) final { static const Op ptx_cp_async_op = Op::Get("tirx.s_tir.cp_async_raw"); if (op->op.same_as(tirx::tensor_data_ptr_op())) { TensorVar tensor = op->args[0].as_or_throw(); @@ -718,14 +718,14 @@ class SharedMemoryRewriter : public StmtExprMutator { Var buffer = buffer_opt.value(); DLDataType dtype; bool is_shared; - if (buffer->ty.as()) { + if (buffer->ty.as()) { TensorVar typed_tensor = buffer.as_or_throw(); dtype = typed_tensor.dtype()->dtype; is_shared = IsAppropriateSharedMemory(typed_tensor); } else { - const auto* ptr_type = buffer->ty.as(); + const auto* ptr_type = buffer->ty.as(); TVM_FFI_ICHECK(ptr_type) << "The buffer should be a pointer type."; - const auto* prim_type = ptr_type->element_type.as(); + const auto* prim_type = ptr_type->element_type.as(); TVM_FFI_ICHECK(prim_type) << "The buffer should be a pointer to a primitive type."; dtype = prim_type->dtype; is_shared = IsAppropriateSharedMemory(buffer); @@ -736,7 +736,7 @@ class SharedMemoryRewriter : public StmtExprMutator { } PrimExpr extra_offset = GetBufferOffset(buffer, dtype); PrimExpr offset = Mutate(op->args[1]).ValueOrUnchanged(op->args[1]).as_or_throw(); - if (buffer->ty.as()) { + if (buffer->ty.as()) { Expr merged_data = GetUpdatedTensor(buffer.as_or_throw()).data(); ffi::Array args = op->args; args.Set(0, merged_data); @@ -768,7 +768,7 @@ class SharedMemoryRewriter : public StmtExprMutator { PrimExpr GetBufferOffset(Var buffer_var, DLDataType dtype) { TVM_FFI_ICHECK(!scope_stack_.empty()); KernelScope& scope = scope_stack_.back(); - const VarNode* allocation = ResolveAllocation(buffer_var.get(), scope); + const VarObj* allocation = ResolveAllocation(buffer_var.get(), scope); TVM_FFI_ICHECK(allocation); auto it = scope.buffer_byte_offsets.find(allocation); TVM_FFI_ICHECK(it != scope.buffer_byte_offsets.end()); @@ -776,7 +776,7 @@ class SharedMemoryRewriter : public StmtExprMutator { return indexdiv(it->second, elem_bytes); } - const VarNode* ResolveAllocation(const VarNode* buffer, const KernelScope& scope) const { + const VarObj* ResolveAllocation(const VarObj* buffer, const KernelScope& scope) const { while (true) { auto it = scope.buffer_alias_sources.find(buffer); if (it == scope.buffer_alias_sources.end()) break; @@ -801,10 +801,10 @@ class SharedMemoryRewriter : public StmtExprMutator { */ void LivenessAnalysis(const std::vector& seq, KernelScope& scope) { // find kill point, do a reverse linear scan. - std::unordered_set touched; + std::unordered_set touched; for (size_t i = seq.size(); i != 0; --i) { const StmtEntry& s = seq[i - 1]; - for (const VarNode* buffer : s.touched) { + for (const VarObj* buffer : s.touched) { if (!touched.count(buffer)) { touched.insert(buffer); scope.event_map[s.stmt].kill.push_back(buffer); @@ -817,7 +817,7 @@ class SharedMemoryRewriter : public StmtExprMutator { int64_t offset = seq[i].scope_pair_offset; if (offset < 0) continue; const StmtEntry& s = seq[i + offset]; - for (const VarNode* buffer : s.touched) { + for (const VarObj* buffer : s.touched) { if (!touched.count(buffer)) { touched.insert(buffer); scope.event_map[s.stmt].gen.push_back(buffer); @@ -832,7 +832,7 @@ class SharedMemoryRewriter : public StmtExprMutator { * \param scope the kernel scope to write results into */ void PlanMemory(const std::vector& seq, KernelScope& scope) { - std::unordered_set inplace_flag; + std::unordered_set inplace_flag; for (size_t i = 0; i < seq.size(); ++i) { auto it = scope.event_map.find(seq[i].stmt); @@ -840,12 +840,12 @@ class SharedMemoryRewriter : public StmtExprMutator { // - leaf stmt(offset = 0) // - end of scope(offset < 0) // In both cases, we need to handle the kill event correctly - auto is_leaf_alloc = [&](const VarNode* var) { + auto is_leaf_alloc = [&](const VarObj* var) { return seq[i].scope_pair_offset == 0 && std::find(it->second.gen.begin(), it->second.gen.end(), var) != it->second.gen.end(); }; if (it != scope.event_map.end() && seq[i].scope_pair_offset <= 0) { - for (const VarNode* var : it->second.kill) { + for (const VarObj* var : it->second.kill) { if (!is_leaf_alloc(var)) this->Free(var, scope); } } @@ -854,7 +854,7 @@ class SharedMemoryRewriter : public StmtExprMutator { // - beginning of scope(offset < 0) // In both cases, we need to handle the gen event correctly if (it != scope.event_map.end() && seq[i].scope_pair_offset >= 0) { - for (const VarNode* var : it->second.gen) { + for (const VarObj* var : it->second.gen) { TVM_FFI_ICHECK(scope.shmem_allocs.count(var)); const TensorVar& buf = scope.shmem_allocs.at(var); StorageEntry* dst_entry = FindAlloc(buf, scope); @@ -862,7 +862,7 @@ class SharedMemoryRewriter : public StmtExprMutator { } } if (it != scope.event_map.end() && seq[i].scope_pair_offset <= 0) { - for (const VarNode* var : it->second.kill) { + for (const VarObj* var : it->second.kill) { if (is_leaf_alloc(var)) this->Free(var, scope); } } @@ -889,7 +889,7 @@ class SharedMemoryRewriter : public StmtExprMutator { std::vector align(max_layer_num, 0); for (const StorageEntry* e : all_entry) { for (int i = 0; i < static_cast(e->allocs.size()); i++) { - for (const VarNode* buffer : e->allocs[i]) { + for (const VarObj* buffer : e->allocs[i]) { const TensorVar& buf = scope.shmem_allocs.at(buffer); int elem_bytes = static_cast(buf.dtype().StorageBytes()); align[i] = std::max(align[i], elem_bytes); @@ -901,7 +901,7 @@ class SharedMemoryRewriter : public StmtExprMutator { PrimExpr max_inner_offset = 0; for (int i = 0; i < static_cast(e->allocs.size()); i++) { PrimExpr inner_offset = 0; - for (const VarNode* buffer : e->allocs[i]) { + for (const VarObj* buffer : e->allocs[i]) { const TensorVar& buf = scope.shmem_allocs.at(buffer); ffi::Array alloc_shape = GetTensorAllocationShape(buf); int elem_bytes = static_cast(buf.dtype().StorageBytes()); @@ -980,7 +980,7 @@ class SharedMemoryRewriter : public StmtExprMutator { // and finally free all these entry in the free list std::vector::iterator> delete_it; // the alloc list for the new entry - std::vector> reuse_allocs; + std::vector> reuse_allocs; uint64_t mem_ct = 0; for (auto it = mid; it != begin;) { --it; @@ -991,7 +991,7 @@ class SharedMemoryRewriter : public StmtExprMutator { reuse_allocs.resize(n, {}); } for (int i = 0; i < n; i++) { - for (const VarNode* alloc : it->second->allocs[i]) { + for (const VarObj* alloc : it->second->allocs[i]) { reuse_allocs[i].push_back(alloc); } } @@ -1026,7 +1026,7 @@ class SharedMemoryRewriter : public StmtExprMutator { * \param var the buffer var * \param scope the kernel scope whose free lists to update */ - void Free(const VarNode* var, KernelScope& scope) { + void Free(const VarObj* var, KernelScope& scope) { auto it = scope.alloc_map.find(var); TVM_FFI_ICHECK(it != scope.alloc_map.end()); StorageEntry* e = it->second; diff --git a/src/s_tir/transform/plan_update_tensor_allocation_location.cc b/src/s_tir/transform/plan_update_tensor_allocation_location.cc index 58bf03f69ff1..e5167a6b167c 100644 --- a/src/s_tir/transform/plan_update_tensor_allocation_location.cc +++ b/src/s_tir/transform/plan_update_tensor_allocation_location.cc @@ -37,7 +37,7 @@ using namespace tvm::tirx; class CollectManagedAllocations : public StmtExprVisitor { public: using StmtExprVisitor::Visit_; - ffi::Optional Visit_(const SBlockNode* op) final { + ffi::Optional Visit_(const SBlockObj* op) final { for (const auto& tensor : op->alloc_tensors) { managed_allocations.insert(tensor.get()); } @@ -49,7 +49,7 @@ class CollectManagedAllocations : public StmtExprVisitor { /*! \brief Tensors that are allocated outside of the BlockNode, and should not be moved by * TensorAllocationLocator. */ - std::unordered_set managed_allocations; + std::unordered_set managed_allocations; }; /*! \brief Collect the allocate tensor order. */ @@ -73,7 +73,7 @@ class TensorAllocateOrderCollector : public StmtExprVisitor { tensor_alloc_recorder_.end(); } - ffi::Optional Visit_(const SBlockNode* op) final { + ffi::Optional Visit_(const SBlockObj* op) final { for (const TensorVar& tensor : op->alloc_tensors) { tensor_alloc_recorder_.push_back(tensor); } @@ -89,14 +89,14 @@ class TensorAllocateOrderCollector : public StmtExprVisitor { return StmtExprVisitor::Visit_(op); } - ffi::Optional Visit_(const TensorLoadNode* op) final { + ffi::Optional Visit_(const TensorLoadObj* op) final { if (!find(op->source.as_or_throw())) { tensor_alloc_recorder_.push_back(op->source.as_or_throw()); } return StmtExprVisitor::Visit_(op); } - ffi::Optional Visit_(const TensorStoreNode* op) final { + ffi::Optional Visit_(const TensorStoreObj* op) final { if (!find(op->dest.as_or_throw())) { tensor_alloc_recorder_.push_back(op->dest.as_or_throw()); } @@ -117,7 +117,7 @@ class TensorAllocationLocator : public StmtExprMutator { // The tensor_alloc_recorder Array is used to keep the tensor allocation order // since the tensor_lca Map is unordered. ffi::Array tensor_alloc_recorder = TensorAllocateOrderCollector::Collect(func); - std::unordered_set arg_tensor_vars; + std::unordered_set arg_tensor_vars; auto collector = ffi::make_object(); collector->Visit(func->body); managed_allocations_ = collector->managed_allocations; @@ -132,7 +132,7 @@ class TensorAllocationLocator : public StmtExprMutator { for (const auto& tensor : tensor_alloc_recorder) { auto it = tensor_lca.find(tensor); if (it != tensor_lca.end()) { - const StmtNode* stmt = (*it).second.has_value() ? (*it).second.value().get() : nullptr; + const StmtObj* stmt = (*it).second.has_value() ? (*it).second.value().get() : nullptr; if (arg_tensor_vars.count(tensor.get())) { continue; } @@ -145,7 +145,7 @@ class TensorAllocationLocator : public StmtExprMutator { } private: - UnchangedOr Mutate_(const ForNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const ForObj* op, InplaceMode inplace_mode) final { auto it = alloc_tensors_.find(op); if (it == alloc_tensors_.end()) { return StmtExprMutator::Mutate_(op, inplace_mode); @@ -172,7 +172,7 @@ class TensorAllocationLocator : public StmtExprMutator { return node; } - UnchangedOr Mutate_(const RegionStmtNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const RegionStmtObj* op, InplaceMode inplace_mode) final { auto it = alloc_tensors_.find(op); if (it == alloc_tensors_.end()) { return StmtExprMutator::Mutate_(op, inplace_mode); @@ -199,7 +199,7 @@ class TensorAllocationLocator : public StmtExprMutator { return node; } - UnchangedOr Mutate_(const SBlockNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const SBlockObj* op, InplaceMode inplace_mode) final { TVM_FFI_ICHECK(!op->init.has_value()); ffi::Array alloc_tensors; auto it = alloc_tensors_.find(op); @@ -218,7 +218,7 @@ class TensorAllocationLocator : public StmtExprMutator { SBlock stmt = StmtExprMutator::Mutate_(op, inplace_mode) .ValueOrUnchanged(ffi::GetRef(op)) .as_or_throw(); - op = stmt.as(); + op = stmt.as(); TVM_FFI_ICHECK(op != nullptr); // No longer consider tensors created by match_tensor inside the block when updating access @@ -234,7 +234,7 @@ class TensorAllocationLocator : public StmtExprMutator { } } - SBlockNode* n = stmt.CopyOnWrite(); + SBlockObj* n = stmt.CopyOnWrite(); n->alloc_tensors = std::move(alloc_tensors); // Erase tensor allocated inside the block from access region. n->reads = RemoveRedundantTensorRegion(n->reads); @@ -251,7 +251,7 @@ class TensorAllocationLocator : public StmtExprMutator { /*body=*/std::move(body), /*init=*/std::nullopt, /*alloc_tensors=*/alloc_tensors); - SBlockNode* n = opaque_block.CopyOnWrite(); + SBlockObj* n = opaque_block.CopyOnWrite(); ffi::Array> access = GetSBlockReadWriteRegion(opaque_block, tensor_data_to_tensor_); n->reads = access[0]; @@ -272,11 +272,11 @@ class TensorAllocationLocator : public StmtExprMutator { } /*! \brief The map from stmt to the tensors to be allocated under it. */ - std::unordered_map> alloc_tensors_; + std::unordered_map> alloc_tensors_; /*! \brief The tensor already allocated during recursive visiting. */ ffi::Map tensor_data_to_tensor_; /*! \brief Tensors that are allocated within a BlockNode, and may be moved. */ - std::unordered_set managed_allocations_; + std::unordered_set managed_allocations_; }; namespace transform { diff --git a/src/s_tir/transform/remove_store_undef.cc b/src/s_tir/transform/remove_store_undef.cc index 7f2872f6801c..88068e067f58 100644 --- a/src/s_tir/transform/remove_store_undef.cc +++ b/src/s_tir/transform/remove_store_undef.cc @@ -36,8 +36,8 @@ namespace s_tir { using namespace tvm::tirx; struct UndefInfo { - std::unordered_set undef_stores; - std::unordered_set undef_bind_vars; + std::unordered_set undef_stores; + std::unordered_set undef_bind_vars; }; class StoreUndefLocator : public StmtExprVisitor { @@ -52,7 +52,7 @@ class StoreUndefLocator : public StmtExprVisitor { StoreUndefLocator() = default; private: - ffi::Optional Visit_(const TensorStoreNode* op) final { + ffi::Optional Visit_(const TensorStoreObj* op) final { // Check the value for undef. bool stash_undef = false; std::swap(has_undef_, stash_undef); @@ -80,7 +80,7 @@ class StoreUndefLocator : public StmtExprVisitor { return std::nullopt; } - ffi::Optional Visit_(const TensorLoadNode* op) final { + ffi::Optional Visit_(const TensorLoadObj* op) final { // Check indices for undef. Undef in tensor indices is always an error. bool idx_undef = false; std::swap(has_undef_, idx_undef); @@ -92,7 +92,7 @@ class StoreUndefLocator : public StmtExprVisitor { return std::nullopt; } - ffi::Optional Visit_(const BindNode* op) final { + ffi::Optional Visit_(const BindObj* op) final { bool stash_undef = false; std::swap(has_undef_, stash_undef); TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit(op->value)); @@ -107,14 +107,14 @@ class StoreUndefLocator : public StmtExprVisitor { return std::nullopt; } - ffi::Optional Visit_(const VarNode* op) final { + ffi::Optional Visit_(const VarObj* op) final { if (var_bindings_with_undef_.count(op)) { has_undef_ = true; } return std::nullopt; } - ffi::Optional Visit_(const CallNode* op) final { + ffi::Optional Visit_(const CallObj* op) final { if (op->op.same_as(tirx::undef_op())) { has_undef_ = true; } @@ -123,8 +123,8 @@ class StoreUndefLocator : public StmtExprVisitor { bool has_undef_{false}; - std::unordered_set var_bindings_with_undef_; - std::unordered_set undef_stores_; + std::unordered_set var_bindings_with_undef_; + std::unordered_set undef_stores_; }; // Remove TensorStores whose value depends on T.undef, and also @@ -149,7 +149,7 @@ class StoreUndefRemover : public StmtExprMutator { : stores_to_remove_(info.undef_stores), bind_vars_to_remove_(info.undef_bind_vars) {} private: - UnchangedOr Mutate_(const TensorStoreNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const TensorStoreObj* op, InplaceMode inplace_mode) final { if (stores_to_remove_.count(op)) { return Evaluate(0); } else { @@ -157,7 +157,7 @@ class StoreUndefRemover : public StmtExprMutator { } } - UnchangedOr Mutate_(const BindNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const BindObj* op, InplaceMode inplace_mode) final { if (bind_vars_to_remove_.count(op->var.get())) { return Evaluate(0); } else { @@ -165,8 +165,8 @@ class StoreUndefRemover : public StmtExprMutator { } } - const std::unordered_set& stores_to_remove_; - const std::unordered_set& bind_vars_to_remove_; + const std::unordered_set& stores_to_remove_; + const std::unordered_set& bind_vars_to_remove_; }; // Check that no tirx::undef_op() remains in the IR. @@ -180,7 +180,7 @@ class ContainsUndefChecker : public StmtExprVisitor { } private: - ffi::Optional Visit_(const CallNode* op) final { + ffi::Optional Visit_(const CallObj* op) final { if (op->op.same_as(tirx::undef_op())) { contains_undef = true; } diff --git a/src/s_tir/transform/remove_weight_layout_rewrite_block.cc b/src/s_tir/transform/remove_weight_layout_rewrite_block.cc index aa3f703b6783..edf040bb809e 100644 --- a/src/s_tir/transform/remove_weight_layout_rewrite_block.cc +++ b/src/s_tir/transform/remove_weight_layout_rewrite_block.cc @@ -41,24 +41,24 @@ class RemoveLayoutRewriteBlock : public StmtExprMutator { using StmtExprMutator::Mutate; using StmtExprMutator::Mutate_; UnchangedOr Mutate(ffi::AnyView value, InplaceMode inplace_mode) override { - if (value.as()) return ffi::Unchanged(); + if (value.as()) return ffi::Unchanged(); return StmtExprMutator::Mutate(value, inplace_mode); } static std::tuple, - std::unordered_map, - std::unordered_map>> + std::unordered_map, + std::unordered_map>> Rewrite(Function f) { auto rewriter = ffi::make_object(); - FunctionNode* n = f.CopyOnWrite(); + FunctionObj* n = f.CopyOnWrite(); n->body = rewriter->Mutate(n->body, InplaceMode::kAllow).ValueOrUnchanged(std::move(n->body)); return std::make_tuple(f, rewriter->tensor_map_, rewriter->tensor_var_to_index_map_, rewriter->tensor_var_to_rewritten_shape_); } private: - UnchangedOr Mutate_(const SBlockNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const SBlockObj* op, InplaceMode inplace_mode) final { SBlock block = StmtExprMutator::Mutate_(op, inplace_mode) .ValueOrUnchanged(ffi::GetRef(op)) .as_or_throw(); @@ -74,7 +74,7 @@ class RemoveLayoutRewriteBlock : public StmtExprMutator { } } if (alloc_tensors.size() < block->alloc_tensors.size()) { - SBlockNode* n = block.CopyOnWrite(); + SBlockObj* n = block.CopyOnWrite(); n->alloc_tensors = std::move(alloc_tensors); return block; } else { @@ -88,11 +88,11 @@ class RemoveLayoutRewriteBlock : public StmtExprMutator { // Step 1. Checking the body is a TensorStore const auto* store = - block->body->size() == 1 ? block->body->seq[0].as() : nullptr; + block->body->size() == 1 ? block->body->seq[0].as() : nullptr; TVM_FFI_ICHECK(store); // Step 2. Checking the rhs of tensor store is a TensorLoad - const auto* load = store->value.as(); + const auto* load = store->value.as(); TVM_FFI_ICHECK(load); // Step 3. Update TensorVar @@ -102,7 +102,7 @@ class RemoveLayoutRewriteBlock : public StmtExprMutator { // Step 4. Set block body as no_op Stmt old_body = block->body; - SBlockNode* n = block.CopyOnWrite(); + SBlockObj* n = block.CopyOnWrite(); n->body = SeqStmt({}); n->reads = {}; n->writes = {}; @@ -128,9 +128,9 @@ class RemoveLayoutRewriteBlock : public StmtExprMutator { std::unordered_set rewritten_tensors_; /*! \brief Maps a tensor load to an index map associated with the load / store in a layout rewrite block. */ - std::unordered_map tensor_var_to_index_map_; + std::unordered_map tensor_var_to_index_map_; /*! \brief Maps a tensor load to the shape of the corresponding rewritten tensor. */ - std::unordered_map> tensor_var_to_rewritten_shape_; + std::unordered_map> tensor_var_to_rewritten_shape_; }; class WeightLayoutRewriteBlockRemover : public StmtExprMutator { @@ -138,7 +138,7 @@ class WeightLayoutRewriteBlockRemover : public StmtExprMutator { using StmtExprMutator::Mutate; using StmtExprMutator::Mutate_; UnchangedOr Mutate(ffi::AnyView value, InplaceMode inplace_mode) override { - if (value.as()) return ffi::Unchanged(); + if (value.as()) return ffi::Unchanged(); return StmtExprMutator::Mutate(value, inplace_mode); } @@ -146,7 +146,7 @@ class WeightLayoutRewriteBlockRemover : public StmtExprMutator { auto [f_, tensor_map, tensor_var_to_index_map, tensor_var_to_rewritten_shape] = RemoveLayoutRewriteBlock::Rewrite(f); - FunctionNode* n = f_.CopyOnWrite(); + FunctionObj* n = f_.CopyOnWrite(); ffi::Array params; for (const tvm::Var& param : f_->params) { diff --git a/src/s_tir/transform/rewrite_unsafe_select.cc b/src/s_tir/transform/rewrite_unsafe_select.cc index 0ee073315a74..10ea661b8683 100644 --- a/src/s_tir/transform/rewrite_unsafe_select.cc +++ b/src/s_tir/transform/rewrite_unsafe_select.cc @@ -41,12 +41,12 @@ class UnsafeExprDetector : public tvm::ExprFunctor { public: // select itself is always considered safe if condition is safe // Because we will issue guard to make sure it is. - bool Dispatch_(const prim::SelectNode* op) { return Dispatch(op->condition); } - bool Dispatch_(const CallNode* op) { + bool Dispatch_(const prim::SelectObj* op) { return Dispatch(op->condition); } + bool Dispatch_(const CallObj* op) { if (op->op.same_as(prim::if_then_else_op())) { return Dispatch(op->args[0].as_or_throw()); } else if (op->op.same_as(tirx::address_of_op())) { - if (const auto* load = op->args[0].as()) { + if (const auto* load = op->args[0].as()) { for (const auto& index : load->indices) { if (Dispatch(index)) { return true; @@ -69,52 +69,50 @@ class UnsafeExprDetector : public tvm::ExprFunctor { return true; } } - bool Dispatch_(const TensorLoadNode* op) { + bool Dispatch_(const TensorLoadObj* op) { // Load is considered unsafe. return true; } - bool Dispatch_(const prim::AddNode* op) final { return BinaryOp(op); } - bool Dispatch_(const prim::SubNode* op) final { return BinaryOp(op); } - bool Dispatch_(const prim::MulNode* op) final { return BinaryOp(op); } - bool Dispatch_(const prim::DivNode* op) final { return BinaryOp(op); } - bool Dispatch_(const prim::ModNode* op) final { return BinaryOp(op); } - bool Dispatch_(const prim::FloorDivNode* op) final { return BinaryOp(op); } - bool Dispatch_(const prim::FloorModNode* op) final { return BinaryOp(op); } - bool Dispatch_(const prim::MinNode* op) final { return BinaryOp(op); } - bool Dispatch_(const prim::MaxNode* op) final { return BinaryOp(op); } - bool Dispatch_(const prim::EQNode* op) final { return BinaryOp(op); } - bool Dispatch_(const prim::NENode* op) final { return BinaryOp(op); } - bool Dispatch_(const prim::LTNode* op) final { return BinaryOp(op); } - bool Dispatch_(const prim::LENode* op) final { return BinaryOp(op); } - bool Dispatch_(const prim::GTNode* op) final { return BinaryOp(op); } - bool Dispatch_(const prim::GENode* op) final { return BinaryOp(op); } - bool Dispatch_(const prim::AndNode* op) final { return BinaryOp(op); } - bool Dispatch_(const prim::OrNode* op) final { return BinaryOp(op); } - bool Dispatch_(const prim::LShiftNode* op) final { return BinaryOp(op); } - bool Dispatch_(const prim::RShiftNode* op) final { return BinaryOp(op); } - bool Dispatch_(const prim::BitwiseAndNode* op) final { return BinaryOp(op); } - bool Dispatch_(const prim::BitwiseOrNode* op) final { return BinaryOp(op); } - bool Dispatch_(const prim::BitwiseXorNode* op) final { return BinaryOp(op); } - bool Dispatch_(const prim::BitwiseNotNode* op) final { return Dispatch(op->a); } - bool Dispatch_(const prim::NotNode* op) final { return Dispatch(op->a); } - bool Dispatch_(const prim::LetNode* op) final { - return Dispatch(op->body) || Dispatch(op->value); - } - bool Dispatch_(const prim::CastNode* op) final { return Dispatch(op->value); } - bool Dispatch_(const prim::BroadcastNode* op) final { return Dispatch(op->value); } - bool Dispatch_(const prim::RampNode* op) final { + bool Dispatch_(const prim::AddObj* op) final { return BinaryOp(op); } + bool Dispatch_(const prim::SubObj* op) final { return BinaryOp(op); } + bool Dispatch_(const prim::MulObj* op) final { return BinaryOp(op); } + bool Dispatch_(const prim::DivObj* op) final { return BinaryOp(op); } + bool Dispatch_(const prim::ModObj* op) final { return BinaryOp(op); } + bool Dispatch_(const prim::FloorDivObj* op) final { return BinaryOp(op); } + bool Dispatch_(const prim::FloorModObj* op) final { return BinaryOp(op); } + bool Dispatch_(const prim::MinObj* op) final { return BinaryOp(op); } + bool Dispatch_(const prim::MaxObj* op) final { return BinaryOp(op); } + bool Dispatch_(const prim::EQObj* op) final { return BinaryOp(op); } + bool Dispatch_(const prim::NEObj* op) final { return BinaryOp(op); } + bool Dispatch_(const prim::LTObj* op) final { return BinaryOp(op); } + bool Dispatch_(const prim::LEObj* op) final { return BinaryOp(op); } + bool Dispatch_(const prim::GTObj* op) final { return BinaryOp(op); } + bool Dispatch_(const prim::GEObj* op) final { return BinaryOp(op); } + bool Dispatch_(const prim::AndObj* op) final { return BinaryOp(op); } + bool Dispatch_(const prim::OrObj* op) final { return BinaryOp(op); } + bool Dispatch_(const prim::LShiftObj* op) final { return BinaryOp(op); } + bool Dispatch_(const prim::RShiftObj* op) final { return BinaryOp(op); } + bool Dispatch_(const prim::BitwiseAndObj* op) final { return BinaryOp(op); } + bool Dispatch_(const prim::BitwiseOrObj* op) final { return BinaryOp(op); } + bool Dispatch_(const prim::BitwiseXorObj* op) final { return BinaryOp(op); } + bool Dispatch_(const prim::BitwiseNotObj* op) final { return Dispatch(op->a); } + bool Dispatch_(const prim::NotObj* op) final { return Dispatch(op->a); } + bool Dispatch_(const prim::LetObj* op) final { return Dispatch(op->body) || Dispatch(op->value); } + bool Dispatch_(const prim::CastObj* op) final { return Dispatch(op->value); } + bool Dispatch_(const prim::BroadcastObj* op) final { return Dispatch(op->value); } + bool Dispatch_(const prim::RampObj* op) final { return Dispatch(op->base) && Dispatch(op->stride); } - bool Dispatch_(const prim::ShuffleNode* op) final { + bool Dispatch_(const prim::ShuffleObj* op) final { for (PrimExpr e : op->vectors) { if (Dispatch(e)) return true; } return false; } - bool Dispatch_(const VarNode* op) final { return false; } - bool Dispatch_(const IntImmNode* op) final { return false; } - bool Dispatch_(const FloatImmNode* op) final { return false; } - bool Dispatch_(const StringImmNode* op) final { return false; } + bool Dispatch_(const VarObj* op) final { return false; } + bool Dispatch_(const IntImmObj* op) final { return false; } + bool Dispatch_(const FloatImmObj* op) final { return false; } + bool Dispatch_(const StringImmObj* op) final { return false; } private: template @@ -130,10 +128,10 @@ class UnsafeSelectRewriter : public StmtExprMutator { using StmtExprMutator::Mutate; using StmtExprMutator::Mutate_; - UnchangedOr Mutate_(const prim::SelectNode* op, InplaceMode inplace_mode) { + UnchangedOr Mutate_(const prim::SelectObj* op, InplaceMode inplace_mode) { PrimExpr expr = StmtExprMutator::Mutate_(op, inplace_mode).ValueOrUnchanged(ffi::GetRef(op)); - op = expr.as(); + op = expr.as(); UnsafeExprDetector unsafe; PrimType cond_ty = op->condition.ty(); bool cond_is_scalar_bool = cond_ty.MatchesCode(DLDataTypeCode::kDLBool) && cond_ty.IsScalar(); diff --git a/src/s_tir/transform/stmt_simplify.cc b/src/s_tir/transform/stmt_simplify.cc index 092c2633ddad..d731398f6434 100644 --- a/src/s_tir/transform/stmt_simplify.cc +++ b/src/s_tir/transform/stmt_simplify.cc @@ -30,10 +30,11 @@ #include #include #include -#include +#include #include #include #include +#include #include #include @@ -44,20 +45,21 @@ using namespace tvm::prim; namespace { -class StmtSimplifier : public StmtExprMutatorWithAnalyzer { +class StmtSimplifier : public StmtExprMutatorWithConstraintContext { public: - using StmtExprMutatorWithAnalyzer::Mutate; - using StmtExprMutatorWithAnalyzer::Mutate_; - explicit StmtSimplifier(const sym::Analyzer& analyzer) : StmtExprMutatorWithAnalyzer(analyzer) {} + using StmtExprMutatorWithConstraintContext::Mutate; + using StmtExprMutatorWithConstraintContext::Mutate_; + explicit StmtSimplifier(const sym::Analyzer& analyzer) + : StmtExprMutatorWithConstraintContext(analyzer->GetConstraintContext()) {} Function Run(Function func); - UnchangedOr Mutate_(const EvaluateNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const EvaluateObj* op, InplaceMode inplace_mode) final { Stmt stmt = Parent::Mutate_(op, inplace_mode).ValueOrUnchanged(ffi::GetRef(op)); - if (const auto* evaluate = stmt.as()) { - if (const auto* call = evaluate->value.as(); + if (const auto* evaluate = stmt.as()) { + if (const auto* call = evaluate->value.as(); call && call->op.same_as(s_tir::async_wait()) && - analyzer_->CanProve(call->args[1].as_or_throw() < 0)) { + GetSyncedAnalyzer()->CanProve(call->args[1].as_or_throw() < 0)) { // Unrolling may make a wait count negative; such a wait is a no-op. return Evaluate(0); } @@ -66,18 +68,18 @@ class StmtSimplifier : public StmtExprMutatorWithAnalyzer { } protected: - using Parent = StmtExprMutatorWithAnalyzer; + using Parent = StmtExprMutatorWithConstraintContext; UnchangedOr Mutate(ffi::AnyView input, InplaceMode inplace_mode) final; - UnchangedOr Mutate_(const ForNode* op, InplaceMode inplace_mode) final; + UnchangedOr Mutate_(const ForObj* op, InplaceMode inplace_mode) final; - UnchangedOr Mutate_(const BindNode* op, InplaceMode inplace_mode) override; + UnchangedOr Mutate_(const BindObj* op, InplaceMode inplace_mode) override; - UnchangedOr Mutate_(const IfNode* op, InplaceMode inplace_mode) override; + UnchangedOr Mutate_(const IfObj* op, InplaceMode inplace_mode) override; // eliminate useless stores - UnchangedOr Mutate_(const TensorStoreNode* op, InplaceMode inplace_mode) override; + UnchangedOr Mutate_(const TensorStoreObj* op, InplaceMode inplace_mode) override; private: bool ArrayDeepEqual(const ffi::Array& lhs, const ffi::Array& rhs); @@ -86,7 +88,7 @@ class StmtSimplifier : public StmtExprMutatorWithAnalyzer { * * Substitutes any known Bind values and then simplifies with the analyzer. */ - ffi::Optional ProveCondition(PrimExpr condition) const; + ffi::Optional ProveCondition(PrimExpr condition); // Pure Bind values kept for substitution into assert conditions. // Grows monotonically under SSA — no scope-based cleanup required. @@ -94,7 +96,7 @@ class StmtSimplifier : public StmtExprMutatorWithAnalyzer { }; Function StmtSimplifier::Run(Function func) { - MarkTensorParamShapes(func); + AddParamConstraints(func); auto* n = func.CopyOnWrite(); // Shared string literals no longer enter the primitive analyzer. Inline their // SSA bindings so existing literal arguments remain constants during lowering. @@ -114,7 +116,9 @@ Function StmtSimplifier::Run(Function func) { return ffi::Unchanged(); }) .as_or_throw>(); - n->body = Mutate(n->body, InplaceMode::kAllow).ValueOrUnchanged(n->body); + n->body = Mutate(n->body, InplaceMode::kAllow) + .as_or_throw>>() + .ValueOrUnchanged(n->body); return func; } @@ -125,36 +129,36 @@ UnchangedOr StmtSimplifier::Mutate(ffi::AnyView input, InplaceMode inp if (def_region_kind() != kTVMFFIDefRegionKindNone && input.as()) { return Parent::Mutate(input, inplace_mode); } + if (auto call = input.as(); call && call.value()->op.same_as(assume_always_op())) { + return Parent::Mutate(input, inplace_mode); + } if (auto expr = input.as()) { - PrimExpr simplified = analyzer_->Simplify(*expr); + PrimExpr simplified = GetSyncedAnalyzer()->Simplify(*expr); if (simplified.same_as(*expr)) return ffi::Unchanged(); return simplified; } return Parent::Mutate(input, inplace_mode); } -UnchangedOr StmtSimplifier::Mutate_(const ForNode* op, InplaceMode inplace_mode) { - analyzer_->Bind(op->var, Range(op->begin, op->end)); - With ctx1(analyzer_, op->var >= op->begin); - With ctx2(analyzer_, static_cast(op->var) < op->end); +UnchangedOr StmtSimplifier::Mutate_(const ForObj* op, InplaceMode inplace_mode) { return Parent::Mutate_(op, inplace_mode); } -UnchangedOr StmtSimplifier::Mutate_(const BindNode* op, InplaceMode inplace_mode) { - if (const auto* call = op->value.as()) { +UnchangedOr StmtSimplifier::Mutate_(const BindObj* op, InplaceMode inplace_mode) { + if (const auto* call = op->value.as()) { // Preserve tensor type metadata; only declaration data is simplified. if (call->op.same_as(tirx::alloc_tensor_op())) return ffi::Unchanged(); if (call->op.same_as(tirx::decl_tensor_op())) { // The Call and its arguments may be shared even when the Bind is writable. auto data = this->Mutate(call->args[0], InplaceMode::kDisallow); if (data.UnchangedOrSameAs(call->args[0])) return ffi::Unchanged(); - auto value = ffi::make_object(*call); + auto value = ffi::make_object(*call); value->args.Set(0, std::move(data).ValueUnchecked()); if (inplace_mode == InplaceMode::kAllow) { - const_cast(op)->value = Expr(std::move(value)); + const_cast(op)->value = Expr(std::move(value)); return ffi::Unchanged(); } - auto copy = ffi::make_object(*op); + auto copy = ffi::make_object(*op); copy->value = Expr(std::move(value)); return Stmt(std::move(copy)); } @@ -165,14 +169,11 @@ UnchangedOr StmtSimplifier::Mutate_(const BindNode* op, InplaceMode inplac } PrimExpr value = this->Mutate(prim_value.value(), inplace_mode).ValueOrUnchanged(prim_value.value()); - // Bind in analyzer for constraint proving and simplification of - // subsequent expressions. Don't remove the Bind statement -- - // with flat Bind there's no body to inspect for usage patterns, - // so we always keep the Bind. + // Keep the statement: the sequence collector records its transformed + // binding for following siblings after this callback returns. if (SideEffect(value) <= CallEffectKind::kPure) { - analyzer_->Bind(op->var, value); // Record the binding so we can substitute it into assert conditions - // (see Mutate_(const AssertStmtNode*, InplaceMode)). Under SSA each var is + // (see Mutate_(const AssertStmtObj*, InplaceMode)). Under SSA each var is // bound exactly once, so the map grows monotonically without key // conflicts. No scope-based cleanup is needed because vars bound // in inner scopes are only referenced within those scopes; stale @@ -184,17 +185,17 @@ UnchangedOr StmtSimplifier::Mutate_(const BindNode* op, InplaceMode inplac return ffi::Unchanged(); } else { if (inplace_mode == InplaceMode::kAllow) { - auto* n = const_cast(op); + auto* n = const_cast(op); n->value = std::move(value); return ffi::Unchanged(); } - auto n = ffi::make_object(*op); + auto n = ffi::make_object(*op); n->value = std::move(value); return Stmt(n); } } -UnchangedOr StmtSimplifier::Mutate_(const IfNode* op, InplaceMode inplace_mode) { +UnchangedOr StmtSimplifier::Mutate_(const IfObj* op, InplaceMode inplace_mode) { if (ffi::Optional cond = ProveCondition(op->condition)) { if (cond.value()) { return this->Mutate(op->then_case, inplace_mode).ValueOrUnchanged(op->then_case); @@ -209,11 +210,11 @@ UnchangedOr StmtSimplifier::Mutate_(const IfNode* op, InplaceMode inplace_ } } -UnchangedOr StmtSimplifier::Mutate_(const TensorStoreNode* op, InplaceMode inplace_mode) { +UnchangedOr StmtSimplifier::Mutate_(const TensorStoreObj* op, InplaceMode inplace_mode) { TensorStore store = Parent::Mutate_(op, inplace_mode) .ValueOrUnchanged(ffi::GetRef(op)) .as_or_throw(); - if (const TensorLoadNode* load = store->value.as()) { + if (const TensorLoadObj* load = store->value.as()) { TensorVar tensor = load->source.as_or_throw(); if (tensor.same_as(store->dest.as_or_throw()) && ArrayDeepEqual(load->indices, store->indices)) { @@ -236,15 +237,15 @@ bool StmtSimplifier::ArrayDeepEqual(const ffi::Array& lhs, return true; } -ffi::Optional StmtSimplifier::ProveCondition(PrimExpr condition) const { +ffi::Optional StmtSimplifier::ProveCondition(PrimExpr condition) { auto f_substitute = [this](const Var& var) -> ffi::Expected> { if (auto repl = non_inlined_bindings_.Get(var)) return ffi::Any(*std::move(repl)); return ffi::Unchanged(); }; condition = ffi::StructuralMap(condition, f_substitute) .as_or_throw(); - condition = analyzer_->Simplify(condition); - if (const auto* as_int = condition.as()) { + condition = GetSyncedAnalyzer()->Simplify(condition); + if (const auto* as_int = condition.as()) { return as_int->value != 0; } else { return std::nullopt; diff --git a/src/s_tir/transform/storage_access.cc b/src/s_tir/transform/storage_access.cc index e47bd8a5fa0c..1d76ff31579f 100644 --- a/src/s_tir/transform/storage_access.cc +++ b/src/s_tir/transform/storage_access.cc @@ -41,19 +41,19 @@ namespace { ffi::Optional GetBufferDataVar(const ffi::Any& data) { if (auto var = data.as()) { - if (var.value()->ty.as() || var.value()->ty.as()) return var; + if (var.value()->ty.as() || var.value()->ty.as()) return var; return std::nullopt; } - if (const auto* call = data.as(); + if (const auto* call = data.as(); call && call->op.same_as(tirx::tensor_data_ptr_op()) && call->args.size() == 1) { return call->args[0].as(); } - if (const auto* call = data.as()) { + if (const auto* call = data.as()) { if (call->op.same_as(tirx::ptr_byte_offset_op()) || call->op.same_as(tirx::reinterpret_op())) { return GetBufferDataVar(call->args[0]); } if (call->op.same_as(tirx::address_of_op())) { - if (const auto* load = call->args[0].as()) { + if (const auto* load = call->args[0].as()) { return load->source.as(); } } @@ -63,7 +63,7 @@ ffi::Optional GetBufferDataVar(const ffi::Any& data) { } // namespace -ffi::Optional StorageAccessVisitor::Visit_(const TensorLoadNode* op) { +ffi::Optional StorageAccessVisitor::Visit_(const TensorLoadObj* op) { Var buf = ResolveBuffer(op->source.as_or_throw()); MemoryScope scope = op->source.as_or_throw().scope(); if (Enabled(buf.get(), scope)) { @@ -84,7 +84,7 @@ ffi::Optional StorageAccessVisitor::Visit_(const TensorLoadNode* return StmtExprVisitor::Visit_(op); } -ffi::Optional StorageAccessVisitor::Visit_(const TensorStoreNode* op) { +ffi::Optional StorageAccessVisitor::Visit_(const TensorStoreObj* op) { allow_append_ = true; TVM_FFI_ICHECK_EQ(curr_stmt_.access.size(), 0U); curr_stmt_.stmt = op; @@ -114,7 +114,7 @@ ffi::Optional StorageAccessVisitor::Visit_(const TensorStoreNode return std::nullopt; } -ffi::Optional StorageAccessVisitor::Visit_(const EvaluateNode* op) { +ffi::Optional StorageAccessVisitor::Visit_(const EvaluateObj* op) { allow_append_ = true; TVM_FFI_ICHECK_EQ(curr_stmt_.access.size(), 0U); curr_stmt_.stmt = op; @@ -128,15 +128,15 @@ ffi::Optional StorageAccessVisitor::Visit_(const EvaluateNode* o return std::nullopt; } -ffi::Optional StorageAccessVisitor::Visit_(const BindNode* op) { - if (op->var->ty.as()) { +ffi::Optional StorageAccessVisitor::Visit_(const BindObj* op) { + if (op->var->ty.as()) { pointer_values_.insert_or_assign(op->var.get(), op->value); if (HasOffsetPointer(op->value)) offset_aliases_.insert(op->var.get()); if (auto source = GetBufferDataVar(op->value)) { buffer_aliases_.insert_or_assign(op->var.get(), ResolveBuffer(source.value())); } } - if (const auto* call = op->value.as(); + if (const auto* call = op->value.as(); call && call->op.same_as(tirx::decl_tensor_op())) { if (HasOffsetPointer(call->args[0])) offset_aliases_.insert(op->var.get()); if (auto source = GetBufferDataVar(call->args[0])) { @@ -145,8 +145,7 @@ ffi::Optional StorageAccessVisitor::Visit_(const BindNode* op) { } return StmtExprVisitor::Visit_(op); } - if (const auto* call = op->value.as(); - call && call->op.same_as(tirx::alloc_tensor_op())) + if (const auto* call = op->value.as(); call && call->op.same_as(tirx::alloc_tensor_op())) return StmtExprVisitor::Visit_(op); allow_append_ = true; TVM_FFI_ICHECK_EQ(curr_stmt_.access.size(), 0U); @@ -160,7 +159,7 @@ ffi::Optional StorageAccessVisitor::Visit_(const BindNode* op) { return std::nullopt; } -ffi::Optional StorageAccessVisitor::Visit_(const RegionStmtNode* op) { +ffi::Optional StorageAccessVisitor::Visit_(const RegionStmtObj* op) { if (op->op.same_as(s_tir::manual_sync())) { // Trust the region's explicit barriers instead of planning access-based synchronization. return std::nullopt; @@ -198,7 +197,7 @@ ffi::Optional StorageAccessVisitor::Visit_(const RegionStmtNode* return std::nullopt; } -ffi::Optional StorageAccessVisitor::Visit_(const ForNode* op) { +ffi::Optional StorageAccessVisitor::Visit_(const ForObj* op) { scope_.push_back(std::vector()); TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(op)); StmtEntry s; @@ -227,7 +226,7 @@ ffi::Optional StorageAccessVisitor::Visit_(const ForNode* op) { } bool IsThreadInvariant(const PrimExpr& cond) { - if (auto call = cond.as()) { + if (auto call = cond.as()) { if (auto opt_call_op = call->op.as()) { auto call_op = opt_call_op.value(); if (call_op.same_as(tirx::gpu_thread_invariant_op())) { @@ -238,7 +237,7 @@ bool IsThreadInvariant(const PrimExpr& cond) { return false; } -ffi::Optional StorageAccessVisitor::Visit_(const IfNode* op) { +ffi::Optional StorageAccessVisitor::Visit_(const IfObj* op) { bool is_thread_invariant = IsThreadInvariant(op->condition); if (!is_thread_invariant) { ++condition_counter_; @@ -264,7 +263,7 @@ ffi::Optional StorageAccessVisitor::Visit_(const IfNode* op) { return std::nullopt; } -ffi::Optional StorageAccessVisitor::Visit_(const WhileNode* op) { +ffi::Optional StorageAccessVisitor::Visit_(const WhileObj* op) { bool is_thread_invariant = IsThreadInvariant(op->condition); if (!is_thread_invariant) { ++condition_counter_; @@ -284,8 +283,8 @@ ffi::Optional StorageAccessVisitor::Visit_(const WhileNode* op) } bool StorageAccessVisitor::HasOffsetPointer(const Expr& value) const { - if (const auto* var = value.as()) return offset_aliases_.count(var); - if (const auto* call = value.as()) { + if (const auto* var = value.as()) return offset_aliases_.count(var); + if (const auto* call = value.as()) { if (call->op.same_as(tirx::ptr_byte_offset_op()) || call->op.same_as(tirx::address_of_op()) || call->op.same_as(tirx::reinterpret_op())) return true; @@ -321,17 +320,17 @@ void StorageAccessVisitor::RecordOpaqueAccess(Var source) { curr_stmt_.access.push_back(std::move(entry)); } -ffi::Optional StorageAccessVisitor::Visit_(const VarNode* op) { +ffi::Optional StorageAccessVisitor::Visit_(const VarObj* op) { if (allow_append_ && in_opaque_call_) { if (auto it = pointer_values_.find(op); it != pointer_values_.end()) { return Visit(it->second); } - if (op->ty.as()) RecordOpaqueAccess(ffi::GetRef(op)); + if (op->ty.as()) RecordOpaqueAccess(ffi::GetRef(op)); } return StmtExprVisitor::Visit_(op); } -ffi::Optional StorageAccessVisitor::Visit_(const CallNode* op) { +ffi::Optional StorageAccessVisitor::Visit_(const CallObj* op) { static const auto scope_warp = runtime::MemoryScope::Get("warp"); Call call = ffi::GetRef(op); @@ -340,7 +339,7 @@ ffi::Optional StorageAccessVisitor::Visit_(const CallNode* op) { if (op->op.same_as(tirx::tensor_data_ptr_op())) { tensor = op->args[0].as(); } else if (op->op.same_as(tirx::address_of_op())) { - if (const auto* load = op->args[0].as()) { + if (const auto* load = op->args[0].as()) { tensor = load->source.as(); } } @@ -371,7 +370,7 @@ ffi::Optional StorageAccessVisitor::Visit_(const CallNode* op) { } TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(op)); } else if (op->op.same_as(tirx::address_of_op())) { - if (const auto* load = op->args[0].as()) { + if (const auto* load = op->args[0].as()) { // Taking an address does not read the tensor value. Visit only the // load's children so index expressions still contribute accesses. TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(load)); @@ -382,7 +381,7 @@ ffi::Optional StorageAccessVisitor::Visit_(const CallNode* op) { } } else if (op->op.same_as(tirx::gpu_storage_sync_op())) { TVM_FFI_ICHECK(allow_append_); - const std::string& s = op->args[0].as()->value; + const std::string& s = op->args[0].as()->value; MemoryScope scope = runtime::MemoryScope::Get(s); if (scope != scope_warp) { AccessEntry e; @@ -410,7 +409,7 @@ MemoryScope StorageAccessVisitor::GetScope(Var buffer_var) const { if (auto tensor_type = buffer_var->ty.as()) { return tensor_type.value()->scope; } - if (const auto* pointer_type = buffer_var->ty.as()) { + if (const auto* pointer_type = buffer_var->ty.as()) { return pointer_type->scope; } return MemoryScope(); // global by default diff --git a/src/s_tir/transform/storage_access.h b/src/s_tir/transform/storage_access.h index 475eedd2ab3d..c445369e973f 100644 --- a/src/s_tir/transform/storage_access.h +++ b/src/s_tir/transform/storage_access.h @@ -81,16 +81,16 @@ class StorageAccessVisitor : public StmtExprVisitor { std::vector access; }; // override visitor pattern - ffi::Optional Visit_(const TensorLoadNode* op) final; - ffi::Optional Visit_(const TensorStoreNode* op) final; - ffi::Optional Visit_(const EvaluateNode* op) final; - ffi::Optional Visit_(const BindNode* op) final; - ffi::Optional Visit_(const RegionStmtNode* op) final; - ffi::Optional Visit_(const ForNode* op) final; - ffi::Optional Visit_(const IfNode* op) final; - ffi::Optional Visit_(const WhileNode* op) final; - ffi::Optional Visit_(const CallNode* op) final; - ffi::Optional Visit_(const VarNode* op) final; + ffi::Optional Visit_(const TensorLoadObj* op) final; + ffi::Optional Visit_(const TensorStoreObj* op) final; + ffi::Optional Visit_(const EvaluateObj* op) final; + ffi::Optional Visit_(const BindObj* op) final; + ffi::Optional Visit_(const RegionStmtObj* op) final; + ffi::Optional Visit_(const ForObj* op) final; + ffi::Optional Visit_(const IfObj* op) final; + ffi::Optional Visit_(const WhileObj* op) final; + ffi::Optional Visit_(const CallObj* op) final; + ffi::Optional Visit_(const VarObj* op) final; protected: StorageAccessVisitor() { scope_.push_back(std::vector()); } @@ -106,7 +106,7 @@ class StorageAccessVisitor : public StmtExprVisitor { * \param scope The scope of the buffer. * \return Whether the analysis of buffer is enabled. */ - virtual bool Enabled(const VarNode* buffer, const MemoryScope& scope) const { return true; } + virtual bool Enabled(const VarObj* buffer, const MemoryScope& scope) const { return true; } /*! * \brief Summarize the sequence of operations into parent. * @@ -118,7 +118,7 @@ class StorageAccessVisitor : public StmtExprVisitor { * \return The summarized sequence that represent access that * the parent should taken care of to synchronize. */ - virtual std::vector Summarize(std::vector seq, const ForNode* loop) = 0; + virtual std::vector Summarize(std::vector seq, const ForObj* loop) = 0; /*! * \brief Get the scope of the buffer array. * \return The scope of the final buffer array. @@ -145,11 +145,11 @@ class StorageAccessVisitor : public StmtExprVisitor { // The involving threads ffi::Array env_threads_; // Physical storage root for each declared logical buffer view. - std::unordered_map buffer_aliases_; + std::unordered_map buffer_aliases_; // Resolve pointer bindings at their use, without treating address construction as access. - std::unordered_map pointer_values_; + std::unordered_map pointer_values_; // Shifted or retyped aliases cannot compare coordinates with their storage root. - std::unordered_set offset_aliases_; + std::unordered_set offset_aliases_; }; } // namespace s_tir } // namespace tvm diff --git a/src/s_tir/transform/tensorcore_infer_fragment.cc b/src/s_tir/transform/tensorcore_infer_fragment.cc index f41f1046a769..1b2a0aec27db 100644 --- a/src/s_tir/transform/tensorcore_infer_fragment.cc +++ b/src/s_tir/transform/tensorcore_infer_fragment.cc @@ -43,13 +43,13 @@ namespace tvm { namespace s_tir { using namespace tvm::tirx; -const VarNode* GetBufferVarFromData(const Expr& data) { - if (const auto* var = data.as()) { +const VarObj* GetBufferVarFromData(const Expr& data) { + if (const auto* var = data.as()) { return var; } - if (const auto* call = data.as(); + if (const auto* call = data.as(); call && call->op.same_as(tirx::tensor_data_ptr_op()) && call->args.size() == 1) { - return call->args[0].as(); + return call->args[0].as(); } return nullptr; } @@ -58,7 +58,7 @@ const VarNode* GetBufferVarFromData(const Expr& data) { class FragmentGetter : public s_tir::StmtExprVisitor { public: using s_tir::StmtExprVisitor::Visit_; - ffi::Optional Visit_(const CallNode* op) final { + ffi::Optional Visit_(const CallObj* op) final { TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(s_tir::StmtExprVisitor::Visit_(op)); static const Op gpu_load_matrix_sync_op = Op::Get("tirx.gpu_load_matrix_sync"); @@ -71,21 +71,21 @@ class FragmentGetter : public s_tir::StmtExprVisitor { // Get shape and layout information from load and store intrinsic TVM_FFI_ICHECK_EQ(op->args.size(), 8U); - const VarNode* buffer_var = GetBufferVarFromData(op->args[0]); + const VarObj* buffer_var = GetBufferVarFromData(op->args[0]); TVM_FFI_ICHECK(buffer_var); // Get shape - const IntImmNode* m = op->args[1].as(); - const IntImmNode* n = op->args[2].as(); - const IntImmNode* k = op->args[3].as(); - const StringImmNode* layout = op->args[7].as(); + const IntImmObj* m = op->args[1].as(); + const IntImmObj* n = op->args[2].as(); + const IntImmObj* k = op->args[3].as(); + const StringImmObj* layout = op->args[7].as(); TVM_FFI_ICHECK(m); TVM_FFI_ICHECK(n); TVM_FFI_ICHECK(k); TVM_FFI_ICHECK(layout); runtime::MemoryScope scope = - (buffer_var->ty.as() ? buffer_var->ty.as()->scope - : buffer_var->ty.as_or_throw()->scope); + (buffer_var->ty.as() ? buffer_var->ty.as()->scope + : buffer_var->ty.as_or_throw()->scope); if (fragments.count(buffer_var)) { // check if the fragment has met before FragmentInfo info = fragments[buffer_var]; @@ -110,19 +110,19 @@ class FragmentGetter : public s_tir::StmtExprVisitor { } else if (op->op.same_as(gpu_fill_fragment_op)) { // Get shape information from fill intrinsic TVM_FFI_ICHECK_EQ(op->args.size(), 6U); - const VarNode* buffer_var = GetBufferVarFromData(op->args[0]); + const VarObj* buffer_var = GetBufferVarFromData(op->args[0]); TVM_FFI_ICHECK(buffer_var); // Get shape - const IntImmNode* m = op->args[1].as(); - const IntImmNode* n = op->args[2].as(); - const IntImmNode* k = op->args[3].as(); + const IntImmObj* m = op->args[1].as(); + const IntImmObj* n = op->args[2].as(); + const IntImmObj* k = op->args[3].as(); TVM_FFI_ICHECK(m); TVM_FFI_ICHECK(n); TVM_FFI_ICHECK(k); runtime::MemoryScope scope = - (buffer_var->ty.as() ? buffer_var->ty.as()->scope - : buffer_var->ty.as_or_throw()->scope); + (buffer_var->ty.as() ? buffer_var->ty.as()->scope + : buffer_var->ty.as_or_throw()->scope); if (fragments.count(buffer_var)) { FragmentInfo info = fragments[buffer_var]; TVM_FFI_ICHECK_EQ(m->value, info.m); @@ -139,10 +139,10 @@ class FragmentGetter : public s_tir::StmtExprVisitor { } // Fragment metadata for all fragments - std::unordered_map fragments; + std::unordered_map fragments; }; -std::unordered_map GetTensorCoreFragmentInfo(const Stmt& stmt) { +std::unordered_map GetTensorCoreFragmentInfo(const Stmt& stmt) { auto getter = ffi::make_object(); getter->Visit(stmt); return std::move(getter->fragments); @@ -154,17 +154,17 @@ class FragmentChecker : public s_tir::StmtExprVisitor { using s_tir::StmtExprVisitor::Visit_; explicit FragmentChecker(const FragmentGetter& getter) : fragment_getter(getter) {} - ffi::Optional Visit_(const CallNode* op) final { + ffi::Optional Visit_(const CallObj* op) final { TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(s_tir::StmtExprVisitor::Visit_(op)); // Check shape when calling gpu_mma_sync static const Op gpu_mma_sync_op = Op::Get("tirx.gpu_mma_sync"); static const Op bmma_sync_op = Op::Get("tirx.cuda.bmma_sync"); if (op->op.same_as(gpu_mma_sync_op) || op->op.same_as(bmma_sync_op)) { TVM_FFI_ICHECK_EQ(op->args.size(), 8U); - const VarNode* buffer_var_d = GetBufferVarFromData(op->args[0]); - const VarNode* buffer_var_a = GetBufferVarFromData(op->args[2]); - const VarNode* buffer_var_b = GetBufferVarFromData(op->args[4]); - const VarNode* buffer_var_c = GetBufferVarFromData(op->args[6]); + const VarObj* buffer_var_d = GetBufferVarFromData(op->args[0]); + const VarObj* buffer_var_a = GetBufferVarFromData(op->args[2]); + const VarObj* buffer_var_b = GetBufferVarFromData(op->args[4]); + const VarObj* buffer_var_c = GetBufferVarFromData(op->args[6]); TVM_FFI_ICHECK(buffer_var_d); TVM_FFI_ICHECK(buffer_var_a); TVM_FFI_ICHECK(buffer_var_b); @@ -180,7 +180,7 @@ class FragmentChecker : public s_tir::StmtExprVisitor { private: // A tool for checking shapes of two fragments - bool CheckShape(const VarNode* buffer1, const VarNode* buffer2) { + bool CheckShape(const VarObj* buffer1, const VarObj* buffer2) { TVM_FFI_ICHECK(fragment_getter.fragments.count(buffer1)) << "Tensorecore fragment " << buffer1->name << " must be filled (with gpu_fill_fragment) or loaded (with gpu_load_matrix_sync) before " @@ -203,14 +203,14 @@ class InferFragmenter : public s_tir::StmtExprMutator { using s_tir::StmtExprMutator::Mutate; using s_tir::StmtExprMutator::Mutate_; UnchangedOr Mutate(ffi::AnyView value, InplaceMode inplace_mode) override { - if (value.as()) return ffi::Unchanged(); + if (value.as()) return ffi::Unchanged(); return s_tir::StmtExprMutator::Mutate(value, inplace_mode); } explicit InferFragmenter(const FragmentGetter& getter) : fragment_getter(getter) {} - UnchangedOr Mutate_(const BindNode* op, InplaceMode inplace_mode) final { - if (const auto* call = op->value.as(); + UnchangedOr Mutate_(const BindObj* op, InplaceMode inplace_mode) final { + if (const auto* call = op->value.as(); call && call->op.same_as(tirx::alloc_tensor_op())) { auto it = fragment_getter.fragments.find(op->var.get()); if (it == fragment_getter.fragments.end()) return ffi::Unchanged(); diff --git a/src/s_tir/transform/thread_storage_sync.cc b/src/s_tir/transform/thread_storage_sync.cc index 87314c8e5f56..1bacd4d14b33 100644 --- a/src/s_tir/transform/thread_storage_sync.cc +++ b/src/s_tir/transform/thread_storage_sync.cc @@ -51,11 +51,11 @@ class ThreadSyncPlanner : public StorageAccessVisitor { std::unordered_set syncs_inserted_; protected: - bool Enabled(const VarNode* buf, const MemoryScope& scope) const final { + bool Enabled(const VarObj* buf, const MemoryScope& scope) const final { return in_device_env() && scope == sync_scope_; } // Plan the sync - std::vector Summarize(std::vector seq, const ForNode* loop) final { + std::vector Summarize(std::vector seq, const ForObj* loop) final { static const auto scope_shared_dyn = runtime::MemoryScope::Get("shared.dyn"); // Redirect all "shared.dyn" buffer access to the same buffer var @@ -223,7 +223,7 @@ class ThreadSyncPlanner : public StorageAccessVisitor { // Even if access has the same index, those indices need to // depend on the innermost thread id to avoid race condition bool depends_on_thread_index = true; - const VarNode* thread_index_var = nullptr; + const VarObj* thread_index_var = nullptr; if (!curr.threads.empty()) { thread_index_var = curr.threads.back()->var.get(); } @@ -237,7 +237,7 @@ class ThreadSyncPlanner : public StorageAccessVisitor { PrimExpr curr_index = curr_intset.PointValue(); has_same_index = prim::ExprDeepEqual()(prev_index, curr_index); if (thread_index_var != nullptr) { - auto f_uses_thread_index = [=](const tvm::VarNode* parameter) { + auto f_uses_thread_index = [=](const tvm::VarObj* parameter) { return parameter == thread_index_var; }; auto walkfn = [&](const Var& var) -> ffi::Expected { @@ -299,9 +299,8 @@ class ThreadSyncAfterWaitQueueInserter : public StmtExprMutator { explicit ThreadSyncAfterWaitQueueInserter(MemoryScope sync_scope) : sync_scope_(sync_scope) {} - UnchangedOr Mutate_(const EvaluateNode* op, InplaceMode inplace_mode) final { - if (const auto* call = op->value.as(); - call && call->op.same_as(s_tir::async_wait())) { + UnchangedOr Mutate_(const EvaluateObj* op, InplaceMode inplace_mode) final { + if (const auto* call = op->value.as(); call && call->op.same_as(s_tir::async_wait())) { auto sync = Evaluate( Call(PrimType::Int(32), tirx::gpu_storage_sync_op(), {StringImm(sync_scope_.name())})); return SeqStmt({ffi::GetRef(op), sync}); @@ -322,7 +321,7 @@ class ThreadSyncInserter : public StmtExprMutator { : sync_scope_(sync_scope), syncs_(syncs) {} UnchangedOr Mutate(ffi::AnyView value, InplaceMode inplace_mode) final { - const auto* stmt = value.as(); + const auto* stmt = value.as(); if (!stmt) return StmtExprMutator::Mutate(value, inplace_mode); if (syncs_.empty()) return ffi::Unchanged(); if (!syncs_.count(stmt)) return StmtExprMutator::Mutate(value, inplace_mode); @@ -364,15 +363,15 @@ class SynchronizationLowerer : public StmtExprMutator { using StmtExprMutator::Mutate_; explicit SynchronizationLowerer(bool is_cuda) : is_cuda_(is_cuda) {} - UnchangedOr Mutate_(const RegionStmtNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const RegionStmtObj* op, InplaceMode inplace_mode) final { if (op->op.same_as(s_tir::manual_sync()) || op->op.same_as(s_tir::async_copy_scope())) { return Mutate(op->body, inplace_mode).ValueOrUnchanged(op->body); } return StmtExprMutator::Mutate_(op, inplace_mode); } - UnchangedOr Mutate_(const EvaluateNode* op, InplaceMode inplace_mode) final { - const auto* call = op->value.as(); + UnchangedOr Mutate_(const EvaluateObj* op, InplaceMode inplace_mode) final { + const auto* call = op->value.as(); if (!call || (!call->op.same_as(s_tir::async_commit()) && !call->op.same_as(s_tir::async_wait()))) { return StmtExprMutator::Mutate_(op, inplace_mode); diff --git a/src/s_tir/transform/transform_mma_tensor_layout.cc b/src/s_tir/transform/transform_mma_tensor_layout.cc index 2990d890992d..e0f66be70cb8 100644 --- a/src/s_tir/transform/transform_mma_tensor_layout.cc +++ b/src/s_tir/transform/transform_mma_tensor_layout.cc @@ -50,7 +50,7 @@ class MmaTensorLayoutTransformer : public StmtExprMutator { using StmtExprMutator::Mutate; using StmtExprMutator::Mutate_; - UnchangedOr Mutate_(const SBlockNode* op, InplaceMode inplace_mode) { + UnchangedOr Mutate_(const SBlockObj* op, InplaceMode inplace_mode) { SBlock block = ffi::GetRef(op); auto* n = block.CopyOnWrite(); auto fmutate = [this](const TensorVar& tensor) { @@ -68,8 +68,8 @@ class MmaTensorLayoutTransformer : public StmtExprMutator { // bi = 16, bj = 8 size_t size = tensor.shape().size(); TVM_FFI_ICHECK_GE(size, 2); - const IntImmNode* dim0 = tensor.shape()[size - 2].as(); - const IntImmNode* dim1 = tensor.shape()[size - 1].as(); + const IntImmObj* dim0 = tensor.shape()[size - 2].as(); + const IntImmObj* dim1 = tensor.shape()[size - 1].as(); TVM_FFI_ICHECK(dim0 != nullptr && dim1 != nullptr); TVM_FFI_ICHECK(dim0->value % 16 == 0 && dim1->value % 8 == 0); @@ -90,8 +90,8 @@ class MmaTensorLayoutTransformer : public StmtExprMutator { // bi = 32, bj = 8 size_t size = tensor.shape().size(); TVM_FFI_ICHECK_GE(size, 2); - const IntImmNode* dim0 = tensor.shape()[size - 2].as(); - const IntImmNode* dim1 = tensor.shape()[size - 1].as(); + const IntImmObj* dim0 = tensor.shape()[size - 2].as(); + const IntImmObj* dim1 = tensor.shape()[size - 1].as(); TVM_FFI_ICHECK(dim0 != nullptr && dim1 != nullptr); TVM_FFI_ICHECK(dim0->value % 32 == 0 && dim1->value % 8 == 0); std::vector new_shape; @@ -111,8 +111,8 @@ class MmaTensorLayoutTransformer : public StmtExprMutator { // bj = 8, bj = 32 size_t size = tensor.shape().size(); TVM_FFI_ICHECK_GE(size, 2); - const IntImmNode* dim0 = tensor.shape()[size - 2].as(); - const IntImmNode* dim1 = tensor.shape()[size - 1].as(); + const IntImmObj* dim0 = tensor.shape()[size - 2].as(); + const IntImmObj* dim1 = tensor.shape()[size - 1].as(); TVM_FFI_ICHECK(dim0 != nullptr && dim1 != nullptr); TVM_FFI_ICHECK(dim0->value % 8 == 0 && dim1->value % 32 == 0); std::vector new_shape; @@ -134,7 +134,7 @@ class MmaTensorLayoutTransformer : public StmtExprMutator { return block; } - UnchangedOr Mutate_(const TensorStoreNode* op, InplaceMode inplace_mode) { + UnchangedOr Mutate_(const TensorStoreObj* op, InplaceMode inplace_mode) { const auto scope_mma_matrixC = GetRegisteredMemoryScope("mma.matrixC"); const auto scope_mma_matrixA = GetRegisteredMemoryScope("mma.matrixA"); const auto scope_mma_matrixB = GetRegisteredMemoryScope("mma.matrixB"); @@ -166,7 +166,7 @@ class MmaTensorLayoutTransformer : public StmtExprMutator { return store; } - UnchangedOr Mutate_(const TensorLoadNode* op, InplaceMode inplace_mode) { + UnchangedOr Mutate_(const TensorLoadObj* op, InplaceMode inplace_mode) { const auto scope_mma_matrixC = GetRegisteredMemoryScope("mma.matrixC"); TensorVar tensor = op->source.as_or_throw(); diff --git a/src/te/operation/compute_op.cc b/src/te/operation/compute_op.cc index 01106dca03ec..e89ae89e325b 100644 --- a/src/te/operation/compute_op.cc +++ b/src/te/operation/compute_op.cc @@ -42,22 +42,22 @@ namespace te { using namespace tirx; TVM_FFI_STATIC_INIT_BLOCK() { - OperationNode::RegisterReflection(); - BaseComputeOpNode::RegisterReflection(); - ComputeOpNode::RegisterReflection(); + OperationObj::RegisterReflection(); + BaseComputeOpObj::RegisterReflection(); + ComputeOpObj::RegisterReflection(); } // Pattern A (RM): auto-default repr from reflection. /// Verify if ComputeOp is valid with respect to Reduce operations. -static void VerifyComputeOp(const ComputeOpNode* op); +static void VerifyComputeOp(const ComputeOpObj* op); -static inline void AssertReduceEqual(const te::ReduceNode* a, const te::ReduceNode* b) { +static inline void AssertReduceEqual(const te::ReduceObj* a, const te::ReduceObj* b) { const char* shared_text = "When a TE compute node produces multiple outputs, " "each of which is a reduction, " "each reduction must be structurally identical, " - "except for the ReduceNode::value_index. "; + "except for the ReduceObj::value_index. "; ffi::StructuralEqual eq; // Compare expression values, not the binder semantics of axis metadata. @@ -82,14 +82,14 @@ static inline void AssertReduceEqual(const te::ReduceNode* a, const te::ReduceNo << shared_text << "However, the initial value " << a->init << " does not match " << b->init; } -int ComputeOpNode::num_outputs() const { return body.size(); } +int ComputeOpObj::num_outputs() const { return body.size(); } -PrimType ComputeOpNode::output_dtype(size_t idx) const { +PrimType ComputeOpObj::output_dtype(size_t idx) const { TVM_FFI_ICHECK_LT(idx, num_outputs()); return body[idx].ty(); } -ffi::Array BaseComputeOpNode::output_shape(size_t idx) const { +ffi::Array BaseComputeOpObj::output_shape(size_t idx) const { TVM_FFI_ICHECK_LT(idx, num_outputs()); // for now, all outputs of a BaseComputeOp have the same shape ffi::Array shape; @@ -145,14 +145,14 @@ ComputeOp::ComputeOp(std::string name, std::string tag, ffi::Map(); } - auto n = ffi::make_object(); + auto n = ffi::make_object(); n->name = std::move(name); n->tag = std::move(tag); n->attrs = std::move(attrs); n->axis = std::move(axis); n->body = std::move(body); - if (n->body[0]->IsInstance()) { - const te::ReduceNode* reduce = n->body[0].as(); + if (n->body[0]->IsInstance()) { + const te::ReduceObj* reduce = n->body[0].as(); n->reduce_axis = reduce->axis; } VerifyComputeOp(n.get()); @@ -169,7 +169,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { } // The schedule related logics -ffi::Array ComputeOpNode::InputTensors() const { +ffi::Array ComputeOpObj::InputTensors() const { ffi::Array ret; std::unordered_set visited; auto visit_call = [&ret, &visited](const Call& call, ffi::StructuralVisitorObj* visitor) @@ -197,7 +197,7 @@ ffi::Array ComputeOpNode::InputTensors() const { visit_call); }; for (const PrimExpr& e : body) { - if (const auto* reduce = e.as()) { + if (const auto* reduce = e.as()) { for (const IterVar& axis : reduce->axis) { visit(axis->dom.value()->begin); visit(axis->dom.value()->extent); @@ -228,7 +228,7 @@ class ComputeVerifier final : public tirx::StmtExprVisitor { public: /// Special member functions //@{ - explicit ComputeVerifier(const ComputeOpNode* compute) + explicit ComputeVerifier(const ComputeOpObj* compute) : tirx::StmtExprVisitor([] { static const VTable table = [] { VTable table; @@ -239,7 +239,7 @@ class ComputeVerifier final : public tirx::StmtExprVisitor { return &table; }()), compute_(compute), - reduce_(compute->body[0].as()) {} + reduce_(compute->body[0].as()) {} ~ComputeVerifier() = default; ComputeVerifier(const ComputeVerifier&) = delete; ComputeVerifier(ComputeVerifier&&) = delete; @@ -251,7 +251,7 @@ class ComputeVerifier final : public tirx::StmtExprVisitor { void Run() { for (const PrimExpr e : compute_->body) { // Check for consistency of top level reductions - const te::ReduceNode* reduce = e.as(); + const te::ReduceObj* reduce = e.as(); TVM_FFI_ICHECK((reduce && reduce_) || (!reduce && !reduce_)) << "All ComputeOp should be consistent " << "with being Reduce operation or not."; @@ -266,7 +266,7 @@ class ComputeVerifier final : public tirx::StmtExprVisitor { } ffi::Optional Visit(ffi::AnyView value) final { - if (!value.as()) return tirx::StmtExprVisitor::Visit(value); + if (!value.as()) return tirx::StmtExprVisitor::Visit(value); ++level_; auto interrupt = tirx::StmtExprVisitor::Visit(value); --level_; @@ -274,8 +274,8 @@ class ComputeVerifier final : public tirx::StmtExprVisitor { } using tirx::StmtExprVisitor::Visit_; - ffi::Optional Visit_(const IterVarNode*) { return std::nullopt; } - ffi::Optional Visit_(const te::ReduceNode* reduce) { + ffi::Optional Visit_(const IterVarObj*) { return std::nullopt; } + ffi::Optional Visit_(const te::ReduceObj* reduce) { TVM_FFI_ICHECK(0 == level_) << "Reductions are only allowed at the top level of compute. " << "Please create another tensor for further composition."; for (const PrimExpr& expr : reduce->combiner->result) { @@ -296,19 +296,19 @@ class ComputeVerifier final : public tirx::StmtExprVisitor { protected: static void InitVTable(VTable* table) { tirx::StmtExprVisitor::InitVTable(table); - SetDispatch(table); - SetDispatch(table); + SetDispatch(table); + SetDispatch(table); } private: - const ComputeOpNode* compute_{nullptr}; ///< ComputeOpNode to verify - const te::ReduceNode* reduce_{nullptr}; ///< Top level Reduce operation - int level_{0}; ///< Level of op being processed + const ComputeOpObj* compute_{nullptr}; ///< ComputeOpObj to verify + const te::ReduceObj* reduce_{nullptr}; ///< Top level Reduce operation + int level_{0}; ///< Level of op being processed }; } // namespace /// Verify if ComputeOp is valid with respect to Reduce operations. -static void VerifyComputeOp(const ComputeOpNode* op) { +static void VerifyComputeOp(const ComputeOpObj* op) { ffi::make_object(op)->Run(); } diff --git a/src/te/operation/create_function.cc b/src/te/operation/create_function.cc index 732a331bfb8f..10964734d717 100644 --- a/src/te/operation/create_function.cc +++ b/src/te/operation/create_function.cc @@ -59,7 +59,7 @@ namespace { class OpaqueArtifactVerifier : public ObjectVisitor { public: ffi::Optional Visit(ffi::AnyView value) final { - if (const auto* block = value.as()) { + if (const auto* block = value.as()) { for (const auto& axis : block->iter_vars) { Visit(axis->dom); Visit(axis->var); @@ -76,11 +76,11 @@ class OpaqueArtifactVerifier : public ObjectVisitor { Visit(block->loc); return std::nullopt; } - if (const auto* expr = value.as()) { + if (const auto* expr = value.as()) { TVM_FFI_THROW(InternalError) << "CreateFunction produced construction-only opaque artifact " << expr->GetTypeKey(); } - if (value.as()) { + if (value.as()) { TVM_FFI_THROW(InternalError) << "CreateFunction produced construction-only opaque artifact ir.OpaqueType"; } @@ -113,20 +113,20 @@ class TensorLoadToTensorTransformer : public s_tir::StmtExprMutator { if (!axis->dom.has_value()) return axis; auto domain = Mutate(axis->dom).ValueOrUnchanged(axis->dom).cast(); if (axis->dom.same_as(domain)) return axis; - auto copy = ffi::make_object(*axis.get()); + auto copy = ffi::make_object(*axis.get()); copy->dom = std::move(domain); return s_tir::IterVar(std::move(copy)); } - UnchangedOr Mutate_(const s_tir::IterVarNode* op, InplaceMode) { return op->var; } + UnchangedOr Mutate_(const s_tir::IterVarObj* op, InplaceMode) { return op->var; } - UnchangedOr Mutate_(const te::ReduceNode* op, InplaceMode inplace_mode) { - return Mutate_(static_cast(op), inplace_mode); + UnchangedOr Mutate_(const te::ReduceObj* op, InplaceMode inplace_mode) { + return Mutate_(static_cast(op), inplace_mode); } - UnchangedOr Mutate_(const OpaqueExprNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const OpaqueExprObj* op, InplaceMode inplace_mode) final { const auto* reduce = - op->IsInstance() ? static_cast(op) : nullptr; + op->IsInstance() ? static_cast(op) : nullptr; if (reduce == nullptr) { return s_tir::StmtExprMutator::Mutate_(op, inplace_mode); } @@ -154,7 +154,7 @@ class TensorLoadToTensorTransformer : public s_tir::StmtExprMutator { reduce->loc); } - UnchangedOr Mutate_(const CallNode* op, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const CallObj* op, InplaceMode inplace_mode) final { Call call = s_tir::StmtExprMutator::Mutate_(op, inplace_mode) .ValueOrUnchanged(ffi::GetRef(op)) .as_or_throw(); @@ -173,8 +173,8 @@ class TensorLoadToTensorTransformer : public s_tir::StmtExprMutator { static const VTable table = [] { VTable table; s_tir::StmtExprMutator::InitVTable(&table); - SetDispatch(&table); - SetDispatch(&table); + SetDispatch(&table); + SetDispatch(&table); table.Finalize(); return table; }(); @@ -187,8 +187,8 @@ class TensorLoadToTensorTransformer : public s_tir::StmtExprMutator { /*! \brief The helper mutator to rewrite tensor and tensor var accessed by block body */ class TensorSubstituter : public s_tir::StmtExprMutator { public: - explicit TensorSubstituter(const std::unordered_map& var_map, - const std::unordered_map& tensor_map) { + explicit TensorSubstituter(const std::unordered_map& var_map, + const std::unordered_map& tensor_map) { for (const auto& [source, target] : tensor_map) { VarRemapSet(ffi::AnyView(source), target); } @@ -227,7 +227,7 @@ class LayoutFreePlaceholdersNormalizer : public s_tir::StmtExprMutator { public: using s_tir::StmtExprMutator::Mutate; UnchangedOr Mutate(ffi::AnyView value, InplaceMode inplace_mode) final { - if (value.as()) return ffi::Unchanged(); + if (value.as()) return ffi::Unchanged(); return s_tir::StmtExprMutator::Mutate(value, inplace_mode); } @@ -237,7 +237,7 @@ class LayoutFreePlaceholdersNormalizer : public s_tir::StmtExprMutator { tensor2index_[tensor.value()] = i; } } - FunctionNode* f = func.CopyOnWrite(); + FunctionObj* f = func.CopyOnWrite(); f->body = Mutate(f->body, InplaceMode::kDisallow).ValueOrUnchanged(f->body); if (this->layout_free_tensor_indices_.empty()) { return func; @@ -250,11 +250,11 @@ class LayoutFreePlaceholdersNormalizer : public s_tir::StmtExprMutator { return WithAttr(std::move(func), tvm::s_tir::attr::kLayoutFreeTensors.name, indices); } - UnchangedOr Mutate_(const s_tir::SBlockNode* _block, InplaceMode inplace_mode) final { + UnchangedOr Mutate_(const s_tir::SBlockObj* _block, InplaceMode inplace_mode) final { s_tir::SBlock block = s_tir::StmtExprMutator::Mutate_(_block, inplace_mode) .ValueOrUnchanged(ffi::GetRef(_block)) .as_or_throw(); - s_tir::SBlockNode* n = block.CopyOnWrite(); + s_tir::SBlockObj* n = block.CopyOnWrite(); if (auto opt_ann = n->annotations.Get(topi_attr)) { ffi::Array new_tensors; for (TensorVar tensor : opt_ann.value().as_or_throw>()) { @@ -352,8 +352,8 @@ NestedIterLevels GenerateNestedIterLevels(const ffi::Array& axes ffi::Array GenerateOutputTensors(const te::ComputeOp& compute_op, CreateFuncInfo* info) { // Step 1. Collect output tensors in TE operation. ffi::Array tensors; - if (compute_op->body[0]->IsInstance()) { - auto f_reducer_equal = [](const te::ReduceNode* a, const te::ReduceNode* b) -> bool { + if (compute_op->body[0]->IsInstance()) { + auto f_reducer_equal = [](const te::ReduceObj* a, const te::ReduceObj* b) -> bool { ffi::StructuralEqual eq; return eq(a->combiner, b->combiner) && // eq(a->source, b->source) && // @@ -363,10 +363,10 @@ ffi::Array GenerateOutputTensors(const te::ComputeOp& compute_op, Cre }; PrimExpr expr_body = compute_op->body[0]; tensors.push_back(compute_op.output(0)); - const te::ReduceNode* reduce = expr_body.as(); + const te::ReduceObj* reduce = expr_body.as(); // specially handle reduction inline for multiplre reductions. for (size_t k = 1; k < compute_op->body.size(); ++k) { - const te::ReduceNode* reduce_ = compute_op->body[k].as(); + const te::ReduceObj* reduce_ = compute_op->body[k].as(); TVM_FFI_ICHECK(reduce_); TVM_FFI_ICHECK(f_reducer_equal(reduce_, reduce)) << "The Reduce inputs of ComputeOp should have the same attribute except value_index, " @@ -441,7 +441,7 @@ ffi::Map GenerateBlockAnnotations(const te::ComputeOp& co * \returns Init stmt. **/ SeqStmt GenerateInitStmt(const ffi::Array& indices, const ffi::Array& tensors, - const te::ReduceNode* reduce, const ffi::Map& var_map, + const te::ReduceObj* reduce, const ffi::Map& var_map, CreateFuncInfo* info) { auto f_substitute = [&var_map](const Var& var) -> ffi::Expected> { if (auto repl = var_map.Get(var)) return ffi::Any(*std::move(repl)); @@ -487,7 +487,7 @@ Stmt GenerateBodyStmt(const ffi::Array& indices, const ffi::Arraytransformer->Mutate(e).ValueOrUnchanged(e), f_substitute) .as_or_throw(); }; - if (const auto* reduce = expr_body.as()) { + if (const auto* reduce = expr_body.as()) { // Case 1. Reduce compute int n_tensors = tensors.size(); @@ -578,7 +578,7 @@ struct NestedScopeInfo { // helper to renew leaf block var defs to ensure SSA. void Renew(const ffi::Array& origin_axes) { block_iters.MutateByApply([](const s_tir::IterVar& itervar) { - auto n = ffi::make_object(*itervar.get()); + auto n = ffi::make_object(*itervar.get()); n->var = n->var.CopyWithSuffix(""); return s_tir::IterVar(n); }); @@ -672,7 +672,7 @@ Stmt GenerateStmtFromCompute(const te::ComputeOp& compute_op, CreateFuncInfo* in ffi::Array seq_stmt; auto leaf = scopes.back(); ffi::Map annotations = GenerateBlockAnnotations(compute_op, info); - const te::ReduceNode* reduce = compute_op->body[0].as(); + const te::ReduceObj* reduce = compute_op->body[0].as(); if (reduce) { PrimExpr expr_body = compute_op->body[0]; @@ -760,8 +760,8 @@ Stmt GenerateStmtFromCompute(const te::ComputeOp& compute_op, CreateFuncInfo* in Stmt GenerateStmtFromExternOp(const te::ExternOp& extern_op, CreateFuncInfo* info, sym::AnalyzerObj* analyzer) { // Step 1. Check all inputs are visited before and update var_map. - std::unordered_map var_map; - std::unordered_map input_tensor_map; + std::unordered_map var_map; + std::unordered_map input_tensor_map; TVM_FFI_ICHECK_EQ(extern_op->inputs.size(), extern_op->input_placeholders.size()); for (size_t i = 0; i < extern_op->inputs.size(); ++i) { const TensorVar& placeholder = extern_op->input_placeholders[i]; @@ -783,8 +783,8 @@ Stmt GenerateStmtFromExternOp(const te::ExternOp& extern_op, CreateFuncInfo* inf tvm::tirx::MatchShapeStridePattern(placeholder.ty()).value().get<2>().as()) { var_map.insert_or_assign(offset_var.value().get(), IntImm(offset_var.value().ty(), 0)); } - ffi::ObjectPtr type = - ffi::make_object(*output_tir_tensor.ty().get()); + ffi::ObjectPtr type = + ffi::make_object(*output_tir_tensor.ty().get()); if (type->layout.has_value()) { auto tile = type->layout.value().as_or_throw(); type->layout = tirx::TileLayout(tile->shard, tile->replica, {}); @@ -843,8 +843,8 @@ ffi::Array CollectOrderedOps(const ffi::Array& arg_li ffi::Array order = te::PostDFSOrder(arg_ops, g); for (const te::Operation& op : order) { - if (!(op->IsInstance() || op->IsInstance() || - op->IsInstance())) + if (!(op->IsInstance() || op->IsInstance() || + op->IsInstance())) TVM_FFI_THROW(InternalError) << "TypeError: Unsupported Operation: " << op->GetTypeKey() << ". " << "Only te.placeholder and te.compute are allowed for now."; @@ -856,7 +856,7 @@ void InitializeTensorBinds(const ffi::Array& ordered_ops, CreateF // Process any TE operations which contain user defined tensors for (const auto& op : ordered_ops) { // Initialize the tensor2tensor binds map with tensors defined by the te.extern - if (const auto* extern_op = op.as()) { + if (const auto* extern_op = op.as()) { TVM_FFI_ICHECK_EQ(extern_op->inputs.size(), extern_op->input_placeholders.size()); for (size_t i = 0; i < extern_op->inputs.size(); ++i) { const te::Tensor& input = extern_op->inputs[i]; @@ -870,7 +870,7 @@ void InitializeTensorBinds(const ffi::Array& ordered_ops, CreateF void RewriteStageToBlock(const te::Operation& op, CreateFuncInfo* info, ffi::Array* root_stmts, sym::AnalyzerObj* analyzer) { - if (const auto* placeholder = op.as()) { + if (const auto* placeholder = op.as()) { // Case 1. PlaceholderOp (te.placeholder) TVM_FFI_ICHECK_EQ(op->num_outputs(), 1); const te::Tensor& tensor = op.output(0); @@ -999,7 +999,7 @@ Function CreateFunction(const ffi::Array& arg_list, std::optional index_dtype_override) { ffi::Array tensor_arg_list; for (const ffi::ObjectRef& x : arg_list) { - if (auto tensor_node = x.as()) { + if (auto tensor_node = x.as()) { te::Tensor tensor = ffi::GetRef(tensor_node); tensor_arg_list.push_back(tensor); } diff --git a/src/te/operation/extern_op.cc b/src/te/operation/extern_op.cc index ea5e01789716..b101973c784c 100644 --- a/src/te/operation/extern_op.cc +++ b/src/te/operation/extern_op.cc @@ -31,15 +31,15 @@ namespace tvm { namespace te { using namespace tirx; -TVM_FFI_STATIC_INIT_BLOCK() { ExternOpNode::RegisterReflection(); } +TVM_FFI_STATIC_INIT_BLOCK() { ExternOpObj::RegisterReflection(); } // Pattern A (RM): auto-default repr from reflection. -int ExternOpNode::num_outputs() const { return static_cast(output_placeholders.size()); } +int ExternOpObj::num_outputs() const { return static_cast(output_placeholders.size()); } -PrimType ExternOpNode::output_dtype(size_t i) const { return output_placeholders[i].dtype(); } +PrimType ExternOpObj::output_dtype(size_t i) const { return output_placeholders[i].dtype(); } -ffi::Array ExternOpNode::output_shape(size_t i) const { +ffi::Array ExternOpObj::output_shape(size_t i) const { return output_placeholders[i].shape(); } @@ -50,7 +50,7 @@ ExternOp::ExternOp(std::string name, std::string tag, ffi::Map(); } - auto n = ffi::make_object(std::move(body)); + auto n = ffi::make_object(std::move(body)); n->name = std::move(name); n->tag = std::move(tag); n->attrs = std::move(attrs); @@ -63,8 +63,9 @@ ExternOp::ExternOp(std::string name, std::string tag, ffi::Maplayout; if (layout.has_value()) { - const auto* tile = layout.as(); - TVM_FFI_ICHECK(tile && tile->IsContiguous()) << "External inputs require contiguous strides"; + const auto* tile = layout.as(); + TVM_FFI_ICHECK(tile && ffi::GetRef(tile).IsContiguous()) + << "External inputs require contiguous strides"; } } n->inputs = std::move(inputs); @@ -85,7 +86,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { }); } -ffi::Array ExternOpNode::InputTensors() const { return inputs; } +ffi::Array ExternOpObj::InputTensors() const { return inputs; } } // namespace te } // namespace tvm diff --git a/src/te/operation/placeholder_op.cc b/src/te/operation/placeholder_op.cc index 52c54b04e5bc..6a5b43e22dc9 100644 --- a/src/te/operation/placeholder_op.cc +++ b/src/te/operation/placeholder_op.cc @@ -29,25 +29,25 @@ namespace tvm { namespace te { -TVM_FFI_STATIC_INIT_BLOCK() { PlaceholderOpNode::RegisterReflection(); } +TVM_FFI_STATIC_INIT_BLOCK() { PlaceholderOpObj::RegisterReflection(); } // Pattern A (RM): auto-default repr from reflection. -int PlaceholderOpNode::num_outputs() const { return 1; } +int PlaceholderOpObj::num_outputs() const { return 1; } -PrimType PlaceholderOpNode::output_dtype(size_t i) const { +PrimType PlaceholderOpObj::output_dtype(size_t i) const { TVM_FFI_ICHECK_EQ(i, 0U); return dtype; } -ffi::Array PlaceholderOpNode::output_shape(size_t i) const { +ffi::Array PlaceholderOpObj::output_shape(size_t i) const { TVM_FFI_ICHECK_EQ(i, 0U); return shape; } PlaceholderOp::PlaceholderOp(std::string name, ffi::Array shape, PrimType dtype) : Operation(ffi::UnsafeInit{}) { - auto n = ffi::make_object(); + auto n = ffi::make_object(); n->name = name; n->shape = shape; n->dtype = dtype; @@ -75,7 +75,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { }); } -ffi::Array PlaceholderOpNode::InputTensors() const { return {}; } +ffi::Array PlaceholderOpObj::InputTensors() const { return {}; } } // namespace te } // namespace tvm diff --git a/src/te/operation/reduce.cc b/src/te/operation/reduce.cc index 1726b5673315..fb10695144ee 100644 --- a/src/te/operation/reduce.cc +++ b/src/te/operation/reduce.cc @@ -34,8 +34,8 @@ namespace te { using namespace tvm::tirx; TVM_FFI_STATIC_INIT_BLOCK() { - CommReducerNode::RegisterReflection(); - ReduceNode::RegisterReflection(); + CommReducerObj::RegisterReflection(); + ReduceObj::RegisterReflection(); } // CommReducer @@ -55,7 +55,7 @@ CommReducer::CommReducer(ffi::Array lhs, ffi::Array rhs, // Change the dtype of input vars to adapt to the dtype of identities ffi::ArrayObj* p_lhs = lhs.CopyOnWrite(); ffi::ArrayObj* p_rhs = rhs.CopyOnWrite(); - std::unordered_map var_map; + std::unordered_map var_map; var_map.reserve(n_group * 2); for (int i = 0; i < static_cast(n_group); ++i) { PrimType dtype = identity_element[i].ty(); @@ -71,10 +71,10 @@ CommReducer::CommReducer(ffi::Array lhs, ffi::Array rhs, auto f_substitute = [&var_map](const Expr& expr) -> ffi::Expected> { // An axis can itself be a reducer parameter value. Substitute it before // structural traversal reaches the metadata's Var-typed binder field. - const VarNode* var = expr.as(); - if (const auto* axis = expr.as()) var = axis->var.get(); + const VarObj* var = expr.as(); + if (const auto* axis = expr.as()) var = axis->var.get(); if (auto it = var_map.find(var); it != var_map.end()) return ffi::Any(it->second); - if (expr.as()) return ffi::Any(ffi::GetRef(var)); + if (expr.as()) return ffi::Any(ffi::GetRef(var)); return ffi::Unchanged(); }; // The replacement variables intentionally adopt each identity element's dtype. @@ -84,7 +84,7 @@ CommReducer::CommReducer(ffi::Array lhs, ffi::Array rhs, .as_or_throw()); } - auto node = ffi::make_object(); + auto node = ffi::make_object(); node->lhs = lhs; node->rhs = rhs; node->result = result; @@ -93,8 +93,8 @@ CommReducer::CommReducer(ffi::Array lhs, ffi::Array rhs, data_ = std::move(node); } -ffi::Array CommReducerNode::operator()(ffi::Array a, - ffi::Array b) const { +ffi::Array CommReducerObj::operator()(ffi::Array a, + ffi::Array b) const { TVM_FFI_ICHECK_EQ(a.size(), b.size()); TVM_FFI_ICHECK_EQ(lhs.size(), a.size()); TVM_FFI_ICHECK_EQ(rhs.size(), b.size()); @@ -104,11 +104,11 @@ ffi::Array CommReducerNode::operator()(ffi::Array a, value_map.Set(rhs[i], b[i]); } auto f_substitute = [&value_map](const Expr& expr) -> ffi::Expected> { - const VarNode* var = expr.as(); - if (const auto* axis = expr.as()) var = axis->var.get(); + const VarObj* var = expr.as(); + if (const auto* axis = expr.as()) var = axis->var.get(); if (var != nullptr) { if (auto repl = value_map.Get(ffi::GetRef(var))) return ffi::Any(*std::move(repl)); - if (expr.as()) return ffi::Any(ffi::GetRef(var)); + if (expr.as()) return ffi::Any(ffi::GetRef(var)); } return ffi::Unchanged(); }; @@ -125,7 +125,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { [](ffi::Array lhs, ffi::Array rhs, ffi::Array result, ffi::Array identity_element, Location loc) { return CommReducer(lhs, rhs, result, identity_element, loc); }) - .def_method("te.CommReducerCombine", &te::CommReducerNode::operator()); + .def_method("te.CommReducerCombine", &te::CommReducerObj::operator()); } // Reduce @@ -137,7 +137,7 @@ Reduce::Reduce(CommReducer combiner, ffi::Array source, ffi::Arrayiter_type, kCommReduce) << "Can only take axis created by reduce_axis"; } - auto n = ffi::make_object(combiner, condition.value_or(IntImm::Bool(true))); + auto n = ffi::make_object(combiner, condition.value_or(IntImm::Bool(true))); TVM_FFI_ICHECK(source.defined()); for (size_t i = 0; i < axis.size(); ++i) { TVM_FFI_ICHECK(axis[i].defined()); @@ -149,13 +149,13 @@ Reduce::Reduce(CommReducer combiner, ffi::Array source, ffi::Array()); } else { - TVM_FFI_ICHECK(init[i]->IsInstance() || init[i]->IsInstance()) + TVM_FFI_ICHECK(init[i]->IsInstance() || init[i]->IsInstance()) << "init can only be an IntImm, FloatImm or Tensor-load Call, " << "but received " << init[i] << " of type " << init[i]->GetTypeKey(); } } } - n->ExprNode::ty = source[value_index].ty(); + n->ExprObj::ty = source[value_index].ty(); n->combiner = std::move(combiner); n->source = std::move(source); n->init = std::move(init); diff --git a/src/te/operation/scan_op.cc b/src/te/operation/scan_op.cc index 77e2471b142f..04ac3988c99f 100644 --- a/src/te/operation/scan_op.cc +++ b/src/te/operation/scan_op.cc @@ -32,15 +32,15 @@ using namespace tvm::prim; using namespace tirx; -TVM_FFI_STATIC_INIT_BLOCK() { ScanOpNode::RegisterReflection(); } +TVM_FFI_STATIC_INIT_BLOCK() { ScanOpObj::RegisterReflection(); } // Pattern A (RM): auto-default repr from reflection. -int ScanOpNode::num_outputs() const { return static_cast(update.size()); } +int ScanOpObj::num_outputs() const { return static_cast(update.size()); } -PrimType ScanOpNode::output_dtype(size_t i) const { return update[i]->GetDataType(); } +PrimType ScanOpObj::output_dtype(size_t i) const { return update[i]->GetDataType(); } -ffi::Array ScanOpNode::output_shape(size_t i) const { +ffi::Array ScanOpObj::output_shape(size_t i) const { TVM_FFI_ICHECK_LT(i, state_placeholder.size()); return state_placeholder[i]->shape; } @@ -53,7 +53,7 @@ ScanOp::ScanOp(std::string name, std::string tag, if (!attrs.has_value()) { attrs = ffi::Map(); } - auto n = ffi::make_object(); + auto n = ffi::make_object(); TVM_FFI_ICHECK_EQ(init.size(), update.size()); TVM_FFI_ICHECK_EQ(init.size(), state_placeholder.size()); sym::Analyzer analyzer; @@ -126,7 +126,7 @@ ffi::Array scan(ffi::Array init, ffi::Array update, return res; } -ffi::Array ScanOpNode::InputTensors() const { +ffi::Array ScanOpObj::InputTensors() const { ffi::Array ret; for (Tensor t : init) { ret.push_back(t); diff --git a/src/te/script/printer/expr.cc b/src/te/script/printer/expr.cc index dbc0e5030257..492145571a67 100644 --- a/src/te/script/printer/expr.cc +++ b/src/te/script/printer/expr.cc @@ -23,14 +23,14 @@ namespace tvm::script::printer::details { TVM_FFI_STATIC_INIT_BLOCK() { - RegisterScriptRepr(); - RegisterScriptRepr(); + RegisterScriptRepr(); + RegisterScriptRepr(); } namespace { ffi::Optional CommReducerDocTranslate(DocTranslatorObj* d, ffi::AnyView input, const ffi::Object*) { const auto* reducer = - ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(input); + ffi::detail::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(input); TVM_FFI_CHECK(reducer->lhs.size() == reducer->rhs.size(), TypeError) << "printer comm reducer arguments must be paired"; ffi::Array args; @@ -50,7 +50,7 @@ ffi::Optional CommReducerDocTranslate(DocTranslatorObj* d, ffi::AnyView } TVM_FFI_STATIC_INIT_BLOCK() { - ffi::reflection::TypeAttrDef().attr( + ffi::reflection::TypeAttrDef().attr( tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&CommReducerDocTranslate>()); } @@ -58,7 +58,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { ffi::Optional ReduceDocTranslate(DocTranslatorObj* d, ffi::AnyView input, const ffi::Object*) { const auto* node = - ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(input); + ffi::detail::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(input); // Reduction axes bind the source and predicate. Their declarations must use // the same identifiers as the IterVars, without becoming external captures. for (const s_tir::IterVar& axis : node->axis) { @@ -85,7 +85,7 @@ ffi::Optional ReduceDocTranslate(DocTranslatorObj* d, ffi::AnyView inpu } TVM_FFI_STATIC_INIT_BLOCK() { - ffi::reflection::TypeAttrDef().attr( + ffi::reflection::TypeAttrDef().attr( tvm::script::printer::type_attr::kDocTranslate.name, FDocTranslate::FromNative<&ReduceDocTranslate>()); } diff --git a/src/te/tensor.cc b/src/te/tensor.cc index 01168cf97bf7..7dcdbe4c4110 100644 --- a/src/te/tensor.cc +++ b/src/te/tensor.cc @@ -43,12 +43,12 @@ const Op& TensorLoadOp() { ffi::Array ValidateTensorLoad(const Call& call, Tensor* tensor_out) { TVM_FFI_ICHECK(call->op.same_as(TensorLoadOp())) << "Expected a te.tensor_load Call"; TVM_FFI_ICHECK(!call->args.empty()) << "Tensor-load Call requires a Tensor argument"; - const auto* tensor_node = call->args[0].as(); + const auto* tensor_node = call->args[0].as(); TVM_FFI_ICHECK(tensor_node != nullptr) << "Tensor-load first argument must be a TE Tensor"; Tensor tensor = ffi::GetRef(tensor_node); TVM_FFI_ICHECK_EQ(call->args.size() - 1, tensor->shape.size()) << "Tensor-load index count must match tensor rank"; - TVM_FFI_ICHECK(call->ty.as() != nullptr && call->ty == tensor->dtype) + TVM_FFI_ICHECK(call->ty.as() != nullptr && call->ty == tensor->dtype) << "Tensor-load result type must match the tensor element type"; ffi::Array indices; @@ -81,13 +81,13 @@ TVM_FFI_INLINE ffi::Expected> TensorMaybeInplaceMutat } // namespace -void TensorNode::RegisterReflection() { +void TensorObj::RegisterReflection() { namespace refl = tvm::ffi::reflection; - refl::ObjectDef() - .def_ro("shape", &TensorNode::shape) - .def_ro("dtype", &TensorNode::dtype) - .def_ro("op", &TensorNode::op) - .def_ro("value_index", &TensorNode::value_index); + refl::ObjectDef() + .def_ro("shape", &TensorObj::shape) + .def_ro("dtype", &TensorObj::dtype) + .def_ro("op", &TensorObj::op) + .def_ro("value_index", &TensorObj::value_index); } IterVar thread_axis(Range dom, std::string tag) { @@ -129,7 +129,7 @@ inline PrimExpr Tensor::IndexTensor(ffi::Array indices, for (const PrimExpr& index : indices) { args.push_back(index); } - return PrimExpr(Call((*this)->dtype, TensorLoadOp(), args)); + return Call((*this)->dtype, TensorLoadOp(), args).as_or_throw(); } PrimExpr Tensor::operator()(ffi::Array indices) const { @@ -152,7 +152,7 @@ PrimExpr Tensor::IndexWithNegativeIndices(ffi::Array indices) const { return IndexTensor(indices, true); } -ffi::String TensorNode::GetNameHint() const { +ffi::String TensorObj::GetNameHint() const { return op.value()->num_outputs() == 1 ? op.value()->name : (op.value()->name + ".v" + std::to_string(value_index)); } @@ -164,8 +164,8 @@ Tensor Operation::output(size_t i) const { Tensor::Tensor(ffi::Array shape, PrimType dtype, ffi::Optional op, int value_index) : OpaqueExpr(ffi::UnsafeInit{}) { - auto n = ffi::make_object(); - n->ExprNode::ty = OpaqueType(); + auto n = ffi::make_object(); + n->ExprObj::ty = OpaqueType(); n->shape = std::move(shape); n->dtype = dtype; n->op = op; @@ -175,8 +175,8 @@ Tensor::Tensor(ffi::Array shape, PrimType dtype, ffi::Optional() + TensorObj::RegisterReflection(); + refl::TypeAttrDef() .attr(ffi::type_attr::kStructuralVisit, ffi::FStructuralVisit::FromNative<&TensorVisit>()) .attr(ffi::type_attr::kStructuralMutate, ffi::FStructuralMutate::FromNative<&TensorMutate>()) .attr(ffi::type_attr::kStructuralMaybeInplaceMutate, @@ -184,7 +184,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { } bool IsTensorLoad(const Expr& expr) { - const auto* call = expr.as(); + const auto* call = expr.as(); return call != nullptr && call->op.same_as(TensorLoadOp()); } @@ -219,8 +219,8 @@ TVM_FFI_STATIC_INIT_BLOCK() { }) .def("te.OpGetOutput", [](Operation op, int64_t output) { return op.output(static_cast(output)); }) - .def_method("te.OpNumOutputs", &OperationNode::num_outputs) - .def_method("te.OpInputTensors", &OperationNode::InputTensors); + .def_method("te.OpNumOutputs", &OperationObj::num_outputs) + .def_method("te.OpInputTensors", &OperationObj::InputTensors); } } // namespace te diff --git a/src/topi/einsum.cc b/src/topi/einsum.cc index 308bce6a0b63..1e96e90f8dbe 100644 --- a/src/topi/einsum.cc +++ b/src/topi/einsum.cc @@ -104,8 +104,8 @@ EinsumEquation EinsumEquation::FromString(const std::string& equation) { } PrimExpr GetBroadcastedExtent(const PrimExpr& extent1, const PrimExpr& extent2) { - const IntImmNode* extent1_imm = extent1.as(); - const IntImmNode* extent2_imm = extent2.as(); + const IntImmObj* extent1_imm = extent1.as(); + const IntImmObj* extent2_imm = extent2.as(); if (extent1_imm != nullptr && extent2_imm != nullptr) { if (extent1_imm->value == extent2_imm->value) { return extent1; diff --git a/tests/python/relax/test_expr_functor.py b/tests/python/relax/test_expr_functor.py index eb114a8bf07b..407c1f80b37b 100644 --- a/tests/python/relax/test_expr_functor.py +++ b/tests/python/relax/test_expr_functor.py @@ -282,7 +282,7 @@ def visit_seq_expr_(self, op: SeqExpr) -> Expr: return op def visit_var_binding_(self, binding: VarBinding) -> None: - """Identical with ExprMutator::VisitBinding_(const VarBindingNode* binding) on the C++ side.""" + """Identical with ExprMutator::VisitBinding_(const VarBindingObj* binding) on the C++ side.""" new_value = self.visit_expr(binding.value) new_var = self.visit_var_def(binding.var) @@ -299,7 +299,7 @@ def visit_var_binding_(self, binding: VarBinding) -> None: self.builder_.emit_normalized(VarBinding(new_var, new_value)) def visit_match_cast_(self, binding: MatchCast) -> None: - """Identical with ExprMutator::VisitBinding_(const MatchCastNode* binding) on the C++ side.""" + """Identical with ExprMutator::VisitBinding_(const MatchCastObj* binding) on the C++ side.""" new_var = self.visit_var_def(binding.var) new_value = self.visit_expr(binding.value) @@ -312,7 +312,7 @@ def visit_match_cast_(self, binding: MatchCast) -> None: self.builder_.emit_normalized(MatchCast(new_var, new_value, binding.ty)) def visit_binding_block_(self, block: BindingBlock) -> BindingBlock: - """Identical with ExprMutator::VisitBindingBlock_(const BindingBlockNode* block) on the C++ side.""" + """Identical with ExprMutator::VisitBindingBlock_(const BindingBlockObj* block) on the C++ side.""" self.builder_._begin_binding_block() for binding in block.bindings: self.visit_binding(binding) @@ -320,7 +320,7 @@ def visit_binding_block_(self, block: BindingBlock) -> BindingBlock: return self.builder_._end_block() def visit_dataflow_block_(self, block: DataflowBlock) -> None: - """Identical with ExprMutator::VisitBindingBlock_(const DataflowBlockNode* block) on the C++ side.""" + """Identical with ExprMutator::VisitBindingBlock_(const DataflowBlockObj* block) on the C++ side.""" self.builder_._begin_dataflow_block() for binding in block.bindings: self.visit_binding(binding) @@ -328,12 +328,12 @@ def visit_dataflow_block_(self, block: DataflowBlock) -> None: return self.builder_._end_block() def visit_var_def_(self, var: Var) -> None: - """Identical with ExprMutator::VisitVarDef_(const VarNode* var) on the C++ side.""" + """Identical with ExprMutator::VisitVarDef_(const VarObj* var) on the C++ side.""" self.log.add("VarDef") return var def visit_dataflow_var_def_(self, var: DataflowVar) -> None: - """Identical with ExprMutator::VisitVarDef_(const DataflowVarNode* var) on the C++ side.""" + """Identical with ExprMutator::VisitVarDef_(const DataflowVarObj* var) on the C++ side.""" self.log.add("DataflowVarDef") return var diff --git a/tests/python/relax/test_op_index.py b/tests/python/relax/test_op_index.py index c2a38eb656ce..4a4b30359262 100644 --- a/tests/python/relax/test_op_index.py +++ b/tests/python/relax/test_op_index.py @@ -564,7 +564,7 @@ def test_strided_slice_infer_ty_symbolic_begin_end_strides(): bb, relax.op.strided_slice(x, axes=[0], begin=[0], end=[8], strides=[var]), relax.TensorType( - [tirx.if_then_else(var < 0, -8 // (0 - var) + 1, (var + 7) // var), 9], + [tirx.if_then_else(var < 0, -8 // (0 - var) + 1, 7 // var + 1), 9], dtype="float32", ), ) diff --git a/tests/python/relax/test_transform_bind_params.py b/tests/python/relax/test_transform_bind_params.py index 5e1fb89d972a..ca5dbd55c210 100644 --- a/tests/python/relax/test_transform_bind_params.py +++ b/tests/python/relax/test_transform_bind_params.py @@ -121,7 +121,7 @@ def main( params_dict = {"w0": w0_tvm, "b0": b0_tvm, "w1": w1_tvm, "b1": b1_tvm} mod = relax.transform.BindParams("main", params_dict)(Before) - # Since it contains ConstantNode, it's hard to check with structural equality. + # Since it contains ConstantObj, it's hard to check with structural equality. func = mod["main"] assert len(func.params) == 1 batch = func.params[0].ty.shape[0] diff --git a/tests/python/relax/test_transform_merge_composite_functions.py b/tests/python/relax/test_transform_merge_composite_functions.py index 7104895291e6..23f8e189e348 100644 --- a/tests/python/relax/test_transform_merge_composite_functions.py +++ b/tests/python/relax/test_transform_merge_composite_functions.py @@ -1015,7 +1015,7 @@ def test_mixed_non_composite(): def test_reshape(): - # Verify that the non-CallNode input (shape in reshape) can be handled properly. + # Verify that the non-CallObj input (shape in reshape) can be handled properly. @I.ir_module class Module: @R.function(private=True) diff --git a/tests/python/relax/test_transform_normalize.py b/tests/python/relax/test_transform_normalize.py index 14a3c096b7a4..85fde63118f2 100644 --- a/tests/python/relax/test_transform_normalize.py +++ b/tests/python/relax/test_transform_normalize.py @@ -57,7 +57,7 @@ def expected(x: R.Tensor((m, n), "float16")) -> R.Tensor(dtype="float16", ndim=2 def test_normalize_if(): cond = relax.Var("cond", R.Tensor([], "bool")) x = relax.Var("x", R.Tensor([1], "float32")) - # TODO(relax-team): add type and shape inference for IfExprNode + # TODO(relax-team): add type and shape inference for IfExprObj y = relax.Var("y") # Note: the parser automatically normalize the IR written in TVMScript, @@ -237,7 +237,7 @@ def expected( def test_normalize_if_condition(): cond = relax.Var("cond", R.Tensor([], "bool")) x = relax.Var("x", R.Tensor([1], "float32")) - # TODO(relax-team): add type and shape inference for IfExprNode + # TODO(relax-team): add type and shape inference for IfExprObj y = relax.Var("y") # The condition is wrapped in a tuple and then indexed diff --git a/tests/python/relax/test_utils.py b/tests/python/relax/test_utils.py index a8b22fdb872e..f75aea9120ce 100644 --- a/tests/python/relax/test_utils.py +++ b/tests/python/relax/test_utils.py @@ -82,7 +82,7 @@ def func_copied(x: R.Tensor((3,), "float32"), y: R.Tensor((3,), "float32")): "global_symbol", "func_copied" ) - # Assertion will fail if the f_copied contains the same VarNode that's used in + # Assertion will fail if the f_copied contains the same VarObj that's used in # the original function, due to var mapping during structural equal. assert_structural_equal(Actual, Expected) diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_feature_extractor_per_store_feature.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_feature_extractor_per_store_feature.py index 0c0796e11bb8..6de9dfede784 100644 --- a/tests/python/s_tir/meta_schedule/test_meta_schedule_feature_extractor_per_store_feature.py +++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_feature_extractor_per_store_feature.py @@ -851,6 +851,9 @@ def _create_schedule(): rtol=1e-5, atol=1e-5, ) + # Region estimates replay each loop-range context, without retaining the + # temporary constant bindings used while relaxing inner loops. In particular, + # the full 512x512 float32 input contributes 2**20 unique bytes. # Group 2.1: Tensor A assert_allclose( actual=f[57:75], @@ -859,19 +862,19 @@ def _create_schedule(): 0.0, 0.0, 25.00000004, - 19.99718086, + 20.000001375860553, 23.00000017, - 13.99726771, + 14.000088052430122, 1.0, 0.0, 0.0, 18.0000055, - 20.00000138, + 20.00562591970089, 2.32192809, 23.00000017, - 17.997185, + 18.00000550343433, 21.00000069, - 11.99753235, + 12.0003521774803, 12.00035218, ], rtol=1e-5, @@ -885,9 +888,9 @@ def _create_schedule(): 1.0, 0.0, 25.00000004, - 11.00070427, + 12.0003521774803, 23.00000017, - 5.04439412, + 9.002815015607053, 1.0, 0.0, 0.0, @@ -895,9 +898,9 @@ def _create_schedule(): 11.98049664, 8.00562455, 17.00001101, - 3.169925, + 4.087462841250339, 15.00004403, - 0.169925, + 1.584962500721156, 4.08746284, ], rtol=1e-5, @@ -1065,19 +1068,19 @@ def _create_schedule(): 0.0, 0.0, 22.00000034, - 19.85798251, + 20.000001375860553, 20.00000138, - 13.85807816, + 14.000088052430122, 1.0, 0.0, 0.0, 15.00004403, - 20.04456622, + 20.17555076886471, 2.32192809, 20.00000138, - 17.85798707, + 18.00000550343433, 18.0000055, - 11.8583696, + 12.0003521774803, 4.08746284, ], rtol=1e-5, @@ -1091,19 +1094,19 @@ def _create_schedule(): 1.0, 0.0, 22.00000034, - 7.01122726, + 9.002815015607053, 20.00000138, - 4.08746284, + 3.169925001442312, 1.0, 0.0, 0.0, 3.169925001442312, - 4.08746284, + 9.61654884377899, 8.005624549193879, 14.000088052430122, - 0.5849625, + 1.584962500721156, 12.00035218, - 0.08746284, + 0.044394119358453436, 4.087462841250339, ], rtol=1e-5, diff --git a/tests/scripts/task_build_arm.sh b/tests/scripts/task_build_arm.sh index 824fb347a6e1..1d1c132d24e4 100755 --- a/tests/scripts/task_build_arm.sh +++ b/tests/scripts/task_build_arm.sh @@ -37,7 +37,7 @@ export CMAKE_ARGS=" -DTVM_USE_GTEST=ON -DBUILD_TESTING=ON -DTVM_TIRX_BUILD_TESTS=ON - -DTVM_TIRX_USE_LLVM=llvm-config-17 + -DTVM_TIRX_USE_LLVM=ON " uv build --wheel --no-build-isolation --offline --python python3 --out-dir dist ./3rdparty/tvm-tirx/3rdparty/tvm-ffi uv build --wheel --no-build-isolation --offline --python python3 --out-dir dist ./3rdparty/tvm-tirx diff --git a/tests/scripts/task_build_cpu.sh b/tests/scripts/task_build_cpu.sh index 5613c4602642..be42ff6c72cd 100755 --- a/tests/scripts/task_build_cpu.sh +++ b/tests/scripts/task_build_cpu.sh @@ -40,7 +40,7 @@ export CMAKE_ARGS=" -DTVM_USE_GTEST=ON -DBUILD_TESTING=ON -DTVM_TIRX_BUILD_TESTS=ON - -DTVM_TIRX_USE_LLVM='/usr/bin/llvm-config-17 --link-static' + -DTVM_TIRX_USE_LLVM=ON " uv build --wheel --no-build-isolation --offline --python python3 --out-dir dist ./3rdparty/tvm-tirx/3rdparty/tvm-ffi uv build --wheel --no-build-isolation --offline --python python3 --out-dir dist ./3rdparty/tvm-tirx diff --git a/tests/scripts/task_build_gpu.sh b/tests/scripts/task_build_gpu.sh index a4b836b442eb..d8cc6737c6a0 100755 --- a/tests/scripts/task_build_gpu.sh +++ b/tests/scripts/task_build_gpu.sh @@ -47,7 +47,7 @@ export CMAKE_ARGS=" -DTVM_TIRX_BUILD_TESTS=OFF -DTVM_TIRX_USE_VULKAN=OFF -DTVM_TIRX_USE_OPENCL=OFF - -DTVM_TIRX_USE_LLVM='/usr/bin/llvm-config-15 --link-static' + -DTVM_TIRX_USE_LLVM=ON " uv build --wheel --no-build-isolation --offline --python python3 --out-dir dist ./3rdparty/tvm-tirx/3rdparty/tvm-ffi uv build --wheel --no-build-isolation --offline --python python3 --out-dir dist ./3rdparty/tvm-tirx diff --git a/tests/scripts/task_build_gpu_other.sh b/tests/scripts/task_build_gpu_other.sh index a70262d6fea9..3f948e0a605f 100755 --- a/tests/scripts/task_build_gpu_other.sh +++ b/tests/scripts/task_build_gpu_other.sh @@ -24,9 +24,9 @@ rm -f 3rdparty/tvm-tirx/3rdparty/tvm-ffi/build/CMakeCache.txt # Keep these libraries separate from the published GPU wheel artifacts. uv pip install --python python3 --no-deps --target build/python ./3rdparty/tvm-tirx/3rdparty/tvm-ffi -export PYTHONPATH="$PWD/build/python" +export PYTHONPATH="$PWD/build/python:$PWD/.local/python" cmake -S . -B build -GNinja -DCMAKE_BUILD_TYPE=RelWithDebInfo \ -DTVM_TIRX_USE_OPENCL=ON -DTVM_TIRX_USE_VULKAN=ON -DTVM_TIRX_USE_ROCM=ON \ - -DTVM_USE_CLML=ON -DTVM_TIRX_USE_LLVM='/usr/bin/llvm-config-15 --link-static' \ + -DTVM_USE_CLML=ON -DTVM_TIRX_USE_LLVM=ON \ -DCMAKE_CXX_FLAGS=-Werror -DTVM_USE_CCACHE=OFF cmake --build build --target tvm_compiler tvm_runtime_opencl tvm_runtime_vulkan tvm_runtime_rocm diff --git a/tests/scripts/task_build_wasm.sh b/tests/scripts/task_build_wasm.sh index 2536b57adf7e..017b8b5dd6f9 100755 --- a/tests/scripts/task_build_wasm.sh +++ b/tests/scripts/task_build_wasm.sh @@ -30,7 +30,7 @@ export CMAKE_ARGS=" -DTVM_FFI_BUILD_TESTS=OFF -DTVM_FFI_ATTACH_DEBUG_SYMBOLS=ON -DTVM_USE_SORT=ON - -DTVM_TIRX_USE_LLVM=llvm-config-15 + -DTVM_TIRX_USE_LLVM=ON -DCMAKE_CXX_FLAGS=-Werror -DTVM_HIDE_PRIVATE_SYMBOLS=ON -DTVM_USE_CCACHE=OFF diff --git a/tests/scripts/task_install_build_dependencies.sh b/tests/scripts/task_install_build_dependencies.sh index 3567e5ef7eef..177dfe6742db 100755 --- a/tests/scripts/task_install_build_dependencies.sh +++ b/tests/scripts/task_install_build_dependencies.sh @@ -23,5 +23,6 @@ rm -f 3rdparty/tvm-tirx/3rdparty/tvm-ffi/build/CMakeCache.txt rm -rf dist uv pip install --python python3 --target .local/python \ - "scikit-build-core>=0.11,!=1.1.0" "setuptools-scm>=8" "cython>=3.2.8" "mlc-z3-static>=4.16.0" + "scikit-build-core>=0.11,!=1.1.0" "setuptools-scm>=8" "cython>=3.2.8" "mlc-z3-static>=4.16.0" \ + "mlc-llvm-static[all]==23.1.1.post0" export PYTHONPATH="$PWD/.local/python" diff --git a/web/emcc/tvmjs_support.cc b/web/emcc/tvmjs_support.cc index 8183e3200feb..e04fd755e863 100644 --- a/web/emcc/tvmjs_support.cc +++ b/web/emcc/tvmjs_support.cc @@ -102,7 +102,7 @@ int TVMFFIWasmFunctionCreate(void* self, TVMFFIObjectHandle* out) { const char* TVMFFIWasmGetLastError() { static thread_local std::string last_error; - last_error = ::tvm::ffi::details::MoveFromSafeCallRaised().what(); + last_error = ::tvm::ffi::detail::MoveFromSafeCallRaised().what(); return last_error.c_str(); } @@ -119,10 +119,10 @@ class AsyncLocalSession : public LocalSession { if (name == "runtime.RPCTimeEvaluator") { return get_time_eval_placeholder_.get(); } else if (auto fp = tvm::ffi::Function::GetGlobal(name)) { - TVMFFIAny val = tvm::ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(tvm::ffi::Any(*fp)); + TVMFFIAny val = tvm::ffi::detail::AnyUnsafe::MoveAnyToTVMFFIAny(tvm::ffi::Any(*fp)); return val.v_obj; } else if (auto fp = tvm::ffi::Function::GetGlobal("__async." + name)) { - TVMFFIAny val = tvm::ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(tvm::ffi::Any(*fp)); + TVMFFIAny val = tvm::ffi::detail::AnyUnsafe::MoveAnyToTVMFFIAny(tvm::ffi::Any(*fp)); async_func_set_.insert(val.v_obj); return val.v_obj; } else { diff --git a/web/emcc/wasm_runtime.cc b/web/emcc/wasm_runtime.cc index bf67e1384c5a..3d62bc9ecd5a 100644 --- a/web/emcc/wasm_runtime.cc +++ b/web/emcc/wasm_runtime.cc @@ -51,7 +51,6 @@ #include "3rdparty/tvm-tirx/src/runtime/logging.cc" #include "3rdparty/tvm-tirx/src/runtime/metadata.cc" #include "3rdparty/tvm-tirx/src/runtime/tensor.cc" -#include "3rdparty/tvm-tirx/src/runtime/workspace_pool.cc" // relax setup #include "3rdparty/tvm-tirx/3rdparty/tvm-ffi/src/ffi/backtrace.cc" #include "3rdparty/tvm-tirx/3rdparty/tvm-ffi/src/ffi/container.cc"