[REFACTOR][Python] Use TVM-FFI APIs directly

apps/android_rpc/tests/android_rpc_test.py

@@ -25,6 +25,7 @@
 import os
 
 import numpy as np
+import tvm_ffi
 
 import tvm
 import tvm_tirx
@@ -49,7 +50,7 @@
 
 def test_rpc_module():
     # graph
-    n = tvm_tirx.runtime.convert(1024)
+    n = tvm_ffi.convert(1024)
     A = te.placeholder((n,), name="A")
     B = te.compute(A.shape, lambda *i: A(*i) + 1.0, name="B")
     a_np = np.random.uniform(size=1024).astype(A.dtype)

apps/ios_rpc/tests/ios_rpc_test.py

@@ -23,6 +23,7 @@
 import argparse
 
 import numpy as np
+import tvm_ffi
 
 import tvm
 import tvm_tirx
@@ -38,14 +39,14 @@
 
 
 # override metal compiler to compile to iphone
-@tvm_tirx.register_global_func("tvm_callback_metal_compile")
+@tvm_ffi.register_global_func("tvm_callback_metal_compile")
 def compile_metal(src, target):
     return xcode.compile_metal(src, sdk=sdk)
 
 
 def test_rpc_module(host, port, key, mode):
     # graph
-    n = tvm_tirx.runtime.convert(1024)
+    n = tvm_ffi.convert(1024)
     A = te.placeholder((n,), name="A")
     B = te.compute(A.shape, lambda *i: A(*i) + 1.0, name="B")
     temp = utils.tempdir()

docs/arch/codegen.rst

@@ -238,7 +238,7 @@ when a device kernel is called.
 Runtime Modules
 ---------------
 
-Each codegen produces a ``runtime.Module`` — the container that holds the generated code and
+Each codegen produces a ``tvm_ffi.Module`` — the container that holds the generated code and
 exposes it as callable ``PackedFunc``\ s.
 
 .. list-table::

docs/arch/device_target_interactions.rst

@@ -54,9 +54,9 @@ The ``DeviceAPI`` represents a handle to a specific hardware device
 API.  (e.g. ``CUDADeviceAPI`` handles all interactions through the
 CUDA framework.)  Most ``DeviceAPI`` methods accept a ``device_id``
 parameter to specify which device should be accessed.  In Python,
-these are typically accessed using the :py:func:`tvm.runtime.device`
+these are typically accessed using the :py:func:`tvm_ffi.device`
 function, which returns a handle to a specific device, accessed
-through a specific API.  (e.g. ``tvm.runtime.device('cuda',0)`` gives
+through a specific API.  (e.g. ``tvm_ffi.device('cuda',0)`` gives
 access to physical device ``0``, accessed through the CUDA API.)
 
 .. _device_api.h: https://github.com/apache/tvm/blob/main/include/tvm/runtime/device_api.h

docs/arch/external_library_dispatch.rst

@@ -61,7 +61,7 @@ Each stage is a Relax transformation pass that operates on the ``IRModule``:
 3. **RunCodegen** — finds all functions with a ``Codegen`` attribute, invokes the corresponding
    backend code generator via FFI, and replaces the original calls with ``call_dps_packed``
    to externally compiled functions.
-4. **Linking** — the resulting external ``runtime.Module``\ s are attached to the ``IRModule``
+4. **Linking** — the resulting external ``tvm_ffi.Module``\ s are attached to the ``IRModule``
    as the ``external_mods`` attribute and bundled into the final executable during
    ``relax.build()``.
 
@@ -199,7 +199,7 @@ code generation:
 2. Group them by backend target name.
 3. For each backend, look up the registered codegen function via FFI key
    ``"relax.ext.<backend>"`` (e.g., ``"relax.ext.cublas"``).
-4. Call the codegen function, which returns an array of compiled ``runtime.Module``\ s.
+4. Call the codegen function, which returns an array of compiled ``tvm_ffi.Module``\ s.
 5. Replace the original function calls with ``call_dps_packed(ExternFunc(...), args)``.
 6. Attach the compiled modules to the ``IRModule`` as the ``external_mods`` attribute.
 
@@ -239,7 +239,7 @@ The codegen function receives:
 - ``options``: backend-specific compilation options.
 - ``constant_names``: mapping from constant values to their names (for weight handling).
 
-It returns an array of ``runtime.Module`` objects — one per function — that contain the
+It returns an array of ``tvm_ffi.Module`` objects — one per function — that contain the
 externally compiled code.
 
 Codegen strategies
@@ -267,7 +267,7 @@ After ``RunCodegen``, the original high-level function calls are replaced with:
    R.call_dps_packed(ExternFunc("fused_relax_matmul_cublas0"), (x, w, bias), ...)
 
 At runtime, ``call_dps_packed`` invokes the externally compiled function through the
-``PackedFunc`` interface. The external ``runtime.Module``\ s (produced by the codegen) are
+``PackedFunc`` interface. The external ``tvm_ffi.Module``\ s (produced by the codegen) are
 imported into the final executable during ``relax.build()`` and are available via the module's
 function lookup mechanism.
 

docs/arch/index.rst

@@ -42,8 +42,8 @@ In this guide, we will study an example compilation flow in the compiler. The fi
   equivalent(e.g. in the case of quantization) IRModule. Many of the transformations are target (backend) independent.
   We also allow target to affect the configuration of the transformation pipeline.
 - **Target Translation**: The compiler translates(codegen) the IRModule to an executable format specified by the target.
-  The target translation result is encapsulated as a `runtime.Module` that can be exported, loaded, and executed on the target runtime environment.
-- **Runtime Execution**: the user loads back a `runtime.Module` and runs the compiled functions in the supported runtime environment.
+  The target translation result is encapsulated as a `tvm_ffi.Module` that can be exported, loaded, and executed on the target runtime environment.
+- **Runtime Execution**: the user loads back a `tvm_ffi.Module` and runs the compiled functions in the supported runtime environment.
 
 
 .. figure:: https://raw.githubusercontent.com/tlc-pack/web-data/main/images/design/tvm_overall_flow.svg
@@ -148,44 +148,46 @@ The main goal of TVM's runtime is to provide a minimal API for loading and execu
 .. code-block:: python
 
     import tvm
+    import tvm_ffi
     # Example runtime execution program in python, with type annotated
-    mod: tvm.runtime.Module = tvm.runtime.load_module("compiled_artifact.so")
+    mod: tvm_ffi.Module = tvm.runtime.load_module("compiled_artifact.so")
     arr: tvm.runtime.Tensor = tvm.runtime.tensor([1, 2, 3], device=tvm.cuda(0))
     fun: tvm_ffi.Function = mod["addone"]
     fun(arr)
     print(arr.numpy())
 
 
-:py:class:`tvm.runtime.Module` encapsulates the result of compilation. A runtime.Module contains a GetFunction method to obtain :py:class:`tvm_ffi.Function` instances by name.
+:py:class:`tvm_ffi.Module` encapsulates the result of compilation. A tvm_ffi.Module contains a GetFunction method to obtain :py:class:`tvm_ffi.Function` instances by name.
 
 :py:class:`tvm_ffi.Function` is a type-erased function interface for both the generated functions. A tvm_ffi.Function can take arguments and return values with the
-following types: POD types(int, float), string, tvm_ffi.Function, runtime.Module, runtime.Tensor, and other sub-classes of runtime.Object.
+following types: POD types(int, float), string, tvm_ffi.Function, tvm_ffi.Module, runtime.Tensor, and other sub-classes of tvm_ffi.Object.
 
-:py:class:`tvm.runtime.Module` and :py:class:`tvm_ffi.Function` are powerful mechanisms to modularize the runtime. For example, to get the above `addone` function on CUDA, we can use LLVM to generate the host-side code to compute the launching parameters(e.g. size of the thread groups) and then call into another tvm_ffi.Function from a CUDAModule that is backed by the CUDA driver API. The same mechanism can be used for OpenCL kernels.
+:py:class:`tvm_ffi.Module` and :py:class:`tvm_ffi.Function` are powerful mechanisms to modularize the runtime. For example, to get the above `addone` function on CUDA, we can use LLVM to generate the host-side code to compute the launching parameters(e.g. size of the thread groups) and then call into another tvm_ffi.Function from a CUDAModule that is backed by the CUDA driver API. The same mechanism can be used for OpenCL kernels.
 
-The above example only deals with a simple `addone` function. The code snippet below gives an example of an end-to-end model execution using the Relax Virtual Machine, which is built on the same runtime.Module and tvm_ffi.Function interface:
+The above example only deals with a simple `addone` function. The code snippet below gives an example of an end-to-end model execution using the Relax Virtual Machine, which is built on the same tvm_ffi.Module and tvm_ffi.Function interface:
 
 .. code-block:: python
 
    import tvm
+   import tvm_ffi
    from tvm import relax
    # Load the compiled artifact
-   mod: tvm.runtime.Module = tvm.runtime.load_module("resnet18.so")
+   mod: tvm_ffi.Module = tvm.runtime.load_module("resnet18.so")
    # Create a VM instance on cuda(0)
    vm = relax.VirtualMachine(mod, tvm.cuda(0))
    data: tvm.runtime.Tensor = get_input_data()
    # Run the model — vm["main"] returns a PackedFunc
    result = vm["main"](data).numpy()
 
-The main take away is that runtime.Module and runtime.PackedFunc are sufficient to encapsulate both operator level programs (such as addone), as well as the end-to-end models.
+The main take away is that tvm_ffi.Module and runtime.PackedFunc are sufficient to encapsulate both operator level programs (such as addone), as well as the end-to-end models.
 
 Summary and Discussions
 ~~~~~~~~~~~~~~~~~~~~~~~
 
 In summary, the key data structures in the compilation flows are:
 
 - IRModule: contains relax.Function and tirx.Function
-- runtime.Module: contains runtime.PackedFunc
+- tvm_ffi.Module: contains runtime.PackedFunc
 
 Most parts of the compilation are transformations among the key data structures.
 
@@ -355,7 +357,7 @@ mechanism, IR builder frame stack, printer pipeline, and syntax reference.
 
 tvm/target
 ----------
-The target module contains all the code generators that translate an IRModule to a target runtime.Module.
+The target module contains all the code generators that translate an IRModule to a target tvm_ffi.Module.
 It also provides a common `Target` class that describes the target.
 
 Targets can be constructed from a registered tag, a configuration dictionary, or a tag with attribute overrides:

docs/arch/pass_infra.rst

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

docs/arch/relax_vm.rst

@@ -48,7 +48,7 @@ The end-to-end flow from model to execution is:
    ExecBuilder (bytecode) + IRModule (TIR only)
         │                        │
         │                        ▼  tirx.compile()
-        │                   runtime.Module (native kernels)
+        │                   tvm_ffi.Module (native kernels)
         │                        │
         ▼  VMLink               ▼
    VMExecutable ◄───────── linked together

docs/arch/runtime.rst

@@ -88,8 +88,9 @@ The following example registers PackedFunc in C++ and calls from python.
 .. code:: python
 
     import tvm
+    import tvm_ffi
 
-    myadd = tvm.get_global_func("myadd")
+    myadd = tvm_ffi.get_global_func("myadd")
     # prints 3
     print(myadd(1, 2))
 
@@ -123,13 +124,14 @@ we can pass functions from python (as PackedFunc) to C++.
 .. code:: python
 
     import tvm
+    import tvm_ffi
 
     def callback(msg):
       print(msg)
 
     # convert to PackedFunc
-    f = tvm.runtime.convert(callback)
-    callhello = tvm.get_global_func("callhello")
+    f = tvm_ffi.convert(callback)
+    callhello = tvm_ffi.get_global_func("callhello")
     # prints hello world
     callhello(f)
 

docs/deep_dive/tensor_ir/tutorials/tir_creation.py

@@ -206,7 +206,7 @@ def mm_relu(A: T.Tensor([M, K], dtype), B: T.Tensor([K, N], dtype), C: T.Tensor(
 # Now let's check the runtime dynamic shape inference:
 
 
-def evaluate_dynamic_shape(lib: tvm_tirx.runtime.Module, m: int, n: int, k: int):
+def evaluate_dynamic_shape(lib: tvm_ffi.Module, m: int, n: int, k: int):
     A = tvm_tirx.runtime.tensor(np.random.uniform(size=(m, k)).astype("float32"))
     B = tvm_tirx.runtime.tensor(np.random.uniform(size=(k, n)).astype("float32"))
     C = tvm_tirx.runtime.tensor(np.zeros((m, n), dtype="float32"))

docs/get_started/tutorials/quick_start.py

@@ -166,7 +166,7 @@ def forward(self, x):
 #   .. code-block:: Python
 #
 #       # Convert PyTorch tensor to TVM Tensor
-#       x_tvm = tvm.runtime.from_dlpack(x_torch)
+#       x_tvm = tvm_ffi.from_dlpack(x_torch)
 #       # Convert TVM Tensor to PyTorch tensor
 #       x_torch = torch.from_dlpack(x_tvm)
 #

docs/how_to/tutorials/bring_your_own_codegen.py

@@ -65,7 +65,6 @@
 # Steps 1 and 2 are pure Python and run anywhere; steps 3 and 4 need the
 # backend's codegen and runtime compiled into TVM, which is why the
 # build-and-run cells below are guarded.
-
 ######################################################################
 # Step 1: Import the backends to register their patterns
 # ------------------------------------------------------
@@ -74,11 +73,11 @@
 # Pattern registration is independent of the C++ build -- only codegen and the
 # runtime require the backend to be compiled in -- so we probe each backend and
 # guard the build-and-run cells accordingly.
-
 import os
 import tempfile
 
 import numpy as np
+import tvm_ffi
 
 import tvm.relax.backend.contrib.example_npu  # noqa: F401
 import tvm_tirx
@@ -88,12 +87,12 @@
 from tvm.relax.backend.pattern_registry import get_patterns_with_prefix
 from tvm.relax.transform import FuseOpsByPattern, MergeCompositeFunctions, RunCodegen
 
-has_example_npu_codegen = tvm_tirx.get_global_func("relax.ext.example_npu", True)
-has_example_npu_runtime = tvm_tirx.get_global_func("runtime.ExampleNPUJSONRuntimeCreate", True)
+has_example_npu_codegen = tvm_ffi.get_global_func("relax.ext.example_npu", True)
+has_example_npu_runtime = tvm_ffi.get_global_func("runtime.ExampleNPUJSONRuntimeCreate", True)
 has_example_npu = has_example_npu_codegen and has_example_npu_runtime
 
-has_tensorrt_codegen = tvm_tirx.get_global_func("relax.ext.tensorrt", True) is not None
-_is_trt_runtime_enabled = tvm_tirx.get_global_func("relax.is_tensorrt_runtime_enabled", True)
+has_tensorrt_codegen = tvm_ffi.get_global_func("relax.ext.tensorrt", True) is not None
+_is_trt_runtime_enabled = tvm_ffi.get_global_func("relax.is_tensorrt_runtime_enabled", True)
 has_tensorrt = (
     has_tensorrt_codegen and _is_trt_runtime_enabled is not None and _is_trt_runtime_enabled()
 )

docs/how_to/tutorials/cross_compilation_and_rpc.py

@@ -102,7 +102,7 @@
 from tvm import rpc, te
 from tvm_tirx.support import utils
 
-n = tvm_tirx.runtime.convert(1024)
+n = tvm_ffi.convert(1024)
 A = te.placeholder((n,), name="A")
 B = te.compute((n,), lambda i: A[i] + 1.0, name="B")
 mod = tvm_tirx.IRModule.from_expr(te.create_function([A, B]).with_attr("global_symbol", "add_one"))

docs/how_to/tutorials/customize_opt.py

@@ -48,11 +48,11 @@
 # - **Build and Universal Deployment**: Build the optimized model to a deployable module to the
 #   universal runtime, and execute it on different devices, such as CPU, GPU, or other accelerators.
 #
-
 import os
 import tempfile
 
 import numpy as np
+import tvm_ffi
 
 import tvm
 import tvm_tirx
@@ -125,7 +125,7 @@ def forward(self, x):
 class CublasDispatch:
     def transform_module(self, mod: IRModule, _ctx: tvm_tirx.ir.transform.PassContext) -> IRModule:
         # Check if CUBLAS is enabled
-        if not tvm_tirx.get_global_func("relax.ext.cublas", True):
+        if not tvm_ffi.get_global_func("relax.ext.cublas", True):
             raise Exception("CUBLAS is not enabled.")
 
         # Get interested patterns

docs/how_to/tutorials/mix_python_and_tvm_with_pymodule.py

@@ -47,9 +47,10 @@
 ######################################################################
 # Preparation
 # -----------
-
 import os
 
+import tvm_ffi
+
 try:
     import torch
     import torch.nn.functional as F
@@ -192,12 +193,13 @@ def forward(self, x, weights):
 # - ``call_dps_packed``: call a TVM packed function (e.g., a third-party library binding)
 # - Direct Python: call any PyTorch function
 #
-# ``call_dps_packed`` is useful for calling functions registered via ``tvm.register_global_func``
+# ``call_dps_packed`` can call functions registered via
+# ``tvm_ffi.register_global_func``
 # — for example, CUBLAS or cuDNN bindings that TVM wraps as packed functions.
 
 if RUN_EXAMPLE:
     # Register a packed function (simulating an external library binding)
-    @tvm_tirx.register_global_func("my_bias_add", override=True)
+    @tvm_ffi.register_global_func("my_bias_add", override=True)
     def my_bias_add(x, bias, out):
         """Packed function: adds bias to each row of x."""
 

docs/how_to/tutorials/optimize_llm.py

@@ -45,22 +45,20 @@
 # - **Build and Universal Deployment**: Build the optimized model to a deployable module to the
 #   universal runtime, and execute it on different devices, such as CPU, GPU, or other accelerators.
 #
-
-
 ######################################################################
 # Construct the model architecture
 # --------------------------------
 # We will use a pre-trained TinyLlama model from Hugging Face. However, usually we only load the
 # pre-trained weight from Hugging Face but not the model architecture. We need to construct the
 # model architecture by ourselves. Apache TVM prepares a PyTorch-liked API to construct the model
 # architecture. We can use the API to construct the model architecture.
-
 import dataclasses
 import enum
 import os
 from pathlib import Path
 from pprint import pprint
 
+import tvm_ffi
 from tvm_ffi import Shape
 
 import tvm
@@ -553,7 +551,7 @@ def _pipeline(
 # The next step is to embed the tokens into the hidden states. We use the `embed` function
 # compiled in the Relax IRModule to embed the tokens into the hidden states.
 
-nd_view_func = tvm_tirx.get_global_func("vm.builtin.reshape")
+nd_view_func = tvm_ffi.get_global_func("vm.builtin.reshape")
 
 
 def embed(tokens, params):
@@ -568,9 +566,9 @@ def embed(tokens, params):
 # ~~~~~~~
 # Before running the forward pass, we first get some help functions for preparation.
 
-add_sequence_func = tvm_tirx.get_global_func("vm.builtin.kv_state_add_sequence")
-begin_forward_func = tvm_tirx.get_global_func("vm.builtin.kv_state_begin_forward")
-end_forward_func = tvm_tirx.get_global_func("vm.builtin.kv_state_end_forward")
+add_sequence_func = tvm_ffi.get_global_func("vm.builtin.kv_state_add_sequence")
+begin_forward_func = tvm_ffi.get_global_func("vm.builtin.kv_state_begin_forward")
+end_forward_func = tvm_ffi.get_global_func("vm.builtin.kv_state_end_forward")
 
 ######################################################################
 # As we are creating a new sequence, we need to call `add_sequence_func` to initialize

docs/tirx/install.rst

@@ -251,8 +251,9 @@ FFI module and ``executable.jit()`` for a runnable module.
 ``tvm_tirx.codegen.load_static_library`` wraps an existing object for inclusion
 in an executable's import graph.
 
-``tvm_tirx.runtime.Module``, ``load_module``, and ``system_lib`` directly alias
-the FFI APIs. Loading no longer invokes a linker for objects or archives; export
+Use ``tvm_ffi.Module`` and ``tvm_ffi.system_lib`` directly.
+``tvm_tirx.runtime.load_module`` remains a convenience alias for
+``tvm_ffi.load_module``. Loading no longer invokes a linker for objects or archives; export
 and link explicitly before loading. Parent TVM's RPC server retains remote
 object/archive linking through its registered server hook.