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 9c0b869 doc(python): Add example coverage to docstrings (#206)
9c0b869 is described below
commit 9c0b86932aaa1abbe033f8cfec916cd35971d91e
Author: Junru Shao <[email protected]>
AuthorDate: Thu Oct 30 05:07:55 2025 -0700
doc(python): Add example coverage to docstrings (#206)
This PR refines Python-side documentation. It includes:
- Adding example code to every public API as much as possible
- Making existing examples copy-pastable as much as possible
- Polish writing
- Clarifies a few API usage, e.g. system_lib
---
docs/reference/python/index.rst | 8 +-
python/tvm_ffi/_convert.py | 39 ++++++++-
python/tvm_ffi/_dtype.py | 157 +++++++++++++++++++++++++++++++++++--
python/tvm_ffi/_tensor.py | 11 +++
python/tvm_ffi/access_path.py | 29 ++++++-
python/tvm_ffi/base.py | 2 +-
python/tvm_ffi/container.py | 26 +++---
python/tvm_ffi/cython/device.pxi | 5 +-
python/tvm_ffi/cython/function.pxi | 6 +-
python/tvm_ffi/cython/object.pxi | 4 +
python/tvm_ffi/cython/tensor.pxi | 4 +
python/tvm_ffi/error.py | 15 ++--
python/tvm_ffi/module.py | 95 +++++++++++++---------
python/tvm_ffi/registry.py | 66 ++++++++++++++--
python/tvm_ffi/serialization.py | 28 +++++--
python/tvm_ffi/stream.py | 54 ++++++++++---
python/tvm_ffi/utils/lockfile.py | 13 +++
src/ffi/testing/testing.cc | 2 +
18 files changed, 468 insertions(+), 96 deletions(-)
diff --git a/docs/reference/python/index.rst b/docs/reference/python/index.rst
index db01ab3..3a5541e 100644
--- a/docs/reference/python/index.rst
+++ b/docs/reference/python/index.rst
@@ -99,6 +99,7 @@ Stream Context
StreamContext
use_torch_stream
use_raw_stream
+ get_raw_stream
Inline Loading
@@ -118,8 +119,11 @@ Misc
.. autosummary::
:toctree: generated/
- serialization
- access_path
+ serialization.from_json_graph_str
+ serialization.to_json_graph_str
+ access_path.AccessKind
+ access_path.AccessPath
+ access_path.AccessStep
convert
ObjectConvertible
diff --git a/python/tvm_ffi/_convert.py b/python/tvm_ffi/_convert.py
index 1ff91af..43a564c 100644
--- a/python/tvm_ffi/_convert.py
+++ b/python/tvm_ffi/_convert.py
@@ -14,7 +14,7 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""Conversion utilities to bring python objects into ffi values."""
+"""Conversion utilities to convert Python objects into TVM FFI values."""
from __future__ import annotations
@@ -38,18 +38,51 @@ except ImportError:
def convert(value: Any) -> Any: # noqa: PLR0911,PLR0912
- """Convert a python object to ffi values.
+ """Convert a Python object into TVM FFI values.
+
+ This helper mirrors the automatic argument conversion that happens when
+ calling FFI functions. It is primarily useful in tests or places where
+ an explicit conversion is desired.
Parameters
----------
value
- The python object to be converted.
+ The Python object to be converted.
Returns
-------
ffi_obj
The converted TVM FFI object.
+ Examples
+ --------
+ .. code-block:: python
+
+ import tvm_ffi
+
+ # Lists and tuples become tvm_ffi.Array
+ a = tvm_ffi.convert([1, 2, 3])
+ assert isinstance(a, tvm_ffi.Array)
+
+ # Dicts become tvm_ffi.Map
+ m = tvm_ffi.convert({"a": 1, "b": 2})
+ assert isinstance(m, tvm_ffi.Map)
+
+ # Strings and bytes become zero-copy FFI-aware types
+ s = tvm_ffi.convert("hello")
+ b = tvm_ffi.convert(b"bytes")
+ assert isinstance(s, tvm_ffi.core.String)
+ assert isinstance(b, tvm_ffi.core.Bytes)
+
+ # Callables are wrapped as tvm_ffi.Function
+ f = tvm_ffi.convert(lambda x: x + 1)
+ assert isinstance(f, tvm_ffi.Function)
+
+ # Array libraries that support DLPack export can be converted to Tensor
+ import numpy as np
+ x = tvm_ffi.convert(np.arange(4, dtype="int32"))
+ assert isinstance(x, tvm_ffi.Tensor)
+
Note
----
Function arguments to ffi function calls are
diff --git a/python/tvm_ffi/_dtype.py b/python/tvm_ffi/_dtype.py
index de2d1f1..5079226 100644
--- a/python/tvm_ffi/_dtype.py
+++ b/python/tvm_ffi/_dtype.py
@@ -14,7 +14,7 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""dtype class."""
+"""Lightweight ``dtype`` wrapper for TVM FFI."""
# pylint: disable=invalid-name
from __future__ import annotations
@@ -47,11 +47,35 @@ class DataTypeCode(IntEnum):
class dtype(str):
- """TVM FFI dtype class.
+ """Lightweight ``dtype`` in TVM FFI.
+
+ ``dtype`` behaves like a Python ``str`` but also carries an internal FFI
+ representation. You can construct it from strings, NumPy/ML dtypes, or
+ via :py:meth:`from_dlpack_data_type`.
Parameters
----------
dtype_str
+ The string representation of the dtype.
+
+ Examples
+ --------
+ .. code-block:: python
+
+ import tvm_ffi
+
+ # Create from string
+ f32 = tvm_ffi.dtype("float32")
+ assert f32.bits == 32
+ assert f32.itemsize == 4
+
+ # Adjust lanes to create vector types
+ v4f32 = f32.with_lanes(4)
+ assert v4f32 == "float32x4"
+
+ # Round-trip from a DLPack (code, bits, lanes) triple
+ f16 = tvm_ffi.dtype.from_dlpack_data_type((2, 16, 1))
+ assert f16 == "float16"
Note
----
@@ -78,11 +102,31 @@ class dtype(str):
Parameters
----------
dltype_data_type
- The DLPack data type tuple (type_code, bits, lanes).
+ The DLPack data type tuple ``(type_code, bits, lanes)``.
Returns
-------
- The created dtype.
+ dtype
+ The created dtype.
+
+ Examples
+ --------
+ .. code-block:: python
+
+ import tvm_ffi
+
+ # Create float16 and int8 directly from DLPack triples
+ f16 = tvm_ffi.dtype.from_dlpack_data_type((2, 16, 1))
+ i8 = tvm_ffi.dtype.from_dlpack_data_type((0, 8, 1))
+ assert f16 == "float16"
+ assert i8 == "int8"
+
+ See Also
+ --------
+ :py:class:`tvm_ffi.dtype`
+ User-facing dtype wrapper.
+ :py:meth:`tvm_ffi.dtype.with_lanes`
+ Create vector dtypes from a scalar base.
"""
cdtype = core._create_dtype_from_tuple(
@@ -104,13 +148,29 @@ class dtype(str):
Parameters
----------
lanes
- The number of lanes.
+ The number of lanes for the resulting vector type.
Returns
-------
dtype
The new dtype with the given number of lanes.
+ Examples
+ --------
+ .. code-block:: python
+
+ import tvm_ffi
+
+ f32 = tvm_ffi.dtype("float32")
+ v4f32 = f32.with_lanes(4)
+ assert v4f32 == "float32x4"
+ assert v4f32.bits == f32.bits and v4f32.lanes == 4
+
+ See Also
+ --------
+ :py:meth:`tvm_ffi.dtype.from_dlpack_data_type`
+ Construct from a DLPack ``(code, bits, lanes)`` triple.
+
"""
cdtype = core._create_dtype_from_tuple(
core.DataType,
@@ -124,18 +184,105 @@ class dtype(str):
@property
def itemsize(self) -> int:
+ """Size of one element in bytes.
+
+ The size is computed as ``bits * lanes // 8``. When the number of
+ lanes is greater than 1, the ``itemsize`` represents the byte size
+ of the vector element.
+
+ Examples
+ --------
+ .. code-block:: python
+
+ import tvm_ffi
+
+ assert tvm_ffi.dtype("float32").itemsize == 4
+ assert tvm_ffi.dtype("float32").with_lanes(4).itemsize == 16
+
+ See Also
+ --------
+ :py:attr:`tvm_ffi.dtype.bits`
+ Bit width of the scalar base type.
+ :py:attr:`tvm_ffi.dtype.lanes`
+ Number of lanes for vector types.
+ :py:meth:`tvm_ffi.dtype.with_lanes`
+ Create a vector dtype from a scalar base.
+
+ """
return self._tvm_ffi_dtype.itemsize
@property
def type_code(self) -> int:
+ """Integer DLDataTypeCode of the scalar base type.
+
+ Examples
+ --------
+ .. code-block:: python
+
+ import tvm_ffi
+
+ f32 = tvm_ffi.dtype("float32")
+ # The type code is an integer following DLPack conventions
+ assert isinstance(f32.type_code, int)
+ # Consistent with constructing from an explicit (code, bits, lanes)
+ assert f32.type_code == tvm_ffi.dtype.from_dlpack_data_type((2,
32, 1)).type_code
+
+ See Also
+ --------
+ :py:meth:`tvm_ffi.dtype.from_dlpack_data_type`
+ Construct a dtype from a DLPack ``(code, bits, lanes)`` triple.
+
+ """
return self._tvm_ffi_dtype.type_code
@property
def bits(self) -> int:
+ """Number of bits of the scalar base type.
+
+ Examples
+ --------
+ .. code-block:: python
+
+ import tvm_ffi
+
+ assert tvm_ffi.dtype("int8").bits == 8
+ v4f32 = tvm_ffi.dtype("float32").with_lanes(4)
+ assert v4f32.bits == 32 # per-lane bit width
+
+ See Also
+ --------
+ :py:attr:`tvm_ffi.dtype.itemsize`
+ Byte size accounting for lanes.
+ :py:attr:`tvm_ffi.dtype.lanes`
+ Number of lanes for vector types.
+
+ """
return self._tvm_ffi_dtype.bits
@property
def lanes(self) -> int:
+ """Number of lanes (for vector types).
+
+ Returns ``1`` for scalar dtypes and the lane count for vector dtypes
+ created via :py:meth:`tvm_ffi.dtype.with_lanes`.
+
+ Examples
+ --------
+ .. code-block:: python
+
+ import tvm_ffi
+
+ assert tvm_ffi.dtype("float32").lanes == 1
+ assert tvm_ffi.dtype("float32").with_lanes(4).lanes == 4
+
+ See Also
+ --------
+ :py:meth:`tvm_ffi.dtype.with_lanes`
+ Create a vector dtype from a scalar base.
+ :py:attr:`tvm_ffi.dtype.itemsize`
+ Byte size accounting for lanes.
+
+ """
return self._tvm_ffi_dtype.lanes
diff --git a/python/tvm_ffi/_tensor.py b/python/tvm_ffi/_tensor.py
index 4c9714b..5005bc6 100644
--- a/python/tvm_ffi/_tensor.py
+++ b/python/tvm_ffi/_tensor.py
@@ -43,6 +43,16 @@ class Shape(tuple, PyNativeObject):
This class subclasses :class:`tuple` so it can be used in most places where
:class:`tuple` is used in Python array APIs.
+ Examples
+ --------
+ .. code-block:: python
+
+ import numpy as np
+ import tvm_ffi
+
+ x = tvm_ffi.from_dlpack(np.arange(6, dtype="int32").reshape(2, 3))
+ assert x.shape == (2, 3)
+
"""
_tvm_ffi_cached_object: Any
@@ -85,6 +95,7 @@ def device(device_type: str | int | DLDeviceType, index: int
| None = None) -> D
.. code-block:: python
+ import tvm_ffi
assert tvm_ffi.device("cuda:0") == tvm_ffi.device("cuda", 0)
assert tvm_ffi.device("cpu:0") == tvm_ffi.device("cpu", 0)
diff --git a/python/tvm_ffi/access_path.py b/python/tvm_ffi/access_path.py
index a048a8a..6b9db9c 100644
--- a/python/tvm_ffi/access_path.py
+++ b/python/tvm_ffi/access_path.py
@@ -53,7 +53,25 @@ class AccessStep(Object):
@register_object("ffi.reflection.AccessPath")
class AccessPath(Object):
- """Access path container."""
+ """Access path container.
+
+ An ``AccessPath`` describes how to reach a nested attribute or item
+ inside a complex FFI object by recording a sequence of steps
+ (attribute, array index, or map key). It is primarily used by
+ diagnostics to pinpoint structural mismatches.
+
+ Examples
+ --------
+ .. code-block:: python
+
+ from tvm_ffi.access_path import AccessPath
+
+ root = AccessPath.root()
+ # Build a path equivalent to obj.layer.weight[2]
+ p = root.attr("layer").attr("weight").array_item(2)
+ assert isinstance(p, AccessPath)
+
+ """
# tvm-ffi-stubgen(begin): object/ffi.reflection.AccessPath
if TYPE_CHECKING:
@@ -86,7 +104,14 @@ class AccessPath(Object):
@staticmethod
def root() -> AccessPath:
- """Create a root access path."""
+ """Create a root access path.
+
+ Returns
+ -------
+ AccessPath
+ A path representing the root of an object graph.
+
+ """
return AccessPath._root()
def __eq__(self, other: Any) -> bool:
diff --git a/python/tvm_ffi/base.py b/python/tvm_ffi/base.py
index 45c9be2..2e2ece3 100644
--- a/python/tvm_ffi/base.py
+++ b/python/tvm_ffi/base.py
@@ -41,7 +41,7 @@ if sys.version_info[:2] < (3, 8): # noqa: UP036
def _load_lib() -> ctypes.CDLL:
- """Load libary by searching possible path."""
+ """Load the tvm_ffi shared library by searching likely paths."""
lib_path = libinfo.find_libtvm_ffi()
# The dll search path need to be added explicitly in windows
if sys.platform.startswith("win32"):
diff --git a/python/tvm_ffi/container.py b/python/tvm_ffi/container.py
index a23fa47..90a84a2 100644
--- a/python/tvm_ffi/container.py
+++ b/python/tvm_ffi/container.py
@@ -124,7 +124,7 @@ def getitem_helper(
@register_object("ffi.Array")
class Array(core.Object, Sequence[T]):
- """Array container that represents a sequence of values in ffi.
+ """Array container that represents a sequence of values in the FFI.
:py:func:`tvm_ffi.convert` will map python list/tuple to this class.
@@ -133,19 +133,18 @@ class Array(core.Object, Sequence[T]):
input_list
The list of values to be stored in the array.
- See Also
- --------
- :py:func:`tvm_ffi.convert`
-
Examples
--------
.. code-block:: python
import tvm_ffi
- a = tvm_ffi.convert([1, 2, 3])
- assert isinstance(a, tvm_ffi.Array)
- assert len(a) == 3
+ a = tvm_ffi.Array([1, 2, 3])
+ assert tuple(a) == (1, 2, 3)
+
+ See Also
+ --------
+ :py:func:`tvm_ffi.convert`
"""
@@ -274,22 +273,21 @@ class Map(core.Object, Mapping[K, V]):
input_dict
The dictionary of values to be stored in the map.
- See Also
- --------
- :py:func:`tvm_ffi.convert`
-
Examples
--------
.. code-block:: python
import tvm_ffi
- amap = tvm_ffi.convert({"a": 1, "b": 2})
- assert isinstance(amap, tvm_ffi.Map)
+ amap = tvm_ffi.Map({"a": 1, "b": 2})
assert len(amap) == 2
assert amap["a"] == 1
assert amap["b"] == 2
+ See Also
+ --------
+ :py:func:`tvm_ffi.convert`
+
"""
def __init__(self, input_dict: Mapping[K, V]) -> None:
diff --git a/python/tvm_ffi/cython/device.pxi b/python/tvm_ffi/cython/device.pxi
index e23fee7..413ef40 100644
--- a/python/tvm_ffi/cython/device.pxi
+++ b/python/tvm_ffi/cython/device.pxi
@@ -35,7 +35,7 @@ def _create_device_from_tuple(cls, device_type, device_id):
class DLDeviceType(IntEnum):
- """Enumeration mirroring DLPack's ``DLDeviceType``.
+ """Enumeration mirroring DLPack's `DLDeviceType
<https://dmlc.github.io/dlpack/latest/c_api.html#c.DLDeviceType>`_
Values can be compared against :py:meth:`Device.dlpack_device_type`.
@@ -43,6 +43,7 @@ class DLDeviceType(IntEnum):
--------
.. code-block:: python
+ import tvm_ffi
dev = tvm_ffi.device("cuda", 0)
assert dev.dlpack_device_type() == tvm_ffi.DLDeviceType.kDLCUDA
@@ -83,6 +84,8 @@ cdef class Device:
--------
.. code-block:: python
+ import tvm_ffi
+
dev = tvm_ffi.device("cuda:0")
assert dev.type == "cuda"
assert dev.index == 0
diff --git a/python/tvm_ffi/cython/function.pxi
b/python/tvm_ffi/cython/function.pxi
index bf84091..4f9fca6 100644
--- a/python/tvm_ffi/cython/function.pxi
+++ b/python/tvm_ffi/cython/function.pxi
@@ -769,8 +769,10 @@ cdef class Function(Object):
See Also
--------
- tvm_ffi.register_global_func: How to register global function.
- tvm_ffi.get_global_func: How to get global function.
+ :py:func:`tvm_ffi.register_global_func`
+ Register a Python callable as a global FFI function.
+ :py:func:`tvm_ffi.get_global_func`
+ Look up a previously registered global FFI function by name.
"""
cdef int c_release_gil
cdef dict __dict__
diff --git a/python/tvm_ffi/cython/object.pxi b/python/tvm_ffi/cython/object.pxi
index 9042cdc..0c5bdac 100644
--- a/python/tvm_ffi/cython/object.pxi
+++ b/python/tvm_ffi/cython/object.pxi
@@ -101,6 +101,8 @@ cdef class Object:
.. code-block:: python
+ import tvm_ffi.testing
+
# Acquire a testing object constructed through FFI
obj = tvm_ffi.testing.create_object("testing.TestObjectBase", v_i64=12)
assert isinstance(obj, tvm_ffi.Object)
@@ -214,6 +216,8 @@ cdef class Object:
--------
.. code-block:: python
+ import tvm_ffi.testing
+
x = tvm_ffi.testing.create_object("testing.TestObjectBase")
y = x
z = tvm_ffi.testing.create_object("testing.TestObjectBase")
diff --git a/python/tvm_ffi/cython/tensor.pxi b/python/tvm_ffi/cython/tensor.pxi
index 0879073..0985056 100644
--- a/python/tvm_ffi/cython/tensor.pxi
+++ b/python/tvm_ffi/cython/tensor.pxi
@@ -177,6 +177,8 @@ def from_dlpack(
.. code-block:: python
import numpy as np
+ import tvm_ffi
+
x_np = np.arange(8, dtype="int32")
x = tvm_ffi.from_dlpack(x_np)
y_np = np.from_dlpack(x)
@@ -218,6 +220,8 @@ cdef class Tensor(Object):
.. code-block:: python
import numpy as np
+ import tvm_ffi
+
x = tvm_ffi.from_dlpack(np.arange(6, dtype="int32"))
assert x.shape == (6,)
assert x.dtype == tvm_ffi.dtype("int32")
diff --git a/python/tvm_ffi/error.py b/python/tvm_ffi/error.py
index 14dd2b0..015fe80 100644
--- a/python/tvm_ffi/error.py
+++ b/python/tvm_ffi/error.py
@@ -178,12 +178,17 @@ def register_error(
--------
.. code-block:: python
- @tvm.error.register_error
- class MyError(RuntimeError):
- pass
+ import tvm_ffi
- err_inst = tvm.error.create_ffi_error("MyError: xyz")
- assert isinstance(err_inst, MyError)
+ # Register a custom Python exception so tvm_ffi.Error maps to it
+ @tvm_ffi.error.register_error
+ class MyError(RuntimeError):
+ pass
+
+ # Convert a Python exception to an FFI Error and back
+ ffi_err = tvm_ffi.convert(MyError("boom"))
+ py_err = ffi_err.py_error()
+ assert isinstance(py_err, MyError)
"""
if callable(name_or_cls):
diff --git a/python/tvm_ffi/module.py b/python/tvm_ffi/module.py
index 3b09df7..73970e0 100644
--- a/python/tvm_ffi/module.py
+++ b/python/tvm_ffi/module.py
@@ -115,13 +115,14 @@ class Module(core.Object):
def implements_function(self, name: str, query_imports: bool = False) ->
bool:
"""Return True if the module defines a global function.
- Note
- ----
- that has_function(name) does not imply get_function(name) is non-null
since the module
- that has_function(name) does not imply get_function(name) is non-null
since the module
- may be, eg, a CSourceModule which cannot supply a packed-func
implementation of the function
- without further compilation. However, get_function(name) non null
should always imply
- has_function(name).
+ Notes
+ -----
+ ``implements_function(name)`` does not guarantee that
+ ``get_function(name)`` will return a callable, because some module
+ kinds (e.g. a source-only module) may not provide a packed function
+ implementation until further compilation occurs. However, a non-null
+ result from ``get_function(name)`` should imply the module implements
+ the function.
Parameters
----------
@@ -133,8 +134,7 @@ class Module(core.Object):
Returns
-------
- b
- True if module (or one of its imports) has a definition for name.
+ True if module (or one of its imports) has a definition for name.
"""
return _ffi_api.ModuleImplementsFunction(self, name, query_imports)
@@ -157,12 +157,11 @@ class Module(core.Object):
The name of the function
query_imports
- Whether also query modules imported by this module.
+ Whether to also query modules imported by this module.
Returns
-------
- f
- The result function.
+ The result function.
"""
func = _ffi_api.ModuleGetFunction(self, name, query_imports)
@@ -203,8 +202,7 @@ class Module(core.Object):
Returns
-------
- source
- The result source code.
+ The result source code.
"""
return _ffi_api.ModuleInspectSource(self, fmt)
@@ -218,8 +216,7 @@ class Module(core.Object):
Returns
-------
- mask
- Bitmask of runtime module property
+ Bitmask of runtime module property
"""
return _ffi_api.ModuleGetPropertyMask(self)
@@ -229,8 +226,7 @@ class Module(core.Object):
Returns
-------
- b
- True if the module is binary serializable.
+ True if the module is binary serializable.
"""
return (self.get_property_mask() &
ModulePropertyMask.BINARY_SERIALIZABLE) != 0
@@ -240,8 +236,7 @@ class Module(core.Object):
Returns
-------
- b
- True if the module is runnable.
+ True if the module is runnable.
"""
return (self.get_property_mask() & ModulePropertyMask.RUNNABLE) != 0
@@ -253,8 +248,7 @@ class Module(core.Object):
Returns
-------
- b
- True if the module is compilation exportable.
+ True if the module is compilation exportable.
"""
return (self.get_property_mask() &
ModulePropertyMask.COMPILATION_EXPORTABLE) != 0
@@ -275,21 +269,20 @@ class Module(core.Object):
See Also
--------
- runtime.Module.export_library : export the module to shared library.
+ tvm.runtime.Module.export_library : export the module to shared
library.
"""
_ffi_api.ModuleWriteToFile(self, file_name, fmt)
def system_lib(symbol_prefix: str = "") -> Module:
- """Get system-wide library module singleton.
+ """Get system-wide library module singleton with functions prefixed by
``__tvm_ffi_{symbol_prefix}``.
- System lib is a global module that contains self register functions in
startup.
- Unlike normal dso modules which need to be loaded explicitly.
- It is useful in environments where dynamic loading api like dlopen is
banned.
+ The library module contains symbols that are registered via
:cpp:func:`TVMFFIEnvModRegisterSystemLibSymbol`.
- The system lib is intended to be linked and loaded during the entire
life-cyle of the program.
- If you want dynamic loading features, use dso modules instead.
+ .. note::
+ The system lib is intended to be statically linked and loaded during
the entire lifecycle of the program.
+ If you want dynamic loading features, use DSO modules instead.
Parameters
----------
@@ -299,8 +292,36 @@ def system_lib(symbol_prefix: str = "") -> Module:
Returns
-------
- module
- The system-wide library module.
+ The system-wide library module.
+
+ Examples
+ --------
+ Register the function ``test_symbol_add_one`` in C++ with the name
``__tvm_ffi_test_symbol_add_one``
+ via :cpp:func:`TVMFFIEnvModRegisterSystemLibSymbol`.
+
+ .. code-block:: cpp
+
+ // A function to be registered in the system lib
+ static int test_symbol_add_one(void*, const TVMFFIAny* args, int32_t
num_args, TVMFFIAny* ret) {
+ TVM_FFI_SAFE_CALL_BEGIN();
+ TVM_FFI_CHECK(num_args == 1, "Expected 1 argument, but got: " +
std::to_string(num_args));
+ int64_t x = reinterpret_cast<const
AnyView*>(args)[0].cast<int64_t>();
+ reinterpret_cast<Any*>(ret)[0] = x + 1;
+ TVM_FFI_SAFE_CALL_END();
+ }
+
+ // Register the function with name `test_symbol_add_one` prefixed by
`__tvm_ffi_`
+ int _ =
TVMFFIEnvModRegisterSystemLibSymbol("__tvm_ffi_testing.add_one",
reinterpret_cast<void*>(test_symbol_add_one));
+
+ Look up and call the function from Python:
+
+ .. code-block:: python
+
+ import tvm_ffi
+
+ mod: tvm_ffi.Module = tvm_ffi.system_lib("testing.") # symbols
prefixed with `__tvm_ffi_testing.`
+ func: tvm_ffi.Function = mod["add_one"] # looks up
`__tvm_ffi_testing.add_one`
+ assert func(10) == 11
"""
return _ffi_api.SystemLib(symbol_prefix)
@@ -311,25 +332,25 @@ def load_module(path: str | PathLike) -> Module:
Parameters
----------
- path : str | PathLike
+ path
The path to the module file.
Returns
-------
- module
- The loaded module
+ The loaded module
Examples
--------
.. code-block:: python
+ import tvm_ffi
+ from pathlib import Path
+
# Works with string paths
mod = tvm_ffi.load_module("path/to/module.so")
- mod.func_name(*args)
-
# Also works with pathlib.Path objects
- from pathlib import Path
mod = tvm_ffi.load_module(Path("path/to/module.so"))
+
mod.func_name(*args)
See Also
diff --git a/python/tvm_ffi/registry.py b/python/tvm_ffi/registry.py
index bd37ebd..d149488 100644
--- a/python/tvm_ffi/registry.py
+++ b/python/tvm_ffi/registry.py
@@ -35,12 +35,13 @@ def register_object(type_key: str | type | None = None) ->
Callable[[type], type
Parameters
----------
type_key
- The type key of the node
+ The type key of the node. It requires ``type_key`` to be registered
already
+ on the C++ side. If not specified, the class name will be used.
Examples
--------
- The following code registers MyObject
- using type key "test.MyObject"
+ The following code registers MyObject using type key "test.MyObject", if
the
+ type key is already registered on the C++ side.
.. code-block:: python
@@ -165,7 +166,20 @@ def get_global_func(name: str, allow_missing: bool =
False) -> core.Function | N
Returns
-------
func
- The function to be returned, None if function is missing.
+ The function to be returned, ``None`` if function is missing.
+
+ Examples
+ --------
+ .. code-block:: python
+
+ import tvm_ffi
+
+ @tvm_ffi.register_global_func("demo.echo")
+ def echo(x):
+ return x
+
+ f = tvm_ffi.get_global_func("demo.echo")
+ assert f(123) == 123
See Also
--------
@@ -195,24 +209,60 @@ def remove_global_func(name: str) -> None:
Parameters
----------
name
- The name of the global function
+ The name of the global function.
+
+ Examples
+ --------
+ .. code-block:: python
+
+ import tvm_ffi
+
+ @tvm_ffi.register_global_func("my.temp")
+ def temp():
+ return 42
+
+ assert tvm_ffi.get_global_func("my.temp", allow_missing=True) is not
None
+ tvm_ffi.remove_global_func("my.temp")
+ assert tvm_ffi.get_global_func("my.temp", allow_missing=True) is None
+
+ See Also
+ --------
+ :py:func:`tvm_ffi.register_global_func`
+ :py:func:`tvm_ffi.get_global_func`
"""
get_global_func("ffi.FunctionRemoveGlobal")(name)
def get_global_func_metadata(name: str) -> dict[str, Any]:
- """Get the type schema string of a global function by name.
+ """Get metadata (including type schema) for a global function.
Parameters
----------
name
- The name of the global function
+ The name of the global function.
Returns
-------
metadata
- The metadata of the function
+ A dictionary containing function metadata. The ``type_schema`` field
+ encodes the callable signature.
+
+ Examples
+ --------
+ .. code-block:: python
+
+ import tvm_ffi
+
+ meta = tvm_ffi.get_global_func_metadata("testing.add_one")
+ print(meta)
+
+ See Also
+ --------
+ :py:func:`tvm_ffi.get_global_func`
+ Retrieve a callable for an existing global function.
+ :py:func:`tvm_ffi.register_global_func`
+ Register a Python callable as a global FFI function.
"""
return json.loads(get_global_func("ffi.GetGlobalFuncMetadata")(name))
diff --git a/python/tvm_ffi/serialization.py b/python/tvm_ffi/serialization.py
index 285d764..66679ff 100644
--- a/python/tvm_ffi/serialization.py
+++ b/python/tvm_ffi/serialization.py
@@ -14,7 +14,12 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""Serialization related utilities to enable some object can be pickled."""
+"""Utilities for serializing and deserializing FFI object graphs.
+
+These helpers produce a stable JSON graph representation that preserves
+object identity and references. It is useful for debugging and for
+lightweight persistence when pickling is not available.
+"""
from __future__ import annotations
@@ -26,9 +31,9 @@ from . import _ffi_api
def to_json_graph_str(obj: Any, metadata: dict[str, Any] | None = None) -> str:
"""Dump an object to a JSON graph string.
- The JSON graph string is a string representation of of the object
- graph includes the reference information of same objects, which can
- be used for serialization and debugging.
+ The JSON graph is a textual representation of the object graph that
+ preserves shared references. It can be used for debugging or simple
+ persistence.
Parameters
----------
@@ -43,6 +48,17 @@ def to_json_graph_str(obj: Any, metadata: dict[str, Any] |
None = None) -> str:
json_str
The JSON graph string.
+ Examples
+ --------
+ .. code-block:: python
+
+ import tvm_ffi
+
+ a = tvm_ffi.convert([1, 2, 3])
+ s = tvm_ffi.serialization.to_json_graph_str(a)
+ b = tvm_ffi.serialization.from_json_graph_str(s)
+ assert list(b) == [1, 2, 3]
+
"""
return _ffi_api.ToJSONGraphString(obj, metadata)
@@ -50,8 +66,8 @@ def to_json_graph_str(obj: Any, metadata: dict[str, Any] |
None = None) -> str:
def from_json_graph_str(json_str: str) -> Any:
"""Load an object from a JSON graph string.
- The JSON graph string is a string representation of of the object
- graph that also includes the reference information.
+ The JSON graph string is produced by :py:func:`to_json_graph_str` and
+ can be converted back into the corresponding FFI-backed objects.
Parameters
----------
diff --git a/python/tvm_ffi/stream.py b/python/tvm_ffi/stream.py
index ee71e5d..a805378 100644
--- a/python/tvm_ffi/stream.py
+++ b/python/tvm_ffi/stream.py
@@ -124,7 +124,8 @@ try:
Note
----
- When working with raw cudaStream_t handle, using
:py:func:`tvm_ffi.use_raw_stream` instead.
+ When working with a raw ``cudaStream_t`` handle, use
+ :py:func:`tvm_ffi.use_raw_stream` instead.
"""
return TorchStreamContext(context)
@@ -137,7 +138,7 @@ except ImportError:
def use_raw_stream(device: core.Device, stream: int | c_void_p) ->
StreamContext:
- """Create a ffi stream context with given device and stream handle.
+ """Create an FFI stream context with the given device and stream handle.
Parameters
----------
@@ -145,28 +146,45 @@ def use_raw_stream(device: core.Device, stream: int |
c_void_p) -> StreamContext
The device to which the stream belongs.
stream
- The stream handle.
+ The stream handle (for example, a CUDA ``cudaStream_t`` as an integer,
or ``0``).
Returns
-------
context
- The ffi stream context.
+ The FFI stream context.
+
+ Examples
+ --------
+ The example below uses a CPU device and a dummy stream handle. On CUDA,
pass a
+ real ``cudaStream_t`` integer.
+
+ .. code-block:: python
- Note
- ----
- When working with torch stram or cuda graph, using
:py:func:`tvm_ffi.use_torch_stream` instead.
+ import tvm_ffi
+
+ dev = tvm_ffi.device("cpu:0")
+ with tvm_ffi.use_raw_stream(dev, 0):
+ # Within the context, the current stream for this device is set
+ assert tvm_ffi.get_raw_stream(dev) == 0
+
+ See Also
+ --------
+ :py:func:`tvm_ffi.use_torch_stream`
+ Use a Torch stream or CUDA graph as the source of truth.
+ :py:func:`tvm_ffi.get_raw_stream`
+ Query the current FFI stream for a device.
"""
if not isinstance(stream, (int, c_void_p)):
raise ValueError(
- "use_raw_stream only accepts int or c_void_p as stram input, "
+ "use_raw_stream only accepts int or c_void_p as stream input, "
"try use_torch_stream when using torch.cuda.Stream or
torch.cuda.graph"
)
return StreamContext(device, stream)
def get_raw_stream(device: core.Device) -> int:
- """Get the current ffi stream of given device.
+ """Get the current FFI stream of a given device.
Parameters
----------
@@ -176,7 +194,23 @@ def get_raw_stream(device: core.Device) -> int:
Returns
-------
stream
- The current ffi stream.
+ The current FFI stream as an integer handle.
+
+ Examples
+ --------
+ .. code-block:: python
+
+ import tvm_ffi
+
+ dev = tvm_ffi.device("cpu:0")
+ # Default stream is implementation-defined; set it explicitly
+ with tvm_ffi.use_raw_stream(dev, 0):
+ assert tvm_ffi.get_raw_stream(dev) == 0
+
+ See Also
+ --------
+ :py:func:`tvm_ffi.use_raw_stream`
+ Set the current stream for a device.
"""
return core._env_get_current_stream(device.dlpack_device_type(),
device.index)
diff --git a/python/tvm_ffi/utils/lockfile.py b/python/tvm_ffi/utils/lockfile.py
index 6efe80f..da98491 100644
--- a/python/tvm_ffi/utils/lockfile.py
+++ b/python/tvm_ffi/utils/lockfile.py
@@ -35,6 +35,19 @@ class FileLock:
This class implements an advisory lock, which must be respected by all
cooperating processes.
+
+ Examples
+ --------
+ .. code-block:: python
+
+ from tvm_ffi.utils import FileLock
+
+ with FileLock("/tmp/my.lock"):
+ # Critical section guarded by the lock.
+ # Other processes attempting to acquire the same lock will block
+ # (or fail, if using ``acquire()``) until this context exits.
+ do_work()
+
"""
def __init__(self, lock_file_path: str) -> None:
diff --git a/src/ffi/testing/testing.cc b/src/ffi/testing/testing.cc
index cccd89e..bccc5c0 100644
--- a/src/ffi/testing/testing.cc
+++ b/src/ffi/testing/testing.cc
@@ -446,6 +446,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.def("testing.schema_no_args", []() { return 1; })
.def("testing.schema_no_return", [](int64_t x) {})
.def("testing.schema_no_args_no_return", []() {});
+ TVMFFIEnvModRegisterSystemLibSymbol("__tvm_ffi_testing.add_one",
+
reinterpret_cast<void*>(__add_one_c_symbol));
}
} // namespace ffi
