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 f2e6f403 [CORE] Unify custom Any hash and equality hooks (#841)
f2e6f403 is described below
commit f2e6f403975e576dcf23080ad468f566a66bacc9
Author: Tianqi Chen <[email protected]>
AuthorDate: Sun Oct 11 13:57:16 2026 -0400
[CORE] Unify custom Any hash and equality hooks (#841)
Use the existing NativeFunction convention for custom Any hash and
equality hooks, retaining packed Function support and strict int64/bool
results. Keep custom dispatch inline with a direct native success path,
preserve the packed fast path, and update native registrations and hook
signatures.
---
include/tvm/ffi/any.h | 36 +++++++++++++++++++++--------------
include/tvm/ffi/reflection/accessor.h | 12 +++++++-----
src/ffi/extra/reflection_extra.cc | 7 +++++--
tests/cpp/testing_object.h | 8 ++++++--
4 files changed, 40 insertions(+), 23 deletions(-)
diff --git a/include/tvm/ffi/any.h b/include/tvm/ffi/any.h
index 6226c749..6b65b8be 100644
--- a/include/tvm/ffi/any.h
+++ b/include/tvm/ffi/any.h
@@ -756,16 +756,20 @@ struct AnyHash {
* \param src The source Any.
* \return The hash value.
*/
- static uint64_t CallCustomAnyHash(const TVMFFIAny& custom_any_hash, const
Any& src) {
- // NOTE: we explicitly use low-level ABI here since we do not want to have
dep on function.h
- // it also keeps the logic simple
+ TVM_FFI_INLINE static uint64_t CallCustomAnyHash(const TVMFFIAny&
custom_any_hash,
+ const Any& src) {
+ // Use the NativeFunction convention without depending on
native_function.h.
if (custom_any_hash.type_index == TypeIndex::kTVMFFIOpaquePtr) {
- // we allow this attribute to be a function pointer for fast path
- using FCustomAnyHashPtr = int64_t (*)(const Any&);
+ using FCustomAnyHashPtr = TVMFFIAny (*)(const Any&) noexcept;
FCustomAnyHashPtr hash_func =
reinterpret_cast<FCustomAnyHashPtr>(custom_any_hash.v_ptr);
TVM_FFI_ICHECK_NOTNULL(hash_func);
- return static_cast<uint64_t>(
- (*hash_func)(src)); // NOLINT(clang-analyzer-core.CallAndMessage)
+ TVMFFIAny result = (*hash_func)(src); //
NOLINT(clang-analyzer-core.CallAndMessage)
+ if (TVM_FFI_PREDICT_TRUE(result.type_index == TypeIndex::kTVMFFIInt)) {
+ return static_cast<uint64_t>(result.v_int64);
+ }
+ Any error = details::AnyUnsafe::MoveTVMFFIAnyRawToAny(result);
+ TVM_FFI_ICHECK_EQ(error.type_index(), TypeIndex::kTVMFFIError);
+ throw details::AnyUnsafe::MoveFromAnyAfterCheck<Error>(std::move(error));
} else {
// alternatively it can be a ffi.Function object.
TVM_FFI_ICHECK_EQ(custom_any_hash.type_index,
TypeIndex::kTVMFFIFunction);
@@ -888,16 +892,20 @@ struct AnyEqual {
* \param rhs The right-hand side Any.
* \return The equality result.
*/
- static bool CallCustomAnyEqual(const TVMFFIAny& custom_any_equal, const Any&
lhs,
- const Any& rhs) {
- // NOTE: we explicitly use low-level ABI here since we do not want to have
dep on function.h
- // it also keeps the logic simple
+ TVM_FFI_INLINE static bool CallCustomAnyEqual(const TVMFFIAny&
custom_any_equal, const Any& lhs,
+ const Any& rhs) {
+ // Use the NativeFunction convention without depending on
native_function.h.
if (custom_any_equal.type_index == TypeIndex::kTVMFFIOpaquePtr) {
- // we allow this attribute to be a function pointer for fast path
- using FCustomAnyEqualPtr = bool (*)(const Any&, const Any&);
+ using FCustomAnyEqualPtr = TVMFFIAny (*)(const Any&, const Any&)
noexcept;
FCustomAnyEqualPtr equal_func =
reinterpret_cast<FCustomAnyEqualPtr>(custom_any_equal.v_ptr);
TVM_FFI_ICHECK_NOTNULL(equal_func);
- return (*equal_func)(lhs, rhs); //
NOLINT(clang-analyzer-core.CallAndMessage)
+ TVMFFIAny result = (*equal_func)(lhs, rhs); //
NOLINT(clang-analyzer-core.CallAndMessage)
+ if (TVM_FFI_PREDICT_TRUE(result.type_index == TypeIndex::kTVMFFIBool)) {
+ return result.v_int64 != 0;
+ }
+ Any error = details::AnyUnsafe::MoveTVMFFIAnyRawToAny(result);
+ TVM_FFI_ICHECK_EQ(error.type_index(), TypeIndex::kTVMFFIError);
+ throw details::AnyUnsafe::MoveFromAnyAfterCheck<Error>(std::move(error));
} else {
// alternatively it can be a ffi.Function object.
TVM_FFI_ICHECK_EQ(custom_any_equal.type_index,
TypeIndex::kTVMFFIFunction);
diff --git a/include/tvm/ffi/reflection/accessor.h
b/include/tvm/ffi/reflection/accessor.h
index 08aa021b..18aa2f66 100644
--- a/include/tvm/ffi/reflection/accessor.h
+++ b/include/tvm/ffi/reflection/accessor.h
@@ -442,10 +442,10 @@ inline constexpr const char* kCompare = "__ffi_compare__";
* 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.
*
- * Can be either a raw C function pointer (fast path, no boxing overhead) or
- * a ``Function`` object.
+ * Register a native hook with ``NativeFunction<int64_t(const
Any&)>::FromNative<&Fn>()``
+ * or a packed ``Function`` object.
*
- * Raw pointer signature: ``int64_t (*)(const Any& src)``
+ * Native adapter signature: ``TVMFFIAny (*)(const Any& src) noexcept``
*
* Function signature: ``(Any src) -> int64``
*/
@@ -456,9 +456,11 @@ inline constexpr const char* kAnyHash = "__any_hash__";
* Unlike ``kEq`` (which operates on ``ObjectRef``), this hook operates at the
* ``Any`` level — called by ``AnyEqual`` for container key comparison.
*
- * Can be either a raw C function pointer (fast path) or a ``Function`` object.
+ * Register a native hook with
+ * ``NativeFunction<bool(const Any&, const Any&)>::FromNative<&Fn>()``
+ * or a packed ``Function`` object.
*
- * Raw pointer signature: ``bool (*)(const Any& lhs, const Any& rhs)``
+ * Native adapter signature: ``TVMFFIAny (*)(const Any& lhs, const Any& rhs)
noexcept``
*
* Function signature: ``(Any lhs, Any rhs) -> bool``
*/
diff --git a/src/ffi/extra/reflection_extra.cc
b/src/ffi/extra/reflection_extra.cc
index 143bfe96..51efcd38 100644
--- a/src/ffi/extra/reflection_extra.cc
+++ b/src/ffi/extra/reflection_extra.cc
@@ -24,6 +24,7 @@
#include <tvm/ffi/extra/structural_key.h>
#include <tvm/ffi/reflection/access_path.h>
#include <tvm/ffi/reflection/accessor.h>
+#include <tvm/ffi/reflection/native_function.h>
#include <tvm/ffi/reflection/registry.h>
#include <sstream>
@@ -206,8 +207,10 @@ inline void StructuralKeyRegisterReflection() {
.def_ro("key", &StructuralKeyObj::key)
.def_ro("hash_i64", &StructuralKeyObj::hash_i64);
refl::TypeAttrDef<StructuralKeyObj>()
- .attr(type_attr::kAnyHash, reinterpret_cast<void*>(&StructuralKeyHash))
- .attr(type_attr::kAnyEqual,
reinterpret_cast<void*>(&StructuralKeyEqual));
+ .attr(type_attr::kAnyHash,
+ NativeFunction<int64_t(const
Any&)>::FromNative<&StructuralKeyHash>())
+ .attr(type_attr::kAnyEqual,
+ NativeFunction<bool(const Any&, const
Any&)>::FromNative<&StructuralKeyEqual>());
}
TVM_FFI_STATIC_INIT_BLOCK() {
diff --git a/tests/cpp/testing_object.h b/tests/cpp/testing_object.h
index 3cd5474d..ad01c4a6 100644
--- a/tests/cpp/testing_object.h
+++ b/tests/cpp/testing_object.h
@@ -29,6 +29,7 @@
#include <tvm/ffi/extra/structural_visit.h>
#include <tvm/ffi/memory.h>
#include <tvm/ffi/object.h>
+#include <tvm/ffi/reflection/native_function.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ffi/string.h>
@@ -99,8 +100,11 @@ inline void TIntObj::RegisterReflection() {
refl::TypeAttrDef<TIntObj>()
.def("test.GetValue", &TIntObj::GetValue)
.attr("test.size", sizeof(TIntObj))
- .attr(refl::type_attr::kAnyHash,
reinterpret_cast<void*>(&TInt::CustomAnyHash))
- .attr(refl::type_attr::kAnyEqual,
reinterpret_cast<void*>(&TInt::CustomAnyEqual));
+ .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>());
// custom json serialization
refl::TypeAttrDef<TIntObj>()
.def(refl::type_attr::kDataToJson,
