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 70e5af34 [PERF] Forward single StructuralMutate callbacks directly
(#774)
70e5af34 is described below
commit 70e5af34b162096340695c583b11c197f3a89595
Author: Tianqi Chen <[email protected]>
AuthorDate: Thu Sep 10 17:37:36 2026 -0400
[PERF] Forward single StructuralMutate callbacks directly (#774)
Forward single StructuralMutate callback results directly and use
bool-match/result-output dispatch for multiple callbacks. Preserve
owning conversions, typed fallback, first-match behavior, and cold
error-context attachment.
---
include/tvm/ffi/extra/structural_mutate.h | 138 ++++++++++++++++++++++--------
tests/cpp/extra/test_structural_mutate.cc | 14 +++
2 files changed, 115 insertions(+), 37 deletions(-)
diff --git a/include/tvm/ffi/extra/structural_mutate.h
b/include/tvm/ffi/extra/structural_mutate.h
index cf980088..a50ea694 100644
--- a/include/tvm/ffi/extra/structural_mutate.h
+++ b/include/tvm/ffi/extra/structural_mutate.h
@@ -579,11 +579,18 @@ class StructuralMutatorObj : public Object {
result = DefaultMutateRawTail(value, attr);
}
if (TVM_FFI_PREDICT_FALSE(result.type_index == TypeIndex::kTVMFFIError)) {
- details::UpdateVisitErrorContext(result, value);
+ return AttachVisitErrorContextRaw(result, value);
}
return result;
}
+ /*! \brief Keep address-taking for error decoration off the successful
raw-result path. */
+ TVM_FFI_COLD_CODE static TVMFFIAny AttachVisitErrorContextRaw(TVMFFIAny
result,
+ AnyView value)
noexcept {
+ details::UpdateVisitErrorContext(result, value);
+ return result;
+ }
+
/*! \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) {
@@ -654,7 +661,7 @@ class StructuralMutatorObj : public Object {
// 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 AttachVisitErrorContextRaw(result, value);
}
return result;
}
@@ -670,7 +677,7 @@ class StructuralMutatorObj : public Object {
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 AttachVisitErrorContextRaw(result, value);
}
return result;
}
@@ -1580,19 +1587,29 @@ class StructuralMutateEngine : public Parent {
/*! \brief Dispatch ordinary mutation from the erased mutator pointer. */
static TVMFFIAny DispatchMutate(StructuralMutatorObj* mutator, AnyView
value) noexcept {
- return static_cast<StructuralMutateEngine*>(mutator)->MutateImplRaw(value);
+ auto* self = static_cast<StructuralMutateEngine*>(mutator);
+ if constexpr (sizeof...(Callbacks) == 1) {
+ return self->template MutateSingleCallbackRaw<false>(value);
+ } else {
+ return self->MutateImplRaw(value);
+ }
}
/*! \brief Dispatch maybe-in-place mutation from the erased mutator pointer.
*/
static TVMFFIAny DispatchMaybeInplaceMutate(StructuralMutatorObj* mutator,
AnyView value) noexcept {
- return
static_cast<StructuralMutateEngine*>(mutator)->MaybeInplaceMutateImplRaw(value);
+ auto* self = static_cast<StructuralMutateEngine*>(mutator);
+ if constexpr (sizeof...(Callbacks) == 1) {
+ return self->template MutateSingleCallbackRaw<true>(value);
+ } else {
+ return self->MaybeInplaceMutateImplRaw(value);
+ }
}
/*! \brief Mutate one value, handing a matched callback ownership of
descent. */
TVMFFIAny MutateImplRaw(AnyView value) noexcept {
- if (std::optional<Expected<Any>> matched = DispatchCallbacks(value,
false)) {
- Expected<Any> result = *std::move(matched);
+ Expected<Any> result{Any()};
+ if (DispatchCallbacks(value, false, &result)) {
if (TVM_FFI_PREDICT_FALSE(result.is_err())) {
// Keep callback-boundary context in addition to the default-descent
// context: a callback may return a rebuilt value, so the two nodes
can differ.
@@ -1605,8 +1622,8 @@ class StructuralMutateEngine : public Parent {
/*! \brief Maybe mutate one value in place, with callback-owned descent. */
TVMFFIAny MaybeInplaceMutateImplRaw(AnyView value) noexcept {
- if (std::optional<Expected<Any>> matched = DispatchCallbacks(value, true))
{
- Expected<Any> result = *std::move(matched);
+ Expected<Any> result{Any()};
+ if (DispatchCallbacks(value, true, &result)) {
if (TVM_FFI_PREDICT_FALSE(result.is_err())) {
// Keep callback-boundary context in addition to the default-descent
// context: a callback may return a rebuilt value, so the two nodes
can differ.
@@ -1618,16 +1635,14 @@ class StructuralMutateEngine : public Parent {
Parent::DefaultMaybeInplaceMutateExpected(value));
}
- /*! \brief Try one typed callback and preserve Error as an expected result.
*/
- template <typename Callback>
- TVM_FFI_INLINE std::optional<Expected<Any>> TryLink(Callback& callback,
AnyView value,
- bool allow_inplace)
noexcept {
+ /*! \brief Invoke a matched callback with the Parent view and preserve
returned/thrown Error. */
+ template <typename Callback, typename Matched>
+ TVM_FFI_INLINE Expected<Any> InvokeCallback(Callback& callback, Matched&&
matched,
+ bool allow_inplace) noexcept {
using FuncInfo = details::FunctionInfo<std::decay_t<Callback>>;
static_assert(FuncInfo::num_args == 2 || FuncInfo::num_args == 3,
"StructuralMutate callback must take (value, mutator) or "
"(value, mutator, allow_inplace)");
- using FirstArg = std::tuple_element_t<0, typename FuncInfo::ArgType>;
- using TSub = std::remove_cv_t<std::remove_reference_t<FirstArg>>;
using SecondArg = std::decay_t<std::tuple_element_t<1, typename
FuncInfo::ArgType>>;
using Second = std::remove_pointer_t<SecondArg>;
static_assert(std::is_same_v<Second, typename Parent::MutatorObjType>,
@@ -1639,39 +1654,88 @@ class StructuralMutateEngine : public Parent {
"third StructuralMutate callback argument must be bool");
}
auto* mutator = static_cast<typename Parent::MutatorObjType*>(this);
- auto invoke = [&](auto&& matched) -> Expected<Any> {
- try {
- if constexpr (FuncInfo::num_args == 3) {
- return callback(std::forward<decltype(matched)>(matched), mutator,
allow_inplace);
- } else {
- return callback(std::forward<decltype(matched)>(matched), mutator);
- }
- } catch (Error& err) {
- return Unexpected(std::move(err));
+ try {
+ if constexpr (FuncInfo::num_args == 3) {
+ return callback(std::forward<Matched>(matched), mutator,
allow_inplace);
+ } else {
+ return callback(std::forward<Matched>(matched), mutator);
}
- };
+ } catch (Error& err) {
+ return Unexpected(std::move(err));
+ }
+ }
+
+ /*! \brief Materialize the callback result only on the cold Parent
error-decoration path. */
+ TVM_FFI_COLD_CODE TVMFFIAny AnnotateCallbackErrorRaw(TVMFFIAny raw, AnyView
value) noexcept {
+ Expected<Any> result =
details::ExpectedUnsafe::MoveFromTVMFFIAny<Any>(raw);
+ Parent::UpdateVisitErrorContext(result, value);
+ return details::ExpectedUnsafe::MoveToTVMFFIAny(std::move(result));
+ }
+
+ // One callback forwards its result directly, without a callback-chain
envelope.
+ // AnyView/Any always match; only a typed miss needs the Parent's default
descent.
+ template <bool kMaybeInplace>
+ TVM_FFI_INLINE TVMFFIAny MutateSingleCallbackRaw(AnyView value) noexcept {
+ auto& callback = std::get<0>(callbacks_);
+ using FuncInfo = details::FunctionInfo<std::decay_t<decltype(callback)>>;
+ using FirstArg = std::tuple_element_t<0, typename FuncInfo::ArgType>;
+ using TSub = std::remove_cv_t<std::remove_reference_t<FirstArg>>;
+ TVMFFIAny result;
if constexpr (std::is_same_v<TSub, AnyView>) {
- return invoke(value);
+ result =
+ details::ExpectedUnsafe::MoveToTVMFFIAny(InvokeCallback(callback,
value, kMaybeInplace));
} else if constexpr (std::is_same_v<TSub, Any>) {
- return invoke(Any(value));
+ result = details::ExpectedUnsafe::MoveToTVMFFIAny(
+ InvokeCallback(callback, Any(value), kMaybeInplace));
} else if (auto matched = value.template as<TSub>()) {
- return invoke(*std::move(matched));
+ result = details::ExpectedUnsafe::MoveToTVMFFIAny(
+ InvokeCallback(callback, *std::move(matched), kMaybeInplace));
+ } else {
+ if constexpr (kMaybeInplace) {
+ return details::ExpectedUnsafe::MoveToTVMFFIAny(
+ Parent::DefaultMaybeInplaceMutateExpected(value));
+ } else {
+ return
details::ExpectedUnsafe::MoveToTVMFFIAny(Parent::DefaultMutateExpected(value));
+ }
}
- return std::nullopt;
+ // Release any owning match before naming the callback boundary, as in the
general path.
+ if (TVM_FFI_PREDICT_FALSE(result.type_index == TypeIndex::kTVMFFIError)) {
+ return AnnotateCallbackErrorRaw(result, value);
+ }
+ return result;
+ }
+
+ /*! \brief Write a matched callback result, leaving out untouched on a type
miss. */
+ template <typename Callback>
+ TVM_FFI_INLINE bool TryLink(Callback& callback, AnyView value, bool
allow_inplace,
+ Expected<Any>* out) noexcept {
+ using FuncInfo = details::FunctionInfo<std::decay_t<Callback>>;
+ 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>) {
+ *out = InvokeCallback(callback, value, allow_inplace);
+ return true;
+ } else if constexpr (std::is_same_v<TSub, Any>) {
+ *out = InvokeCallback(callback, Any(value), allow_inplace);
+ return true;
+ } else if (auto matched = value.template as<TSub>()) {
+ *out = InvokeCallback(callback, *std::move(matched), allow_inplace);
+ return true;
+ }
+ return false;
}
/*! \brief Fold callbacks in declaration order, stopping at the first match.
*/
template <size_t... Is>
- TVM_FFI_INLINE std::optional<Expected<Any>> TryLinks(AnyView value, bool
allow_inplace,
-
std::index_sequence<Is...>) noexcept {
- std::optional<Expected<Any>> result;
- (... || (result = TryLink(std::get<Is>(callbacks_), value,
allow_inplace)).has_value());
- return result;
+ TVM_FFI_INLINE bool TryLinks(AnyView value, bool allow_inplace,
Expected<Any>* out,
+ std::index_sequence<Is...>) noexcept {
+ return (TryLink(std::get<Is>(callbacks_), value, allow_inplace, out) ||
...);
}
- /*! \brief Run the callback chain, or return empty when no callback matched.
*/
- std::optional<Expected<Any>> DispatchCallbacks(AnyView value, bool
allow_inplace) noexcept {
- return TryLinks(value, allow_inplace,
std::index_sequence_for<Callbacks...>{});
+ /*! \brief Run the callback chain, returning whether a callback matched. */
+ TVM_FFI_INLINE bool DispatchCallbacks(AnyView value, bool allow_inplace,
+ Expected<Any>* out) noexcept {
+ return TryLinks(value, allow_inplace, out,
std::index_sequence_for<Callbacks...>{});
}
/*! \brief Typed callbacks tested in declaration order, first match wins. */
diff --git a/tests/cpp/extra/test_structural_mutate.cc
b/tests/cpp/extra/test_structural_mutate.cc
index f4c80a72..dd254c80 100644
--- a/tests/cpp/extra/test_structural_mutate.cc
+++ b/tests/cpp/extra/test_structural_mutate.cc
@@ -432,6 +432,20 @@ TEST(StructuralMutate, CallbackControlsRecursion) {
EXPECT_TRUE(mapped->rhs.same_as(original_rhs));
}
+TEST(StructuralMutate, SingleCallbackCanDelegateToDefault) {
+ auto mutate = [](AnyView value, StructuralMutatorObj* mutator) ->
Expected<UnchangedOr<Any>> {
+ if (auto integer = value.as<int64_t>()) {
+ return Any(*integer + 1);
+ }
+ return mutator->DefaultMutateExpected(value);
+ };
+ AnyArray root{int64_t{1}, AnyArray{int64_t{2}}};
+ AnyArray result = StructuralMutate(root, mutate).cast<AnyArray>();
+ EXPECT_EQ(result[0].cast<int64_t>(), 2);
+ EXPECT_EQ(result[1].cast<AnyArray>()[0].cast<int64_t>(), 3);
+ EXPECT_EQ(root[0].cast<int64_t>(), 1);
+}
+
TEST(StructuralMutate, PreservesUniqueContainerIdentity) {
AnyArray inner{int64_t{1}};
const Object* inner_address = inner.get();
