[REFACTOR][PYTHON] Phase out tvm._ffi and Limited API support

docs/arch/pass_infra.rst

@@ -550,7 +550,7 @@ a certain scope.
 
 .. code:: python
 
-    @tvm._ffi.register_object("transform.PassContext")
+    @tvm.ffi.register_object("transform.PassContext")
     class PassContext(tvm.runtime.Object):
         def __enter__(self):
             _transform.EnterPassContext(self)

python/setup.py

@@ -20,6 +20,7 @@
 import pathlib
 import shutil
 import sys
+import sys
 
 from setuptools import find_packages
 from setuptools.dist import Distribution
@@ -42,7 +43,7 @@ def get_lib_path():
     """Get library path, name and version"""
     # We can not import `libinfo.py` in setup.py directly since __init__.py
     # Will be invoked which introduces dependencies
-    libinfo_py = os.path.join(CURRENT_DIR, "./tvm/_ffi/libinfo.py")
+    libinfo_py = os.path.join(CURRENT_DIR, "./tvm/libinfo.py")
     libinfo = {"__file__": libinfo_py}
     exec(compile(open(libinfo_py, "rb").read(), libinfo_py, "exec"), libinfo, libinfo)
     version = libinfo["__version__"]
@@ -145,7 +146,15 @@ def config_cython():
     try:
         from Cython.Build import cythonize
 
-        subdir = "_cy3"
+        # for python 3.12+, use limited API for future compact
+        limited_api_kwargs = {}
+        if sys.version_info >= (3, 12):
+            limited_api_kwargs = {
+                "define_macros": [
+                    ("Py_LIMITED_API", 0x030C0000),
+                ],
+                "py_limited_api": True,
+            }
 
         ret = []
         extra_compile_args = ["-std=c++17", "-DDMLC_USE_LOGGING_LIBRARY=<tvm/runtime/logging.h>"]
@@ -179,6 +188,7 @@ def config_cython():
                     library_dirs=library_dirs,
                     libraries=libraries,
                     language="c++",
+                    **limited_api_kwargs,
                 )
             )
         return cythonize(ret, compiler_directives={"language_level": 3})

python/tvm/__init__.py

@@ -19,13 +19,12 @@
 import multiprocessing
 import sys
 import os
-import traceback
 
 # top-level alias
 # tvm._ffi
-from ._ffi.base import TVMError, __version__, _RUNTIME_ONLY
+from .base import TVMError, __version__, _RUNTIME_ONLY
 
-from ._ffi import register_object, register_func, get_global_func
+from .ffi import register_object, register_func, get_global_func
 
 # top-level alias
 # tvm.runtime

python/tvm/arith/_ffi_api.py

@@ -15,7 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """FFI APIs for tvm.arith"""
-import tvm._ffi
+import tvm.ffi
 
 
-tvm._ffi._init_api("arith", __name__)
+tvm.ffi._init_api("arith", __name__)

python/tvm/arith/analyzer.py

@@ -19,7 +19,7 @@
 import enum
 from typing import Union
 
-import tvm._ffi
+import tvm.ffi
 from tvm import tir, ir
 from tvm.runtime import Object
 
@@ -46,15 +46,15 @@ class Extension(enum.Flag):
     ComparisonOfProductAndSum = 1 << 3
 
 
-@tvm._ffi.register_object("arith.ModularSet")
+@tvm.ffi.register_object("arith.ModularSet")
 class ModularSet(Object):
     """Represent range of (coeff * x + base) for x in Z"""
 
     def __init__(self, coeff, base):
         self.__init_handle_by_constructor__(_ffi_api.ModularSet, coeff, base)
 
 
-@tvm._ffi.register_object("arith.ConstIntBound")
+@tvm.ffi.register_object("arith.ConstIntBound")
 class ConstIntBound(Object):
     """Represent constant integer bound
 

python/tvm/arith/int_set.py

@@ -15,7 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """Integer set."""
-import tvm._ffi
+import tvm.ffi
 from tvm.runtime import Object
 from . import _ffi_api
 
@@ -64,7 +64,7 @@ def single_point(point):
         return _ffi_api.intset_single_point(point)
 
 
-@tvm._ffi.register_object("arith.IntervalSet")
+@tvm.ffi.register_object("arith.IntervalSet")
 class IntervalSet(IntSet):
     """Represent set of continuous interval [min_value, max_value]
 
@@ -81,7 +81,7 @@ def __init__(self, min_value, max_value):
         self.__init_handle_by_constructor__(_ffi_api.IntervalSet, min_value, max_value)
 
 
-@tvm._ffi.register_object("arith.PresburgerSet")
+@tvm.ffi.register_object("arith.PresburgerSet")
 class PresburgerSet(IntSet):
     """Represent of Presburger Set"""
 

python/tvm/arith/int_solver.py

@@ -15,12 +15,12 @@
 # specific language governing permissions and limitations
 # under the License.
 """integer constraints data structures and solvers"""
-import tvm._ffi
+import tvm.ffi
 from tvm.runtime import Object
 from . import _ffi_api
 
 
-@tvm._ffi.register_object("arith.IntGroupBounds")
+@tvm.ffi.register_object("arith.IntGroupBounds")
 class IntGroupBounds(Object):
     """Represent integer grouped bounds which are classified into
        lower bounds (include), upper bounds (include) and equalities.
@@ -66,7 +66,7 @@ def find_best_range(self):
         return _ffi_api.IntGroupBounds_FindBestRange(self)
 
 
-@tvm._ffi.register_object("arith.IntConstraints")
+@tvm.ffi.register_object("arith.IntConstraints")
 class IntConstraints(Object):
     """Represent a set of integer constraints including variables, their ranges and
        the relations between them (either equations or inequalities)
@@ -85,7 +85,7 @@ def __init__(self, variables, ranges, relations):
         self.__init_handle_by_constructor__(_ffi_api.IntConstraints, variables, ranges, relations)
 
 
-@tvm._ffi.register_object("arith.IntConstraintsTransform")
+@tvm.ffi.register_object("arith.IntConstraintsTransform")
 class IntConstraintsTransform(Object):
     """We can have different set of variables to represent the same integer constraints.
        For example, the following two constrains are equivalent,

python/tvm/arith/iter_affine_map.py

@@ -14,9 +14,9 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-""" Iterator (quasi)affine mapping patterns."""
+"""Iterator (quasi)affine mapping patterns."""
 from enum import IntEnum
-import tvm._ffi
+import tvm.ffi
 from tvm.runtime import Object
 from tvm.ir import PrimExpr
 from . import _ffi_api
@@ -26,7 +26,7 @@ class IterMapExpr(PrimExpr):
     """Base class of all IterMap expressions."""
 
 
-@tvm._ffi.register_object("arith.IterMark")
+@tvm.ffi.register_object("arith.IterMark")
 class IterMark(Object):
     """Mark the source as an iterator in [0, extent).
 
@@ -43,7 +43,7 @@ def __init__(self, source, extent):
         self.__init_handle_by_constructor__(_ffi_api.IterMark, source, extent)
 
 
-@tvm._ffi.register_object("arith.IterSplitExpr")
+@tvm.ffi.register_object("arith.IterSplitExpr")
 class IterSplitExpr(IterMapExpr):
     """Split of an iterator.
 
@@ -70,7 +70,7 @@ def __init__(self, source, lower_factor, extent, scale):
         )
 
 
-@tvm._ffi.register_object("arith.IterSumExpr")
+@tvm.ffi.register_object("arith.IterSumExpr")
 class IterSumExpr(IterMapExpr):
     """Fuse multiple iterators by summing them with scaling.
 

python/tvm/_ffi/base.py → python/tvm/base.py

@@ -16,22 +16,23 @@
 # under the License.
 # coding: utf-8
 # pylint: disable=invalid-name, import-outside-toplevel
-"""Base library for TVM FFI."""
+"""Base library for TVM."""
 import ctypes
 import os
 import sys
 
-
-import numpy as np
-
 from . import libinfo
 
+# ----------------------------
+# Python3 version.
+# ----------------------------
+if not (sys.version_info[0] >= 3 and sys.version_info[1] >= 8):
+    PY3STATEMENT = "The minimal Python requirement is Python 3.8"
+    raise Exception(PY3STATEMENT)
+
 # ----------------------------
 # library loading
 # ----------------------------
-string_types = (str,)
-integer_types = (int, np.int32)
-numeric_types = integer_types + (float, np.float16, np.float32)
 
 
 def _load_lib():
@@ -62,7 +63,7 @@ def _load_lib():
 
 
 if _RUNTIME_ONLY:
-    from ..ffi import registry as _tvm_ffi_registry
+    from .ffi import registry as _tvm_ffi_registry
 
     _tvm_ffi_registry._SKIP_UNKNOWN_OBJECTS = True
 

python/tvm/contrib/cc.py

@@ -23,7 +23,7 @@
 import sys
 from typing import Dict
 
-from .._ffi.base import py_str
+from ..base import py_str
 from . import tar as _tar
 from . import utils as _utils
 

python/tvm/contrib/clang.py

@@ -18,7 +18,7 @@
 # pylint: disable=invalid-name
 import subprocess
 
-from tvm._ffi.base import py_str
+from tvm.base import py_str
 import tvm.target
 from . import utils
 

python/tvm/contrib/coreml_runtime.py

@@ -15,7 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """CoreML runtime that load and run coreml models."""
-import tvm._ffi
+import tvm.ffi
 from ..rpc import base as rpc_base
 
 
@@ -41,7 +41,7 @@ def create(symbol, compiled_model_path, device):
     if device_type >= rpc_base.RPC_SESS_MASK:
         fcreate = device._rpc_sess.get_function(runtime_func)
     else:
-        fcreate = tvm._ffi.get_global_func(runtime_func)
+        fcreate = tvm.ffi.get_global_func(runtime_func)
         assert fcreate, "Cannot find `tvm.coreml_runtime.create` function."
 
     return CoreMLModule(fcreate(symbol, compiled_model_path))

python/tvm/contrib/cudnn.py

@@ -20,7 +20,7 @@
 import numpy as np
 import tvm
 
-import tvm._ffi
+import tvm.ffi
 from tvm import te
 
 # algos can be read from cudnn.h
@@ -349,7 +349,7 @@ def _conv_find_algo(
         dims - 2, pad, stride, dilation, x_shape, w_shape
     )
     yshape = np.array(y_shape, dtype=np.int32)
-    func = tvm._ffi.get_global_func(func_name)
+    func = tvm.ffi.get_global_func(func_name)
     return func(
         tensor_format,
         dims - 2,

python/tvm/contrib/cutlass/_ffi_api.py

@@ -15,6 +15,6 @@
 # specific language governing permissions and limitations
 # under the License.
 """FFI API for CUTLASS BYOC."""
-import tvm._ffi
+import tvm.ffi
 
-tvm._ffi._init_api("contrib.cutlass", __name__)
+tvm.ffi._init_api("contrib.cutlass", __name__)