[FFI][REFACTOR] Cleanup tvm_ffi python API and types (apache#18277)

This PR cleans up the python API to make things more consistent
with existing python array api and torch.

Device update
- device_id => index, to be consistent with torch
- device_type => dlpack_device_type() returns int
- added type property same as torch.device

API updates:

- Move the convenient method like cpu() out into tvm runtime to keep device minimal
- tvm_ffi._init_api => tvm_ffi.init_ffi_api
- tvm_ffi.register_func => tvm_ffi.register_global_func

Author: tqchen

docs/get_started/quick_start.md

@@ -144,8 +144,9 @@ TVM_FFI_DLL_EXPORT_TYPED_FUNC(add_one_cuda, tvm_ffi_example::AddOneCUDA);
 ### Working with PyTorch
 
 Atfer build, we will create library such as `build/add_one_cuda.so`, that can be loaded by
-with api `tvm_ffi.load_module`. Then the function will become available as property of the loaded module.
-The tensor arguments in the ffi functions automatically consumes torch.Tensor. The following code shows how
+with api {py:func}`tvm_ffi.load_module` that returns a {py:class}`tvm_ffi.Module`
+Then the function will become available as property of the loaded module.
+The tensor arguments in the ffi functions automatically consumes `torch.Tensor`. The following code shows how
 to use the function in torch.
 
 ```python

docs/guides/packaging.md

@@ -204,7 +204,7 @@ _LIB = _load_lib()
 
 Effectively, it leverages the `tvm_ffi.load_module` call to load the library
 extension DLL shipped along with the package. The `_ffi_api.py` contains a function
-call to `tvm_ffi._init_api` that registers all global functions prefixed
+call to `tvm_ffi.init_ffi_api` that registers all global functions prefixed
 with `my_ffi_extension` into the module.
 
 ```python
@@ -214,7 +214,7 @@ from .base import _LIB
 
 # Register all global functions prefixed with 'my_ffi_extension.'
 # This makes functions registered via TVM_FFI_STATIC_INIT_BLOCK available
-tvm_ffi._init_api("my_ffi_extension", __name__)
+tvm_ffi.init_ffi_api("my_ffi_extension", __name__)
 ```
 
 Then we can redirect the calls to the related functions.

docs/guides/python_guide.md

@@ -47,7 +47,7 @@ y = np.empty_like(x)
 mod.add_one_cpu(x, y)
 ```
 
-In this case, `tvm_ffi.load_module` will return a `tvm_ffi.Module` class that contains
+In this case, {py:func}`tvm_ffi.load_module` will return a {py:class}`tvm_ffi.Module` class that contains
 the exported functions. You can access the functions by their names.
 
 ## Tensor
@@ -67,12 +67,12 @@ np_result = np.from_dlpack(tvm_array)
 
 In most cases, however, you do not have to explicitly create Tensors.
 The Python interface can take in `torch.Tensor` and `numpy.ndarray` objects
-and automatically convert them to `tvm_ffi.Tensor`.
+and automatically convert them to {py:class}`tvm_ffi.Tensor`.
 
 ## Functions and Callbacks
 
-`tvm_ffi.Function` provides the Python interface for `ffi::Function` in the C++.
-You can retrieve globally registered functions via `tvm_ffi.get_global_func()`.
+{py:class}`tvm_ffi.Function` provides the Python interface for `ffi::Function` in the C++.
+You can retrieve globally registered functions via {py:func}`tvm_ffi.get_global_func`.
 
 ```python
 import tvm_ffi
@@ -84,8 +84,8 @@ assert fecho(1) == 1
 ```
 
 You can pass a Python function as an argument to another FFI function as callbacks.
-Under the hood, `tvm_ffi.convert` is called to convert the Python function into a
-`tvm_ffi.Function`.
+Under the hood, {py:func}`tvm_ffi.convert` is called to convert the Python function into a
+{py:class}`tvm_ffi.Function`.
 
 ```python
 import tvm_ffi
@@ -103,7 +103,7 @@ You can also register a Python callback as a global function.
 ```python
 import tvm_ffi
 
-@tvm_ffi.register_func("example.add_one")
+@tvm_ffi.register_global_func("example.add_one")
 def add_one(a):
     return a + 1
 
@@ -112,7 +112,7 @@ assert tvm_ffi.get_global_func("example.add_one")(1) == 2
 
 ## Container Types
 
-When an FFI function takes arguments from lists/tuples, they will be converted into `tvm_ffi.Array`.
+When an FFI function takes arguments from lists/tuples, they will be converted into {py:class}`tvm_ffi.Array`.
 
 ```python
 import tvm_ffi
@@ -124,7 +124,7 @@ assert len(arr) == 4
 assert arr[0] == 1
 ```
 
-Dictionaries will be converted to `tvm_ffi.Map`
+Dictionaries will be converted to {py:class}`tvm_ffi.Map`
 
 ```python
 import tvm_ffi
@@ -167,7 +167,7 @@ File "src/ffi/extra/testing.cc", line 60, in void tvm::ffi::TestRaiseError(tvm::
   throw ffi::Error(kind, msg, TVMFFITraceback(__FILE__, __LINE__, TVM_FFI_FUNC_SIG, 0));
 ```
 
-We register common error kinds. You can also register extra error dispatch via the `tvm_ffi.register_error` function.
+We register common error kinds. You can also register extra error dispatch via the {py:func}`tvm_ffi.register_error` function.
 
 ## Advanced: Register Your Own Object
 
@@ -239,5 +239,5 @@ assert test_int_pair.b == 2
 Under the hood, we leverage the information registered through the reflection registry to
 generate efficient field accessors and methods for each class.
 
-Importantly, when you have multiple inheritance, you need to call `tvm_ffi.register_object`
+Importantly, when you have multiple inheritance, you need to call {py:func}`tvm_ffi.register_object`
 on both the base class and the child class.

docs/index.rst

@@ -39,3 +39,9 @@ Apache TVM FFI Documentation
    :caption: Concepts
 
    concepts/abi_overview.md
+
+.. toctree::
+   :maxdepth: 1
+   :caption: Reference
+
+   reference/python/index.rst

docs/reference/python/index.rst

@@ -0,0 +1,69 @@
+..  Licensed to the Apache Software Foundation (ASF) under one
+    or more contributor license agreements.  See the NOTICE file
+    distributed with this work for additional information
+    regarding copyright ownership.  The ASF licenses this file
+    to you under the Apache License, Version 2.0 (the
+    "License"); you may not use this file except in compliance
+    with the License.  You may obtain a copy of the License at
+
+..    http://www.apache.org/licenses/LICENSE-2.0
+
+..  Unless required by applicable law or agreed to in writing,
+    software distributed under the License is distributed on an
+    "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+    KIND, either express or implied.  See the License for the
+    specific language governing permissions and limitations
+    under the License.
+
+Python API
+==========
+
+.. automodule:: tvm_ffi
+  :no-members:
+
+.. currentmodule:: tvm_ffi
+
+Object
+------
+.. autosummary::
+  :toctree: generated/
+
+  Object
+  register_object
+
+
+Function and Module
+-------------------
+.. autosummary::
+  :toctree: generated/
+
+
+  Function
+  Module
+  register_global_func
+  get_global_func
+  system_lib
+  load_module
+  init_ffi_api
+  register_error
+  convert
+
+
+Tensor
+------
+.. autosummary::
+  :toctree: generated/
+
+  Shape
+  Tensor
+  Device
+  from_dlpack
+
+
+Containers
+----------
+.. autosummary::
+  :toctree: generated/
+
+  Array
+  Map

examples/packaging/python/my_ffi_extension/_ffi_api.py

@@ -21,4 +21,4 @@
 
 # this is a short cut to register all the global functions
 # prefixed by `my_ffi_extension.` to this module
-tvm_ffi._init_api("my_ffi_extension", __name__)
+tvm_ffi.init_ffi_api("my_ffi_extension", __name__)

python/tvm_ffi/__init__.py

@@ -20,50 +20,43 @@
 from . import libinfo
 
 # package init part
-from .registry import register_object, register_func, get_global_func, _init_api
-from .dtype import dtype, DataTypeCode
-from .core import String, Bytes
-from .core import Object, ObjectGeneric, Function
-from .convert import convert
+from .registry import (
+    register_object,
+    register_global_func,
+    get_global_func,
+    remove_global_func,
+    init_ffi_api,
+)
+from ._dtype import dtype
+from .core import Object, ObjectConvertible, Function
+from ._convert import convert
 from .error import register_error
-from .tensor import Device, device
-from .tensor import cpu, cuda, rocm, opencl, metal, vpi, vulkan, ext_dev, hexagon, webgpu
-from .tensor import from_dlpack, Tensor, Shape
+from ._tensor import Device, device, DLDeviceType
+from ._tensor import from_dlpack, Tensor, Shape
 from .container import Array, Map
-from .module import Module, ModulePropertyMask, system_lib, load_module
+from .module import Module, system_lib, load_module
 from . import serialization
 from . import access_path
 from . import testing
 
 
 __all__ = [
     "dtype",
-    "DataTypeCode",
     "Device",
     "Object",
     "register_object",
-    "register_func",
+    "register_global_func",
     "get_global_func",
-    "_init_api",
+    "remove_global_func",
+    "init_ffi_api",
     "Object",
-    "ObjectGeneric",
+    "ObjectConvertible",
     "Function",
     "convert",
-    "String",
-    "Bytes",
     "register_error",
     "Device",
     "device",
-    "cpu",
-    "cuda",
-    "rocm",
-    "opencl",
-    "metal",
-    "vpi",
-    "vulkan",
-    "ext_dev",
-    "hexagon",
-    "webgpu",
+    "DLDeviceType",
     "from_dlpack",
     "Tensor",
     "Shape",
@@ -73,7 +66,6 @@
     "access_path",
     "serialization",
     "Module",
-    "ModulePropertyMask",
     "system_lib",
     "load_module",
 ]

python/tvm_ffi/convert.py renamed to python/tvm_ffi/_convert.py

@@ -33,6 +33,12 @@ def convert(value: Any) -> Any:
     -------
     ffi_obj : Any
         The converted TVM FFI object.
+
+    Note
+    ----
+    Function arguments to ffi function calls are
+    automatically converted. So this function is mainly
+    only used in internal or testing scenarios.
     """
     if isinstance(value, core.Object):
         return value
@@ -48,7 +54,7 @@ def convert(value: Any) -> Any:
         return core.String(value)
     elif isinstance(value, (bytes, bytearray)):
         return core.Bytes(value)
-    elif isinstance(value, core.ObjectGeneric):
+    elif isinstance(value, core.ObjectConvertible):
         return value.asobject()
     elif callable(value):
         return core._convert_to_ffi_func(value)

python/tvm_ffi/dtype.py renamed to python/tvm_ffi/_dtype.py

@@ -22,7 +22,7 @@
 
 
 class DataTypeCode(IntEnum):
-    """DataType code in DLTensor."""
+    """DLDataTypeCode code in DLTensor."""
 
     INT = 0
     UINT = 1
@@ -57,7 +57,7 @@ class dtype(str):
 
     __slots__ = ["__tvm_ffi_dtype__"]
 
-    NUMPY_DTYPE_TO_STR = {}
+    _NUMPY_DTYPE_TO_STR = {}
 
     def __new__(cls, content):
         content = str(content)
@@ -111,30 +111,30 @@ def lanes(self):
     # although almost in all cases we want numpy
     import numpy as np
 
-    dtype.NUMPY_DTYPE_TO_STR[np.dtype(np.bool_)] = "bool"
-    dtype.NUMPY_DTYPE_TO_STR[np.dtype(np.int8)] = "int8"
-    dtype.NUMPY_DTYPE_TO_STR[np.dtype(np.int16)] = "int16"
-    dtype.NUMPY_DTYPE_TO_STR[np.dtype(np.int32)] = "int32"
-    dtype.NUMPY_DTYPE_TO_STR[np.dtype(np.int64)] = "int64"
-    dtype.NUMPY_DTYPE_TO_STR[np.dtype(np.uint8)] = "uint8"
-    dtype.NUMPY_DTYPE_TO_STR[np.dtype(np.uint16)] = "uint16"
-    dtype.NUMPY_DTYPE_TO_STR[np.dtype(np.uint32)] = "uint32"
-    dtype.NUMPY_DTYPE_TO_STR[np.dtype(np.uint64)] = "uint64"
-    dtype.NUMPY_DTYPE_TO_STR[np.dtype(np.float16)] = "float16"
-    dtype.NUMPY_DTYPE_TO_STR[np.dtype(np.float32)] = "float32"
-    dtype.NUMPY_DTYPE_TO_STR[np.dtype(np.float64)] = "float64"
+    dtype._NUMPY_DTYPE_TO_STR[np.dtype(np.bool_)] = "bool"
+    dtype._NUMPY_DTYPE_TO_STR[np.dtype(np.int8)] = "int8"
+    dtype._NUMPY_DTYPE_TO_STR[np.dtype(np.int16)] = "int16"
+    dtype._NUMPY_DTYPE_TO_STR[np.dtype(np.int32)] = "int32"
+    dtype._NUMPY_DTYPE_TO_STR[np.dtype(np.int64)] = "int64"
+    dtype._NUMPY_DTYPE_TO_STR[np.dtype(np.uint8)] = "uint8"
+    dtype._NUMPY_DTYPE_TO_STR[np.dtype(np.uint16)] = "uint16"
+    dtype._NUMPY_DTYPE_TO_STR[np.dtype(np.uint32)] = "uint32"
+    dtype._NUMPY_DTYPE_TO_STR[np.dtype(np.uint64)] = "uint64"
+    dtype._NUMPY_DTYPE_TO_STR[np.dtype(np.float16)] = "float16"
+    dtype._NUMPY_DTYPE_TO_STR[np.dtype(np.float32)] = "float32"
+    dtype._NUMPY_DTYPE_TO_STR[np.dtype(np.float64)] = "float64"
     if hasattr(np, "float_"):
-        dtype.NUMPY_DTYPE_TO_STR[np.dtype(np.float_)] = "float64"
+        dtype._NUMPY_DTYPE_TO_STR[np.dtype(np.float_)] = "float64"
 except ImportError:
     pass
 
 try:
     import ml_dtypes
 
-    dtype.NUMPY_DTYPE_TO_STR[np.dtype(ml_dtypes.bfloat16)] = "bfloat16"
-    dtype.NUMPY_DTYPE_TO_STR[np.dtype(ml_dtypes.float8_e4m3fn)] = "float8_e4m3fn"
-    dtype.NUMPY_DTYPE_TO_STR[np.dtype(ml_dtypes.float8_e5m2)] = "float8_e5m2"
-    dtype.NUMPY_DTYPE_TO_STR[np.dtype(ml_dtypes.float4_e2m1fn)] = "float4_e2m1fn"
+    dtype._NUMPY_DTYPE_TO_STR[np.dtype(ml_dtypes.bfloat16)] = "bfloat16"
+    dtype._NUMPY_DTYPE_TO_STR[np.dtype(ml_dtypes.float8_e4m3fn)] = "float8_e4m3fn"
+    dtype._NUMPY_DTYPE_TO_STR[np.dtype(ml_dtypes.float8_e5m2)] = "float8_e5m2"
+    dtype._NUMPY_DTYPE_TO_STR[np.dtype(ml_dtypes.float4_e2m1fn)] = "float4_e2m1fn"
 except ImportError:
     pass
 

python/tvm_ffi/_ffi_api.py

@@ -15,7 +15,6 @@
 # specific language governing permissions and limitations
 # under the License.
 """FFI API."""
-from .registry import _init_api
+from . import registry
 
-
-_init_api("ffi", __name__)
+registry.init_ffi_api("ffi", __name__)