This is an automated email from the ASF dual-hosted git repository.
ruihangl 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 935a5a0 [ABI] Stablize doc and metadata for functions (#79)
935a5a0 is described below
commit 935a5a074686839ae42a9bc52581232beeb5b1fc
Author: Tianqi Chen <[email protected]>
AuthorDate: Wed Oct 1 09:20:53 2025 -0400
[ABI] Stablize doc and metadata for functions (#79)
This PR stablizes the doc and metadata abi for functions.
- doc will store optional unstructured docstring about the function
- metadata will store structured json about extra information about the
function
The PR changes the ModuleObj c++ ABI, but given it is part of extra and
python is directly accessing function through packed registry, it won't
impact python side users. For python users it remains ABI compatible.
---
include/tvm/ffi/c_api.h | 11 +++++++----
include/tvm/ffi/extra/module.h | 16 ++++++++++++++++
include/tvm/ffi/reflection/registry.h | 6 +++---
python/tvm_ffi/cython/base.pxi | 4 ++--
src/ffi/extra/module.cc | 18 ++++++++++++++++++
src/ffi/function.cc | 6 +++---
src/ffi/object.cc | 4 ++--
7 files changed, 51 insertions(+), 14 deletions(-)
diff --git a/include/tvm/ffi/c_api.h b/include/tvm/ffi/c_api.h
index c688aee..619dd5b 100644
--- a/include/tvm/ffi/c_api.h
+++ b/include/tvm/ffi/c_api.h
@@ -810,8 +810,8 @@ typedef struct {
TVMFFIByteArray name;
/*! \brief The docstring about the field. */
TVMFFIByteArray doc;
- /*! \brief The type schema of the field in JSON string. */
- TVMFFIByteArray type_schema;
+ /*! \brief The structured metadata of the field in JSON string. */
+ TVMFFIByteArray metadata;
/*!
* \brief bitmask flags of the field.
*/
@@ -864,8 +864,11 @@ typedef struct {
TVMFFIByteArray name;
/*! \brief The docstring about the method. */
TVMFFIByteArray doc;
- /*! \brief Optional type schema of the method in JSON string. */
- TVMFFIByteArray type_schema;
+ // Rationale: We separate the docstring from the metadata since docstrings
+ // can be unstructured and sometimes large, while metadata can be focused
+ // on storing structured information.
+ /*! \brief Optional structured metadata of the method in JSON string. */
+ TVMFFIByteArray metadata;
/*! \brief bitmask flags of the method. */
int64_t flags;
/*!
diff --git a/include/tvm/ffi/extra/module.h b/include/tvm/ffi/extra/module.h
index fd6bf19..500a4b2 100644
--- a/include/tvm/ffi/extra/module.h
+++ b/include/tvm/ffi/extra/module.h
@@ -69,6 +69,15 @@ class TVM_FFI_EXTRA_CXX_API ModuleObj : public Object {
* \return True if the module implements the function, false otherwise.
*/
virtual bool ImplementsFunction(const String& name) { return
GetFunction(name).defined(); }
+ /*!
+ * \brief Get the docstring of the function, if available.
+ * \param name The name of the function.
+ * \return The docstring of the function.
+ */
+ virtual Optional<String> GetFunctionDoc(const String& name) { return
std::nullopt; }
+ // Rationale: We separate the docstring from the metadata since docstrings
+ // can be unstructured and sometimes large, while metadata can be focused
+ // on storing structured information.
/*!
* \brief Get the metadata of the function, if available.
* \param name The name of the function.
@@ -128,6 +137,13 @@ class TVM_FFI_EXTRA_CXX_API ModuleObj : public Object {
* \return True if the module implements the function, false otherwise.
*/
bool ImplementsFunction(const String& name, bool query_imports);
+ /*!
+ * \brief Get the function docstring of the function if available.
+ * \param name The name of the function.
+ * \param query_imports Whether to query imported modules.
+ * \return The function docstring of the function.
+ */
+ Optional<String> GetFunctionDoc(const String& name, bool query_imports);
/*!
* \brief Get the function metadata of the function if available.
* \param name The name of the function.
diff --git a/include/tvm/ffi/reflection/registry.h
b/include/tvm/ffi/reflection/registry.h
index e989530..30f3a11 100644
--- a/include/tvm/ffi/reflection/registry.h
+++ b/include/tvm/ffi/reflection/registry.h
@@ -300,7 +300,7 @@ class GlobalDef : public ReflectionDefBase {
TVMFFIMethodInfo info;
info.name = TVMFFIByteArray{name, std::char_traits<char>::length(name)};
info.doc = TVMFFIByteArray{nullptr, 0};
- info.type_schema = TVMFFIByteArray{nullptr, 0};
+ info.metadata = TVMFFIByteArray{nullptr, 0};
info.flags = 0;
// obtain the method function
info.method = AnyView(func).CopyToTVMFFIAny();
@@ -447,7 +447,7 @@ class ObjectDef : public ReflectionDefBase {
// initialize default value to nullptr
info.default_value = AnyView(nullptr).CopyToTVMFFIAny();
info.doc = TVMFFIByteArray{nullptr, 0};
- info.type_schema = TVMFFIByteArray{nullptr, 0};
+ info.metadata = TVMFFIByteArray{nullptr, 0};
// apply field info traits
((ApplyFieldInfoTrait(&info, std::forward<ExtraArgs>(extra_args)), ...));
// call register
@@ -460,7 +460,7 @@ class ObjectDef : public ReflectionDefBase {
TVMFFIMethodInfo info;
info.name = TVMFFIByteArray{name, std::char_traits<char>::length(name)};
info.doc = TVMFFIByteArray{nullptr, 0};
- info.type_schema = TVMFFIByteArray{nullptr, 0};
+ info.metadata = TVMFFIByteArray{nullptr, 0};
info.flags = 0;
if (is_static) {
info.flags |= kTVMFFIFieldFlagBitMaskIsStaticMethod;
diff --git a/python/tvm_ffi/cython/base.pxi b/python/tvm_ffi/cython/base.pxi
index a3ab73e..633ace4 100644
--- a/python/tvm_ffi/cython/base.pxi
+++ b/python/tvm_ffi/cython/base.pxi
@@ -166,7 +166,7 @@ cdef extern from "tvm/ffi/c_api.h":
ctypedef struct TVMFFIFieldInfo:
TVMFFIByteArray name
TVMFFIByteArray doc
- TVMFFIByteArray type_schema
+ TVMFFIByteArray metadata
int64_t flags
int64_t size
int64_t alignment
@@ -179,7 +179,7 @@ cdef extern from "tvm/ffi/c_api.h":
ctypedef struct TVMFFIMethodInfo:
TVMFFIByteArray name
TVMFFIByteArray doc
- TVMFFIByteArray type_schema
+ TVMFFIByteArray metadata
int64_t flags
TVMFFIAny method
diff --git a/src/ffi/extra/module.cc b/src/ffi/extra/module.cc
index d2ebcd1..9601ef1 100644
--- a/src/ffi/extra/module.cc
+++ b/src/ffi/extra/module.cc
@@ -58,6 +58,20 @@ Optional<String> ModuleObj::GetFunctionMetadata(const
String& name, bool query_i
return std::nullopt;
}
+Optional<String> ModuleObj::GetFunctionDoc(const String& name, bool
query_imports) {
+ if (auto opt_str = this->GetFunctionDoc(name)) {
+ return opt_str;
+ }
+ if (query_imports) {
+ for (const Any& import : imports_) {
+ if (auto opt_str = import.cast<Module>()->GetFunctionDoc(name,
query_imports)) {
+ return *opt_str;
+ }
+ }
+ }
+ return std::nullopt;
+}
+
void ModuleObj::ImportModule(const Module& other) {
std::unordered_set<const ModuleObj*> visited{other.operator->()};
std::vector<const ModuleObj*> stack{other.operator->()};
@@ -133,6 +147,10 @@ TVM_FFI_STATIC_INIT_BLOCK() {
[](Module mod, String name, bool query_imports) {
return mod->GetFunctionMetadata(name, query_imports);
})
+ .def_method("ffi.ModuleGetFunctionDoc",
+ [](Module mod, String name, bool query_imports) {
+ return mod->GetFunctionDoc(name, query_imports);
+ })
.def_method("ffi.ModuleGetFunction",
[](Module mod, String name, bool query_imports) {
return mod->GetFunction(name, query_imports);
diff --git a/src/ffi/function.cc b/src/ffi/function.cc
index b1bee7e..954608f 100644
--- a/src/ffi/function.cc
+++ b/src/ffi/function.cc
@@ -54,14 +54,14 @@ class GlobalFunctionTable {
public:
String name_data;
String doc_data;
- String type_schema_data;
+ String metadata_data;
ffi::Function func_data;
explicit Entry(const TVMFFIMethodInfo* method_info) {
// make copy of the metadata
name_data = String(method_info->name.data, method_info->name.size);
doc_data = String(method_info->doc.data, method_info->doc.size);
- type_schema_data = String(method_info->type_schema.data,
method_info->type_schema.size);
+ metadata_data = String(method_info->metadata.data,
method_info->metadata.size);
func_data =
AnyView::CopyFromTVMFFIAny(method_info->method).cast<ffi::Function>();
this->SyncMethodInfo(method_info->flags);
// no need to update method pointer as it would remain the same as func
and we retained
@@ -75,7 +75,7 @@ class GlobalFunctionTable {
this->flags = flags;
this->name = TVMFFIByteArray{name_data.data(), name_data.size()};
this->doc = TVMFFIByteArray{doc_data.data(), doc_data.size()};
- this->type_schema = TVMFFIByteArray{type_schema_data.data(),
type_schema_data.size()};
+ this->metadata = TVMFFIByteArray{metadata_data.data(),
metadata_data.size()};
}
};
diff --git a/src/ffi/object.cc b/src/ffi/object.cc
index aa7edbb..2f25506 100644
--- a/src/ffi/object.cc
+++ b/src/ffi/object.cc
@@ -198,7 +198,7 @@ class TypeTable {
TVMFFIFieldInfo field_data = *info;
field_data.name = this->CopyString(info->name);
field_data.doc = this->CopyString(info->doc);
- field_data.type_schema = this->CopyString(info->type_schema);
+ field_data.metadata = this->CopyString(info->metadata);
if (info->flags & kTVMFFIFieldFlagBitMaskHasDefault) {
field_data.default_value =
this->CopyAny(AnyView::CopyFromTVMFFIAny(info->default_value)).CopyToTVMFFIAny();
@@ -216,7 +216,7 @@ class TypeTable {
TVMFFIMethodInfo method_data = *info;
method_data.name = this->CopyString(info->name);
method_data.doc = this->CopyString(info->doc);
- method_data.type_schema = this->CopyString(info->type_schema);
+ method_data.metadata = this->CopyString(info->metadata);
method_data.method =
this->CopyAny(AnyView::CopyFromTVMFFIAny(info->method)).CopyToTVMFFIAny();
entry->type_methods_data.push_back(method_data);
entry->methods = entry->type_methods_data.data();
