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 897ece64 [PERF][EXTRA] Split lazy structural mutation loops (#756)
897ece64 is described below
commit 897ece64d6ad0857f803e68221375021867e81a5
Author: Tianqi Chen <[email protected]>
AuthorDate: Sun Sep 6 14:38:26 2026 -0400
[PERF][EXTRA] Split lazy structural mutation loops (#756)
The ordinary sequence and map structural mutation paths currently keep
their lazy-output null check inside the per-element loop, even after the
first changed child has allocated the output. Split each path into an
unchanged-prefix scan and a post-allocation suffix loop so the suffix
does strictly less bookkeeping.
The no-change return, prefix copy, stable map iterator alignment,
copy-on-write behavior, and error propagation remain unchanged. A
focused regression covers retained arrays and maps with unchanged and
changed suffix elements, plus an error after allocation.
---
src/ffi/extra/structural_mutate.cc | 105 ++++++++++++++++++++----------
tests/cpp/extra/test_structural_mutate.cc | 37 +++++++++++
2 files changed, 106 insertions(+), 36 deletions(-)
diff --git a/src/ffi/extra/structural_mutate.cc
b/src/ffi/extra/structural_mutate.cc
index 71511a94..6060fb5e 100644
--- a/src/ffi/extra/structural_mutate.cc
+++ b/src/ffi/extra/structural_mutate.cc
@@ -149,6 +149,34 @@ Expected<Any> StructuralMutateExpected(
// Built-in container structural mutation.
// ---------------------------------------------------------------------------
+/*!
+ * \brief Finish mutating a sequence after its first changed element.
+ *
+ * \tparam SeqObj The underlying sequence object type.
+ * \param mutator The active structural mutator.
+ * \param self The source sequence object.
+ * \param index The index of the first changed element.
+ * \param first The mapped value for the first changed element.
+ * \return The mutated sequence, or an Error.
+ */
+template <typename SeqObj>
+TVM_FFI_INLINE TVMFFIAny MutateSeqContainerChanged(StructuralMutatorObj*
mutator,
+ const SeqObj* self, int64_t
index,
+ Any first) noexcept {
+ int64_t size = static_cast<int64_t>(self->size());
+ const Any* items = self->begin();
+ ObjectPtr<SeqObj> output = SeqObj::CreateRepeated(size, Any());
+ output->InitRange(0, items, items + index);
+ output->SetItemAfterCheck(index, std::move(first));
+
+ for (int64_t i = index + 1; i < size; ++i) {
+ const Any& item = items[i];
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(Any, mapped_value,
mutator->MutateExpected(item));
+ output->SetItemAfterCheck(i, std::move(mapped_value));
+ }
+ return AnyUnsafe::MoveAnyToTVMFFIAny(Any(std::move(output)));
+}
+
/*!
* \brief Structurally mutate the elements of a sequence container.
*
@@ -163,26 +191,15 @@ 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));
-
- if (output == nullptr) {
- if (item.same_as(mapped_value)) {
- continue;
- }
- output = SeqObj::CreateRepeated(size, Any());
- output->InitRange(0, items, items + i);
+ if (!item.same_as(mapped_value)) {
+ return MutateSeqContainerChanged(mutator, self, i,
std::move(mapped_value));
}
- output->SetItemAfterCheck(i, std::move(mapped_value));
}
-
- if (output == nullptr) {
- return AnyUnsafe::MoveAnyToTVMFFIAny(Any(value));
- }
- return AnyUnsafe::MoveAnyToTVMFFIAny(Any(std::move(output)));
+ return AnyUnsafe::MoveAnyToTVMFFIAny(Any(value));
}
/*!
@@ -209,6 +226,41 @@ TVMFFIAny
MaybeInplaceMutateSeqContainerRaw(StructuralMutatorObj* mutator, AnyVi
return AnyUnsafe::MoveAnyToTVMFFIAny(Any(value));
}
+/*!
+ * \brief Finish mutating map values after the first changed value.
+ *
+ * \tparam MapObjType The underlying map object type.
+ * \param mutator The active structural mutator.
+ * \param self The source map object.
+ * \param source_it Iterator at the first changed value.
+ * \param index Iteration index of the first changed value.
+ * \param first The mapped value for the first changed entry.
+ * \return The mutated map, or an Error.
+ */
+template <typename MapObjType>
+TVM_FFI_INLINE TVMFFIAny MutateMapValuesChanged(StructuralMutatorObj* mutator,
+ const MapObjType* self,
+ MapBaseObj::iterator
source_it, size_t index,
+ Any first) noexcept {
+ ObjectPtr<Object> output = MapObjType::ShallowCopy(self);
+ auto output_it = static_cast<MapBaseObj*>(output.get())->begin();
+ for (size_t i = 0; i < index; ++i) {
+ ++output_it;
+ }
+ output_it->second = std::move(first);
+ ++source_it;
+ ++output_it;
+
+ for (; source_it != self->end(); ++source_it, ++output_it) {
+ const Any& old_value = source_it->second;
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(Any, new_value,
mutator->MutateExpected(old_value));
+ if (!old_value.same_as(new_value)) {
+ output_it->second = std::move(new_value);
+ }
+ }
+ return AnyUnsafe::MoveAnyToTVMFFIAny(Any(std::move(output)));
+}
+
/*!
* \brief Structurally mutate the values of a map container.
*
@@ -221,34 +273,15 @@ TVMFFIAny
MaybeInplaceMutateSeqContainerRaw(StructuralMutatorObj* mutator, AnyVi
template <typename MapObjType>
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));
- bool changed = !old_value.same_as(new_value);
- if (output == nullptr) {
- if (!changed) {
- continue;
- }
- 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 = std::move(new_value);
+ if (!old_value.same_as(new_value)) {
+ return MutateMapValuesChanged(mutator, self, source_it, index,
std::move(new_value));
}
- ++output_it;
- }
-
- if (output == nullptr) {
- return AnyUnsafe::MoveAnyToTVMFFIAny(Any(value));
}
- return AnyUnsafe::MoveAnyToTVMFFIAny(Any(std::move(output)));
+ return AnyUnsafe::MoveAnyToTVMFFIAny(Any(value));
}
/*!
diff --git a/tests/cpp/extra/test_structural_mutate.cc
b/tests/cpp/extra/test_structural_mutate.cc
index 0170c45d..b99b0a3c 100644
--- a/tests/cpp/extra/test_structural_mutate.cc
+++ b/tests/cpp/extra/test_structural_mutate.cc
@@ -560,6 +560,43 @@ TEST(StructuralMap, PreservesSharedArrayAndMapInputs) {
}
}
+TEST(StructuralMap, CopyOnWriteProcessesSuffixAfterFirstChange) {
+ AnyArray array_root{String("prefix"), int64_t{1}, String("middle"),
int64_t{2}};
+ AnyArray array_owner = array_root; //
NOLINT(performance-unnecessary-copy-initialization)
+ AnyArray mapped_array =
+ StructuralMap<WalkOrder::kPostOrder>(array_root,
Increment).cast<AnyArray>();
+
+ EXPECT_TRUE(mapped_array[0].same_as(array_root[0]));
+ EXPECT_EQ(mapped_array[1].cast<int64_t>(), 2);
+ EXPECT_TRUE(mapped_array[2].same_as(array_root[2]));
+ EXPECT_EQ(mapped_array[3].cast<int64_t>(), 3);
+ EXPECT_TRUE(array_owner.same_as(array_root));
+
+ StringMap map_root{{"prefix", String("unchanged")},
+ {"first", int64_t{1}},
+ {"middle", String("also-unchanged")},
+ {"second", int64_t{2}}};
+ StringMap map_owner = map_root; //
NOLINT(performance-unnecessary-copy-initialization)
+ StringMap mapped_map =
+ StructuralMap<WalkOrder::kPostOrder>(map_root,
Increment).cast<StringMap>();
+
+ EXPECT_TRUE(mapped_map["prefix"].same_as(map_root["prefix"]));
+ EXPECT_EQ(mapped_map["first"].cast<int64_t>(), 2);
+ EXPECT_TRUE(mapped_map["middle"].same_as(map_root["middle"]));
+ EXPECT_EQ(mapped_map["second"].cast<int64_t>(), 3);
+ EXPECT_TRUE(map_owner.same_as(map_root));
+
+ Expected<Any> suffix_error = StructuralMapExpected<WalkOrder::kPostOrder>(
+ AnyArray{int64_t{1}, int64_t{2}}, [](int64_t value) -> Expected<Any> {
+ if (value == 2) {
+ return Unexpected(Error("ValueError", "suffix mutate failed", ""));
+ }
+ return Any(value + 1);
+ });
+ ASSERT_TRUE(suffix_error.is_err());
+ EXPECT_EQ(suffix_error.error().message(), "suffix mutate failed");
+}
+
TEST(StructuralMap, PreOrderRecursivelyMapsCallbackResult) {
StringMap root{{"value", AnyArray{int64_t{1}}}};
AnyArray replacement{int64_t{10}};
