diff --git a/src/ir/expr.cc b/src/ir/expr.cc index 1af350c70d23..b9b51da08963 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() + .attr(refl::type_attr::kStructuralVisit, reinterpret_cast(&GlobalVarVisit)) + .attr(refl::type_attr::kStructuralMutate, reinterpret_cast(&GlobalVarMutate)) + .attr(refl::type_attr::kStructuralMaybeInplaceMutate, + reinterpret_cast(&GlobalVarMaybeInplaceMutate)); +} // Call Call::Call(Type ret_ty, Expr op, ffi::Array args, Attrs attrs, ffi::Array ty_args, diff --git a/src/ir/op.cc b/src/ir/op.cc index 560cee6c9d89..c300806b74a1 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 +#include #include #include #include @@ -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() + .attr(refl::type_attr::kStructuralVisit, reinterpret_cast(&OpVisit)) + .attr(refl::type_attr::kStructuralMutate, reinterpret_cast(&OpMutate)) + .attr(refl::type_attr::kStructuralMaybeInplaceMutate, + reinterpret_cast(&OpMaybeInplaceMutate)); } using ffi::Any; diff --git a/src/ir/type.cc b/src/ir/type.cc index d4a56cb4765e..1947f489856c 100644 --- a/src/ir/type.cc +++ b/src/ir/type.cc @@ -69,6 +69,32 @@ ffi::ObjectPtr 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(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(value); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, mapped_element_type, + mutator->MutateExpected(self->element_type)); + if (mapped_element_type.UnchangedOrSameAs(self->element_type)) { + return ffi::Unchanged().CopyToTVMFFIAny(); + } + ffi::ObjectPtr copy = ffi::make_object(*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( + ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value)); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN( + ffi::UnchangedOr, 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(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(value); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr>, mapped_arg_types, + mutator->MutateExpected(self->arg_types)); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, 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 copy = ffi::make_object(*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( + ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value)); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr>, mapped_arg_types, + mutator->MaybeInplaceMutateIfUniqueExpected(self->arg_types)); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, 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(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(value); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr>, mapped_fields, + mutator->MutateExpected(self->fields)); + if (mapped_fields.UnchangedOrSameAs(self->fields)) { + return ffi::Unchanged().CopyToTVMFFIAny(); + } + ffi::ObjectPtr copy = ffi::make_object(*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( + ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value)); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr>, 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(); } -TVM_FFI_STATIC_INIT_BLOCK() { TypeNode::RegisterReflection(); } +TVM_FFI_STATIC_INIT_BLOCK() { + namespace refl = tvm::ffi::reflection; + TypeNode::RegisterReflection(); + refl::TypeAttrDef() + .attr(refl::type_attr::kStructuralVisit, reinterpret_cast(&TypeVisit)) + .attr(refl::type_attr::kStructuralMutate, reinterpret_cast(&TypeMutate)) + .attr(refl::type_attr::kStructuralMaybeInplaceMutate, + reinterpret_cast(&TypeMaybeInplaceMutate)); +} -TVM_FFI_STATIC_INIT_BLOCK() { OpaqueTypeNode::RegisterReflection(); } +TVM_FFI_STATIC_INIT_BLOCK() { + namespace refl = tvm::ffi::reflection; + OpaqueTypeNode::RegisterReflection(); + refl::TypeAttrDef() + .attr(refl::type_attr::kStructuralVisit, reinterpret_cast(&OpaqueTypeVisit)) + .attr(refl::type_attr::kStructuralMutate, reinterpret_cast(&OpaqueTypeMutate)) + .attr(refl::type_attr::kStructuralMaybeInplaceMutate, + reinterpret_cast(&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() + .attr(refl::type_attr::kStructuralVisit, reinterpret_cast(&PointerTypeVisit)) + .attr(refl::type_attr::kStructuralMutate, reinterpret_cast(&PointerTypeMutate)) + .attr(refl::type_attr::kStructuralMaybeInplaceMutate, + reinterpret_cast(&PointerTypeMaybeInplaceMutate)); } TVM_FFI_STATIC_INIT_BLOCK() { @@ -235,6 +405,11 @@ TVM_FFI_STATIC_INIT_BLOCK() { refl::GlobalDef().def("ir.FuncType", [](tvm::ffi::Array arg_types, Type ret_type) { return FuncType(arg_types, ret_type); }); + refl::TypeAttrDef() + .attr(refl::type_attr::kStructuralVisit, reinterpret_cast(&FuncTypeVisit)) + .attr(refl::type_attr::kStructuralMutate, reinterpret_cast(&FuncTypeMutate)) + .attr(refl::type_attr::kStructuralMaybeInplaceMutate, + reinterpret_cast(&FuncTypeMaybeInplaceMutate)); } TVM_FFI_STATIC_INIT_BLOCK() { @@ -245,6 +420,16 @@ TVM_FFI_STATIC_INIT_BLOCK() { .def("ir.TupleType", [](ffi::Array fields, Span span) { return TupleType(fields, span); }) .def("ir.TensorMapType", [](Span span) { return TensorMapType(span); }); + refl::TypeAttrDef() + .attr(refl::type_attr::kStructuralVisit, reinterpret_cast(&TupleTypeVisit)) + .attr(refl::type_attr::kStructuralMutate, reinterpret_cast(&TupleTypeMutate)) + .attr(refl::type_attr::kStructuralMaybeInplaceMutate, + reinterpret_cast(&TupleTypeMaybeInplaceMutate)); + refl::TypeAttrDef() + .attr(refl::type_attr::kStructuralVisit, reinterpret_cast(&TensorMapTypeVisit)) + .attr(refl::type_attr::kStructuralMutate, reinterpret_cast(&TensorMapTypeMutate)) + .attr(refl::type_attr::kStructuralMaybeInplaceMutate, + reinterpret_cast(&TensorMapTypeMaybeInplaceMutate)); } } // namespace tvm diff --git a/src/relax/distributed/type.cc b/src/relax/distributed/type.cc index 8fb164c2a86e..0a48faad2c72 100644 --- a/src/relax/distributed/type.cc +++ b/src/relax/distributed/type.cc @@ -22,16 +22,80 @@ * \brief Relax DTensor type. */ +#include +#include #include #include namespace tvm { namespace relax { namespace distributed { +namespace { + +TVMFFIAny DTensorTypeVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept { + const DTensorTypeNode* self = + ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value); + TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->device_mesh)); + TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->placement)); + TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->tensor_ty)); + return ffi::AnyView(nullptr).CopyToTVMFFIAny(); +} + +TVMFFIAny DTensorTypeMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept { + const DTensorTypeNode* self = + ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, mapped_device_mesh, + mutator->MutateExpected(self->device_mesh)); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, mapped_placement, + mutator->MutateExpected(self->placement)); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, 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 copy = ffi::make_object(*self); + copy->device_mesh = std::move(mapped_device_mesh).ValueOrUnchanged(std::move(copy->device_mesh)); + copy->placement = std::move(mapped_placement).ValueOrUnchanged(std::move(copy->placement)); + copy->tensor_ty = std::move(mapped_tensor_ty).ValueOrUnchanged(std::move(copy->tensor_ty)); + return ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(copy))); +} + +TVMFFIAny DTensorTypeMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator, + ffi::AnyView value) noexcept { + DTensorTypeNode* self = const_cast( + ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value)); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, mapped_device_mesh, + mutator->MaybeInplaceMutateIfUniqueExpected(self->device_mesh)); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, mapped_placement, + mutator->MaybeInplaceMutateIfUniqueExpected(self->placement)); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, 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() + .attr(refl::type_attr::kStructuralVisit, reinterpret_cast(&DTensorTypeVisit)) + .attr(refl::type_attr::kStructuralMutate, reinterpret_cast(&DTensorTypeMutate)) + .attr(refl::type_attr::kStructuralMaybeInplaceMutate, + reinterpret_cast(&DTensorTypeMaybeInplaceMutate)); } PlacementSpec PlacementSpec::Sharding(int axis) { diff --git a/src/relax/ir/dependent_type.cc b/src/relax/ir/dependent_type.cc index ff0bce1fb6d9..14755ca573e9 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 +#include #include #include #include @@ -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(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(value); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr>>, + mapped_values, mutator->MutateExpected(self->values)); + if (mapped_values.UnchangedOrSameAs(self->values)) return ffi::Unchanged().CopyToTVMFFIAny(); + ffi::ObjectPtr copy = ffi::make_object(*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( + ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value)); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr>>, + 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(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(value); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr>, mapped_shape, + mutator->MutateExpected(self->shape)); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr>, mapped_dtype, + mutator->MutateExpected(self->dtype)); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr>, 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 copy = ffi::make_object(*self); + copy->shape = std::move(mapped_shape).ValueOrUnchanged(std::move(copy->shape)); + copy->dtype = std::move(mapped_dtype).ValueOrUnchanged(std::move(copy->dtype)); + copy->vdevice = std::move(mapped_vdevice).ValueOrUnchanged(std::move(copy->vdevice)); + 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( + ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value)); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr>, mapped_shape, + mutator->MaybeInplaceMutateIfUniqueExpected(self->shape)); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr>, mapped_dtype, + mutator->MaybeInplaceMutateIfUniqueExpected(self->dtype)); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr>, 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(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(value); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN( + ffi::UnchangedOr>>, mapped_params, + mutator->WithDefRegionKind(kTVMFFIDefRegionKindPattern, + [&]() { return mutator->MutateExpected(self->params); })); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, mapped_ret, + mutator->MutateExpected(self->ret)); + if (mapped_params.UnchangedOrSameAs(self->params) && mapped_ret.UnchangedOrSameAs(self->ret)) { + return ffi::Unchanged().CopyToTVMFFIAny(); + } + ffi::ObjectPtr copy = ffi::make_object(*self); + copy->params = std::move(mapped_params).ValueOrUnchanged(std::move(copy->params)); + copy->ret = std::move(mapped_ret).ValueOrUnchanged(std::move(copy->ret)); + return ffi::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( + ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value)); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN( + ffi::UnchangedOr>>, mapped_params, + mutator->WithDefRegionKind(kTVMFFIDefRegionKindPattern, [&]() { + return mutator->MaybeInplaceMutateIfUniqueExpected(self->params); + })); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, 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().def( - "__subscript_expr_realize__", - [](Expr value, - ffi::Array, ffi::Optional, ffi::Optional>, - 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(); - TVM_FFI_CHECK(index.has_value(), TypeError) << "A Relax expression requires a point index"; - const auto* imm = index.value().as(); - TVM_FFI_CHECK(imm != nullptr, TypeError) - << "A Relax expression requires a constant integer index"; - return TupleGetItem(value, static_cast(imm->value), span); - }); + refl::TypeAttrDef() + .attr(refl::type_attr::kStructuralVisit, reinterpret_cast(&AnyTypeVisit)) + .attr(refl::type_attr::kStructuralMutate, reinterpret_cast(&AnyTypeMutate)) + .attr(refl::type_attr::kStructuralMaybeInplaceMutate, + reinterpret_cast(&AnyTypeMaybeInplaceMutate)); + refl::TypeAttrDef() + .attr(refl::type_attr::kStructuralVisit, reinterpret_cast(&ShapeTypeVisit)) + .attr(refl::type_attr::kStructuralMutate, reinterpret_cast(&ShapeTypeMutate)) + .attr(refl::type_attr::kStructuralMaybeInplaceMutate, + reinterpret_cast(&ShapeTypeMaybeInplaceMutate)); + refl::TypeAttrDef() + .attr(refl::type_attr::kStructuralVisit, reinterpret_cast(&TensorTypeVisit)) + .attr(refl::type_attr::kStructuralMutate, reinterpret_cast(&TensorTypeMutate)) + .attr(refl::type_attr::kStructuralMaybeInplaceMutate, + reinterpret_cast(&TensorTypeMaybeInplaceMutate)) + .def("__subscript_expr_realize__", + [](Expr value, + ffi::Array, ffi::Optional, + ffi::Optional>, + 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(); + TVM_FFI_CHECK(index.has_value(), TypeError) + << "A Relax expression requires a point index"; + const auto* imm = index.value().as(); + TVM_FFI_CHECK(imm != nullptr, TypeError) + << "A Relax expression requires a constant integer index"; + return TupleGetItem(value, static_cast(imm->value), span); + }); + refl::TypeAttrDef() + .attr(refl::type_attr::kStructuralVisit, reinterpret_cast(&FuncTypeVisit)) + .attr(refl::type_attr::kStructuralMutate, reinterpret_cast(&FuncTypeMutate)) + .attr(refl::type_attr::kStructuralMaybeInplaceMutate, + reinterpret_cast(&FuncTypeMaybeInplaceMutate)); } AnyType::AnyType(Span span) : Type(ffi::UnsafeInit{}) { diff --git a/src/relax/ir/expr.cc b/src/relax/ir/expr.cc index b21cdbad879c..d841d7817a6c 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 +#include #include #include #include @@ -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 +TVMFFIAny TypeOnlyExprVisit(ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept { + const TNode* self = + ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value); + TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->ty)); + return ffi::AnyView(nullptr).CopyToTVMFFIAny(); +} + +template +TVMFFIAny TypeOnlyExprMutate(ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept { + const TNode* self = + ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, mapped_ty, + mutator->MutateExpected(self->ty)); + if (mapped_ty.UnchangedOrSameAs(self->ty)) return ffi::Unchanged().CopyToTVMFFIAny(); + ffi::ObjectPtr copy = ffi::make_object(*self); + copy->ty = std::move(mapped_ty).ValueOrUnchanged(std::move(copy->ty)); + return ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(copy))); +} + +template +TVMFFIAny TypeOnlyExprMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator, + ffi::AnyView value) noexcept { + TNode* self = const_cast( + ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value)); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, 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(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(value); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, mapped_ty, + mutator->MutateExpected(self->ty)); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr>, mapped_values, + mutator->MutateExpected(self->values)); + if (mapped_ty.UnchangedOrSameAs(self->ty) && mapped_values.UnchangedOrSameAs(self->values)) { + return ffi::Unchanged().CopyToTVMFFIAny(); + } + ffi::ObjectPtr copy = ffi::make_object(*self); + copy->ty = std::move(mapped_ty).ValueOrUnchanged(std::move(copy->ty)); + copy->values = std::move(mapped_values).ValueOrUnchanged(std::move(copy->values)); + return ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(copy))); +} + +TVMFFIAny ShapeExprMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator, + ffi::AnyView value) noexcept { + ShapeExprNode* self = const_cast( + ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value)); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, mapped_ty, + mutator->MaybeInplaceMutateIfUniqueExpected(self->ty)); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr>, 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(value); + if (!self->ty.as()) { + if (visitor->def_region_kind() == kTVMFFIDefRegionKindSimple) { + TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->WithDefRegionKind( + kTVMFFIDefRegionKindNone, [&]() { return visitor->VisitExpected(self->ty); })); + } 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(value); + ffi::Expected remap_result = mutator->VarRemapGetExpected(value); + TVM_FFI_S_MUTATE_MAYBE_EARLY_RETURN(remap_result); + if (ffi::details::ExpectedUnsafe::GetData(remap_result).type_index() != + ffi::TypeIndex::kTVMFFINone) { + return ffi::details::ExpectedUnsafe::MoveToTVMFFIAny(std::move(remap_result)); + } + if (mutator->def_region_kind() == kTVMFFIDefRegionKindNone) { + return ffi::Unchanged().CopyToTVMFFIAny(); + } + ffi::UnchangedOr result = ffi::Unchanged(); + ffi::Any mapped_value; + if (!self->ty.as()) { + ffi::Expected> 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, mapped_ty, + std::move(mapped_ty_result)); + if (!mapped_ty.UnchangedOrSameAs(self->ty)) { + ffi::ObjectPtr copy = ffi::make_object(*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( + ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value)); + ffi::Expected remap_result = mutator->VarRemapGetExpected(value); + TVM_FFI_S_MUTATE_MAYBE_EARLY_RETURN(remap_result); + if (ffi::details::ExpectedUnsafe::GetData(remap_result).type_index() != + ffi::TypeIndex::kTVMFFINone) { + return ffi::details::ExpectedUnsafe::MoveToTVMFFIAny(std::move(remap_result)); + } + if (mutator->def_region_kind() == kTVMFFIDefRegionKindNone) { + return ffi::Unchanged().CopyToTVMFFIAny(); + } + ffi::UnchangedOr result = ffi::Unchanged(); + ffi::Any mapped_value; + if (!self->ty.as()) { + ffi::Expected> 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, 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(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(value); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr>, mapped_blocks, + mutator->MutateExpected(self->blocks)); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, mapped_ty, + mutator->MutateExpected(self->ty)); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, 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 copy = ffi::make_object(*self); + copy->blocks = std::move(mapped_blocks).ValueOrUnchanged(std::move(copy->blocks)); + copy->ty = std::move(mapped_ty).ValueOrUnchanged(std::move(copy->ty)); + copy->body = std::move(mapped_body).ValueOrUnchanged(std::move(copy->body)); + return ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(ffi::Any(std::move(copy))); +} + +TVMFFIAny SeqExprMaybeInplaceMutate(ffi::StructuralMutatorObj* mutator, + ffi::AnyView value) noexcept { + SeqExprNode* self = const_cast( + ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value)); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr>, mapped_blocks, + mutator->MaybeInplaceMutateIfUniqueExpected(self->blocks)); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, mapped_ty, + mutator->MaybeInplaceMutateIfUniqueExpected(self->ty)); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, 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(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(value); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, mapped_ty, + mutator->MutateExpected(self->ty)); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, mapped_cond, + mutator->MutateExpected(self->cond)); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, mapped_true_branch, + mutator->MutateExpected(self->true_branch)); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, 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 copy = ffi::make_object(*self); + copy->ty = std::move(mapped_ty).ValueOrUnchanged(std::move(copy->ty)); + copy->cond = std::move(mapped_cond).ValueOrUnchanged(std::move(copy->cond)); + copy->true_branch = std::move(mapped_true_branch).ValueOrUnchanged(std::move(copy->true_branch)); + 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( + ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value)); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, mapped_ty, + mutator->MaybeInplaceMutateIfUniqueExpected(self->ty)); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, mapped_cond, + mutator->MaybeInplaceMutateIfUniqueExpected(self->cond)); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, mapped_true_branch, + mutator->MaybeInplaceMutateIfUniqueExpected(self->true_branch)); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN( + ffi::UnchangedOr, 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(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(value); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr>, mapped_params, + mutator->WithDefRegionKind(kTVMFFIDefRegionKindPattern, [&]() { + return mutator->MutateExpected(self->params); + })); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, mapped_ty, + mutator->MutateExpected(self->ty)); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, mapped_body, + mutator->MutateExpected(self->body)); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, 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 copy = ffi::make_object(*self); + copy->params = std::move(mapped_params).ValueOrUnchanged(std::move(copy->params)); + copy->ty = std::move(mapped_ty).ValueOrUnchanged(std::move(copy->ty)); + copy->body = std::move(mapped_body).ValueOrUnchanged(std::move(copy->body)); + 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( + ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value)); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN( + ffi::UnchangedOr>, mapped_params, + mutator->WithDefRegionKind(kTVMFFIDefRegionKindPattern, [&]() { + return mutator->MaybeInplaceMutateIfUniqueExpected(self->params); + })); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, mapped_ty, + mutator->MaybeInplaceMutateIfUniqueExpected(self->ty)); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, mapped_body, + mutator->MaybeInplaceMutateIfUniqueExpected(self->body)); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, 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() + .attr(refl::type_attr::kStructuralVisit, reinterpret_cast(&ShapeExprVisit)) + .attr(refl::type_attr::kStructuralMutate, reinterpret_cast(&ShapeExprMutate)) + .attr(refl::type_attr::kStructuralMaybeInplaceMutate, + reinterpret_cast(&ShapeExprMaybeInplaceMutate)); + refl::TypeAttrDef() + .attr(refl::type_attr::kStructuralVisit, reinterpret_cast(&DataflowVarVisit)) + .attr(refl::type_attr::kStructuralMutate, reinterpret_cast(&DataflowVarMutate)) + .attr(refl::type_attr::kStructuralMaybeInplaceMutate, + reinterpret_cast(&DataflowVarMaybeInplaceMutate)); + refl::TypeAttrDef() + .attr(refl::type_attr::kStructuralVisit, + reinterpret_cast(&TypeOnlyExprVisit)) + .attr(refl::type_attr::kStructuralMutate, + reinterpret_cast(&TypeOnlyExprMutate)) + .attr(refl::type_attr::kStructuralMaybeInplaceMutate, + reinterpret_cast(&TypeOnlyExprMaybeInplaceMutate)); + refl::TypeAttrDef() + .attr(refl::type_attr::kStructuralVisit, + reinterpret_cast(&TypeOnlyExprVisit)) + .attr(refl::type_attr::kStructuralMutate, + reinterpret_cast(&TypeOnlyExprMutate)) + .attr(refl::type_attr::kStructuralMaybeInplaceMutate, + reinterpret_cast(&TypeOnlyExprMaybeInplaceMutate)); + refl::TypeAttrDef() + .attr(refl::type_attr::kStructuralVisit, + reinterpret_cast(&TypeOnlyExprVisit)) + .attr(refl::type_attr::kStructuralMutate, + reinterpret_cast(&TypeOnlyExprMutate)) + .attr(refl::type_attr::kStructuralMaybeInplaceMutate, + reinterpret_cast(&TypeOnlyExprMaybeInplaceMutate)); + refl::TypeAttrDef() + .attr(refl::type_attr::kStructuralVisit, reinterpret_cast(&SeqExprVisit)) + .attr(refl::type_attr::kStructuralMutate, reinterpret_cast(&SeqExprMutate)) + .attr(refl::type_attr::kStructuralMaybeInplaceMutate, + reinterpret_cast(&SeqExprMaybeInplaceMutate)); + refl::TypeAttrDef() + .attr(refl::type_attr::kStructuralVisit, reinterpret_cast(&IfVisit)) + .attr(refl::type_attr::kStructuralMutate, reinterpret_cast(&IfMutate)) + .attr(refl::type_attr::kStructuralMaybeInplaceMutate, + reinterpret_cast(&IfMaybeInplaceMutate)); + refl::TypeAttrDef() + .attr(refl::type_attr::kStructuralVisit, reinterpret_cast(&FunctionVisit)) + .attr(refl::type_attr::kStructuralMutate, reinterpret_cast(&FunctionMutate)) + .attr(refl::type_attr::kStructuralMaybeInplaceMutate, + reinterpret_cast(&FunctionMaybeInplaceMutate)); + refl::TypeAttrDef() + .attr(refl::type_attr::kStructuralVisit, + reinterpret_cast(&TypeOnlyExprVisit)) + .attr(refl::type_attr::kStructuralMutate, + reinterpret_cast(&TypeOnlyExprMutate)) + .attr(refl::type_attr::kStructuralMaybeInplaceMutate, + reinterpret_cast(&TypeOnlyExprMaybeInplaceMutate)); } 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 d6fa7ada9cc4..6664474ce978 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 +#include #include #include #include @@ -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() + .attr(refl::type_attr::kStructuralVisit, reinterpret_cast(&PackedFuncTypeVisit)) + .attr(refl::type_attr::kStructuralMutate, reinterpret_cast(&PackedFuncTypeMutate)) + .attr(refl::type_attr::kStructuralMaybeInplaceMutate, + reinterpret_cast(&PackedFuncTypeMaybeInplaceMutate)); +} PackedFuncType::PackedFuncType(Span span) : Type(ffi::UnsafeInit{}) { ffi::ObjectPtr n = ffi::make_object(); diff --git a/src/tirx/ir/function.cc b/src/tirx/ir/function.cc index 463a15084f7d..e0f295c93118 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 +#include #include #include #include @@ -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(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(value); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr>, mapped_params, + mutator->WithDefRegionKind(kTVMFFIDefRegionKindPattern, [&]() { + return mutator->MutateExpected(self->params); + })); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, mapped_ret_type, + mutator->MutateExpected(self->ret_type)); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, 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 copy = ffi::make_object(*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( + ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck(value)); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN( + ffi::UnchangedOr>, mapped_params, + mutator->WithDefRegionKind(kTVMFFIDefRegionKindPattern, [&]() { + return mutator->MaybeInplaceMutateIfUniqueExpected(self->params); + })); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, mapped_ret_type, + mutator->MaybeInplaceMutateIfUniqueExpected(self->ret_type)); + TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr, 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() + .attr(refl::type_attr::kStructuralVisit, reinterpret_cast(&PrimFuncVisit)) + .attr(refl::type_attr::kStructuralMutate, reinterpret_cast(&PrimFuncMutate)) + .attr(refl::type_attr::kStructuralMaybeInplaceMutate, + reinterpret_cast(&PrimFuncMaybeInplaceMutate)); } namespace { diff --git a/src/tirx/ir/stmt.cc b/src/tirx/ir/stmt.cc index bd774356b2e5..644cdd1775fa 100644 --- a/src/tirx/ir/stmt.cc +++ b/src/tirx/ir/stmt.cc @@ -46,6 +46,18 @@ using SubscriptSlice = ffi::Array, ffi::Optional, ffi::Optional>, 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(); 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().def("__subscript_expr_realize__", - RealizeBufferRegionSubscript); + refl::TypeAttrDef() + .attr(refl::type_attr::kStructuralVisit, reinterpret_cast(&BufferRegionTypeVisit)) + .attr(refl::type_attr::kStructuralMutate, reinterpret_cast(&BufferRegionTypeMutate)) + .attr(refl::type_attr::kStructuralMaybeInplaceMutate, + reinterpret_cast(&BufferRegionTypeMaybeInplaceMutate)) + .def("__subscript_expr_realize__", RealizeBufferRegionSubscript); } BufferRegion::BufferRegion(BufferVar buffer, ffi::Array region, Span span) {