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 141a187f35 [REFACTOR][IR] Add structural hooks to remaining Type and
Relax Expr nodes (#20302)
141a187f35 is described below
commit 141a187f35e9b6263db6665bc99eb03146d9b23f
Author: Tianqi Chen <[email protected]>
AuthorDate: Thu Sep 10 10:37:26 2026 -0400
[REFACTOR][IR] Add structural hooks to remaining Type and Relax Expr nodes
(#20302)
Add explicit structural traversal hooks to the remaining concrete Type
and Relax Expr families so generic structural map and walk operations
preserve node identity and field semantics.
- Add visit, mutate, and maybe-in-place mutate hooks for the remaining
reachable concrete Type and Expr nodes in scope, including
DataflowVarNode.
- Preserve unchanged-node identity with ffi::Unchanged() and retain
existing definition-region semantics.
- Keep internal arithmetic temporaries and other excluded families out
of scope; concrete Stmt registrations were already complete.
- Rebase onto current upstream main.
Validation:
- pre-commit checks pass for all modified source files.
- Full existing C++ cpptest suite passes 132/132.
---
src/ir/expr.cc | 28 ++-
src/ir/op.cc | 25 +++
src/ir/type.cc | 189 ++++++++++++++++++-
src/relax/distributed/type.cc | 64 +++++++
src/relax/ir/dependent_type.cc | 210 +++++++++++++++++++--
src/relax/ir/expr.cc | 408 +++++++++++++++++++++++++++++++++++++++++
src/relax/ir/type.cc | 28 ++-
src/tirx/ir/function.cc | 73 ++++++++
src/tirx/ir/stmt.cc | 20 +-
9 files changed, 1022 insertions(+), 23 deletions(-)
diff --git a/src/ir/expr.cc b/src/ir/expr.cc
index 1af350c70d..b9b51da089 100644
--- a/src/ir/expr.cc
+++ b/src/ir/expr.cc
@@ -278,6 +278,8 @@ TVMFFIAny
RangeMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyVi
return ffi::Unchanged().CopyToTVMFFIAny();
}
+// DataflowVarNode duplicates this protocol because structural hooks do not
inherit. Keep the two
+// hook triples in lockstep when changing remap, PrimType-skip, or
definition-region behavior.
TVMFFIAny VarVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value)
noexcept {
// skips: name
const VarNode* self =
@@ -381,6 +383,22 @@ TVMFFIAny VarMaybeInplaceMutate(ffi::StructuralMutatorObj*
mutator, ffi::AnyView
return ffi::details::UnchangedOrUnsafe::MoveToTVMFFIAny(std::move(result));
}
+TVMFFIAny GlobalVarVisit(ffi::StructuralVisitorObj*, ffi::AnyView) noexcept {
+ // GlobalVar is a module-level symbol. name_hint is scalar identity and ty
is derived from the
+ // referenced function, matching GlobalVarNode's custom structural
equality/hash definition.
+ // It has no definition site where this hook could establish a VarRemap. A
callback that renames
+ // GlobalVars is therefore responsible for returning one stable replacement
per module symbol.
+ return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+}
+
+TVMFFIAny GlobalVarMutate(ffi::StructuralMutatorObj*, ffi::AnyView) noexcept {
+ return ffi::Unchanged().CopyToTVMFFIAny();
+}
+
+TVMFFIAny GlobalVarMaybeInplaceMutate(ffi::StructuralMutatorObj*,
ffi::AnyView) noexcept {
+ return ffi::Unchanged().CopyToTVMFFIAny();
+}
+
TVMFFIAny CallVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value)
noexcept {
// skips: attrs, constant metadata left untouched like the classic Expr
functors.
const CallNode* self =
@@ -859,7 +877,15 @@ GlobalVar::GlobalVar(ffi::String name_hint, Span span) {
data_ = std::move(n);
}
-TVM_FFI_STATIC_INIT_BLOCK() { GlobalVarNode::RegisterReflection(); }
+TVM_FFI_STATIC_INIT_BLOCK() {
+ namespace refl = tvm::ffi::reflection;
+ GlobalVarNode::RegisterReflection();
+ refl::TypeAttrDef<GlobalVarNode>()
+ .attr(refl::type_attr::kStructuralVisit,
reinterpret_cast<void*>(&GlobalVarVisit))
+ .attr(refl::type_attr::kStructuralMutate,
reinterpret_cast<void*>(&GlobalVarMutate))
+ .attr(refl::type_attr::kStructuralMaybeInplaceMutate,
+ reinterpret_cast<void*>(&GlobalVarMaybeInplaceMutate));
+}
// Call
Call::Call(Type ret_ty, Expr op, ffi::Array<Expr> args, Attrs attrs,
ffi::Array<Type> ty_args,
diff --git a/src/ir/op.cc b/src/ir/op.cc
index 560cee6c9d..c300806b74 100644
--- a/src/ir/op.cc
+++ b/src/ir/op.cc
@@ -21,6 +21,8 @@
* \file src/ir/op.cc
* \brief Primitive operators and intrinsics.
*/
+#include <tvm/ffi/extra/structural_mutate.h>
+#include <tvm/ffi/extra/structural_visit.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/op.h>
@@ -33,9 +35,32 @@
namespace tvm {
+namespace {
+
+TVMFFIAny OpVisit(ffi::StructuralVisitorObj*, ffi::AnyView) noexcept {
+ // Ops are unique registry atoms. Avoid reflecting through their registry
metadata.
+ return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+}
+
+TVMFFIAny OpMutate(ffi::StructuralMutatorObj*, ffi::AnyView) noexcept {
+ return ffi::Unchanged().CopyToTVMFFIAny();
+}
+
+TVMFFIAny OpMaybeInplaceMutate(ffi::StructuralMutatorObj*, ffi::AnyView)
noexcept {
+ return ffi::Unchanged().CopyToTVMFFIAny();
+}
+
+} // namespace
+
TVM_FFI_STATIC_INIT_BLOCK() {
+ namespace refl = tvm::ffi::reflection;
ArgumentInfoNode::RegisterReflection();
OpNode::RegisterReflection();
+ refl::TypeAttrDef<OpNode>()
+ .attr(refl::type_attr::kStructuralVisit,
reinterpret_cast<void*>(&OpVisit))
+ .attr(refl::type_attr::kStructuralMutate,
reinterpret_cast<void*>(&OpMutate))
+ .attr(refl::type_attr::kStructuralMaybeInplaceMutate,
+ reinterpret_cast<void*>(&OpMaybeInplaceMutate));
}
using ffi::Any;
diff --git a/src/ir/type.cc b/src/ir/type.cc
index d4a56cb476..1947f48985 100644
--- a/src/ir/type.cc
+++ b/src/ir/type.cc
@@ -69,6 +69,32 @@ ffi::ObjectPtr<PrimTypeNode>
GetCachedPrimTypeNode(DLDataType dtype) {
// Structural traversal hooks
+TVMFFIAny TypeVisit(ffi::StructuralVisitorObj*, ffi::AnyView) noexcept {
+ // Type::Missing() is the only concrete TypeNode value; span is ignored
debug metadata.
+ return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+}
+
+TVMFFIAny TypeMutate(ffi::StructuralMutatorObj*, ffi::AnyView) noexcept {
+ return ffi::Unchanged().CopyToTVMFFIAny();
+}
+
+TVMFFIAny TypeMaybeInplaceMutate(ffi::StructuralMutatorObj*, ffi::AnyView)
noexcept {
+ return ffi::Unchanged().CopyToTVMFFIAny();
+}
+
+TVMFFIAny OpaqueTypeVisit(ffi::StructuralVisitorObj*, ffi::AnyView) noexcept {
+ // OpaqueType is a field-less construction-time marker; span is ignored
debug metadata.
+ return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+}
+
+TVMFFIAny OpaqueTypeMutate(ffi::StructuralMutatorObj*, ffi::AnyView) noexcept {
+ return ffi::Unchanged().CopyToTVMFFIAny();
+}
+
+TVMFFIAny OpaqueTypeMaybeInplaceMutate(ffi::StructuralMutatorObj*,
ffi::AnyView) noexcept {
+ return ffi::Unchanged().CopyToTVMFFIAny();
+}
+
TVMFFIAny PrimTypeVisit(ffi::StructuralVisitorObj*, ffi::AnyView) noexcept {
// dtype is a constant: reflected for StructuralEqual/Hash,
// not traversed by the visitor/mutator contract.
@@ -87,6 +113,129 @@ TVMFFIAny
PrimTypeMaybeInplaceMutate(ffi::StructuralMutatorObj*, ffi::AnyView) n
return ffi::Unchanged().CopyToTVMFFIAny();
}
+TVMFFIAny PointerTypeVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView
value) noexcept {
+ // skips: storage_scope (scalar)
+ const PointerTypeNode* self =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
PointerTypeNode>(value);
+
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->element_type));
+ return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+}
+
+TVMFFIAny PointerTypeMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView
value) noexcept {
+ // skips: storage_scope (scalar)
+ const PointerTypeNode* self =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
PointerTypeNode>(value);
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Type>,
mapped_element_type,
+
mutator->MutateExpected(self->element_type));
+ if (mapped_element_type.UnchangedOrSameAs(self->element_type)) {
+ return ffi::Unchanged().CopyToTVMFFIAny();
+ }
+ ffi::ObjectPtr<PointerTypeNode> copy =
ffi::make_object<PointerTypeNode>(*self);
+ copy->element_type =
+
std::move(mapped_element_type).ValueOrUnchanged(std::move(copy->element_type));
+ return
ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(copy)));
+}
+
+TVMFFIAny PointerTypeMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
+ ffi::AnyView value) noexcept {
+ // skips: storage_scope (scalar)
+ PointerTypeNode* self = const_cast<PointerTypeNode*>(
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
PointerTypeNode>(value));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(
+ ffi::UnchangedOr<Type>, mapped_element_type,
+ mutator->MaybeInplaceMutateIfUniqueExpected(self->element_type));
+ if (!mapped_element_type.UnchangedOrSameAs(self->element_type)) {
+ self->element_type = std::move(mapped_element_type).ValueUnchecked();
+ }
+ return ffi::Unchanged().CopyToTVMFFIAny();
+}
+
+TVMFFIAny FuncTypeVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView
value) noexcept {
+ const FuncTypeNode* self =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
FuncTypeNode>(value);
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->arg_types));
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->ret_type));
+ return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+}
+
+TVMFFIAny FuncTypeMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView
value) noexcept {
+ const FuncTypeNode* self =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
FuncTypeNode>(value);
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<ffi::Array<Type>>,
mapped_arg_types,
+ mutator->MutateExpected(self->arg_types));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Type>, mapped_ret_type,
+ mutator->MutateExpected(self->ret_type));
+ if (mapped_arg_types.UnchangedOrSameAs(self->arg_types) &&
+ mapped_ret_type.UnchangedOrSameAs(self->ret_type)) {
+ return ffi::Unchanged().CopyToTVMFFIAny();
+ }
+ ffi::ObjectPtr<FuncTypeNode> copy = ffi::make_object<FuncTypeNode>(*self);
+ copy->arg_types =
std::move(mapped_arg_types).ValueOrUnchanged(std::move(copy->arg_types));
+ copy->ret_type =
std::move(mapped_ret_type).ValueOrUnchanged(std::move(copy->ret_type));
+ return
ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(copy)));
+}
+
+TVMFFIAny FuncTypeMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
+ ffi::AnyView value) noexcept {
+ FuncTypeNode* self = const_cast<FuncTypeNode*>(
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
FuncTypeNode>(value));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<ffi::Array<Type>>,
mapped_arg_types,
+
mutator->MaybeInplaceMutateIfUniqueExpected(self->arg_types));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Type>, mapped_ret_type,
+
mutator->MaybeInplaceMutateIfUniqueExpected(self->ret_type));
+ if (!mapped_arg_types.UnchangedOrSameAs(self->arg_types)) {
+ self->arg_types = std::move(mapped_arg_types).ValueUnchecked();
+ }
+ if (!mapped_ret_type.UnchangedOrSameAs(self->ret_type)) {
+ self->ret_type = std::move(mapped_ret_type).ValueUnchecked();
+ }
+ return ffi::Unchanged().CopyToTVMFFIAny();
+}
+
+TVMFFIAny TupleTypeVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView
value) noexcept {
+ const TupleTypeNode* self =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
TupleTypeNode>(value);
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->fields));
+ return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+}
+
+TVMFFIAny TupleTypeMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView
value) noexcept {
+ const TupleTypeNode* self =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
TupleTypeNode>(value);
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<ffi::Array<Type>>,
mapped_fields,
+ mutator->MutateExpected(self->fields));
+ if (mapped_fields.UnchangedOrSameAs(self->fields)) {
+ return ffi::Unchanged().CopyToTVMFFIAny();
+ }
+ ffi::ObjectPtr<TupleTypeNode> copy = ffi::make_object<TupleTypeNode>(*self);
+ copy->fields =
std::move(mapped_fields).ValueOrUnchanged(std::move(copy->fields));
+ return
ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(copy)));
+}
+
+TVMFFIAny TupleTypeMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
+ ffi::AnyView value) noexcept {
+ TupleTypeNode* self = const_cast<TupleTypeNode*>(
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
TupleTypeNode>(value));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<ffi::Array<Type>>,
mapped_fields,
+
mutator->MaybeInplaceMutateIfUniqueExpected(self->fields));
+ if (!mapped_fields.UnchangedOrSameAs(self->fields)) {
+ self->fields = std::move(mapped_fields).ValueUnchecked();
+ }
+ return ffi::Unchanged().CopyToTVMFFIAny();
+}
+
+TVMFFIAny TensorMapTypeVisit(ffi::StructuralVisitorObj*, ffi::AnyView)
noexcept {
+ return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+}
+
+TVMFFIAny TensorMapTypeMutate(ffi::StructuralMutatorObj*, ffi::AnyView)
noexcept {
+ return ffi::Unchanged().CopyToTVMFFIAny();
+}
+
+TVMFFIAny TensorMapTypeMaybeInplaceMutate(ffi::StructuralMutatorObj*,
ffi::AnyView) noexcept {
+ return ffi::Unchanged().CopyToTVMFFIAny();
+}
+
} // namespace
Type Type::Missing() {
@@ -102,9 +251,25 @@ bool Type::IsMissing() const { return
this->same_as(Type::Missing()); }
OpaqueType::OpaqueType() : Type(ffi::UnsafeInit{}) { data_ =
ffi::make_object<OpaqueTypeNode>(); }
-TVM_FFI_STATIC_INIT_BLOCK() { TypeNode::RegisterReflection(); }
+TVM_FFI_STATIC_INIT_BLOCK() {
+ namespace refl = tvm::ffi::reflection;
+ TypeNode::RegisterReflection();
+ refl::TypeAttrDef<TypeNode>()
+ .attr(refl::type_attr::kStructuralVisit,
reinterpret_cast<void*>(&TypeVisit))
+ .attr(refl::type_attr::kStructuralMutate,
reinterpret_cast<void*>(&TypeMutate))
+ .attr(refl::type_attr::kStructuralMaybeInplaceMutate,
+ reinterpret_cast<void*>(&TypeMaybeInplaceMutate));
+}
-TVM_FFI_STATIC_INIT_BLOCK() { OpaqueTypeNode::RegisterReflection(); }
+TVM_FFI_STATIC_INIT_BLOCK() {
+ namespace refl = tvm::ffi::reflection;
+ OpaqueTypeNode::RegisterReflection();
+ refl::TypeAttrDef<OpaqueTypeNode>()
+ .attr(refl::type_attr::kStructuralVisit,
reinterpret_cast<void*>(&OpaqueTypeVisit))
+ .attr(refl::type_attr::kStructuralMutate,
reinterpret_cast<void*>(&OpaqueTypeMutate))
+ .attr(refl::type_attr::kStructuralMaybeInplaceMutate,
+ reinterpret_cast<void*>(&OpaqueTypeMaybeInplaceMutate));
+}
// PrimType
PrimType::PrimType(DLDataType dtype) : Type(ffi::UnsafeInit{}) {
@@ -227,6 +392,11 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("ir.PointerType", [](Type element_type, ffi::String
storage_scope = "") {
return PointerType(element_type, storage_scope);
});
+ refl::TypeAttrDef<PointerTypeNode>()
+ .attr(refl::type_attr::kStructuralVisit,
reinterpret_cast<void*>(&PointerTypeVisit))
+ .attr(refl::type_attr::kStructuralMutate,
reinterpret_cast<void*>(&PointerTypeMutate))
+ .attr(refl::type_attr::kStructuralMaybeInplaceMutate,
+ reinterpret_cast<void*>(&PointerTypeMaybeInplaceMutate));
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -235,6 +405,11 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def("ir.FuncType", [](tvm::ffi::Array<Type> arg_types,
Type ret_type) {
return FuncType(arg_types, ret_type);
});
+ refl::TypeAttrDef<FuncTypeNode>()
+ .attr(refl::type_attr::kStructuralVisit,
reinterpret_cast<void*>(&FuncTypeVisit))
+ .attr(refl::type_attr::kStructuralMutate,
reinterpret_cast<void*>(&FuncTypeMutate))
+ .attr(refl::type_attr::kStructuralMaybeInplaceMutate,
+ reinterpret_cast<void*>(&FuncTypeMaybeInplaceMutate));
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -245,6 +420,16 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.def("ir.TupleType",
[](ffi::Array<Type> fields, Span span) { return TupleType(fields,
span); })
.def("ir.TensorMapType", [](Span span) { return TensorMapType(span); });
+ refl::TypeAttrDef<TupleTypeNode>()
+ .attr(refl::type_attr::kStructuralVisit,
reinterpret_cast<void*>(&TupleTypeVisit))
+ .attr(refl::type_attr::kStructuralMutate,
reinterpret_cast<void*>(&TupleTypeMutate))
+ .attr(refl::type_attr::kStructuralMaybeInplaceMutate,
+ reinterpret_cast<void*>(&TupleTypeMaybeInplaceMutate));
+ refl::TypeAttrDef<TensorMapTypeNode>()
+ .attr(refl::type_attr::kStructuralVisit,
reinterpret_cast<void*>(&TensorMapTypeVisit))
+ .attr(refl::type_attr::kStructuralMutate,
reinterpret_cast<void*>(&TensorMapTypeMutate))
+ .attr(refl::type_attr::kStructuralMaybeInplaceMutate,
+ reinterpret_cast<void*>(&TensorMapTypeMaybeInplaceMutate));
}
} // namespace tvm
diff --git a/src/relax/distributed/type.cc b/src/relax/distributed/type.cc
index 8fb164c2a8..0a48faad2c 100644
--- a/src/relax/distributed/type.cc
+++ b/src/relax/distributed/type.cc
@@ -22,16 +22,80 @@
* \brief Relax DTensor type.
*/
+#include <tvm/ffi/extra/structural_mutate.h>
+#include <tvm/ffi/extra/structural_visit.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/relax/distributed/type.h>
namespace tvm {
namespace relax {
namespace distributed {
+namespace {
+
+TVMFFIAny DTensorTypeVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView
value) noexcept {
+ const DTensorTypeNode* self =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
DTensorTypeNode>(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));
+ return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+}
+
+TVMFFIAny DTensorTypeMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView
value) noexcept {
+ const DTensorTypeNode* self =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
DTensorTypeNode>(value);
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<DeviceMesh>,
mapped_device_mesh,
+
mutator->MutateExpected(self->device_mesh));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Placement>,
mapped_placement,
+ mutator->MutateExpected(self->placement));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<TensorType>,
mapped_tensor_ty,
+ mutator->MutateExpected(self->tensor_ty));
+ if (mapped_device_mesh.UnchangedOrSameAs(self->device_mesh) &&
+ mapped_placement.UnchangedOrSameAs(self->placement) &&
+ mapped_tensor_ty.UnchangedOrSameAs(self->tensor_ty)) {
+ return ffi::Unchanged().CopyToTVMFFIAny();
+ }
+ ffi::ObjectPtr<DTensorTypeNode> copy =
ffi::make_object<DTensorTypeNode>(*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));
+ return
ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(copy)));
+}
+
+TVMFFIAny DTensorTypeMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
+ ffi::AnyView value) noexcept {
+ DTensorTypeNode* self = const_cast<DTensorTypeNode*>(
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
DTensorTypeNode>(value));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<DeviceMesh>,
mapped_device_mesh,
+
mutator->MaybeInplaceMutateIfUniqueExpected(self->device_mesh));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Placement>,
mapped_placement,
+
mutator->MaybeInplaceMutateIfUniqueExpected(self->placement));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<TensorType>,
mapped_tensor_ty,
+
mutator->MaybeInplaceMutateIfUniqueExpected(self->tensor_ty));
+ if (!mapped_device_mesh.UnchangedOrSameAs(self->device_mesh)) {
+ self->device_mesh = std::move(mapped_device_mesh).ValueUnchecked();
+ }
+ if (!mapped_placement.UnchangedOrSameAs(self->placement)) {
+ self->placement = std::move(mapped_placement).ValueUnchecked();
+ }
+ if (!mapped_tensor_ty.UnchangedOrSameAs(self->tensor_ty)) {
+ self->tensor_ty = std::move(mapped_tensor_ty).ValueUnchecked();
+ }
+ return ffi::Unchanged().CopyToTVMFFIAny();
+}
+
+} // namespace
+
TVM_FFI_STATIC_INIT_BLOCK() {
+ namespace refl = tvm::ffi::reflection;
DTensorTypeNode::RegisterReflection();
PlacementNode::RegisterReflection();
PlacementSpecNode::RegisterReflection();
+ refl::TypeAttrDef<DTensorTypeNode>()
+ .attr(refl::type_attr::kStructuralVisit,
reinterpret_cast<void*>(&DTensorTypeVisit))
+ .attr(refl::type_attr::kStructuralMutate,
reinterpret_cast<void*>(&DTensorTypeMutate))
+ .attr(refl::type_attr::kStructuralMaybeInplaceMutate,
+ reinterpret_cast<void*>(&DTensorTypeMaybeInplaceMutate));
}
PlacementSpec PlacementSpec::Sharding(int axis) {
diff --git a/src/relax/ir/dependent_type.cc b/src/relax/ir/dependent_type.cc
index ff0bce1fb6..14755ca573 100644
--- a/src/relax/ir/dependent_type.cc
+++ b/src/relax/ir/dependent_type.cc
@@ -21,6 +21,8 @@
* \file src/relax/ir/dependent_type.cc
* \brief Relax type nodes.
*/
+#include <tvm/ffi/extra/structural_mutate.h>
+#include <tvm/ffi/extra/structural_visit.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/relax/analysis.h>
@@ -30,29 +32,203 @@
namespace tvm {
namespace relax {
+namespace {
+
+TVMFFIAny AnyTypeVisit(ffi::StructuralVisitorObj*, ffi::AnyView) noexcept {
+ return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+}
+
+TVMFFIAny AnyTypeMutate(ffi::StructuralMutatorObj*, ffi::AnyView) noexcept {
+ return ffi::Unchanged().CopyToTVMFFIAny();
+}
+
+TVMFFIAny AnyTypeMaybeInplaceMutate(ffi::StructuralMutatorObj*, ffi::AnyView)
noexcept {
+ return ffi::Unchanged().CopyToTVMFFIAny();
+}
+
+TVMFFIAny ShapeTypeVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView
value) noexcept {
+ // skips: ndim (scalar)
+ const ShapeTypeNode* self =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
ShapeTypeNode>(value);
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->values));
+ return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+}
+
+TVMFFIAny ShapeTypeMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView
value) noexcept {
+ // skips: ndim (scalar)
+ const ShapeTypeNode* self =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
ShapeTypeNode>(value);
+
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<ffi::Optional<ffi::Array<PrimExpr>>>,
+ mapped_values,
mutator->MutateExpected(self->values));
+ if (mapped_values.UnchangedOrSameAs(self->values)) return
ffi::Unchanged().CopyToTVMFFIAny();
+ ffi::ObjectPtr<ShapeTypeNode> copy = ffi::make_object<ShapeTypeNode>(*self);
+ copy->values =
std::move(mapped_values).ValueOrUnchanged(std::move(copy->values));
+ return
ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(copy)));
+}
+
+TVMFFIAny ShapeTypeMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
+ ffi::AnyView value) noexcept {
+ // skips: ndim (scalar)
+ ShapeTypeNode* self = const_cast<ShapeTypeNode*>(
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
ShapeTypeNode>(value));
+
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<ffi::Optional<ffi::Array<PrimExpr>>>,
+ mapped_values,
+
mutator->MaybeInplaceMutateIfUniqueExpected(self->values));
+ if (!mapped_values.UnchangedOrSameAs(self->values)) {
+ self->values = std::move(mapped_values).ValueUnchecked();
+ }
+ return ffi::Unchanged().CopyToTVMFFIAny();
+}
+
+TVMFFIAny TensorTypeVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView
value) noexcept {
+ // skips: ndim (scalar)
+ const TensorTypeNode* self =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
TensorTypeNode>(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));
+ return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+}
+
+TVMFFIAny TensorTypeMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView
value) noexcept {
+ // skips: ndim (scalar)
+ const TensorTypeNode* self =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
TensorTypeNode>(value);
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<ffi::Optional<Expr>>,
mapped_shape,
+ mutator->MutateExpected(self->shape));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<ffi::Optional<PrimType>>,
mapped_dtype,
+ mutator->MutateExpected(self->dtype));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<ffi::Optional<VDevice>>,
mapped_vdevice,
+ mutator->MutateExpected(self->vdevice));
+ if (mapped_shape.UnchangedOrSameAs(self->shape) &&
mapped_dtype.UnchangedOrSameAs(self->dtype) &&
+ mapped_vdevice.UnchangedOrSameAs(self->vdevice)) {
+ return ffi::Unchanged().CopyToTVMFFIAny();
+ }
+ ffi::ObjectPtr<TensorTypeNode> copy =
ffi::make_object<TensorTypeNode>(*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));
+ return
ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(copy)));
+}
+
+TVMFFIAny TensorTypeMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
+ ffi::AnyView value) noexcept {
+ // skips: ndim (scalar)
+ TensorTypeNode* self = const_cast<TensorTypeNode*>(
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
TensorTypeNode>(value));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<ffi::Optional<Expr>>,
mapped_shape,
+
mutator->MaybeInplaceMutateIfUniqueExpected(self->shape));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<ffi::Optional<PrimType>>,
mapped_dtype,
+
mutator->MaybeInplaceMutateIfUniqueExpected(self->dtype));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<ffi::Optional<VDevice>>,
mapped_vdevice,
+
mutator->MaybeInplaceMutateIfUniqueExpected(self->vdevice));
+ if (!mapped_shape.UnchangedOrSameAs(self->shape)) {
+ self->shape = std::move(mapped_shape).ValueUnchecked();
+ }
+ if (!mapped_dtype.UnchangedOrSameAs(self->dtype)) {
+ self->dtype = std::move(mapped_dtype).ValueUnchecked();
+ }
+ if (!mapped_vdevice.UnchangedOrSameAs(self->vdevice)) {
+ self->vdevice = std::move(mapped_vdevice).ValueUnchecked();
+ }
+ return ffi::Unchanged().CopyToTVMFFIAny();
+}
+
+TVMFFIAny FuncTypeVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView
value) noexcept {
+ // skips: derive_func (environment-backed callable metadata), purity (scalar)
+ const FuncTypeNode* self =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
FuncTypeNode>(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));
+ return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+}
+
+TVMFFIAny FuncTypeMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView
value) noexcept {
+ // skips: derive_func (environment-backed callable metadata), purity (scalar)
+ const FuncTypeNode* self =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
FuncTypeNode>(value);
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(
+ ffi::UnchangedOr<ffi::Optional<ffi::Array<Type>>>, mapped_params,
+ mutator->WithDefRegionKind(kTVMFFIDefRegionKindPattern,
+ [&]() { return
mutator->MutateExpected(self->params); }));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Type>, mapped_ret,
+ mutator->MutateExpected(self->ret));
+ if (mapped_params.UnchangedOrSameAs(self->params) &&
mapped_ret.UnchangedOrSameAs(self->ret)) {
+ return ffi::Unchanged().CopyToTVMFFIAny();
+ }
+ ffi::ObjectPtr<FuncTypeNode> copy = ffi::make_object<FuncTypeNode>(*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::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(copy)));
+}
+
+TVMFFIAny FuncTypeMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
+ ffi::AnyView value) noexcept {
+ // skips: derive_func (environment-backed callable metadata), purity (scalar)
+ FuncTypeNode* self = const_cast<FuncTypeNode*>(
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
FuncTypeNode>(value));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(
+ ffi::UnchangedOr<ffi::Optional<ffi::Array<Type>>>, mapped_params,
+ mutator->WithDefRegionKind(kTVMFFIDefRegionKindPattern, [&]() {
+ return mutator->MaybeInplaceMutateIfUniqueExpected(self->params);
+ }));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Type>, mapped_ret,
+
mutator->MaybeInplaceMutateIfUniqueExpected(self->ret));
+ if (!mapped_params.UnchangedOrSameAs(self->params)) {
+ self->params = std::move(mapped_params).ValueUnchecked();
+ }
+ if (!mapped_ret.UnchangedOrSameAs(self->ret)) {
+ self->ret = std::move(mapped_ret).ValueUnchecked();
+ }
+ return ffi::Unchanged().CopyToTVMFFIAny();
+}
+
+} // namespace
+
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
AnyTypeNode::RegisterReflection();
ShapeTypeNode::RegisterReflection();
TensorTypeNode::RegisterReflection();
FuncTypeNode::RegisterReflection();
- refl::TypeAttrDef<TensorTypeNode>().def(
- "__subscript_expr_realize__",
- [](Expr value,
- ffi::Array<ffi::Variant<
- ffi::Tuple<ffi::Optional<PrimExpr>, ffi::Optional<PrimExpr>,
ffi::Optional<PrimExpr>>,
- PrimExpr>>
- slice,
- Span span) -> ffi::ObjectRef {
- TVM_FFI_CHECK_EQ(slice.size(), 1, IndexError)
- << "A Relax expression requires exactly one index";
- auto index = slice[0].as<PrimExpr>();
- TVM_FFI_CHECK(index.has_value(), TypeError) << "A Relax expression
requires a point index";
- const auto* imm = index.value().as<IntImmNode>();
- TVM_FFI_CHECK(imm != nullptr, TypeError)
- << "A Relax expression requires a constant integer index";
- return TupleGetItem(value, static_cast<int>(imm->value), span);
- });
+ refl::TypeAttrDef<AnyTypeNode>()
+ .attr(refl::type_attr::kStructuralVisit,
reinterpret_cast<void*>(&AnyTypeVisit))
+ .attr(refl::type_attr::kStructuralMutate,
reinterpret_cast<void*>(&AnyTypeMutate))
+ .attr(refl::type_attr::kStructuralMaybeInplaceMutate,
+ reinterpret_cast<void*>(&AnyTypeMaybeInplaceMutate));
+ refl::TypeAttrDef<ShapeTypeNode>()
+ .attr(refl::type_attr::kStructuralVisit,
reinterpret_cast<void*>(&ShapeTypeVisit))
+ .attr(refl::type_attr::kStructuralMutate,
reinterpret_cast<void*>(&ShapeTypeMutate))
+ .attr(refl::type_attr::kStructuralMaybeInplaceMutate,
+ reinterpret_cast<void*>(&ShapeTypeMaybeInplaceMutate));
+ refl::TypeAttrDef<TensorTypeNode>()
+ .attr(refl::type_attr::kStructuralVisit,
reinterpret_cast<void*>(&TensorTypeVisit))
+ .attr(refl::type_attr::kStructuralMutate,
reinterpret_cast<void*>(&TensorTypeMutate))
+ .attr(refl::type_attr::kStructuralMaybeInplaceMutate,
+ reinterpret_cast<void*>(&TensorTypeMaybeInplaceMutate))
+ .def("__subscript_expr_realize__",
+ [](Expr value,
+ ffi::Array<ffi::Variant<ffi::Tuple<ffi::Optional<PrimExpr>,
ffi::Optional<PrimExpr>,
+ ffi::Optional<PrimExpr>>,
+ PrimExpr>>
+ slice,
+ Span span) -> ffi::ObjectRef {
+ TVM_FFI_CHECK_EQ(slice.size(), 1, IndexError)
+ << "A Relax expression requires exactly one index";
+ auto index = slice[0].as<PrimExpr>();
+ TVM_FFI_CHECK(index.has_value(), TypeError)
+ << "A Relax expression requires a point index";
+ const auto* imm = index.value().as<IntImmNode>();
+ TVM_FFI_CHECK(imm != nullptr, TypeError)
+ << "A Relax expression requires a constant integer index";
+ return TupleGetItem(value, static_cast<int>(imm->value), span);
+ });
+ refl::TypeAttrDef<FuncTypeNode>()
+ .attr(refl::type_attr::kStructuralVisit,
reinterpret_cast<void*>(&FuncTypeVisit))
+ .attr(refl::type_attr::kStructuralMutate,
reinterpret_cast<void*>(&FuncTypeMutate))
+ .attr(refl::type_attr::kStructuralMaybeInplaceMutate,
+ reinterpret_cast<void*>(&FuncTypeMaybeInplaceMutate));
}
AnyType::AnyType(Span span) : Type(ffi::UnsafeInit{}) {
diff --git a/src/relax/ir/expr.cc b/src/relax/ir/expr.cc
index b21cdbad87..d841d7817a 100644
--- a/src/relax/ir/expr.cc
+++ b/src/relax/ir/expr.cc
@@ -16,6 +16,8 @@
* specific language governing permissions and limitations
* under the License.
*/
+#include <tvm/ffi/extra/structural_mutate.h>
+#include <tvm/ffi/extra/structural_visit.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/relax/analysis.h>
#include <tvm/relax/block_builder.h>
@@ -27,7 +29,360 @@
namespace tvm {
namespace relax {
+namespace {
+
+// Traverses only ExprNode::ty. The per-node payload is an intentional
constant leaf:
+// ConstantNode::data is tensor data; StringImmNode::value and
DataTypeImmNode::value are scalars;
+// ExternFuncNode::global_symbol is scalar and BaseFuncNode::attrs is metadata.
+template <typename TNode>
+TVMFFIAny TypeOnlyExprVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView
value) noexcept {
+ const TNode* self =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
TNode>(value);
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->ty));
+ return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+}
+
+template <typename TNode>
+TVMFFIAny TypeOnlyExprMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView
value) noexcept {
+ const TNode* self =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
TNode>(value);
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Type>, mapped_ty,
+ mutator->MutateExpected(self->ty));
+ if (mapped_ty.UnchangedOrSameAs(self->ty)) return
ffi::Unchanged().CopyToTVMFFIAny();
+ ffi::ObjectPtr<TNode> copy = ffi::make_object<TNode>(*self);
+ copy->ty = std::move(mapped_ty).ValueOrUnchanged(std::move(copy->ty));
+ return
ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(copy)));
+}
+
+template <typename TNode>
+TVMFFIAny TypeOnlyExprMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
+ ffi::AnyView value) noexcept {
+ TNode* self = const_cast<TNode*>(
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
TNode>(value));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Type>, mapped_ty,
+
mutator->MaybeInplaceMutateIfUniqueExpected(self->ty));
+ if (!mapped_ty.UnchangedOrSameAs(self->ty)) {
+ self->ty = std::move(mapped_ty).ValueUnchecked();
+ }
+ return ffi::Unchanged().CopyToTVMFFIAny();
+}
+
+TVMFFIAny ShapeExprVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView
value) noexcept {
+ const ShapeExprNode* self =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
ShapeExprNode>(value);
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->ty));
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->values));
+ return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+}
+
+TVMFFIAny ShapeExprMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView
value) noexcept {
+ const ShapeExprNode* self =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
ShapeExprNode>(value);
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Type>, mapped_ty,
+ mutator->MutateExpected(self->ty));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<ffi::Array<PrimExpr>>,
mapped_values,
+ mutator->MutateExpected(self->values));
+ if (mapped_ty.UnchangedOrSameAs(self->ty) &&
mapped_values.UnchangedOrSameAs(self->values)) {
+ return ffi::Unchanged().CopyToTVMFFIAny();
+ }
+ ffi::ObjectPtr<ShapeExprNode> copy = ffi::make_object<ShapeExprNode>(*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::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(copy)));
+}
+
+TVMFFIAny ShapeExprMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
+ ffi::AnyView value) noexcept {
+ ShapeExprNode* self = const_cast<ShapeExprNode*>(
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
ShapeExprNode>(value));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Type>, mapped_ty,
+
mutator->MaybeInplaceMutateIfUniqueExpected(self->ty));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<ffi::Array<PrimExpr>>,
mapped_values,
+
mutator->MaybeInplaceMutateIfUniqueExpected(self->values));
+ if (!mapped_ty.UnchangedOrSameAs(self->ty)) self->ty =
std::move(mapped_ty).ValueUnchecked();
+ if (!mapped_values.UnchangedOrSameAs(self->values)) {
+ self->values = std::move(mapped_values).ValueUnchecked();
+ }
+ return ffi::Unchanged().CopyToTVMFFIAny();
+}
+
+// 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.
+TVMFFIAny DataflowVarVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView
value) noexcept {
+ const DataflowVarNode* self =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
DataflowVarNode>(value);
+ if (!self->ty.as<PrimTypeNode>()) {
+ if (visitor->def_region_kind() == kTVMFFIDefRegionKindSimple) {
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->WithDefRegionKind(
+ kTVMFFIDefRegionKindNone, [&]() { return
visitor->VisitExpected(self->ty); }));
+ } else {
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->ty));
+ }
+ }
+ return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+}
+
+TVMFFIAny DataflowVarMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView
value) noexcept {
+ const DataflowVarNode* self =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
DataflowVarNode>(value);
+ ffi::Expected<ffi::Any> remap_result = mutator->VarRemapGetExpected(value);
+ TVM_FFI_S_MUTATE_MAYBE_EARLY_RETURN(remap_result);
+ if (ffi::details::ExpectedUnsafe::GetData(remap_result).type_index() !=
+ ffi::TypeIndex::kTVMFFINone) {
+ return
ffi::details::ExpectedUnsafe::MoveToTVMFFIAny(std::move(remap_result));
+ }
+ if (mutator->def_region_kind() == kTVMFFIDefRegionKindNone) {
+ return ffi::Unchanged().CopyToTVMFFIAny();
+ }
+ ffi::UnchangedOr<ffi::Any> result = ffi::Unchanged();
+ ffi::Any mapped_value;
+ if (!self->ty.as<PrimTypeNode>()) {
+ ffi::Expected<ffi::UnchangedOr<ffi::Any>> mapped_ty_result =
+ mutator->def_region_kind() == kTVMFFIDefRegionKindSimple
+ ? mutator->WithDefRegionKind(kTVMFFIDefRegionKindNone,
+ [&]() { return
mutator->MutateExpected(self->ty); })
+ : mutator->MutateExpected(self->ty);
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Type>, mapped_ty,
+ std::move(mapped_ty_result));
+ if (!mapped_ty.UnchangedOrSameAs(self->ty)) {
+ ffi::ObjectPtr<DataflowVarNode> copy =
ffi::make_object<DataflowVarNode>(*self);
+ copy->ty = std::move(mapped_ty).ValueUnchecked();
+ mapped_value = ffi::Any(std::move(copy));
+ result = mapped_value;
+ }
+ }
+ if (!result.IsUnchanged() || mutator->def_region_kind() ==
kTVMFFIDefRegionKindPattern) {
+ ffi::AnyView value_to_store = result.IsUnchanged() ? value :
ffi::AnyView(mapped_value);
+ auto set_result = mutator->VarRemapSetExpected(value, value_to_store);
+ if (TVM_FFI_PREDICT_FALSE(set_result.is_err())) {
+ return
ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(set_result).error()));
+ }
+ }
+ return ffi::details::UnchangedOrUnsafe::MoveToTVMFFIAny(std::move(result));
+}
+
+TVMFFIAny DataflowVarMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
+ ffi::AnyView value) noexcept {
+ DataflowVarNode* self = const_cast<DataflowVarNode*>(
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
DataflowVarNode>(value));
+ ffi::Expected<ffi::Any> remap_result = mutator->VarRemapGetExpected(value);
+ TVM_FFI_S_MUTATE_MAYBE_EARLY_RETURN(remap_result);
+ if (ffi::details::ExpectedUnsafe::GetData(remap_result).type_index() !=
+ ffi::TypeIndex::kTVMFFINone) {
+ return
ffi::details::ExpectedUnsafe::MoveToTVMFFIAny(std::move(remap_result));
+ }
+ if (mutator->def_region_kind() == kTVMFFIDefRegionKindNone) {
+ return ffi::Unchanged().CopyToTVMFFIAny();
+ }
+ ffi::UnchangedOr<ffi::Any> result = ffi::Unchanged();
+ ffi::Any mapped_value;
+ if (!self->ty.as<PrimTypeNode>()) {
+ ffi::Expected<ffi::UnchangedOr<ffi::Any>> mapped_ty_result =
+ mutator->def_region_kind() == kTVMFFIDefRegionKindSimple
+ ? mutator->WithDefRegionKind(
+ kTVMFFIDefRegionKindNone,
+ [&]() { return
mutator->MaybeInplaceMutateIfUniqueExpected(self->ty); })
+ : mutator->MaybeInplaceMutateIfUniqueExpected(self->ty);
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Type>, mapped_ty,
+ std::move(mapped_ty_result));
+ if (!mapped_ty.UnchangedOrSameAs(self->ty)) {
+ self->ty = std::move(mapped_ty).ValueUnchecked();
+ mapped_value = ffi::Any(self);
+ result = mapped_value;
+ }
+ }
+ if (!result.IsUnchanged() || mutator->def_region_kind() ==
kTVMFFIDefRegionKindPattern) {
+ ffi::AnyView value_to_store = result.IsUnchanged() ? value :
ffi::AnyView(mapped_value);
+ auto set_result = mutator->VarRemapSetExpected(value, value_to_store);
+ if (TVM_FFI_PREDICT_FALSE(set_result.is_err())) {
+ return
ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(set_result).error()));
+ }
+ }
+ return ffi::details::UnchangedOrUnsafe::MoveToTVMFFIAny(std::move(result));
+}
+
+TVMFFIAny SeqExprVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value)
noexcept {
+ const SeqExprNode* self =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
SeqExprNode>(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));
+ return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+}
+
+TVMFFIAny SeqExprMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView
value) noexcept {
+ const SeqExprNode* self =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
SeqExprNode>(value);
+
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<ffi::Array<BindingBlock>>,
mapped_blocks,
+ mutator->MutateExpected(self->blocks));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Type>, mapped_ty,
+ mutator->MutateExpected(self->ty));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Expr>, mapped_body,
+ mutator->MutateExpected(self->body));
+ if (mapped_blocks.UnchangedOrSameAs(self->blocks) &&
mapped_ty.UnchangedOrSameAs(self->ty) &&
+ mapped_body.UnchangedOrSameAs(self->body)) {
+ return ffi::Unchanged().CopyToTVMFFIAny();
+ }
+ ffi::ObjectPtr<SeqExprNode> copy = ffi::make_object<SeqExprNode>(*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));
+ return
ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(copy)));
+}
+
+TVMFFIAny SeqExprMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
+ ffi::AnyView value) noexcept {
+ SeqExprNode* self = const_cast<SeqExprNode*>(
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
SeqExprNode>(value));
+
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<ffi::Array<BindingBlock>>,
mapped_blocks,
+
mutator->MaybeInplaceMutateIfUniqueExpected(self->blocks));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Type>, mapped_ty,
+
mutator->MaybeInplaceMutateIfUniqueExpected(self->ty));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Expr>, mapped_body,
+
mutator->MaybeInplaceMutateIfUniqueExpected(self->body));
+ if (!mapped_blocks.UnchangedOrSameAs(self->blocks)) {
+ self->blocks = std::move(mapped_blocks).ValueUnchecked();
+ }
+ if (!mapped_ty.UnchangedOrSameAs(self->ty)) self->ty =
std::move(mapped_ty).ValueUnchecked();
+ if (!mapped_body.UnchangedOrSameAs(self->body)) {
+ self->body = std::move(mapped_body).ValueUnchecked();
+ }
+ return ffi::Unchanged().CopyToTVMFFIAny();
+}
+
+TVMFFIAny IfVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value)
noexcept {
+ const IfNode* self =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
IfNode>(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));
+
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->false_branch));
+ return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+}
+
+TVMFFIAny IfMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value)
noexcept {
+ const IfNode* self =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
IfNode>(value);
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Type>, mapped_ty,
+ mutator->MutateExpected(self->ty));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Expr>, mapped_cond,
+ mutator->MutateExpected(self->cond));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<SeqExpr>,
mapped_true_branch,
+
mutator->MutateExpected(self->true_branch));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<SeqExpr>,
mapped_false_branch,
+
mutator->MutateExpected(self->false_branch));
+ if (mapped_ty.UnchangedOrSameAs(self->ty) &&
mapped_cond.UnchangedOrSameAs(self->cond) &&
+ mapped_true_branch.UnchangedOrSameAs(self->true_branch) &&
+ mapped_false_branch.UnchangedOrSameAs(self->false_branch)) {
+ return ffi::Unchanged().CopyToTVMFFIAny();
+ }
+ ffi::ObjectPtr<IfNode> copy = ffi::make_object<IfNode>(*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));
+ copy->false_branch =
+
std::move(mapped_false_branch).ValueOrUnchanged(std::move(copy->false_branch));
+ return
ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(copy)));
+}
+
+TVMFFIAny IfMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
ffi::AnyView value) noexcept {
+ IfNode* self = const_cast<IfNode*>(
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
IfNode>(value));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Type>, mapped_ty,
+
mutator->MaybeInplaceMutateIfUniqueExpected(self->ty));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Expr>, mapped_cond,
+
mutator->MaybeInplaceMutateIfUniqueExpected(self->cond));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<SeqExpr>,
mapped_true_branch,
+
mutator->MaybeInplaceMutateIfUniqueExpected(self->true_branch));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(
+ ffi::UnchangedOr<SeqExpr>, mapped_false_branch,
+ mutator->MaybeInplaceMutateIfUniqueExpected(self->false_branch));
+ if (!mapped_ty.UnchangedOrSameAs(self->ty)) self->ty =
std::move(mapped_ty).ValueUnchecked();
+ if (!mapped_cond.UnchangedOrSameAs(self->cond)) {
+ self->cond = std::move(mapped_cond).ValueUnchecked();
+ }
+ if (!mapped_true_branch.UnchangedOrSameAs(self->true_branch)) {
+ self->true_branch = std::move(mapped_true_branch).ValueUnchecked();
+ }
+ if (!mapped_false_branch.UnchangedOrSameAs(self->false_branch)) {
+ self->false_branch = std::move(mapped_false_branch).ValueUnchecked();
+ }
+ return ffi::Unchanged().CopyToTVMFFIAny();
+}
+
+// Parameters precede the reflected ty field in this hook triple so their
pattern region establishes
+// the remap before the derived function type can refer to those symbolic
definitions.
+TVMFFIAny FunctionVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView
value) noexcept {
+ // skips: attrs (metadata), is_pure (scalar)
+ const FunctionNode* self =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
FunctionNode>(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));
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->body));
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->ret_ty));
+ return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+}
+
+TVMFFIAny FunctionMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView
value) noexcept {
+ // skips: attrs (metadata), is_pure (scalar)
+ const FunctionNode* self =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
FunctionNode>(value);
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<ffi::Array<Var>>,
mapped_params,
+
mutator->WithDefRegionKind(kTVMFFIDefRegionKindPattern, [&]() {
+ return
mutator->MutateExpected(self->params);
+ }));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Type>, mapped_ty,
+ mutator->MutateExpected(self->ty));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<SeqExpr>, mapped_body,
+ mutator->MutateExpected(self->body));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Type>, mapped_ret_ty,
+ mutator->MutateExpected(self->ret_ty));
+ if (mapped_params.UnchangedOrSameAs(self->params) &&
mapped_ty.UnchangedOrSameAs(self->ty) &&
+ mapped_body.UnchangedOrSameAs(self->body) &&
mapped_ret_ty.UnchangedOrSameAs(self->ret_ty)) {
+ return ffi::Unchanged().CopyToTVMFFIAny();
+ }
+ ffi::ObjectPtr<FunctionNode> copy = ffi::make_object<FunctionNode>(*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));
+ copy->ret_ty =
std::move(mapped_ret_ty).ValueOrUnchanged(std::move(copy->ret_ty));
+ return
ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(copy)));
+}
+
+TVMFFIAny FunctionMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
+ ffi::AnyView value) noexcept {
+ // skips: attrs (metadata), is_pure (scalar)
+ FunctionNode* self = const_cast<FunctionNode*>(
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
FunctionNode>(value));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(
+ ffi::UnchangedOr<ffi::Array<Var>>, mapped_params,
+ mutator->WithDefRegionKind(kTVMFFIDefRegionKindPattern, [&]() {
+ return mutator->MaybeInplaceMutateIfUniqueExpected(self->params);
+ }));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Type>, mapped_ty,
+
mutator->MaybeInplaceMutateIfUniqueExpected(self->ty));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<SeqExpr>, mapped_body,
+
mutator->MaybeInplaceMutateIfUniqueExpected(self->body));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Type>, mapped_ret_ty,
+
mutator->MaybeInplaceMutateIfUniqueExpected(self->ret_ty));
+ if (!mapped_params.UnchangedOrSameAs(self->params)) {
+ self->params = std::move(mapped_params).ValueUnchecked();
+ }
+ if (!mapped_ty.UnchangedOrSameAs(self->ty)) self->ty =
std::move(mapped_ty).ValueUnchecked();
+ if (!mapped_body.UnchangedOrSameAs(self->body)) {
+ self->body = std::move(mapped_body).ValueUnchecked();
+ }
+ if (!mapped_ret_ty.UnchangedOrSameAs(self->ret_ty)) {
+ self->ret_ty = std::move(mapped_ret_ty).ValueUnchecked();
+ }
+ return ffi::Unchanged().CopyToTVMFFIAny();
+}
+
+} // namespace
+
TVM_FFI_STATIC_INIT_BLOCK() {
+ namespace refl = tvm::ffi::reflection;
ShapeExprNode::RegisterReflection();
BindingNode::RegisterReflection();
DataflowVarNode::RegisterReflection();
@@ -42,6 +397,59 @@ TVM_FFI_STATIC_INIT_BLOCK() {
IfNode::RegisterReflection();
FunctionNode::RegisterReflection();
ExternFuncNode::RegisterReflection();
+ refl::TypeAttrDef<ShapeExprNode>()
+ .attr(refl::type_attr::kStructuralVisit,
reinterpret_cast<void*>(&ShapeExprVisit))
+ .attr(refl::type_attr::kStructuralMutate,
reinterpret_cast<void*>(&ShapeExprMutate))
+ .attr(refl::type_attr::kStructuralMaybeInplaceMutate,
+ reinterpret_cast<void*>(&ShapeExprMaybeInplaceMutate));
+ refl::TypeAttrDef<DataflowVarNode>()
+ .attr(refl::type_attr::kStructuralVisit,
reinterpret_cast<void*>(&DataflowVarVisit))
+ .attr(refl::type_attr::kStructuralMutate,
reinterpret_cast<void*>(&DataflowVarMutate))
+ .attr(refl::type_attr::kStructuralMaybeInplaceMutate,
+ reinterpret_cast<void*>(&DataflowVarMaybeInplaceMutate));
+ refl::TypeAttrDef<ConstantNode>()
+ .attr(refl::type_attr::kStructuralVisit,
+ reinterpret_cast<void*>(&TypeOnlyExprVisit<ConstantNode>))
+ .attr(refl::type_attr::kStructuralMutate,
+ reinterpret_cast<void*>(&TypeOnlyExprMutate<ConstantNode>))
+ .attr(refl::type_attr::kStructuralMaybeInplaceMutate,
+
reinterpret_cast<void*>(&TypeOnlyExprMaybeInplaceMutate<ConstantNode>));
+ refl::TypeAttrDef<StringImmNode>()
+ .attr(refl::type_attr::kStructuralVisit,
+ reinterpret_cast<void*>(&TypeOnlyExprVisit<StringImmNode>))
+ .attr(refl::type_attr::kStructuralMutate,
+ reinterpret_cast<void*>(&TypeOnlyExprMutate<StringImmNode>))
+ .attr(refl::type_attr::kStructuralMaybeInplaceMutate,
+
reinterpret_cast<void*>(&TypeOnlyExprMaybeInplaceMutate<StringImmNode>));
+ refl::TypeAttrDef<DataTypeImmNode>()
+ .attr(refl::type_attr::kStructuralVisit,
+ reinterpret_cast<void*>(&TypeOnlyExprVisit<DataTypeImmNode>))
+ .attr(refl::type_attr::kStructuralMutate,
+ reinterpret_cast<void*>(&TypeOnlyExprMutate<DataTypeImmNode>))
+ .attr(refl::type_attr::kStructuralMaybeInplaceMutate,
+
reinterpret_cast<void*>(&TypeOnlyExprMaybeInplaceMutate<DataTypeImmNode>));
+ refl::TypeAttrDef<SeqExprNode>()
+ .attr(refl::type_attr::kStructuralVisit,
reinterpret_cast<void*>(&SeqExprVisit))
+ .attr(refl::type_attr::kStructuralMutate,
reinterpret_cast<void*>(&SeqExprMutate))
+ .attr(refl::type_attr::kStructuralMaybeInplaceMutate,
+ reinterpret_cast<void*>(&SeqExprMaybeInplaceMutate));
+ refl::TypeAttrDef<IfNode>()
+ .attr(refl::type_attr::kStructuralVisit,
reinterpret_cast<void*>(&IfVisit))
+ .attr(refl::type_attr::kStructuralMutate,
reinterpret_cast<void*>(&IfMutate))
+ .attr(refl::type_attr::kStructuralMaybeInplaceMutate,
+ reinterpret_cast<void*>(&IfMaybeInplaceMutate));
+ refl::TypeAttrDef<FunctionNode>()
+ .attr(refl::type_attr::kStructuralVisit,
reinterpret_cast<void*>(&FunctionVisit))
+ .attr(refl::type_attr::kStructuralMutate,
reinterpret_cast<void*>(&FunctionMutate))
+ .attr(refl::type_attr::kStructuralMaybeInplaceMutate,
+ reinterpret_cast<void*>(&FunctionMaybeInplaceMutate));
+ refl::TypeAttrDef<ExternFuncNode>()
+ .attr(refl::type_attr::kStructuralVisit,
+ reinterpret_cast<void*>(&TypeOnlyExprVisit<ExternFuncNode>))
+ .attr(refl::type_attr::kStructuralMutate,
+ reinterpret_cast<void*>(&TypeOnlyExprMutate<ExternFuncNode>))
+ .attr(refl::type_attr::kStructuralMaybeInplaceMutate,
+
reinterpret_cast<void*>(&TypeOnlyExprMaybeInplaceMutate<ExternFuncNode>));
}
If::If(Expr cond, Expr true_branch, Expr false_branch, Span span) {
diff --git a/src/relax/ir/type.cc b/src/relax/ir/type.cc
index d6fa7ada9c..6664474ce9 100644
--- a/src/relax/ir/type.cc
+++ b/src/relax/ir/type.cc
@@ -21,6 +21,8 @@
* \file src/relax/ir/type.cc
* \brief Relax type system.
*/
+#include <tvm/ffi/extra/structural_mutate.h>
+#include <tvm/ffi/extra/structural_visit.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/relax/type.h>
@@ -28,7 +30,31 @@
namespace tvm {
namespace relax {
-TVM_FFI_STATIC_INIT_BLOCK() { PackedFuncTypeNode::RegisterReflection(); }
+namespace {
+
+TVMFFIAny PackedFuncTypeVisit(ffi::StructuralVisitorObj*, ffi::AnyView)
noexcept {
+ return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+}
+
+TVMFFIAny PackedFuncTypeMutate(ffi::StructuralMutatorObj*, ffi::AnyView)
noexcept {
+ return ffi::Unchanged().CopyToTVMFFIAny();
+}
+
+TVMFFIAny PackedFuncTypeMaybeInplaceMutate(ffi::StructuralMutatorObj*,
ffi::AnyView) noexcept {
+ return ffi::Unchanged().CopyToTVMFFIAny();
+}
+
+} // namespace
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+ namespace refl = tvm::ffi::reflection;
+ PackedFuncTypeNode::RegisterReflection();
+ refl::TypeAttrDef<PackedFuncTypeNode>()
+ .attr(refl::type_attr::kStructuralVisit,
reinterpret_cast<void*>(&PackedFuncTypeVisit))
+ .attr(refl::type_attr::kStructuralMutate,
reinterpret_cast<void*>(&PackedFuncTypeMutate))
+ .attr(refl::type_attr::kStructuralMaybeInplaceMutate,
+ reinterpret_cast<void*>(&PackedFuncTypeMaybeInplaceMutate));
+}
PackedFuncType::PackedFuncType(Span span) : Type(ffi::UnsafeInit{}) {
ffi::ObjectPtr<PackedFuncTypeNode> n =
ffi::make_object<PackedFuncTypeNode>();
diff --git a/src/tirx/ir/function.cc b/src/tirx/ir/function.cc
index 463a15084f..e0f295c931 100644
--- a/src/tirx/ir/function.cc
+++ b/src/tirx/ir/function.cc
@@ -21,6 +21,8 @@
* \file src/tirx/ir/function.cc
* \brief The function data structure.
*/
+#include <tvm/ffi/extra/structural_mutate.h>
+#include <tvm/ffi/extra/structural_visit.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/relax/expr.h>
@@ -32,9 +34,80 @@
namespace tvm {
namespace tirx {
+namespace {
+
+TVMFFIAny PrimFuncVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView
value) noexcept {
+ // skips: attrs (metadata), ty (derived by InferType)
+ const PrimFuncNode* self =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
PrimFuncNode>(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_type));
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->body));
+ return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+}
+
+TVMFFIAny PrimFuncMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView
value) noexcept {
+ // skips: attrs (metadata), ty (derived by InferType)
+ const PrimFuncNode* self =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
PrimFuncNode>(value);
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<ffi::Array<Var>>,
mapped_params,
+
mutator->WithDefRegionKind(kTVMFFIDefRegionKindPattern, [&]() {
+ return
mutator->MutateExpected(self->params);
+ }));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Type>, mapped_ret_type,
+ mutator->MutateExpected(self->ret_type));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Stmt>, mapped_body,
+ mutator->MutateExpected(self->body));
+ if (mapped_params.UnchangedOrSameAs(self->params) &&
+ mapped_ret_type.UnchangedOrSameAs(self->ret_type) &&
+ mapped_body.UnchangedOrSameAs(self->body)) {
+ return ffi::Unchanged().CopyToTVMFFIAny();
+ }
+ ffi::ObjectPtr<PrimFuncNode> copy = ffi::make_object<PrimFuncNode>(*self);
+ copy->params =
std::move(mapped_params).ValueOrUnchanged(std::move(copy->params));
+ copy->ret_type =
std::move(mapped_ret_type).ValueOrUnchanged(std::move(copy->ret_type));
+ copy->body = std::move(mapped_body).ValueOrUnchanged(std::move(copy->body));
+ return
ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(copy)));
+}
+
+TVMFFIAny PrimFuncMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator,
+ ffi::AnyView value) noexcept {
+ // skips: attrs (metadata), ty (derived by InferType)
+ PrimFuncNode* self = const_cast<PrimFuncNode*>(
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
PrimFuncNode>(value));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(
+ ffi::UnchangedOr<ffi::Array<Var>>, mapped_params,
+ mutator->WithDefRegionKind(kTVMFFIDefRegionKindPattern, [&]() {
+ return mutator->MaybeInplaceMutateIfUniqueExpected(self->params);
+ }));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Type>, mapped_ret_type,
+
mutator->MaybeInplaceMutateIfUniqueExpected(self->ret_type));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Stmt>, mapped_body,
+
mutator->MaybeInplaceMutateIfUniqueExpected(self->body));
+ if (!mapped_params.UnchangedOrSameAs(self->params)) {
+ self->params = std::move(mapped_params).ValueUnchecked();
+ }
+ if (!mapped_ret_type.UnchangedOrSameAs(self->ret_type)) {
+ self->ret_type = std::move(mapped_ret_type).ValueUnchecked();
+ }
+ if (!mapped_body.UnchangedOrSameAs(self->body)) {
+ self->body = std::move(mapped_body).ValueUnchecked();
+ }
+ return ffi::Unchanged().CopyToTVMFFIAny();
+}
+
+} // namespace
+
TVM_FFI_STATIC_INIT_BLOCK() {
+ namespace refl = tvm::ffi::reflection;
PrimFuncNode::RegisterReflection();
TensorIntrinNode::RegisterReflection();
+ refl::TypeAttrDef<PrimFuncNode>()
+ .attr(refl::type_attr::kStructuralVisit,
reinterpret_cast<void*>(&PrimFuncVisit))
+ .attr(refl::type_attr::kStructuralMutate,
reinterpret_cast<void*>(&PrimFuncMutate))
+ .attr(refl::type_attr::kStructuralMaybeInplaceMutate,
+ reinterpret_cast<void*>(&PrimFuncMaybeInplaceMutate));
}
namespace {
diff --git a/src/tirx/ir/stmt.cc b/src/tirx/ir/stmt.cc
index bd774356b2..644cdd1775 100644
--- a/src/tirx/ir/stmt.cc
+++ b/src/tirx/ir/stmt.cc
@@ -46,6 +46,18 @@ using SubscriptSlice = ffi::Array<ffi::Variant<
ffi::Tuple<ffi::Optional<PrimExpr>, ffi::Optional<PrimExpr>,
ffi::Optional<PrimExpr>>,
PrimExpr>>;
+TVMFFIAny BufferRegionTypeVisit(ffi::StructuralVisitorObj*, ffi::AnyView)
noexcept {
+ return ffi::AnyView(nullptr).CopyToTVMFFIAny();
+}
+
+TVMFFIAny BufferRegionTypeMutate(ffi::StructuralMutatorObj*, ffi::AnyView)
noexcept {
+ return ffi::Unchanged().CopyToTVMFFIAny();
+}
+
+TVMFFIAny BufferRegionTypeMaybeInplaceMutate(ffi::StructuralMutatorObj*,
ffi::AnyView) noexcept {
+ return ffi::Unchanged().CopyToTVMFFIAny();
+}
+
ffi::ObjectRef RealizeBufferRegionSubscript(Expr value, SubscriptSlice slice,
Span span) {
BufferRegion source = value.as_or_throw<BufferRegion>();
TVM_FFI_CHECK_LE(slice.size(), source->region.size(), IndexError)
@@ -1648,8 +1660,12 @@ BufferRegionType::BufferRegionType() :
Type(ffi::UnsafeInit{}) {
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
BufferRegionTypeNode::RegisterReflection();
- refl::TypeAttrDef<BufferRegionTypeNode>().def("__subscript_expr_realize__",
- RealizeBufferRegionSubscript);
+ refl::TypeAttrDef<BufferRegionTypeNode>()
+ .attr(refl::type_attr::kStructuralVisit,
reinterpret_cast<void*>(&BufferRegionTypeVisit))
+ .attr(refl::type_attr::kStructuralMutate,
reinterpret_cast<void*>(&BufferRegionTypeMutate))
+ .attr(refl::type_attr::kStructuralMaybeInplaceMutate,
+ reinterpret_cast<void*>(&BufferRegionTypeMaybeInplaceMutate))
+ .def("__subscript_expr_realize__", RealizeBufferRegionSubscript);
}
BufferRegion::BufferRegion(BufferVar buffer, ffi::Array<Range> region, Span
span) {
