[FFI][REFACTOR] Cleanup tvm_ffi python API and types

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

apps/ios_rpc/tests/ios_rpc_test.py

@@ -39,7 +39,7 @@
 
 
 # override metal compiler to compile to iphone
-@tvm.register_func("tvm_callback_metal_compile")
+@tvm.register_global_func("tvm_callback_metal_compile")
 def compile_metal(src, target):
     return xcode.compile_metal(src, sdk=sdk)
 

docs/arch/device_target_interactions.rst

@@ -169,7 +169,7 @@ then be registered with the following steps.
    enum value to a string representation.  This string representation
    should match the name given to ``GlobalDef().def``.
 
-#. Add entries to the ``DEVICE_TYPE_TO_NAME`` and ``DEVICE_NAME_TO_TYPE`` dictionaries of
+#. Add entries to the ``_DEVICE_TYPE_TO_NAME`` and ``_DEVICE_NAME_TO_TYPE`` dictionaries of
    :py:class:`tvm.runtime.Device` for the new enum value.
 
 

docs/get_started/tutorials/quick_start.py

@@ -164,9 +164,9 @@ def forward(self, x):
 #   .. code-block:: Python
 #
 #       # Convert PyTorch tensor to TVM Tensor
-#       x_tvm = tvm.runtime.from_dlpack(x_torch.to_dlpack())
+#       x_tvm = tvm.runtime.from_dlpack(x_torch)
 #       # Convert TVM Tensor to PyTorch tensor
-#       x_torch = torch.from_dlpack(x_tvm.to_dlpack())
+#       x_torch = torch.from_dlpack(x_tvm)
 #
 # - TVM runtime works in non-python environments, so it works on settings such as mobile
 #

ffi/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

ffi/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.

ffi/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.

ffi/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

ffi/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

ffi/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__)

ffi/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",
 ]