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 0e51eddf [FIX] Preserve post-order structural map descent (#866)
0e51eddf is described below
commit 0e51eddf5976ea5b031a82bacfe6a16c6e5fd1eb
Author: Tianqi Chen <[email protected]>
AuthorDate: Thu Oct 22 18:02:37 2026 -0400
[FIX] Preserve post-order structural map descent (#866)
---
include/tvm/ffi/extra/structural_mutate.h | 23 ++++++++++----
tests/cpp/extra/test_structural_mutate.cc | 53 +++++++++++++++++++++++++------
2 files changed, 61 insertions(+), 15 deletions(-)
diff --git a/include/tvm/ffi/extra/structural_mutate.h
b/include/tvm/ffi/extra/structural_mutate.h
index d768c3a0..975dd333 100644
--- a/include/tvm/ffi/extra/structural_mutate.h
+++ b/include/tvm/ffi/extra/structural_mutate.h
@@ -1350,7 +1350,7 @@ class StructuralMapEngine : public Parent {
matched = mapped_view.template as<TSub>();
if (!matched.has_value()) return false;
}
- *out = [&]() -> Expected<Any> {
+ Expected<Any> callback_result = [&]() -> Expected<Any> {
if constexpr (std::is_same_v<TSub, AnyView>) {
return InvokeTypedCallbackLink(callback, mapped_view,
StateIndices{});
} else if constexpr (std::is_same_v<TSub, Any>) {
@@ -1359,10 +1359,15 @@ class StructuralMapEngine : public Parent {
return InvokeTypedCallbackLink(callback, *std::move(matched),
StateIndices{});
}
}();
- if (TVM_FFI_PREDICT_FALSE(out->is_err())) {
- this->UpdateVisitErrorContext(*out, mapped_view);
+ if (TVM_FFI_PREDICT_FALSE(callback_result.is_err())) {
+ this->UpdateVisitErrorContext(callback_result, mapped_view);
+ *out = std::move(callback_result);
return true;
}
+ // Unchanged refers to the descended node, whose owner stays live
through the callback.
+ if (ExpectedUnsafe::GetData(callback_result).type_index() !=
TypeIndex::kTVMFFIUnchanged) {
+ *out = std::move(callback_result);
+ }
return true;
}
}
@@ -1570,11 +1575,17 @@ class StructuralMapDynEngine : public Parent {
if (!matched.has_value()) return false;
// WithDefRegionKind restores its state through RAII, so this late read
is equivalent to
// the typed engine's invocation-time read even after recursive descent.
- *out = InvokeLink(*matched, with_kind, mapped_view,
this->def_region_kind());
- if (TVM_FFI_PREDICT_FALSE(out->is_err())) {
- this->UpdateVisitErrorContext(*out, mapped_view);
+ Expected<Any> callback_result =
+ InvokeLink(*matched, with_kind, mapped_view,
this->def_region_kind());
+ if (TVM_FFI_PREDICT_FALSE(callback_result.is_err())) {
+ this->UpdateVisitErrorContext(callback_result, mapped_view);
+ *out = std::move(callback_result);
return true;
}
+ // Preserve descent and its lifetime just as in the typed engine.
+ if (ExpectedUnsafe::GetData(callback_result).type_index() !=
TypeIndex::kTVMFFIUnchanged) {
+ *out = std::move(callback_result);
+ }
return true;
}
}
diff --git a/tests/cpp/extra/test_structural_mutate.cc
b/tests/cpp/extra/test_structural_mutate.cc
index 6973eb94..02ea9f86 100644
--- a/tests/cpp/extra/test_structural_mutate.cc
+++ b/tests/cpp/extra/test_structural_mutate.cc
@@ -733,7 +733,11 @@ TEST(StructuralMap,
RootByValueProtectsSharedParentSubvalue) {
AnyArray outer{Any(std::move(child))};
const Object* child_address = outer[0].cast<AnyArray>().get();
- AnyArray mapped = StructuralMap<WalkOrder::kPostOrder>(outer[0],
Increment).cast<AnyArray>();
+ AnyArray mapped =
+ StructuralMap<WalkOrder::kPostOrder>(outer[0], Increment, [](const
AnyArray& array) {
+ EXPECT_EQ(array[0].cast<int64_t>(), 2);
+ return Unchanged();
+ }).cast<AnyArray>();
EXPECT_NE(mapped.get(), child_address);
EXPECT_EQ(outer[0].cast<AnyArray>()[0].cast<int64_t>(), 1);
@@ -1137,18 +1141,49 @@ Any CallDynStructuralMap(AnyView root, const
Array<Tuple<int32_t, Function>>& ca
return fn(root, callbacks, Array<Tuple<int32_t, Function>>(),
static_cast<int32_t>(order));
}
-TEST(StructuralMapDyn, PreOrderDescendsOriginalAfterUnchangedCallback) {
+TEST(StructuralMapDyn, UnchangedCallbackPreservesDescent) {
AnyArray root{int64_t{1}};
Function unchanged = Function::FromTyped([](const AnyArray&) -> Any { return
Unchanged(); });
Function increment = Function::FromTyped([](int64_t value) -> Any { return
Any(value + 1); });
- AnyArray mapped =
- CallDynStructuralMap(root,
- {Tuple<int32_t, Function>(TypeIndex::kTVMFFIArray,
unchanged),
- Tuple<int32_t, Function>(TypeIndex::kTVMFFIInt,
increment)},
- WalkOrder::kPreOrder)
- .cast<AnyArray>();
- EXPECT_EQ(mapped[0].cast<int64_t>(), 2);
+ for (WalkOrder order : {WalkOrder::kPreOrder, WalkOrder::kPostOrder}) {
+ AnyArray mapped =
+ CallDynStructuralMap(root,
+ {Tuple<int32_t,
Function>(TypeIndex::kTVMFFIArray, unchanged),
+ Tuple<int32_t, Function>(TypeIndex::kTVMFFIInt,
increment)},
+ order)
+ .cast<AnyArray>();
+ EXPECT_EQ(mapped[0].cast<int64_t>(), 2);
+ EXPECT_EQ(root[0].cast<int64_t>(), 1);
+ }
+}
+
+TEST(StructuralMap, PostOrderCallbackErrorKeepsDescendedNode) {
+ AnyArray root{int64_t{1}};
+ auto fail = [](const ArrayObj*) -> Any {
+ return Any(Error("ValueError", "post-order callback failed", ""));
+ };
+ auto check_error = [&](const Expected<Any>& result) {
+ ASSERT_TRUE(result.is_err());
+ EXPECT_EQ(result.error().message(), "post-order callback failed");
+ Optional<VisitErrorContext> context =
VisitErrorContext::TryGetFromError(result.error());
+ ASSERT_TRUE(context.has_value());
+ ASSERT_EQ(context.value()->reverse_visit_pattern.size(), 1U);
+ AnyArray descended =
context.value()->reverse_visit_pattern[0].as_or_throw<AnyArray>();
+ EXPECT_FALSE(descended.same_as(root));
+ EXPECT_EQ(descended[0].cast<int64_t>(), 2);
+ };
+
+ check_error(StructuralMapExpected<WalkOrder::kPostOrder>(root, Increment,
fail));
+ Function map = Function::GetGlobalRequired("ffi.StructuralMap");
+ check_error(map.CallExpected<Any>(
+ root,
+ Array<Tuple<int32_t, Function>>{
+ Tuple<int32_t, Function>(TypeIndex::kTVMFFIArray,
Function::FromTyped(fail)),
+ Tuple<int32_t, Function>(TypeIndex::kTVMFFIInt,
+ Function::FromTyped([](int64_t value) {
return value + 1; }))},
+ Array<Tuple<int32_t, Function>>(),
static_cast<int32_t>(WalkOrder::kPostOrder)));
+ EXPECT_EQ(root[0].cast<int64_t>(), 1);
}
TEST(StructuralMapDyn, InvokesCallbackForEveryFreeVarOccurrence) {
