This is an automated email from the ASF dual-hosted git repository.
cyx-6 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 bdd4989f [REFACTOR][EXTRA] Optimize StructuralMap (#733)
bdd4989f is described below
commit bdd4989f9c9ad4a5e8fd61db01c185577570dc51
Author: Tianqi Chen <[email protected]>
AuthorDate: Thu Sep 3 04:58:17 2026 -0400
[REFACTOR][EXTRA] Optimize StructuralMap (#733)
This PR refactor optimizes `StructuralMap`.
- Move successful `Expected` payloads through structural-mutation hooks
without redundant reference-count traffic, while returning conversion
failures through `Expected`/`Unexpected`.
- Require structural visit and mutation early returns to carry node
context, and document the public mutation-unwrapping macro.
- Match callbacks before recursive descent so unmatched nodes skip
identity-remap lookup, error-frame setup, and callback-chain work.
- Keep remap lookup and callback invocation inside each matched
continuation, preserving callback ordering and post-order target
revalidation.
- Add registered-hook and error-context coverage for successful,
nullable, failing, and wrong-type mutation results.
---
include/tvm/ffi/expected.h | 44 ++++-
include/tvm/ffi/extra/structural_mutate.h | 293 ++++++++++++++++++++----------
include/tvm/ffi/extra/structural_visit.h | 53 +++---
src/ffi/extra/structural_mutate.cc | 80 ++++----
src/ffi/extra/structural_visit.cc | 12 +-
tests/cpp/extra/test_structural_mutate.cc | 50 +++++
tests/cpp/testing_object.h | 65 ++++++-
7 files changed, 411 insertions(+), 186 deletions(-)
diff --git a/include/tvm/ffi/expected.h b/include/tvm/ffi/expected.h
index 4d31fdce..bd7417fa 100644
--- a/include/tvm/ffi/expected.h
+++ b/include/tvm/ffi/expected.h
@@ -348,6 +348,29 @@ 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.
*
@@ -380,10 +403,25 @@ struct ExpectedUnsafe {
}
/*!
- * \brief Return the underlying Any storage.
+ * \brief Return the underlying Any storage as an xvalue; by-value
initialization moves from it.
+ * \tparam T The Expected success type.
+ * \param result The Expected value whose storage will be exposed as an
xvalue.
+ * \return An xvalue reference to the underlying Any storage.
+ *
+ * \note The const overload returns ``const Any&``, which remains const
under ``std::move`` and
+ * therefore selects the copy constructor. The assign-or-return macro
selects this overload
+ * from a non-const result, so its outer ``std::move`` only makes the
move intent explicit.
+ */
+ template <typename T>
+ TVM_FFI_INLINE static Any&& GetData(Expected<T>& result) noexcept {
+ return std::move(result.data_);
+ }
+
+ /*!
+ * \brief Return a const reference to the underlying Any storage.
* \tparam T The Expected success type.
- * \param result The Expected value to inspect.
- * \return Const reference to the raw Any storage.
+ * \param result The Expected value whose storage will be viewed.
+ * \return A const reference to the underlying Any storage.
*/
template <typename T>
TVM_FFI_INLINE static const Any& GetData(const Expected<T>& result) noexcept
{
diff --git a/include/tvm/ffi/extra/structural_mutate.h
b/include/tvm/ffi/extra/structural_mutate.h
index c7916b79..56eff51c 100644
--- a/include/tvm/ffi/extra/structural_mutate.h
+++ b/include/tvm/ffi/extra/structural_mutate.h
@@ -532,27 +532,66 @@ namespace details {
/// \cond Doxygen_Suppress
// Return from the current mutation function if Result is an Error.
-// Append Node to the mutate error context before returning.
-#define TVM_FFI_S_MUTATE_MAYBE_EARLY_RETURN_WITH_ERROR_CONTEXT(Result, Node)
\
+// 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)) { \
- if ((Node).type_index() >=
::tvm::ffi::TypeIndex::kTVMFFIStaticObjectBegin) { \
+ 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_,
\
-
(Node).cast<::tvm::ffi::ObjectRef>()); \
+ ::tvm::ffi::details::UpdateVisitErrorContext(
\
+ tvm_ffi_mutate_err_,
tvm_ffi_mutate_node_.cast<::tvm::ffi::ObjectRef>()); \
}
\
return ::std::move(tvm_ffi_res_);
\
}
\
} while (0)
+
+#define TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN_IMPL_(Result, Converted, Type, Name,
ResultExpr, Node) \
+ 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)))
/// \endcond
+/*!
+ * \brief Unwrap a successful mutation result into a newly declared value or
return its error.
+ *
+ * ``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.
+ *
+ * Example:
+ * \code{.cpp}
+ * 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)
+
/*!
* \brief Structural mutator that invokes typed callbacks during recursive
mapping.
*
* \tparam order Callback placement relative to child mapping.
- * \tparam Dispatch Callback dispatcher with signature
- * ``Expected<Any>(AnyView, TVMFFIDefRegionKind)``.
+ * \tparam Dispatch Callback dispatcher.
* \sa StructuralMapCallbackChain
*/
template <WalkOrder order, typename Dispatch>
@@ -679,38 +718,60 @@ class StructuralMapMutatorObj : public
StructuralMutatorObj {
* \return The mutated value or an Error.
*/
Expected<Any> MaybeInplaceMutateImpl(AnyView value) noexcept {
- return MutateWithIdentityRemapExpected(value, [&]() -> Expected<Any> {
- if constexpr (order == WalkOrder::kPreOrder) {
- Expected<Any> callback_result = dispatch_(value, def_region_kind());
-
TVM_FFI_S_MUTATE_MAYBE_EARLY_RETURN_WITH_ERROR_CONTEXT(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_WITH_ERROR_CONTEXT(result,
mapped_value);
- return result;
- }
- Expected<Any> result = DefaultMaybeInplaceMutateExpected(value);
- TVM_FFI_S_MUTATE_MAYBE_EARLY_RETURN_WITH_ERROR_CONTEXT(result, value);
- return result;
- } else {
- Expected<Any> result = DefaultMaybeInplaceMutateExpected(value);
- TVM_FFI_S_MUTATE_MAYBE_EARLY_RETURN_WITH_ERROR_CONTEXT(result, value);
-
- const Any& mapped_value = ExpectedUnsafe::GetData(result);
- Expected<Any> callback_result = dispatch_(mapped_value,
def_region_kind());
-
TVM_FFI_S_MUTATE_MAYBE_EARLY_RETURN_WITH_ERROR_CONTEXT(callback_result,
mapped_value);
- return callback_result;
- }
- });
+ // 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.
+ 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); });
+ }
}
/*!
@@ -719,25 +780,46 @@ class StructuralMapMutatorObj : public
StructuralMutatorObj {
* \return The mutated value or an Error.
*/
Expected<Any> MutateImpl(AnyView value) noexcept {
- return MutateWithIdentityRemapExpected(value, [&]() -> Expected<Any> {
- if constexpr (order == WalkOrder::kPreOrder) {
- Expected<Any> callback_result = dispatch_(value, def_region_kind());
-
TVM_FFI_S_MUTATE_MAYBE_EARLY_RETURN_WITH_ERROR_CONTEXT(callback_result, value);
-
- const Any& mapped_value = ExpectedUnsafe::GetData(callback_result);
- Expected<Any> result = DefaultMutateExpected(mapped_value);
- TVM_FFI_S_MUTATE_MAYBE_EARLY_RETURN_WITH_ERROR_CONTEXT(result,
mapped_value);
- return result;
- } else {
- Expected<Any> result = DefaultMutateExpected(value);
- TVM_FFI_S_MUTATE_MAYBE_EARLY_RETURN_WITH_ERROR_CONTEXT(result, value);
-
- const Any& mapped_value = ExpectedUnsafe::GetData(result);
- Expected<Any> callback_result = dispatch_(mapped_value,
def_region_kind());
-
TVM_FFI_S_MUTATE_MAYBE_EARLY_RETURN_WITH_ERROR_CONTEXT(callback_result,
mapped_value);
- return callback_result;
- }
- });
+ // 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.
+ 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); });
+ } 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); });
+ }
}
/*! \brief Composed callback dispatcher owned by this mutator. */
@@ -747,9 +829,13 @@ class StructuralMapMutatorObj : public
StructuralMutatorObj {
Map<ObjectRef, Any> var_remap_;
};
-/*!
- * \brief Build a callback dispatcher from a typed callback chain.
- */
+// 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:
/*!
@@ -761,37 +847,32 @@ struct StructuralMapCallbackChain {
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) mutable ->
Expected<Any> {
- try {
- std::optional<Expected<Any>> result;
- // Fold expression: each TryCallLink returns empty std::optional on
no-match
- // (falsy) or a result on match (truthy); || short-circuits on first
match.
- std::apply(
- [&](auto&... callback) { (... || (result = TryCallLink(callback,
value, kind))); },
- callbacks);
- if (result.has_value()) {
- return *std::move(result);
- }
- return Any(value);
- } catch (const Error& err) {
- return Unexpected(err);
+ 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();
};
}
private:
- /*!
- * \brief Invoke \p callback when \p value matches its first argument.
- * \tparam Callback Callback type whose first argument selects the value
type.
- * \param callback The callback under test.
- * \param value The value to match and pass to the callback.
- * \param kind The active def-region kind.
- * \return The callback result on match, or an empty ``std::optional``
otherwise.
- */
- template <typename Callback>
- TVM_FFI_INLINE static std::optional<Expected<Any>> TryCallLink(Callback&
callback, AnyView value,
-
TVMFFIDefRegionKind kind) {
+ // 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 "
@@ -799,15 +880,27 @@ struct StructuralMapCallbackChain {
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>) {
- return InvokeCallbackLink(callback, value, kind);
+ on_match([&](AnyView target) -> Expected<Any> {
+ return InvokeCallbackLink(callback, target, kind);
+ });
+ return true;
} else if constexpr (std::is_same_v<TSub, Any>) {
- return InvokeCallbackLink(callback, Any(value), kind);
- } else {
- if (auto opt = value.template as<TSub>()) {
- return InvokeCallbackLink(callback, *std::move(opt), kind);
- }
+ 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;
}
- return std::nullopt;
+ return false;
}
/*!
@@ -826,10 +919,14 @@ struct StructuralMapCallbackChain {
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>");
- if constexpr (FuncInfo::num_args == 1) {
- return callback(std::forward<Value>(value));
- } else {
- return callback(std::forward<Value>(value), kind);
+ 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);
}
}
};
diff --git a/include/tvm/ffi/extra/structural_visit.h
b/include/tvm/ffi/extra/structural_visit.h
index 463843b6..89f53fac 100644
--- a/include/tvm/ffi/extra/structural_visit.h
+++ b/include/tvm/ffi/extra/structural_visit.h
@@ -448,32 +448,26 @@ namespace details {
/// \cond Doxygen_Suppress
// Return from the current ABI visit function if Result stops traversal.
// Result must evaluate to Expected whose raw storage can be moved to
TVMFFIAny.
-#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)
-
-// Return from the current ABI visit function if Result stops traversal.
-// If Result is an Error, append Node to the visit error context before
returning.
-#define TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN_WITH_ERROR_CONTEXT(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)) {
\
- if ((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_,
\
-
(Node).cast<::tvm::ffi::ObjectRef>()); \
- }
\
- }
\
- return
::tvm::ffi::details::ExpectedUnsafe::MoveToTVMFFIAny(std::move(tvm_ffi_res_));
\
- }
\
+// 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_));
\
+ }
\
} while (0)
/// \endcond
@@ -529,7 +523,7 @@ class StructuralWalkVisitorObj : public
StructuralVisitorObj {
}
if constexpr (order == WalkOrder::kPreOrder) {
auto result = dispatch_(value, this->def_region_kind());
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN_WITH_ERROR_CONTEXT(result, value);
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(result, value);
// 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
@@ -543,11 +537,10 @@ class StructuralWalkVisitorObj : public
StructuralVisitorObj {
}
}
-
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN_WITH_ERROR_CONTEXT(DefaultVisitExpected(value),
value);
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(DefaultVisitExpected(value), value);
if constexpr (order == WalkOrder::kPostOrder) {
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN_WITH_ERROR_CONTEXT(
- dispatch_(value, this->def_region_kind()), value);
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(dispatch_(value,
this->def_region_kind()), value);
}
return details::ExpectedUnsafe::MoveToTVMFFIAny(
diff --git a/src/ffi/extra/structural_mutate.cc
b/src/ffi/extra/structural_mutate.cc
index c240392c..b1b74c0f 100644
--- a/src/ffi/extra/structural_mutate.cc
+++ b/src/ffi/extra/structural_mutate.cc
@@ -50,15 +50,17 @@ namespace details {
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) -> Expected<Any> {
+ 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;
}
Function fn = entry.template get<1>();
- return fn.CallExpected<Any>(x);
+ 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>();
@@ -66,9 +68,10 @@ Expected<Any> StructuralMapExpected(
continue;
}
Function fn = entry.template get<1>();
- return fn.CallExpected<Any>(x, kind);
+ return on_match(
+ [&](AnyView target) -> Expected<Any> { return
fn.CallExpected<Any>(target, kind); });
}
- return Any(x);
+ return on_no_match();
};
if (order == static_cast<int>(WalkOrder::kPreOrder)) {
@@ -92,23 +95,19 @@ Expected<Any> StructuralMapExpected(
* \tparam SeqObj The underlying sequence object type.
* \param mutator The active structural mutator.
* \param value The borrowed sequence container.
- * \param source The sequence object stored in \p value.
+ * \param self The sequence object stored in \p value.
* \return The mutated sequence, or an Error.
*/
template <typename SeqObj>
Expected<Any> MutateSeqContainerExpected(StructuralMutatorObj* mutator,
AnyView value,
- const SeqObj* source) noexcept {
+ const SeqObj* self) noexcept {
try {
- int64_t size = static_cast<int64_t>(source->size());
+ int64_t size = static_cast<int64_t>(self->size());
ObjectPtr<SeqObj> output = nullptr;
for (int64_t i = 0; i < size; ++i) {
- const Any& item = source->at(i);
- Expected<Any> mapped_item = mutator->MutateExpected(item);
- if (TVM_FFI_PREDICT_FALSE(mapped_item.is_err())) {
- return Unexpected(std::move(mapped_item).error());
- }
- const Any& mapped_value = details::ExpectedUnsafe::GetData(mapped_item);
+ const Any& item = self->at(i);
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(Any, mapped_value,
mutator->MutateExpected(item), self);
if (output == nullptr) {
if (item.same_as(mapped_value)) {
@@ -116,10 +115,10 @@ Expected<Any>
MutateSeqContainerExpected(StructuralMutatorObj* mutator, AnyView
}
output = SeqObj::CreateRepeated(size, Any());
for (int64_t j = 0; j < i; ++j) {
- output->SetItem(j, source->at(j));
+ output->SetItem(j, self->at(j));
}
}
- output->SetItem(i, mapped_value);
+ output->SetItem(i, std::move(mapped_value));
}
if (output == nullptr) {
@@ -137,23 +136,20 @@ Expected<Any>
MutateSeqContainerExpected(StructuralMutatorObj* mutator, AnyView
* \tparam SeqObj The underlying sequence object type.
* \param mutator The active structural mutator.
* \param value The borrowed sequence container, which must be safe to mutate
in place.
- * \param target The sequence object stored in \p value.
+ * \param self The sequence object stored in \p value.
* \return The mutated sequence, or an Error.
*/
template <typename SeqObj>
Expected<Any> MaybeInplaceMutateSeqContainerExpected(StructuralMutatorObj*
mutator, AnyView value,
- SeqObj* target) noexcept {
+ SeqObj* self) noexcept {
try {
- for (int64_t i = 0; i < static_cast<int64_t>(target->size()); ++i) {
- const Any& item = target->at(i);
- Expected<Any> mapped_item =
mutator->MaybeInplaceMutateIfUniqueExpected(item);
- if (TVM_FFI_PREDICT_FALSE(mapped_item.is_err())) {
- return Unexpected(std::move(mapped_item).error());
- }
- const Any& mapped_value = details::ExpectedUnsafe::GetData(mapped_item);
+ for (int64_t i = 0; i < static_cast<int64_t>(self->size()); ++i) {
+ const Any& item = self->at(i);
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(Any, mapped_value,
+
mutator->MaybeInplaceMutateIfUniqueExpected(item), self);
if (!item.same_as(mapped_value)) {
- target->SetItem(i, mapped_value);
+ self->SetItem(i, std::move(mapped_value));
}
}
return Any(value);
@@ -168,38 +164,33 @@ Expected<Any>
MaybeInplaceMutateSeqContainerExpected(StructuralMutatorObj* mutat
* \tparam MapObjType The underlying map object type.
* \param mutator The active structural mutator.
* \param value The borrowed map container.
- * \param source The map object stored in \p value.
+ * \param self The map object stored in \p value.
* \return The mutated map, or an Error.
*/
template <typename MapObjType>
Expected<Any> MutateMapValuesExpected(StructuralMutatorObj* mutator, AnyView
value,
- const MapObjType* source) noexcept {
+ const MapObjType* self) noexcept {
try {
ObjectPtr<Object> output = nullptr;
MapBaseObj::iterator output_it;
size_t index = 0;
- for (auto source_it = source->begin(); source_it != source->end();
++source_it, ++index) {
+ for (auto source_it = self->begin(); source_it != self->end();
++source_it, ++index) {
const Any& old_value = source_it->second;
- Expected<Any> mapped_value = mutator->MutateExpected(old_value);
- if (TVM_FFI_PREDICT_FALSE(mapped_value.is_err())) {
- return Unexpected(std::move(mapped_value).error());
- }
-
- const Any& new_value = details::ExpectedUnsafe::GetData(mapped_value);
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(Any, new_value,
mutator->MutateExpected(old_value), self);
bool changed = !old_value.same_as(new_value);
if (output == nullptr) {
if (!changed) {
continue;
}
- output = MapObjType::ShallowCopy(source);
+ output = MapObjType::ShallowCopy(self);
output_it = static_cast<MapBaseObj*>(output.get())->begin();
for (size_t i = 0; i < index; ++i) {
++output_it;
}
}
if (changed) {
- output_it->second = new_value;
+ output_it->second = std::move(new_value);
}
++output_it;
}
@@ -219,23 +210,20 @@ Expected<Any>
MutateMapValuesExpected(StructuralMutatorObj* mutator, AnyView val
* \tparam MapObjType The underlying map object type.
* \param mutator The active structural mutator.
* \param value The borrowed map container, which must be safe to mutate in
place.
- * \param target The map object stored in \p value.
+ * \param self The map object stored in \p value.
* \return The mutated map, or an Error.
*/
template <typename MapObjType>
Expected<Any> MaybeInplaceMutateMapValuesExpected(StructuralMutatorObj*
mutator, AnyView value,
- MapObjType* target) noexcept
{
+ MapObjType* self) noexcept {
try {
- for (auto it = target->begin(); it != target->end(); ++it) {
+ for (auto it = self->begin(); it != self->end(); ++it) {
const Any& old_value = it->second;
- Expected<Any> mapped_value =
mutator->MaybeInplaceMutateIfUniqueExpected(old_value);
- if (TVM_FFI_PREDICT_FALSE(mapped_value.is_err())) {
- return Unexpected(std::move(mapped_value).error());
- }
- const Any& new_value = details::ExpectedUnsafe::GetData(mapped_value);
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(
+ Any, new_value,
mutator->MaybeInplaceMutateIfUniqueExpected(old_value), self);
if (!old_value.same_as(new_value)) {
- it->second = new_value;
+ it->second = std::move(new_value);
}
}
return Any(value);
diff --git a/src/ffi/extra/structural_visit.cc
b/src/ffi/extra/structural_visit.cc
index bcd4ce41..95bf80d8 100644
--- a/src/ffi/extra/structural_visit.cc
+++ b/src/ffi/extra/structural_visit.cc
@@ -85,17 +85,17 @@ Expected<Optional<VisitInterrupt>> StructuralWalkExpected(
}
/*! \brief Visit entries in a sequence container. */
-TVMFFIAny VisitSeqContainer(StructuralVisitorObj* visitor, const SeqBaseObj*
seq) noexcept {
- for (const Any& item : *seq) {
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(item));
+TVMFFIAny VisitSeqContainer(StructuralVisitorObj* visitor, const SeqBaseObj*
self) noexcept {
+ for (const Any& item : *self) {
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(item), self);
}
return
ExpectedUnsafe::MoveToTVMFFIAny(Expected<Optional<VisitInterrupt>>(std::nullopt));
}
/*! \brief Visit values in a map container while treating keys as structural
anchors. */
-TVMFFIAny VisitMapContainer(StructuralVisitorObj* visitor, const MapBaseObj*
map) noexcept {
- for (const auto& kv : *map) {
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(kv.second));
+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);
}
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 d2feb251..8d8f87a9 100644
--- a/tests/cpp/extra/test_structural_mutate.cc
+++ b/tests/cpp/extra/test_structural_mutate.cc
@@ -38,6 +38,8 @@ using namespace tvm::ffi::testing;
using AnyArray = Array<Any>;
using StringMap = Map<String, Any>;
+TVM_FFI_STATIC_INIT_BLOCK() { TMutatePairObj::RegisterReflection(); }
+
Expected<Any> Increment(int64_t value) { return Any(value + 1); }
template <WalkOrder order>
@@ -88,6 +90,54 @@ TEST(StructuralMap, MapsNestedArrayAndMapInConfiguredOrder) {
CheckNestedArrayMapOrder<WalkOrder::kPostOrder>({"int", "inner-array",
"map", "outer-array"});
}
+TEST(StructuralMap, RegisteredMutateHookUsesAssignOrReturn) {
+ TVar lhs("lhs");
+ TVar rhs("rhs");
+ TMutatePair root(lhs, rhs);
+ TMutatePairObj::StructuralMutateCallCount() = 0;
+
+ TMutatePair mapped =
+ StructuralMap<WalkOrder::kPostOrder>(root, [](const TVarObj* var) ->
Expected<Any> {
+ return Any(TVar(var->name + "-mapped"));
+ }).cast<TMutatePair>();
+
+ EXPECT_EQ(TMutatePairObj::StructuralMutateCallCount(), 1);
+ Optional<TVar> mapped_lhs = mapped->lhs.as<TVar>();
+ Optional<TVar> mapped_rhs = mapped->rhs.as<TVar>();
+ ASSERT_TRUE(mapped_lhs.has_value());
+ ASSERT_TRUE(mapped_rhs.has_value());
+ EXPECT_EQ(mapped_lhs.value()->name, "lhs-mapped");
+ EXPECT_EQ(mapped_rhs.value()->name, "rhs-mapped");
+
+ Expected<Any> failed =
+ StructuralMapExpected<WalkOrder::kPostOrder>(root, [](const TVarObj*
var) -> Expected<Any> {
+ if (var->name == "lhs") {
+ return Unexpected(Error("ValueError", "registered hook child
failed", ""));
+ }
+ return Any(TVar(var->name + "-mapped"));
+ });
+ ASSERT_TRUE(failed.is_err());
+ Optional<VisitErrorContext> context =
VisitErrorContext::TryGetFromError(failed.error());
+ ASSERT_TRUE(context.has_value());
+ const List<ObjectRef>& reverse_pattern =
context.value()->reverse_visit_pattern;
+ ASSERT_EQ(reverse_pattern.size(), 2U);
+ EXPECT_TRUE(reverse_pattern[0].same_as(lhs));
+ EXPECT_TRUE(reverse_pattern[1].same_as(root));
+
+ TMutatePair nullable(ObjectRef(nullptr), rhs);
+ TMutatePair nullable_mapped =
+ StructuralMap<WalkOrder::kPostOrder>(nullable, [](const String& value)
-> Expected<Any> {
+ return Any(value);
+ }).cast<TMutatePair>();
+ EXPECT_FALSE(nullable_mapped->lhs.defined());
+ EXPECT_TRUE(nullable_mapped->rhs.same_as(rhs));
+
+ Expected<Any> wrong_type = StructuralMapExpected<WalkOrder::kPostOrder>(
+ root, [](const TVarObj*) -> Expected<Any> { return Any(int64_t{1}); });
+ ASSERT_TRUE(wrong_type.is_err());
+ EXPECT_EQ(wrong_type.error().kind(), "TypeError");
+}
+
TEST(StructuralMap, PreservesSharedArrayAndMapInputs) {
// A shared Array is copied when one of its elements changes.
{
diff --git a/tests/cpp/testing_object.h b/tests/cpp/testing_object.h
index b5a6a394..32e1e917 100644
--- a/tests/cpp/testing_object.h
+++ b/tests/cpp/testing_object.h
@@ -25,6 +25,7 @@
#include <tvm/ffi/container/map.h>
#include <tvm/ffi/device.h>
#include <tvm/ffi/dtype.h>
+#include <tvm/ffi/extra/structural_mutate.h>
#include <tvm/ffi/extra/structural_visit.h>
#include <tvm/ffi/memory.h>
#include <tvm/ffi/object.h>
@@ -234,6 +235,60 @@ class TPair : public ObjectRef {
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(TPair, ObjectRef, TPairObj);
};
+class TMutatePairObj : public Object {
+ public:
+ ObjectRef lhs;
+ ObjectRef rhs;
+
+ TMutatePairObj(ObjectRef lhs, ObjectRef rhs) : lhs(std::move(lhs)),
rhs(std::move(rhs)) {}
+ explicit TMutatePairObj(UnsafeInit) {}
+
+ // Every test that reads this process-global counter must reset it first.
+ static int& StructuralMutateCallCount() {
+ static int count = 0;
+ return count;
+ }
+
+ static Expected<Any> StructuralMutateExpected(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);
+ if (lhs.same_as(self->lhs) && rhs.same_as(self->rhs)) {
+ return 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));
+ }
+
+ static void RegisterReflection() {
+ namespace refl = tvm::ffi::reflection;
+ refl::ObjectDef<TMutatePairObj>()
+ .def_ro("lhs", &TMutatePairObj::lhs)
+ .def_ro("rhs", &TMutatePairObj::rhs);
+ refl::EnsureTypeAttrColumn(refl::type_attr::kStructuralMutate);
+ refl::TypeAttrDef<TMutatePairObj>().attr(
+ refl::type_attr::kStructuralMutate,
+
reinterpret_cast<void*>(static_cast<FStructuralMutate>(&TMutatePairObj::StructuralMutate)));
+ }
+
+ static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind =
kTVMFFISEqHashKindTreeNode;
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("test.MutatePair", TMutatePairObj, Object);
+};
+
+class TMutatePair : public ObjectRef {
+ public:
+ TMutatePair(ObjectRef lhs, ObjectRef rhs) {
+ data_ = make_object<TMutatePairObj>(std::move(lhs), std::move(rhs));
+ }
+
+ TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(TMutatePair, ObjectRef,
TMutatePairObj);
+};
+
class TObjectPtrHolderObj : public Object {
public:
Arc<TIntObj> value;
@@ -358,10 +413,14 @@ 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); }));
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(
+ visitor->WithDefRegionKind(kTVMFFIDefRegionKindRecursive,
+ [&]() { return
visitor->VisitExpected(self->params); }),
+ self);
- return
details::ExpectedUnsafe::MoveToTVMFFIAny(visitor->VisitExpected(self->body));
+ auto body_result = visitor->VisitExpected(self->body);
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(body_result, self);
+ return details::ExpectedUnsafe::MoveToTVMFFIAny(std::move(body_result));
}
static void RegisterReflection() {
