This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new e0ed4aad41 Simplify instruction kind and script repr registration
(#20472)
e0ed4aad41 is described below
commit e0ed4aad411c69019d13956d9aa3ae3f452d6d7c
Author: Tianqi Chen <[email protected]>
AuthorDate: Sat Sep 26 13:27:39 2026 -0400
Simplify instruction kind and script repr registration (#20472)
---
include/tvm/s_tir/schedule/instruction.h | 86 ++++++++--------------
include/tvm/script/printer/config.h | 2 +-
include/tvm/script/printer/printer.h | 43 ++++++-----
src/ir/expr.cc | 2 +-
src/ir/prim/expr.cc | 2 +-
src/relax/script/printer/binding.cc | 8 +-
src/relax/script/printer/call.cc | 2 +-
src/relax/script/printer/dependent_type.cc | 10 ++-
src/relax/script/printer/distributed.cc | 8 +-
src/relax/script/printer/expr.cc | 16 ++--
src/relax/script/printer/function.cc | 6 +-
src/relax/script/printer/region.cc | 8 +-
src/relax/script/printer/type.cc | 4 +-
src/s_tir/schedule/instruction.cc | 44 ++++++++---
src/s_tir/schedule/instruction_traits.h | 58 ++++-----------
src/s_tir/schedule/primitive/annotate.cc | 6 +-
.../schedule/primitive/annotate_buffer_access.cc | 2 +-
src/s_tir/schedule/primitive/block_annotate.cc | 8 +-
src/s_tir/schedule/primitive/blockize_tensorize.cc | 6 +-
src/s_tir/schedule/primitive/cache_index.cc | 2 +-
src/s_tir/schedule/primitive/cache_read_write.cc | 14 ++--
src/s_tir/schedule/primitive/compute_at.cc | 6 +-
src/s_tir/schedule/primitive/compute_inline.cc | 8 +-
src/s_tir/schedule/primitive/decompose_padding.cc | 2 +-
src/s_tir/schedule/primitive/for_kind.cc | 10 ++-
src/s_tir/schedule/primitive/get_block_loop.cc | 14 ++--
src/s_tir/schedule/primitive/hide_buffer_access.cc | 2 +-
.../schedule/primitive/layout_transformation.cc | 6 +-
.../schedule/primitive/loop_transformation.cc | 14 ++--
src/s_tir/schedule/primitive/pad_einsum.cc | 2 +-
src/s_tir/schedule/primitive/read_write_at.cc | 6 +-
src/s_tir/schedule/primitive/reduction.cc | 6 +-
.../schedule/primitive/reorder_block_iter_var.cc | 2 +-
src/s_tir/schedule/primitive/rolling_buffer.cc | 2 +-
src/s_tir/schedule/primitive/sampling.cc | 10 ++-
src/s_tir/schedule/trace.cc | 2 +-
src/script/printer/ir/ir.cc | 12 +--
src/tirx/script/printer/block.cc | 12 +--
src/tirx/script/printer/buffer.cc | 18 +++--
src/tirx/script/printer/expr.cc | 74 ++++++++++---------
src/tirx/script/printer/for_loop.cc | 2 +-
src/tirx/script/printer/function.cc | 2 +-
src/tirx/script/printer/ir.cc | 12 +--
src/tirx/script/printer/stmt.cc | 30 ++++----
44 files changed, 310 insertions(+), 281 deletions(-)
diff --git a/include/tvm/s_tir/schedule/instruction.h
b/include/tvm/s_tir/schedule/instruction.h
index 10a14e1468..7412090745 100644
--- a/include/tvm/s_tir/schedule/instruction.h
+++ b/include/tvm/s_tir/schedule/instruction.h
@@ -26,10 +26,6 @@
namespace tvm {
-// Forward declaration
-template <typename, typename>
-class AttrRegistry;
-
namespace s_tir {
// Forward declaration
@@ -201,85 +197,65 @@ class Instruction : public ffi::ObjectRef {
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Instruction, ffi::ObjectRef,
InstructionNode);
};
-/*!
- * \brief A helper macro to register InstructionKind, only used in
`TVM_REGISTER_INST_KIND`
- * \note This macro is not user-facing.
- * \sa TVM_REGISTER_INST_KIND
- */
-#define TVM_INST_KIND_REGISTER_VAR_DEF \
- [[maybe_unused]] static ::tvm::s_tir::InstructionKindRegEntry&
__make_##InstructionKind
-
-/*!
- * \brief Register an InstructionKind
- * \param InstructionKindName The name of the InstructionKind
- *
- * Example:
- *
- * \code
- *
- * TVM_REGISTER_INST_KIND("ComputeInline")
- * .set_is_pure(false)
- * .set_apply_to_schedule(ApplyToSchedule)
- * .set_attrs_as_json(AttrsAsJSON)
- * .set_attrs_from_json(AttrsFromJSON)
- * .set_as_python(AsPython);
- *
- * \endcode
- */
-#define TVM_REGISTER_INST_KIND(InstructionKindName) \
- TVM_FFI_STR_CONCAT(TVM_INST_KIND_REGISTER_VAR_DEF, __COUNTER__) = \
-
::tvm::s_tir::InstructionKindRegEntry::RegisterOrGet(InstructionKindName).set_name()
-
-/*! \brief An entry in the registry of InstructionKind */
-class InstructionKindRegEntry {
+/*! \brief Temporary builder for a canonical named instruction kind. */
+class InstructionKindDef {
public:
- static InstructionKindRegEntry& RegisterOrGet(const ffi::String& name);
-
- InstructionKindRegEntry& set_name() {
- get_mutable()->name = this->name;
- return *this;
- }
+ /*! \brief Find or create the canonical kind and assign its name.
+ * \param name The instruction kind name.
+ */
+ explicit InstructionKindDef(const ffi::String& name);
- InstructionKindRegEntry& set_is_pure(bool is_pure) {
+ /*! \brief Set whether removing this instruction leaves schedule state
unchanged.
+ * \param is_pure Whether the instruction is pure.
+ * \return This builder for chaining.
+ */
+ InstructionKindDef& set_is_pure(bool is_pure) {
get_mutable()->is_pure = is_pure;
return *this;
}
- InstructionKindRegEntry& set_apply_to_schedule(FInstructionApply
f_apply_to_schedule) {
+ /*! \brief Set the schedule application callback.
+ * \param f_apply_to_schedule Callback that applies the instruction.
+ * \return This builder for chaining.
+ */
+ InstructionKindDef& set_apply_to_schedule(FInstructionApply
f_apply_to_schedule) {
get_mutable()->f_apply_to_schedule = std::move(f_apply_to_schedule);
return *this;
}
- InstructionKindRegEntry& set_as_python(FInstructionAsPython f_as_python) {
+ /*! \brief Set the Python trace formatting callback.
+ * \param f_as_python Callback that formats an instruction.
+ * \return This builder for chaining.
+ */
+ InstructionKindDef& set_as_python(FInstructionAsPython f_as_python) {
get_mutable()->f_as_python = std::move(f_as_python);
return *this;
}
- InstructionKindRegEntry& set_attrs_as_json(FInstructionAttrsAsJSON
f_attrs_as_json) {
+ /*! \brief Set the optional JSON serialization callback.
+ * \param f_attrs_as_json Callback, or null when conversion is unnecessary.
+ * \return This builder for chaining.
+ */
+ InstructionKindDef& set_attrs_as_json(FInstructionAttrsAsJSON
f_attrs_as_json) {
get_mutable()->f_attrs_as_json = std::move(f_attrs_as_json);
return *this;
}
- InstructionKindRegEntry& set_attrs_from_json(FInstructionAttrsFromJSON
f_attrs_from_json) {
+ /*! \brief Set the optional JSON deserialization callback.
+ * \param f_attrs_from_json Callback, or null when conversion is unnecessary.
+ * \return This builder for chaining.
+ */
+ InstructionKindDef& set_attrs_from_json(FInstructionAttrsFromJSON
f_attrs_from_json) {
get_mutable()->f_attrs_from_json = std::move(f_attrs_from_json);
return *this;
}
private:
- /*! \brief Private constructor, used only by AttrRegistry */
- explicit InstructionKindRegEntry(uint32_t reg_index);
- /*! \brief Get the mutable reference to the internal InstructionKind */
InstructionKindNode* get_mutable() const {
return const_cast<InstructionKindNode*>(inst_kind_.get());
}
- /*! \brief The name of the registry entry */
- ffi::String name;
- /*! \brief The instruction kind */
InstructionKind inst_kind_;
- template <typename, typename>
- friend class ::tvm::AttrRegistry;
- friend class InstructionKind;
};
} // namespace s_tir
diff --git a/include/tvm/script/printer/config.h
b/include/tvm/script/printer/config.h
index 7f45b34cd5..02bd96bba7 100644
--- a/include/tvm/script/printer/config.h
+++ b/include/tvm/script/printer/config.h
@@ -152,7 +152,7 @@ class TVM_DLL PrinterConfig : public ffi::ObjectRef {
};
/*!
- * \brief The fallback body used by TVM_REGISTER_SCRIPT_AS_REPR (defined in
printer.h).
+ * \brief The fallback body used by TVMScriptPrinter::Register.
*
* Tries to format \p obj via tvm::Script; on error falls back to a plain
* address string. Defined in src/script/printer/script_printer.cc so that
diff --git a/include/tvm/script/printer/printer.h
b/include/tvm/script/printer/printer.h
index a68c3b9e5d..93d5a0d64e 100644
--- a/include/tvm/script/printer/printer.h
+++ b/include/tvm/script/printer/printer.h
@@ -40,32 +40,35 @@ namespace tvm {
TVM_DLL std::string Script(const ffi::ObjectRef& node,
const ffi::Optional<PrinterConfig>& config =
std::nullopt);
-/*! \brief Dispatch vtable used by per-dialect printers to register their
- * object-type printing functions. Internal, but exposed here because
- * TVM_REGISTER_SCRIPT_AS_REPR refers to it.
- */
+/*! \brief Dispatch table for TVMScript printing and repr registration. */
class TVMScriptPrinter {
public:
using FType = ObjectFunctor<std::string(const ffi::ObjectRef&, const
PrinterConfig&)>;
TVM_DLL static FType& vtable();
-};
-/*!
- * \brief Register Script as the kRepr callback for ObjectType and install
- * the per-type dispatch entry in TVMScriptPrinter::vtable().
- *
- * \param ObjectType The concrete object node type (e.g. tirx::VarNode).
- * \param Method The TVMScriptPrinter vtable dispatch function.
- */
-#define TVM_REGISTER_SCRIPT_AS_REPR(ObjectType, Method)
\
- TVM_FFI_STATIC_INIT_BLOCK() {
\
- namespace refl = tvm::ffi::reflection;
\
- refl::TypeAttrDef<ObjectType>().def(refl::type_attr::kRepr,
\
- [](ffi::ObjectRef obj, ffi::Function)
-> ffi::String { \
- return
RedirectedReprPrinterMethod(obj); \
- });
\
- TVMScriptPrinter::vtable().SetDispatch<ObjectType>(Method);
\
+ /*! \brief Register a printer method for script dispatch and FFI repr.
+ * \tparam ObjectType Concrete object node type.
+ * \tparam Method Callable printer method type.
+ * \param method Printer dispatch method for that type.
+ *
+ * Example:
+ * \code
+ * TVM_FFI_STATIC_INIT_BLOCK() {
+ * TVMScriptPrinter::Register<tirx::ForNode>(ReprPrintTIR);
+ * TVMScriptPrinter::Register<tirx::WhileNode>(ReprPrintTIR);
+ * }
+ * \endcode
+ */
+ template <typename ObjectType, typename Method>
+ static void Register(Method method) {
+ namespace refl = tvm::ffi::reflection;
+ refl::TypeAttrDef<ObjectType>().def(refl::type_attr::kRepr,
+ [](ffi::ObjectRef obj, ffi::Function)
-> ffi::String {
+ return
RedirectedReprPrinterMethod(obj);
+ });
+ vtable().SetDispatch<ObjectType>(method);
}
+};
} // namespace tvm
#endif // TVM_SCRIPT_PRINTER_PRINTER_H_
diff --git a/src/ir/expr.cc b/src/ir/expr.cc
index fca78bafc0..189e039d9a 100644
--- a/src/ir/expr.cc
+++ b/src/ir/expr.cc
@@ -1166,7 +1166,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
return ss.str();
});
// Note: kRepr for GlobalVarNode is registered in script/printer/ir/ir.cc
- // via TVM_REGISTER_SCRIPT_AS_REPR(GlobalVarNode, ReprPrintIR).
+ // via TVMScriptPrinter::Register<GlobalVarNode>(ReprPrintIR).
}
} // namespace tvm
diff --git a/src/ir/prim/expr.cc b/src/ir/prim/expr.cc
index 93aa5bdf75..78ca9b032c 100644
--- a/src/ir/prim/expr.cc
+++ b/src/ir/prim/expr.cc
@@ -392,7 +392,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
refl::GlobalDef().def("prim.convert",
[](ffi::Variant<PrimExpr, ffi::Array<PrimExpr>> expr)
{ return expr; });
- // Note: kRepr for VarNode is registered via TVM_REGISTER_SCRIPT_AS_REPR in
+ // Note: kRepr for VarNode is registered via TVMScriptPrinter::Register in
// src/script/printer/tirx/expr.cc (-> ReprPrintTIR which delegates to
TVMScriptPrinter).
}
diff --git a/src/relax/script/printer/binding.cc
b/src/relax/script/printer/binding.cc
index 701c97ebf7..0896d2e5a8 100644
--- a/src/relax/script/printer/binding.cc
+++ b/src/relax/script/printer/binding.cc
@@ -99,9 +99,11 @@ TVM_FFI_STATIC_INIT_BLOCK() {
});
}
-TVM_REGISTER_SCRIPT_AS_REPR(relax::MatchCastNode, ReprPrintRelax);
-TVM_REGISTER_SCRIPT_AS_REPR(relax::VarBindingNode, ReprPrintRelax);
-TVM_REGISTER_SCRIPT_AS_REPR(relax::IfNode, ReprPrintRelax);
+TVM_FFI_STATIC_INIT_BLOCK() {
+ TVMScriptPrinter::Register<relax::MatchCastNode>(ReprPrintRelax);
+ TVMScriptPrinter::Register<relax::VarBindingNode>(ReprPrintRelax);
+ TVMScriptPrinter::Register<relax::IfNode>(ReprPrintRelax);
+}
} // namespace printer
} // namespace script
diff --git a/src/relax/script/printer/call.cc b/src/relax/script/printer/call.cc
index 628cbfa5f8..c93efa6cf1 100644
--- a/src/relax/script/printer/call.cc
+++ b/src/relax/script/printer/call.cc
@@ -352,7 +352,7 @@ std::string ReprPrintCall(const ffi::ObjectRef& obj, const
PrinterConfig& cfg) {
return ReprPrintRelax(obj, cfg);
}
-TVM_REGISTER_SCRIPT_AS_REPR(CallNode, ReprPrintCall);
+TVM_FFI_STATIC_INIT_BLOCK() {
TVMScriptPrinter::Register<CallNode>(ReprPrintCall); }
} // namespace printer
} // namespace script
diff --git a/src/relax/script/printer/dependent_type.cc
b/src/relax/script/printer/dependent_type.cc
index 433932502a..ec8115897a 100644
--- a/src/relax/script/printer/dependent_type.cc
+++ b/src/relax/script/printer/dependent_type.cc
@@ -151,10 +151,12 @@ TVM_FFI_STATIC_INIT_BLOCK() {
});
}
-TVM_REGISTER_SCRIPT_AS_REPR(AnyTypeNode, ReprPrintRelax);
-TVM_REGISTER_SCRIPT_AS_REPR(relax::ShapeTypeNode, ReprPrintRelax);
-TVM_REGISTER_SCRIPT_AS_REPR(relax::TensorTypeNode, ReprPrintRelax);
-TVM_REGISTER_SCRIPT_AS_REPR(relax::FuncTypeNode, ReprPrintRelax);
+TVM_FFI_STATIC_INIT_BLOCK() {
+ TVMScriptPrinter::Register<AnyTypeNode>(ReprPrintRelax);
+ TVMScriptPrinter::Register<relax::ShapeTypeNode>(ReprPrintRelax);
+ TVMScriptPrinter::Register<relax::TensorTypeNode>(ReprPrintRelax);
+ TVMScriptPrinter::Register<relax::FuncTypeNode>(ReprPrintRelax);
+}
} // namespace printer
} // namespace script
diff --git a/src/relax/script/printer/distributed.cc
b/src/relax/script/printer/distributed.cc
index cfc8dc59e8..5cac35ef22 100644
--- a/src/relax/script/printer/distributed.cc
+++ b/src/relax/script/printer/distributed.cc
@@ -137,9 +137,11 @@ TVM_FFI_STATIC_INIT_BLOCK() {
});
}
-TVM_REGISTER_SCRIPT_AS_REPR(relax::distributed::DeviceMeshNode,
ReprPrintRelax);
-TVM_REGISTER_SCRIPT_AS_REPR(relax::distributed::PlacementNode, ReprPrintRelax);
-TVM_REGISTER_SCRIPT_AS_REPR(relax::distributed::DTensorTypeNode,
ReprPrintRelax);
+TVM_FFI_STATIC_INIT_BLOCK() {
+
TVMScriptPrinter::Register<relax::distributed::DeviceMeshNode>(ReprPrintRelax);
+
TVMScriptPrinter::Register<relax::distributed::PlacementNode>(ReprPrintRelax);
+
TVMScriptPrinter::Register<relax::distributed::DTensorTypeNode>(ReprPrintRelax);
+}
} // namespace printer
} // namespace script
} // namespace tvm
diff --git a/src/relax/script/printer/expr.cc b/src/relax/script/printer/expr.cc
index 176d01ae83..c03be24742 100644
--- a/src/relax/script/printer/expr.cc
+++ b/src/relax/script/printer/expr.cc
@@ -173,13 +173,15 @@ std::string ReprPrintVar(const ffi::ObjectRef& obj, const
PrinterConfig& cfg) {
return ReprPrintRelax(obj, cfg);
}
-TVM_REGISTER_SCRIPT_AS_REPR(relax::TupleNode, ReprPrintRelax);
-TVM_REGISTER_SCRIPT_AS_REPR(relax::TupleGetItemNode, ReprPrintRelax);
-TVM_REGISTER_SCRIPT_AS_REPR(relax::ShapeExprNode, ReprPrintRelax);
-TVM_REGISTER_SCRIPT_AS_REPR(VarNode, ReprPrintVar);
-TVM_REGISTER_SCRIPT_AS_REPR(relax::DataflowVarNode, ReprPrintRelax);
-TVM_REGISTER_SCRIPT_AS_REPR(::tvm::GenericConstNode, ReprPrintRelax);
-TVM_REGISTER_SCRIPT_AS_REPR(::tvm::DataTypeImmNode, ReprPrintRelax);
+TVM_FFI_STATIC_INIT_BLOCK() {
+ TVMScriptPrinter::Register<relax::TupleNode>(ReprPrintRelax);
+ TVMScriptPrinter::Register<relax::TupleGetItemNode>(ReprPrintRelax);
+ TVMScriptPrinter::Register<relax::ShapeExprNode>(ReprPrintRelax);
+ TVMScriptPrinter::Register<VarNode>(ReprPrintVar);
+ TVMScriptPrinter::Register<relax::DataflowVarNode>(ReprPrintRelax);
+ TVMScriptPrinter::Register<::tvm::GenericConstNode>(ReprPrintRelax);
+ TVMScriptPrinter::Register<::tvm::DataTypeImmNode>(ReprPrintRelax);
+}
} // namespace printer
} // namespace script
diff --git a/src/relax/script/printer/function.cc
b/src/relax/script/printer/function.cc
index 7e5cc28b8f..53037d5d75 100644
--- a/src/relax/script/printer/function.cc
+++ b/src/relax/script/printer/function.cc
@@ -151,8 +151,10 @@ TVM_FFI_STATIC_INIT_BLOCK() {
});
}
-TVM_REGISTER_SCRIPT_AS_REPR(relax::FunctionNode, ReprPrintRelax);
-TVM_REGISTER_SCRIPT_AS_REPR(relax::ExternFuncNode, ReprPrintRelax);
+TVM_FFI_STATIC_INIT_BLOCK() {
+ TVMScriptPrinter::Register<relax::FunctionNode>(ReprPrintRelax);
+ TVMScriptPrinter::Register<relax::ExternFuncNode>(ReprPrintRelax);
+}
} // namespace printer
} // namespace script
diff --git a/src/relax/script/printer/region.cc
b/src/relax/script/printer/region.cc
index 31b1068f00..42fb43666b 100644
--- a/src/relax/script/printer/region.cc
+++ b/src/relax/script/printer/region.cc
@@ -98,9 +98,11 @@ TVM_FFI_STATIC_INIT_BLOCK() {
});
}
-TVM_REGISTER_SCRIPT_AS_REPR(relax::SeqExprNode, ReprPrintRelax);
-TVM_REGISTER_SCRIPT_AS_REPR(relax::BindingBlockNode, ReprPrintRelax);
-TVM_REGISTER_SCRIPT_AS_REPR(relax::DataflowBlockNode, ReprPrintRelax);
+TVM_FFI_STATIC_INIT_BLOCK() {
+ TVMScriptPrinter::Register<relax::SeqExprNode>(ReprPrintRelax);
+ TVMScriptPrinter::Register<relax::BindingBlockNode>(ReprPrintRelax);
+ TVMScriptPrinter::Register<relax::DataflowBlockNode>(ReprPrintRelax);
+}
} // namespace printer
} // namespace script
diff --git a/src/relax/script/printer/type.cc b/src/relax/script/printer/type.cc
index ec217e9c06..65408d855f 100644
--- a/src/relax/script/printer/type.cc
+++ b/src/relax/script/printer/type.cc
@@ -61,7 +61,9 @@ TVM_FFI_STATIC_INIT_BLOCK() {
});
}
-TVM_REGISTER_SCRIPT_AS_REPR(relax::PackedFuncTypeNode, ReprPrintRelax);
+TVM_FFI_STATIC_INIT_BLOCK() {
+ TVMScriptPrinter::Register<relax::PackedFuncTypeNode>(ReprPrintRelax);
+}
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
refl::GlobalDef().def("script.printer.ReprPrintRelax", ReprPrintRelax);
diff --git a/src/s_tir/schedule/instruction.cc
b/src/s_tir/schedule/instruction.cc
index 81c20c0f37..bb8f78b219 100644
--- a/src/s_tir/schedule/instruction.cc
+++ b/src/s_tir/schedule/instruction.cc
@@ -17,6 +17,7 @@
* under the License.
*/
#include <tvm/ffi/cast.h>
+#include <tvm/ffi/container/dict.h>
#include <tvm/ffi/extra/structural_mutate.h>
#include <tvm/ffi/reflection/registry.h>
@@ -99,21 +100,42 @@ TVM_FFI_STATIC_INIT_BLOCK() {
});
}
-using InstructionKindRegistry = AttrRegistry<InstructionKindRegEntry,
InstructionKind>;
+namespace {
+class InstructionKindRegistry {
+ public:
+ static InstructionKindRegistry* Global() {
+ static InstructionKindRegistry registry;
+ return ®istry;
+ }
-InstructionKind InstructionKind::Get(const ffi::String& name) {
- const InstructionKindRegEntry* reg =
InstructionKindRegistry::Global()->Get(name);
- TVM_FFI_CHECK(reg != nullptr, AttributeError)
- << "Instruction kind " << name << " is not registered";
- return reg->inst_kind_;
-}
+ InstructionKind Get(const ffi::String& name) const {
+ auto kind = kinds_.Get(name);
+ TVM_FFI_CHECK(kind.has_value(), AttributeError)
+ << "Instruction kind " << name << " is not registered";
+ return *kind;
+ }
-InstructionKindRegEntry::InstructionKindRegEntry(uint32_t reg_index) {
- this->inst_kind_ = InstructionKind(ffi::make_object<InstructionKindNode>());
+ InstructionKind RegisterOrGet(const ffi::String& name) {
+ if (auto kind = kinds_.Get(name)) {
+ return *kind;
+ }
+ InstructionKind kind(ffi::make_object<InstructionKindNode>());
+ kinds_.Set(name, kind);
+ return kind;
+ }
+
+ private:
+ ffi::Dict<ffi::String, InstructionKind> kinds_;
+};
+} // namespace
+
+InstructionKind InstructionKind::Get(const ffi::String& name) {
+ return InstructionKindRegistry::Global()->Get(name);
}
-InstructionKindRegEntry& InstructionKindRegEntry::RegisterOrGet(const
ffi::String& name) {
- return InstructionKindRegistry::Global()->RegisterOrGet(name);
+InstructionKindDef::InstructionKindDef(const ffi::String& name)
+ : inst_kind_(InstructionKindRegistry::Global()->RegisterOrGet(name)) {
+ get_mutable()->name = name;
}
/**************** Repr ****************/
diff --git a/src/s_tir/schedule/instruction_traits.h
b/src/s_tir/schedule/instruction_traits.h
index 24d91d9140..8a6e9ebd90 100644
--- a/src/s_tir/schedule/instruction_traits.h
+++ b/src/s_tir/schedule/instruction_traits.h
@@ -33,52 +33,26 @@ namespace tvm {
namespace s_tir {
using namespace tvm::tirx;
-/*!
- * \brief Register an InstructionKind using a trait class
- * \param InstructionKindTraits A traits class of an InstructionKind
+/*! \brief Register a named instruction kind using its typed traits.
+ * \tparam Traits Traits with kName, kIsPure, and the instruction callbacks.
*
* Example:
- *
* \code
- *
- * struct SomeInstructionKindTraits {
- * static constexpr const char* kName = "name-of-the-instruction";
- * static constexpr bool kIsPure = false;
- *
- * // Convertible to `InstructionKindNode::FInstructionApply`
- * static ffi::Array<ffi::ObjectRef> ApplyToSchedule(
- * const s_tir::Schedule& sch,
- * const ffi::Array<ffi::ObjectRef>& inputs,
- * const ffi::Array<ffi::ObjectRef>& attrs,
- * const ffi::Optional<ffi::ObjectRef>& decision);
- *
- * // Convertible to `InstructionKindNode::FInstructionAsPython`
- * static ffi::String AsPython(
- * const ffi::Array<ffi::String>& inputs,
- * const ffi::Array<ffi::ObjectRef>& attrs,
- * const ffi::Optional<ffi::ObjectRef>& decision,
- * const ffi::Array<ffi::String>& outputs);
- *
- * // Convertible to `InstructionKindNode::FInstructionAttrsAsJSON`
- * static ffi::ObjectRef AttrsAsJSON(
- * const ffi::Array<ffi::ObjectRef>& attrs);
- *
- * // Convertible to `InstructionKindNode::FInstructionAttrsFromJSON`
- * static ffi::Array<ffi::ObjectRef> AttrsFromJSON(
- * const ffi::ObjectRef& attrs_record);
- * };
- *
- * TVM_REGISTER_INST_KIND_TRAITS(SomeInstructionKindTraits);
- *
+ * TVM_FFI_STATIC_INIT_BLOCK() {
+ * RegisterInstructionKind<ComputeInlineTraits>();
+ * RegisterInstructionKind<ReverseComputeInlineTraits>();
+ * }
* \endcode
*/
-#define TVM_REGISTER_INST_KIND_TRAITS(InstructionKindTraits) \
- TVM_REGISTER_INST_KIND(InstructionKindTraits::kName) \
- .set_is_pure(InstructionKindTraits::kIsPure) \
- .set_apply_to_schedule(InstructionKindTraits::ApplyToSchedule) \
- .set_attrs_as_json(InstructionKindTraits::AttrsAsJSON) \
- .set_attrs_from_json(InstructionKindTraits::AttrsFromJSON) \
- .set_as_python(InstructionKindTraits::AsPython)
+template <typename Traits>
+void RegisterInstructionKind() {
+ InstructionKindDef(Traits::kName)
+ .set_is_pure(Traits::kIsPure)
+ .set_apply_to_schedule(Traits::ApplyToSchedule)
+ .set_attrs_as_json(Traits::AttrsAsJSON)
+ .set_attrs_from_json(Traits::AttrsFromJSON)
+ .set_as_python(Traits::AsPython);
+}
/*!
* \brief A helper to conveniently register an InstructionKind. When inherited
in curiously
@@ -144,7 +118,7 @@ using namespace tvm::tirx;
* friend struct UnpackedInstTraits;
* };
*
- * TVM_REGISTER_INST_KIND(SamplePerfectTileTraits);
+ * RegisterInstructionKind<SamplePerfectTileTraits>();
* \endcode
*/
template <class TTraits>
diff --git a/src/s_tir/schedule/primitive/annotate.cc
b/src/s_tir/schedule/primitive/annotate.cc
index 9fd1bc2752..170c308ea9 100644
--- a/src/s_tir/schedule/primitive/annotate.cc
+++ b/src/s_tir/schedule/primitive/annotate.cc
@@ -160,8 +160,10 @@ struct UnannotateTraits : public
UnpackedInstTraits<UnannotateTraits> {
friend struct ::tvm::s_tir::UnpackedInstTraits;
};
-TVM_REGISTER_INST_KIND_TRAITS(AnnotateTraits);
-TVM_REGISTER_INST_KIND_TRAITS(UnannotateTraits);
+TVM_FFI_STATIC_INIT_BLOCK() {
+ RegisterInstructionKind<AnnotateTraits>();
+ RegisterInstructionKind<UnannotateTraits>();
+}
} // namespace s_tir
} // namespace tvm
diff --git a/src/s_tir/schedule/primitive/annotate_buffer_access.cc
b/src/s_tir/schedule/primitive/annotate_buffer_access.cc
index ff24b191c0..625338ebc4 100644
--- a/src/s_tir/schedule/primitive/annotate_buffer_access.cc
+++ b/src/s_tir/schedule/primitive/annotate_buffer_access.cc
@@ -190,7 +190,7 @@ struct AnnotateBufferAccessTraits : public
UnpackedInstTraits<AnnotateBufferAcce
friend struct ::tvm::s_tir::UnpackedInstTraits;
};
-TVM_REGISTER_INST_KIND_TRAITS(AnnotateBufferAccessTraits);
+TVM_FFI_STATIC_INIT_BLOCK() {
RegisterInstructionKind<AnnotateBufferAccessTraits>(); }
} // namespace s_tir
} // namespace tvm
diff --git a/src/s_tir/schedule/primitive/block_annotate.cc
b/src/s_tir/schedule/primitive/block_annotate.cc
index 6e6c0f4804..3ebb1fb9a8 100644
--- a/src/s_tir/schedule/primitive/block_annotate.cc
+++ b/src/s_tir/schedule/primitive/block_annotate.cc
@@ -473,9 +473,11 @@ struct UnsafeSetDTypeTraits : public
UnpackedInstTraits<UnsafeSetDTypeTraits> {
friend struct ::tvm::s_tir::UnpackedInstTraits;
};
-TVM_REGISTER_INST_KIND_TRAITS(StorageAlignTraits);
-TVM_REGISTER_INST_KIND_TRAITS(SetScopeTraits);
-TVM_REGISTER_INST_KIND_TRAITS(UnsafeSetDTypeTraits);
+TVM_FFI_STATIC_INIT_BLOCK() {
+ RegisterInstructionKind<StorageAlignTraits>();
+ RegisterInstructionKind<SetScopeTraits>();
+ RegisterInstructionKind<UnsafeSetDTypeTraits>();
+}
} // namespace s_tir
} // namespace tvm
diff --git a/src/s_tir/schedule/primitive/blockize_tensorize.cc
b/src/s_tir/schedule/primitive/blockize_tensorize.cc
index 6b81582787..6d713ae327 100644
--- a/src/s_tir/schedule/primitive/blockize_tensorize.cc
+++ b/src/s_tir/schedule/primitive/blockize_tensorize.cc
@@ -985,8 +985,10 @@ struct TensorizeTraits : public
UnpackedInstTraits<TensorizeTraits> {
friend struct ::tvm::s_tir::UnpackedInstTraits;
};
-TVM_REGISTER_INST_KIND_TRAITS(BlockizeTraits);
-TVM_REGISTER_INST_KIND_TRAITS(TensorizeTraits);
+TVM_FFI_STATIC_INIT_BLOCK() {
+ RegisterInstructionKind<BlockizeTraits>();
+ RegisterInstructionKind<TensorizeTraits>();
+}
} // namespace s_tir
} // namespace tvm
diff --git a/src/s_tir/schedule/primitive/cache_index.cc
b/src/s_tir/schedule/primitive/cache_index.cc
index 51a1ba7fbf..88fc3cbf6e 100644
--- a/src/s_tir/schedule/primitive/cache_index.cc
+++ b/src/s_tir/schedule/primitive/cache_index.cc
@@ -559,7 +559,7 @@ struct CacheIndexTraits : public
UnpackedInstTraits<CacheIndexTraits> {
friend struct ::tvm::s_tir::UnpackedInstTraits;
};
-TVM_REGISTER_INST_KIND_TRAITS(CacheIndexTraits);
+TVM_FFI_STATIC_INIT_BLOCK() { RegisterInstructionKind<CacheIndexTraits>(); }
} // namespace s_tir
} // namespace tvm
diff --git a/src/s_tir/schedule/primitive/cache_read_write.cc
b/src/s_tir/schedule/primitive/cache_read_write.cc
index de27c902ed..b2cbb1a6c1 100644
--- a/src/s_tir/schedule/primitive/cache_read_write.cc
+++ b/src/s_tir/schedule/primitive/cache_read_write.cc
@@ -2814,12 +2814,14 @@ struct ReindexCacheWriteTraits : public
UnpackedInstTraits<ReindexCacheWriteTrai
friend struct ::tvm::s_tir::UnpackedInstTraits;
};
-TVM_REGISTER_INST_KIND_TRAITS(CacheReadTraits);
-TVM_REGISTER_INST_KIND_TRAITS(CacheWriteTraits);
-TVM_REGISTER_INST_KIND_TRAITS(CacheInplaceTraits);
-TVM_REGISTER_INST_KIND_TRAITS(ReIndexTraits);
-TVM_REGISTER_INST_KIND_TRAITS(ReindexCacheReadTraits);
-TVM_REGISTER_INST_KIND_TRAITS(ReindexCacheWriteTraits);
+TVM_FFI_STATIC_INIT_BLOCK() {
+ RegisterInstructionKind<CacheReadTraits>();
+ RegisterInstructionKind<CacheWriteTraits>();
+ RegisterInstructionKind<CacheInplaceTraits>();
+ RegisterInstructionKind<ReIndexTraits>();
+ RegisterInstructionKind<ReindexCacheReadTraits>();
+ RegisterInstructionKind<ReindexCacheWriteTraits>();
+}
} // namespace s_tir
} // namespace tvm
diff --git a/src/s_tir/schedule/primitive/compute_at.cc
b/src/s_tir/schedule/primitive/compute_at.cc
index 3790a26b44..d4e2be14c8 100644
--- a/src/s_tir/schedule/primitive/compute_at.cc
+++ b/src/s_tir/schedule/primitive/compute_at.cc
@@ -893,8 +893,10 @@ struct ReverseComputeAtTraits : public
UnpackedInstTraits<ReverseComputeAtTraits
friend struct ::tvm::s_tir::UnpackedInstTraits;
};
-TVM_REGISTER_INST_KIND_TRAITS(ComputeAtTraits);
-TVM_REGISTER_INST_KIND_TRAITS(ReverseComputeAtTraits);
+TVM_FFI_STATIC_INIT_BLOCK() {
+ RegisterInstructionKind<ComputeAtTraits>();
+ RegisterInstructionKind<ReverseComputeAtTraits>();
+}
} // namespace s_tir
} // namespace tvm
diff --git a/src/s_tir/schedule/primitive/compute_inline.cc
b/src/s_tir/schedule/primitive/compute_inline.cc
index b93e29f5b8..416e96d3f4 100644
--- a/src/s_tir/schedule/primitive/compute_inline.cc
+++ b/src/s_tir/schedule/primitive/compute_inline.cc
@@ -1937,8 +1937,10 @@ struct ReverseComputeInlineTraits : public
UnpackedInstTraits<ReverseComputeInli
friend struct ::tvm::s_tir::UnpackedInstTraits;
};
-TVM_REGISTER_INST_KIND_TRAITS(ComputeInlineTraits);
-TVM_REGISTER_INST_KIND_TRAITS(ReverseComputeInlineTraits);
+TVM_FFI_STATIC_INIT_BLOCK() {
+ RegisterInstructionKind<ComputeInlineTraits>();
+ RegisterInstructionKind<ReverseComputeInlineTraits>();
+}
struct FuseReductionEpilogueTraits : public
UnpackedInstTraits<FuseReductionEpilogueTraits> {
static constexpr const char* kName = "FuseReductionEpilogue";
@@ -1967,7 +1969,7 @@ struct FuseReductionEpilogueTraits : public
UnpackedInstTraits<FuseReductionEpil
friend struct ::tvm::s_tir::UnpackedInstTraits;
};
-TVM_REGISTER_INST_KIND_TRAITS(FuseReductionEpilogueTraits);
+TVM_FFI_STATIC_INIT_BLOCK() {
RegisterInstructionKind<FuseReductionEpilogueTraits>(); }
} // namespace s_tir
} // namespace tvm
diff --git a/src/s_tir/schedule/primitive/decompose_padding.cc
b/src/s_tir/schedule/primitive/decompose_padding.cc
index 0ad27fc00f..f184c872ff 100644
--- a/src/s_tir/schedule/primitive/decompose_padding.cc
+++ b/src/s_tir/schedule/primitive/decompose_padding.cc
@@ -609,7 +609,7 @@ struct DecomposPaddingTraits : public
UnpackedInstTraits<DecomposPaddingTraits>
friend struct ::tvm::s_tir::UnpackedInstTraits;
};
-TVM_REGISTER_INST_KIND_TRAITS(DecomposPaddingTraits);
+TVM_FFI_STATIC_INIT_BLOCK() {
RegisterInstructionKind<DecomposPaddingTraits>(); }
} // namespace s_tir
} // namespace tvm
diff --git a/src/s_tir/schedule/primitive/for_kind.cc
b/src/s_tir/schedule/primitive/for_kind.cc
index 558aff410d..392ea8f080 100644
--- a/src/s_tir/schedule/primitive/for_kind.cc
+++ b/src/s_tir/schedule/primitive/for_kind.cc
@@ -304,10 +304,12 @@ struct UnrollTraits : public
UnpackedInstTraits<UnrollTraits> {
friend struct ::tvm::s_tir::UnpackedInstTraits;
};
-TVM_REGISTER_INST_KIND_TRAITS(ParallelTraits);
-TVM_REGISTER_INST_KIND_TRAITS(VectorizeTraits);
-TVM_REGISTER_INST_KIND_TRAITS(BindTraits);
-TVM_REGISTER_INST_KIND_TRAITS(UnrollTraits);
+TVM_FFI_STATIC_INIT_BLOCK() {
+ RegisterInstructionKind<ParallelTraits>();
+ RegisterInstructionKind<VectorizeTraits>();
+ RegisterInstructionKind<BindTraits>();
+ RegisterInstructionKind<UnrollTraits>();
+}
} // namespace s_tir
} // namespace tvm
diff --git a/src/s_tir/schedule/primitive/get_block_loop.cc
b/src/s_tir/schedule/primitive/get_block_loop.cc
index b2e0823679..1539e56fa9 100644
--- a/src/s_tir/schedule/primitive/get_block_loop.cc
+++ b/src/s_tir/schedule/primitive/get_block_loop.cc
@@ -274,12 +274,14 @@ struct GetOutputBlocksTraits : public
UnpackedInstTraits<GetOutputBlocksTraits>
friend struct ::tvm::s_tir::UnpackedInstTraits;
};
-TVM_REGISTER_INST_KIND_TRAITS(GetSBlockTraits);
-TVM_REGISTER_INST_KIND_TRAITS(GetLoopsTraits);
-TVM_REGISTER_INST_KIND_TRAITS(GetChildBlocksTraits);
-TVM_REGISTER_INST_KIND_TRAITS(GetProducersTraits);
-TVM_REGISTER_INST_KIND_TRAITS(GetConsumersTraits);
-TVM_REGISTER_INST_KIND_TRAITS(GetOutputBlocksTraits);
+TVM_FFI_STATIC_INIT_BLOCK() {
+ RegisterInstructionKind<GetSBlockTraits>();
+ RegisterInstructionKind<GetLoopsTraits>();
+ RegisterInstructionKind<GetChildBlocksTraits>();
+ RegisterInstructionKind<GetProducersTraits>();
+ RegisterInstructionKind<GetConsumersTraits>();
+ RegisterInstructionKind<GetOutputBlocksTraits>();
+}
} // namespace s_tir
} // namespace tvm
diff --git a/src/s_tir/schedule/primitive/hide_buffer_access.cc
b/src/s_tir/schedule/primitive/hide_buffer_access.cc
index 43649044ea..58755a0ae0 100644
--- a/src/s_tir/schedule/primitive/hide_buffer_access.cc
+++ b/src/s_tir/schedule/primitive/hide_buffer_access.cc
@@ -172,7 +172,7 @@ struct UnsafeHideBufferAccessTraits : public
UnpackedInstTraits<UnsafeHideBuffer
friend struct ::tvm::s_tir::UnpackedInstTraits;
};
-TVM_REGISTER_INST_KIND_TRAITS(UnsafeHideBufferAccessTraits);
+TVM_FFI_STATIC_INIT_BLOCK() {
RegisterInstructionKind<UnsafeHideBufferAccessTraits>(); }
} // namespace s_tir
} // namespace tvm
diff --git a/src/s_tir/schedule/primitive/layout_transformation.cc
b/src/s_tir/schedule/primitive/layout_transformation.cc
index bb56fa6bb7..5099f435c1 100644
--- a/src/s_tir/schedule/primitive/layout_transformation.cc
+++ b/src/s_tir/schedule/primitive/layout_transformation.cc
@@ -1703,8 +1703,10 @@ struct TransformBlockLayoutTraits : public
UnpackedInstTraits<TransformBlockLayo
friend struct ::tvm::s_tir::UnpackedInstTraits;
};
-TVM_REGISTER_INST_KIND_TRAITS(TransformLayoutTraits);
-TVM_REGISTER_INST_KIND_TRAITS(TransformBlockLayoutTraits);
+TVM_FFI_STATIC_INIT_BLOCK() {
+ RegisterInstructionKind<TransformLayoutTraits>();
+ RegisterInstructionKind<TransformBlockLayoutTraits>();
+}
} // namespace s_tir
} // namespace tvm
diff --git a/src/s_tir/schedule/primitive/loop_transformation.cc
b/src/s_tir/schedule/primitive/loop_transformation.cc
index 4f6280f6d3..e342d24be6 100644
--- a/src/s_tir/schedule/primitive/loop_transformation.cc
+++ b/src/s_tir/schedule/primitive/loop_transformation.cc
@@ -1476,12 +1476,14 @@ struct AddUnitLoopTraits : public
UnpackedInstTraits<AddUnitLoopTraits> {
friend struct ::tvm::s_tir::UnpackedInstTraits;
};
-TVM_REGISTER_INST_KIND_TRAITS(SplitTraits);
-TVM_REGISTER_INST_KIND_TRAITS(LoopPartitionTraits);
-TVM_REGISTER_INST_KIND_TRAITS(MergeTraits);
-TVM_REGISTER_INST_KIND_TRAITS(FuseTraits);
-TVM_REGISTER_INST_KIND_TRAITS(ReorderTraits);
-TVM_REGISTER_INST_KIND_TRAITS(AddUnitLoopTraits);
+TVM_FFI_STATIC_INIT_BLOCK() {
+ RegisterInstructionKind<SplitTraits>();
+ RegisterInstructionKind<LoopPartitionTraits>();
+ RegisterInstructionKind<MergeTraits>();
+ RegisterInstructionKind<FuseTraits>();
+ RegisterInstructionKind<ReorderTraits>();
+ RegisterInstructionKind<AddUnitLoopTraits>();
+}
} // namespace s_tir
} // namespace tvm
diff --git a/src/s_tir/schedule/primitive/pad_einsum.cc
b/src/s_tir/schedule/primitive/pad_einsum.cc
index c7758f3348..6cafdc3536 100644
--- a/src/s_tir/schedule/primitive/pad_einsum.cc
+++ b/src/s_tir/schedule/primitive/pad_einsum.cc
@@ -530,7 +530,7 @@ struct PadEinsumTraits : public
UnpackedInstTraits<PadEinsumTraits> {
friend struct ::tvm::s_tir::UnpackedInstTraits;
};
-TVM_REGISTER_INST_KIND_TRAITS(PadEinsumTraits);
+TVM_FFI_STATIC_INIT_BLOCK() { RegisterInstructionKind<PadEinsumTraits>(); }
} // namespace s_tir
} // namespace tvm
diff --git a/src/s_tir/schedule/primitive/read_write_at.cc
b/src/s_tir/schedule/primitive/read_write_at.cc
index 7c8f2f134b..3c664879f5 100644
--- a/src/s_tir/schedule/primitive/read_write_at.cc
+++ b/src/s_tir/schedule/primitive/read_write_at.cc
@@ -452,8 +452,10 @@ struct WriteAtTraits : public
UnpackedInstTraits<WriteAtTraits> {
friend struct ::tvm::s_tir::UnpackedInstTraits;
};
-TVM_REGISTER_INST_KIND_TRAITS(ReadAtTraits);
-TVM_REGISTER_INST_KIND_TRAITS(WriteAtTraits);
+TVM_FFI_STATIC_INIT_BLOCK() {
+ RegisterInstructionKind<ReadAtTraits>();
+ RegisterInstructionKind<WriteAtTraits>();
+}
} // namespace s_tir
} // namespace tvm
diff --git a/src/s_tir/schedule/primitive/reduction.cc
b/src/s_tir/schedule/primitive/reduction.cc
index 5caa524fb9..33388020dc 100644
--- a/src/s_tir/schedule/primitive/reduction.cc
+++ b/src/s_tir/schedule/primitive/reduction.cc
@@ -1530,8 +1530,10 @@ struct RFactorTraits : public
UnpackedInstTraits<RFactorTraits> {
friend struct ::tvm::s_tir::UnpackedInstTraits;
};
-TVM_REGISTER_INST_KIND_TRAITS(RFactorTraits);
-TVM_REGISTER_INST_KIND_TRAITS(DecomposeReductionTraits);
+TVM_FFI_STATIC_INIT_BLOCK() {
+ RegisterInstructionKind<RFactorTraits>();
+ RegisterInstructionKind<DecomposeReductionTraits>();
+}
/******** FFI ********/
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 58bc95e036..6e07f0af20 100644
--- a/src/s_tir/schedule/primitive/reorder_block_iter_var.cc
+++ b/src/s_tir/schedule/primitive/reorder_block_iter_var.cc
@@ -171,7 +171,7 @@ struct ReorderBlockIterVarTraits : public
UnpackedInstTraits<ReorderBlockIterVar
friend struct ::tvm::s_tir::UnpackedInstTraits;
};
-TVM_REGISTER_INST_KIND_TRAITS(ReorderBlockIterVarTraits);
+TVM_FFI_STATIC_INIT_BLOCK() {
RegisterInstructionKind<ReorderBlockIterVarTraits>(); }
} // namespace s_tir
} // namespace tvm
diff --git a/src/s_tir/schedule/primitive/rolling_buffer.cc
b/src/s_tir/schedule/primitive/rolling_buffer.cc
index 49aa13771f..e1a96a657e 100644
--- a/src/s_tir/schedule/primitive/rolling_buffer.cc
+++ b/src/s_tir/schedule/primitive/rolling_buffer.cc
@@ -512,6 +512,6 @@ struct RollingBufferTraits : public
UnpackedInstTraits<RollingBufferTraits> {
friend struct ::tvm::s_tir::UnpackedInstTraits;
};
-TVM_REGISTER_INST_KIND_TRAITS(RollingBufferTraits);
+TVM_FFI_STATIC_INIT_BLOCK() { RegisterInstructionKind<RollingBufferTraits>(); }
} // namespace s_tir
} // namespace tvm
diff --git a/src/s_tir/schedule/primitive/sampling.cc
b/src/s_tir/schedule/primitive/sampling.cc
index 4c862a3745..37020ec203 100644
--- a/src/s_tir/schedule/primitive/sampling.cc
+++ b/src/s_tir/schedule/primitive/sampling.cc
@@ -581,10 +581,12 @@ struct SampleComputeLocationTraits : public
UnpackedInstTraits<SampleComputeLoca
friend struct ::tvm::s_tir::UnpackedInstTraits;
};
-TVM_REGISTER_INST_KIND_TRAITS(SampleCategoricalTraits);
-TVM_REGISTER_INST_KIND_TRAITS(SamplePerfectTileTraits);
-TVM_REGISTER_INST_KIND_TRAITS(SamplePartitionedTileTraits);
-TVM_REGISTER_INST_KIND_TRAITS(SampleComputeLocationTraits);
+TVM_FFI_STATIC_INIT_BLOCK() {
+ RegisterInstructionKind<SampleCategoricalTraits>();
+ RegisterInstructionKind<SamplePerfectTileTraits>();
+ RegisterInstructionKind<SamplePartitionedTileTraits>();
+ RegisterInstructionKind<SampleComputeLocationTraits>();
+}
} // namespace s_tir
} // namespace tvm
diff --git a/src/s_tir/schedule/trace.cc b/src/s_tir/schedule/trace.cc
index e766ea4f44..471523ee60 100644
--- a/src/s_tir/schedule/trace.cc
+++ b/src/s_tir/schedule/trace.cc
@@ -677,7 +677,7 @@ struct EnterPostprocTraits : public
UnpackedInstTraits<EnterPostprocTraits> {
friend struct ::tvm::s_tir::UnpackedInstTraits;
};
-TVM_REGISTER_INST_KIND_TRAITS(EnterPostprocTraits);
+TVM_FFI_STATIC_INIT_BLOCK() { RegisterInstructionKind<EnterPostprocTraits>(); }
/**************** FFI ****************/
diff --git a/src/script/printer/ir/ir.cc b/src/script/printer/ir/ir.cc
index d4d41342b6..bc60d69694 100644
--- a/src/script/printer/ir/ir.cc
+++ b/src/script/printer/ir/ir.cc
@@ -198,11 +198,13 @@ TVM_FFI_STATIC_INIT_BLOCK() {
});
}
-TVM_REGISTER_SCRIPT_AS_REPR(GlobalVarNode, ReprPrintIR);
-TVM_REGISTER_SCRIPT_AS_REPR(DictAttrsNode, ReprPrintIR);
-TVM_REGISTER_SCRIPT_AS_REPR(FuncTypeNode, ReprPrintIR);
-TVM_REGISTER_SCRIPT_AS_REPR(RangeNode, ReprPrintIR);
-TVM_REGISTER_SCRIPT_AS_REPR(IRModuleNode, ReprPrintIR);
+TVM_FFI_STATIC_INIT_BLOCK() {
+ TVMScriptPrinter::Register<GlobalVarNode>(ReprPrintIR);
+ TVMScriptPrinter::Register<DictAttrsNode>(ReprPrintIR);
+ TVMScriptPrinter::Register<FuncTypeNode>(ReprPrintIR);
+ TVMScriptPrinter::Register<RangeNode>(ReprPrintIR);
+ TVMScriptPrinter::Register<IRModuleNode>(ReprPrintIR);
+}
} // namespace printer
} // namespace script
diff --git a/src/tirx/script/printer/block.cc b/src/tirx/script/printer/block.cc
index e6b392f3cd..d3e491d395 100644
--- a/src/tirx/script/printer/block.cc
+++ b/src/tirx/script/printer/block.cc
@@ -237,8 +237,10 @@ TVM_FFI_STATIC_INIT_BLOCK() {
});
}
-TVM_REGISTER_SCRIPT_AS_REPR(s_tir::SBlockNode, ReprPrintTIR);
-TVM_REGISTER_SCRIPT_AS_REPR(s_tir::SBlockRealizeNode, ReprPrintTIR);
+TVM_FFI_STATIC_INIT_BLOCK() {
+ TVMScriptPrinter::Register<s_tir::SBlockNode>(ReprPrintTIR);
+ TVMScriptPrinter::Register<s_tir::SBlockRealizeNode>(ReprPrintTIR);
+}
TVM_FFI_STATIC_INIT_BLOCK() {
IRDocsifier::vtable().set_dispatch<tirx::ScopeIdDefStmt>(
@@ -283,7 +285,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
});
}
-TVM_REGISTER_SCRIPT_AS_REPR(tirx::ScopeIdDefStmtNode, ReprPrintTIR);
+TVM_FFI_STATIC_INIT_BLOCK() {
TVMScriptPrinter::Register<tirx::ScopeIdDefStmtNode>(ReprPrintTIR); }
TVM_FFI_STATIC_INIT_BLOCK() {
IRDocsifier::vtable().set_dispatch<tirx::ExecScope>(
@@ -292,7 +294,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
return doc;
});
}
-TVM_REGISTER_SCRIPT_AS_REPR(tirx::ExecScopeNode, ReprPrintTIR);
+TVM_FFI_STATIC_INIT_BLOCK() {
TVMScriptPrinter::Register<tirx::ExecScopeNode>(ReprPrintTIR); }
TVM_FFI_STATIC_INIT_BLOCK() {
IRDocsifier::vtable().set_dispatch<tirx::ScopeIdDef>(
@@ -308,7 +310,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
return doc;
});
}
-TVM_REGISTER_SCRIPT_AS_REPR(tirx::ScopeIdDefNode, ReprPrintTIR);
+TVM_FFI_STATIC_INIT_BLOCK() {
TVMScriptPrinter::Register<tirx::ScopeIdDefNode>(ReprPrintTIR); }
} // namespace printer
} // namespace script
diff --git a/src/tirx/script/printer/buffer.cc
b/src/tirx/script/printer/buffer.cc
index 76f1a93611..272fb044d2 100644
--- a/src/tirx/script/printer/buffer.cc
+++ b/src/tirx/script/printer/buffer.cc
@@ -559,14 +559,16 @@ TVM_FFI_STATIC_INIT_BLOCK() {
});
}
-TVM_REGISTER_SCRIPT_AS_REPR(tvm::TensorRegionNode, ReprPrintTIR);
-TVM_REGISTER_SCRIPT_AS_REPR(TensorLoadNode, ReprPrintTIR);
-TVM_REGISTER_SCRIPT_AS_REPR(tirx::BufferStoreNode, ReprPrintTIR);
-TVM_REGISTER_SCRIPT_AS_REPR(tirx::BufferTypeNode, ReprPrintTIR);
-TVM_REGISTER_SCRIPT_AS_REPR(tirx::IterNode, ReprPrintTIR);
-TVM_REGISTER_SCRIPT_AS_REPR(tirx::TileLayoutNode, ReprPrintTIR);
-TVM_REGISTER_SCRIPT_AS_REPR(tirx::ComposeLayoutNode, ReprPrintTIR);
-TVM_REGISTER_SCRIPT_AS_REPR(s_tir::MatchBufferRegionNode, ReprPrintTIR);
+TVM_FFI_STATIC_INIT_BLOCK() {
+ TVMScriptPrinter::Register<tvm::TensorRegionNode>(ReprPrintTIR);
+ TVMScriptPrinter::Register<TensorLoadNode>(ReprPrintTIR);
+ TVMScriptPrinter::Register<tirx::BufferStoreNode>(ReprPrintTIR);
+ TVMScriptPrinter::Register<tirx::BufferTypeNode>(ReprPrintTIR);
+ TVMScriptPrinter::Register<tirx::IterNode>(ReprPrintTIR);
+ TVMScriptPrinter::Register<tirx::TileLayoutNode>(ReprPrintTIR);
+ TVMScriptPrinter::Register<tirx::ComposeLayoutNode>(ReprPrintTIR);
+ TVMScriptPrinter::Register<s_tir::MatchBufferRegionNode>(ReprPrintTIR);
+}
} // namespace printer
} // namespace script
diff --git a/src/tirx/script/printer/expr.cc b/src/tirx/script/printer/expr.cc
index 7bfb7fd50e..526a9971ee 100644
--- a/src/tirx/script/printer/expr.cc
+++ b/src/tirx/script/printer/expr.cc
@@ -422,42 +422,44 @@ TVM_FFI_STATIC_INIT_BLOCK() {
});
}
-TVM_REGISTER_SCRIPT_AS_REPR(tirx::IterVarNode, ReprPrintTIR);
-TVM_REGISTER_SCRIPT_AS_REPR(StringImmNode, ReprPrintTIR);
-TVM_REGISTER_SCRIPT_AS_REPR(prim::CastNode, ReprPrintTIR);
-TVM_REGISTER_SCRIPT_AS_REPR(prim::AddNode, ReprPrintTIR);
-TVM_REGISTER_SCRIPT_AS_REPR(prim::SubNode, ReprPrintTIR);
-TVM_REGISTER_SCRIPT_AS_REPR(prim::MulNode, ReprPrintTIR);
-TVM_REGISTER_SCRIPT_AS_REPR(prim::DivNode, ReprPrintTIR);
-TVM_REGISTER_SCRIPT_AS_REPR(prim::ModNode, ReprPrintTIR);
-TVM_REGISTER_SCRIPT_AS_REPR(prim::FloorDivNode, ReprPrintTIR);
-TVM_REGISTER_SCRIPT_AS_REPR(prim::FloorModNode, ReprPrintTIR);
-TVM_REGISTER_SCRIPT_AS_REPR(prim::LShiftNode, ReprPrintTIR);
-TVM_REGISTER_SCRIPT_AS_REPR(prim::RShiftNode, ReprPrintTIR);
-TVM_REGISTER_SCRIPT_AS_REPR(prim::BitwiseAndNode, ReprPrintTIR);
-TVM_REGISTER_SCRIPT_AS_REPR(prim::BitwiseOrNode, ReprPrintTIR);
-TVM_REGISTER_SCRIPT_AS_REPR(prim::BitwiseXorNode, ReprPrintTIR);
-TVM_REGISTER_SCRIPT_AS_REPR(prim::BitwiseNotNode, ReprPrintTIR);
-TVM_REGISTER_SCRIPT_AS_REPR(prim::MinNode, ReprPrintTIR);
-TVM_REGISTER_SCRIPT_AS_REPR(prim::MaxNode, ReprPrintTIR);
-TVM_REGISTER_SCRIPT_AS_REPR(prim::LTNode, ReprPrintTIR);
-TVM_REGISTER_SCRIPT_AS_REPR(prim::LENode, ReprPrintTIR);
-TVM_REGISTER_SCRIPT_AS_REPR(prim::EQNode, ReprPrintTIR);
-TVM_REGISTER_SCRIPT_AS_REPR(prim::NENode, ReprPrintTIR);
-TVM_REGISTER_SCRIPT_AS_REPR(prim::GTNode, ReprPrintTIR);
-TVM_REGISTER_SCRIPT_AS_REPR(prim::GENode, ReprPrintTIR);
-TVM_REGISTER_SCRIPT_AS_REPR(prim::AndNode, ReprPrintTIR);
-TVM_REGISTER_SCRIPT_AS_REPR(prim::OrNode, ReprPrintTIR);
-TVM_REGISTER_SCRIPT_AS_REPR(prim::NotNode, ReprPrintTIR);
-TVM_REGISTER_SCRIPT_AS_REPR(prim::SelectNode, ReprPrintTIR);
-TVM_REGISTER_SCRIPT_AS_REPR(prim::RampNode, ReprPrintTIR);
-TVM_REGISTER_SCRIPT_AS_REPR(prim::BroadcastNode, ReprPrintTIR);
-TVM_REGISTER_SCRIPT_AS_REPR(prim::LetNode, ReprPrintTIR);
-TVM_REGISTER_SCRIPT_AS_REPR(prim::ShuffleNode, ReprPrintTIR);
-TVM_REGISTER_SCRIPT_AS_REPR(te::CommReducerNode, ReprPrintTIR);
-TVM_REGISTER_SCRIPT_AS_REPR(tirx::IndexMapNode, ReprPrintTIR);
-TVM_REGISTER_SCRIPT_AS_REPR(te::ReduceNode, ReprPrintTIR);
-TVM_REGISTER_SCRIPT_AS_REPR(tirx::LambdaExprNode, ReprPrintTIR);
+TVM_FFI_STATIC_INIT_BLOCK() {
+ TVMScriptPrinter::Register<tirx::IterVarNode>(ReprPrintTIR);
+ TVMScriptPrinter::Register<StringImmNode>(ReprPrintTIR);
+ TVMScriptPrinter::Register<prim::CastNode>(ReprPrintTIR);
+ TVMScriptPrinter::Register<prim::AddNode>(ReprPrintTIR);
+ TVMScriptPrinter::Register<prim::SubNode>(ReprPrintTIR);
+ TVMScriptPrinter::Register<prim::MulNode>(ReprPrintTIR);
+ TVMScriptPrinter::Register<prim::DivNode>(ReprPrintTIR);
+ TVMScriptPrinter::Register<prim::ModNode>(ReprPrintTIR);
+ TVMScriptPrinter::Register<prim::FloorDivNode>(ReprPrintTIR);
+ TVMScriptPrinter::Register<prim::FloorModNode>(ReprPrintTIR);
+ TVMScriptPrinter::Register<prim::LShiftNode>(ReprPrintTIR);
+ TVMScriptPrinter::Register<prim::RShiftNode>(ReprPrintTIR);
+ TVMScriptPrinter::Register<prim::BitwiseAndNode>(ReprPrintTIR);
+ TVMScriptPrinter::Register<prim::BitwiseOrNode>(ReprPrintTIR);
+ TVMScriptPrinter::Register<prim::BitwiseXorNode>(ReprPrintTIR);
+ TVMScriptPrinter::Register<prim::BitwiseNotNode>(ReprPrintTIR);
+ TVMScriptPrinter::Register<prim::MinNode>(ReprPrintTIR);
+ TVMScriptPrinter::Register<prim::MaxNode>(ReprPrintTIR);
+ TVMScriptPrinter::Register<prim::LTNode>(ReprPrintTIR);
+ TVMScriptPrinter::Register<prim::LENode>(ReprPrintTIR);
+ TVMScriptPrinter::Register<prim::EQNode>(ReprPrintTIR);
+ TVMScriptPrinter::Register<prim::NENode>(ReprPrintTIR);
+ TVMScriptPrinter::Register<prim::GTNode>(ReprPrintTIR);
+ TVMScriptPrinter::Register<prim::GENode>(ReprPrintTIR);
+ TVMScriptPrinter::Register<prim::AndNode>(ReprPrintTIR);
+ TVMScriptPrinter::Register<prim::OrNode>(ReprPrintTIR);
+ TVMScriptPrinter::Register<prim::NotNode>(ReprPrintTIR);
+ TVMScriptPrinter::Register<prim::SelectNode>(ReprPrintTIR);
+ TVMScriptPrinter::Register<prim::RampNode>(ReprPrintTIR);
+ TVMScriptPrinter::Register<prim::BroadcastNode>(ReprPrintTIR);
+ TVMScriptPrinter::Register<prim::LetNode>(ReprPrintTIR);
+ TVMScriptPrinter::Register<prim::ShuffleNode>(ReprPrintTIR);
+ TVMScriptPrinter::Register<te::CommReducerNode>(ReprPrintTIR);
+ TVMScriptPrinter::Register<tirx::IndexMapNode>(ReprPrintTIR);
+ TVMScriptPrinter::Register<te::ReduceNode>(ReprPrintTIR);
+ TVMScriptPrinter::Register<tirx::LambdaExprNode>(ReprPrintTIR);
+}
} // namespace printer
} // namespace script
diff --git a/src/tirx/script/printer/for_loop.cc
b/src/tirx/script/printer/for_loop.cc
index df1ec3a933..cd8d4da693 100644
--- a/src/tirx/script/printer/for_loop.cc
+++ b/src/tirx/script/printer/for_loop.cc
@@ -155,7 +155,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
});
}
-TVM_REGISTER_SCRIPT_AS_REPR(tirx::ForNode, ReprPrintTIR);
+TVM_FFI_STATIC_INIT_BLOCK() {
TVMScriptPrinter::Register<tirx::ForNode>(ReprPrintTIR); }
} // namespace printer
} // namespace script
diff --git a/src/tirx/script/printer/function.cc
b/src/tirx/script/printer/function.cc
index ab613b46e0..603e03ede8 100644
--- a/src/tirx/script/printer/function.cc
+++ b/src/tirx/script/printer/function.cc
@@ -229,7 +229,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
});
}
-TVM_REGISTER_SCRIPT_AS_REPR(tirx::PrimFuncNode, ReprPrintTIR);
+TVM_FFI_STATIC_INIT_BLOCK() {
TVMScriptPrinter::Register<tirx::PrimFuncNode>(ReprPrintTIR); }
TVM_FFI_STATIC_INIT_BLOCK() {
IRDocsifier::vtable().set_dispatch<tvm::GlobalVar>( //
diff --git a/src/tirx/script/printer/ir.cc b/src/tirx/script/printer/ir.cc
index 3b5044b463..c299869c6c 100644
--- a/src/tirx/script/printer/ir.cc
+++ b/src/tirx/script/printer/ir.cc
@@ -81,11 +81,13 @@ TVM_FFI_STATIC_INIT_BLOCK() {
});
}
-TVM_REGISTER_SCRIPT_AS_REPR(IntImmNode, ReprPrintTIR);
-TVM_REGISTER_SCRIPT_AS_REPR(FloatImmNode, ReprPrintTIR);
-TVM_REGISTER_SCRIPT_AS_REPR(PrimTypeNode, ReprPrintTIR);
-TVM_REGISTER_SCRIPT_AS_REPR(PointerTypeNode, ReprPrintTIR);
-TVM_REGISTER_SCRIPT_AS_REPR(TupleTypeNode, ReprPrintTIR);
+TVM_FFI_STATIC_INIT_BLOCK() {
+ TVMScriptPrinter::Register<IntImmNode>(ReprPrintTIR);
+ TVMScriptPrinter::Register<FloatImmNode>(ReprPrintTIR);
+ TVMScriptPrinter::Register<PrimTypeNode>(ReprPrintTIR);
+ TVMScriptPrinter::Register<PointerTypeNode>(ReprPrintTIR);
+ TVMScriptPrinter::Register<TupleTypeNode>(ReprPrintTIR);
+}
} // namespace printer
} // namespace script
diff --git a/src/tirx/script/printer/stmt.cc b/src/tirx/script/printer/stmt.cc
index b0d6e6159c..1c2a14b81e 100644
--- a/src/tirx/script/printer/stmt.cc
+++ b/src/tirx/script/printer/stmt.cc
@@ -144,7 +144,9 @@ TVM_FFI_STATIC_INIT_BLOCK() {
d->AsDoc<DictDoc>(op_call->config,
p->Attr("config")), disp);
});
}
-TVM_REGISTER_SCRIPT_AS_REPR(tirx::TilePrimitiveCallNode, ReprPrintTIR);
+TVM_FFI_STATIC_INIT_BLOCK() {
+ TVMScriptPrinter::Register<tirx::TilePrimitiveCallNode>(ReprPrintTIR);
+}
TVM_FFI_STATIC_INIT_BLOCK() {
IRDocsifier::vtable().set_dispatch<tirx::Evaluate>(
@@ -840,18 +842,20 @@ TVM_FFI_STATIC_INIT_BLOCK() {
});
}
-TVM_REGISTER_SCRIPT_AS_REPR(tirx::BindNode, ReprPrintTIR);
-TVM_REGISTER_SCRIPT_AS_REPR(tirx::AttrStmtNode, ReprPrintTIR);
-TVM_REGISTER_SCRIPT_AS_REPR(tirx::AssertStmtNode, ReprPrintTIR);
-TVM_REGISTER_SCRIPT_AS_REPR(tirx::WhileNode, ReprPrintTIR);
-TVM_REGISTER_SCRIPT_AS_REPR(tirx::AllocBufferNode, ReprPrintTIR);
-TVM_REGISTER_SCRIPT_AS_REPR(tirx::ReturnNode, ReprPrintTIR);
-TVM_REGISTER_SCRIPT_AS_REPR(tirx::BreakNode, ReprPrintTIR);
-TVM_REGISTER_SCRIPT_AS_REPR(tirx::ContinueNode, ReprPrintTIR);
-TVM_REGISTER_SCRIPT_AS_REPR(tirx::DeclBufferNode, ReprPrintTIR);
-TVM_REGISTER_SCRIPT_AS_REPR(tirx::SeqStmtNode, ReprPrintTIR);
-TVM_REGISTER_SCRIPT_AS_REPR(tirx::IfThenElseNode, ReprPrintTIR);
-TVM_REGISTER_SCRIPT_AS_REPR(tirx::EvaluateNode, ReprPrintTIR);
+TVM_FFI_STATIC_INIT_BLOCK() {
+ TVMScriptPrinter::Register<tirx::BindNode>(ReprPrintTIR);
+ TVMScriptPrinter::Register<tirx::AttrStmtNode>(ReprPrintTIR);
+ TVMScriptPrinter::Register<tirx::AssertStmtNode>(ReprPrintTIR);
+ TVMScriptPrinter::Register<tirx::WhileNode>(ReprPrintTIR);
+ TVMScriptPrinter::Register<tirx::AllocBufferNode>(ReprPrintTIR);
+ TVMScriptPrinter::Register<tirx::ReturnNode>(ReprPrintTIR);
+ TVMScriptPrinter::Register<tirx::BreakNode>(ReprPrintTIR);
+ TVMScriptPrinter::Register<tirx::ContinueNode>(ReprPrintTIR);
+ TVMScriptPrinter::Register<tirx::DeclBufferNode>(ReprPrintTIR);
+ TVMScriptPrinter::Register<tirx::SeqStmtNode>(ReprPrintTIR);
+ TVMScriptPrinter::Register<tirx::IfThenElseNode>(ReprPrintTIR);
+ TVMScriptPrinter::Register<tirx::EvaluateNode>(ReprPrintTIR);
+}
} // namespace printer
} // namespace script
} // namespace tvm
