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-ffi.git
The following commit(s) were added to refs/heads/main by this push:
new 42d4be73 [REFACTOR][EXTRA] Optimize StructuralMap engine (#741)
42d4be73 is described below
commit 42d4be739fc452651b808f3cc6f2c87c17847b27
Author: Tianqi Chen <[email protected]>
AuthorDate: Fri Sep 4 20:48:27 2026 -0400
[REFACTOR][EXTRA] Optimize StructuralMap engine (#741)
This PR improves the `StructuralMap` engine by bringing lambda callback
style to direct calling style. We also internally used the raw
`TVMFFIAny` ABI path so engine overhead is minimized in the default
path.
Main changes:
- **Direct calling style.** The callback chain was continuation-passing,
a lambda nested per link. The mutator now owns its callbacks and tests
them with a `(... || ...)` fold; selection, identity remap, descent, and
invocation are one straight-line function.
- **Raw `TVMFFIAny` on the ABI boundary.** Hooks are C-ABI function
pointers returning `TVMFFIAny`, a 16-byte POD that passes in registers,
while `Expected<Any>` is classified MEMORY and forces the result to the
stack. Descent through an unmatched node is that boundary, so it now
stays raw end to end. Hook bodies are unchanged:
`TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN` still takes a typed `Expected<Any>`
and moves out only on return.
- **Split static and dynamic selection.** `StructuralMapMutatorObj`
(header) and `StructuralMapDynMutatorObj` (`.cc`) now share only the
identity remap. This also removes a hazard: the dynamic form kept its
selected `Function` in mutable state, where a post-order walk could
descend into a matching child and overwrite the parent's selection.
- **Error context is attached by the engine, not by hooks.** The engine
names a node where it dispatches into that node, giving one frame per
node instead of two on some paths. Both `MAYBE_EARLY_RETURN` macros lose
their node argument, and the visit engine no longer routes its own walk
through the hook-facing macro.
- **API.** Removes `TVM_FFI_S_MUTATE_ASSIGN_FROM`,
`SMutateResultUnchanged`, `AssignOrReturnHelper`, and
`MutateWithIdentityRemapExpected`. Adds
`TVM_FFI_UNSAFE_S_MUTATE_ASSIGN_OR_RETURN_SKIP_CHECK` for hooks that
have measured the type check to matter; it is UB on a wrong-typed result
and has no call sites here. Also fixes `AnyUnsafe::CheckAnyStrict`,
which was missing the `Any` special case both of its siblings had.
- **Behavior change.** When descent hands back a node whose type no
longer matches the link that selected it, the engine reports a
`TypeError` instead of passing the node through. Default mutation must
preserve a node's type, so this only fires when a hook has broken that
invariant. Match-before-descent ordering, identity-remap entry and exit,
and the pre-order in-place rules are unchanged.
---
include/tvm/ffi/any.h | 7 +-
include/tvm/ffi/expected.h | 23 -
include/tvm/ffi/extra/structural_mutate.h | 782 ++++++++++++++++------------
include/tvm/ffi/extra/structural_visit.h | 66 ++-
include/tvm/ffi/extra/visit_error_context.h | 38 ++
src/ffi/extra/structural_mutate.cc | 297 ++++++++---
src/ffi/extra/structural_visit.cc | 4 +-
tests/cpp/extra/test_structural_mutate.cc | 25 +
tests/cpp/extra/test_structural_visit.cc | 25 +
tests/cpp/testing_object.h | 24 +-
10 files changed, 838 insertions(+), 453 deletions(-)
diff --git a/include/tvm/ffi/any.h b/include/tvm/ffi/any.h
index 3d203df3..219ab47a 100644
--- a/include/tvm/ffi/any.h
+++ b/include/tvm/ffi/any.h
@@ -609,7 +609,12 @@ struct AnyUnsafe : public ObjectUnsafe {
template <typename T>
TVM_FFI_INLINE static bool CheckAnyStrict(const Any& ref) {
- return TypeTraits<T>::CheckAnyStrict(&(ref.data_));
+ if constexpr (!std::is_same_v<T, Any>) {
+ return TypeTraits<T>::CheckAnyStrict(&(ref.data_));
+ } else {
+ // Any holds any value, so there is nothing to check against.
+ return true;
+ }
}
template <typename T>
diff --git a/include/tvm/ffi/expected.h b/include/tvm/ffi/expected.h
index bd7417fa..6ebd9ab5 100644
--- a/include/tvm/ffi/expected.h
+++ b/include/tvm/ffi/expected.h
@@ -348,29 +348,6 @@ class Expected<void> {
namespace details {
-// Helper for TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN.
-class AssignOrReturnHelper {
- public:
- TVM_FFI_INLINE explicit AssignOrReturnHelper(Any&& data) :
data_(std::move(data)) {}
-
- template <typename T>
- TVM_FFI_INLINE Expected<T> TryMove() && {
- if constexpr (!std::is_same_v<T, Any>) {
- const TVMFFIAny* data = AnyUnsafe::TVMFFIAnyPtrFromAny(data_);
- if (TVM_FFI_PREDICT_FALSE(!TypeTraits<T>::CheckAnyStrict(data))) {
- return Unexpected(Error("TypeError",
- "Cannot treat type `" +
TypeTraits<T>::GetMismatchTypeInfo(data) +
- "` as type `" + TypeTraits<T>::TypeStr() +
"`",
- ""));
- }
- }
- return AnyUnsafe::MoveFromAnyAfterCheck<T>(std::move(data_));
- }
-
- private:
- Any data_;
-};
-
/*!
* \brief Unsafe raw-storage helpers for Expected.
*
diff --git a/include/tvm/ffi/extra/structural_mutate.h
b/include/tvm/ffi/extra/structural_mutate.h
index d2d04aa1..f2164c09 100644
--- a/include/tvm/ffi/extra/structural_mutate.h
+++ b/include/tvm/ffi/extra/structural_mutate.h
@@ -212,27 +212,9 @@ class StructuralMutatorObj : public Object {
* responsible for variable-remap lookup and insertion when it
represents a FreeVar or DAG
* identity. Automatic remapping applies only to the reflected
fallback.
*/
+
TVM_FFI_INLINE Expected<Any> DefaultMutateExpected(AnyView value) noexcept {
- int32_t type_index = value.type_index();
- static reflection::TypeAttrColumn
column(reflection::type_attr::kStructuralMutate);
- AnyView attr = column[type_index];
- if (attr.type_index() != TypeIndex::kTVMFFINone) {
- if (attr.type_index() == TypeIndex::kTVMFFIOpaquePtr) {
- auto* hook = reinterpret_cast<FStructuralMutate>(attr.cast<void*>());
- return details::ExpectedUnsafe::MoveFromTVMFFIAny<Any>((*hook)(this,
value));
- }
- if (attr.type_index() == TypeIndex::kTVMFFIFunction) {
- return attr.cast<Function>().CallExpected<Any>(this, value);
- }
- return Unexpected(Error(
- "TypeError", "__s_mutate__ must be an opaque function pointer or
ffi.Function", ""));
- }
- if (type_index < TypeIndex::kTVMFFIStaticObjectBegin) {
- return Any(value);
- }
- return MutateWithIdentityRemapExpected(value, [&]() -> Expected<Any> {
- return details::MutateReflectedFieldsExpected(this, value);
- });
+ return
details::ExpectedUnsafe::MoveFromTVMFFIAny<Any>(DefaultMutateRaw(value));
}
/*!
@@ -248,18 +230,7 @@ class StructuralMutatorObj : public Object {
* \ref DefaultMutateExpected.
*/
TVM_FFI_INLINE Expected<Any> DefaultMaybeInplaceMutateExpected(AnyView
value) noexcept {
- int32_t type_index = value.type_index();
- static reflection::TypeAttrColumn maybe_inplace_mutate_column(
- reflection::type_attr::kStructuralMaybeInplaceMutate);
- AnyView maybe_inplace_mutate_attr =
maybe_inplace_mutate_column[type_index];
- if (maybe_inplace_mutate_attr.type_index() == TypeIndex::kTVMFFIOpaquePtr)
{
- auto* hook =
reinterpret_cast<FStructuralMutate>(maybe_inplace_mutate_attr.cast<void*>());
- return details::ExpectedUnsafe::MoveFromTVMFFIAny<Any>((*hook)(this,
value));
- }
- if (maybe_inplace_mutate_attr.type_index() == TypeIndex::kTVMFFIFunction) {
- return
maybe_inplace_mutate_attr.cast<Function>().CallExpected<Any>(this, value);
- }
- return DefaultMutateExpected(value);
+ return
details::ExpectedUnsafe::MoveFromTVMFFIAny<Any>(DefaultMaybeInplaceMutateRaw(value));
}
/*!
@@ -332,48 +303,139 @@ class StructuralMutatorObj : public Object {
protected:
/*!
- * \brief Mutate a FreeVar or DAG identity once and reuse its final result.
+ * \brief Diagnostic for a malformed ``__s_mutate__`` registration.
*
- * \tparam Mutation Nullary callable returning ``Expected<Any>``.
- * \param value The borrowed value to mutate.
- * \param mutation The complete mutation to apply on a cache miss.
- * \return The cached or newly computed owning result, or an Error.
+ * Kept out of line and cold: it can only fire for a type whose registered
attribute is
+ * neither an opaque function pointer nor an ffi.Function, so it is
unreachable for any
+ * correctly registered type. Inlined, its three string literals and Error
construction
+ * land in the traversal's hot path for no reason.
*/
- template <typename Mutation>
- TVM_FFI_INLINE Expected<Any> MutateWithIdentityRemapExpected(AnyView value,
- Mutation&&
mutation) noexcept {
- int32_t type_index = value.type_index();
- if (type_index < TypeIndex::kTVMFFIStaticObjectBegin) {
- return mutation();
+ TVM_FFI_COLD_CODE static Expected<Any> BadStructuralMutateHookError()
noexcept {
+ return Unexpected(
+ Error("TypeError", "__s_mutate__ must be an opaque function pointer or
ffi.Function", ""));
+ }
+
+ // Convention: the ABI boundary is a raw TVMFFIAny; everything inside a
callback or hook body
+ // works in Expected<Any> and moves out to TVMFFIAny at that boundary.
+ //
+ // The Raw forms below exist because that boundary is also the default path.
A hook is a C-ABI
+ // function pointer returning TVMFFIAny, a 16-byte POD that stays in
registers; wrapping the
+ // result in Expected<Any> would force it to memory, since Expected<Any> is
not trivially
+ // destructible and is therefore classified MEMORY. Descent through an
unmatched node calls a
+ // hook and returns its result unchanged, so keeping that path raw removes
the round trip
+ // entirely. Only a matched callback pays for an Expected<Any>, and it
materializes an Any
+ // anyway to record its remap entry.
+ //
+ // Engine-internal: subclasses call the Expected forms above.
+ /*! \brief Raw default mutation: attr lookup then hook, favouring the fn-ptr
case. */
+ TVM_FFI_INLINE TVMFFIAny DefaultMutateRaw(AnyView value) noexcept {
+ static reflection::TypeAttrColumn
column(reflection::type_attr::kStructuralMutate);
+ AnyView attr = column[value.type_index()];
+ // Exactly one frame per node: hooks propagate errors untouched, and this
is the engine
+ // dispatching into `value`, so both exits below name it here and nowhere
else.
+ TVMFFIAny result;
+ if (TVM_FFI_PREDICT_TRUE(attr.type_index() ==
TypeIndex::kTVMFFIOpaquePtr)) {
+ result =
(*reinterpret_cast<FStructuralMutate>(attr.cast<void*>()))(this, value);
+ } else {
+ result = DefaultMutateRawTail(value, attr);
}
+ if (TVM_FFI_PREDICT_FALSE(result.type_index == TypeIndex::kTVMFFIError)) {
+ details::UpdateVisitErrorContext(result, value);
+ }
+ return result;
+ }
+ // The cold remainder of DefaultMutateRaw, out of line so that
always-inlined caller stays
+ // small at every inlining site. `attr` is the attribute the caller already
read, so the
+ // column is never looked up twice.
+ /*!
+ * \brief Whether a node's identity is remappable, so it maps once and
reuses that result.
+ * \param type_index The node's runtime type index.
+ * \return True for a FreeVar or DAG node.
+ */
+ TVM_FFI_INLINE static bool IsRemappableIdentity(int32_t type_index) noexcept
{
+ if (type_index < TypeIndex::kTVMFFIStaticObjectBegin) return false;
const TVMFFITypeInfo* type_info = TVMFFIGetTypeInfo(type_index);
- bool is_remappable_identity =
- type_info->metadata != nullptr &&
- (type_info->metadata->structural_eq_hash_kind ==
kTVMFFISEqHashKindFreeVar ||
- type_info->metadata->structural_eq_hash_kind ==
kTVMFFISEqHashKindDAGNode);
- if (!is_remappable_identity) {
- return mutation();
- }
+ return type_info->metadata != nullptr &&
+ (type_info->metadata->structural_eq_hash_kind ==
kTVMFFISEqHashKindFreeVar ||
+ type_info->metadata->structural_eq_hash_kind ==
kTVMFFISEqHashKindDAGNode);
+ }
- Expected<Any> mapped_value = VarRemapGetExpected(value);
- if (TVM_FFI_PREDICT_FALSE(mapped_value.is_err())) {
- return Unexpected(std::move(mapped_value).error());
+ /*! \brief The cold remainder of DefaultMutateRaw: an ffi.Function hook, or
no hook at all. */
+ TVMFFIAny DefaultMutateRawTail(AnyView value, AnyView attr) noexcept {
+ if (attr.type_index() != TypeIndex::kTVMFFINone) {
+ // Registered, but as an ffi.Function rather than an opaque pointer.
+ if (attr.type_index() == TypeIndex::kTVMFFIFunction) {
+ return details::ExpectedUnsafe::MoveToTVMFFIAny(
+ attr.cast<Function>().CallExpected<Any>(this, value));
+ }
+ // Registered as neither: a malformed hook.
+ return
details::ExpectedUnsafe::MoveToTVMFFIAny(BadStructuralMutateHookError());
}
- if (details::ExpectedUnsafe::GetData(mapped_value).type_index() !=
TypeIndex::kTVMFFINone) {
- return mapped_value;
+ // No hook at all. A POD carries through unchanged; an object walks its
reflected fields.
+ if (value.type_index() < TypeIndex::kTVMFFIStaticObjectBegin) {
+ return details::AnyUnsafe::MoveAnyToTVMFFIAny(Any(value));
}
-
- Expected<Any> result = mutation();
+ // A FreeVar or DAG node maps once and every later occurrence reuses that
result, so the
+ // reflected walk runs under an identity remap.
+ const bool remappable = IsRemappableIdentity(value.type_index());
+ if (remappable) {
+ Expected<Any> mapped = VarRemapGetExpected(value);
+ if (TVM_FFI_PREDICT_FALSE(mapped.is_err()) ||
+ details::ExpectedUnsafe::GetData(mapped).type_index() !=
TypeIndex::kTVMFFINone) {
+ return details::ExpectedUnsafe::MoveToTVMFFIAny(std::move(mapped));
+ }
+ }
+ Expected<Any> result = details::MutateReflectedFieldsExpected(this, value);
if (TVM_FFI_PREDICT_FALSE(result.is_err())) {
+ return details::ExpectedUnsafe::MoveToTVMFFIAny(std::move(result));
+ }
+ if (remappable) {
+ Expected<void> set_result =
+ VarRemapSetExpected(value, details::ExpectedUnsafe::GetData(result));
+ if (TVM_FFI_PREDICT_FALSE(set_result.is_err())) {
+ return details::ExpectedUnsafe::MoveToTVMFFIAny(
+ Expected<Any>(Unexpected(std::move(set_result).error())));
+ }
+ }
+ return details::ExpectedUnsafe::MoveToTVMFFIAny(std::move(result));
+ }
+ /*!
+ * \brief Raw default maybe-in-place mutation.
+ *
+ * \note A registered opaque hook is the expected case here too, so the
attribute is read once
+ * and every other shape is handed to the out-of-line remainder.
+ */
+ TVM_FFI_INLINE TVMFFIAny DefaultMaybeInplaceMutateRaw(AnyView value)
noexcept {
+ static reflection::TypeAttrColumn
column(reflection::type_attr::kStructuralMaybeInplaceMutate);
+ AnyView attr = column[value.type_index()];
+ if (TVM_FFI_PREDICT_TRUE(attr.type_index() ==
TypeIndex::kTVMFFIOpaquePtr)) {
+ // This is the engine dispatching into `value`; hooks propagate errors
untouched, so the
+ // node is named here. The fall-through re-dispatches the same node
through
+ // DefaultMutateRaw, which names it there instead -- exactly one frame
either way.
+ TVMFFIAny result =
(*reinterpret_cast<FStructuralMutate>(attr.cast<void*>()))(this, value);
+ if (TVM_FFI_PREDICT_FALSE(result.type_index == TypeIndex::kTVMFFIError))
{
+ details::UpdateVisitErrorContext(result, value);
+ }
return result;
}
- Expected<void> set_result =
- VarRemapSetExpected(value, details::ExpectedUnsafe::GetData(result));
- if (TVM_FFI_PREDICT_FALSE(set_result.is_err())) {
- return Unexpected(std::move(set_result).error());
+ return DefaultMaybeInplaceMutateRawTail(value, attr);
+ }
+
+ /*!
+ * \brief The cold remainder of DefaultMaybeInplaceMutateRaw: an
ffi.Function in-place hook, or
+ * no in-place hook at all, in which case the ordinary mutate path
runs.
+ */
+ TVMFFIAny DefaultMaybeInplaceMutateRawTail(AnyView value, AnyView attr)
noexcept {
+ if (attr.type_index() == TypeIndex::kTVMFFIFunction) {
+ TVMFFIAny result = details::ExpectedUnsafe::MoveToTVMFFIAny(
+ attr.cast<Function>().CallExpected<Any>(this, value));
+ if (TVM_FFI_PREDICT_FALSE(result.type_index == TypeIndex::kTVMFFIError))
{
+ details::UpdateVisitErrorContext(result, value);
+ }
+ return result;
}
- return result;
+ return DefaultMutateRaw(value);
}
/*!
@@ -540,31 +602,28 @@ namespace details {
// Append Node to the mutate error context before returning. Node is required:
dropping it
// silently degrades every error message produced below this frame.
// A raw pointer Node must be non-null; pass nullable nodes as ObjectRef or
Any so None is skipped.
-#define TVM_FFI_S_MUTATE_MAYBE_EARLY_RETURN(Result, Node)
\
- do {
\
- auto&& tvm_ffi_res_ = (Result);
\
- if (TVM_FFI_PREDICT_FALSE(tvm_ffi_res_.type_index() ==
::tvm::ffi::TypeIndex::kTVMFFIError)) { \
- auto&& tvm_ffi_mutate_node_owner_ = (Node);
\
- ::tvm::ffi::AnyView tvm_ffi_mutate_node_ = tvm_ffi_mutate_node_owner_;
\
- if (tvm_ffi_mutate_node_.type_index() >=
::tvm::ffi::TypeIndex::kTVMFFIStaticObjectBegin) { \
- ::tvm::ffi::Error tvm_ffi_mutate_err_ = tvm_ffi_res_.error();
\
- ::tvm::ffi::details::UpdateVisitErrorContext(
\
- tvm_ffi_mutate_err_,
tvm_ffi_mutate_node_.cast<::tvm::ffi::ObjectRef>()); \
- }
\
- return ::std::move(tvm_ffi_res_);
\
- }
\
+#define TVM_FFI_S_MUTATE_MAYBE_EARLY_RETURN(Result)
\
+ do {
\
+ auto&& tvm_ffi_res_ = (Result);
\
+ if (TVM_FFI_PREDICT_FALSE(tvm_ffi_res_.is_err())) {
\
+ return
::tvm::ffi::details::ExpectedUnsafe::MoveToTVMFFIAny(::std::move(tvm_ffi_res_));
\
+ }
\
} while (0)
-#define TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN_IMPL_(Result, Converted, Type, Name,
ResultExpr, Node) \
+/// \endcond
+
+/// \cond Doxygen_Suppress
+#define TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN_IMPL_(Result, Type, Name,
ResultExpr) \
auto Result = (ResultExpr); /* NOLINT(bugprone-macro-parentheses) */
\
- TVM_FFI_S_MUTATE_MAYBE_EARLY_RETURN(Result, Node);
\
- auto Converted = /* NOLINT(bugprone-macro-parentheses) */
\
- ::tvm::ffi::details::AssignOrReturnHelper(
\
- ::std::move(::tvm::ffi::details::ExpectedUnsafe::GetData(Result)))
\
- .template TryMove<Type>();
\
- TVM_FFI_S_MUTATE_MAYBE_EARLY_RETURN(Converted, Node);
\
- Type Name = ::tvm::ffi::details::AnyUnsafe::MoveFromAnyAfterCheck<Type>(
\
- ::std::move(::tvm::ffi::details::ExpectedUnsafe::GetData(Converted)))
+ TVM_FFI_S_MUTATE_MAYBE_EARLY_RETURN(Result);
\
+ if
(TVM_FFI_PREDICT_FALSE(!::tvm::ffi::details::AnyUnsafe::CheckAnyStrict<Type>(
\
+ ::tvm::ffi::details::ExpectedUnsafe::GetData(Result)))) {
\
+ return
::tvm::ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(::tvm::ffi::Any(::tvm::ffi::Error(
\
+ "TypeError", "structural mutate result does not match the declared
type", ""))); \
+ }
\
+ Type Name = /* NOLINT(bugprone-macro-parentheses) */
\
+ ::tvm::ffi::details::AnyUnsafe::MoveFromAnyAfterCheck<Type>(
\
+ ::std::move(::tvm::ffi::details::ExpectedUnsafe::GetData(Result)))
/// \endcond
/*!
@@ -572,25 +631,56 @@ namespace details {
*
* ``Type`` must be concrete; use a type alias when it contains a top-level
comma. A type mismatch
* returns ``Unexpected(TypeError)`` through the surrounding ``Expected``
function without
- * throwing. This macro declares ``Name`` into the enclosing scope and must be
used in a braced
- * block, never as an unbraced control-flow body. A raw pointer node must be
non-null; pass nullable
- * nodes as ``ObjectRef`` or ``Any`` so ``None`` is skipped when constructing
error context.
+ * throwing, reported with a fixed string so a correct hook pays only one
predicted-not-taken
+ * branch per field. This macro declares ``Name`` into the enclosing scope and
must be used in a
+ * braced block, never as an unbraced control-flow body. A raw pointer node
must be non-null; pass
+ * nullable nodes as ``ObjectRef`` or ``Any`` so ``None`` is skipped when
constructing error
+ * context.
*
* Example:
* \code{.cpp}
- * TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(
- * ObjectRef, child, mutator->MutateExpected(self->child), self);
+ * TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ObjectRef, child,
mutator->MutateExpected(self->child), self);
* \endcode
*
* \param Type The concrete type of the successful value.
* \param Name The name of the value declared in the enclosing scope.
* \param ResultExpr An expression producing the ``Expected`` value to unwrap.
- * \param Node The current mutation node appended to the error context on
failure.
*/
-#define TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(Type, Name, ResultExpr, Node) \
- TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN_IMPL_( \
- TVM_FFI_STR_CONCAT(tvm_ffi_mutate_result_, __COUNTER__), \
- TVM_FFI_STR_CONCAT(tvm_ffi_mutate_converted_, __COUNTER__), Type, Name,
ResultExpr, Node)
+#define TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(Type, Name, ResultExpr)
\
+
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN_IMPL_(TVM_FFI_STR_CONCAT(tvm_ffi_mutate_result_,
__COUNTER__), \
+ Type, Name, ResultExpr)
+
+/// \cond Doxygen_Suppress
+#define TVM_FFI_UNSAFE_S_MUTATE_ASSIGN_OR_RETURN_SKIP_CHECK_IMPL_(Result,
Type, Name, ResultExpr) \
+ auto Result = (ResultExpr); /* NOLINT(bugprone-macro-parentheses) */
\
+ TVM_FFI_S_MUTATE_MAYBE_EARLY_RETURN(Result);
\
+ TVM_FFI_DCHECK(::tvm::ffi::details::AnyUnsafe::CheckAnyStrict<Type>(
\
+ ::tvm::ffi::details::ExpectedUnsafe::GetData(Result)))
\
+ << "unchecked structural-mutate assign: result is not of the declared
type"; \
+ Type Name = /* NOLINT(bugprone-macro-parentheses) */
\
+ ::tvm::ffi::details::AnyUnsafe::MoveFromAnyAfterCheck<Type>(
\
+ ::std::move(::tvm::ffi::details::ExpectedUnsafe::GetData(Result)))
+/// \endcond
+
+/*!
+ * \brief \ref TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN without the type check.
+ *
+ * Same signature and same error propagation; the difference is only what
happens to a
+ * successful result that is not of type \p Type.
+ *
+ * The caller must guarantee the result has the declared type; a mismatch is
undefined behavior
+ * in a release build, and debug builds catch it with ``TVM_FFI_DCHECK``.
+ *
+ * \param Type The concrete type of the successful value.
+ * \param Name The name of the value declared in the enclosing scope.
+ * \param ResultExpr An expression producing the ``Expected`` value to unwrap.
+ */
+#define TVM_FFI_UNSAFE_S_MUTATE_ASSIGN_OR_RETURN_SKIP_CHECK(Type, Name,
ResultExpr) \
+ TVM_FFI_UNSAFE_S_MUTATE_ASSIGN_OR_RETURN_SKIP_CHECK_IMPL_(
\
+ TVM_FFI_STR_CONCAT(tvm_ffi_mutate_result_, __COUNTER__), Type, Name,
ResultExpr)
+
+/// \cond Doxygen_Suppress
+/// \endcond
/*!
* \brief Structural mutator that invokes typed callbacks during recursive
mapping.
@@ -599,45 +689,40 @@ namespace details {
* \tparam Dispatch Callback dispatcher.
* \sa StructuralMapCallbackChain
*/
-template <WalkOrder order, typename Dispatch>
-class StructuralMapMutatorObj : public StructuralMutatorObj {
- public:
- /*!
- * \brief Construct a callback-aware mutator.
- * \param dispatch The composed callback dispatcher.
- */
- explicit StructuralMapMutatorObj(Dispatch dispatch)
- : StructuralMutatorObj(VTable()), dispatch_(std::move(dispatch)) {}
+/*!
+ * \brief A runtime table of Function callbacks, usable as a single link.
+ *
+ * The typed links are matched at compile time from their argument type. The
Python-driven API
+ * instead carries a runtime list keyed by type index, so it appears to the
mutator as one link
+ * that performs its own lookup. That keeps a single traversal for both
dispatch strategies.
+ */
- private:
- /*!
- * \brief Return the shared callback-aware mutator vtable.
- * \return Pointer to the immutable mutator vtable for this specialization.
- */
- static const StructuralMutatorVTable* VTable() {
- static const StructuralMutatorVTable vtable{
- &StructuralMapMutatorObj::DispatchMutate,
- &StructuralMapMutatorObj::DispatchMaybeInplaceMutate,
- &StructuralMapMutatorObj::DispatchVarRemapGet,
- &StructuralMapMutatorObj::DispatchVarRemapSet,
- };
- return &vtable;
- }
+/*!
+ * \brief Base of both the static and the dynamic StructuralMapMutator.
+ *
+ * Carries the identity-substitution environment they share. The dispatch
thunks downcast only
+ * as far as this class, so both reuse them regardless of how they store their
callbacks.
+ *
+ */
+class StructuralMapMutatorBaseObj : public StructuralMutatorObj {
+ public:
+ explicit StructuralMapMutatorBaseObj(const StructuralMutatorVTable* vtable)
+ : StructuralMutatorObj(vtable) {}
+ protected:
/*!
- * \brief Dispatch variable-remap lookup through this mutator's vtable.
+ * \brief Dispatch variable-remap lookup through the mutator vtable.
* \param mutator The erased callback-aware mutator.
* \param var The borrowed variable identity to look up.
* \return Raw ``TVMFFIAny`` containing the owning replacement, FFI None, or
Error.
*/
static TVMFFIAny DispatchVarRemapGet(StructuralMutatorObj* mutator, AnyView
var) noexcept {
- auto* self = static_cast<StructuralMapMutatorObj*>(mutator);
- Expected<Any> result = self->VarRemapGetImpl(var);
- return ExpectedUnsafe::MoveToTVMFFIAny(std::move(result));
+ auto* self = static_cast<StructuralMapMutatorBaseObj*>(mutator);
+ return ExpectedUnsafe::MoveToTVMFFIAny(self->VarRemapGetImpl(var));
}
/*!
- * \brief Dispatch variable-remap insertion through this mutator's vtable.
+ * \brief Dispatch variable-remap insertion through the mutator vtable.
* \param mutator The erased callback-aware mutator.
* \param var The borrowed variable identity to bind.
* \param mapped_value The borrowed replacement value.
@@ -645,13 +730,57 @@ class StructuralMapMutatorObj : public
StructuralMutatorObj {
*/
static TVMFFIAny DispatchVarRemapSet(StructuralMutatorObj* mutator, AnyView
var,
AnyView mapped_value) noexcept {
- auto* self = static_cast<StructuralMapMutatorObj*>(mutator);
- Expected<void> result = self->VarRemapSetImpl(var, mapped_value);
- return ExpectedUnsafe::MoveToTVMFFIAny(std::move(result));
+ auto* self = static_cast<StructuralMapMutatorBaseObj*>(mutator);
+ return ExpectedUnsafe::MoveToTVMFFIAny(self->VarRemapSetImpl(var,
mapped_value));
}
/*!
- * \brief Look up a replacement in this mutator's identity-substitution
environment.
+ * \brief Invoke a matched callback with optional def-region context.
+ * \tparam Callback Callable returning a value implicitly convertible to
``Expected<Any>``.
+ * \tparam Value Type of the converted value passed to the callback.
+ * \param callback The matched callback.
+ * \param value The converted value passed to the callback.
+ * \param kind The active def-region kind.
+ * \return The callback result normalized to ``Expected<Any>``.
+ *
+ * The callback may return a different type than the one that selected the
link; only the
+ * caller holding a field's static type can check that.
+ */
+ template <typename Callback, typename Value>
+ TVM_FFI_INLINE static Expected<Any> InvokeCallbackLink(Callback& callback,
Value&& value,
+ TVMFFIDefRegionKind
kind) {
+ using FuncInfo = FunctionInfo<std::decay_t<Callback>>;
+ static_assert(std::is_convertible_v<typename FuncInfo::RetType,
Expected<Any>>,
+ "StructuralMap callbacks must return a replacement value,
Error, Unexpected, "
+ "or Expected<U> implicitly convertible to Expected<Any>");
+ try {
+ if constexpr (FuncInfo::num_args == 1) {
+ return callback(std::forward<Value>(value));
+ } else {
+ return callback(std::forward<Value>(value), kind);
+ }
+ } catch (const Error& err) {
+ return Unexpected(err);
+ }
+ }
+
+ /*!
+ * \brief Append \p node to a failed result's mutate error context.
+ * \param result The failed result whose Error is annotated.
+ * \param node The borrowed node to name in the context.
+ */
+ TVM_FFI_COLD_CODE static void UpdateVisitErrorContext(const Expected<Any>&
result,
+ AnyView node) noexcept
{
+ // The Error is refcounted, so annotating the local handle annotates the
object the result
+ // holds. A non-object node has no context to add.
+ if (node.type_index() >= TypeIndex::kTVMFFIStaticObjectBegin) {
+ Error err = result.error();
+ ::tvm::ffi::details::UpdateVisitErrorContext(err,
node.cast<ObjectRef>());
+ }
+ }
+
+ /*!
+ * \brief Look up a replacement in the identity-substitution environment.
* \param var The borrowed variable identity to look up.
* \return The owning replacement, FFI None on a miss, or an Error.
*/
@@ -673,7 +802,7 @@ class StructuralMapMutatorObj : public StructuralMutatorObj
{
}
/*!
- * \brief Record a replacement in this mutator's identity-substitution
environment.
+ * \brief Record a replacement in the identity-substitution environment.
* \param var The borrowed variable identity to bind.
* \param mapped_value The borrowed replacement value.
* \return Successful completion, or an Error if the binding cannot be
stored.
@@ -693,6 +822,45 @@ class StructuralMapMutatorObj : public
StructuralMutatorObj {
}
}
+ private:
+ /*! \brief Identity-substitution environment, shared across the whole
traversal. */
+ Map<ObjectRef, Any> var_remap_;
+};
+
+/*!
+ * \brief Structural mutator that invokes statically typed callbacks during
recursive mapping.
+ *
+ * Each callback is an ordinary callable and the engine selects it on its
first argument's
+ * type, so selection is a compile-time-known ``as<TSub>()`` on the input node.
+ *
+ * \tparam order Callback placement relative to child mapping.
+ * \tparam Callbacks The callbacks, tested in declaration order.
+ */
+template <WalkOrder order, typename... Callbacks>
+class StructuralMapMutatorObj : public StructuralMapMutatorBaseObj {
+ public:
+ /*!
+ * \brief Construct a callback-aware mutator that owns its callbacks.
+ * \param callbacks The typed callback links, tested in declaration order.
+ */
+ explicit StructuralMapMutatorObj(Callbacks... callbacks)
+ : StructuralMapMutatorBaseObj(VTable()),
callbacks_(std::move(callbacks)...) {}
+
+ private:
+ /*!
+ * \brief Return the shared callback-aware mutator vtable.
+ * \return Pointer to the immutable mutator vtable for this specialization.
+ */
+ static const StructuralMutatorVTable* VTable() {
+ static const StructuralMutatorVTable vtable{
+ &StructuralMapMutatorObj::DispatchMutate,
+ &StructuralMapMutatorObj::DispatchMaybeInplaceMutate,
+ &StructuralMapMutatorObj::DispatchVarRemapGet,
+ &StructuralMapMutatorObj::DispatchVarRemapSet,
+ };
+ return &vtable;
+ }
+
/*!
* \brief Dispatch callback-aware optional in-place mutation through the ABI
vtable.
* \param mutator The erased callback-aware mutator.
@@ -702,7 +870,7 @@ class StructuralMapMutatorObj : public StructuralMutatorObj
{
static TVMFFIAny DispatchMaybeInplaceMutate(StructuralMutatorObj* mutator,
AnyView value) noexcept {
auto* self = static_cast<StructuralMapMutatorObj*>(mutator);
- return
ExpectedUnsafe::MoveToTVMFFIAny(self->MaybeInplaceMutateImpl(value));
+ return self->MaybeInplaceMutateImplRaw(value);
}
/*!
@@ -713,227 +881,185 @@ class StructuralMapMutatorObj : public
StructuralMutatorObj {
*/
static TVMFFIAny DispatchMutate(StructuralMutatorObj* mutator, AnyView
value) noexcept {
auto* self = static_cast<StructuralMapMutatorObj*>(mutator);
- return ExpectedUnsafe::MoveToTVMFFIAny(self->MutateImpl(value));
+ return self->MutateImplRaw(value);
}
/*!
- * \brief Invoke the callback and select mutation according to its result
and walk order.
+ * \brief Test one link against \p value and, if it matches, mutate the node
through it.
*
- * \param value The borrowed value to mutate.
- * \return The mutated value or an Error.
+ * \tparam kMaybeInplace Whether the caller may mutate a uniquely owned node
in place.
+ * \tparam Callback The link's callback type.
+ * \param callback The link's callback.
+ * \param value The borrowed value to test and mutate.
+ * \param out Receives the mutated value or Error when the link matched.
+ * \return Whether the link matched, in which case \p out was written.
*/
- Expected<Any> MaybeInplaceMutateImpl(AnyView value) noexcept {
- // The link match is tested before any descent, and the identity remap
lives
- // inside the match continuation, so the var-type detection only runs when
a
- // callback signature actually matches. Unmatched nodes skip both it and
the
- // callback chain.
+ template <bool kMaybeInplace, typename Callback>
+ TVM_FFI_INLINE bool TryLink(Callback& callback, AnyView value,
Expected<Any>* out) noexcept {
+ using FuncInfo = FunctionInfo<std::decay_t<Callback>>;
+ static_assert(FuncInfo::num_args == 1 || FuncInfo::num_args == 2,
+ "StructuralMap callbacks must take one argument (value) or
two arguments "
+ "(value, def-region kind)");
+ using FirstArg = std::tuple_element_t<0, typename FuncInfo::ArgType>;
+ using TSub = std::remove_cv_t<std::remove_reference_t<FirstArg>>;
+
+ // Deliberately duplicated by StructuralMapDynMutatorObj::TryLink in
structural_mutate.cc,
+ // which differs only in how a link is found and called; keep the two in
step.
//
- // The ordering is load-bearing: DefaultMutateExpected's reflected fallback
- // writes its own remap entry for this node, so testing the match after
descent
- // would let the remap read that entry, treat the node as already handled,
and
- // never run the callback.
- if constexpr (order == WalkOrder::kPreOrder) {
- return dispatch_(
- value, def_region_kind(),
- [&](auto&& invoke_callback) -> Expected<Any> {
- return MutateWithIdentityRemapExpected(value, [&]() ->
Expected<Any> {
- Expected<Any> callback_result = invoke_callback(value);
- TVM_FFI_S_MUTATE_MAYBE_EARLY_RETURN(callback_result, value);
- // A pre-order result can be mutated in place if unchanged or
uniquely owned.
- const Any& mapped_value =
ExpectedUnsafe::GetData(callback_result);
- const TVMFFIAny* mapped_data =
AnyUnsafe::TVMFFIAnyPtrFromAny(mapped_value);
- const TVMFFIAny input_data = value.CopyToTVMFFIAny();
- if (mapped_data->type_index != input_data.type_index ||
- mapped_data->zero_padding != input_data.zero_padding ||
- mapped_data->v_int64 != input_data.v_int64) {
- const Object* mapped_obj = mapped_value.as<Object>();
- bool can_mutate_mapped_value_inplace =
- mapped_obj != nullptr && mapped_obj->unique();
- Expected<Any> result = can_mutate_mapped_value_inplace
- ?
DefaultMaybeInplaceMutateExpected(mapped_value)
- :
DefaultMutateExpected(mapped_value);
- TVM_FFI_S_MUTATE_MAYBE_EARLY_RETURN(result, mapped_value);
- return result;
- }
- Expected<Any> result = DefaultMaybeInplaceMutateExpected(value);
- TVM_FFI_S_MUTATE_MAYBE_EARLY_RETURN(result, value);
- return result;
- });
- },
- [&]() -> Expected<Any> { return
DefaultMaybeInplaceMutateExpected(value); });
- } else {
- return dispatch_(
- value, def_region_kind(),
- [&](auto&& invoke_callback) -> Expected<Any> {
- return MutateWithIdentityRemapExpected(value, [&]() ->
Expected<Any> {
- Expected<Any> result = DefaultMaybeInplaceMutateExpected(value);
- TVM_FFI_S_MUTATE_MAYBE_EARLY_RETURN(result, value);
-
- const Any& mapped_value = ExpectedUnsafe::GetData(result);
- Expected<Any> callback_result = invoke_callback(mapped_value);
- TVM_FFI_S_MUTATE_MAYBE_EARLY_RETURN(callback_result,
mapped_value);
- return callback_result;
- });
- },
- [&]() -> Expected<Any> { return
DefaultMaybeInplaceMutateExpected(value); });
+ // The match test and the matched-node path live together rather than in
separate functions:
+ // selecting the link already computes value.as<TSub>(), and a pre-order
walk invokes on that
+ // very node, so the converted value is reused instead of converted twice.
+ //
+ // Selection is on the input node, before any descent. A post-order walk
must not select on
+ // the descended node: the reflected fallback writes its own remap entry,
so testing after
+ // descent lets that entry swallow the callback.
+ std::optional<TSub> matched;
+ if constexpr (!std::is_same_v<TSub, AnyView> && !std::is_same_v<TSub,
Any>) {
+ matched = value.template as<TSub>();
+ if (!matched.has_value()) return false;
}
- }
- /*!
- * \brief Mutate one value and invoke its matching callback in the
configured order.
- * \param value The borrowed input value.
- * \return The mutated value or an Error.
- */
- Expected<Any> MutateImpl(AnyView value) noexcept {
- // The link match is tested before any descent, and the identity remap
lives
- // inside the match continuation, so the var-type detection only runs when
a
- // callback signature actually matches. Unmatched nodes skip both it and
the
- // callback chain.
- //
- // The ordering is load-bearing: DefaultMutateExpected's reflected fallback
- // writes its own remap entry for this node, so testing the match after
descent
- // would let the remap read that entry, treat the node as already handled,
and
- // never run the callback.
+ // A FreeVar or DAG node maps once and every later occurrence reuses that
result, so if this
+ // node already has a cached remap entry, return it instead of mutating it
again.
+ const bool remappable = IsRemappableIdentity(value.type_index());
+ if (remappable) {
+ Expected<Any> mapped = VarRemapGetExpected(value);
+ if (mapped.is_err()) {
+ *out = std::move(mapped);
+ return true;
+ }
+ if (ExpectedUnsafe::GetData(mapped).type_index() !=
TypeIndex::kTVMFFINone) {
+ *out = std::move(mapped);
+ return true;
+ }
+ }
+
+ const TVMFFIDefRegionKind kind = def_region_kind();
if constexpr (order == WalkOrder::kPreOrder) {
- return dispatch_(
- value, def_region_kind(),
- [&](auto&& invoke_callback) -> Expected<Any> {
- return MutateWithIdentityRemapExpected(value, [&]() ->
Expected<Any> {
- Expected<Any> callback_result = invoke_callback(value);
- TVM_FFI_S_MUTATE_MAYBE_EARLY_RETURN(callback_result, value);
-
- const Any& mapped_value =
ExpectedUnsafe::GetData(callback_result);
- Expected<Any> result = DefaultMutateExpected(mapped_value);
- TVM_FFI_S_MUTATE_MAYBE_EARLY_RETURN(result, mapped_value);
- return result;
- });
- },
- [&]() -> Expected<Any> { return DefaultMutateExpected(value); });
+ // Pre-order: the callback rewrites this node first, then descent runs
over whatever it
+ // produced, so a replacement subtree is itself mapped.
+ Expected<Any> callback_result = [&]() -> Expected<Any> {
+ if constexpr (std::is_same_v<TSub, AnyView>) {
+ return InvokeCallbackLink(callback, value, kind);
+ } else if constexpr (std::is_same_v<TSub, Any>) {
+ return InvokeCallbackLink(callback, Any(value), kind);
+ } else {
+ // Reuses the conversion the match already performed.
+ return InvokeCallbackLink(callback, *std::move(matched), kind);
+ }
+ }();
+ if (TVM_FFI_PREDICT_FALSE(callback_result.is_err())) {
+ UpdateVisitErrorContext(callback_result, value);
+ *out = std::move(callback_result);
+ return true;
+ }
+ // Own the callback's result: it is the only reference from here on, and
moving it out
+ // beats holding a reference into an Expected that stays alive across
the descent below.
+ Any mapped_value = ExpectedUnsafe::GetData(callback_result);
+ // Each descent names the node it actually ran on in the error context.
+ *out = [&]() -> Expected<Any> {
+ if constexpr (kMaybeInplace) {
+ // A pre-order result can be mutated in place if unchanged or
uniquely owned.
+ const TVMFFIAny* mapped_data =
AnyUnsafe::TVMFFIAnyPtrFromAny(mapped_value);
+ const TVMFFIAny input_data = value.CopyToTVMFFIAny();
+ if (mapped_data->type_index == input_data.type_index &&
+ mapped_data->zero_padding == input_data.zero_padding &&
+ mapped_data->v_int64 == input_data.v_int64) {
+ return DefaultMaybeInplaceMutateExpected(value);
+ }
+ const Object* mapped_obj = mapped_value.as<Object>();
+ bool can_inplace = mapped_obj != nullptr && mapped_obj->unique();
+ return can_inplace ? DefaultMaybeInplaceMutateExpected(mapped_value)
+ : DefaultMutateExpected(mapped_value);
+ } else {
+ return DefaultMutateExpected(mapped_value);
+ }
+ }();
+ if (TVM_FFI_PREDICT_FALSE(out->is_err())) return true;
} else {
- return dispatch_(
- value, def_region_kind(),
- [&](auto&& invoke_callback) -> Expected<Any> {
- return MutateWithIdentityRemapExpected(value, [&]() ->
Expected<Any> {
- Expected<Any> result = DefaultMutateExpected(value);
- TVM_FFI_S_MUTATE_MAYBE_EARLY_RETURN(result, value);
-
- const Any& mapped_value = ExpectedUnsafe::GetData(result);
- Expected<Any> callback_result = invoke_callback(mapped_value);
- TVM_FFI_S_MUTATE_MAYBE_EARLY_RETURN(callback_result,
mapped_value);
- return callback_result;
- });
- },
- [&]() -> Expected<Any> { return DefaultMutateExpected(value); });
+ // Post-order: children are mapped first and the callback sees the
rebuilt node, so it
+ // observes its operands already substituted.
+ Expected<Any> descended =
+ kMaybeInplace ? DefaultMaybeInplaceMutateExpected(value) :
DefaultMutateExpected(value);
+ if (TVM_FFI_PREDICT_FALSE(descended.is_err())) {
+ *out = std::move(descended);
+ return true;
+ }
+ // Held by reference, not moved out: the error path below names this
node, so it has to
+ // survive the callback. Only a storage-enabled TSub could gain from a
move, and that is
+ // exactly the case whose move would empty it.
+ const Any& mapped_value = ExpectedUnsafe::GetData(descended);
+ *out = [&]() -> Expected<Any> {
+ if constexpr (std::is_same_v<TSub, AnyView>) {
+ return InvokeCallbackLink(callback, AnyView(mapped_value), kind);
+ } else if constexpr (std::is_same_v<TSub, Any>) {
+ return InvokeCallbackLink(callback, Any(mapped_value), kind);
+ } else {
+ // Re-converted rather than reusing the match: the callback is
invoked on the node
+ // descent handed back, and must only see the type it asked for.
Default mutation is
+ // required to preserve the type, so failing here means some hook
broke that.
+ std::optional<TSub> descended_sub = mapped_value.template as<TSub>();
+ if (TVM_FFI_PREDICT_FALSE(!descended_sub.has_value())) {
+ return Unexpected(
+ Error("TypeError", "structural mutate: descent changed the
node type", ""));
+ }
+ return InvokeCallbackLink(callback, *std::move(descended_sub), kind);
+ }
+ }();
+ if (TVM_FFI_PREDICT_FALSE(out->is_err())) {
+ UpdateVisitErrorContext(*out, mapped_value);
+ return true;
+ }
}
- }
-
- /*! \brief Composed callback dispatcher owned by this mutator. */
- Dispatch dispatch_;
- /*! \brief Identity-substitution table. */
- Map<ObjectRef, Any> var_remap_;
-};
+ // Bind this node's identity to its final result, so every later
occurrence reuses it.
+ if (remappable) {
+ Expected<void> set_result = VarRemapSetExpected(value,
ExpectedUnsafe::GetData(*out));
+ if (TVM_FFI_PREDICT_FALSE(set_result.is_err())) {
+ *out = Unexpected(std::move(set_result).error());
+ }
+ }
+ return true;
+ }
-// Build a callback dispatcher from a typed callback chain. The dispatcher
answers "which link
-// matches this node" as control flow rather than as a value. It takes the
node, the active
-// def-region kind, an on-match continuation, and an on-no-match continuation,
and returns whichever
-// one applies. The on-match continuation receives a callable
`invoke_callback(AnyView target)` that
-// converts target to the matched link's argument type and invokes it, so the
caller decides what
-// wraps the invocation. A node that matches no link therefore pays one type
test per link and
-// nothing else.
-struct StructuralMapCallbackChain {
- public:
/*!
- * \brief Construct a dispatcher owning \p callbacks.
- * \tparam Callbacks Callback types.
- * \param callbacks Callbacks tested in declaration order.
- * \return A callback-aware structural-map dispatcher.
+ * \brief Test every link in declaration order, stopping at the first that
matches.
+ * \return Whether some link matched, in which case \p out was written.
*/
- template <typename... Callbacks>
- static auto FromChain(Callbacks... callbacks) {
- auto callback_tuple = std::make_tuple(std::move(callbacks)...);
- return [callbacks = std::move(callback_tuple)](AnyView value,
TVMFFIDefRegionKind kind,
- auto&& on_match,
- auto&& on_no_match) mutable
-> Expected<Any> {
- std::optional<Expected<Any>> result;
- auto run_match = [&](auto&& invoke_callback) {
-
result.emplace(on_match(std::forward<decltype(invoke_callback)>(invoke_callback)));
- };
- // Selection is control flow: the first matching link hands its
invocation thunk to
- // run_match, while a node matching no link goes directly to on_no_match.
- bool matched = std::apply(
- [&](auto&... callback) { return (... || TryCallLink(callback, value,
kind, run_match)); },
- callbacks);
- if (matched) {
- return *std::move(result);
- }
- return on_no_match();
- };
+ template <bool kMaybeInplace, size_t... Is>
+ TVM_FFI_INLINE bool TryLinks(AnyView value, Expected<Any>* out,
+ std::index_sequence<Is...>) noexcept {
+ return (TryLink<kMaybeInplace>(std::get<Is>(callbacks_), value, out) ||
...);
}
- private:
- // Hand the matched link's invocation thunk to on_match, or return false on
no match. The
- // target conversion stays in the thunk so post-order can select on the
input and invoke on the
- // node whose children have already been mapped.
- template <typename Callback, typename OnMatch>
- TVM_FFI_INLINE static bool TryCallLink(Callback& callback, AnyView value,
- TVMFFIDefRegionKind kind, OnMatch&&
on_match) {
- using FuncInfo = FunctionInfo<std::decay_t<Callback>>;
- static_assert(FuncInfo::num_args == 1 || FuncInfo::num_args == 2,
- "StructuralMap callbacks must take one argument (value) or
two arguments "
- "(value, def-region kind)");
- using FirstArg = std::tuple_element_t<0, typename FuncInfo::ArgType>;
- using TSub = std::remove_cv_t<std::remove_reference_t<FirstArg>>;
- if constexpr (std::is_same_v<TSub, AnyView>) {
- on_match([&](AnyView target) -> Expected<Any> {
- return InvokeCallbackLink(callback, target, kind);
- });
- return true;
- } else if constexpr (std::is_same_v<TSub, Any>) {
- on_match([&](AnyView target) -> Expected<Any> {
- return InvokeCallbackLink(callback, Any(target), kind);
- });
- return true;
- } else if (auto opt = value.template as<TSub>()) {
- on_match([&](AnyView target) -> Expected<Any> {
- std::optional<TSub> converted = target.template as<TSub>();
- if (!converted.has_value()) {
- // A custom post-order hook may change the node type after the input
selected this link.
- return Any(target);
- }
- return InvokeCallbackLink(callback, *std::move(converted), kind);
- });
- return true;
+ /*!
+ * \brief Mutate a value, invoking the first matching callback link.
+ * \param value The borrowed value to mutate.
+ * \return Raw ``TVMFFIAny`` containing the mutated value or Error.
+ */
+ TVM_FFI_INLINE TVMFFIAny MutateImplRaw(AnyView value) noexcept {
+ Expected<Any> out{Any()};
+ if (TryLinks<false>(value, &out, std::index_sequence_for<Callbacks...>{}))
{
+ return ExpectedUnsafe::MoveToTVMFFIAny(std::move(out));
}
- return false;
+ return DefaultMutateRaw(value);
}
/*!
- * \brief Invoke a matched callback with optional def-region context.
- * \tparam Callback Callable returning a value implicitly convertible to
``Expected<Any>``.
- * \tparam Value Type of the converted value passed to the callback.
- * \param callback The matched callback.
- * \param value The converted value passed to the callback.
- * \param kind The active def-region kind.
- * \return The callback result normalized to ``Expected<Any>``.
+ * \brief Mutate a value in place when safe, invoking the first matching
callback link.
+ * \param value The borrowed value to mutate.
+ * \return Raw ``TVMFFIAny`` containing the mutated value or Error.
*/
- template <typename Callback, typename Value>
- TVM_FFI_INLINE static Expected<Any> InvokeCallbackLink(Callback& callback,
Value&& value,
- TVMFFIDefRegionKind
kind) {
- using FuncInfo = FunctionInfo<std::decay_t<Callback>>;
- static_assert(std::is_convertible_v<typename FuncInfo::RetType,
Expected<Any>>,
- "StructuralMap callbacks must return a replacement value,
Error, Unexpected, "
- "or Expected<U> implicitly convertible to Expected<Any>");
- try {
- if constexpr (FuncInfo::num_args == 1) {
- return callback(std::forward<Value>(value));
- } else {
- return callback(std::forward<Value>(value), kind);
- }
- } catch (const Error& err) {
- return Unexpected(err);
+ TVM_FFI_INLINE TVMFFIAny MaybeInplaceMutateImplRaw(AnyView value) noexcept {
+ Expected<Any> out{Any()};
+ if (TryLinks<true>(value, &out, std::index_sequence_for<Callbacks...>{})) {
+ return ExpectedUnsafe::MoveToTVMFFIAny(std::move(out));
}
+ return DefaultMaybeInplaceMutateRaw(value);
}
+
+ /*! \brief The callback links, tested in declaration order. */
+ std::tuple<Callbacks...> callbacks_;
};
} // namespace details
@@ -993,10 +1119,8 @@ struct StructuralMapCallbackChain {
template <WalkOrder order, typename... Callbacks>
Expected<Any> StructuralMapExpected(AnyView root, Callbacks&&... callbacks)
noexcept {
static_assert(sizeof...(Callbacks) != 0, "StructuralMap requires at least
one callback");
- auto dispatch =
-
details::StructuralMapCallbackChain::FromChain(std::forward<Callbacks>(callbacks)...);
- using Mutator = details::StructuralMapMutatorObj<order, decltype(dispatch)>;
- StructuralMutator mutator(make_object<Mutator>(std::move(dispatch)));
+ using Mutator = details::StructuralMapMutatorObj<order,
std::decay_t<Callbacks>...>;
+ StructuralMutator
mutator(make_object<Mutator>(std::forward<Callbacks>(callbacks)...));
return mutator->MaybeInplaceMutateIfUniqueExpected(root);
}
diff --git a/include/tvm/ffi/extra/structural_visit.h
b/include/tvm/ffi/extra/structural_visit.h
index 89f53fac..468bfd4f 100644
--- a/include/tvm/ffi/extra/structural_visit.h
+++ b/include/tvm/ffi/extra/structural_visit.h
@@ -205,16 +205,27 @@ class StructuralVisitorObj : public Object {
static reflection::TypeAttrColumn
column(reflection::type_attr::kStructuralVisit);
AnyView attr = column[type_index];
+ // Hooks propagate errors untouched; this is the engine dispatching into
`value`, so the
+ // node is named here and nowhere else -- exactly one frame per node.
// case 1: Type-specific override registered as an opaque ABI visit
function pointer.
if (attr.type_index() == TypeIndex::kTVMFFIOpaquePtr) {
auto* visit_fn = reinterpret_cast<FStructuralVisit>(attr.cast<void*>());
- return
details::ExpectedUnsafe::MoveFromTVMFFIAny<Optional<VisitInterrupt>>(
- (*visit_fn)(this, value));
+ TVMFFIAny raw = (*visit_fn)(this, value);
+ if (TVM_FFI_PREDICT_FALSE(raw.type_index == TypeIndex::kTVMFFIError)) {
+ details::UpdateVisitErrorContext(raw, value);
+ }
+ return
details::ExpectedUnsafe::MoveFromTVMFFIAny<Optional<VisitInterrupt>>(raw);
}
// case 2: Type-specific override registered as an ffi::Function.
if (attr.type_index() == TypeIndex::kTVMFFIFunction) {
- return
attr.cast<Function>().CallExpected<Optional<VisitInterrupt>>(this, value);
+ Expected<Optional<VisitInterrupt>> result =
+ attr.cast<Function>().CallExpected<Optional<VisitInterrupt>>(this,
value);
+ if (TVM_FFI_PREDICT_FALSE(result.is_err())) {
+ Error err = result.error();
+ details::UpdateVisitErrorContext(err, value);
+ }
+ return result;
}
if (TVM_FFI_PREDICT_FALSE(attr.type_index() != TypeIndex::kTVMFFINone)) {
@@ -451,23 +462,13 @@ namespace details {
// If Result is an Error, append Node to the visit error context before
returning. Node is
// required: dropping it silently degrades every error message produced below
this frame.
// A raw pointer Node must be non-null; pass nullable nodes as ObjectRef or
Any so None is skipped.
-#define TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Result, Node)
\
- do {
\
- auto&& tvm_ffi_res_ = (Result);
\
- if (TVM_FFI_PREDICT_FALSE(
\
-
::tvm::ffi::details::StructuralVisitNeedEarlyReturn(tvm_ffi_res_))) {
\
- if (TVM_FFI_PREDICT_FALSE(tvm_ffi_res_.type_index() ==
\
- ::tvm::ffi::TypeIndex::kTVMFFIError)) {
\
- auto&& tvm_ffi_visit_node_owner_ = (Node);
\
- ::tvm::ffi::AnyView tvm_ffi_visit_node_ = tvm_ffi_visit_node_owner_;
\
- if (tvm_ffi_visit_node_.type_index() >=
::tvm::ffi::TypeIndex::kTVMFFIStaticObjectBegin) { \
- ::tvm::ffi::Error tvm_ffi_visit_err_ = tvm_ffi_res_.error();
\
- ::tvm::ffi::details::UpdateVisitErrorContext(
\
- tvm_ffi_visit_err_,
tvm_ffi_visit_node_.cast<::tvm::ffi::ObjectRef>()); \
- }
\
- }
\
- return
::tvm::ffi::details::ExpectedUnsafe::MoveToTVMFFIAny(std::move(tvm_ffi_res_));
\
- }
\
+#define TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Result)
\
+ do {
\
+ auto&& tvm_ffi_res_ = (Result);
\
+ if (TVM_FFI_PREDICT_FALSE(
\
+
::tvm::ffi::details::StructuralVisitNeedEarlyReturn(tvm_ffi_res_))) {
\
+ return
::tvm::ffi::details::ExpectedUnsafe::MoveToTVMFFIAny(std::move(tvm_ffi_res_)); \
+ }
\
} while (0)
/// \endcond
@@ -523,7 +524,13 @@ class StructuralWalkVisitorObj : public
StructuralVisitorObj {
}
if constexpr (order == WalkOrder::kPreOrder) {
auto result = dispatch_(value, this->def_region_kind());
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(result, value);
+ if
(TVM_FFI_PREDICT_FALSE(details::StructuralVisitNeedEarlyReturn(result))) {
+ if (TVM_FFI_PREDICT_FALSE(result.is_err())) {
+ Error err = result.error();
+ details::UpdateVisitErrorContext(err, value);
+ }
+ return details::ExpectedUnsafe::MoveToTVMFFIAny(std::move(result));
+ }
// Hoist the call out of TVM_FFI_UNSAFE_ASSUME: clang's -Wassume rejects
// arguments that contain a call expression (its potential side effects
// would be discarded), while [[maybe_unused]] keeps -Wunused-variable
@@ -537,10 +544,23 @@ class StructuralWalkVisitorObj : public
StructuralVisitorObj {
}
}
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(DefaultVisitExpected(value), value);
+ {
+ // DefaultVisitExpected already named `value` if a hook it dispatched
failed.
+ auto result = DefaultVisitExpected(value);
+ if
(TVM_FFI_PREDICT_FALSE(details::StructuralVisitNeedEarlyReturn(result))) {
+ return details::ExpectedUnsafe::MoveToTVMFFIAny(std::move(result));
+ }
+ }
if constexpr (order == WalkOrder::kPostOrder) {
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(dispatch_(value,
this->def_region_kind()), value);
+ auto result = dispatch_(value, this->def_region_kind());
+ if
(TVM_FFI_PREDICT_FALSE(details::StructuralVisitNeedEarlyReturn(result))) {
+ if (TVM_FFI_PREDICT_FALSE(result.is_err())) {
+ Error err = result.error();
+ details::UpdateVisitErrorContext(err, value);
+ }
+ return details::ExpectedUnsafe::MoveToTVMFFIAny(std::move(result));
+ }
}
return details::ExpectedUnsafe::MoveToTVMFFIAny(
diff --git a/include/tvm/ffi/extra/visit_error_context.h
b/include/tvm/ffi/extra/visit_error_context.h
index e8b5c482..76c49375 100644
--- a/include/tvm/ffi/extra/visit_error_context.h
+++ b/include/tvm/ffi/extra/visit_error_context.h
@@ -264,6 +264,44 @@ inline void UpdateVisitErrorContext(Error& err, const
ObjectRef& node) { // NOL
error_obj->extra_context =
details::ObjectUnsafe::MoveObjectPtrToTVMFFIObjectPtr(std::move(new_context));
}
+
+/*!
+ * \brief Name \p node in a raw failed result's visit error context.
+ *
+ * The overload the visit and mutate engines use at a hook boundary: hooks
propagate errors
+ * untouched, and the engine, which knows the node it dispatched on, adds that
frame. The Error
+ * is refcounted, so annotating this handle annotates the object \p result
carries.
+ *
+ * \param result A raw result already known to hold an Error.
+ * \param node The node to name; a non-object node has no frame to add.
+ */
+TVM_FFI_COLD_CODE inline void UpdateVisitErrorContext(const TVMFFIAny& result,
+ AnyView node) noexcept {
+ if (node.type_index() >= TypeIndex::kTVMFFIStaticObjectBegin) {
+ Error err = AnyView::CopyFromTVMFFIAny(result).cast<Error>();
+ UpdateVisitErrorContext(err, node.cast<ObjectRef>());
+ }
+}
+
+/*!
+ * \brief Name \p node in a failed result's visit error context.
+ *
+ * The guarded counterpart to the ObjectRef overload, for engine call sites
that hold a borrowed
+ * node of unknown kind. A visited node may be a primitive -- a container
element, a reflected
+ * field -- and `AnyView::cast<ObjectRef>()` throws on one. The engines that
call this are
+ * `noexcept`, so an unguarded cast would terminate rather than propagate;
skipping the frame is
+ * correct because a non-object node has no identity to name in the context.
+ *
+ * \param err The Error to annotate. Refcounted, so annotating this handle
annotates the shared
+ * object the caller's result carries.
+ * \param node The node to name; skipped when it is not object-backed.
+ */
+TVM_FFI_COLD_CODE inline void UpdateVisitErrorContext(Error& err,
+ AnyView node) noexcept {
// NOLINT(*)
+ if (node.type_index() >= TypeIndex::kTVMFFIStaticObjectBegin) {
+ UpdateVisitErrorContext(err, node.cast<ObjectRef>());
+ }
+}
} // namespace details
} // namespace ffi
diff --git a/src/ffi/extra/structural_mutate.cc
b/src/ffi/extra/structural_mutate.cc
index b8e5a072..1705046c 100644
--- a/src/ffi/extra/structural_mutate.cc
+++ b/src/ffi/extra/structural_mutate.cc
@@ -47,40 +47,225 @@ namespace details {
* \param order Integer value of \ref WalkOrder.
* \return The mapped owning value, or an Error.
*/
-Expected<Any> StructuralMapExpected(
- AnyView root, const Array<Tuple<int32_t, Function>>& callbacks,
- const Array<Tuple<int32_t, Function>>& callbacks_with_def_region_kind, int
order) noexcept {
- auto dispatch = [callbacks, callbacks_with_def_region_kind](AnyView x,
TVMFFIDefRegionKind kind,
- auto&& on_match,
- auto&&
on_no_match) -> Expected<Any> {
- for (const auto& entry : callbacks) {
- int32_t type_index = entry.template get<0>();
- if (!RuntimeTypeIndexMatch(x.type_index(), type_index)) {
- continue;
+/*!
+ * \brief Structural mutator whose links are runtime ffi.Functions keyed by
type index.
+ *
+ * The dynamic counterpart of \ref StructuralMapMutatorObj. It selects a link
by scanning
+ * registered type indices where the static version does a compile-time
``as<TSub>()``. The two
+ * are kept apart, rather than sharing one template with a mode flag, so each
version reads
+ * straight through; only the identity remap they share lives in the common
base. Keeping this
+ * one in the .cc also keeps the link table out of the public header.
+ *
+ * \tparam order Callback placement relative to child mapping.
+ */
+template <WalkOrder order>
+class StructuralMapDynMutatorObj : public StructuralMapMutatorBaseObj {
+ public:
+ StructuralMapDynMutatorObj(Array<Tuple<int32_t, Function>> callbacks,
+ Array<Tuple<int32_t, Function>>
callbacks_with_def_region_kind)
+ : StructuralMapMutatorBaseObj(VTable()),
+ callbacks_(std::move(callbacks)),
+
callbacks_with_def_region_kind_(std::move(callbacks_with_def_region_kind)) {}
+
+ private:
+ static const StructuralMutatorVTable* VTable() {
+ static const StructuralMutatorVTable vtable{
+ &StructuralMapDynMutatorObj::DispatchMutate,
+ &StructuralMapDynMutatorObj::DispatchMaybeInplaceMutate,
+ &StructuralMapDynMutatorObj::DispatchVarRemapGet,
+ &StructuralMapDynMutatorObj::DispatchVarRemapSet,
+ };
+ return &vtable;
+ }
+
+ static TVMFFIAny DispatchMaybeInplaceMutate(StructuralMutatorObj* mutator,
+ AnyView value) noexcept {
+ return
static_cast<StructuralMapDynMutatorObj*>(mutator)->MaybeInplaceMutateImplRaw(value);
+ }
+
+ static TVMFFIAny DispatchMutate(StructuralMutatorObj* mutator, AnyView
value) noexcept {
+ return
static_cast<StructuralMapDynMutatorObj*>(mutator)->MutateImplRaw(value);
+ }
+
+ /*!
+ * \brief Find the first link registered for \p type_index.
+ *
+ * \param type_index The input node's runtime type index.
+ * \param with_kind Set when the matched link also takes a def-region kind.
+ * \param link_type_index Set to the registered type index the link matched
on, so a post-order
+ * walk can recheck the descended node against the same target.
+ * \return The matched Function, or nullopt when no link applies.
+ */
+ Optional<Function> FindLink(int32_t type_index, bool* with_kind,
+ int32_t* link_type_index) const noexcept {
+ for (const Tuple<int32_t, Function>& entry : callbacks_) {
+ if (RuntimeTypeIndexMatch(type_index, entry.get<0>())) {
+ *with_kind = false;
+ *link_type_index = entry.get<0>();
+ return entry.get<1>();
}
- Function fn = entry.template get<1>();
- return on_match(
- [&](AnyView target) -> Expected<Any> { return
fn.CallExpected<Any>(target); });
}
- for (const auto& entry : callbacks_with_def_region_kind) {
- int32_t type_index = entry.template get<0>();
- if (!RuntimeTypeIndexMatch(x.type_index(), type_index)) {
- continue;
+ for (const Tuple<int32_t, Function>& entry :
callbacks_with_def_region_kind_) {
+ if (RuntimeTypeIndexMatch(type_index, entry.get<0>())) {
+ *with_kind = true;
+ *link_type_index = entry.get<0>();
+ return entry.get<1>();
}
- Function fn = entry.template get<1>();
- return on_match(
- [&](AnyView target) -> Expected<Any> { return
fn.CallExpected<Any>(target, kind); });
}
- return on_no_match();
- };
+ return std::nullopt;
+ }
+
+ /*!
+ * \brief Invoke a matched link, threading the live def-region kind through.
+ *
+ * \p kind is passed rather than stashed at selection time: the engine only
knows the true
+ * def-region kind of the node handed back after descent (post-order) or of
the matched node
+ * itself (pre-order).
+ */
+ TVM_FFI_INLINE static Expected<Any> InvokeLink(const Function& fn, bool
with_kind, AnyView target,
+ TVMFFIDefRegionKind kind)
noexcept {
+ // CallExpected is exception-free: it goes through the safe-call path and
returns any raised
+ // error as Unexpected, unlike a directly invoked C++ callback.
+ return with_kind ? fn.CallExpected<Any>(target, kind) :
fn.CallExpected<Any>(target);
+ }
+
+ /*!
+ * \brief Test the link table against \p value and mutate through the first
match.
+ *
+ * \tparam kMaybeInplace Whether a uniquely owned node may be mutated in
place.
+ * \param value The borrowed value to test and mutate.
+ * \param out Receives the mutated value or Error when a link matched.
+ * \return Whether a link matched, in which case \p out was written.
+ */
+ template <bool kMaybeInplace>
+ TVM_FFI_INLINE bool TryLink(AnyView value, Expected<Any>* out) noexcept {
+ // Step for step the same walk as StructuralMapMutatorObj::TryLink, and
deliberately so:
+ // only link detection and invocation differ between the two, and keeping
them as separate
+ // straight-line copies lets each specialize on its own selection strategy
and keeps both
+ // readable. Everything below except finding and calling the link is
shared semantics, so a
+ // change to either copy belongs in both.
+
+ bool with_kind = false;
+ int32_t link_type_index = TypeIndex::kTVMFFINone;
+ // A local, so descending into a matching child cannot change what this
node invokes.
+ Optional<Function> matched = FindLink(value.type_index(), &with_kind,
&link_type_index);
+ if (!matched.has_value()) return false;
+ // --- identity remap, entry half -----------------------------------------
+ // A FreeVar or DAG node maps once and every later occurrence reuses that
result, so if this
+ // node already has a cached remap entry, return it instead of mutating it
again.
+ const bool remappable = IsRemappableIdentity(value.type_index());
+ if (remappable) {
+ Expected<Any> mapped = VarRemapGetExpected(value);
+ if (mapped.is_err()) {
+ *out = std::move(mapped);
+ return true;
+ }
+ if (ExpectedUnsafe::GetData(mapped).type_index() !=
TypeIndex::kTVMFFINone) {
+ *out = std::move(mapped);
+ return true;
+ }
+ }
+
+ // --- callback and descent, in walk order --------------------------------
+ const TVMFFIDefRegionKind kind = def_region_kind();
+ if constexpr (order == WalkOrder::kPreOrder) {
+ // Pre-order: the callback rewrites this node first, then descent runs
over what it made.
+ Expected<Any> callback_result = InvokeLink(*matched, with_kind, value,
kind);
+ if (TVM_FFI_PREDICT_FALSE(callback_result.is_err())) {
+ UpdateVisitErrorContext(callback_result, value);
+ *out = std::move(callback_result);
+ return true;
+ }
+ // Own the callback's result: it is the only reference from here on.
+ Any mapped_value = ExpectedUnsafe::GetData(callback_result);
+ // Each descent names the node it actually ran on in the error context.
+ *out = [&]() -> Expected<Any> {
+ if constexpr (kMaybeInplace) {
+ const TVMFFIAny* mapped_data =
AnyUnsafe::TVMFFIAnyPtrFromAny(mapped_value);
+ const TVMFFIAny input_data = value.CopyToTVMFFIAny();
+ if (mapped_data->type_index == input_data.type_index &&
+ mapped_data->zero_padding == input_data.zero_padding &&
+ mapped_data->v_int64 == input_data.v_int64) {
+ return DefaultMaybeInplaceMutateExpected(value);
+ }
+ const Object* mapped_obj = mapped_value.as<Object>();
+ bool can_inplace = mapped_obj != nullptr && mapped_obj->unique();
+ return can_inplace ? DefaultMaybeInplaceMutateExpected(mapped_value)
+ : DefaultMutateExpected(mapped_value);
+ } else {
+ return DefaultMutateExpected(mapped_value);
+ }
+ }();
+ if (TVM_FFI_PREDICT_FALSE(out->is_err())) return true;
+ } else {
+ // Post-order: children are mapped first, so the callback sees the
rebuilt node.
+ Expected<Any> descended =
+ kMaybeInplace ? DefaultMaybeInplaceMutateExpected(value) :
DefaultMutateExpected(value);
+ if (TVM_FFI_PREDICT_FALSE(descended.is_err())) {
+ *out = std::move(descended);
+ return true;
+ }
+ // Held by reference, not moved out: the error path below names this
node, so it has to
+ // survive the callback.
+ const Any& mapped_value = ExpectedUnsafe::GetData(descended);
+ // The link was selected on the input node, and the callback must only
see the type it
+ // registered for. The typed mutator gets this from its
`mapped_value.as<TSub>()`; here the
+ // registered type index is the same target, so recheck against it.
+ if (TVM_FFI_PREDICT_FALSE(
+ !RuntimeTypeIndexMatch(mapped_value.type_index(),
link_type_index))) {
+ *out =
+ Unexpected(Error("TypeError", "structural mutate: descent changed
the node type", ""));
+ UpdateVisitErrorContext(*out, mapped_value);
+ return true;
+ }
+ *out = InvokeLink(*matched, with_kind, mapped_value, kind);
+ if (TVM_FFI_PREDICT_FALSE(out->is_err())) {
+ UpdateVisitErrorContext(*out, mapped_value);
+ return true;
+ }
+ }
+
+ // Bind this node's identity to its final result, so every later
occurrence reuses it.
+ if (remappable) {
+ Expected<void> set_result = VarRemapSetExpected(value,
ExpectedUnsafe::GetData(*out));
+ if (TVM_FFI_PREDICT_FALSE(set_result.is_err())) {
+ *out = Unexpected(std::move(set_result).error());
+ }
+ }
+ return true;
+ }
+
+ TVM_FFI_INLINE TVMFFIAny MutateImplRaw(AnyView value) noexcept {
+ Expected<Any> out{Any()};
+ if (TryLink<false>(value, &out)) {
+ return ExpectedUnsafe::MoveToTVMFFIAny(std::move(out));
+ }
+ return DefaultMutateRaw(value);
+ }
+
+ TVM_FFI_INLINE TVMFFIAny MaybeInplaceMutateImplRaw(AnyView value) noexcept {
+ Expected<Any> out{Any()};
+ if (TryLink<true>(value, &out)) {
+ return ExpectedUnsafe::MoveToTVMFFIAny(std::move(out));
+ }
+ return DefaultMaybeInplaceMutateRaw(value);
+ }
+
+ Array<Tuple<int32_t, Function>> callbacks_;
+ Array<Tuple<int32_t, Function>> callbacks_with_def_region_kind_;
+};
+
+Expected<Any> StructuralMapExpected(
+ AnyView root, const Array<Tuple<int32_t, Function>>& callbacks,
+ const Array<Tuple<int32_t, Function>>& callbacks_with_def_region_kind, int
order) noexcept {
if (order == static_cast<int>(WalkOrder::kPreOrder)) {
- using Mutator = StructuralMapMutatorObj<WalkOrder::kPreOrder,
decltype(dispatch)>;
- StructuralMutator mutator(make_object<Mutator>(std::move(dispatch)));
+ using Mutator = StructuralMapDynMutatorObj<WalkOrder::kPreOrder>;
+ StructuralMutator mutator(make_object<Mutator>(callbacks,
callbacks_with_def_region_kind));
return mutator->MaybeInplaceMutateIfUniqueExpected(root);
} else {
- using Mutator = StructuralMapMutatorObj<WalkOrder::kPostOrder,
decltype(dispatch)>;
- StructuralMutator mutator(make_object<Mutator>(std::move(dispatch)));
+ using Mutator = StructuralMapDynMutatorObj<WalkOrder::kPostOrder>;
+ StructuralMutator mutator(make_object<Mutator>(callbacks,
callbacks_with_def_region_kind));
return mutator->MaybeInplaceMutateIfUniqueExpected(root);
}
}
@@ -99,15 +284,15 @@ Expected<Any> StructuralMapExpected(
* \return The mutated sequence, or an Error.
*/
template <typename SeqObj>
-Expected<Any> MutateSeqContainerExpected(StructuralMutatorObj* mutator,
AnyView value,
- const SeqObj* self) noexcept {
+TVMFFIAny MutateSeqContainerRaw(StructuralMutatorObj* mutator, AnyView value,
+ const SeqObj* self) noexcept {
int64_t size = static_cast<int64_t>(self->size());
const Any* items = self->begin();
ObjectPtr<SeqObj> output = nullptr;
for (int64_t i = 0; i < size; ++i) {
const Any& item = items[i];
- TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(Any, mapped_value,
mutator->MutateExpected(item), self);
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(Any, mapped_value,
mutator->MutateExpected(item));
if (output == nullptr) {
if (item.same_as(mapped_value)) {
@@ -120,9 +305,9 @@ Expected<Any>
MutateSeqContainerExpected(StructuralMutatorObj* mutator, AnyView
}
if (output == nullptr) {
- return Any(value);
+ return AnyUnsafe::MoveAnyToTVMFFIAny(Any(value));
}
- return Any(ObjectRef(std::move(output)));
+ return AnyUnsafe::MoveAnyToTVMFFIAny(Any(std::move(output)));
}
/*!
@@ -135,18 +320,18 @@ Expected<Any>
MutateSeqContainerExpected(StructuralMutatorObj* mutator, AnyView
* \return The mutated sequence, or an Error.
*/
template <typename SeqObj>
-Expected<Any> MaybeInplaceMutateSeqContainerExpected(StructuralMutatorObj*
mutator, AnyView value,
- SeqObj* self) noexcept {
+TVMFFIAny MaybeInplaceMutateSeqContainerRaw(StructuralMutatorObj* mutator,
AnyView value,
+ SeqObj* self) noexcept {
for (int64_t i = 0; i < static_cast<int64_t>(self->size()); ++i) {
const Any& item = self->begin()[i];
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(Any, mapped_value,
-
mutator->MaybeInplaceMutateIfUniqueExpected(item), self);
+
mutator->MaybeInplaceMutateIfUniqueExpected(item));
if (!item.same_as(mapped_value)) {
self->SetItemAfterCheck(i, std::move(mapped_value));
}
}
- return Any(value);
+ return AnyUnsafe::MoveAnyToTVMFFIAny(Any(value));
}
/*!
@@ -159,15 +344,15 @@ Expected<Any>
MaybeInplaceMutateSeqContainerExpected(StructuralMutatorObj* mutat
* \return The mutated map, or an Error.
*/
template <typename MapObjType>
-Expected<Any> MutateMapValuesExpected(StructuralMutatorObj* mutator, AnyView
value,
- const MapObjType* self) noexcept {
+TVMFFIAny MutateMapValuesRaw(StructuralMutatorObj* mutator, AnyView value,
+ const MapObjType* self) noexcept {
ObjectPtr<Object> output = nullptr;
MapBaseObj::iterator output_it;
size_t index = 0;
for (auto source_it = self->begin(); source_it != self->end(); ++source_it,
++index) {
const Any& old_value = source_it->second;
- TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(Any, new_value,
mutator->MutateExpected(old_value), self);
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(Any, new_value,
mutator->MutateExpected(old_value));
bool changed = !old_value.same_as(new_value);
if (output == nullptr) {
if (!changed) {
@@ -186,9 +371,9 @@ Expected<Any> MutateMapValuesExpected(StructuralMutatorObj*
mutator, AnyView val
}
if (output == nullptr) {
- return Any(value);
+ return AnyUnsafe::MoveAnyToTVMFFIAny(Any(value));
}
- return Any(ObjectRef(std::move(output)));
+ return AnyUnsafe::MoveAnyToTVMFFIAny(Any(std::move(output)));
}
/*!
@@ -201,18 +386,18 @@ Expected<Any>
MutateMapValuesExpected(StructuralMutatorObj* mutator, AnyView val
* \return The mutated map, or an Error.
*/
template <typename MapObjType>
-Expected<Any> MaybeInplaceMutateMapValuesExpected(StructuralMutatorObj*
mutator, AnyView value,
- MapObjType* self) noexcept {
+TVMFFIAny MaybeInplaceMutateMapValuesRaw(StructuralMutatorObj* mutator,
AnyView value,
+ MapObjType* self) noexcept {
for (auto it = self->begin(); it != self->end(); ++it) {
const Any& old_value = it->second;
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(Any, new_value,
-
mutator->MaybeInplaceMutateIfUniqueExpected(old_value), self);
+
mutator->MaybeInplaceMutateIfUniqueExpected(old_value));
if (!old_value.same_as(new_value)) {
it->second = std::move(new_value);
}
}
- return Any(value);
+ return AnyUnsafe::MoveAnyToTVMFFIAny(Any(value));
}
/*! \brief Identity structural mutation hook for immutable String and Bytes
leaves. */
@@ -223,58 +408,50 @@ TVMFFIAny MutateImmutableLeaf(StructuralMutatorObj*,
AnyView value) noexcept {
/*! \brief Structural mutation hook for ArrayObj. */
TVMFFIAny MutateArray(StructuralMutatorObj* mutator, AnyView value) noexcept {
- Expected<Any> result = MutateSeqContainerExpected(
+ return MutateSeqContainerRaw(
mutator, value,
details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const ArrayObj>(value));
- return ExpectedUnsafe::MoveToTVMFFIAny(std::move(result));
}
/*! \brief Maybe-in-place structural mutation hook for ArrayObj. */
TVMFFIAny MaybeInplaceMutateArray(StructuralMutatorObj* mutator, AnyView
value) noexcept {
- Expected<Any> result = MaybeInplaceMutateSeqContainerExpected(
+ return MaybeInplaceMutateSeqContainerRaw(
mutator, value,
details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<ArrayObj>(value));
- return ExpectedUnsafe::MoveToTVMFFIAny(std::move(result));
}
/*! \brief Structural mutation hook for ListObj. */
TVMFFIAny MutateList(StructuralMutatorObj* mutator, AnyView value) noexcept {
- Expected<Any> result = MutateSeqContainerExpected(
+ return MutateSeqContainerRaw(
mutator, value,
details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const ListObj>(value));
- return ExpectedUnsafe::MoveToTVMFFIAny(std::move(result));
}
/*! \brief Maybe-in-place structural mutation hook for ListObj. */
TVMFFIAny MaybeInplaceMutateList(StructuralMutatorObj* mutator, AnyView value)
noexcept {
- Expected<Any> result = MaybeInplaceMutateSeqContainerExpected(
+ return MaybeInplaceMutateSeqContainerRaw(
mutator, value,
details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<ListObj>(value));
- return ExpectedUnsafe::MoveToTVMFFIAny(std::move(result));
}
/*! \brief Structural mutation hook for MapObj. */
TVMFFIAny MutateMap(StructuralMutatorObj* mutator, AnyView value) noexcept {
- Expected<Any> result = MutateMapValuesExpected(
+ return MutateMapValuesRaw(
mutator, value,
details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const MapObj>(value));
- return ExpectedUnsafe::MoveToTVMFFIAny(std::move(result));
}
/*! \brief Maybe-in-place structural mutation hook for MapObj. */
TVMFFIAny MaybeInplaceMutateMap(StructuralMutatorObj* mutator, AnyView value)
noexcept {
- Expected<Any> result = MaybeInplaceMutateMapValuesExpected(
+ return MaybeInplaceMutateMapValuesRaw(
mutator, value,
details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<MapObj>(value));
- return ExpectedUnsafe::MoveToTVMFFIAny(std::move(result));
}
/*! \brief Structural mutation hook for DictObj. */
TVMFFIAny MutateDict(StructuralMutatorObj* mutator, AnyView value) noexcept {
- Expected<Any> result = MutateMapValuesExpected(
+ return MutateMapValuesRaw(
mutator, value,
details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const DictObj>(value));
- return ExpectedUnsafe::MoveToTVMFFIAny(std::move(result));
}
/*! \brief Maybe-in-place structural mutation hook for DictObj. */
TVMFFIAny MaybeInplaceMutateDict(StructuralMutatorObj* mutator, AnyView value)
noexcept {
- Expected<Any> result = MaybeInplaceMutateMapValuesExpected(
+ return MaybeInplaceMutateMapValuesRaw(
mutator, value,
details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<DictObj>(value));
- return ExpectedUnsafe::MoveToTVMFFIAny(std::move(result));
}
} // namespace details
diff --git a/src/ffi/extra/structural_visit.cc
b/src/ffi/extra/structural_visit.cc
index 95bf80d8..a0ca4d9d 100644
--- a/src/ffi/extra/structural_visit.cc
+++ b/src/ffi/extra/structural_visit.cc
@@ -87,7 +87,7 @@ Expected<Optional<VisitInterrupt>> StructuralWalkExpected(
/*! \brief Visit entries in a sequence container. */
TVMFFIAny VisitSeqContainer(StructuralVisitorObj* visitor, const SeqBaseObj*
self) noexcept {
for (const Any& item : *self) {
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(item), self);
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(item));
}
return
ExpectedUnsafe::MoveToTVMFFIAny(Expected<Optional<VisitInterrupt>>(std::nullopt));
}
@@ -95,7 +95,7 @@ TVMFFIAny VisitSeqContainer(StructuralVisitorObj* visitor,
const SeqBaseObj* sel
/*! \brief Visit values in a map container while treating keys as structural
anchors. */
TVMFFIAny VisitMapContainer(StructuralVisitorObj* visitor, const MapBaseObj*
self) noexcept {
for (const auto& kv : *self) {
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(kv.second),
self);
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(kv.second));
}
return
ExpectedUnsafe::MoveToTVMFFIAny(Expected<Optional<VisitInterrupt>>(std::nullopt));
}
diff --git a/tests/cpp/extra/test_structural_mutate.cc
b/tests/cpp/extra/test_structural_mutate.cc
index c5ae8599..162cb38d 100644
--- a/tests/cpp/extra/test_structural_mutate.cc
+++ b/tests/cpp/extra/test_structural_mutate.cc
@@ -401,4 +401,29 @@ TEST(StructuralMap,
HandlesInlineAndHeapStringAndBytesLeaves) {
CheckStringAndBytesLeaves<WalkOrder::kPostOrder>();
}
+// The dynamic mutator duplicates the static one's walk deliberately, so the
semantics they
+// share need coverage on this copy too. Driven through the same entry point
Python uses.
+Any CallDynStructuralMap(AnyView root, const Array<Tuple<int32_t, Function>>&
callbacks,
+ WalkOrder order) {
+ Function fn = Function::GetGlobalRequired("ffi.StructuralMap");
+ return fn(root, callbacks, Array<Tuple<int32_t, Function>>(),
static_cast<int32_t>(order));
+}
+
+TEST(StructuralMapDyn, ReusesRemapResultForRepeatedVar) {
+ // A FreeVar maps once and every later occurrence reuses that result. Both
mutators share this
+ // half of the walk, so it must hold identically here.
+ TVar var("x");
+ AnyArray root{Any(var), Any(var)};
+ int64_t calls = 0;
+ Function remap = Function::FromTyped([&](AnyView v) -> Any {
+ ++calls;
+ return Any(TVar(v.cast<TVar>()->name + "-mapped"));
+ });
+ Any mapped = CallDynStructuralMap(
+ root, {Tuple<int32_t, Function>(TVarObj::RuntimeTypeIndex(), remap)},
WalkOrder::kPostOrder);
+ auto arr = mapped.cast<AnyArray>();
+ EXPECT_EQ(calls, 1);
+ EXPECT_TRUE(arr[0].cast<TVar>().same_as(arr[1].cast<TVar>()));
+}
+
} // namespace
diff --git a/tests/cpp/extra/test_structural_visit.cc
b/tests/cpp/extra/test_structural_visit.cc
index 6ac34a43..eeb71fc6 100644
--- a/tests/cpp/extra/test_structural_visit.cc
+++ b/tests/cpp/extra/test_structural_visit.cc
@@ -429,6 +429,31 @@ TEST(StructuralVisitor, WalkReturnsError) {
EXPECT_EQ(result.error().message(), "walk callback failed");
}
+// A callback that fails on a primitive node must surface the error, not abort.
+//
+// On failure the engine names the node it dispatched on in the visit error
context. A visited
+// node is not always object-backed -- a container element or reflected field
may be a primitive --
+// and `AnyView::cast<ObjectRef>()` throws on one. The walk is `noexcept`, so
naming the node
+// without a type-index guard terminates the process instead of returning the
error. Every other
+// error test above fails on an ObjectRef node, so only a primitive exercises
that guard.
+TEST(StructuralVisitor, WalkReturnsErrorOnPrimitiveNode) {
+ Array<Any> root = {Any(static_cast<int64_t>(1))};
+ auto fail = [](int64_t) -> Expected<WalkResult> {
+ return Unexpected(Error("ValueError", "walk callback failed on primitive",
""));
+ };
+
+ Expected<Optional<VisitInterrupt>> pre =
StructuralWalkExpected<WalkOrder::kPreOrder>(root, fail);
+ ASSERT_TRUE(pre.is_err());
+ EXPECT_EQ(pre.error().kind(), "ValueError");
+ EXPECT_EQ(pre.error().message(), "walk callback failed on primitive");
+
+ Expected<Optional<VisitInterrupt>> post =
+ StructuralWalkExpected<WalkOrder::kPostOrder>(root, fail);
+ ASSERT_TRUE(post.is_err());
+ EXPECT_EQ(post.error().kind(), "ValueError");
+ EXPECT_EQ(post.error().message(), "walk callback failed on primitive");
+}
+
TEST(StructuralVisitor, WalkCatchesError) {
ObjectRef root = TVar("root");
diff --git a/tests/cpp/testing_object.h b/tests/cpp/testing_object.h
index 32e1e917..b5cf64e0 100644
--- a/tests/cpp/testing_object.h
+++ b/tests/cpp/testing_object.h
@@ -249,20 +249,16 @@ class TMutatePairObj : public Object {
return count;
}
- static Expected<Any> StructuralMutateExpected(StructuralMutatorObj* mutator,
- AnyView value) noexcept {
+ static TVMFFIAny StructuralMutate(StructuralMutatorObj* mutator, AnyView
value) noexcept {
++StructuralMutateCallCount();
const auto* self = value.cast<const TMutatePairObj*>();
- TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ObjectRef, lhs,
mutator->MutateExpected(self->lhs), self);
- TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ObjectRef, rhs,
mutator->MutateExpected(self->rhs), self);
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ObjectRef, lhs,
mutator->MutateExpected(self->lhs));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ObjectRef, rhs,
mutator->MutateExpected(self->rhs));
if (lhs.same_as(self->lhs) && rhs.same_as(self->rhs)) {
- return Any(value);
+ return details::AnyUnsafe::MoveAnyToTVMFFIAny(Any(value));
}
- return Any(ObjectRef(make_object<TMutatePairObj>(std::move(lhs),
std::move(rhs))));
- }
-
- static TVMFFIAny StructuralMutate(StructuralMutatorObj* mutator, AnyView
value) noexcept {
- return
details::ExpectedUnsafe::MoveToTVMFFIAny(StructuralMutateExpected(mutator,
value));
+ return details::AnyUnsafe::MoveAnyToTVMFFIAny(
+ Any(make_object<TMutatePairObj>(std::move(lhs), std::move(rhs))));
}
static void RegisterReflection() {
@@ -413,13 +409,11 @@ class TFuncObj : public Object {
static TVMFFIAny StructuralVisit(StructuralVisitorObj* visitor, AnyView
value) noexcept {
const auto* self = value.cast<const TFuncObj*>();
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(
- visitor->WithDefRegionKind(kTVMFFIDefRegionKindRecursive,
- [&]() { return
visitor->VisitExpected(self->params); }),
- self);
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->WithDefRegionKind(
+ kTVMFFIDefRegionKindRecursive, [&]() { return
visitor->VisitExpected(self->params); }));
auto body_result = visitor->VisitExpected(self->body);
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(body_result, self);
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(body_result);
return details::ExpectedUnsafe::MoveToTVMFFIAny(std::move(body_result));
}
