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 88d9f706 [FEAT] Add typed reflection attribute keys (#846)
88d9f706 is described below
commit 88d9f7068717a8dc0df7d2ad06e1567dfc24472f
Author: Tianqi Chen <[email protected]>
AuthorDate: Tue Oct 13 22:37:30 2026 -0400
[FEAT] Add typed reflection attribute keys (#846)
Introduce generic `reflection::AttrKey<T>` descriptors with strict
optional extraction and typed value registration alongside dynamic
attribute names.
Co-locate canonical `tvm::ffi::type_attr` contracts with their owning
APIs while preserving erased columns and name-based packed registration.
Add `native_fn_attr<&Fn>(key)` and enum `set_native_fn_attr<&Fn>(key)`
for explicit key-directed native callbacks.
---
docs/concepts/structural_eq_hash.rst | 11 +-
docs/guides/dataclass_reflection.rst | 8 +-
include/tvm/ffi/enum.h | 11 +-
include/tvm/ffi/extra/dataclass.h | 47 +++++
include/tvm/ffi/extra/serialization.h | 23 +++
include/tvm/ffi/extra/structural_equal.h | 20 +++
include/tvm/ffi/extra/structural_hash.h | 19 ++
include/tvm/ffi/extra/structural_mutate.h | 52 +++++-
include/tvm/ffi/extra/structural_visit.h | 23 ++-
include/tvm/ffi/reflection/accessor.h | 282 ++++++++----------------------
include/tvm/ffi/reflection/creator.h | 8 +-
include/tvm/ffi/reflection/enum_def.h | 14 +-
include/tvm/ffi/reflection/overload.h | 3 +-
include/tvm/ffi/reflection/registry.h | 45 +++--
src/ffi/container.cc | 66 ++++---
src/ffi/extra/dataclass.cc | 77 ++++----
src/ffi/extra/reflection_extra.cc | 15 +-
src/ffi/extra/serialization.cc | 16 +-
src/ffi/extra/structural_equal.cc | 13 +-
src/ffi/extra/structural_hash.cc | 15 +-
src/ffi/extra/structural_mutate.cc | 46 ++---
src/ffi/extra/structural_visit.cc | 22 +--
src/ffi/object.cc | 2 +-
src/ffi/testing/testing.cc | 27 +--
tests/cpp/extra/test_serialization.cc | 13 +-
tests/cpp/extra/test_structural_mutate.cc | 11 +-
tests/cpp/test_reflection.cc | 25 ++-
tests/cpp/testing_object.h | 45 ++---
28 files changed, 523 insertions(+), 436 deletions(-)
diff --git a/docs/concepts/structural_eq_hash.rst
b/docs/concepts/structural_eq_hash.rst
index 51840d5d..1e869b1e 100644
--- a/docs/concepts/structural_eq_hash.rst
+++ b/docs/concepts/structural_eq_hash.rst
@@ -807,17 +807,18 @@ Example (C++)
.def_ro("params", &MyNodeObj::params)
.def_ro("body", &MyNodeObj::body);
refl::TypeAttrDef<MyNodeObj>()
- .def(refl::type_attr::kSEqual, &MyNodeObj::SEqual)
- .def(refl::type_attr::kSHash, &MyNodeObj::SHash);
+ .def(type_attr::kSEqual.name, &MyNodeObj::SEqual)
+ .def(type_attr::kSHash.name, &MyNodeObj::SHash);
}
static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind =
kTVMFFISEqHashKindTreeNode;
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("my.Node", MyNodeObj, Object);
};
-See :cpp:var:`tvm::ffi::reflection::type_attr::kSEqual` and
-:cpp:var:`tvm::ffi::reflection::type_attr::kSHash` in
-``include/tvm/ffi/reflection/accessor.h`` for the full reference.
+See :cpp:var:`tvm::ffi::type_attr::kSEqual` and
+:cpp:var:`tvm::ffi::type_attr::kSHash` in
+``include/tvm/ffi/extra/structural_equal.h`` and
+``include/tvm/ffi/extra/structural_hash.h`` for the full reference.
All Kinds at a Glance
diff --git a/docs/guides/dataclass_reflection.rst
b/docs/guides/dataclass_reflection.rst
index 803ada55..3e27ed6e 100644
--- a/docs/guides/dataclass_reflection.rst
+++ b/docs/guides/dataclass_reflection.rst
@@ -320,7 +320,7 @@ type-level attributes via
:cpp:class:`~tvm::ffi::reflection::TypeAttrDef`:
// Custom hash: only hash the "key" field
refl::TypeAttrDef<MyObj>().def(
- refl::type_attr::kHash,
+ type_attr::kHash.name,
[](const Object* self, const Function& fn_hash) -> int64_t {
auto* obj = static_cast<const MyObj*>(self);
return fn_hash(AnyView(obj->key)).cast<int64_t>();
@@ -328,7 +328,7 @@ type-level attributes via
:cpp:class:`~tvm::ffi::reflection::TypeAttrDef`:
// Custom equality: only compare "key" fields
refl::TypeAttrDef<MyObj>().def(
- refl::type_attr::kEq,
+ type_attr::kEq.name,
[](const Object* lhs, const Object* rhs, const Function& fn_eq) -> bool
{
auto* a = static_cast<const MyObj*>(lhs);
auto* b = static_cast<const MyObj*>(rhs);
@@ -337,7 +337,7 @@ type-level attributes via
:cpp:class:`~tvm::ffi::reflection::TypeAttrDef`:
// Custom three-way comparison
refl::TypeAttrDef<MyObj>().def(
- refl::type_attr::kCompare,
+ type_attr::kCompare.name,
[](const Object* lhs, const Object* rhs, const Function& fn_cmp) ->
int32_t {
auto* a = static_cast<const MyObj*>(lhs);
auto* b = static_cast<const MyObj*>(rhs);
@@ -346,7 +346,7 @@ type-level attributes via
:cpp:class:`~tvm::ffi::reflection::TypeAttrDef`:
// Custom repr
refl::TypeAttrDef<MyObj>().def(
- refl::type_attr::kRepr,
+ type_attr::kRepr.name,
[](const Object* self, const Function& fn_repr) -> String {
auto* obj = static_cast<const MyObj*>(self);
return "MyObj(key=" + fn_repr(AnyView(obj->key)).cast<String>() + ")";
diff --git a/include/tvm/ffi/enum.h b/include/tvm/ffi/enum.h
index 97d12442..6661e949 100644
--- a/include/tvm/ffi/enum.h
+++ b/include/tvm/ffi/enum.h
@@ -184,11 +184,16 @@ class StrEnum : public Enum {
/// \endcond
};
+namespace type_attr {
+/*! \brief Per-class enum state: ordered entries, canonical indices, and
extensible attrs. */
+inline constexpr reflection::AttrKey<EnumState> kEnumState{"__ffi_enum__"};
+} // namespace type_attr
+
template <typename Index>
inline Enum EnumObj::GetByIndex(int32_t type_index, const Index& index) {
- static reflection::TypeAttrColumn
state_column(reflection::type_attr::kEnumState);
- if (AnyView value = state_column[type_index]; value != nullptr) {
- EnumState state = value.cast<EnumState>();
+ static reflection::TypeAttrColumn state_column(type_attr::kEnumState);
+ if (auto value = type_attr::kEnumState.ValueFrom(state_column[type_index])) {
+ const EnumState& state = *value;
if (auto entry = state->indexes.Get(Any(index))) return
entry->as_or_throw<Enum>();
}
TVM_FFI_THROW(ValueError) << "Enum `" << TypeIndexToTypeKey(type_index)
diff --git a/include/tvm/ffi/extra/dataclass.h
b/include/tvm/ffi/extra/dataclass.h
index 6c849725..c2409a97 100644
--- a/include/tvm/ffi/extra/dataclass.h
+++ b/include/tvm/ffi/extra/dataclass.h
@@ -26,7 +26,9 @@
#include <tvm/ffi/any.h>
#include <tvm/ffi/container/variant.h>
#include <tvm/ffi/extra/base.h>
+#include <tvm/ffi/function.h>
#include <tvm/ffi/optional.h>
+#include <tvm/ffi/reflection/accessor.h>
#include <tvm/ffi/string.h>
#include <ostream>
@@ -136,6 +138,51 @@ inline std::ostream& operator<<(std::ostream& os, const
Optional<T>& value) {
return os << ReprPrint(Any(value));
}
+namespace type_attr {
+/*!
+ * \brief Custom recursive repr hook.
+ *
+ * If registered, ``RecursiveRepr`` (Python ``__repr__``) calls this instead
+ * of the default field-by-field formatting. The hook receives a callback
+ * ``fn_repr`` to recursively format nested values, avoiding infinite loops.
+ */
+inline constexpr reflection::AttrKey<
+ TypedFunction<String(const Object*, const
TypedFunction<String(AnyView)>&)>>
+ kRepr{"__ffi_repr__"};
+/*!
+ * \brief Custom recursive hash hook.
+ *
+ * If registered, ``RecursiveHash`` (Python ``hash()``) calls this instead
+ * of the default field-by-field hashing. The hook receives a callback
+ * ``fn_hash`` to recursively hash nested values, with cycle detection.
+ */
+inline constexpr reflection::AttrKey<
+ TypedFunction<int64_t(const Object*, const
TypedFunction<int64_t(AnyView)>&)>>
+ kHash{"__ffi_hash__"};
+/*!
+ * \brief Custom recursive equality hook.
+ *
+ * If registered, ``RecursiveEq`` (Python ``==``) calls this instead of the
+ * default field-by-field comparison. Falls back to ``kCompare`` if only the
+ * compare hook is present. The hook receives a callback ``fn_eq`` to
+ * recursively compare nested values.
+ */
+inline constexpr reflection::AttrKey<
+ TypedFunction<bool(const Object*, const Object*, const
TypedFunction<bool(AnyView, AnyView)>&)>>
+ kEq{"__ffi_eq__"};
+/*!
+ * \brief Custom recursive three-way comparison hook.
+ *
+ * If registered, ``RecursiveCompare`` (Python ``<``, ``<=``, ``>``, ``>=``)
+ * calls this. Also used as fallback for ``RecursiveEq`` when ``kEq`` is not
+ * registered. The hook receives a callback ``fn_cmp`` to recursively
+ * compare nested values, returning < 0, == 0, or > 0.
+ */
+inline constexpr reflection::AttrKey<TypedFunction<int32_t(
+ const Object*, const Object*, const TypedFunction<int32_t(AnyView,
AnyView)>&)>>
+ kCompare{"__ffi_compare__"};
+} // namespace type_attr
+
} // namespace ffi
} // namespace tvm
#endif // TVM_FFI_EXTRA_DATACLASS_H_
diff --git a/include/tvm/ffi/extra/serialization.h
b/include/tvm/ffi/extra/serialization.h
index b5aa2891..f8baa6cb 100644
--- a/include/tvm/ffi/extra/serialization.h
+++ b/include/tvm/ffi/extra/serialization.h
@@ -25,6 +25,8 @@
#include <tvm/ffi/extra/base.h>
#include <tvm/ffi/extra/json.h>
+#include <tvm/ffi/function.h>
+#include <tvm/ffi/reflection/accessor.h>
namespace tvm {
namespace ffi {
@@ -67,6 +69,27 @@ TVM_FFI_EXTRA_CXX_API json::Value ToJSONGraph(const Any&
value, const Any& metad
*/
TVM_FFI_EXTRA_CXX_API Any FromJSONGraph(const json::Value& value);
+namespace type_attr {
+/*!
+ * \brief Serialize object data to a JSON-compatible value.
+ *
+ * If registered, ``ToJSONGraph`` calls this instead of the default
+ * field-by-field serialization. Allows types with non-trivial internal
+ * state (e.g. ``TInt`` storing a plain ``int64_t``) to define a compact
+ * custom JSON representation.
+ */
+inline constexpr reflection::AttrKey<TypedFunction<Any(AnyView)>>
kDataToJson{"__data_to_json__"};
+/*!
+ * \brief Deserialize object data from a JSON-compatible value.
+ *
+ * Counterpart to ``kDataToJson``. If registered, ``FromJSONGraph`` calls
+ * this to reconstruct the object from its serialized representation
+ * instead of using field-by-field deserialization.
+ */
+inline constexpr reflection::AttrKey<TypedFunction<Any(AnyView)>>
kDataFromJson{
+ "__data_from_json__"};
+} // namespace type_attr
+
} // namespace ffi
} // namespace tvm
#endif // TVM_FFI_EXTRA_SERIALIZATION_H_
diff --git a/include/tvm/ffi/extra/structural_equal.h
b/include/tvm/ffi/extra/structural_equal.h
index 15876a56..13003fad 100644
--- a/include/tvm/ffi/extra/structural_equal.h
+++ b/include/tvm/ffi/extra/structural_equal.h
@@ -25,11 +25,14 @@
#include <tvm/ffi/any.h>
#include <tvm/ffi/extra/base.h>
+#include <tvm/ffi/function.h>
#include <tvm/ffi/optional.h>
#include <tvm/ffi/reflection/access_path.h>
+#include <tvm/ffi/reflection/accessor.h>
namespace tvm {
namespace ffi {
+
/*!
* \brief Structural equality comparators
*/
@@ -73,6 +76,23 @@ class StructuralEqual {
}
};
+namespace type_attr {
+/*!
+ * \brief Custom structural equality hook (used by ``StructuralEqual``).
+ *
+ * Unlike ``kEq`` (which is for ``RecursiveEq`` / Python ``==``), this is for
+ * the ``StructuralEqual`` system that supports def/use region semantics and
+ * alpha-equivalence of bound variables. The hook receives a callback to
+ * recursively compare sub-values.
+ *
+ * ``eq_cb`` recursively compares sub-values; its integer
``TVMFFIDefRegionKind``
+ * argument controls free-variable mapping; ``field_name`` is used for
mismatch path reporting.
+ */
+inline constexpr reflection::AttrKey<TypedFunction<bool(
+ const Object*, const Object*, const TypedFunction<bool(AnyView, AnyView,
int, AnyView)>&)>>
+ kSEqual{"__s_equal__"};
+} // namespace type_attr
+
} // namespace ffi
} // namespace tvm
#endif // TVM_FFI_EXTRA_STRUCTURAL_EQUAL_H_
diff --git a/include/tvm/ffi/extra/structural_hash.h
b/include/tvm/ffi/extra/structural_hash.h
index bfe023c3..41c07a52 100644
--- a/include/tvm/ffi/extra/structural_hash.h
+++ b/include/tvm/ffi/extra/structural_hash.h
@@ -25,6 +25,8 @@
#include <tvm/ffi/any.h>
#include <tvm/ffi/extra/base.h>
+#include <tvm/ffi/function.h>
+#include <tvm/ffi/reflection/accessor.h>
namespace tvm {
namespace ffi {
@@ -52,6 +54,23 @@ class StructuralHash {
TVM_FFI_INLINE uint64_t operator()(const Any& value) const { return
Hash(value); }
};
+namespace type_attr {
+/*!
+ * \brief Custom structural hash hook (used by ``StructuralHash``).
+ *
+ * Unlike ``kHash`` (which is for ``RecursiveHash`` / Python ``hash()``),
+ * this is for the ``StructuralHash`` system that supports def/use region
+ * semantics (``map_free_vars``). The hook receives the current accumulated
+ * hash and a callback to recursively hash sub-values.
+ *
+ * ``hash_cb`` recursively hashes a sub-value; its integer
``TVMFFIDefRegionKind``
+ * argument controls free-variable mapping.
+ */
+inline constexpr reflection::AttrKey<TypedFunction<int64_t(
+ const Object*, int64_t, const TypedFunction<int64_t(AnyView, int64_t,
int)>&)>>
+ kSHash{"__s_hash__"};
+} // namespace type_attr
+
} // namespace ffi
} // namespace tvm
#endif // TVM_FFI_EXTRA_STRUCTURAL_HASH_H_
diff --git a/include/tvm/ffi/extra/structural_mutate.h
b/include/tvm/ffi/extra/structural_mutate.h
index b8db9197..d768c3a0 100644
--- a/include/tvm/ffi/extra/structural_mutate.h
+++ b/include/tvm/ffi/extra/structural_mutate.h
@@ -164,10 +164,6 @@ struct StructuralMutatorVTable {
FStructuralVarRemapSet var_remap_set = nullptr;
};
-/*! \brief Borrowed structural mutation hook at attribute lookup and call. */
-using FStructuralMutate =
- reflection::NativeFunctionView<UnchangedOr<Any>(StructuralMutatorObj*,
AnyView)>;
-
namespace details {
template <typename Parent>
class StructuralMutateDynEngine;
@@ -433,6 +429,45 @@ TVM_FFI_COLD_CODE inline Unexpected<Error>
SMutateDeclaredTypeError() noexcept {
}
} // namespace details
+/*! \brief Borrowed structural mutation hook at attribute lookup and call. */
+using FStructuralMutate =
+ reflection::NativeFunctionView<UnchangedOr<Any>(StructuralMutatorObj*,
AnyView)>;
+
+namespace type_attr {
+/*!
+ * \brief Custom structural mutation hook used by ``StructuralMutator``.
+ *
+ * The hook receives the active mutator and the input value. It should
recursively mutate structural
+ * children through the mutator, return ``Unchanged`` if no fields are
changed, and avoid
+ * intentionally mutating the source object. The input is borrowed and remains
valid for the
+ * duration of the call.
+ *
+ * Value type: either an opaque native hook pointer returning raw
+ * ``Expected<UnchangedOr<Any>>`` storage, or an ``ffi::Function``.
+ *
+ * This is the canonical mutation hook. A type that defines
+ * ``kStructuralMaybeInplaceMutate`` must also define this attribute.
+ */
+inline constexpr reflection::AttrKey<FStructuralMutate>
kStructuralMutate{"__s_mutate__"};
+/*!
+ * \brief Optional custom mutation hook that may reuse its input object.
+ *
+ * The hook receives the active mutator and a borrowed input value that is
safe to mutate in place.
+ * Callers must route shared objects through non-in-place mutation instead.
The hook may mutate and
+ * return the source object, delegate to non-in-place mutation, or return
another replacement.
+ *
+ * Value type: either an opaque native hook pointer returning raw
+ * ``Expected<UnchangedOr<Any>>`` storage, or an ``ffi::Function``.
+ *
+ * This hook is optional. When it is absent, ``DefaultMutateExpected`` with
+ * ``inplace_mode=InplaceMode::kAllow`` falls back to non-in-place mutation
through
+ * ``kStructuralMutate`` or reflected structural fields. In-place mutation is
therefore explicitly
+ * opt-in and is never inferred from ownership by the reflected fallback.
+ */
+inline constexpr reflection::AttrKey<FStructuralMutate>
kStructuralMaybeInplaceMutate{
+ "__s_maybe_inplace_mutate__"};
+} // namespace type_attr
+
/*!
* \brief Object node of a structural mutator.
*/
@@ -620,7 +655,7 @@ class StructuralMutatorObj : public Object {
// Engine-internal: subclasses call the Expected forms above.
/*! \brief Raw default mutation: attr lookup then hook, favouring the fn-ptr
case. */
TVM_FFI_INLINE TVMFFIAny DefaultMutateRaw(AnyView value) noexcept {
- static reflection::TypeAttrColumn
column(reflection::type_attr::kStructuralMutate);
+ static reflection::TypeAttrColumn column(type_attr::kStructuralMutate);
AnyView attr = column[value.type_index()];
// Exactly one frame per node: hooks propagate errors untouched, and this
is the engine
// dispatching into `value`, so both exits below name it here and nowhere
else.
@@ -705,7 +740,7 @@ class StructuralMutatorObj : public Object {
* and every other shape is handed to the out-of-line remainder.
*/
TVM_FFI_INLINE TVMFFIAny DefaultMaybeInplaceMutateRaw(AnyView value)
noexcept {
- static reflection::TypeAttrColumn
column(reflection::type_attr::kStructuralMaybeInplaceMutate);
+ static reflection::TypeAttrColumn
column(type_attr::kStructuralMaybeInplaceMutate);
AnyView attr = column[value.type_index()];
if (TVM_FFI_PREDICT_TRUE(attr.type_index() ==
TypeIndex::kTVMFFIOpaquePtr)) {
// This is the engine dispatching into `value`; hooks propagate errors
untouched, so the
@@ -790,7 +825,7 @@ TVM_FFI_INLINE static Expected<Any>
MutateReflectedFieldsExpected(StructuralMuta
const Object* obj = value.as<Object>();
int32_t type_index = obj->type_index();
- static reflection::TypeAttrColumn
column(reflection::type_attr::kShallowCopy);
+ static reflection::TypeAttrColumn column(type_attr::kShallowCopy);
AnyView attr = column[type_index];
if (TVM_FFI_PREDICT_FALSE(attr.type_index() != TypeIndex::kTVMFFIFunction)) {
return Unexpected(Error("TypeError", "__ffi_shallow_copy__ must be an
ffi.Function", ""));
@@ -957,7 +992,8 @@ namespace details {
* copy->b = std::move(b).ValueOrUnchanged(std::move(copy->b));
* return copy;
* }
- * FStructuralMutate hook = FStructuralMutate::FromNative<&FooMutate>();
+ * FStructuralMutate hook =
+ * FStructuralMutate::FromNative<&FooMutate>();
* Expected<UnchangedOr<Any>> result = hook.CallExpected(mutator, value);
* \endcode
*
diff --git a/include/tvm/ffi/extra/structural_visit.h
b/include/tvm/ffi/extra/structural_visit.h
index de3eba98..055226c4 100644
--- a/include/tvm/ffi/extra/structural_visit.h
+++ b/include/tvm/ffi/extra/structural_visit.h
@@ -132,6 +132,24 @@ struct StructuralVisitorVTable {
using FStructuralVisit =
reflection::NativeFunctionView<Optional<VisitInterrupt>(StructuralVisitorObj*,
AnyView)>;
+namespace type_attr {
+/*!
+ * \brief Custom structural visitor hook (used by ``StructuralVisitor``).
+ *
+ * The hook receives the active visitor and the current object value. Opaque
+ * pointer hooks return a ``TVMFFIAny`` that stores an
+ * ``Expected<Optional<VisitInterrupt>>`` result.
+ *
+ * Value type: either an opaque native hook pointer or an ``ffi::Function``.
+ *
+ * On success, the hook returns ``None`` for no interrupt, or a
+ * ``VisitInterrupt`` to halt traversal. On failure, it returns an ``Error``.
+ * The ``visitor`` parameter is the active traversal context used to
recursively
+ * visit structural children.
+ */
+inline constexpr reflection::AttrKey<FStructuralVisit>
kStructuralVisit{"__s_visit__"};
+} // namespace type_attr
+
/*!
* \brief Object node of a structural visitor.
*
@@ -238,7 +256,7 @@ class StructuralVisitorObj : public Object {
private:
/*! \brief Internal raw ABI descent, avoiding a typed result around compiled
hooks. */
TVM_FFI_INLINE TVMFFIAny DefaultVisitRaw(AnyView value) noexcept {
- static reflection::TypeAttrColumn
column(reflection::type_attr::kStructuralVisit);
+ static reflection::TypeAttrColumn column(type_attr::kStructuralVisit);
AnyView attr = column[value.type_index()];
if (TVM_FFI_PREDICT_TRUE(attr.type_index() ==
TypeIndex::kTVMFFIOpaquePtr)) {
return
(*reinterpret_cast<FStructuralVisitABI>(attr.cast<void*>()))(this, value);
@@ -532,7 +550,8 @@ namespace details {
* TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->b));
* return Optional<VisitInterrupt>(std::nullopt);
* }
- * FStructuralVisit hook = FStructuralVisit::FromNative<&FooVisit>();
+ * FStructuralVisit hook =
+ * FStructuralVisit::FromNative<&FooVisit>();
* Expected<Optional<VisitInterrupt>> result = hook.CallExpected(visitor,
value);
* \endcode
*
diff --git a/include/tvm/ffi/reflection/accessor.h
b/include/tvm/ffi/reflection/accessor.h
index 18aa2f66..5c251787 100644
--- a/include/tvm/ffi/reflection/accessor.h
+++ b/include/tvm/ffi/reflection/accessor.h
@@ -26,13 +26,17 @@
#include <tvm/ffi/any.h>
#include <tvm/ffi/c_api.h>
#include <tvm/ffi/function.h>
+#include <tvm/ffi/reflection/native_function.h>
#include <tvm/ffi/type_traits.h>
+#include <optional>
#include <string>
namespace tvm {
namespace ffi {
namespace reflection {
+template <typename T>
+struct AttrKey;
/*!
* \brief helper function to get reflection field info by type key and field
name
@@ -177,6 +181,10 @@ class TypeAttrColumn {
TVM_FFI_THROW(RuntimeError) << "Cannot find type attribute " <<
attr_name;
}
}
+ /*! \brief Look up a typed key's column while preserving erased entry
access. */
+ template <typename T>
+ explicit TypeAttrColumn(AttrKey<T> key) : TypeAttrColumn(key.name) {}
+
/*!
* \brief Get the type attribute column by type index.
* \param type_index The type index.
@@ -325,13 +333,53 @@ inline bool ForEachFieldInfoWithEarlyStop(const TypeInfo*
type_info,
return false;
}
+/*!
+ * \brief C++17 constexpr helper: convert a ``const char*`` to
``TVMFFIByteArray``.
+ * \param s A null-terminated string literal.
+ * \return A ``TVMFFIByteArray`` whose ``data`` points to *s* and ``size``
equals ``strlen(s)``.
+ */
+inline constexpr TVMFFIByteArray AsByteArray(const char* s) {
+ return {s, std::char_traits<char>::length(s)};
+}
+
+/*!
+ * \brief A typed attribute name. Null-terminated name storage must outlive
the key.
+ *
+ * \code{.cpp}
+ * // In namespace tvm::ffi:
+ * namespace type_attr {
+ * inline constexpr reflection::AttrKey<bool>
kSerializable{"example.serializable"};
+ * }
+ * reflection::TypeAttrDef<MyObject>().attr(type_attr::kSerializable, true);
+ * \endcode
+ */
+template <typename T>
+struct AttrKey {
+ /*! \brief The value type associated with this key. */
+ using ValueType = T;
+ /*! \brief Borrowed null-terminated attribute name; its storage must outlive
the key. */
+ const char* name;
+
+ /*!
+ * \brief Extract a present value strictly; None denotes an absent entry.
+ * \throws Error if a present value has the wrong type. No fallback
conversion is run.
+ * \note Borrowed results remain valid only while the source value's owner
remains alive.
+ */
+ TVM_FFI_INLINE std::optional<ValueType> ValueFrom(AnyView value) const {
+ if (value == nullptr) return std::nullopt;
+ return value.template as_or_throw<ValueType>();
+ }
+};
+
+} // namespace reflection
+
/*!
* \brief Type attribute names used by the reflection system.
*
- * Each constant names a TypeAttrColumn — a sparse, type-indexed slot that
+ * Each key names a TypeAttrColumn — a sparse, type-indexed slot that
* stores a ``Function`` or any other Any-compliant value.
*
- * - ``TypeAttrDef<T>.def(type_attr::kFoo, ...)`` to register a hook,
+ * - ``TypeAttrDef<T>.def(type_attr::kFoo.name, ...)`` to register a hook,
* - ``TypeAttrColumn(type_attr::kFoo)[type_index]`` to look one up,
*
*/
@@ -342,10 +390,8 @@ namespace type_attr {
* Allocates a zero-initialised object of the registered type. For C++ types
* this wraps ``metadata->creator``; for Python ``@py_class`` types it is a
* ``calloc``-based allocator.
- *
- * Signature: ``() -> TSelf``, where ``TSelf`` is a subclass of ObjectRef.
*/
-inline constexpr const char* kNew = "__ffi_new__";
+inline constexpr reflection::AttrKey<TypedFunction<ObjectRef()>>
kNew{"__ffi_new__"};
/*!
* \brief Packed constructor — creates **and** initialises a new object.
*
@@ -356,21 +402,17 @@ inline constexpr const char* kNew = "__ffi_new__";
* For C++ types with ``refl::init<Args...>()``, the function is the explicit
* init; otherwise ``ffi._RegisterFFIInit`` generates a default one using
reflection.
*
- * Signature: ``(*args, **kwargs) -> TSelf``, where ``TSelf`` is a subclass of
ObjectRef.
- *
* Keyword arguments are packed as ``[KWARGS, key0, val0, key1, val1, ...]``.
*/
-inline constexpr const char* kInit = "__ffi_init__";
+inline constexpr reflection::AttrKey<Function> kInit{"__ffi_init__"};
/*!
* \brief Convert ``AnyView`` to a specific reflected ``TSelf`` type.
*
* Registered via ``TypeAttrDef<TObj>().def_convert<TSelf>()`` or
* ``ObjectDef<TObj>().def_convert<TSelf>()``. Used by the Python type
converter
* to marshal values into the correct ``TSelf`` subclass.
- *
- * Signature: ``(AnyView src) -> TSelf``, where ``TSelf`` has a registered
TypeTraits converter.
*/
-inline constexpr const char* kConvert = "__ffi_convert__";
+inline constexpr reflection::AttrKey<TypedFunction<Any(AnyView)>>
kConvert{"__ffi_convert__"};
/*!
* \brief Type schema accepted by ``kConvert`` before it returns ``TSelf``.
*
@@ -378,232 +420,52 @@ inline constexpr const char* kConvert =
"__ffi_convert__";
* generation uses this attribute to render widened input annotations while
* keeping output annotations precise.
*/
-inline constexpr const char* kConvertTypeSchema =
"__ffi_convert_type_schema__";
+inline constexpr reflection::AttrKey<Any>
kConvertTypeSchema{"__ffi_convert_type_schema__"};
/*!
* \brief Shallow-copy factory.
*
* Allocates a new object and copies all reflected field values from the
* source. Used by Python ``copy.copy()`` via ``__copy__`` and by
* ``copy.replace()`` via ``__replace__``.
- *
- * Signature: ``(TSelf self) -> TSelf``, where ``TSelf`` is a subclass of
ObjectRef.
*/
-inline constexpr const char* kShallowCopy = "__ffi_shallow_copy__";
+inline constexpr reflection::AttrKey<TypedFunction<ObjectRef(const Object*)>>
kShallowCopy{
+ "__ffi_shallow_copy__"};
/*!
- * \brief Custom recursive repr hook.
- *
- * If registered, ``RecursiveRepr`` (Python ``__repr__``) calls this instead
- * of the default field-by-field formatting. The hook receives a callback
- * ``fn_repr`` to recursively format nested values, avoiding infinite loops.
- *
- * Signature: ``(TSelf self, fn_repr: FnRepr) -> String``, where ``TSelf`` is
a subclass of
- * ObjectRef, and ``FnRepr: (AnyView value) -> String`` formats a nested value.
- */
-inline constexpr const char* kRepr = "__ffi_repr__";
-/*!
- * \brief Custom recursive hash hook.
- *
- * If registered, ``RecursiveHash`` (Python ``hash()``) calls this instead
- * of the default field-by-field hashing. The hook receives a callback
- * ``fn_hash`` to recursively hash nested values, with cycle detection.
- *
- * Signature: ``(TSelf self, fn_hash: FnHash) -> int64``, where ``TSelf`` is a
subclass of
- * ObjectRef, and ``FnHash: (AnyView value) -> int64`` hashes a nested value.
- */
-inline constexpr const char* kHash = "__ffi_hash__";
-/*!
- * \brief Custom recursive equality hook.
- *
- * If registered, ``RecursiveEq`` (Python ``==``) calls this instead of the
- * default field-by-field comparison. Falls back to ``kCompare`` if only the
- * compare hook is present. The hook receives a callback ``fn_eq`` to
- * recursively compare nested values.
- *
- * Signature: ``(TSelf lhs, TSelf rhs, fn_eq: FnEq) -> bool``, where ``TSelf``
is a subclass of
- * ObjectRef, and ``FnEq: (AnyView lhs, AnyView rhs) -> bool`` compares nested
values.
- */
-inline constexpr const char* kEq = "__ffi_eq__";
-/*!
- * \brief Custom recursive three-way comparison hook.
- *
- * If registered, ``RecursiveCompare`` (Python ``<``, ``<=``, ``>``, ``>=``)
- * calls this. Also used as fallback for ``RecursiveEq`` when ``kEq`` is not
- * registered. The hook receives a callback ``fn_cmp`` to recursively
- * compare nested values.
+ * \brief Whether the class is declared final (``Class::_type_final``).
*
- * Signature: ``(TSelf lhs, TSelf rhs, fn_cmp: FnCmp) -> int32``, where
``TSelf`` is a subclass of
- * ObjectRef, and ``FnCmp: (TSelf lhs, TSelf rhs) -> int32`` returns < 0, ==
0, or > 0.
+ * Registered as a ``bool`` by every ``ObjectDef<Class>``. A final class can
+ * never gain a subtype, so a binding generator may store it by value.
*/
-inline constexpr const char* kCompare = "__ffi_compare__";
+inline constexpr reflection::AttrKey<bool> kTypeFinal{"__ffi_type_final__"};
/*!
* \brief Custom Any-level hash for use as ``Map``/``Dict`` key.
*
- * Unlike ``kHash`` (which operates on ``ObjectRef`` and is recursive), this
+ * Unlike ``kHash`` (which operates on objects and is recursive), this
* hook operates at the ``Any`` level — it is called by ``AnyHash`` whenever
* an ``Any`` holding an object of this type is used as a container key.
*
- * Register a native hook with ``NativeFunction<int64_t(const
Any&)>::FromNative<&Fn>()``
- * or a packed ``Function`` object.
+ * Register a native hook with ``TypeAttrDef<T>::native_fn_attr<&Fn>(key)``
+ * or a packed ``Function`` object through the name-based registration API.
*
- * Native adapter signature: ``TVMFFIAny (*)(const Any& src) noexcept``
- *
- * Function signature: ``(Any src) -> int64``
+ * Native adapters use a context-free noexcept ABI and return raw
``TVMFFIAny`` storage.
*/
-inline constexpr const char* kAnyHash = "__any_hash__";
+inline constexpr
reflection::AttrKey<reflection::NativeFunctionView<int64_t(const Any&)>>
kAnyHash{
+ "__any_hash__"};
/*!
* \brief Custom Any-level equality for use as ``Map``/``Dict`` key.
*
- * Unlike ``kEq`` (which operates on ``ObjectRef``), this hook operates at the
+ * Unlike ``kEq`` (which operates on objects), this hook operates at the
* ``Any`` level — called by ``AnyEqual`` for container key comparison.
*
- * Register a native hook with
- * ``NativeFunction<bool(const Any&, const Any&)>::FromNative<&Fn>()``
- * or a packed ``Function`` object.
- *
- * Native adapter signature: ``TVMFFIAny (*)(const Any& lhs, const Any& rhs)
noexcept``
- *
- * Function signature: ``(Any lhs, Any rhs) -> bool``
- */
-inline constexpr const char* kAnyEqual = "__any_equal__";
-/*!
- * \brief Custom structural hash hook (used by ``StructuralHash``).
- *
- * Unlike ``kHash`` (which is for ``RecursiveHash`` / Python ``hash()``),
- * this is for the ``StructuralHash`` system that supports def/use region
- * semantics (``map_free_vars``). The hook receives the current accumulated
- * hash and a callback to recursively hash sub-values.
- *
- * Signature: ``(TSelf self, int64 init_hash, Function hash_cb) -> int64``,
where
- * ``TSelf`` is a subclass of ``ObjectRef``.
- *
- * ``hash_cb(AnyView val, int64 init_hash, bool def_region) -> int64``
- * recursively hashes a sub-value; ``def_region`` controls free-variable
- * mapping.
- */
-inline constexpr const char* kSHash = "__s_hash__";
-/*!
- * \brief Custom structural equality hook (used by ``StructuralEqual``).
- *
- * Unlike ``kEq`` (which is for ``RecursiveEq`` / Python ``==``), this is for
- * the ``StructuralEqual`` system that supports def/use region semantics and
- * alpha-equivalence of bound variables. The hook receives a callback to
- * recursively compare sub-values.
- *
- * Signature: ``(TSelf lhs, TSelf rhs, Function eq_cb) -> bool``, where
``TSelf`` is a
- * subclass of ``ObjectRef``.
+ * Register a native hook with ``TypeAttrDef<T>::native_fn_attr<&Fn>(key)``
+ * or a packed ``Function`` object through the name-based registration API.
*
- * ``eq_cb(AnyView lhs, AnyView rhs, bool def_region, AnyView field_name) ->
bool``
- * recursively compares sub-values; ``def_region`` controls free-variable
- * mapping; ``field_name`` is used for mismatch path reporting.
+ * Native adapters use a context-free noexcept ABI and return raw
``TVMFFIAny`` storage.
*/
-inline constexpr const char* kSEqual = "__s_equal__";
-/*!
- * \brief Custom structural visitor hook (used by ``StructuralVisitor``).
- *
- * The hook receives the active visitor and the current object value. Opaque
- * pointer hooks return a ``TVMFFIAny`` that stores an
- * ``Expected<Optional<VisitInterrupt>>`` result.
- *
- * Value type: either an opaque function pointer to a C++ structural visit hook
- *
- * ``TVMFFIAny (*)(StructuralVisitorObj* visitor, AnyView value) noexcept``
- *
- * or an ``ffi::Function`` with signature
- *
- * ``(StructuralVisitor visitor, Any value) -> Optional<VisitInterrupt>``.
- *
- * On success, the hook returns ``None`` for no interrupt, or a
- * ``VisitInterrupt`` to halt traversal. On failure, it returns an ``Error``.
- * The ``visitor`` parameter is the active traversal context used to
recursively
- * visit structural children.
- */
-inline constexpr const char* kStructuralVisit = "__s_visit__";
-
-/*!
- * \brief Custom structural mutation hook used by ``StructuralMutator``.
- *
- * The hook receives the active mutator and the input value. It should
recursively mutate structural
- * children through the mutator, return the input unchanged if no fields are
changed, and avoid
- * intentionally mutating the source object. The input is borrowed and remains
valid for the
- * duration of the call.
- *
- * Value type: either an opaque function pointer to a C++ structural mutation
hook
- *
- * ``TVMFFIAny (*)(StructuralMutatorObj* mutator, AnyView value) noexcept``
- *
- * returning raw ``Expected<Any>`` storage, or an ``ffi::Function`` with
signature
- *
- * ``(StructuralMutator mutator, Any value) -> Any``.
- *
- * This is the canonical mutation hook. A type that defines
- * ``kStructuralMaybeInplaceMutate`` must also define this attribute.
- */
-inline constexpr const char* kStructuralMutate = "__s_mutate__";
-/*!
- * \brief Optional custom mutation hook that may reuse its input object.
- *
- * The hook receives the active mutator and a borrowed input value that is
safe to mutate in place.
- * Callers must route shared objects through non-in-place mutation instead.
The hook may mutate and
- * return the source object, delegate to non-in-place mutation, or return
another replacement.
- *
- * Value type: either an opaque function pointer to a C++ structural mutation
hook
- *
- * ``TVMFFIAny (*)(StructuralMutatorObj* mutator, AnyView value) noexcept``
- *
- * returning raw ``Expected<Any>`` storage, or an
- * ``ffi::Function`` with signature
- *
- * ``(StructuralMutator mutator, Any value) -> Any``.
- *
- * This hook is optional. When it is absent, ``DefaultMutateExpected`` with
- * ``inplace_mode=InplaceMode::kAllow`` falls back to non-in-place mutation
through
- * ``kStructuralMutate`` or reflected structural fields. In-place mutation is
therefore explicitly
- * opt-in and is never inferred from ownership by the reflected fallback.
- */
-inline constexpr const char* kStructuralMaybeInplaceMutate =
"__s_maybe_inplace_mutate__";
-
-/*!
- * \brief Serialize object data to a JSON-compatible value.
- *
- * If registered, ``ToJSONGraph`` calls this instead of the default
- * field-by-field serialization. Allows types with non-trivial internal
- * state (e.g. ``TInt`` storing a plain ``int64_t``) to define a compact
- * custom JSON representation.
- *
- * Signature: ``(TSelf self) -> Any``, where ``TSelf`` is a subclass of
``ObjectRef``.
- */
-inline constexpr const char* kDataToJson = "__data_to_json__";
-/*!
- * \brief Deserialize object data from a JSON-compatible value.
- *
- * Counterpart to ``kDataToJson``. If registered, ``FromJSONGraph`` calls
- * this to reconstruct the object from its serialized representation
- * instead of using field-by-field deserialization.
- *
- * Signature: ``(Any json_data) -> TSelf``, where ``TSelf`` is a subclass of
``ObjectRef``.
- */
-inline constexpr const char* kDataFromJson = "__data_from_json__";
-/*! \brief Per-class enum state: ordered entries, canonical indices, and
extensible attrs. */
-inline constexpr const char* kEnumState = "__ffi_enum__";
-/*!
- * \brief Whether the class is declared final (``Class::_type_final``).
- *
- * Registered as a ``bool`` by every ``ObjectDef<Class>``. A final class can
- * never gain a subtype, so a binding generator may store it by value.
- */
-inline constexpr const char* kTypeFinal = "__ffi_type_final__";
+inline constexpr reflection::AttrKey<reflection::NativeFunctionView<bool(const
Any&, const Any&)>>
+ kAnyEqual{"__any_equal__"};
} // namespace type_attr
-/*!
- * \brief C++17 constexpr helper: convert a ``const char*`` to
``TVMFFIByteArray``.
- * \param s A null-terminated string literal.
- * \return A ``TVMFFIByteArray`` whose ``data`` points to *s* and ``size``
equals ``strlen(s)``.
- */
-inline constexpr TVMFFIByteArray AsByteArray(const char* s) {
- return {s, std::char_traits<char>::length(s)};
-}
-
-} // namespace reflection
} // namespace ffi
} // namespace tvm
#endif // TVM_FFI_REFLECTION_ACCESSOR_H_
diff --git a/include/tvm/ffi/reflection/creator.h
b/include/tvm/ffi/reflection/creator.h
index 588a035a..1873d7b5 100644
--- a/include/tvm/ffi/reflection/creator.h
+++ b/include/tvm/ffi/reflection/creator.h
@@ -48,14 +48,14 @@ inline ObjectPtr<Object> CreateEmptyObject(const
TVMFFITypeInfo* type_info) {
return
details::ObjectUnsafe::ObjectPtrFromOwned<Object>(static_cast<TVMFFIObject*>(handle));
}
// Fallback: __ffi_new__ type attr (Python-defined types)
- constexpr TVMFFIByteArray kFFINewAttrName =
reflection::AsByteArray(reflection::type_attr::kNew);
+ constexpr TVMFFIByteArray kFFINewAttrName =
reflection::AsByteArray(type_attr::kNew.name);
const TVMFFITypeAttrColumn* column =
TVMFFIGetTypeAttrColumn(&kFFINewAttrName);
if (column != nullptr) {
int32_t offset = type_info->type_index - column->begin_index;
if (offset >= 0 && offset < column->size) {
AnyView attr_view = AnyView::CopyFromTVMFFIAny(column->data[offset]);
- if (auto opt_func = attr_view.try_cast<Function>()) {
- ObjectRef obj_ref = (*opt_func)().cast<ObjectRef>();
+ if (auto opt_func = type_attr::kNew.ValueFrom(attr_view)) {
+ ObjectRef obj_ref = (*opt_func)();
return
details::ObjectUnsafe::ObjectPtrFromObjectRef<Object>(std::move(obj_ref));
}
}
@@ -78,7 +78,7 @@ inline bool HasCreator(const TVMFFITypeInfo* type_info) {
if (type_info->metadata != nullptr && type_info->metadata->creator !=
nullptr) {
return true;
}
- constexpr TVMFFIByteArray kFFINewAttrName =
reflection::AsByteArray(reflection::type_attr::kNew);
+ constexpr TVMFFIByteArray kFFINewAttrName =
reflection::AsByteArray(type_attr::kNew.name);
const TVMFFITypeAttrColumn* column =
TVMFFIGetTypeAttrColumn(&kFFINewAttrName);
if (column != nullptr) {
int32_t offset = type_info->type_index - column->begin_index;
diff --git a/include/tvm/ffi/reflection/enum_def.h
b/include/tvm/ffi/reflection/enum_def.h
index 4b72fe3a..f444014a 100644
--- a/include/tvm/ffi/reflection/enum_def.h
+++ b/include/tvm/ffi/reflection/enum_def.h
@@ -93,6 +93,18 @@ class EnumDef : public ReflectionDefBase {
return *this;
}
+ /*! \brief Set an enum attribute using its key's declared value type. */
+ template <typename T, typename Value>
+ EnumDef& set_attr(AttrKey<T> key, Value&& value) {
+ return set_attr(key.name, T(std::forward<Value>(value)));
+ }
+
+ /*! \brief Set a named native callback using its key's declared adapter. */
+ template <auto Fn, typename T>
+ EnumDef& set_native_fn_attr(AttrKey<T> key) {
+ return set_attr(key, T::template FromNative<Fn>());
+ }
+
private:
void Register(const EnumState& state, int64_t int_index, String str_index) {
List<ObjectRef> entries = state->entries;
@@ -133,7 +145,7 @@ class EnumDef : public ReflectionDefBase {
::tvm::ffi::details::ObjectUnsafe::GetHeader(state_obj.get())->type_index
= state_type_index;
EnumState state =
::tvm::ffi::details::ObjectUnsafe::ObjectRefFromObjectPtr<EnumState>(std::move(state_obj));
- constexpr TVMFFIByteArray name = AsByteArray(type_attr::kEnumState);
+ constexpr TVMFFIByteArray name = AsByteArray(type_attr::kEnumState.name);
TVMFFIAny value_any = AnyView(state).CopyToTVMFFIAny();
TVM_FFI_CHECK_SAFE_CALL(TVMFFITypeRegisterAttr(type_index_, &name,
&value_any));
return state;
diff --git a/include/tvm/ffi/reflection/overload.h
b/include/tvm/ffi/reflection/overload.h
index c9193c40..324704cc 100644
--- a/include/tvm/ffi/reflection/overload.h
+++ b/include/tvm/ffi/reflection/overload.h
@@ -30,6 +30,7 @@
#include <tvm/ffi/function.h>
#include <tvm/ffi/function_details.h>
#include <tvm/ffi/optional.h>
+#include <tvm/ffi/reflection/accessor.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ffi/string.h>
#include <tvm/ffi/type_traits.h>
@@ -398,7 +399,7 @@ class OverloadObjectDef : private ObjectDef<Class> {
// via Register() are automatically visible through the TypeAttrColumn.
if (is_first) {
const TVMFFITypeInfo* tinfo = TVMFFIGetTypeInfo(type_index_);
- constexpr TVMFFIByteArray attr_name = AsByteArray(type_attr::kInit);
+ constexpr TVMFFIByteArray attr_name = AsByteArray(type_attr::kInit.name);
for (int32_t i = 0; i < tinfo->num_methods; ++i) {
if (tinfo->methods[i].name.size == attr_name.size &&
std::strncmp(tinfo->methods[i].name.data, attr_name.data,
attr_name.size) == 0) {
diff --git a/include/tvm/ffi/reflection/registry.h
b/include/tvm/ffi/reflection/registry.h
index 6b9f8be4..e702fa64 100644
--- a/include/tvm/ffi/reflection/registry.h
+++ b/include/tvm/ffi/reflection/registry.h
@@ -29,7 +29,6 @@
#include <tvm/ffi/container/variant.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/function_details.h>
-#include <tvm/ffi/optional.h>
#include <tvm/ffi/reflection/creator.h>
#include <tvm/ffi/string.h>
#include <tvm/ffi/type_traits.h>
@@ -47,8 +46,8 @@ namespace tvm {
namespace ffi {
/*! \brief Reflection namespace */
namespace reflection {
-namespace details {
+namespace details {
template <typename T, typename = void>
struct HasContainerType : public std::false_type {};
@@ -755,8 +754,8 @@ class ObjectDef : public ReflectionDefBase {
if constexpr (std::is_copy_constructible_v<Class>) {
Function fn = Function::FromTyped(
[](const Class* self) { return
ObjectRef(ffi::make_object<Class>(*self)); },
- std::string(type_key_) + "." + type_attr::kShallowCopy);
- TVMFFIByteArray attr = AsByteArray(type_attr::kShallowCopy);
+ std::string(type_key_) + "." + type_attr::kShallowCopy.name);
+ TVMFFIByteArray attr = AsByteArray(type_attr::kShallowCopy.name);
TVMFFIAny fn_any = AnyView(fn).CopyToTVMFFIAny();
TVM_FFI_CHECK_SAFE_CALL(TVMFFITypeRegisterAttr(type_index_, &attr,
&fn_any));
}
@@ -770,8 +769,8 @@ class ObjectDef : public ReflectionDefBase {
return
ObjectRef(::tvm::ffi::details::ObjectUnsafe::ObjectPtrFromOwned<Object>(
static_cast<TVMFFIObject*>(handle)));
},
- std::string(type_key_) + "." + type_attr::kNew);
- TVMFFIByteArray attr = AsByteArray(type_attr::kNew);
+ std::string(type_key_) + "." + type_attr::kNew.name);
+ TVMFFIByteArray attr = AsByteArray(type_attr::kNew.name);
TVMFFIAny fn_any = AnyView(fn).CopyToTVMFFIAny();
TVM_FFI_CHECK_SAFE_CALL(TVMFFITypeRegisterAttr(type_index_, &attr,
&fn_any));
// Step 3. Register `__ffi_init__` if no explicit init is defined.
@@ -784,7 +783,7 @@ class ObjectDef : public ReflectionDefBase {
}
}
// Step 4. Publish finality, which no other registered metadata exposes.
- RegisterTypeAttrValue(type_index_, type_attr::kTypeFinal,
Class::_type_final);
+ RegisterTypeAttrValue(type_index_, type_attr::kTypeFinal.name,
Class::_type_final);
}
/*!
@@ -918,7 +917,8 @@ class ObjectDef : public ReflectionDefBase {
static_assert(std::is_same_v<typename TSelf::ContainerType, Class>,
"TSelf::ContainerType must match ObjectDef's object type");
}
- return def_type_attr(type_attr::kConvert,
&details::FFIConvertFromAnyViewToObjectRef<TSelf>);
+ return def_type_attr(type_attr::kConvert.name,
+ &details::FFIConvertFromAnyViewToObjectRef<TSelf>);
}
/*!
@@ -951,7 +951,7 @@ class ObjectDef : public ReflectionDefBase {
std::forward<Extra>(extra)...);
// Also mirror into __ffi_init__ TypeAttrColumn for runtime dispatch.
const TVMFFITypeInfo* tinfo = TVMFFIGetTypeInfo(type_index_);
- constexpr TVMFFIByteArray attr_name = AsByteArray(type_attr::kInit);
+ constexpr TVMFFIByteArray attr_name = AsByteArray(type_attr::kInit.name);
for (int32_t i = 0; i < tinfo->num_methods; ++i) {
if (tinfo->methods[i].name.size == attr_name.size &&
std::strncmp(tinfo->methods[i].name.data, attr_name.data,
attr_name.size) == 0) {
@@ -1050,7 +1050,7 @@ class ObjectDef : public ReflectionDefBase {
int32_t type_index_;
const char* type_key_;
bool has_explicit_init_{false};
- static constexpr const char* kInitMethodName = type_attr::kInit;
+ static constexpr const char* kInitMethodName = type_attr::kInit.name;
};
/*!
@@ -1075,6 +1075,12 @@ class TypeAttrDef : public ReflectionDefBase {
explicit TypeAttrDef(ExtraArgs&&... extra_args)
: type_index_(Class::_GetOrAllocRuntimeTypeIndex()),
type_key_(Class::_type_key) {}
+ /*! \brief Register a named native callback using its key's declared
adapter. */
+ template <auto Fn, typename T>
+ TypeAttrDef& native_fn_attr(AttrKey<T> key) {
+ return attr(key, T::template FromNative<Fn>());
+ }
+
/*!
* \brief Define a value type attribute.
*
@@ -1131,7 +1137,15 @@ class TypeAttrDef : public ReflectionDefBase {
static_assert(std::is_same_v<typename TSelf::ContainerType, Class>,
"TSelf::ContainerType must match TypeAttrDef's object
type");
}
- return def(type_attr::kConvert,
&details::FFIConvertFromAnyViewToObjectRef<TSelf>);
+ return def(type_attr::kConvert.name,
&details::FFIConvertFromAnyViewToObjectRef<TSelf>);
+ }
+
+ /*! \brief Register a value using its key's declared type. */
+ template <typename T, typename Value>
+ TypeAttrDef& attr(AttrKey<T> key, Value&& value) {
+ T typed_value(std::forward<Value>(value));
+ RegisterTypeAttrValue(type_index_, key.name, std::move(typed_value));
+ return *this;
}
/*!
@@ -1145,7 +1159,8 @@ class TypeAttrDef : public ReflectionDefBase {
* \return The TypeAttrDef object.
*/
template <typename T>
- TypeAttrDef& attr(const char* name, T value) { //
NOLINT(performance-unnecessary-value-param)
+ TypeAttrDef& attr(const char* name,
+ T value) { // NOLINT(performance-unnecessary-value-param)
RegisterTypeAttrValue(type_index_, name, std::move(value));
return *this;
}
@@ -1167,6 +1182,12 @@ inline void EnsureTypeAttrColumn(std::string_view name) {
reinterpret_cast<const
TVMFFIAny*>(&any_view)));
}
+/*! \brief Ensure the erased column named by a typed key exists. */
+template <typename T>
+inline void EnsureTypeAttrColumn(AttrKey<T> key) {
+ EnsureTypeAttrColumn(key.name);
+}
+
/// \cond Doxygen_Suppress
namespace details {
diff --git a/src/ffi/container.cc b/src/ffi/container.cc
index 37f000ef..f36f5e74 100644
--- a/src/ffi/container.cc
+++ b/src/ffi/container.cc
@@ -26,7 +26,10 @@
#include <tvm/ffi/container/map.h>
#include <tvm/ffi/container/tensor.h>
#include <tvm/ffi/device.h>
+#include <tvm/ffi/extra/dataclass.h>
+#include <tvm/ffi/extra/serialization.h>
#include <tvm/ffi/function.h>
+#include <tvm/ffi/reflection/native_function.h>
#include <tvm/ffi/reflection/registry.h>
#include <sstream>
@@ -103,44 +106,57 @@ class MapForwardIterFunctor {
ffi::MapObj::iterator end_;
};
+int64_t BigIntAnyHash(const Any& value) {
+ return details::int_ops::BitcastToInt64(value.as_or_throw<BigInt>().hash());
+}
+
+bool BigIntAnyEqual(const Any& lhs, const Any& rhs) {
+ return lhs.as_or_throw<BigInt>() == rhs.as_or_throw<BigInt>();
+}
+
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
- refl::EnsureTypeAttrColumn(refl::type_attr::kAnyHash);
- refl::EnsureTypeAttrColumn(refl::type_attr::kAnyEqual);
+ refl::EnsureTypeAttrColumn(type_attr::kAnyHash);
+ refl::EnsureTypeAttrColumn(type_attr::kAnyEqual);
refl::ObjectDef<details::BigIntObj>()
- .def_type_attr(refl::type_attr::kAnyHash,
- [](const BigInt& value) -> int64_t {
- return details::int_ops::BitcastToInt64(value.hash());
+ .def_type_attr(
+ type_attr::kAnyHash.name,
+ reflection::NativeFunctionView<int64_t(const
Any&)>::FromNative<&BigIntAnyHash>())
+ .def_type_attr(type_attr::kAnyEqual.name,
+ reflection::NativeFunctionView<bool(
+ const Any&, const
Any&)>::FromNative<&BigIntAnyEqual>())
+ .def_type_attr(type_attr::kEq.name,
+ [](const Object* lhs, const Object* rhs, const Function&)
{
+ return AnyView(lhs).as_or_throw<BigInt>() ==
+ AnyView(rhs).as_or_throw<BigInt>();
})
- .def_type_attr(refl::type_attr::kAnyEqual,
- [](const BigInt& lhs, const BigInt& rhs) { return lhs ==
rhs; })
- .def_type_attr(refl::type_attr::kEq, [](const BigInt& lhs, const BigInt&
rhs,
- const Function&) { return lhs ==
rhs; })
- .def_type_attr(refl::type_attr::kCompare,
- [](const BigInt& lhs, const BigInt& rhs, const Function&)
{
+ .def_type_attr(type_attr::kCompare.name,
+ [](const Object* lhs, const Object* rhs, const Function&)
{
return details::int_ops::CompareFallback(
- details::BigIntUnsafe::GetArrayView(lhs),
- details::BigIntUnsafe::GetArrayView(rhs));
- })
- .def_type_attr(refl::type_attr::kHash,
- [](const BigInt& value, const Function&) -> int64_t {
- return details::int_ops::BitcastToInt64(value.hash());
+
details::BigIntUnsafe::GetArrayView(AnyView(lhs).as_or_throw<BigInt>()),
+
details::BigIntUnsafe::GetArrayView(AnyView(rhs).as_or_throw<BigInt>()));
})
- .def_type_attr(refl::type_attr::kDataToJson,
- [](const BigInt& value) -> String {
+ .def_type_attr(
+ type_attr::kHash.name,
+ [](const Object* value, const Function&) -> int64_t {
+ return
details::int_ops::BitcastToInt64(AnyView(value).as_or_throw<BigInt>().hash());
+ })
+ .def_type_attr(type_attr::kDataToJson.name,
+ [](AnyView value) -> Any {
std::ostringstream stream;
- stream << value;
- return stream.str();
+ stream << value.as_or_throw<BigInt>();
+ return String(stream.str());
})
- .def_type_attr(refl::type_attr::kDataFromJson, [](const String& data) ->
BigInt {
+ .def_type_attr(type_attr::kDataFromJson.name, [](AnyView value) -> Any {
+ String data = value.as_or_throw<String>();
std::istringstream stream(static_cast<std::string>(data));
- BigInt value;
- stream >> std::noskipws >> value;
+ BigInt result;
+ stream >> std::noskipws >> result;
// A graph payload must be one complete decimal integer, without
whitespace or junk.
if (stream.fail() || !stream.eof()) {
TVM_FFI_THROW(ValueError) << "Invalid decimal string for ffi.BigInt";
}
- return value;
+ return result;
});
refl::GlobalDef()
.def_packed("ffi.Array",
diff --git a/src/ffi/extra/dataclass.cc b/src/ffi/extra/dataclass.cc
index b1d7ffbc..170adb3a 100644
--- a/src/ffi/extra/dataclass.cc
+++ b/src/ffi/extra/dataclass.cc
@@ -340,12 +340,12 @@ class ObjectDeepCopier : public
ObjectGraphDFS<ObjectDeepCopier, CopyFrame, Any>
default: {
// Reflected object: shallow-copy and register
const TVMFFITypeInfo* type_info = TVMFFIGetTypeInfo(ti);
- TVM_FFI_ICHECK((*column_)[ti] != nullptr)
+ auto copy_fn = type_attr::kShallowCopy.ValueFrom((*column_)[ti]);
+ TVM_FFI_ICHECK(copy_fn.has_value())
<< "Cannot deep copy object of type \""
<< std::string_view(type_info->type_key.data,
type_info->type_key.size)
<< "\" because it is not copy-constructible";
- Function copy_fn = (*column_)[ti].cast<Function>();
- f.copy = copy_fn(f.obj);
+ f.copy = (*copy_fn)(f.obj);
copy_map_[f.obj] = f.copy;
break;
}
@@ -884,13 +884,11 @@ class ReprPrinter : public ObjectGraphDFS<ReprPrinter,
ReprFrame, std::string> {
return true;
}
// Custom __ffi_repr__ hook
- static refl::TypeAttrColumn repr_column(refl::type_attr::kRepr);
- AnyView custom_repr = repr_column[ti];
- if (custom_repr != nullptr) {
+ static refl::TypeAttrColumn repr_column(type_attr::kRepr);
+ if (auto repr_fn = type_attr::kRepr.ValueFrom(repr_column[ti])) {
state_[obj] = State::kInProgress;
- Function repr_fn = custom_repr.cast<Function>();
- Function fn_repr = CreateFnRepr();
- String r = repr_fn(obj, fn_repr).cast<String>();
+ TypedFunction<String(AnyView)> fn_repr = CreateFnRepr();
+ String r = (*repr_fn)(obj, fn_repr);
std::string result(r.data(), r.size());
if (show_addr_ && (ti == TypeIndex::kTVMFFIArray || ti ==
TypeIndex::kTVMFFIList ||
ti == TypeIndex::kTVMFFIMap || ti ==
TypeIndex::kTVMFFIDict)) {
@@ -1095,13 +1093,11 @@ class RecursiveHasher : public
ObjectGraphDFS<RecursiveHasher, HashFrame, uint64
return true;
}
// Check for custom __ffi_hash__ hook
- static refl::TypeAttrColumn hash_column(refl::type_attr::kHash);
- AnyView custom = hash_column[obj->type_index()];
- if (custom != nullptr) {
+ static refl::TypeAttrColumn hash_column(type_attr::kHash);
+ if (auto hook =
type_attr::kHash.ValueFrom(hash_column[obj->type_index()])) {
on_stack_.insert(obj);
- Function hook = custom.cast<Function>();
- Function fn_hash = CreateFnHash();
- int64_t r = hook(obj, fn_hash).cast<int64_t>();
+ TypedFunction<int64_t(AnyView)> fn_hash = CreateFnHash();
+ int64_t r = (*hook)(obj, fn_hash);
uint64_t h = static_cast<uint64_t>(r);
memo_[obj] = h;
on_stack_.erase(obj);
@@ -1111,8 +1107,8 @@ class RecursiveHasher : public
ObjectGraphDFS<RecursiveHasher, HashFrame, uint64
// For reflected types (not built-in containers), error if the type has
// __ffi_eq__ or __ffi_compare__ but no __ffi_hash__.
if (ti >= TypeIndex::kTVMFFIStaticObjectEnd) {
- static refl::TypeAttrColumn eq_column(refl::type_attr::kEq);
- static refl::TypeAttrColumn cmp_column(refl::type_attr::kCompare);
+ static refl::TypeAttrColumn eq_column(type_attr::kEq);
+ static refl::TypeAttrColumn cmp_column(type_attr::kCompare);
if (eq_column[obj->type_index()] != nullptr ||
cmp_column[obj->type_index()] != nullptr) {
const TVMFFITypeInfo* info = TVMFFIGetTypeInfo(ti);
TVM_FFI_THROW(ValueError)
@@ -1402,39 +1398,36 @@ class RecursiveComparer : public
ObjectGraphDFS<RecursiveComparer, CompareFrame,
<< String(lhs_info->type_key) << " vs "
<< String(rhs_info->type_key);
}
- static refl::TypeAttrColumn eq_column(refl::type_attr::kEq);
- static refl::TypeAttrColumn cmp_column(refl::type_attr::kCompare);
+ static refl::TypeAttrColumn eq_column(type_attr::kEq);
+ static refl::TypeAttrColumn cmp_column(type_attr::kCompare);
int32_t ti = lhs_obj->type_index();
AnyView custom_eq = eq_column[ti];
AnyView custom_cmp = cmp_column[ti];
if (eq_only_) {
- if (custom_eq != nullptr) {
+ if (auto hook = type_attr::kEq.ValueFrom(custom_eq)) {
auto pair = std::make_pair(lhs_obj, rhs_obj);
on_stack_.insert(pair);
- Function hook = custom_eq.cast<Function>();
- Function fn_eq = CreateFnEq();
- bool result = hook(lhs_obj, rhs_obj, fn_eq).cast<bool>();
+ TypedFunction<bool(AnyView, AnyView)> fn_eq = CreateFnEq();
+ bool result = (*hook)(lhs_obj, rhs_obj, fn_eq);
on_stack_.erase(pair);
*out = result ? 0 : 1;
return true;
}
- if (custom_cmp != nullptr) {
+ if (auto hook = type_attr::kCompare.ValueFrom(custom_cmp)) {
auto pair = std::make_pair(lhs_obj, rhs_obj);
on_stack_.insert(pair);
- Function hook = custom_cmp.cast<Function>();
- Function fn_cmp = CreateFnCompare();
- int32_t result = hook(lhs_obj, rhs_obj, fn_cmp).cast<int32_t>();
+ TypedFunction<int32_t(AnyView, AnyView)> fn_cmp = CreateFnCompare();
+ int32_t result = (*hook)(lhs_obj, rhs_obj, fn_cmp);
on_stack_.erase(pair);
*out = result;
return true;
}
} else {
- if (custom_cmp != nullptr) {
+ if (auto hook = type_attr::kCompare.ValueFrom(custom_cmp)) {
auto pair = std::make_pair(lhs_obj, rhs_obj);
on_stack_.insert(pair);
- Function hook = custom_cmp.cast<Function>();
- Function fn_cmp = CreateFnCompare();
- int32_t result = hook(lhs_obj, rhs_obj, fn_cmp).cast<int32_t>();
+ TypedFunction<int32_t(AnyView, AnyView)> fn_cmp = CreateFnCompare();
+ int32_t result = (*hook)(lhs_obj, rhs_obj, fn_cmp);
on_stack_.erase(pair);
*out = result;
return true;
@@ -2138,7 +2131,7 @@ Function MakeInit(int32_t type_index) {
void RegisterFFIInit(int32_t type_index) {
namespace refl = ::tvm::ffi::reflection;
Function auto_init_fn = MakeInit(type_index);
- TVMFFIByteArray attr_name = refl::AsByteArray(refl::type_attr::kInit);
+ TVMFFIByteArray attr_name = refl::AsByteArray(type_attr::kInit.name);
TVMFFIAny attr_value = AnyView(auto_init_fn).CopyToTVMFFIAny();
TVM_FFI_CHECK_SAFE_CALL(TVMFFITypeRegisterAttr(type_index, &attr_name,
&attr_value));
}
@@ -2167,7 +2160,7 @@ void PyClassRegisterTypeAttrColumns(int32_t type_index,
int32_t total_size) {
ConstructPyClassFields(obj, type_info);
return ObjectRef(details::ObjectUnsafe::ObjectPtrFromOwned<Object>(obj));
});
- TVMFFIByteArray attr_name = refl::AsByteArray(refl::type_attr::kNew);
+ TVMFFIByteArray attr_name = refl::AsByteArray(type_attr::kNew.name);
TVMFFIAny attr_value = AnyView(new_fn).CopyToTVMFFIAny();
TVM_FFI_CHECK_SAFE_CALL(TVMFFITypeRegisterAttr(type_index, &attr_name,
&attr_value));
// Step 3. Register `__ffi_shallow_copy__`
@@ -2197,7 +2190,7 @@ void PyClassRegisterTypeAttrColumns(int32_t type_index,
int32_t total_size) {
Object* obj = reinterpret_cast<Object*>(obj_ptr);
return
ObjectRef(details::ObjectUnsafe::ObjectPtrFromOwned<Object>(obj));
});
- TVMFFIByteArray copy_attr_name =
refl::AsByteArray(refl::type_attr::kShallowCopy);
+ TVMFFIByteArray copy_attr_name =
refl::AsByteArray(type_attr::kShallowCopy.name);
TVMFFIAny copy_attr_value = AnyView(copy_fn).CopyToTVMFFIAny();
TVM_FFI_CHECK_SAFE_CALL(TVMFFITypeRegisterAttr(type_index, ©_attr_name,
©_attr_value));
}
@@ -2209,7 +2202,7 @@ void PyClassRegisterTypeAttrColumns(int32_t type_index,
int32_t total_size) {
// ============================================================================
Any DeepCopy(const Any& value) {
- static refl::TypeAttrColumn column(refl::type_attr::kShallowCopy);
+ static refl::TypeAttrColumn column(type_attr::kShallowCopy);
ObjectDeepCopier copier(&column);
return copier.Run(value);
}
@@ -2263,13 +2256,13 @@ bool RecursiveGe(const Any& lhs, const Any& rhs) {
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
- refl::EnsureTypeAttrColumn(refl::type_attr::kShallowCopy);
- refl::EnsureTypeAttrColumn(refl::type_attr::kRepr);
- refl::EnsureTypeAttrColumn(refl::type_attr::kHash);
- refl::EnsureTypeAttrColumn(refl::type_attr::kEq);
- refl::EnsureTypeAttrColumn(refl::type_attr::kCompare);
- refl::EnsureTypeAttrColumn(refl::type_attr::kNew);
- refl::EnsureTypeAttrColumn(refl::type_attr::kInit);
+ refl::EnsureTypeAttrColumn(type_attr::kShallowCopy);
+ refl::EnsureTypeAttrColumn(type_attr::kRepr);
+ refl::EnsureTypeAttrColumn(type_attr::kHash);
+ refl::EnsureTypeAttrColumn(type_attr::kEq);
+ refl::EnsureTypeAttrColumn(type_attr::kCompare);
+ refl::EnsureTypeAttrColumn(type_attr::kNew);
+ refl::EnsureTypeAttrColumn(type_attr::kInit);
refl::GlobalDef()
.def("ffi._RegisterFFIInit", RegisterFFIInit)
.def("ffi.MakeFieldSetter", MakeFieldSetter)
diff --git a/src/ffi/extra/reflection_extra.cc
b/src/ffi/extra/reflection_extra.cc
index 51efcd38..28f86628 100644
--- a/src/ffi/extra/reflection_extra.cc
+++ b/src/ffi/extra/reflection_extra.cc
@@ -21,6 +21,7 @@
*
* \brief Extra reflection registrations. *
*/
+#include <tvm/ffi/extra/dataclass.h>
#include <tvm/ffi/extra/structural_key.h>
#include <tvm/ffi/reflection/access_path.h>
#include <tvm/ffi/reflection/accessor.h>
@@ -117,8 +118,8 @@ inline void AccessStepRegisterReflection() {
// kArrayItemMissing -> "[<missing:index>]"
// kMapItemMissing -> "[<missing:<repr(key)>>]"
refl::TypeAttrDef<AccessStepObj>().def(
- refl::type_attr::kRepr,
- [](const AccessStep& step, const ffi::Function& fn_repr) -> ffi::String {
+ type_attr::kRepr.name, [](const Object* self, const ffi::Function&
fn_repr) -> ffi::String {
+ const auto* step = AnyView(self).as_or_throw<const AccessStepObj*>();
std::ostringstream os;
switch (step->kind) {
case AccessKind::kAttr:
@@ -168,8 +169,8 @@ inline void AccessPathRegisterReflection() {
// Register __ffi_repr__ for AccessPath: flatten via ToSteps() and walk the
resulting vector.
// Root (depth == 0) emits "<root>"; non-root nodes concatenate each step's
fragment.
refl::TypeAttrDef<AccessPathObj>().def(
- refl::type_attr::kRepr,
- [](const AccessPath& path, const ffi::Function& fn_repr) -> ffi::String {
+ type_attr::kRepr.name, [](const Object* self, const ffi::Function&
fn_repr) -> ffi::String {
+ const auto* path = AnyView(self).as_or_throw<const AccessPathObj*>();
Array<AccessStep> steps = path->ToSteps();
std::ostringstream os;
os << "<root>";
@@ -207,10 +208,8 @@ inline void StructuralKeyRegisterReflection() {
.def_ro("key", &StructuralKeyObj::key)
.def_ro("hash_i64", &StructuralKeyObj::hash_i64);
refl::TypeAttrDef<StructuralKeyObj>()
- .attr(type_attr::kAnyHash,
- NativeFunction<int64_t(const
Any&)>::FromNative<&StructuralKeyHash>())
- .attr(type_attr::kAnyEqual,
- NativeFunction<bool(const Any&, const
Any&)>::FromNative<&StructuralKeyEqual>());
+ .native_fn_attr<&StructuralKeyHash>(type_attr::kAnyHash)
+ .native_fn_attr<&StructuralKeyEqual>(type_attr::kAnyEqual);
}
TVM_FFI_STATIC_INIT_BLOCK() {
diff --git a/src/ffi/extra/serialization.cc b/src/ffi/extra/serialization.cc
index 046ab455..076ffb13 100644
--- a/src/ffi/extra/serialization.cc
+++ b/src/ffi/extra/serialization.cc
@@ -201,9 +201,9 @@ class ObjectGraphSerializer {
if (obj->IsInstance<IntEnumObj>()) return static_cast<const
EnumObj*>(obj)->_int_index;
if (obj->IsInstance<EnumObj>()) return static_cast<const
EnumObj*>(obj)->_str_index;
static reflection::TypeAttrColumn data_to_json =
- reflection::TypeAttrColumn(reflection::type_attr::kDataToJson);
- if (data_to_json[value.type_index()] != nullptr) {
- return data_to_json[value.type_index()].cast<Function>()(value);
+ reflection::TypeAttrColumn(type_attr::kDataToJson);
+ if (auto hook =
type_attr::kDataToJson.ValueFrom(data_to_json[value.type_index()])) {
+ return (*hook)(value);
}
// NOTE: invariant: lhs and rhs are already the same type
const TVMFFITypeInfo* type_info = TVMFFIGetTypeInfo(value.type_index());
@@ -394,9 +394,9 @@ class ObjectGraphDeserializer {
return EnumObj::_GetByStrIndex(type_index, data.cast<String>());
}
static reflection::TypeAttrColumn data_from_json =
- reflection::TypeAttrColumn(reflection::type_attr::kDataFromJson);
- if (data_from_json[type_index] != nullptr) {
- return data_from_json[type_index].cast<Function>()(data);
+ reflection::TypeAttrColumn(type_attr::kDataFromJson);
+ if (auto hook =
type_attr::kDataFromJson.ValueFrom(data_from_json[type_index])) {
+ return (*hook)(data);
}
// otherwise, we go over the fields and create the data.
const TVMFFITypeInfo* type_info = TVMFFIGetTypeInfo(type_index);
@@ -489,8 +489,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.def("ffi.ToJSONGraphString", ToJSONGraphString)
.def("ffi.FromJSONGraph", FromJSONGraph)
.def("ffi.FromJSONGraphString", FromJSONGraphString);
- refl::EnsureTypeAttrColumn(refl::type_attr::kDataToJson);
- refl::EnsureTypeAttrColumn(refl::type_attr::kDataFromJson);
+ refl::EnsureTypeAttrColumn(type_attr::kDataToJson);
+ refl::EnsureTypeAttrColumn(type_attr::kDataFromJson);
}
} // namespace ffi
diff --git a/src/ffi/extra/structural_equal.cc
b/src/ffi/extra/structural_equal.cc
index 21934b39..d61ab456 100644
--- a/src/ffi/extra/structural_equal.cc
+++ b/src/ffi/extra/structural_equal.cc
@@ -213,10 +213,11 @@ class StructEqualHandler {
// clamp — the caller (CompareObject) handles that for the FreeVar case.
bool CompareFields(const ObjectRef& lhs, const ObjectRef& rhs, const
TVMFFITypeInfo* type_info) {
static reflection::TypeAttrColumn custom_s_equal =
- reflection::TypeAttrColumn(reflection::type_attr::kSEqual);
+ reflection::TypeAttrColumn(type_attr::kSEqual);
bool success = true;
- if (custom_s_equal[type_info->type_index] == nullptr) {
+ auto hook =
type_attr::kSEqual.ValueFrom(custom_s_equal[type_info->type_index]);
+ if (!hook) {
// We recursively compare the fields the object
reflection::ForEachFieldInfoWithEarlyStop(type_info, [&](const
TVMFFIFieldInfo* field_info) {
// skip fields that are marked as structural eq hash ignore
@@ -291,9 +292,7 @@ class StructEqualHandler {
return sub_success;
});
}
- success = custom_s_equal[type_info->type_index]
- .cast<ffi::Function>()(lhs, rhs, s_equal_callback_)
- .cast<bool>();
+ success = (*hook)(lhs.get(), rhs.get(), s_equal_callback_);
}
return success;
}
@@ -456,7 +455,7 @@ class StructEqualHandler {
std::vector<reflection::AccessStep>* mismatch_lhs_reverse_path_ = nullptr;
std::vector<reflection::AccessStep>* mismatch_rhs_reverse_path_ = nullptr;
// lazily initialize custom equal function
- ffi::Function s_equal_callback_ = nullptr;
+ ffi::TypedFunction<bool(AnyView, AnyView, int, AnyView)> s_equal_callback_ =
nullptr;
// map from lhs to rhs
std::unordered_map<ObjectRef, ObjectRef, ObjectPtrHash, ObjectPtrEqual>
equal_map_lhs_;
// map from rhs to lhs
@@ -499,7 +498,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.def("ffi.StructuralEqual", StructuralEqual::Equal)
.def("ffi.GetFirstStructuralMismatch",
StructuralEqual::GetFirstMismatch);
// ensure the type attribute column is presented in the system even if it is
empty.
- refl::EnsureTypeAttrColumn(refl::type_attr::kSEqual);
+ refl::EnsureTypeAttrColumn(type_attr::kSEqual);
}
} // namespace ffi
diff --git a/src/ffi/extra/structural_hash.cc b/src/ffi/extra/structural_hash.cc
index 34d1c0b5..ced299d3 100644
--- a/src/ffi/extra/structural_hash.cc
+++ b/src/ffi/extra/structural_hash.cc
@@ -179,10 +179,10 @@ class StructuralHashHandler {
// callback). Does not touch the FreeVar def-region clamp — that lives
// inline in HashObject's FreeVar branch, which wraps this helper.
uint64_t HashFields(const ObjectRef& obj, const TVMFFITypeInfo* type_info,
uint64_t init_hash) {
- static reflection::TypeAttrColumn custom_s_hash =
- reflection::TypeAttrColumn(reflection::type_attr::kSHash);
+ static reflection::TypeAttrColumn custom_s_hash =
reflection::TypeAttrColumn(type_attr::kSHash);
- if (custom_s_hash[type_info->type_index] == nullptr) {
+ auto hook =
type_attr::kSHash.ValueFrom(custom_s_hash[type_info->type_index]);
+ if (!hook) {
// go over the content and hash the fields
reflection::ForEachFieldInfo(type_info, [&](const TVMFFIFieldInfo*
field_info) {
// skip fields that are marked as structural eq hash ignore
@@ -233,9 +233,8 @@ class StructuralHashHandler {
return
static_cast<int64_t>(details::StableHashCombine(inner_init, hv));
});
}
- init_hash = custom_s_hash[type_info->type_index]
- .cast<ffi::Function>()(obj,
static_cast<int64_t>(init_hash), s_hash_callback_)
- .cast<uint64_t>();
+ init_hash = static_cast<uint64_t>(
+ (*hook)(obj.get(), static_cast<int64_t>(init_hash),
s_hash_callback_));
}
return init_hash;
}
@@ -382,7 +381,7 @@ class StructuralHashHandler {
// graph node counter.
uint32_t graph_node_counter_{0};
// lazily initialize custom hash function
- ffi::Function s_hash_callback_ = nullptr;
+ ffi::TypedFunction<int64_t(AnyView, int64_t, int)> s_hash_callback_ =
nullptr;
// map from lhs to rhs
std::unordered_map<ObjectRef, uint64_t, ObjectPtrHash, ObjectPtrEqual>
hash_memo_;
};
@@ -405,7 +404,7 @@ static int64_t FFIStructuralHash(const Any& value, bool
map_free_vars, bool skip
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
refl::GlobalDef().def("ffi.StructuralHash", FFIStructuralHash);
- refl::EnsureTypeAttrColumn(refl::type_attr::kSHash);
+ refl::EnsureTypeAttrColumn(type_attr::kSHash);
}
} // namespace ffi
diff --git a/src/ffi/extra/structural_mutate.cc
b/src/ffi/extra/structural_mutate.cc
index f3361bc4..56bc3b18 100644
--- a/src/ffi/extra/structural_mutate.cc
+++ b/src/ffi/extra/structural_mutate.cc
@@ -433,43 +433,29 @@ TVM_FFI_STATIC_INIT_BLOCK() {
const Array<Tuple<int32_t, Function, bool>>& callbacks) -> Any {
return details::StructuralMutateExpected(std::move(root),
callbacks).value();
});
- refl::EnsureTypeAttrColumn(refl::type_attr::kStructuralMutate);
- refl::EnsureTypeAttrColumn(refl::type_attr::kStructuralMaybeInplaceMutate);
+ refl::EnsureTypeAttrColumn(type_attr::kStructuralMutate);
+ refl::EnsureTypeAttrColumn(type_attr::kStructuralMaybeInplaceMutate);
refl::TypeAttrDef<details::BigIntObj>()
- .attr(refl::type_attr::kStructuralMutate,
- FStructuralMutate::FromNative<&details::MutateImmutableLeaf>())
- .attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- FStructuralMutate::FromNative<&details::MutateImmutableLeaf>());
+
.native_fn_attr<&details::MutateImmutableLeaf>(type_attr::kStructuralMutate)
+
.native_fn_attr<&details::MutateImmutableLeaf>(type_attr::kStructuralMaybeInplaceMutate);
refl::TypeAttrDef<details::StringObj>()
- .attr(refl::type_attr::kStructuralMutate,
- FStructuralMutate::FromNative<&details::MutateImmutableLeaf>())
- .attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- FStructuralMutate::FromNative<&details::MutateImmutableLeaf>());
+
.native_fn_attr<&details::MutateImmutableLeaf>(type_attr::kStructuralMutate)
+
.native_fn_attr<&details::MutateImmutableLeaf>(type_attr::kStructuralMaybeInplaceMutate);
refl::TypeAttrDef<details::BytesObj>()
- .attr(refl::type_attr::kStructuralMutate,
- FStructuralMutate::FromNative<&details::MutateImmutableLeaf>())
- .attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- FStructuralMutate::FromNative<&details::MutateImmutableLeaf>());
+
.native_fn_attr<&details::MutateImmutableLeaf>(type_attr::kStructuralMutate)
+
.native_fn_attr<&details::MutateImmutableLeaf>(type_attr::kStructuralMaybeInplaceMutate);
refl::TypeAttrDef<ArrayObj>()
- .attr(refl::type_attr::kStructuralMutate,
- FStructuralMutate::FromNative<&details::MutateArray>())
- .attr(refl::type_attr::kStructuralMaybeInplaceMutate,
-
FStructuralMutate::FromNative<&details::MaybeInplaceMutateArray>());
+ .native_fn_attr<&details::MutateArray>(type_attr::kStructuralMutate)
+
.native_fn_attr<&details::MaybeInplaceMutateArray>(type_attr::kStructuralMaybeInplaceMutate);
refl::TypeAttrDef<ListObj>()
- .attr(refl::type_attr::kStructuralMutate,
- FStructuralMutate::FromNative<&details::MutateList>())
- .attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- FStructuralMutate::FromNative<&details::MaybeInplaceMutateList>());
+ .native_fn_attr<&details::MutateList>(type_attr::kStructuralMutate)
+
.native_fn_attr<&details::MaybeInplaceMutateList>(type_attr::kStructuralMaybeInplaceMutate);
refl::TypeAttrDef<MapObj>()
- .attr(refl::type_attr::kStructuralMutate,
- FStructuralMutate::FromNative<&details::MutateMap>())
- .attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- FStructuralMutate::FromNative<&details::MaybeInplaceMutateMap>());
+ .native_fn_attr<&details::MutateMap>(type_attr::kStructuralMutate)
+
.native_fn_attr<&details::MaybeInplaceMutateMap>(type_attr::kStructuralMaybeInplaceMutate);
refl::TypeAttrDef<DictObj>()
- .attr(refl::type_attr::kStructuralMutate,
- FStructuralMutate::FromNative<&details::MutateDict>())
- .attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- FStructuralMutate::FromNative<&details::MaybeInplaceMutateDict>());
+ .native_fn_attr<&details::MutateDict>(type_attr::kStructuralMutate)
+
.native_fn_attr<&details::MaybeInplaceMutateDict>(type_attr::kStructuralMaybeInplaceMutate);
}
} // namespace ffi
diff --git a/src/ffi/extra/structural_visit.cc
b/src/ffi/extra/structural_visit.cc
index 729a4a7c..0d2124e9 100644
--- a/src/ffi/extra/structural_visit.cc
+++ b/src/ffi/extra/structural_visit.cc
@@ -169,19 +169,15 @@ TVM_FFI_STATIC_INIT_BLOCK() {
const Array<Tuple<int32_t, Function>>& callbacks) ->
Optional<VisitInterrupt> {
return details::StructuralVisitExpected(root, callbacks).value();
});
- refl::EnsureTypeAttrColumn(refl::type_attr::kStructuralVisit);
- refl::TypeAttrDef<ArrayObj>().attr(
- refl::type_attr::kStructuralVisit,
- FStructuralVisit::FromNative<&details::VisitSeqContainer<ArrayObj>>());
- refl::TypeAttrDef<ListObj>().attr(
- refl::type_attr::kStructuralVisit,
- FStructuralVisit::FromNative<&details::VisitSeqContainer<ListObj>>());
- refl::TypeAttrDef<MapObj>().attr(
- refl::type_attr::kStructuralVisit,
- FStructuralVisit::FromNative<&details::VisitMapContainer<MapObj>>());
- refl::TypeAttrDef<DictObj>().attr(
- refl::type_attr::kStructuralVisit,
- FStructuralVisit::FromNative<&details::VisitMapContainer<DictObj>>());
+ refl::EnsureTypeAttrColumn(type_attr::kStructuralVisit);
+
refl::TypeAttrDef<ArrayObj>().native_fn_attr<&details::VisitSeqContainer<ArrayObj>>(
+ type_attr::kStructuralVisit);
+
refl::TypeAttrDef<ListObj>().native_fn_attr<&details::VisitSeqContainer<ListObj>>(
+ type_attr::kStructuralVisit);
+
refl::TypeAttrDef<MapObj>().native_fn_attr<&details::VisitMapContainer<MapObj>>(
+ type_attr::kStructuralVisit);
+
refl::TypeAttrDef<DictObj>().native_fn_attr<&details::VisitMapContainer<DictObj>>(
+ type_attr::kStructuralVisit);
}
} // namespace ffi
diff --git a/src/ffi/object.cc b/src/ffi/object.cc
index 324fe1fd..004e5c5d 100644
--- a/src/ffi/object.cc
+++ b/src/ffi/object.cc
@@ -686,7 +686,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.def_ro("value", &ffi::EnumObj::_str_index, "Canonical string index.",
refl::init(false),
refl::AttachFieldFlag::SEqHashIgnore())
.def_convert<ffi::StrEnum>();
- refl::EnsureTypeAttrColumn(refl::type_attr::kEnumState);
+ refl::EnsureTypeAttrColumn(ffi::type_attr::kEnumState);
refl::GlobalDef()
.def_method("ffi.GetRegisteredTypeKeys",
[]() -> ffi::Array<ffi::String> {
diff --git a/src/ffi/testing/testing.cc b/src/ffi/testing/testing.cc
index 319c8347..31f1b9a4 100644
--- a/src/ffi/testing/testing.cc
+++ b/src/ffi/testing/testing.cc
@@ -30,6 +30,7 @@
#include <tvm/ffi/dtype.h>
#include <tvm/ffi/enum.h>
#include <tvm/ffi/extra/c_env_api.h>
+#include <tvm/ffi/extra/dataclass.h>
#include <tvm/ffi/extra/structural_visit.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/optional.h>
@@ -113,7 +114,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
// silence clang-tidy's bugprone-unused-raii since the RAII finalisation is
the point.
refl::ObjectDef<TestEnumVariantObj>(refl::init(false)); //
NOLINT(bugprone-unused-raii)
refl::TypeAttrDef<TestEnumVariantObj>().def_convert<TestEnumVariant>();
- refl::EnumDef<TestEnumVariantObj>("Alpha").set_attr("code", int64_t{10});
+ constexpr refl::AttrKey<int64_t> code{"code"};
+ refl::EnumDef<TestEnumVariantObj>("Alpha").set_attr(code, int64_t{10});
refl::EnumDef<TestEnumVariantObj>("Beta").set_attr("code", int64_t{20});
}
@@ -176,16 +178,17 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.def_ro("pattern", &TestVisitRegionsObj::pattern,
refl::AttachFieldFlag::SEqHashDefPattern())
.def_ro("ordinary", &TestVisitRegionsObj::ordinary);
refl::ObjectDef<TestVisitHookObj>().def_ro("child",
&TestVisitHookObj::child);
- refl::EnsureTypeAttrColumn(refl::type_attr::kStructuralVisit);
- refl::TypeAttrDef<TestVisitHookObj>().def(
- refl::type_attr::kStructuralVisit,
- [](StructuralVisitor visitor, const TestVisitHookObj* self) {
+ refl::EnsureTypeAttrColumn(type_attr::kStructuralVisit);
+ TypedFunction<Optional<VisitInterrupt>(StructuralVisitorObj*, AnyView)> hook
=
+ [](StructuralVisitorObj* visitor, AnyView value) {
+ const auto* self = value.as_or_throw<const TestVisitHookObj*>();
auto result = visitor->WithDefRegionKind(kTVMFFIDefRegionKindSimple,
[&] { return
visitor->Visit(self->child); });
if (result.has_value()) return result;
// This sibling observes the region restored after the scoped child
visit.
return visitor->Visit(int64_t{4});
- });
+ };
+ refl::TypeAttrDef<TestVisitHookObj>().attr(type_attr::kStructuralVisit,
FStructuralVisit(hook));
refl::GlobalDef().def("testing.make_visit_region_graph", [](bool with_hook)
-> Any {
ObjectRef fields(make_object<TestVisitRegionsObj>());
if (with_hook) return ObjectRef(make_object<TestVisitHookObj>(fields));
@@ -1021,27 +1024,26 @@ TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
// Register __ffi_hash__ for TestCustomHash: only hashes `key`, ignores
`label`.
refl::TypeAttrDef<TestCustomHashObj>().def(
- refl::type_attr::kHash, [](const Object* self, const Function& fn_hash)
-> int64_t {
+ type_attr::kHash.name, [](const Object* self, const Function& fn_hash)
-> int64_t {
auto* obj = static_cast<const TestCustomHashObj*>(self);
return fn_hash(AnyView(obj->key)).cast<int64_t>();
});
// Register __ffi_hash__ for TestCustomCompare: only hashes `key`,
consistent with eq/compare.
refl::TypeAttrDef<TestCustomCompareObj>().def(
- refl::type_attr::kHash, [](const Object* self, const Function& fn_hash)
-> int64_t {
+ type_attr::kHash.name, [](const Object* self, const Function& fn_hash)
-> int64_t {
auto* obj = static_cast<const TestCustomCompareObj*>(self);
return fn_hash(AnyView(obj->key)).cast<int64_t>();
});
// Register __ffi_eq__ for TestCustomCompare: compares only `key`.
refl::TypeAttrDef<TestCustomCompareObj>().def(
- refl::type_attr::kEq,
- [](const Object* lhs, const Object* rhs, const Function& fn_eq) -> bool {
+ type_attr::kEq.name, [](const Object* lhs, const Object* rhs, const
Function& fn_eq) -> bool {
auto* a = static_cast<const TestCustomCompareObj*>(lhs);
auto* b = static_cast<const TestCustomCompareObj*>(rhs);
return fn_eq(AnyView(a->key), AnyView(b->key)).cast<bool>();
});
// Register __ffi_compare__ for TestCustomCompare: three-way ordering on
`key`.
refl::TypeAttrDef<TestCustomCompareObj>().def(
- refl::type_attr::kCompare,
+ type_attr::kCompare.name,
[](const Object* lhs, const Object* rhs, const Function& fn_cmp) ->
int32_t {
auto* a = static_cast<const TestCustomCompareObj*>(lhs);
auto* b = static_cast<const TestCustomCompareObj*>(rhs);
@@ -1050,8 +1052,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
// Register __ffi_eq__ for TestEqWithoutHash: deliberately no __ffi_hash__.
// This exercises the RecursiveHash guard that rejects eq-without-hash types.
refl::TypeAttrDef<TestEqWithoutHashObj>().def(
- refl::type_attr::kEq,
- [](const Object* lhs, const Object* rhs, const Function& fn_eq) -> bool {
+ type_attr::kEq.name, [](const Object* lhs, const Object* rhs, const
Function& fn_eq) -> bool {
auto* a = static_cast<const TestEqWithoutHashObj*>(lhs);
auto* b = static_cast<const TestEqWithoutHashObj*>(rhs);
return fn_eq(AnyView(a->key), AnyView(b->key)).cast<bool>();
diff --git a/tests/cpp/extra/test_serialization.cc
b/tests/cpp/extra/test_serialization.cc
index 4cf086ba..e130e1a2 100644
--- a/tests/cpp/extra/test_serialization.cc
+++ b/tests/cpp/extra/test_serialization.cc
@@ -74,11 +74,15 @@ class SerializationStrEnum : public StrEnum {
SerializationStrEnumObj);
};
+constexpr reflection::AttrKey<reflection::NativeFunctionView<int(int)>>
kEnumIncrement{
+ "test.enum.increment"};
+int EnumIncrement(int value) { return value + 1; }
+
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<SerializationEnumObj>(refl::init(false)); //
NOLINT(bugprone-unused-raii)
refl::TypeAttrDef<SerializationEnumObj>().def_convert<SerializationEnum>();
- refl::EnumDef<SerializationEnumObj>("Alpha"); //
NOLINT(bugprone-unused-raii)
+
refl::EnumDef<SerializationEnumObj>("Alpha").set_native_fn_attr<&EnumIncrement>(kEnumIncrement);
refl::ObjectDef<SerializationIntEnumObj>(refl::init(false)).def_convert<SerializationIntEnum>();
refl::EnumDef<SerializationIntEnumObj>("Seven", 7); //
NOLINT(bugprone-unused-raii)
@@ -100,6 +104,13 @@ TEST(Serialization, EnumSingletonIdentity) {
EXPECT_TRUE(StructuralEqual()(ToJSONGraph(original), expected));
Enum restored = FromJSONGraph(expected).cast<Enum>();
EXPECT_TRUE(restored.same_as(original));
+ reflection::TypeAttrColumn states(type_attr::kEnumState);
+ EnumState state = states[original->type_index()].as_or_throw<EnumState>();
+ auto hook =
kEnumIncrement.ValueFrom(state->attrs.at(kEnumIncrement.name).at(restored));
+ EXPECT_TRUE(hook.has_value());
+ if (hook) {
+ EXPECT_EQ((*hook)(4), 5);
+ }
}
TEST(Serialization, PayloadEnumSingletonIdentity) {
diff --git a/tests/cpp/extra/test_structural_mutate.cc
b/tests/cpp/extra/test_structural_mutate.cc
index 348f94be..6973eb94 100644
--- a/tests/cpp/extra/test_structural_mutate.cc
+++ b/tests/cpp/extra/test_structural_mutate.cc
@@ -298,13 +298,12 @@ class TNestedMapHookObj : public Object {
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<TNestedMapHookObj>().def_rw("field",
&TNestedMapHookObj::field);
- refl::EnsureTypeAttrColumn(refl::type_attr::kStructuralMutate);
- refl::EnsureTypeAttrColumn(refl::type_attr::kStructuralMaybeInplaceMutate);
+ refl::EnsureTypeAttrColumn(type_attr::kStructuralMutate);
+ refl::EnsureTypeAttrColumn(type_attr::kStructuralMaybeInplaceMutate);
refl::TypeAttrDef<TNestedMapHookObj>()
- .attr(refl::type_attr::kStructuralMutate,
-
FStructuralMutate::FromNative<&TNestedMapHookObj::StructuralMutate>())
- .attr(refl::type_attr::kStructuralMaybeInplaceMutate,
-
FStructuralMutate::FromNative<&TNestedMapHookObj::MaybeInplaceMutate>());
+
.native_fn_attr<&TNestedMapHookObj::StructuralMutate>(type_attr::kStructuralMutate)
+ .native_fn_attr<&TNestedMapHookObj::MaybeInplaceMutate>(
+ type_attr::kStructuralMaybeInplaceMutate);
}
static constexpr bool _type_mutable = true;
diff --git a/tests/cpp/test_reflection.cc b/tests/cpp/test_reflection.cc
index 2ea12f98..6fd600c9 100644
--- a/tests/cpp/test_reflection.cc
+++ b/tests/cpp/test_reflection.cc
@@ -25,6 +25,7 @@
#include <tvm/ffi/reflection/access_path.h>
#include <tvm/ffi/reflection/accessor.h>
#include <tvm/ffi/reflection/creator.h>
+#include <tvm/ffi/reflection/enum_def.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ffi/string.h>
@@ -34,7 +35,7 @@
/*! \brief Look up __ffi_init__ from the TypeAttrColumn (not the method
table). */
static tvm::ffi::Function GetInitAttr(const char* type_key) {
- static tvm::ffi::reflection::TypeAttrColumn
col(tvm::ffi::reflection::type_attr::kInit);
+ static tvm::ffi::reflection::TypeAttrColumn col(tvm::ffi::type_attr::kInit);
return col[tvm::ffi::TypeKeyToIndex(type_key)].cast<tvm::ffi::Function>();
}
@@ -43,6 +44,8 @@ namespace {
using namespace tvm::ffi;
using namespace tvm::ffi::testing;
+constexpr reflection::AttrKey<bool> kFlag{"test.attr.typed.flag"};
+
struct TestObjA : public Object {
int64_t x;
int64_t y;
@@ -94,6 +97,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.def_type_attr("test.attr.object.int", int64_t{7})
.def_type_attr("test.attr.object.bool", true)
.def_type_attr("test.attr.object.func", [](int64_t value) -> int64_t {
return value + 1; });
+ refl::TypeAttrDef<TestObjA>().attr(kFlag, false);
refl::ObjectDef<TestObjADerived>()
.def(refl::init<int64_t, int64_t, int64_t>())
.def_ro("z", &TestObjADerived::z);
@@ -328,6 +332,23 @@ TEST(Reflection, TypeAttrColumn) {
EXPECT_EQ(size_attr[TIntObj::RuntimeTypeIndex()].cast<int>(),
sizeof(TIntObj));
}
+TEST(Reflection, AttrKey) {
+ namespace refl = reflection;
+ refl::TypeAttrColumn column(kFlag);
+ auto flag = kFlag.ValueFrom(column[TestObjA::RuntimeTypeIndex()]);
+ EXPECT_EQ(flag, std::optional<bool>(false));
+ EXPECT_FALSE(kFlag.ValueFrom(column[-1]).has_value());
+ EXPECT_THROW(kFlag.ValueFrom(AnyView(1)), Error);
+
+ // Raw names retain dynamic value storage alongside typed registration.
+ constexpr refl::AttrKey<int64_t> count{"test.attr.typed.count"};
+ refl::TypeAttrDef<TestObjA>().attr(count, int64_t{7});
+ refl::TypeAttrDef<TestObjADerived>().attr(count.name, Any("dynamic"));
+ refl::TypeAttrColumn counts(count);
+ EXPECT_EQ(count.ValueFrom(counts[TestObjA::RuntimeTypeIndex()]), 7);
+ EXPECT_THROW(count.ValueFrom(counts[TestObjADerived::RuntimeTypeIndex()]),
Error);
+}
+
TEST(Reflection, TypeAttrColumnBeginIndex) {
// Get the column and verify begin_index
TVMFFIByteArray attr_name = {"test.size",
std::char_traits<char>::length("test.size")};
@@ -377,7 +398,7 @@ TEST(Reflection, TypeAttrDefDirectValues) {
}
TEST(Reflection, ObjectDefPublishesTypeFinal) {
- reflection::TypeAttrColumn final_attr(reflection::type_attr::kTypeFinal);
+ reflection::TypeAttrColumn final_attr(type_attr::kTypeFinal);
EXPECT_FALSE(final_attr[TestObjA::RuntimeTypeIndex()].cast<bool>());
EXPECT_TRUE(final_attr[TestObjADerived::RuntimeTypeIndex()].cast<bool>());
diff --git a/tests/cpp/testing_object.h b/tests/cpp/testing_object.h
index ad01c4a6..d301e309 100644
--- a/tests/cpp/testing_object.h
+++ b/tests/cpp/testing_object.h
@@ -25,6 +25,9 @@
#include <tvm/ffi/container/map.h>
#include <tvm/ffi/device.h>
#include <tvm/ffi/dtype.h>
+#include <tvm/ffi/extra/serialization.h>
+#include <tvm/ffi/extra/structural_equal.h>
+#include <tvm/ffi/extra/structural_hash.h>
#include <tvm/ffi/extra/structural_mutate.h>
#include <tvm/ffi/extra/structural_visit.h>
#include <tvm/ffi/memory.h>
@@ -100,18 +103,17 @@ inline void TIntObj::RegisterReflection() {
refl::TypeAttrDef<TIntObj>()
.def("test.GetValue", &TIntObj::GetValue)
.attr("test.size", sizeof(TIntObj))
- .attr(refl::type_attr::kAnyHash,
- refl::NativeFunction<int64_t(const
Any&)>::FromNative<&TInt::CustomAnyHash>())
- .attr(
- refl::type_attr::kAnyEqual,
- refl::NativeFunction<bool(const Any&, const
Any&)>::FromNative<&TInt::CustomAnyEqual>());
+ .native_fn_attr<&TInt::CustomAnyHash>(type_attr::kAnyHash)
+ .native_fn_attr<&TInt::CustomAnyEqual>(type_attr::kAnyEqual);
// custom json serialization
refl::TypeAttrDef<TIntObj>()
- .def(refl::type_attr::kDataToJson,
- [](const TIntObj* self) -> Map<String, Any> {
+ .def(type_attr::kDataToJson.name,
+ [](AnyView value) -> Any {
+ const auto* self = value.as_or_throw<const TIntObj*>();
return Map<String, Any>{{"value", self->value}};
})
- .def(refl::type_attr::kDataFromJson, [](Map<String, Any> json_obj) ->
TInt {
+ .def(type_attr::kDataFromJson.name, [](AnyView value) -> Any {
+ auto json_obj = value.as_or_throw<Map<String, Any>>();
return TInt(json_obj["value"].cast<int64_t>());
});
}
@@ -152,8 +154,8 @@ inline void TFloatObj::RegisterReflection() {
.def("sub", [](const TFloatObj* self, double other) -> double { return
self->value - other; })
.def("add", &TFloatObj::Add, "add method");
refl::TypeAttrDef<TFloatObj>()
- .def(refl::type_attr::kAnyHash, &TFloat::CustomAnyHash)
- .def(refl::type_attr::kAnyEqual, &TFloat::CustomAnyEqual);
+ .def(type_attr::kAnyHash.name, &TFloat::CustomAnyHash)
+ .def(type_attr::kAnyEqual.name, &TFloat::CustomAnyEqual);
}
class TPrimExprObj : public Object {
@@ -276,10 +278,9 @@ class TMutatePairObj : public Object {
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,
-
FStructuralMutate::FromNative<&TMutatePairObj::StructuralMutateHook>());
+ refl::EnsureTypeAttrColumn(type_attr::kStructuralMutate);
+
refl::TypeAttrDef<TMutatePairObj>().native_fn_attr<&TMutatePairObj::StructuralMutateHook>(
+ type_attr::kStructuralMutate);
}
static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind =
kTVMFFISEqHashKindTreeNode;
@@ -432,10 +433,9 @@ class TFuncObj : public Object {
.def_ro("params", &TFuncObj::params,
refl::AttachFieldFlag::SEqHashDefPattern())
.def_ro("body", &TFuncObj::body)
.def_ro("comment", &TFuncObj::comment,
refl::AttachFieldFlag::SEqHashIgnore());
- refl::EnsureTypeAttrColumn(refl::type_attr::kStructuralVisit);
- refl::TypeAttrDef<TFuncObj>().attr(
- refl::type_attr::kStructuralVisit,
- FStructuralVisit::FromNative<&TFuncObj::StructuralVisitHook>());
+ refl::EnsureTypeAttrColumn(type_attr::kStructuralVisit);
+
refl::TypeAttrDef<TFuncObj>().native_fn_attr<&TFuncObj::StructuralVisitHook>(
+ type_attr::kStructuralVisit);
}
static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind =
kTVMFFISEqHashKindTreeNode;
@@ -461,7 +461,7 @@ class TCustomFuncObj : public Object {
: params(params), body(body), comment(comment) {}
bool SEqual(const TCustomFuncObj* other,
- ffi::TypedFunction<bool(AnyView, AnyView, bool, AnyView)> cmp)
const {
+ const ffi::TypedFunction<bool(AnyView, AnyView, int, AnyView)>&
cmp) const {
if (!cmp(params, other->params, true, "params")) {
return false;
}
@@ -471,7 +471,8 @@ class TCustomFuncObj : public Object {
return true;
}
- int64_t SHash(int64_t init_hash, ffi::TypedFunction<int64_t(AnyView,
int64_t, bool)> hash) const {
+ int64_t SHash(int64_t init_hash,
+ const ffi::TypedFunction<int64_t(AnyView, int64_t, int)>&
hash) const {
int64_t hash_value = init_hash;
hash_value = hash(params, hash_value, true);
hash_value = hash(body, hash_value, false);
@@ -485,8 +486,8 @@ class TCustomFuncObj : public Object {
.def_ro("body", &TCustomFuncObj::body)
.def_ro("comment", &TCustomFuncObj::comment);
refl::TypeAttrDef<TCustomFuncObj>()
- .def(refl::type_attr::kSEqual, &TCustomFuncObj::SEqual)
- .def(refl::type_attr::kSHash, &TCustomFuncObj::SHash);
+ .def(type_attr::kSEqual.name, &TCustomFuncObj::SEqual)
+ .def(type_attr::kSHash.name, &TCustomFuncObj::SHash);
}
static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind =
kTVMFFISEqHashKindTreeNode;
