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 3b0e30b1 [IR] Expose protected remap state and skip empty lookups
(#779)
3b0e30b1 is described below
commit 3b0e30b1f425c554168ab5cd91b497a0b5ecaaf4
Author: Tianqi Chen <[email protected]>
AuthorDate: Sat Sep 12 09:26:52 2026 -0400
[IR] Expose protected remap state and skip empty lookups (#779)
Expose `var_remap_` as protected state so derived structural map engines
can inspect the remap environment. Skip the identity-map lookup when the
environment is empty, after validating the key.
---
include/tvm/ffi/extra/structural_mutate.h | 5 ++---
tests/cpp/extra/test_structural_mutate.cc | 5 +++++
2 files changed, 7 insertions(+), 3 deletions(-)
diff --git a/include/tvm/ffi/extra/structural_mutate.h
b/include/tvm/ffi/extra/structural_mutate.h
index 127e078f..358866c2 100644
--- a/include/tvm/ffi/extra/structural_mutate.h
+++ b/include/tvm/ffi/extra/structural_mutate.h
@@ -1031,7 +1031,6 @@ class StructuralMapEngineBase : public
StructuralMutatorObj {
/*! \brief Return the empty state tuple exposed to typed map callbacks. */
TVM_FFI_INLINE StateTupleType StateTuple() const noexcept { return {}; }
- private:
/// \cond Doxygen_Suppress
// Out of line so its strings stay out of the per-node dispatch function,
which TryLink inlines
// into. Shared by the typed and dynamic engines below.
@@ -1069,7 +1068,6 @@ class StructuralMapEngineBase : public
StructuralMutatorObj {
return details::ExpectedUnsafe::MoveToTVMFFIAny(self->VarRemapSetImpl(var,
mapped_value));
}
- protected:
/*!
* \brief Append \p node to a failed result's mutate error context.
* \param result The failed result whose Error is annotated.
@@ -1094,6 +1092,7 @@ class StructuralMapEngineBase : public
StructuralMutatorObj {
if (TVM_FFI_PREDICT_FALSE(var.type_index() <
TypeIndex::kTVMFFIStaticObjectBegin)) {
return VarRemapKeyTypeError();
}
+ if (var_remap_.empty()) return Any(nullptr);
const Object* var_ptr =
details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
Object>(var);
auto it = var_remap_.find(var_ptr);
@@ -1123,7 +1122,6 @@ class StructuralMapEngineBase : public
StructuralMutatorObj {
return Expected<void>();
}
- private:
template <typename Parent, WalkOrder order, typename... Callbacks>
friend class StructuralMapEngine;
template <typename Parent, WalkOrder order>
@@ -1133,6 +1131,7 @@ class StructuralMapEngineBase : public
StructuralMutatorObj {
template <typename Parent>
friend class details::StructuralMutateDynEngine;
+ /*! \brief Identity-substitution environment keyed by object identity. */
// Raw-pointer key: IncRef once on first insert, DecRef all keys in the
destructor.
std::unordered_map<const Object*, Any> var_remap_;
};
diff --git a/tests/cpp/extra/test_structural_mutate.cc
b/tests/cpp/extra/test_structural_mutate.cc
index 912efb28..ad9724d4 100644
--- a/tests/cpp/extra/test_structural_mutate.cc
+++ b/tests/cpp/extra/test_structural_mutate.cc
@@ -373,6 +373,11 @@ TEST(StructuralMap,
ParentLayerOwnsBothDescentsAndProvidesState) {
decltype(identity), decltype(map_var)>;
auto engine = make_object<Mutator>(std::move(identity), std::move(map_var));
StructuralMutator mutator(engine);
+ TVar key("key");
+ EXPECT_EQ(mutator->VarRemapGetExpected(key).value(), nullptr);
+ EXPECT_EQ(mutator->VarRemapGetExpected(nullptr).error().kind(), "TypeError");
+ EXPECT_EQ(mutator->VarRemapGetExpected(1).error().kind(), "TypeError");
+ mutator->VarRemapSetExpected(key, Any(Unchanged())).value();
ASSERT_FALSE(mutator->MutateExpected(String("unmatched")).is_err());
AnyArray rebuild_root{int64_t{1}};
