This is an automated email from the ASF dual-hosted git repository.
junrushao 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 ac7bf68 feat: Embed metadata for FFI exported functions (#230)
ac7bf68 is described below
commit ac7bf68058b392c3a8475619016fde48bd08334b
Author: Kathryn (Jinqi) Chen <[email protected]>
AuthorDate: Thu Nov 20 17:40:59 2025 -0800
feat: Embed metadata for FFI exported functions (#230)
## Summary
This PR adds reflection metadata support for functions exported via
`TVM_FFI_DLL_EXPORT_TYPED_FUNC`, enabling signature validation and
memory effect analysis without invoking the function. This addresses the
need for compiler frameworks (e.g., JAX/XLA, MLIR-based compilers) to
validate third-party FFI modules at compile time.
## Core Changes
**1. `TVM_FFI_DLL_EXPORT_TYPED_FUNC` macro**
- Exports two symbols per function:
- `__tvm_ffi_<name>`
- `__tvm_ffi__metadata_<name>`
**Metadata format:**
```
{
"type_schema": "{\"type\":\"ffi.Function\",\"args\":[...]}",
}
```
Based on discussion, we choose to not include C++ const information in
metadata in this PR, and think about it as a followup.
**2. `TVM_FFI_DLL_EXPORT_TYPED_FUNC_DOC` macro**
- Exports three symbols per function:
- `__tvm_ffi_<name>`
- `__tvm_ffi__metadata_<name>`
- `__tvm_ffi__doc_<name>`
**3. Module API** (C++)
- `Module::GetFunctionMetadata(name)` - Query metadata from loaded
modules
- `Module::GetFunctionDoc(name)` - Query documentation
- Works with both static and dynamic linking
**4. Module API** (Python)
- `mod.get_function_metadata("func_name")` - Query metadata from loaded
modules
- `mod.get_function_doc("func_name")` - Query documentation
- Might be good to keep methods private
## Tests
The tests include both unit tests and integration tests.
**C++ Tests:**
- `tests/cpp/test_metadata.cc` - Direct symbol tests
- `tests/cpp/test_function.cc` - Static linking scenario
- `tests/cpp/extra/test_module.cc` - Module API
**Python Tests:**
- `tests/python/test_metadata.py` - Unit test
- `tests/python/test_build.py` - Integration test
---
docs/get_started/quickstart.rst | 3 +-
docs/guides/compiler_integration.md | 10 +-
docs/guides/cpp_guide.md | 47 +++++-
include/tvm/ffi/extra/module.h | 31 +++-
include/tvm/ffi/function.h | 131 ++++++++++++++---
python/tvm_ffi/cpp/extension.py | 46 ++++--
python/tvm_ffi/module.py | 101 ++++++++++++-
src/ffi/extra/library_module.cc | 24 ++++
src/ffi/testing/testing.cc | 190 ++++++++++++++++++-------
tests/cpp/test_function.cc | 15 ++
tests/cpp/test_metadata.cc | 42 ++++++
tests/python/test_build.cc | 2 +
tests/python/test_build.py | 276 ++++++++++++++++++++++++++++++++++++
13 files changed, 822 insertions(+), 96 deletions(-)
diff --git a/docs/get_started/quickstart.rst b/docs/get_started/quickstart.rst
index 650b750..1ebb533 100644
--- a/docs/get_started/quickstart.rst
+++ b/docs/get_started/quickstart.rst
@@ -76,7 +76,8 @@ Suppose we implement a C++ function ``AddOne`` that performs
elementwise ``y = x
The macro :c:macro:`TVM_FFI_DLL_EXPORT_TYPED_FUNC` exports the C++ function
``AddOne``
-as a TVM FFI compatible symbol with the name ``__tvm_ffi_add_one_cpu/cuda`` in
the resulting library.
+as a TVM FFI compatible symbol ``__tvm_ffi_add_one_cpu/cuda``. If
:c:macro:`TVM_FFI_DLL_EXPORT_INCLUDE_METADATA` is set to 1,
+it also exports the function's metadata as a symbol
``__tvm_ffi__metadata_add_one_cpu/cuda`` for type checking and stub generation.
The class :cpp:class:`tvm::ffi::TensorView` allows zero-copy interop with
tensors from different ML frameworks:
diff --git a/docs/guides/compiler_integration.md
b/docs/guides/compiler_integration.md
index aff6f90..254df38 100644
--- a/docs/guides/compiler_integration.md
+++ b/docs/guides/compiler_integration.md
@@ -30,9 +30,13 @@ following options:
- For compilers that generate host functions via codegen (e.g., LLVM), one can
generate the symbol `__tvm_ffi_<func_name>`, where `<funcname>` is the
exported
- function.
-- For kernel generators that generate C++ host code, one can directly
- use {c:macro}`TVM_FFI_DLL_EXPORT_TYPED_FUNC` to expose the symbol.
+ function. Optionally, also generate `__tvm_ffi__metadata_<func_name>` for
reflection.
+- For kernel generators that generate C++ host code, use
{c:macro}`TVM_FFI_DLL_EXPORT_TYPED_FUNC`.
+ This macro automatically exports function metadata when
{c:macro}`TVM_FFI_DLL_EXPORT_INCLUDE_METADATA`
+ is set to 1.
+- To export documentation strings, use
{c:macro}`TVM_FFI_DLL_EXPORT_TYPED_FUNC_DOC` separately
+ after exporting the function. This enhances tooling support (stub
generation, IDE tooltips).
+ Documentation export is also controlled by
{c:macro}`TVM_FFI_DLL_EXPORT_INCLUDE_METADATA`.
The following code snippet shows C code that corresponds to a
function performing `add_one_c` under the ABI. It is reasonably
straightforward for
diff --git a/docs/guides/cpp_guide.md b/docs/guides/cpp_guide.md
index ef69ed3..e5f09b7 100644
--- a/docs/guides/cpp_guide.md
+++ b/docs/guides/cpp_guide.md
@@ -292,12 +292,55 @@ void AddOne(DLTensor* x, DLTensor* y) {
TVM_FFI_DLL_EXPORT_TYPED_FUNC(add_one, my_ffi_extension::AddOne);
```
-The new `add_one` takes the signature of `TVMFFISafeCallType` and can be
wrapped as `ffi::Function`
-through the C++ `ffi::Module` API.
+The new `add_one` takes the signature of `TVMFFISafeCallType`, and can loaded
and queried through the C++ `ffi::Module` API.
+
+When flag `TVM_FFI_DLL_EXPORT_TYPED_FUNC_METADATA` is on, the macro exports
both the function and its type metadata, enabling
+signature validation without calling the function.
+
+The metadata contains:
+
+- **type_schema**: JSON string describing function signature (return type and
argument types)
```cpp
ffi::Module mod = ffi::Module::LoadFromFile("path/to/export_lib.so");
+
+// Get the function
ffi::Function func = mod->GetFunction("add_one").value();
+
+// Query metadata (type schema information)
+ffi::Optional<ffi::String> metadata = mod->GetFunctionMetadata("add_one");
+if (metadata.has_value()) {
+ // Parse JSON metadata for validation
+ // Contains: {"type_schema": "..."}
+}
+```
+
+For functions that need documentation, use the
`TVM_FFI_DLL_EXPORT_TYPED_FUNC_DOC` macro separately:
+
+```cpp
+#define TVM_FFI_DLL_EXPORT_INCLUDE_METADATA 1
+
+void ProcessBatch(ffi::TensorView input, ffi::TensorView output) {
+ // ... implementation
+}
+
+// Export the function
+TVM_FFI_DLL_EXPORT_TYPED_FUNC(process_batch, ProcessBatch);
+
+// Export documentation separately (make sure
TVM_FFI_DLL_EXPORT_INCLUDE_METADATA is set to 1)
+TVM_FFI_DLL_EXPORT_TYPED_FUNC_DOC(
+ process_batch,
+ R"(Process a batch of inputs and write results to output tensor.
+
+Parameters
+----------
+input : TensorView
+ Input tensor to process
+output : TensorView
+ Output tensor for results)");
+
+// Query documentation
+ffi::Optional<ffi::String> doc = mod->GetFunctionDoc("process_batch");
```
## Error Handling
diff --git a/include/tvm/ffi/extra/module.h b/include/tvm/ffi/extra/module.h
index ea6cf3f..6af26c2 100644
--- a/include/tvm/ffi/extra/module.h
+++ b/include/tvm/ffi/extra/module.h
@@ -74,7 +74,9 @@ class TVM_FFI_EXTRA_CXX_API ModuleObj : public Object {
/*!
* \brief Get the docstring of the function, if available.
* \param name The name of the function.
- * \return The docstring of the function.
+ * \return The documentation string if available, nullopt otherwise.
+ *
+ * \sa GetFunctionMetadata, TVM_FFI_DLL_EXPORT_TYPED_FUNC_DOC
*/
virtual Optional<String> GetFunctionDoc(const String& name) { return
std::nullopt; }
// Rationale: We separate the docstring from the metadata since docstrings
@@ -83,7 +85,18 @@ class TVM_FFI_EXTRA_CXX_API ModuleObj : public Object {
/*!
* \brief Get the metadata of the function, if available.
* \param name The name of the function.
- * \return The metadata stored in json string format.
+ * \return The metadata as JSON string if available, nullopt otherwise.
+ *
+ * \code
+ * Module mod = Module::LoadFromFile("lib.so");
+ * Optional<String> metadata = mod->GetFunctionMetadata("my_func");
+ * if (metadata.has_value()) {
+ * // Parse JSON: {"type_schema": "..."}
+ * validate_signature(*metadata);
+ * }
+ * \endcode
+ *
+ * \sa GetFunctionDoc, TVM_FFI_DLL_EXPORT_TYPED_FUNC
*/
virtual Optional<String> GetFunctionMetadata(const String& name) { return
std::nullopt; }
/*!
@@ -142,15 +155,19 @@ class TVM_FFI_EXTRA_CXX_API ModuleObj : public Object {
/*!
* \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.
+ * \param query_imports Whether to also query modules imported by this
module.
+ * \return The documentation string if available, nullopt otherwise.
+ *
+ * \sa GetFunctionMetadata
*/
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.
- * \param query_imports Whether to query imported modules.
- * \return The function metadata of the function in json format.
+ * \param query_imports Whether to also query modules imported by this
module.
+ * \return The metadata as JSON string if available, nullopt otherwise.
+ *
+ * \sa GetFunctionDoc
*/
Optional<String> GetFunctionMetadata(const String& name, bool query_imports);
/*!
@@ -275,6 +292,8 @@ constexpr const char* tvm_ffi_library_ctx =
"__tvm_ffi__library_ctx";
constexpr const char* tvm_ffi_library_bin = "__tvm_ffi__library_bin";
/*! \brief Optional metadata prefix of a symbol. */
constexpr const char* tvm_ffi_metadata_prefix = "__tvm_ffi__metadata_";
+/*! \brief Optional documentation prefix of a symbol. */
+constexpr const char* tvm_ffi_doc_prefix = "__tvm_ffi__doc_";
} // namespace symbol
} // namespace ffi
} // namespace tvm
diff --git a/include/tvm/ffi/function.h b/include/tvm/ffi/function.h
index b339c8c..3f6fe9b 100644
--- a/include/tvm/ffi/function.h
+++ b/include/tvm/ffi/function.h
@@ -23,6 +23,16 @@
#ifndef TVM_FFI_FUNCTION_H_
#define TVM_FFI_FUNCTION_H_
+/*!
+ * \brief Controls whether DLL exports should include metadata.
+ *
+ * When set to 1, exported functions will include additional metadata.
+ * When set to 0 (default), exports are minimal without metadata.
+ */
+#ifndef TVM_FFI_DLL_EXPORT_INCLUDE_METADATA
+#define TVM_FFI_DLL_EXPORT_INCLUDE_METADATA 0
+#endif
+
#include <tvm/ffi/any.h>
#include <tvm/ffi/base_details.h>
#include <tvm/ffi/c_api.h>
@@ -30,6 +40,7 @@
#include <tvm/ffi/function_details.h>
#include <functional>
+#include <sstream>
#include <string>
#include <type_traits>
#include <utility>
@@ -858,13 +869,37 @@ inline int32_t TypeKeyToIndex(std::string_view type_key) {
return type_index;
}
+/// \cond Doxygen_Suppress
+// Internal implementation macros used by TVM_FFI_DLL_EXPORT_TYPED_FUNC and
related macros.
+// These should not be used directly; use the public macros instead.
+
+// Internal implementation macro that generates the C ABI wrapper function
+#define TVM_FFI_DLL_EXPORT_TYPED_FUNC_IMPL_(ExportName, Function)
\
+ extern "C" {
\
+ TVM_FFI_DLL_EXPORT int __tvm_ffi_##ExportName(void* self, const TVMFFIAny*
args, \
+ int32_t num_args, TVMFFIAny*
result) { \
+ TVM_FFI_SAFE_CALL_BEGIN();
\
+ using FuncInfo = ::tvm::ffi::details::FunctionInfo<decltype(Function)>;
\
+ static std::string name = #ExportName;
\
+ ::tvm::ffi::details::unpack_call<typename FuncInfo::RetType>(
\
+ std::make_index_sequence<FuncInfo::num_args>{}, &name, Function,
\
+ reinterpret_cast<const ::tvm::ffi::AnyView*>(args), num_args,
\
+ reinterpret_cast<::tvm::ffi::Any*>(result));
\
+ TVM_FFI_SAFE_CALL_END();
\
+ }
\
+ }
+/// \endcond
+
/*!
* \brief Export typed function as a SafeCallType symbol that follows the FFI
ABI.
*
+ * This macro exports the function and automatically exports metadata when
+ * TVM_FFI_DLL_EXPORT_INCLUDE_METADATA is defined.
+ *
* \param ExportName The symbol name to be exported.
* \param Function The typed function.
*
- * \sa ffi::TypedFunction
+ * \sa ffi::TypedFunction, TVM_FFI_DLL_EXPORT_TYPED_FUNC_DOC
*
* \code
*
@@ -881,22 +916,86 @@ inline int32_t TypeKeyToIndex(std::string_view type_key) {
* });
* \endcode
*
- * \note The final symbol name is `__tvm_ffi_<ExportName>`.
+ * \note The final symbol names are:
+ * - `__tvm_ffi_<ExportName>` (function)
+ * - `__tvm_ffi__metadata_<ExportName>` (metadata - only when
+ * TVM_FFI_DLL_EXPORT_INCLUDE_METADATA is defined)
*/
-#define TVM_FFI_DLL_EXPORT_TYPED_FUNC(ExportName, Function)
\
- extern "C" {
\
- TVM_FFI_DLL_EXPORT int __tvm_ffi_##ExportName(void* self, const TVMFFIAny*
args, \
- int32_t num_args, TVMFFIAny*
result) { \
- TVM_FFI_SAFE_CALL_BEGIN();
\
- using FuncInfo = ::tvm::ffi::details::FunctionInfo<decltype(Function)>;
\
- static std::string name = #ExportName;
\
- ::tvm::ffi::details::unpack_call<typename FuncInfo::RetType>(
\
- std::make_index_sequence<FuncInfo::num_args>{}, &name, Function,
\
- reinterpret_cast<const ::tvm::ffi::AnyView*>(args), num_args,
\
- reinterpret_cast<::tvm::ffi::Any*>(result));
\
- TVM_FFI_SAFE_CALL_END();
\
- }
\
- }
+#if TVM_FFI_DLL_EXPORT_INCLUDE_METADATA
+#define TVM_FFI_DLL_EXPORT_TYPED_FUNC(ExportName, Function)
\
+ TVM_FFI_DLL_EXPORT_TYPED_FUNC_IMPL_(ExportName, Function)
\
+ extern "C" {
\
+ TVM_FFI_DLL_EXPORT int __tvm_ffi__metadata_##ExportName(void* self, const
TVMFFIAny* args, \
+ int32_t num_args,
TVMFFIAny* result) { \
+ TVM_FFI_SAFE_CALL_BEGIN();
\
+ using FuncInfo = ::tvm::ffi::details::FunctionInfo<decltype(Function)>;
\
+ std::ostringstream os;
\
+ os << R"({"type_schema":)"
\
+ << ::tvm::ffi::EscapeString(::tvm::ffi::String(FuncInfo::TypeSchema()))
<< R"(})"; \
+ ::tvm::ffi::String str(os.str());
\
+ ::tvm::ffi::TypeTraits<::tvm::ffi::String>::MoveToAny(std::move(str),
result); \
+ TVM_FFI_SAFE_CALL_END();
\
+ }
\
+ }
+#else
+#define TVM_FFI_DLL_EXPORT_TYPED_FUNC(ExportName, Function) \
+ TVM_FFI_DLL_EXPORT_TYPED_FUNC_IMPL_(ExportName, Function)
+#endif
+
+/*!
+ * \brief Export documentation string for a typed function.
+ *
+ * This macro exports a documentation string associated with a function export
name.
+ * The docstring can be used by stub generators and documentation tools.
+ * This macro only exports the docstring; it does not export the function
itself.
+ *
+ * \param ExportName The symbol name that the docstring is associated with.
+ * \param DocString The documentation string (C string literal).
+ *
+ * \sa ffi::TypedFunction, TVM_FFI_DLL_EXPORT_TYPED_FUNC
+ *
+ * \code
+ *
+ * int Add(int a, int b) {
+ * return a + b;
+ * }
+ *
+ * TVM_FFI_DLL_EXPORT_TYPED_FUNC(add, Add);
+ * TVM_FFI_DLL_EXPORT_TYPED_FUNC_DOC(
+ * add,
+ * R"(Add two integers and return the sum.
+ *
+ * Parameters
+ * ----------
+ * a : int
+ * First integer
+ * b : int
+ * Second integer
+ *
+ * Returns
+ * -------
+ * result : int
+ * Sum of a and b)");
+ *
+ * \endcode
+ *
+ * \note The exported symbol name is `__tvm_ffi__doc_<ExportName>` (docstring
getter function).
+ * This symbol is only exported when TVM_FFI_DLL_EXPORT_INCLUDE_METADATA
is defined.
+ */
+#if TVM_FFI_DLL_EXPORT_INCLUDE_METADATA
+#define TVM_FFI_DLL_EXPORT_TYPED_FUNC_DOC(ExportName, DocString)
\
+ extern "C" {
\
+ TVM_FFI_DLL_EXPORT int __tvm_ffi__doc_##ExportName(void* self, const
TVMFFIAny* args, \
+ int32_t num_args,
TVMFFIAny* result) { \
+ TVM_FFI_SAFE_CALL_BEGIN();
\
+ ::tvm::ffi::String str(DocString);
\
+ ::tvm::ffi::TypeTraits<::tvm::ffi::String>::MoveToAny(std::move(str),
result); \
+ TVM_FFI_SAFE_CALL_END();
\
+ }
\
+ }
+#else
+#define TVM_FFI_DLL_EXPORT_TYPED_FUNC_DOC(ExportName, DocString)
+#endif
} // namespace ffi
} // namespace tvm
#endif // TVM_FFI_FUNCTION_H_
diff --git a/python/tvm_ffi/cpp/extension.py b/python/tvm_ffi/cpp/extension.py
index ff4daff..4b1af8f 100644
--- a/python/tvm_ffi/cpp/extension.py
+++ b/python/tvm_ffi/cpp/extension.py
@@ -353,6 +353,23 @@ def build_ninja(build_dir: str) -> None:
raise RuntimeError("\n".join(msg))
+# Translation table for escaping C++ string literals
+_CPP_ESCAPE_TABLE = str.maketrans(
+ {
+ "\\": "\\\\",
+ '"': '\\"',
+ "\n": "\\n",
+ "\r": "\\r",
+ "\t": "\\t",
+ }
+)
+
+
+def _escape_cpp_string_literal(s: str) -> str:
+ """Escape special characters for C++ string literals."""
+ return s.translate(_CPP_ESCAPE_TABLE)
+
+
def _decorate_with_tvm_ffi(source: str, functions: Mapping[str, str]) -> str:
"""Decorate the given source code with TVM FFI export macros."""
sources = [
@@ -367,7 +384,11 @@ def _decorate_with_tvm_ffi(source: str, functions:
Mapping[str, str]) -> str:
for func_name, func_doc in functions.items():
sources.append(f"TVM_FFI_DLL_EXPORT_TYPED_FUNC({func_name},
{func_name});")
- _ = func_doc # todo: add support to embed function docstring to the
tvm ffi functions.
+
+ if func_doc:
+ # Escape the docstring for C++ string literal
+ escaped_doc = _escape_cpp_string_literal(func_doc)
+ sources.append(f'TVM_FFI_DLL_EXPORT_TYPED_FUNC_DOC({func_name},
"{escaped_doc}");')
sources.append("")
@@ -496,11 +517,12 @@ def build_inline(
The CUDA source code. It can be a list of sources or a single source.
functions
The functions in cpp_sources or cuda_source that will be exported to
the tvm ffi module. When a mapping is
- given, the keys are the names of the exported functions, and the
values are docstrings for the functions. When
- a sequence or a single string is given, they are the functions needed
to be exported, and the docstrings are set
- to empty strings. A single function name can also be given as a
string. When cpp_sources is given, the functions
- must be declared (not necessarily defined) in the cpp_sources. When
cpp_sources is not given, the functions
- must be defined in the cuda_sources. If not specified, no function
will be exported.
+ given, the keys are the names of the exported functions, and the
values are docstrings for the functions
+ (use an empty string to skip documentation for specific functions).
When a sequence or a single string is given, they are
+ the functions needed to be exported, and the docstrings are set to
empty strings. A single function name can
+ also be given as a string. When cpp_sources is given, the functions
must be declared (not necessarily defined)
+ in the cpp_sources. When cpp_sources is not given, the functions must
be defined in the cuda_sources. If not
+ specified, no function will be exported.
extra_cflags
The extra compiler flags for C++ compilation.
The default flags are:
@@ -604,7 +626,6 @@ def build_inline(
else:
cpp_source = _decorate_with_tvm_ffi(cpp_source, {})
cuda_source = _decorate_with_tvm_ffi(cuda_source, function_map)
-
# determine the cache dir for the built module
build_dir: Path
if build_directory is None:
@@ -694,11 +715,12 @@ def load_inline(
The CUDA source code. It can be a list of sources or a single source.
functions: Mapping[str, str] | Sequence[str] | str, optional
The functions in cpp_sources or cuda_source that will be exported to
the tvm ffi module. When a mapping is
- given, the keys are the names of the exported functions, and the
values are docstrings for the functions. When
- a sequence or a single string is given, they are the functions needed
to be exported, and the docstrings are set
- to empty strings. A single function name can also be given as a
string. When cpp_sources is given, the functions
- must be declared (not necessarily defined) in the cpp_sources. When
cpp_sources is not given, the functions
- must be defined in the cuda_sources. If not specified, no function
will be exported.
+ given, the keys are the names of the exported functions, and the
values are docstrings for the functions
+ (use an empty string to skip documentation for specific functions).
When a sequence or a single string is given, they are
+ the functions needed to be exported, and the docstrings are set to
empty strings. A single function name can
+ also be given as a string. When cpp_sources is given, the functions
must be declared (not necessarily defined)
+ in the cpp_sources. When cpp_sources is not given, the functions must
be defined in the cuda_sources. If not
+ specified, no function will be exported.
extra_cflags: Sequence[str], optional
The extra compiler flags for C++ compilation.
The default flags are:
diff --git a/python/tvm_ffi/module.py b/python/tvm_ffi/module.py
index 080d431..7b548d6 100644
--- a/python/tvm_ffi/module.py
+++ b/python/tvm_ffi/module.py
@@ -26,6 +26,7 @@ if TYPE_CHECKING:
# isort: on
# fmt: on
# tvm-ffi-stubgen(end)
+import json
from enum import IntEnum
from os import PathLike, fspath
from typing import ClassVar, cast
@@ -54,14 +55,23 @@ class Module(core.Object):
import tvm_ffi
- # load the module from a tvm-ffi shared library
- mod : tvm_ffi.Module = tvm_ffi.load_module("path/to/library.so")
- # you can use mod.func_name to call the exported function
+ # Load the module from a shared library
+ mod = tvm_ffi.load_module("path/to/library.so")
+
+ # Call exported function
mod.func_name(*args)
+ # Query function metadata (type signature)
+ metadata = mod.get_function_metadata("func_name")
+
+ # Query function documentation (if available)
+ doc = mod.get_function_doc("func_name")
+
See Also
--------
:py:func:`tvm_ffi.load_module`
+ :py:meth:`get_function_metadata`
+ :py:meth:`get_function_doc`
Notes
-----
@@ -175,6 +185,91 @@ class Module(core.Object):
raise AttributeError(f"Module has no function '{name}'")
return func
+ def get_function_metadata(
+ self, name: str, query_imports: bool = False
+ ) -> dict[str, Any] | None:
+ """Get metadata for a function exported from the module.
+
+ This retrieves metadata for functions exported via
c:macro:`TVM_FFI_DLL_EXPORT_TYPED_FUNC`
+ and when c:macro:`TVM_FFI_DLL_EXPORT_INCLUDE_METADATA` is on, which
includes type schema
+ information.
+
+ Parameters
+ ----------
+ name
+ The name of the function
+
+ query_imports
+ Whether to also query modules imported by this module.
+
+ Returns
+ -------
+ metadata
+ A dictionary containing function metadata. The ``type_schema``
field
+ encodes the callable signature.
+
+ Examples
+ --------
+ .. code-block:: python
+
+ import tvm_ffi
+ from tvm_ffi.core import TypeSchema
+ import json
+
+ mod = tvm_ffi.load_module("add_one_cpu.so")
+ metadata = mod.get_function_metadata("add_one_cpu")
+ schema = TypeSchema.from_json_str(metadata["type_schema"])
+ print(schema) # Shows function signature
+
+ See Also
+ --------
+ :py:func:`tvm_ffi.get_global_func_metadata`
+ Get metadata for global registry functions.
+
+ """
+ metadata_str = _ffi_api.ModuleGetFunctionMetadata(self, name,
query_imports)
+ if metadata_str is None:
+ return None
+ return json.loads(metadata_str)
+
+ def get_function_doc(self, name: str, query_imports: bool = False) -> str
| None:
+ """Get documentation string for a function exported from the module.
+
+ This retrieves documentation for functions exported via
c:macro:`TVM_FFI_DLL_EXPORT_TYPED_FUNC_DOC`.
+
+ Parameters
+ ----------
+ name
+ The name of the function
+
+ query_imports
+ Whether to also query modules imported by this module.
+
+ Returns
+ -------
+ doc : str or None
+ The documentation string if available, None otherwise.
+
+ Examples
+ --------
+ .. code-block:: python
+
+ import tvm_ffi
+
+ mod = tvm_ffi.load_module("mylib.so")
+ doc = mod.get_function_doc("process_batch")
+ if doc:
+ print(doc)
+
+ See Also
+ --------
+ :py:meth:`get_function_metadata`
+ Get metadata including type schema.
+
+ """
+ doc_str = _ffi_api.ModuleGetFunctionDoc(self, name, query_imports)
+ return doc_str if doc_str else None
+
def import_module(self, module: Module) -> None:
"""Add module to the import list of current one.
diff --git a/src/ffi/extra/library_module.cc b/src/ffi/extra/library_module.cc
index f707c72..cc992aa 100644
--- a/src/ffi/extra/library_module.cc
+++ b/src/ffi/extra/library_module.cc
@@ -58,6 +58,30 @@ class LibraryModuleObj final : public ModuleObj {
return std::nullopt;
}
+ Optional<String> GetFunctionMetadata(const String& name) final {
+ // Look for __tvm_ffi__metadata_<name> symbol
+ String metadata_symbol = symbol::tvm_ffi_metadata_prefix + name;
+ void* symbol = lib_->GetSymbol(metadata_symbol);
+ if (symbol != nullptr) {
+ using MetadataGetter = int (*)(void*, const TVMFFIAny*, int32_t,
TVMFFIAny*);
+ auto metadata_getter = reinterpret_cast<MetadataGetter>(symbol);
+ return Function::InvokeExternC(nullptr, metadata_getter).cast<String>();
+ }
+ return std::nullopt;
+ }
+
+ Optional<String> GetFunctionDoc(const String& name) final {
+ // Look for __tvm_ffi__doc_<name> symbol
+ String doc_symbol = symbol::tvm_ffi_doc_prefix + name;
+ void* symbol = lib_->GetSymbol(doc_symbol);
+ if (symbol != nullptr) {
+ using DocGetter = int (*)(void*, const TVMFFIAny*, int32_t, TVMFFIAny*);
+ auto doc_getter = reinterpret_cast<DocGetter>(symbol);
+ return Function::InvokeExternC(nullptr, doc_getter).cast<String>();
+ }
+ return std::nullopt;
+ }
+
private:
ObjectPtr<Library> lib_;
};
diff --git a/src/ffi/testing/testing.cc b/src/ffi/testing/testing.cc
index 8d068b4..876801f 100644
--- a/src/ffi/testing/testing.cc
+++ b/src/ffi/testing/testing.cc
@@ -17,6 +17,8 @@
* under the License.
*/
// This file is used for testing the FFI API.
+#define TVM_FFI_DLL_EXPORT_INCLUDE_METADATA 1
+
#include <dlpack/dlpack.h>
#include <tvm/ffi/any.h>
#include <tvm/ffi/container/array.h>
@@ -284,6 +286,82 @@ TVM_FFI_STATIC_INIT_BLOCK() {
namespace tvm {
namespace ffi {
+//
-----------------------------------------------------------------------------
+// Implementation functions for schema testing
+//
-----------------------------------------------------------------------------
+namespace schema_test_impl {
+
+// Simple types
+int64_t schema_id_int(int64_t x) { return x; }
+double schema_id_float(double x) { return x; }
+bool schema_id_bool(bool x) { return x; }
+DLDevice schema_id_device(DLDevice d) { return d; }
+DLDataType schema_id_dtype(DLDataType dt) { return dt; }
+String schema_id_string(String s) { return s; }
+Bytes schema_id_bytes(Bytes b) { return b; }
+Function schema_id_func(Function f) { return f; }
+TypedFunction<void(int64_t, float, Function)> schema_id_func_typed(
+ TypedFunction<void(int64_t, float, Function)> f) {
+ return f;
+}
+Any schema_id_any(Any a) { return a; }
+ObjectRef schema_id_object(ObjectRef o) { return o; }
+DLTensor* schema_id_dltensor(DLTensor* t) { return t; }
+Tensor schema_id_tensor(Tensor t) { return t; }
+void schema_tensor_view_input(TensorView t) {}
+
+// Optional types
+Optional<int64_t> schema_id_opt_int(Optional<int64_t> o) { return o; }
+Optional<String> schema_id_opt_str(Optional<String> o) { return o; }
+Optional<ObjectRef> schema_id_opt_obj(Optional<ObjectRef> o) { return o; }
+
+// Array types
+Array<int64_t> schema_id_arr_int(Array<int64_t> arr) { return arr; }
+Array<String> schema_id_arr_str(Array<String> arr) { return arr; }
+Array<ObjectRef> schema_id_arr_obj(Array<ObjectRef> arr) { return arr; }
+const ArrayObj* schema_id_arr(const ArrayObj* arr) { return arr; }
+
+// Map types
+Map<String, int64_t> schema_id_map_str_int(Map<String, int64_t> m) { return m;
}
+Map<String, String> schema_id_map_str_str(Map<String, String> m) { return m; }
+Map<String, ObjectRef> schema_id_map_str_obj(Map<String, ObjectRef> m) {
return m; }
+const MapObj* schema_id_map(const MapObj* m) { return m; }
+
+// Variant types
+Variant<int64_t, String> schema_id_variant_int_str(Variant<int64_t, String> v)
{ return v; }
+Variant<int64_t, String, Array<int64_t>> schema_variant_mix(
+ Variant<int64_t, String, Array<int64_t>> v) {
+ return v;
+}
+
+// Complex nested types
+Map<String, Array<int64_t>> schema_arr_map_opt(Array<Optional<int64_t>> arr,
+ Map<String, Array<int64_t>> mp,
+ Optional<String> os) {
+ // no-op combine
+ if (os.has_value()) {
+ Array<int64_t> extra;
+ for (const auto& i : arr) {
+ if (i.has_value()) extra.push_back(i.value());
+ }
+ mp.Set(os.value(), extra);
+ }
+ return mp;
+}
+
+// Edge cases
+int64_t schema_no_args() { return 1; }
+void schema_no_return(int64_t x) {}
+void schema_no_args_no_return() {}
+
+// Member function pattern
+int64_t test_int_pair_sum_wrapper(const TestIntPair& target) { return
target.Sum(); }
+
+// Documentation export
+int64_t test_add_with_docstring(int64_t a, int64_t b) { return a + b; }
+
+} // namespace schema_test_impl
+
// A class with a wide variety of field types and method signatures
class SchemaAllTypesObj : public Object {
public:
@@ -394,60 +472,42 @@ TVM_FFI_STATIC_INIT_BLOCK() {
// Global typed functions to exercise RegisterFunc with various schemas
refl::GlobalDef()
- .def(
- "testing.schema_id_int", [](int64_t x) { return x; },
- refl::Metadata{{"bool_attr", true}, //
- {"int_attr", 1}, //
- {"str_attr", "hello"}})
- .def("testing.schema_id_float", [](double x) { return x; })
- .def("testing.schema_id_bool", [](bool x) { return x; })
- .def("testing.schema_id_device", [](DLDevice d) { return d; })
- .def("testing.schema_id_dtype", [](DLDataType dt) { return dt; })
- .def("testing.schema_id_string", [](String s) { return s; })
- .def("testing.schema_id_bytes", [](Bytes b) { return b; })
- .def("testing.schema_id_func", [](Function f) -> Function { return f; })
- .def("testing.schema_id_func_typed",
- [](TypedFunction<void(int64_t, float, Function)> f)
- -> TypedFunction<void(int64_t, float, Function)> { return f; })
- .def("testing.schema_id_any", [](Any a) { return a; })
- .def("testing.schema_id_object", [](ObjectRef o) { return o; })
- .def("testing.schema_id_dltensor", [](DLTensor* t) { return t; })
- .def("testing.schema_id_tensor", [](Tensor t) { return t; })
- .def("testing.schema_tensor_view_input", [](TensorView t) -> void {})
- .def("testing.schema_id_opt_int", [](Optional<int64_t> o) { return o; })
- .def("testing.schema_id_opt_str", [](Optional<String> o) { return o; })
- .def("testing.schema_id_opt_obj", [](Optional<ObjectRef> o) { return o;
})
- .def("testing.schema_id_arr_int", [](Array<int64_t> arr) { return arr; })
- .def("testing.schema_id_arr_str", [](Array<String> arr) { return arr; })
- .def("testing.schema_id_arr_obj", [](Array<ObjectRef> arr) { return arr;
})
- .def("testing.schema_id_arr", [](const ArrayObj* arr) { return arr; })
- .def("testing.schema_id_map_str_int", [](Map<String, int64_t> m) {
return m; })
- .def("testing.schema_id_map_str_str", [](Map<String, String> m) { return
m; })
- .def("testing.schema_id_map_str_obj", [](Map<String, ObjectRef> m) {
return m; })
- .def("testing.schema_id_map", [](const MapObj* m) { return m; })
- .def("testing.schema_id_variant_int_str", [](Variant<int64_t, String> v)
{ return v; })
+ .def("testing.schema_id_int", schema_test_impl::schema_id_int,
+ refl::Metadata{{"bool_attr", true}, //
+ {"int_attr", 1}, //
+ {"str_attr", "hello"}})
+ .def("testing.schema_id_float", schema_test_impl::schema_id_float)
+ .def("testing.schema_id_bool", schema_test_impl::schema_id_bool)
+ .def("testing.schema_id_device", schema_test_impl::schema_id_device)
+ .def("testing.schema_id_dtype", schema_test_impl::schema_id_dtype)
+ .def("testing.schema_id_string", schema_test_impl::schema_id_string)
+ .def("testing.schema_id_bytes", schema_test_impl::schema_id_bytes)
+ .def("testing.schema_id_func", schema_test_impl::schema_id_func)
+ .def("testing.schema_id_func_typed",
schema_test_impl::schema_id_func_typed)
+ .def("testing.schema_id_any", schema_test_impl::schema_id_any)
+ .def("testing.schema_id_object", schema_test_impl::schema_id_object)
+ .def("testing.schema_id_dltensor", schema_test_impl::schema_id_dltensor)
+ .def("testing.schema_id_tensor", schema_test_impl::schema_id_tensor)
+ .def("testing.schema_tensor_view_input",
schema_test_impl::schema_tensor_view_input)
+ .def("testing.schema_id_opt_int", schema_test_impl::schema_id_opt_int)
+ .def("testing.schema_id_opt_str", schema_test_impl::schema_id_opt_str)
+ .def("testing.schema_id_opt_obj", schema_test_impl::schema_id_opt_obj)
+ .def("testing.schema_id_arr_int", schema_test_impl::schema_id_arr_int)
+ .def("testing.schema_id_arr_str", schema_test_impl::schema_id_arr_str)
+ .def("testing.schema_id_arr_obj", schema_test_impl::schema_id_arr_obj)
+ .def("testing.schema_id_arr", schema_test_impl::schema_id_arr)
+ .def("testing.schema_id_map_str_int",
schema_test_impl::schema_id_map_str_int)
+ .def("testing.schema_id_map_str_str",
schema_test_impl::schema_id_map_str_str)
+ .def("testing.schema_id_map_str_obj",
schema_test_impl::schema_id_map_str_obj)
+ .def("testing.schema_id_map", schema_test_impl::schema_id_map)
+ .def("testing.schema_id_variant_int_str",
schema_test_impl::schema_id_variant_int_str)
.def_packed("testing.schema_packed", [](PackedArgs args, Any* ret) {})
- .def("testing.schema_arr_map_opt",
- // NOLINTNEXTLINE(performance-unnecessary-value-param)
- [](Array<Optional<int64_t>> arr, Map<String, Array<int64_t>> mp,
- // NOLINTNEXTLINE(performance-unnecessary-value-param)
- Optional<String> os) -> Map<String, Array<int64_t>> {
- // no-op combine
- if (os.has_value()) {
- Array<int64_t> extra;
- for (const auto& i : arr) {
- if (i.has_value()) extra.push_back(i.value());
- }
- mp.Set(os.value(), extra);
- }
- return mp;
- })
- .def(
- "testing.schema_variant_mix",
- [](Variant<int64_t, String, Array<int64_t>> v) { return v; },
"variant passthrough")
- .def("testing.schema_no_args", []() { return 1; })
- .def("testing.schema_no_return", [](int64_t x) {})
- .def("testing.schema_no_args_no_return", []() {});
+ .def("testing.schema_arr_map_opt", schema_test_impl::schema_arr_map_opt)
+ .def("testing.schema_variant_mix", schema_test_impl::schema_variant_mix,
+ "variant passthrough")
+ .def("testing.schema_no_args", schema_test_impl::schema_no_args)
+ .def("testing.schema_no_return", schema_test_impl::schema_no_return)
+ .def("testing.schema_no_args_no_return",
schema_test_impl::schema_no_args_no_return);
TVMFFIEnvModRegisterSystemLibSymbol("__tvm_ffi_testing.add_one",
reinterpret_cast<void*>(__add_one_c_symbol));
}
@@ -455,4 +515,28 @@ TVM_FFI_STATIC_INIT_BLOCK() {
} // namespace ffi
} // namespace tvm
+//
-----------------------------------------------------------------------------
+// Exported symbols for metadata testing on DLL-exported functions
+//
-----------------------------------------------------------------------------
+// We keep minimal DLL exports here to verify the export mechanism.
+TVM_FFI_DLL_EXPORT_TYPED_FUNC(testing_dll_schema_id_int,
tvm::ffi::schema_test_impl::schema_id_int);
+
+// Documentation export
+TVM_FFI_DLL_EXPORT_TYPED_FUNC(testing_dll_test_add_with_docstring,
+
tvm::ffi::schema_test_impl::test_add_with_docstring);
+TVM_FFI_DLL_EXPORT_TYPED_FUNC_DOC(testing_dll_test_add_with_docstring,
+ R"(Add two integers and return the sum.
+
+Parameters
+----------
+a : int
+ First integer
+b : int
+ Second integer
+
+Returns
+-------
+result : int
+ Sum of a and b)");
+
extern "C" TVM_FFI_DLL_EXPORT int TVMFFITestingDummyTarget() { return 0; }
diff --git a/tests/cpp/test_function.cc b/tests/cpp/test_function.cc
index c274e51..b3879d7 100644
--- a/tests/cpp/test_function.cc
+++ b/tests/cpp/test_function.cc
@@ -16,9 +16,13 @@
* specific language governing permissions and limitations
* under the License.
*/
+#define TVM_FFI_DLL_EXPORT_INCLUDE_METADATA 1
+
#include <gtest/gtest.h>
#include <tvm/ffi/any.h>
#include <tvm/ffi/container/array.h>
+#include <tvm/ffi/container/map.h>
+#include <tvm/ffi/extra/json.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/memory.h>
@@ -285,6 +289,17 @@ int invoke_testing_add1(int x) {
TEST(Func, InvokeExternC) { EXPECT_EQ(invoke_testing_add1(1), 2); }
+extern "C" int __tvm_ffi__metadata_testing_add1(void*, const TVMFFIAny*,
int32_t, TVMFFIAny*);
+
+TEST(Func, StaticLinkingMetadata) {
+ String metadata_str =
+ Function::InvokeExternC(nullptr,
__tvm_ffi__metadata_testing_add1).cast<String>();
+ Map<String, Any> metadata = json::Parse(metadata_str).cast<Map<String,
Any>>();
+ EXPECT_TRUE(metadata.count("type_schema"));
+ std::string type_schema_str = metadata["type_schema"].cast<String>();
+ EXPECT_TRUE(type_schema_str.find("int") != std::string::npos);
+}
+
extern "C" TVM_FFI_DLL int TVMFFITestingDummyTarget();
TEST(Func, DummyCFunc) {
diff --git a/tests/cpp/test_metadata.cc b/tests/cpp/test_metadata.cc
index 982893f..d0dfebd 100644
--- a/tests/cpp/test_metadata.cc
+++ b/tests/cpp/test_metadata.cc
@@ -29,11 +29,25 @@
#include <tvm/ffi/reflection/accessor.h>
#include <tvm/ffi/string.h>
+// Forward declarations for exported FFI functions
+extern "C" {
+int __tvm_ffi__metadata_testing_dll_schema_id_int(void*, const TVMFFIAny*,
int32_t, TVMFFIAny*);
+int __tvm_ffi__metadata_testing_dll_test_add_with_docstring(void*, const
TVMFFIAny*, int32_t,
+ TVMFFIAny*);
+int __tvm_ffi__doc_testing_dll_test_add_with_docstring(void*, const
TVMFFIAny*, int32_t,
+ TVMFFIAny*);
+}
+
namespace {
using namespace tvm::ffi;
using namespace tvm::ffi::reflection;
+// Helper to call metadata FFI functions and return the String result
+static String CallMetadataFunc(int (*func)(void*, const TVMFFIAny*, int32_t,
TVMFFIAny*)) {
+ return Function::InvokeExternC(nullptr, func).cast<String>();
+}
+
static std::string ParseMetadataToSchema(const String& metadata) {
return json::Parse(metadata)
.cast<Map<String, Any>>()["type_schema"] //
@@ -200,4 +214,32 @@ TEST(Schema, MethodTypeSchemas) {
R"({"type":"ffi.Function","args":[{"type":"testing.SchemaAllTypes"},{"type":"int"},{"type":"float"},{"type":"ffi.String"}]})");
}
+TEST(Schema, DLLExportedFuncMetadata) {
+ // Minimal sanity check that DLL export metadata mechanism works.
+
EXPECT_EQ(ParseMetadataToSchema(CallMetadataFunc(__tvm_ffi__metadata_testing_dll_schema_id_int)),
+
R"({"type":"ffi.Function","args":[{"type":"int"},{"type":"int"}]})");
+}
+
+TEST(Schema, DLLExportedFuncDocumentation) {
+ EXPECT_EQ(ParseMetadataToSchema(
+
CallMetadataFunc(__tvm_ffi__metadata_testing_dll_test_add_with_docstring)),
+
R"({"type":"ffi.Function","args":[{"type":"int"},{"type":"int"},{"type":"int"}]})");
+ String doc =
CallMetadataFunc(__tvm_ffi__doc_testing_dll_test_add_with_docstring);
+ std::string doc_str(doc);
+ EXPECT_EQ(doc_str,
+ R"(Add two integers and return the sum.
+
+Parameters
+----------
+a : int
+ First integer
+b : int
+ Second integer
+
+Returns
+-------
+result : int
+ Sum of a and b)");
+}
+
} // namespace
diff --git a/tests/python/test_build.cc b/tests/python/test_build.cc
index e25e8cf..bc6acf1 100644
--- a/tests/python/test_build.cc
+++ b/tests/python/test_build.cc
@@ -16,6 +16,8 @@
* specific language governing permissions and limitations
* under the License.
*/
+#define TVM_FFI_DLL_EXPORT_INCLUDE_METADATA 1
+
#include "test_build.h"
#include <tvm/ffi/container/tensor.h>
diff --git a/tests/python/test_build.py b/tests/python/test_build.py
index 30153dc..7f22669 100644
--- a/tests/python/test_build.py
+++ b/tests/python/test_build.py
@@ -14,11 +14,14 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
+import gc
import pathlib
import numpy
import pytest
+import tvm_ffi
import tvm_ffi.cpp
+from tvm_ffi.core import TypeSchema
from tvm_ffi.module import Module
@@ -31,11 +34,284 @@ def test_build_cpp() -> None:
mod: Module = tvm_ffi.load_module(output_lib_path)
+ metadata = mod.get_function_metadata("add_one_cpu")
+ assert metadata is not None, "add_one_cpu should have metadata"
+ assert "type_schema" in metadata, f"{'add_one_cpu'}: {metadata}"
+ schema = TypeSchema.from_json_str(metadata["type_schema"])
+ assert str(schema) == "Callable[[Tensor, Tensor], None]",
f"{'add_one_cpu'}: {schema}"
+ doc = mod.get_function_doc("add_one_cpu")
+ assert doc is None
+
x = numpy.array([1, 2, 3, 4, 5], dtype=numpy.float32)
y = numpy.empty_like(x)
mod.add_one_cpu(x, y)
numpy.testing.assert_equal(x + 1, y)
+def test_build_inline_with_metadata() -> None: # noqa: PLR0915
+ """Test functions with various input and output types."""
+ # Keep module alive until all returned objects are destroyed
+ mod: Module = tvm_ffi.cpp.load_inline(
+ name="test_io_types",
+ cpp_sources=r"""
+ // int input -> int output
+ int square(int x) {
+ return x * x;
+ }
+
+ // float input -> float output
+ float reciprocal(float x) {
+ return 1.0f / x;
+ }
+
+ // bool input -> bool output
+ bool negate(bool x) {
+ return !x;
+ }
+
+ // String input -> String output
+ tvm::ffi::String uppercase_first(tvm::ffi::String s) {
+ std::string result(s.c_str());
+ if (!result.empty()) {
+ result[0] = std::toupper(result[0]);
+ }
+ return tvm::ffi::String(result);
+ }
+
+ // Multiple inputs: int, float -> float
+ float weighted_sum(int count, float weight) {
+ return static_cast<float>(count) * weight;
+ }
+
+ // Multiple inputs: String, int -> String
+ tvm::ffi::String repeat_string(tvm::ffi::String s, int times) {
+ std::string result;
+ for (int i = 0; i < times; ++i) {
+ result += s.c_str();
+ }
+ return tvm::ffi::String(result);
+ }
+
+ // Mixed types: bool, int, float, String -> String
+ tvm::ffi::String format_data(bool flag, int count, float value,
tvm::ffi::String label) {
+ std::ostringstream oss;
+ oss << label.c_str() << ": flag=" << (flag ? "true" : "false")
+ << ", count=" << count << ", value=" << value;
+ return tvm::ffi::String(oss.str());
+ }
+
+ // Tensor input/output
+ void double_tensor(tvm::ffi::TensorView input,
tvm::ffi::TensorView output) {
+ TVM_FFI_ICHECK(input.ndim() == 1);
+ TVM_FFI_ICHECK(output.ndim() == 1);
+ TVM_FFI_ICHECK(input.size(0) == output.size(0));
+ DLDataType f32_dtype{kDLFloat, 32, 1};
+ TVM_FFI_ICHECK(input.dtype() == f32_dtype);
+ TVM_FFI_ICHECK(output.dtype() == f32_dtype);
+
+ for (int i = 0; i < input.size(0); ++i) {
+ static_cast<float*>(output.data_ptr())[i] =
+ static_cast<const float*>(input.data_ptr())[i] * 2.0f;
+ }
+ }
+ """,
+ functions=[
+ "square",
+ "reciprocal",
+ "negate",
+ "uppercase_first",
+ "weighted_sum",
+ "repeat_string",
+ "format_data",
+ "double_tensor",
+ ],
+ extra_cflags=["-DTVM_FFI_DLL_EXPORT_INCLUDE_METADATA=1"],
+ )
+
+ # Test square: int -> int
+ assert mod.square(5) == 25
+ metadata = mod.get_function_metadata("square")
+ assert metadata is not None
+ schema = TypeSchema.from_json_str(metadata["type_schema"])
+ assert str(schema) == "Callable[[int], int]"
+
+ # Test reciprocal: float -> float
+ result = mod.reciprocal(2.0)
+ assert abs(result - 0.5) < 0.001
+ metadata = mod.get_function_metadata("reciprocal")
+ assert metadata is not None
+ schema = TypeSchema.from_json_str(metadata["type_schema"])
+ assert str(schema) == "Callable[[float], float]"
+
+ # Test negate: bool -> bool
+ assert mod.negate(True) is False
+ assert mod.negate(False) is True
+ metadata = mod.get_function_metadata("negate")
+ assert metadata is not None
+ schema = TypeSchema.from_json_str(metadata["type_schema"])
+ assert str(schema) == "Callable[[bool], bool]"
+
+ # Test uppercase_first: String -> String
+ result = mod.uppercase_first("hello")
+ assert result == "Hello"
+ metadata = mod.get_function_metadata("uppercase_first")
+ assert metadata is not None
+ schema = TypeSchema.from_json_str(metadata["type_schema"])
+ assert str(schema) == "Callable[[str], str]"
+
+ # Test weighted_sum: int, float -> float
+ result = mod.weighted_sum(10, 2.5)
+ assert abs(result - 25.0) < 0.001
+ metadata = mod.get_function_metadata("weighted_sum")
+ assert metadata is not None
+ schema = TypeSchema.from_json_str(metadata["type_schema"])
+ assert str(schema) == "Callable[[int, float], float]"
+
+ # Test repeat_string: String, int -> String
+ result = mod.repeat_string("ab", 3)
+ assert result == "ababab"
+ metadata = mod.get_function_metadata("repeat_string")
+ assert metadata is not None
+ schema = TypeSchema.from_json_str(metadata["type_schema"])
+ assert str(schema) == "Callable[[str, int], str]"
+
+ # Test format_data: bool, int, float, String -> String
+ result = mod.format_data(True, 42, 3.14, "test")
+ assert "test:" in result
+ assert "flag=true" in result
+ assert "count=42" in result
+ assert "value=3.14" in result
+ metadata = mod.get_function_metadata("format_data")
+ assert metadata is not None
+ schema = TypeSchema.from_json_str(metadata["type_schema"])
+ assert str(schema) == "Callable[[bool, int, float, str], str]"
+
+ # Test double_tensor: Tensor, Tensor -> None
+ x = numpy.array([1.0, 2.0, 3.0], dtype=numpy.float32)
+ y = numpy.empty_like(x)
+ mod.double_tensor(x, y)
+ numpy.testing.assert_allclose(y, x * 2.0)
+ metadata = mod.get_function_metadata("double_tensor")
+ assert metadata is not None
+ schema = TypeSchema.from_json_str(metadata["type_schema"])
+ assert str(schema) == "Callable[[Tensor, Tensor], None]"
+
+ # Explicitly cleanup all objects before module unload to avoid
use-after-free
+ del metadata, schema, result, x, y, mod
+ gc.collect()
+
+
+def test_build_inline_with_docstrings() -> None:
+ """Test building functions with documentation using the functions dict."""
+ # Keep module alive until all returned objects are destroyed
+ add_docstring = (
+ "Add two integers and return the sum.\n"
+ "\n"
+ "Parameters\n"
+ "----------\n"
+ "a : int\n"
+ " First integer\n"
+ "b : int\n"
+ " Second integer\n"
+ "\n"
+ "Returns\n"
+ "-------\n"
+ "result : int\n"
+ " Sum of a and b"
+ )
+
+ divide_docstring = "Divides two floats. Returns a/b."
+
+ mod: Module = tvm_ffi.cpp.load_inline(
+ name="test_docs",
+ cpp_sources=r"""
+ int add(int a, int b) {
+ return a + b;
+ }
+
+ int subtract(int a, int b) {
+ return a - b;
+ }
+
+ float divide(float a, float b) {
+ TVM_FFI_ICHECK(b != 0.0f) << "Division by zero";
+ return a / b;
+ }
+ """,
+ functions={
+ "add": add_docstring,
+ "subtract": "", # No documentation
+ "divide": divide_docstring,
+ },
+ extra_cflags=["-DTVM_FFI_DLL_EXPORT_INCLUDE_METADATA=1"],
+ )
+
+ # Test add function with full documentation
+ assert mod.add(10, 5) == 15
+ metadata = mod.get_function_metadata("add")
+ assert metadata is not None
+ schema = TypeSchema.from_json_str(metadata["type_schema"])
+ assert str(schema) == "Callable[[int, int], int]"
+
+ doc = mod.get_function_doc("add")
+ assert doc is not None, "add should have documentation"
+ assert doc == add_docstring
+
+ # Test subtract function without documentation
+ assert mod.subtract(10, 5) == 5
+ metadata = mod.get_function_metadata("subtract")
+ assert metadata is not None
+ schema = TypeSchema.from_json_str(metadata["type_schema"])
+ assert str(schema) == "Callable[[int, int], int]"
+
+ doc = mod.get_function_doc("subtract")
+ assert doc is None, "subtract should not have documentation"
+
+ # Test divide function with short documentation
+ result = mod.divide(10.0, 2.0)
+ assert abs(result - 5.0) < 0.001
+ metadata = mod.get_function_metadata("divide")
+ assert metadata is not None
+ schema = TypeSchema.from_json_str(metadata["type_schema"])
+ assert str(schema) == "Callable[[float, float], float]"
+
+ doc = mod.get_function_doc("divide")
+ assert doc is not None, "divide should have documentation"
+ assert doc == divide_docstring
+
+ # Explicitly cleanup all objects before module unload to avoid
use-after-free
+ del metadata, schema, doc, result, mod
+ gc.collect()
+
+
+def test_build_without_metadata() -> None:
+ """Test building without metadata export."""
+ mod: Module = tvm_ffi.cpp.load_inline(
+ name="test_no_meta",
+ cpp_sources=r"""
+ // Note: NOT defining TVM_FFI_DLL_EXPORT_INCLUDE_METADATA
+
+ int simple_add(int a, int b) {
+ return a + b;
+ }
+ """,
+ functions=["simple_add"],
+ )
+
+ # Function should still work
+ result = mod.simple_add(10, 20)
+ assert result == 30
+
+ # But metadata should not be available
+ metadata = mod.get_function_metadata("simple_add")
+ assert metadata is None, (
+ "Metadata should not be available without
TVM_FFI_DLL_EXPORT_INCLUDE_METADATA"
+ )
+
+ # Doc should also not be available
+ doc = mod.get_function_doc("simple_add")
+ assert doc is None
+
+
if __name__ == "__main__":
pytest.main([__file__])
