diff --git a/apps/android_rpc/tests/android_rpc_test.py b/apps/android_rpc/tests/android_rpc_test.py index a17490782dd2..6e29151df23b 100644 --- a/apps/android_rpc/tests/android_rpc_test.py +++ b/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) diff --git a/apps/ios_rpc/tests/ios_rpc_test.py b/apps/ios_rpc/tests/ios_rpc_test.py index d17c484c7cd9..062c117f1683 100644 --- a/apps/ios_rpc/tests/ios_rpc_test.py +++ b/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() diff --git a/docs/arch/codegen.rst b/docs/arch/codegen.rst index 335f4b8447cf..9282483617d5 100644 --- a/docs/arch/codegen.rst +++ b/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:: diff --git a/docs/arch/device_target_interactions.rst b/docs/arch/device_target_interactions.rst index 0bba9cde2059..39b38e9a8dda 100644 --- a/docs/arch/device_target_interactions.rst +++ b/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 diff --git a/docs/arch/external_library_dispatch.rst b/docs/arch/external_library_dispatch.rst index ed5139aa546b..a0db1f715871 100644 --- a/docs/arch/external_library_dispatch.rst +++ b/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."`` (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. diff --git a/docs/arch/index.rst b/docs/arch/index.rst index 77d953846508..0c189d824921 100644 --- a/docs/arch/index.rst +++ b/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,36 +148,38 @@ 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 ~~~~~~~~~~~~~~~~~~~~~~~ @@ -185,7 +187,7 @@ 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: diff --git a/docs/arch/pass_infra.rst b/docs/arch/pass_infra.rst index a0b4b28388eb..79871807ce90 100644 --- a/docs/arch/pass_infra.rst +++ b/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 diff --git a/docs/arch/relax_vm.rst b/docs/arch/relax_vm.rst index 79353680d5fc..f715a9641c2b 100644 --- a/docs/arch/relax_vm.rst +++ b/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 diff --git a/docs/arch/runtime.rst b/docs/arch/runtime.rst index 03e559266155..34f698291277 100644 --- a/docs/arch/runtime.rst +++ b/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) diff --git a/docs/deep_dive/tensor_ir/tutorials/tir_creation.py b/docs/deep_dive/tensor_ir/tutorials/tir_creation.py index 15ab2f0a4352..0f0f12dbb297 100644 --- a/docs/deep_dive/tensor_ir/tutorials/tir_creation.py +++ b/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")) diff --git a/docs/get_started/tutorials/quick_start.py b/docs/get_started/tutorials/quick_start.py index 5f6ecd291951..a99cb92dc398 100644 --- a/docs/get_started/tutorials/quick_start.py +++ b/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) # diff --git a/docs/how_to/tutorials/bring_your_own_codegen.py b/docs/how_to/tutorials/bring_your_own_codegen.py index ba0a34fd06fe..e5978aaaa51f 100644 --- a/docs/how_to/tutorials/bring_your_own_codegen.py +++ b/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() ) diff --git a/docs/how_to/tutorials/cross_compilation_and_rpc.py b/docs/how_to/tutorials/cross_compilation_and_rpc.py index 58caaf450d69..c975224b9c39 100644 --- a/docs/how_to/tutorials/cross_compilation_and_rpc.py +++ b/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")) diff --git a/docs/how_to/tutorials/customize_opt.py b/docs/how_to/tutorials/customize_opt.py index fc52128e6a71..4741355993d1 100644 --- a/docs/how_to/tutorials/customize_opt.py +++ b/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 diff --git a/docs/how_to/tutorials/mix_python_and_tvm_with_pymodule.py b/docs/how_to/tutorials/mix_python_and_tvm_with_pymodule.py index e195f431ce30..7532f63e271e 100644 --- a/docs/how_to/tutorials/mix_python_and_tvm_with_pymodule.py +++ b/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.""" diff --git a/docs/how_to/tutorials/optimize_llm.py b/docs/how_to/tutorials/optimize_llm.py index 2e03930fd244..c8f4cc186892 100644 --- a/docs/how_to/tutorials/optimize_llm.py +++ b/docs/how_to/tutorials/optimize_llm.py @@ -45,8 +45,6 @@ # - **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 # -------------------------------- @@ -54,13 +52,13 @@ # 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 diff --git a/docs/tirx/install.rst b/docs/tirx/install.rst index ad05a9b859f3..c2572cd926bd 100644 --- a/docs/tirx/install.rst +++ b/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. diff --git a/python/tvm/contrib/cudnn.py b/python/tvm/contrib/cudnn.py index 28be5ba194b1..a1ce63a34f31 100644 --- a/python/tvm/contrib/cudnn.py +++ b/python/tvm/contrib/cudnn.py @@ -51,7 +51,7 @@ def exists(): True if CuDNN support is enabled and a CuDNN-capable GPU exists. Otherwise, False. """ - func = tvm_tirx.get_global_func("tvm.contrib.cudnn.exists", allow_missing=True) + func = tvm_ffi.get_global_func("tvm.contrib.cudnn.exists", allow_missing=True) if func is None: return False @@ -199,7 +199,7 @@ def conv_output_shape( c_output = w_shape[0] x_chan = x_shape[1] w_chan_input = w_shape[1] - w_lanes = tvm_tirx.runtime.DataType(conv_dtype).lanes + w_lanes = tvm_ffi.dtype(conv_dtype).lanes assert w_lanes == 1 x_shape = x_shape[2:] w_shape = w_shape[2:] @@ -207,7 +207,7 @@ def conv_output_shape( else: raise ValueError(f"Unknown CuDNN tensor format: '{tensor_format}'") - x_lanes = tvm_tirx.runtime.DataType(data_dtype).lanes + x_lanes = tvm_ffi.dtype(data_dtype).lanes assert x_chan * x_lanes == w_chan_input * groups, ( f"Mismatched dimensions, data has {x_chan // groups} channels/group " f"(dimension {x_chan} with {x_lanes} lanes/value, {groups} groups), " diff --git a/python/tvm/contrib/cutlass/build.py b/python/tvm/contrib/cutlass/build.py index 253732a5eb99..3fd469e458b1 100644 --- a/python/tvm/contrib/cutlass/build.py +++ b/python/tvm/contrib/cutlass/build.py @@ -26,6 +26,7 @@ from collections.abc import Sequence from functools import reduce +import tvm_ffi from tvm_ffi import register_global_func import tvm_tirx @@ -43,7 +44,7 @@ def has_cutlass(): """Returns true if the CUTLASS custom codegen is available""" - return tvm_tirx.get_global_func("relax.ext.cutlass", True) is not None + return tvm_ffi.get_global_func("relax.ext.cutlass", True) is not None def _get_cutlass_path(): @@ -400,7 +401,7 @@ def _extract_relax_function_signature(f): def _extract_arg_idx(pattern_name, f): - extract_func = tvm_tirx.get_global_func("relax.contrib.extract_arg_idx") + extract_func = tvm_ffi.get_global_func("relax.contrib.extract_arg_idx") arg_indices = extract_func(pattern_name, f) return {k: int(v) for k, v in arg_indices.items()} @@ -850,7 +851,7 @@ def compile_cutlass_module(c_source_module, options): Parameters ---------- - c_source_module: runtime.Module + c_source_module: tvm_ffi.Module A C-source module containing CUTLASS kernels. options: dict @@ -863,7 +864,7 @@ def compile_cutlass_module(c_source_module, options): Returns ------- - rt_mod : runtime.Module + rt_mod : tvm_ffi.Module A runtime module where all cutlass kernels have been compiled. """ tmp_dir = options.get("tmp_dir", "./tmp") @@ -890,7 +891,7 @@ def finalize_modules(lib, lib_path="compile.so", tmp_dir="./tmp"): Parameters ---------- - lib : runtime.Module + lib : tvm_ffi.Module The output from build. lib_path : string @@ -901,7 +902,7 @@ def finalize_modules(lib, lib_path="compile.so", tmp_dir="./tmp"): Returns ------- - updated_lib : runtime.Module + updated_lib : tvm_ffi.Module The updated library with all compilation and linking completed. """ diff --git a/python/tvm/contrib/cutlass/gen_tensor_op.py b/python/tvm/contrib/cutlass/gen_tensor_op.py index d609c2f7640c..d167ce31a16a 100644 --- a/python/tvm/contrib/cutlass/gen_tensor_op.py +++ b/python/tvm/contrib/cutlass/gen_tensor_op.py @@ -29,7 +29,6 @@ import tvm_tirx from tvm_tirx.ir import IntImm -from tvm_tirx.runtime import Object from . import _ffi_api as ffi from .attention_operation import ( @@ -450,7 +449,7 @@ def evaluate(self, op, args): return rt -class CodegenResult(Object): +class CodegenResult(tvm_ffi.Object): """The holder for the generated code and required headers.""" def __init__(self, code, headers): diff --git a/python/tvm/contrib/dlpack.py b/python/tvm/contrib/dlpack.py index c38621966e67..5ce7dd021977 100644 --- a/python/tvm/contrib/dlpack.py +++ b/python/tvm/contrib/dlpack.py @@ -16,7 +16,7 @@ # under the License. """Wrapping functions to bridge frameworks with DLPack support to TVM""" -import tvm_tirx.runtime +import tvm_ffi def convert_func(tvm_func, tensor_type, to_dlpack_func): @@ -38,9 +38,7 @@ def convert_func(tvm_func, tensor_type, to_dlpack_func): def _wrapper(*args): args = tuple( - tvm_tirx.runtime.from_dlpack(to_dlpack_func(arg)) - if isinstance(arg, tensor_type) - else arg + tvm_ffi.from_dlpack(to_dlpack_func(arg)) if isinstance(arg, tensor_type) else arg for arg in args ) return tvm_func(*args) diff --git a/python/tvm/contrib/hexagon/tools.py b/python/tvm/contrib/hexagon/tools.py index 97c92b3ceb02..195eaa760ce2 100644 --- a/python/tvm/contrib/hexagon/tools.py +++ b/python/tvm/contrib/hexagon/tools.py @@ -27,6 +27,7 @@ import tarfile import numpy +import tvm_ffi from tvm_ffi import register_global_func import tvm_tirx @@ -227,7 +228,7 @@ def create_aot_shared(so_name: str | pathlib.Path, files, hexagon_arch: str, opt def pack_imports( - module: tvm_tirx.runtime.Module, + module: tvm_ffi.Module, is_system_lib: bool, # pylint: disable=unused-argument c_symbol_prefix: str, workspace_dir: str, @@ -238,7 +239,7 @@ def pack_imports( Parameters ---------- - module: tvm_tirx.runtime.Module + module: tvm_ffi.Module Module whose imported modules need to be serialized. is_system_lib: bool Flag whether the exported module will be used as a system library. @@ -255,12 +256,12 @@ def pack_imports( path_bin = os.path.join(workspace_dir, "imports.bin") pack_to_bin_f_name = "runtime.ModulePackImportsToTensor" - fpack_to_bin = tvm_tirx.get_global_func(pack_to_bin_f_name) + fpack_to_bin = tvm_ffi.get_global_func(pack_to_bin_f_name) assert fpack_to_bin, f"Expecting {pack_to_bin_f_name} in registry" fpack_to_bin(module).numpy().tofile(path_bin) - mblob_symbol = c_symbol_prefix + tvm_tirx.get_global_func("runtime.ModuleImportsBlobName")() + mblob_symbol = c_symbol_prefix + tvm_ffi.get_global_func("runtime.ModuleImportsBlobName")() binary_size = os.path.getsize(path_bin) hexagon_toolchain = os.environ.get("HEXAGON_TOOLCHAIN") diff --git a/python/tvm/contrib/tvmjs.py b/python/tvm/contrib/tvmjs.py index 8a20a6bfadc7..655a50280724 100644 --- a/python/tvm/contrib/tvmjs.py +++ b/python/tvm/contrib/tvmjs.py @@ -30,6 +30,7 @@ from typing import Any, Optional, Union import numpy as np +import tvm_ffi try: import ml_dtypes @@ -38,7 +39,6 @@ import tvm import tvm_tirx -from tvm_tirx.runtime import DataType from .emcc import create_wasm as create_tvmjs_wasm from .emcc import find_wasm_lib @@ -298,8 +298,8 @@ def resolve_encode_format(name): # prefer to preserve original dtype, especially if the format was bfloat16 dtype = origin_v.dtype if isinstance(origin_v, tvm_tirx.runtime.Tensor) else v.dtype - if dtype in DataType._NUMPY_DTYPE_TO_STR: - dtype = DataType._NUMPY_DTYPE_TO_STR[dtype] + if dtype in tvm_ffi.dtype._NUMPY_DTYPE_TO_STR: + dtype = tvm_ffi.dtype._NUMPY_DTYPE_TO_STR[dtype] else: dtype = str(dtype) diff --git a/python/tvm/exec/disco_remote_socket_session.py b/python/tvm/exec/disco_remote_socket_session.py index 23f326cb21ee..cd713f25ab20 100644 --- a/python/tvm/exec/disco_remote_socket_session.py +++ b/python/tvm/exec/disco_remote_socket_session.py @@ -20,6 +20,8 @@ import sys +import tvm_ffi + import tvm import tvm_tirx @@ -33,5 +35,5 @@ server_host = sys.argv[1] server_port = int(sys.argv[2]) num_workers = int(sys.argv[3]) - func = tvm_tirx.get_global_func("runtime.disco.RemoteSocketSession") + func = tvm_ffi.get_global_func("runtime.disco.RemoteSocketSession") func(server_host, server_port, num_workers) diff --git a/python/tvm/relax/analysis/estimate_memory_usage.py b/python/tvm/relax/analysis/estimate_memory_usage.py index 8df5f4ed69dc..855e280f3591 100644 --- a/python/tvm/relax/analysis/estimate_memory_usage.py +++ b/python/tvm/relax/analysis/estimate_memory_usage.py @@ -17,6 +17,8 @@ # pylint: disable=abstract-method,unused-argument # pylint: disable=missing-function-docstring,missing-module-docstring +import tvm_ffi + import tvm_tirx from tvm_tirx.ir import Call, Op from tvm_tirx.ir.module import IRModule @@ -128,7 +130,7 @@ def calculate_size(self, shape: Expr, dtype_str: str) -> int: self.total_dyn_size_tensor_num += 1 return -1 size *= dim_len.value - dtype = tvm_tirx.DataType(dtype_str) + dtype = tvm_ffi.dtype(dtype_str) return size * ((dtype.bits + 7) // 8) * dtype.lanes def accumulate_builtin_tensor_alloc(self, shape: Expr, dtype_str: str) -> None: diff --git a/python/tvm/relax/backend/adreno/clml.py b/python/tvm/relax/backend/adreno/clml.py index 159e9699fa8e..e01c8cc79a9e 100644 --- a/python/tvm/relax/backend/adreno/clml.py +++ b/python/tvm/relax/backend/adreno/clml.py @@ -17,6 +17,8 @@ # pylint: disable=invalid-name, unused-argument, pointless-exception-statement """Pattern table for CLML backend""" +import tvm_ffi + import tvm_tirx from tvm import relax from tvm.relax import transform @@ -99,7 +101,7 @@ def clml_sdk_version(): # Registry: "relax.get_openclml_version" — returns the CLML SDK version # that TVM was built against; registered unconditionally in codegen.cc. # Grep hint: grep -rn 'relax.get_openclml_version' src/ - get_version = tvm_tirx.get_global_func("relax.get_openclml_version", allow_missing=True) + get_version = tvm_ffi.get_global_func("relax.get_openclml_version", allow_missing=True) if get_version is None: return 2 return int(get_version()) @@ -113,7 +115,7 @@ def is_clml_runtime_enabled(): ret: bool True if present, False if not. """ - check_enabled = tvm_tirx.get_global_func("relax.op.is_openclml_runtime_enabled", True) + check_enabled = tvm_ffi.get_global_func("relax.op.is_openclml_runtime_enabled", True) if check_enabled: return check_enabled() return False diff --git a/python/tvm/relax/backend/contrib/example_npu/README.md b/python/tvm/relax/backend/contrib/example_npu/README.md index 769cc60c3b18..7588ba3f526f 100644 --- a/python/tvm/relax/backend/contrib/example_npu/README.md +++ b/python/tvm/relax/backend/contrib/example_npu/README.md @@ -173,8 +173,9 @@ in the TVM build: ```python import tvm -has_codegen = tvm.get_global_func("relax.ext.example_npu", True) -has_runtime = tvm.get_global_func("runtime.ExampleNPUJSONRuntimeCreate", True) +import tvm_ffi +has_codegen = tvm_ffi.get_global_func("relax.ext.example_npu", True) +has_runtime = tvm_ffi.get_global_func("runtime.ExampleNPUJSONRuntimeCreate", True) has_example_npu = has_codegen and has_runtime ``` diff --git a/python/tvm/relax/backend/cuda/cublas.py b/python/tvm/relax/backend/cuda/cublas.py index 811ae99888cd..b5c780c3b429 100644 --- a/python/tvm/relax/backend/cuda/cublas.py +++ b/python/tvm/relax/backend/cuda/cublas.py @@ -20,11 +20,12 @@ import operator from functools import reduce +import tvm_ffi + import tvm import tvm_tirx from tvm.relax import transform from tvm.relax.transform import PatternCheckContext -from tvm_tirx import DataType from tvm_tirx.sym import Analyzer from ..pattern_registry import get_patterns_with_prefix, register_patterns @@ -104,12 +105,12 @@ def _check_matmul(context: PatternCheckContext) -> bool: # cuBLAS FP8 operations require all tensors being aligned to 16 bytes. if ( not isinstance(rhs_shape[-1], tvm_tirx.ir.IntImm | int) - or rhs_shape[-1] % (16 // DataType(lhs_dtype).itemsize) != 0 + or rhs_shape[-1] % (16 // tvm_ffi.dtype(lhs_dtype).itemsize) != 0 ): return False if ( not isinstance(rhs_shape[-2], tvm_tirx.ir.IntImm | int) - or rhs_shape[-2] % (16 // DataType(out_dtype).itemsize) != 0 + or rhs_shape[-2] % (16 // tvm_ffi.dtype(out_dtype).itemsize) != 0 ): return False diff --git a/python/tvm/relax/backend/cuda/flashinfer.py b/python/tvm/relax/backend/cuda/flashinfer.py index 4f4a81591cb8..e895e5354000 100644 --- a/python/tvm/relax/backend/cuda/flashinfer.py +++ b/python/tvm/relax/backend/cuda/flashinfer.py @@ -20,6 +20,8 @@ import re from pathlib import Path +import tvm_ffi + import tvm_tirx from tvm_tirx.target import Target @@ -48,7 +50,7 @@ def _rename_exported_func_names(source_paths: list[Path], prefix: str): source_path.write_text("".join(lines), encoding="utf-8") -def _load_flashinfer_modules(object_files: list[Path]) -> list[tvm_tirx.runtime.Module]: +def _load_flashinfer_modules(object_files: list[Path]) -> list[tvm_ffi.Module]: return [ tvm_tirx.codegen.load_static_library(str(obj_path.absolute()), func_names=[]) for obj_path in object_files @@ -63,7 +65,7 @@ def gen_flashinfer_prefill_module( v_head_dim: int, enable_inline_rope: bool, return_static_libs: bool = False, -) -> list[tvm_tirx.runtime.Module]: +) -> list[tvm_ffi.Module]: """Generate a FlashInfer module for prefill. Parameters @@ -158,7 +160,7 @@ def gen_flashinfer_decode_module( v_head_dim: int, enable_inline_rope: bool, return_static_libs: bool = False, -) -> list[tvm_tirx.runtime.Module]: +) -> list[tvm_ffi.Module]: """Generate a FlashInfer module for decode. Parameters @@ -236,7 +238,7 @@ def gen_flashinfer_mla_module( head_dim_ckv: int, head_dim_kpe: int, return_static_libs: bool = False, -) -> list[tvm_tirx.runtime.Module]: +) -> list[tvm_ffi.Module]: """Generate a FlashInfer module for MLA. Parameters @@ -300,7 +302,7 @@ def gen_flashinfer_mla_module( def gen_grouped_gemm_module( target: Target, return_static_libs: bool = False -) -> list[tvm_tirx.runtime.Module]: +) -> list[tvm_ffi.Module]: """Generate a FlashInfer module for FP8 grouped GEMM. Parameters @@ -314,7 +316,7 @@ def gen_grouped_gemm_module( Returns ------- - List[tvm_tirx.runtime.Module] + List[tvm_ffi.Module] A list of compiled static library modules for FlashInfer FP8 grouped GEMM kernels. Note diff --git a/python/tvm/relax/backend/dispatch_sort_scan.py b/python/tvm/relax/backend/dispatch_sort_scan.py index bd57cafe6992..ed3e2b1479f2 100644 --- a/python/tvm/relax/backend/dispatch_sort_scan.py +++ b/python/tvm/relax/backend/dispatch_sort_scan.py @@ -20,10 +20,11 @@ from functools import reduce from operator import mul +import tvm_ffi + from tvm import relax, topi from tvm.contrib.thrust import can_use_thrust from tvm.relax import expr_functor -from tvm_tirx import DataType from tvm_tirx.ir import GlobalVar, Op from tvm_tirx.ir.module import IRModule from tvm_tirx.ir.transform import PassContext, module_pass @@ -240,9 +241,9 @@ def estimate_thrust_workspace_size(self, call: relax.Call) -> int: Estimate the workspace size for thrust sort/argsort/topk/cumsum """ input_shape = call.args[0].ty.shape - input_byte_per_elem = DataType(call.args[0].ty.dtype.dtype).bits // 8 - int64_byte_per_elem = DataType("int64").bits // 8 - int32_byte_per_elem = DataType("int32").bits // 8 + input_byte_per_elem = tvm_ffi.dtype(call.args[0].ty.dtype.dtype).bits // 8 + int64_byte_per_elem = tvm_ffi.dtype("int64").bits // 8 + int32_byte_per_elem = tvm_ffi.dtype("int32").bits // 8 num_elem = reduce(mul, input_shape, 1) input_size = num_elem * input_byte_per_elem # Most GPU algorithms take O(n) space or less, we choose 8N + 8MB as a safe estimation diff --git a/python/tvm/relax/backend/metal/coreml.py b/python/tvm/relax/backend/metal/coreml.py index 64f3c21caf9d..5b076b2ab620 100644 --- a/python/tvm/relax/backend/metal/coreml.py +++ b/python/tvm/relax/backend/metal/coreml.py @@ -320,7 +320,7 @@ def __init__(self, model_name, function): self.model_inputs_ = [] self.buf_idx_ = 0 - getter = tvm_tirx.get_global_func("relax.analysis.get_var2val") + getter = tvm_ffi.get_global_func("relax.analysis.get_var2val") assert getter, "Cannot find `relax.analysis.get_var2val` function." self.var2val = getter(function) diff --git a/python/tvm/relax/base_py_module.py b/python/tvm/relax/base_py_module.py index b851fba313ea..aa4b1eeb3865 100644 --- a/python/tvm/relax/base_py_module.py +++ b/python/tvm/relax/base_py_module.py @@ -24,6 +24,7 @@ from typing import Any, Optional, Union import numpy as np +import tvm_ffi from tvm_ffi import Function import tvm @@ -84,7 +85,7 @@ def __del__(self): # globals this needs may already be cleared. if not owned or sys.is_finalizing(): return - unregister_py_func = tvm_tirx.get_global_func("vm.builtin.unregister_py_func") + unregister_py_func = tvm_ffi.get_global_func("vm.builtin.unregister_py_func") for func_name in owned: unregister_py_func(func_name) except Exception: # pylint: disable=broad-except @@ -224,7 +225,7 @@ def _register_python_functions(self): return try: - register_py_func = tvm_tirx.get_global_func("vm.builtin.register_py_func") + register_py_func = tvm_ffi.get_global_func("vm.builtin.register_py_func") except ValueError: return @@ -304,7 +305,7 @@ def call_dps_packed(self, func_name: str, args, out_ty): if func_name not in self.extern_funcs: try: - self.extern_funcs[func_name] = tvm_tirx.get_global_func(func_name) + self.extern_funcs[func_name] = tvm_ffi.get_global_func(func_name) except ValueError as error: raise ValueError( f"Function '{func_name}' not found as a global function. " @@ -424,14 +425,14 @@ def _convert_single_pytorch_to_tvm(self, tensor: Any) -> Tensor: if _FASTER_DLPACK_EXTENSION is not None: try: dlpack = torch.to_dlpack(tensor) - return tvm_tirx.runtime.from_dlpack(dlpack) + return tvm_ffi.from_dlpack(dlpack) except (AttributeError, ValueError): pass # Fall through to the next method # 2. Try modern `torch.to_dlpack` (preferred for PyTorch >= 1.7) try: dlpack = torch.to_dlpack(tensor) - return tvm_tirx.runtime.from_dlpack(dlpack) + return tvm_ffi.from_dlpack(dlpack) except (AttributeError, ValueError): pass # Fall through to the next method @@ -439,7 +440,7 @@ def _convert_single_pytorch_to_tvm(self, tensor: Any) -> Tensor: if to_dlpack_legacy: try: dlpack = to_dlpack_legacy(tensor) - return tvm_tirx.runtime.from_dlpack(dlpack) + return tvm_ffi.from_dlpack(dlpack) except (AttributeError, ValueError) as error_legacy: print( f"Warning: Legacy DLPack conversion failed ({error_legacy}), " diff --git a/python/tvm/relax/binding_rewrite.py b/python/tvm/relax/binding_rewrite.py index f67481961b56..2c9a3eb69e28 100644 --- a/python/tvm/relax/binding_rewrite.py +++ b/python/tvm/relax/binding_rewrite.py @@ -20,13 +20,12 @@ import tvm_ffi import tvm_tirx -from tvm_tirx.runtime import Object from . import Binding, DataflowBlock, Expr, Function, Var, _ffi_api @tvm_ffi.register_object("relax.DataflowBlockRewrite") -class DataflowBlockRewrite(Object): +class DataflowBlockRewrite(tvm_ffi.Object): """ A binding/statement-level dataflow block rewriter. diff --git a/python/tvm/relax/block_builder.py b/python/tvm/relax/block_builder.py index 8bfcc8b4c310..943f912f1d36 100644 --- a/python/tvm/relax/block_builder.py +++ b/python/tvm/relax/block_builder.py @@ -27,7 +27,6 @@ from tvm import relax as rx from tvm_tirx import tirx from tvm_tirx.ir.module import IRModule -from tvm_tirx.runtime import Object from . import _ffi_api from .expr import BaseFunc, Binding, BindingBlock, Expr, GlobalVar, Tuple, Var @@ -105,7 +104,7 @@ def __exit__(self, ptype, value, trace): @tvm_ffi.register_object("relax.BlockBuilder") -class BlockBuilder(Object): +class BlockBuilder(tvm_ffi.Object): """A builder to build Relax IR for testing and dev. Examples @@ -214,7 +213,7 @@ def function( self, name: str, params: Var | Tuple | list[Var] | None = None, - attrs: dict[str, Object] | None = None, + attrs: dict[str, tvm_ffi.Object] | None = None, pure: bool = True, private: bool = False, ) -> FunctionScope: @@ -230,7 +229,7 @@ def function( If params is None, it means deferring initialization of function parameters until emit_func_output. - attrs : Dict[str, Object], optional + attrs : Dict[str, tvm_ffi.Object], optional The function attrs pure : bool, optional @@ -374,7 +373,7 @@ def call_te_with_grad( func: Callable, *args: Any, te_grad_name: str, - te_grad_kwargs: dict[str, Object] | None = None, + te_grad_kwargs: dict[str, tvm_ffi.Object] | None = None, **kwargs: Any, ) -> Expr: """Generate a call node according to the te function. @@ -393,7 +392,7 @@ def call_te_with_grad( The registered name of the te gradient function associated with the call_tir_with_grad node. Must be provided as a keyword argument. - te_grad_kwargs : Dict[str, Object], optional + te_grad_kwargs : Dict[str, tvm_ffi.Object], optional The keyword arguments passed to the te gradient function. Optionally provided as a keyword argument. Default: {}. diff --git a/python/tvm/relax/distributed/type.py b/python/tvm/relax/distributed/type.py index f7aa67095126..4ef6af71b336 100644 --- a/python/tvm/relax/distributed/type.py +++ b/python/tvm/relax/distributed/type.py @@ -24,7 +24,6 @@ from tvm.relax.type import TensorType, Type from tvm_tirx.ir import Location from tvm_tirx.ir.location import UNKNOWN_LOC -from tvm_tirx.runtime import Object from . import _ffi_api from .global_info import DeviceMesh @@ -36,7 +35,7 @@ class PlacementSpecKind(enum.IntEnum): @tvm_ffi.register_object("relax.distributed.PlacementSpec") -class PlacementSpec(Object): +class PlacementSpec(tvm_ffi.Object): """Describes how data is distributed in one dimension of the device mesh Parameters @@ -83,7 +82,7 @@ def replica() -> "PlacementSpec": @tvm_ffi.register_object("relax.distributed.Placement") -class Placement(Object): +class Placement(tvm_ffi.Object): """Describes how data is distributed in each dimension of the device mesh Parameters diff --git a/python/tvm/relax/dpl/context.py b/python/tvm/relax/dpl/context.py index 748ae7154513..e9d71dd2b137 100644 --- a/python/tvm/relax/dpl/context.py +++ b/python/tvm/relax/dpl/context.py @@ -17,14 +17,14 @@ """The Graph Matching Context Manager for Dataflow Pattern Language.""" -import tvm_tirx +import tvm_ffi from ..expr import DataflowBlock, Var from . import _ffi as ffi from .pattern import DFPattern -class PatternContext(tvm_tirx.runtime.Object): +class PatternContext(tvm_ffi.Object): """A context object for doing graph (topogical) pattern matching.""" def __init__(self, incremental=False): diff --git a/python/tvm/relax/dpl/pattern.py b/python/tvm/relax/dpl/pattern.py index 38ce8e46362b..f9e0c78018fd 100644 --- a/python/tvm/relax/dpl/pattern.py +++ b/python/tvm/relax/dpl/pattern.py @@ -29,7 +29,6 @@ from tvm_tirx.ir import make_node from tvm_tirx.ir.base import Node from tvm_tirx.ir.op import Op -from tvm_tirx.runtime import Object from ..expr import Expr, Var from . import _ffi as ffi @@ -106,13 +105,13 @@ def __invert__(self) -> "NotPattern": """ return reject(self) - def has_attr(self, attrs: dict[str, Object]) -> "AttrPattern": + def has_attr(self, attrs: dict[str, tvm_ffi.Object]) -> "AttrPattern": """ Add an attribute constraint to this pattern Parameters ---------- - attrs: Dict[str, Object] + attrs: Dict[str, tvm_ffi.Object] Returns ------- @@ -976,7 +975,7 @@ def has_attr(attrs, pattern=None) -> AttrPattern: Parameters ---------- - attrs: Dict[str, Object] + attrs: Dict[str, tvm_ffi.Object] The attributes to match pattern: Optional[tvm.relax.dpl.DFPattern] diff --git a/python/tvm/relax/dpl/rewrite.py b/python/tvm/relax/dpl/rewrite.py index c2a848c3ca95..fff1f23d6d29 100644 --- a/python/tvm/relax/dpl/rewrite.py +++ b/python/tvm/relax/dpl/rewrite.py @@ -18,10 +18,10 @@ from collections.abc import Callable +import tvm_ffi from tvm_ffi import register_object from tvm_tirx.ir import IRModule -from tvm_tirx.runtime import Object from ..expr import Expr, Function, Var from . import _ffi as ffi @@ -30,7 +30,7 @@ @register_object("relax.dpl.PatternMatchingRewriter") -class PatternMatchingRewriter(Object): +class PatternMatchingRewriter(tvm_ffi.Object): """A pattern-matching rewriter for Relax""" @staticmethod diff --git a/python/tvm/relax/exec_builder.py b/python/tvm/relax/exec_builder.py index 04ce0099fd26..7b9e85d8893b 100644 --- a/python/tvm/relax/exec_builder.py +++ b/python/tvm/relax/exec_builder.py @@ -110,7 +110,7 @@ def convert_constant(self, const: object) -> int: def emit_call( self, name: str, - args: list[tvm_tirx.runtime.Tensor | tvm_tirx.DataType] | None = None, + args: list[tvm_tirx.runtime.Tensor | tvm_ffi.dtype] | None = None, dst: int | None = None, ) -> None: """emit a call instruction which calls a packed function.""" @@ -124,7 +124,7 @@ def emit_call( shape_tuple = Shape(arg) new_arg = self.convert_constant(shape_tuple) args_.append(new_arg) - elif isinstance(arg, tvm_tirx.runtime.Tensor | tvm_tirx.DataType | Shape): + elif isinstance(arg, tvm_tirx.runtime.Tensor | tvm_ffi.dtype | Shape): new_arg = self.convert_constant(arg) args_.append(new_arg) else: diff --git a/python/tvm/relax/expr.py b/python/tvm/relax/expr.py index 86ed1eefc151..cecf6e9656c0 100644 --- a/python/tvm/relax/expr.py +++ b/python/tvm/relax/expr.py @@ -30,7 +30,6 @@ import tvm.relax import tvm_tirx.ir import tvm_tirx.runtime -from tvm_tirx import DataType from tvm_tirx.ir import BaseFunc, Location, Node, Scriptable from tvm_tirx.ir.expr import _CallableExprWithOp from tvm_tirx.ir.location import UNKNOWN_LOC, UnknownLoc @@ -250,7 +249,7 @@ def __init__( raise TypeError("name must be a str") if ty is not None: - ty = tvm_tirx.runtime.convert(ty) + ty = tvm_ffi.convert(ty) if not isinstance(ty, Type): raise TypeError( "ty needs to be an instance of Type. " diff --git a/python/tvm/relax/expr_functor.py b/python/tvm/relax/expr_functor.py index 945998267598..9341fdaea5d4 100644 --- a/python/tvm/relax/expr_functor.py +++ b/python/tvm/relax/expr_functor.py @@ -28,7 +28,6 @@ from tvm_tirx import tirx as _tirx from tvm_tirx.ir import Call, DataTypeImm, GenericConst, Op, StringImm, TensorLoad, is_prim_expr from tvm_tirx.ir.module import IRModule -from tvm_tirx.runtime import Object from . import _ffi_api from .block_builder import BlockBuilder @@ -944,7 +943,7 @@ def visit_loc(self, loc: Location) -> None: @tvm_ffi.register_object("expr_functor.PyExprMutator") -class _PyExprMutator(Object): +class _PyExprMutator(tvm_ffi.Object): """ A TVM object to support customization of ExprMutator on the python side. This is the decorated result returned from mutator decorator. diff --git a/python/tvm/relax/frontend/nn/core.py b/python/tvm/relax/frontend/nn/core.py index cd8bc3308859..5db354a7f861 100644 --- a/python/tvm/relax/frontend/nn/core.py +++ b/python/tvm/relax/frontend/nn/core.py @@ -33,6 +33,7 @@ ) import numpy as np # type: ignore +import tvm_ffi import tvm_tirx.runtime from tvm.runtime.vm import VirtualMachine @@ -40,7 +41,6 @@ from tvm_tirx.ir import IRModule from tvm_tirx.ir.transform import Pass from tvm_tirx.runtime import Device -from tvm_tirx.runtime import device as as_device from tvm_tirx.target import Target from .... import relax as rx @@ -564,7 +564,7 @@ def _compile(spec, device, pipeline, debug): params = _param_to_tensor(params, device) return spec, vm, params - device = as_device(device) + device = tvm_ffi.device(device) spec, vm, params = _compile(spec, device, pipeline, debug) # pylint: disable=invalid-name if out_format == "torch": @@ -849,7 +849,7 @@ def _attribute_finder(root: Module, prefix: str, condition_yield: Callable[[Any] def _from_dlpack(tensor) -> tvm_tirx.runtime.Tensor: try: - return tvm_tirx.runtime.from_dlpack(tensor) + return tvm_ffi.from_dlpack(tensor) except RuntimeError: pass # special logic for PyTorch diff --git a/python/tvm/relax/frontend/nn/extern.py b/python/tvm/relax/frontend/nn/extern.py index 74f386ea2fd5..f33441880484 100644 --- a/python/tvm/relax/frontend/nn/extern.py +++ b/python/tvm/relax/frontend/nn/extern.py @@ -30,7 +30,6 @@ import tvm_tirx from tvm_tirx import libinfo from tvm_tirx.codegen import load_static_library -from tvm_tirx.runtime import Module from tvm_tirx.support import cc as _cc from ...op import call_dps_packed @@ -80,10 +79,10 @@ def _convert(arg, name: str): return _call - def _load(self, path: Path) -> Module: + def _load(self, path: Path) -> tvm_ffi.Module: return load_static_library(str(path), func_names=list(self._symbols.keys())) - def load(self) -> Module: + def load(self) -> tvm_ffi.Module: """Loads the external module into a TVM runtime module.""" raise NotImplementedError @@ -106,7 +105,7 @@ def __init__( self.filepath = filepath super().__init__(symbols) - def load(self) -> Module: + def load(self) -> tvm_ffi.Module: return self._load(self.filepath) @@ -167,7 +166,7 @@ def shape_dtype_inference(a, b): **A compiler pass `AttachExternModules`.** It is introduced to attach a list of `nn.ExternModule`s into an IRModule at any stage of the compilation pipeline, - and attach the compiled external modules as `runtime.Module`s into IRModule's `external_mods` + and attach the compiled external modules as `tvm_ffi.Module`s into IRModule's `external_mods` attribute. It is required by linking in `tvm.compile`, but with the existence of this pass, source compilation can be deferred to arbitrary stage of TVM compilation. @@ -403,7 +402,7 @@ def compile(self, output_path: Path) -> None: ) shutil.move(str(object_path), str(output_path)) - def load(self) -> Module: + def load(self) -> tvm_ffi.Module: with tempfile.TemporaryDirectory() as temp_dir_str: output_path = Path(temp_dir_str) / f"main{self.output_suffix}" self.compile(output_path) diff --git a/python/tvm/relax/frontend/nn/llm/_kernel_common.py b/python/tvm/relax/frontend/nn/llm/_kernel_common.py index 834ce7f89247..b5f07184db22 100644 --- a/python/tvm/relax/frontend/nn/llm/_kernel_common.py +++ b/python/tvm/relax/frontend/nn/llm/_kernel_common.py @@ -31,17 +31,17 @@ - ``_make_prefill_macros`` — the ``@T.macro`` bundle invoked by the prefill kernels - Tiling config (``_get_prefill_kernel_config``) and scheduling (``_schedule_prefill_kernel``) """ - # pylint: disable=too-many-statements,too-many-arguments,invalid-name,line-too-long import enum import math from typing import Any +import tvm_ffi + import tvm from tvm import s_tir from tvm.s_tir import script as Ts from tvm_tirx import tirx -from tvm_tirx.runtime import DataType from tvm_tirx.script import tirx as T from tvm_tirx.target import Target @@ -439,7 +439,7 @@ def _get_prefill_shared_memory_usage( """ if d_v is None: d_v = d - dtype_bytes = (DataType(dtype).bits + 7) // 8 + dtype_bytes = (tvm_ffi.dtype(dtype).bits + 7) // 8 kv_elements = tile_z * d if merged_kv else tile_z * (d + d_v) qkv_bytes = (tile_x * d + kv_elements) * dtype_bytes softmax_bytes = (tile_x * tile_z + 3 * tile_x) * 4 @@ -607,7 +607,7 @@ def _get_prefill_kernel_config( if d_v is None: d_v = d NUM_BLKS = 16 - dtype_bytes = (DataType(dtype).bits + 7) // 8 + dtype_bytes = (tvm_ffi.dtype(dtype).bits + 7) // 8 LOAD_VEC = 8 // dtype_bytes # 8 bytes group_size = h_q // h_kv @@ -631,7 +631,7 @@ def _get_prefill_kernel_config( # Preserve the established WebGPU config, which targets WebGPU's portable limits. if ( target.kind.name == "webgpu" - and ((d + 127) // 128) * ((DataType(dtype).bits + 15) // 16) >= 4 + and ((d + 127) // 128) * ((tvm_ffi.dtype(dtype).bits + 15) // 16) >= 4 ): tile_z = 8 num_warps = 2 @@ -662,7 +662,7 @@ def _get_prefill_kernel_config( if target.kind.name == "opencl" and ( ("android" in str(target.host)) or ("adreno" in str(target.attrs)) ): - LOAD_VEC = 16 // ((DataType(dtype).bits + 7) // 8) # 16 bytes + LOAD_VEC = 16 // ((tvm_ffi.dtype(dtype).bits + 7) // 8) # 16 bytes NUM_BLKS = group_size * 8 check_thread_limits(target, bdx=bdx, bdy=num_warps, bdz=1, gdz=1) diff --git a/python/tvm/relax/frontend/nn/llm/kv_cache.py b/python/tvm/relax/frontend/nn/llm/kv_cache.py index d86acfd734c3..87b2400a7b13 100644 --- a/python/tvm/relax/frontend/nn/llm/kv_cache.py +++ b/python/tvm/relax/frontend/nn/llm/kv_cache.py @@ -32,11 +32,12 @@ The private-named kernel factories are re-exported from this module so the test suite can continue to import them via ``tvm.relax.frontend.nn.llm.kv_cache``. """ - # pylint: disable=too-many-statements,too-many-arguments,invalid-name,line-too-long import math from typing import Any, Literal +import tvm_ffi + import tvm_tirx from tvm import relax as rx from tvm.relax import ExternFunc as _ExternFunc @@ -121,7 +122,7 @@ class PagedKVCache(Object): # pylint: disable=too-few-public-methods """The Paged KV Cache used in LLM batching for efficient attention computation.""" - extern_mods: list[tvm_tirx.runtime.Module] = [] + extern_mods: list[tvm_ffi.Module] = [] def attention_with_fused_qkv( self, diff --git a/python/tvm/relax/frontend/nn/op.py b/python/tvm/relax/frontend/nn/op.py index 332ef984154f..94728c782d54 100644 --- a/python/tvm/relax/frontend/nn/op.py +++ b/python/tvm/relax/frontend/nn/op.py @@ -331,7 +331,7 @@ def matmul(a: Tensor, b: Tensor, out_dtype: str | None = None, name: str = "matm b : Tensor The second input tensor. - out_dtype: Optional[str | DataType] + out_dtype: Optional[str | tvm_ffi.dtype] The data type of the matmul result. When it is not specified, the output dtype will be the same as input dtype. @@ -627,7 +627,7 @@ def conv1d_transpose( out_layout : Optional[str] Layout of the output. If not specified, it is the same as data_layout - out_dtype : Optional[str | DataType] + out_dtype : Optional[str | tvm_ffi.dtype] Specifies the output data type for mixed precision conv2d. Returns @@ -2182,7 +2182,7 @@ def extern( ) -> OutType: """Invoke an extern function during runtime. The extern function must be registered with the " TVM runtime using `reflection::GlobalDef().def` (C++), or - `tvm_tirx.register_global_func` (Python). + `tvm_ffi.register_global_func` (Python). Parameters ---------- @@ -2239,7 +2239,7 @@ def debug_func( .. code-block:: python - @tvm_tirx.register_global_func(name_of_debug_func) + @tvm_ffi.register_global_func(name_of_debug_func) def debug_func(lineno: str, arg_0, arg_1, ...) -> None: ... diff --git a/python/tvm/relax/frontend/onnx/onnx_frontend.py b/python/tvm/relax/frontend/onnx/onnx_frontend.py index f326edd6b27d..8866019c5523 100644 --- a/python/tvm/relax/frontend/onnx/onnx_frontend.py +++ b/python/tvm/relax/frontend/onnx/onnx_frontend.py @@ -63,7 +63,7 @@ from tvm_tirx.ir import Expr as TvmExpr from tvm_tirx.ir import IRModule from tvm_tirx.ir.unique_name_supply import UniqueNameSupply -from tvm_tirx.runtime import DataType, DataTypeCode +from tvm_tirx.runtime import DataTypeCode from ..common import autopad @@ -71,7 +71,7 @@ def _relax_dtype_is_floating_point(dtype: str) -> bool: """Whether a Relax dtype string is a floating point type.""" try: - code = DataType(dtype).type_code + code = tvm_ffi.dtype(dtype).type_code except (ValueError, TypeError, RuntimeError): return False return ( @@ -699,8 +699,8 @@ def _impl_v7(cls, bb, inputs, attr, params): if tvm_tirx.ir.is_prim_expr(inputs[1]) else inputs[1].ty.dtype.dtype ) - lhs_code = DataType(lhs_dtype).type_code - rhs_code = DataType(rhs_dtype).type_code + lhs_code = tvm_ffi.dtype(lhs_dtype).type_code + rhs_code = tvm_ffi.dtype(rhs_dtype).type_code except (AttributeError, ValueError, TypeError, RuntimeError): return cls.base_impl(bb, inputs, attr, params) diff --git a/python/tvm/relax/frontend/torch/base_fx_graph_translator.py b/python/tvm/relax/frontend/torch/base_fx_graph_translator.py index d7b8d160f3e0..9bb3b77d73e9 100644 --- a/python/tvm/relax/frontend/torch/base_fx_graph_translator.py +++ b/python/tvm/relax/frontend/torch/base_fx_graph_translator.py @@ -603,7 +603,7 @@ def _round(self, node: fx.Node) -> relax.Expr: scale_exp = abs(decimals) # Largest exponent for which 10**n is still finite in `dtype`. # (float16/bfloat16 are upcast to float32 above, so dtype is float32 or - # float64 here.) Note that `dtype` may be a tvm.DataType-like object whose + # float64 here.) Note that `dtype` may be a tvm_ffi.dtype-like object whose # repr is "T.float32" instead of a plain str, so select via == rather than # indexing a str-keyed dict. max_scale_exp = 308 if dtype == "float64" else 38 diff --git a/python/tvm/relax/frontend/torch/exported_program_translator.py b/python/tvm/relax/frontend/torch/exported_program_translator.py index 92d2b8843905..a521bfa4b7c7 100644 --- a/python/tvm/relax/frontend/torch/exported_program_translator.py +++ b/python/tvm/relax/frontend/torch/exported_program_translator.py @@ -25,6 +25,8 @@ from collections.abc import Callable from functools import partial +import tvm_ffi + try: import torch from torch import fx @@ -65,7 +67,7 @@ def _convert_pytorch_tensor_to_tvm(tensor_value: torch.Tensor) -> tvm_tirx.runti # Try DLPack conversion first (faster) try: - return tvm_tirx.runtime.from_dlpack(tensor_detached) + return tvm_ffi.from_dlpack(tensor_detached) except (RuntimeError, BufferError): # Fallback: convert to numpy and then to TVM tensor # This handles cases where DLPack conversion fails diff --git a/python/tvm/relax/global_info.py b/python/tvm/relax/global_info.py index 24af98dcc13d..e271d340e740 100644 --- a/python/tvm/relax/global_info.py +++ b/python/tvm/relax/global_info.py @@ -19,7 +19,7 @@ import tvm_ffi from tvm_tirx.ir import GlobalInfo -from tvm_tirx.runtime import Device, MemoryScope, convert +from tvm_tirx.runtime import Device, MemoryScope from tvm_tirx.target import Target from . import _ffi_api @@ -46,7 +46,7 @@ def __init__( memory_scope: str | MemoryScope = MemoryScope("global"), ) -> None: if isinstance(target, dict | str): - target = Target(convert(target)) + target = Target(tvm_ffi.convert(target)) if isinstance(target, Device): target = Target.from_device(target) self.__init_handle_by_constructor__(_ffi_api.VDevice, target, vdevice_id, memory_scope) diff --git a/python/tvm/relax/op/base.py b/python/tvm/relax/op/base.py index 3c6d676952a5..22915461d5cc 100644 --- a/python/tvm/relax/op/base.py +++ b/python/tvm/relax/op/base.py @@ -28,7 +28,7 @@ from tvm_tirx.ir.attrs import make_node as _make_attrs from tvm_tirx.ir.location import UNKNOWN_LOC, Location from tvm_tirx.ir.op import _make_op_api -from tvm_tirx.runtime import DataTypeCode, Object +from tvm_tirx.runtime import DataTypeCode from ..expr import Expr, ExternFunc, GlobalVar, Var from ..type import Type @@ -367,7 +367,7 @@ def call_builtin_with_ctx( ) -def make_closure(func: Expr, args: Expr, *, ty=None, loc: Location = UNKNOWN_LOC) -> Object: +def make_closure(func: Expr, args: Expr, *, ty=None, loc: Location = UNKNOWN_LOC) -> tvm_ffi.Object: """ Create a closure with free variables and return the closure. @@ -382,7 +382,7 @@ def make_closure(func: Expr, args: Expr, *, ty=None, loc: Location = UNKNOWN_LOC Returns ------- - ret: Object + ret: tvm_ffi.Object The VMClosure. """ @@ -427,13 +427,13 @@ def invoke_closure( ) # type: ignore -def render_object(val: tvm_tirx.Object) -> str: +def render_object(val: tvm_ffi.Object) -> str: """ Given a TVM Object, renders it in string form. Used for Relax printing and assertions. Parameters ---------- - val: tvm_tirx.Object + val: tvm_ffi.Object An object to render Returns @@ -449,7 +449,7 @@ def render_object(val: tvm_tirx.Object) -> str: return str(val) -@tvm_tirx.register_global_func("relax.run.shape_to_tensor") +@tvm_ffi.register_global_func("relax.run.shape_to_tensor") def relax_shape_to_tensor(shape_tuple: tvm_ffi.Shape) -> tvm_tirx.runtime.Tensor: """ Takes a Shape and convert it to Tensor. @@ -462,8 +462,8 @@ def relax_shape_to_tensor(shape_tuple: tvm_ffi.Shape) -> tvm_tirx.runtime.Tensor return tvm_tirx.runtime.tensor([int(v) for v in shape_tuple]) -@tvm_tirx.register_global_func("relax.run.print") -def relax_print(format_str: str, *format_args: tvm_tirx.Object) -> None: +@tvm_ffi.register_global_func("relax.run.print") +def relax_print(format_str: str, *format_args: tvm_ffi.Object) -> None: """ Takes a list of values to print, formats with the given format string. If the format string is empty, simply prints. @@ -478,7 +478,7 @@ def relax_print(format_str: str, *format_args: tvm_tirx.Object) -> None: format_str: str The first argument is a Python-style format string for printing the values - format_args: List[Object] + format_args: List[tvm_ffi.Object] The values to print. """ val_strs = map(render_object, format_args) @@ -493,9 +493,9 @@ def print(format: str | Expr, *values: Expr, ty=None, loc: Location = UNKNOWN_LO return Call("relax.print", [format, *values], loc=loc, ty=ty) -@tvm_tirx.register_global_func("relax.run.assert_op") +@tvm_ffi.register_global_func("relax.run.assert_op") def relax_assert_op( - condition: tvm_tirx.Object, format_str: str, *format_args: tvm_tirx.Object + condition: tvm_ffi.Object, format_str: str, *format_args: tvm_ffi.Object ) -> None: """ A variadic function. The first value serves as the assertion condition: @@ -510,13 +510,13 @@ def relax_assert_op( Parameters ---------- - condition: tvm_tirx.Object + condition: tvm_ffi.Object The assertion condition. Must be a boolean scalar. format_str: str A Python-style format string for the error message. - format_args: List[tvm_tirx.Object] + format_args: List[tvm_ffi.Object] Values used for formatting the string. """ if not isinstance(format_str, str): diff --git a/python/tvm/relax/op/create.py b/python/tvm/relax/op/create.py index 0fac041ca1d1..595324022a72 100644 --- a/python/tvm/relax/op/create.py +++ b/python/tvm/relax/op/create.py @@ -18,7 +18,7 @@ import tvm_ffi -from tvm_tirx import DataType, DataTypeCode +from tvm_tirx import DataTypeCode from tvm_tirx.ir import Attrs, PrimType, is_prim_expr from tvm_tirx.ir import Call as _Call from tvm_tirx.ir.attrs import make_node as _make_attrs @@ -49,7 +49,7 @@ class InitAttrs(Attrs): def full( shape: tuple[PrimExprLike] | Expr, fill_value: Expr, - dtype: str | DataType | None = None, + dtype: str | tvm_ffi.dtype | None = None, *, ty=None, loc: Location = UNKNOWN_LOC, @@ -64,7 +64,7 @@ def full( fill_value : relax.Expr The value to fill. Must be a scalar tensor. - dtype : Optional[str | DataType] + dtype : Optional[str | tvm_ffi.dtype] The data type of the created tensor. If dtype is not given, it will by default use the dtype of fill_value. @@ -86,7 +86,7 @@ def full( def full_like( x: Expr, fill_value: Expr, - dtype: str | DataType | None = None, + dtype: str | tvm_ffi.dtype | None = None, *, ty=None, loc: Location = UNKNOWN_LOC, @@ -104,7 +104,7 @@ def full_like( fill_value : relax.Expr The value to fill. Must be a scalar tensor. - dtype : Optional[str | DataType] + dtype : Optional[str | tvm_ffi.dtype] The data type of the created tensor. If dtype is not given, it will by default use the dtype of the input tensor. @@ -124,7 +124,7 @@ def full_like( def ones( shape: tuple[PrimExprLike] | Expr, - dtype: str | DataType, + dtype: str | tvm_ffi.dtype, *, ty=None, loc: Location = UNKNOWN_LOC, @@ -136,7 +136,7 @@ def ones( shape : Union[Tuple[PrimExprLike], Expr] The shape of the created tensor. - dtype : str | DataType + dtype : str | tvm_ffi.dtype The data type of the created tensor. Returns @@ -155,7 +155,7 @@ def ones( def ones_like( - x: Expr, dtype: str | DataType | None = None, *, ty=None, loc: Location = UNKNOWN_LOC + x: Expr, dtype: str | tvm_ffi.dtype | None = None, *, ty=None, loc: Location = UNKNOWN_LOC ) -> Expr: """Construct a tensor with all ones, with shape of the input tensor shape. @@ -165,7 +165,7 @@ def ones_like( The input tensor, which provides the shape, and dtype when the `dtype` field is not specified. - dtype : Optional[str | DataType] + dtype : Optional[str | tvm_ffi.dtype] The data type of the created tensor. If dtype is not given, it will by default use the dtype of the input tensor. @@ -185,7 +185,7 @@ def ones_like( def zeros( shape: tuple[PrimExprLike] | Expr, - dtype: str | DataType, + dtype: str | tvm_ffi.dtype, *, ty=None, loc: Location = UNKNOWN_LOC, @@ -197,7 +197,7 @@ def zeros( shape : Union[Tuple[PrimExprLike], Expr] The shape of the created tensor. - dtype : str | DataType + dtype : str | tvm_ffi.dtype The data type of the created tensor. Returns @@ -216,7 +216,7 @@ def zeros( def zeros_like( - x: Expr, dtype: str | DataType | None = None, *, ty=None, loc: Location = UNKNOWN_LOC + x: Expr, dtype: str | tvm_ffi.dtype | None = None, *, ty=None, loc: Location = UNKNOWN_LOC ) -> Expr: """Construct a tensor with all zeros, with shape of the input tensor shape. @@ -226,7 +226,7 @@ def zeros_like( The input tensor, which provides the shape, and dtype when the `dtype` field is not specified. - dtype : Optional[str | DataType] + dtype : Optional[str | tvm_ffi.dtype] The data type of the created tensor. If dtype is not given, it will by default use the dtype of the input tensor. @@ -248,7 +248,7 @@ def eye( n: PrimExprLike, m: PrimExprLike | None = None, k: PrimExprLike = 0, - dtype: str | DataType = "float32", + dtype: str | tvm_ffi.dtype = "float32", *, ty=None, loc: Location = UNKNOWN_LOC, @@ -268,7 +268,7 @@ def eye( a positive value refers to an upper diagonal, and a negative value to a lower diagonal. - dtype : str | DataType + dtype : str | tvm_ffi.dtype The data type of the created tensor. Returns @@ -292,7 +292,7 @@ def eye( def eye_like( x: Expr, k: PrimExprLike = 0, - dtype: str | DataType | None = None, + dtype: str | tvm_ffi.dtype | None = None, *, ty=None, loc: Location = UNKNOWN_LOC, @@ -311,7 +311,7 @@ def eye_like( a positive value refers to an upper diagonal, and a negative value to a lower diagonal. - dtype : Optional[str | DataType] + dtype : Optional[str | tvm_ffi.dtype] The data type of the created tensor. If dtype is not given, it will by default use the dtype of the input tensor. @@ -334,7 +334,7 @@ def arange( start: PrimExprLike, end: PrimExprLike | None = None, step: PrimExprLike = 1, - dtype: str | DataType | PrimType | None = None, + dtype: str | tvm_ffi.dtype | PrimType | None = None, *, ty=None, loc: Location = UNKNOWN_LOC, @@ -353,7 +353,7 @@ def arange( step : PrimExprLike The step size. - dtype : Optional[str | DataType | PrimType] + dtype : Optional[str | tvm_ffi.dtype | PrimType] The data type of the created tensor. Returns diff --git a/python/tvm/relax/op/datatype.py b/python/tvm/relax/op/datatype.py index 1de6855675e0..89d0d5693218 100644 --- a/python/tvm/relax/op/datatype.py +++ b/python/tvm/relax/op/datatype.py @@ -18,7 +18,6 @@ import tvm_ffi -from tvm_tirx import DataType from tvm_tirx.ir import Attrs, PrimType from tvm_tirx.ir import Call as _Call from tvm_tirx.ir.attrs import make_node as _make_attrs @@ -37,7 +36,7 @@ class AstypeAttrs(Attrs): def astype( - x: Expr, dtype: str | DataType | PrimType, *, ty=None, loc: Location = UNKNOWN_LOC + x: Expr, dtype: str | tvm_ffi.dtype | PrimType, *, ty=None, loc: Location = UNKNOWN_LOC ) -> Expr: """Cast input tensor to the given data type. @@ -46,7 +45,7 @@ def astype( x : relax.Expr The input data to the operator. - dtype: Union[str, DataType] + dtype: Union[str, tvm_ffi.dtype] The target data type Returns @@ -70,7 +69,7 @@ class WrapParamAttrs(Attrs): def wrap_param( data: Expr, - dtype: str | DataType | PrimType = "float32", + dtype: str | tvm_ffi.dtype | PrimType = "float32", *, ty=None, loc: Location = UNKNOWN_LOC, @@ -81,7 +80,7 @@ def wrap_param( ---------- data : relax.Expr The input data to the operator. - dtype : Union[str, DataType] + dtype : Union[str, tvm_ffi.dtype] The target data type Returns ------- diff --git a/python/tvm/relax/op/image/image.py b/python/tvm/relax/op/image/image.py index a3fde5f2db22..f2a159497600 100644 --- a/python/tvm/relax/op/image/image.py +++ b/python/tvm/relax/op/image/image.py @@ -20,7 +20,6 @@ import tvm_ffi -from tvm_tirx import DataType from tvm_tirx.ir import Attrs, is_prim_expr from tvm_tirx.ir import Call as _Call from tvm_tirx.ir import PrimType as _PrimType @@ -49,7 +48,7 @@ def resize2d( cubic_alpha: float = -0.75, cubic_exclude: int = 0, extrapolation_value: float = 0.0, - out_dtype: str | DataType | None = None, + out_dtype: str | tvm_ffi.dtype | None = None, *, ty=None, loc: Location = UNKNOWN_LOC, @@ -105,7 +104,7 @@ def resize2d( extrapolation_value: float Fill value to use when roi is outside of the image - out_dtype : Optional[str | DataType] + out_dtype : Optional[str | tvm_ffi.dtype] The dtype of the output tensor. It it is not specified, the output will have the same dtype as input if not specified. @@ -162,7 +161,7 @@ def resize3d( cubic_alpha: float = -0.75, cubic_exclude: int = 0, extrapolation_value: float = 0.0, - out_dtype: str | DataType | None = None, + out_dtype: str | tvm_ffi.dtype | None = None, *, ty=None, loc: Location = UNKNOWN_LOC, diff --git a/python/tvm/relax/op/linear_algebra.py b/python/tvm/relax/op/linear_algebra.py index e4cc50debb20..c1a6724136fa 100644 --- a/python/tvm/relax/op/linear_algebra.py +++ b/python/tvm/relax/op/linear_algebra.py @@ -19,7 +19,6 @@ import tvm_ffi -from tvm_tirx import DataType from tvm_tirx.ir import Attrs from tvm_tirx.ir import Call as _Call from tvm_tirx.ir import PrimType as _PrimType @@ -39,7 +38,7 @@ class MatmulAttrs(Attrs): def matmul( x1: Expr, x2: Expr, - out_dtype: str | DataType | None = None, + out_dtype: str | tvm_ffi.dtype | None = None, *, ty=None, loc: Location = UNKNOWN_LOC, @@ -57,7 +56,7 @@ def matmul( x2 : relax.Expr The second input tensor. - out_dtype: Optional[Union[str, DataType]] + out_dtype: Optional[Union[str, tvm_ffi.dtype]] The data type of the matmul result. When it is not specified, the output dtype will be the same as input dtype. @@ -82,7 +81,7 @@ def linear( data: Expr, weight: Expr, bias: Expr | None = None, - out_dtype: str | DataType | None = None, + out_dtype: str | tvm_ffi.dtype | None = None, ) -> Expr: """Applies a linear transformation to the incoming data: y = xA^T + b @@ -97,7 +96,7 @@ def linear( bias : Optional[Expr] The bias tensor. - out_dtype: Optional[Union[str, DataType]] + out_dtype: Optional[Union[str, tvm_ffi.dtype]] The data type of the matmul result. When it is not specified, the output dtype will be the same as input dtype. diff --git a/python/tvm/relax/op/nn/nn.py b/python/tvm/relax/op/nn/nn.py index 8193abc8da01..2ea810d1908c 100644 --- a/python/tvm/relax/op/nn/nn.py +++ b/python/tvm/relax/op/nn/nn.py @@ -19,7 +19,6 @@ import tvm_ffi from tvm import relax -from tvm_tirx import DataType from tvm_tirx.ir import Attrs, FloatImm from tvm_tirx.ir import Call as _Call from tvm_tirx.ir import PrimType as _PrimType @@ -44,7 +43,7 @@ def conv1d( data_layout: str = "NCW", kernel_layout: str = "OIW", out_layout: str | None = None, - out_dtype: str | DataType | None = None, + out_dtype: str | tvm_ffi.dtype | None = None, *, ty_args=None, ty=None, @@ -107,7 +106,7 @@ def conv1d( out_layout : Optional[str] Layout of the output. If not specified, it is the same as data_layout - out_dtype : Optional[Union[str, DataType]] + out_dtype : Optional[Union[str, tvm_ffi.dtype]] Specifies the output data type for mixed precision conv1d. Returns @@ -167,7 +166,7 @@ def conv2d( data_layout: str = "NCHW", kernel_layout: str = "OIHW", out_layout: str | None = None, - out_dtype: str | DataType | None = None, + out_dtype: str | tvm_ffi.dtype | None = None, *, ty_args=None, ty=None, @@ -230,7 +229,7 @@ def conv2d( out_layout : Optional[str] Layout of the output. If not specified, it is the same as data_layout - out_dtype : Optional[Union[str, DataType]] + out_dtype : Optional[Union[str, tvm_ffi.dtype]] Specifies the output data type for mixed precision conv2d. Returns @@ -294,7 +293,7 @@ def conv3d( data_layout: str = "NCDHW", kernel_layout: str = "OIDHW", out_layout: str | None = None, - out_dtype: str | DataType | None = None, + out_dtype: str | tvm_ffi.dtype | None = None, *, ty_args=None, ty=None, @@ -359,7 +358,7 @@ def conv3d( out_layout : Optional[str] Layout of the output. If not specified, it is the same as data_layout - out_dtype : Optional[Union[str, DataType]] + out_dtype : Optional[Union[str, tvm_ffi.dtype]] Specifies the output data type for mixed precision conv2d. See Also @@ -428,7 +427,7 @@ def conv1d_transpose( data_layout: str = "NCW", kernel_layout: str = "IOW", out_layout: str | None = None, - out_dtype: str | DataType | None = None, + out_dtype: str | tvm_ffi.dtype | None = None, *, ty_args=None, ty=None, @@ -480,7 +479,7 @@ def conv1d_transpose( out_layout : Optional[str] Layout of the output. If not specified, it is the same as data_layout - out_dtype : Optional[Union[str, DataType]] + out_dtype : Optional[Union[str, tvm_ffi.dtype]] Specifies the output data type for mixed precision conv2d. Returns @@ -546,7 +545,7 @@ def conv2d_transpose( data_layout: str = "NCHW", kernel_layout: str = "IOHW", out_layout: str | None = None, - out_dtype: str | DataType | None = None, + out_dtype: str | tvm_ffi.dtype | None = None, *, ty_args=None, ty=None, @@ -608,7 +607,7 @@ def conv2d_transpose( out_layout : Optional[str] Layout of the output. If not specified, it is the same as data_layout - out_dtype : Optional[Union[str, DataType]] + out_dtype : Optional[Union[str, tvm_ffi.dtype]] Specifies the output data type for mixed precision conv2d. Returns @@ -681,7 +680,7 @@ def conv3d_transpose( data_layout: str = "NCDHW", kernel_layout: str = "IODHW", out_layout: str | None = None, - out_dtype: str | DataType | None = None, + out_dtype: str | tvm_ffi.dtype | None = None, *, ty_args=None, ty=None, @@ -739,7 +738,7 @@ def conv3d_transpose( out_layout : Optional[str] Layout of the output. If not specified, it is the same as data_layout - out_dtype : Optional[Union[str, DataType]] + out_dtype : Optional[Union[str, tvm_ffi.dtype]] Specifies the output data type for mixed precision conv3d_transpose. See Also diff --git a/python/tvm/relax/op/set.py b/python/tvm/relax/op/set.py index 4fa33d3b7cf5..51a60f641ecc 100644 --- a/python/tvm/relax/op/set.py +++ b/python/tvm/relax/op/set.py @@ -18,6 +18,7 @@ """Set operators.""" import numpy as np # type: ignore +import tvm_ffi import tvm_tirx from tvm_tirx.ir import Call as _Call @@ -91,7 +92,7 @@ def unique( ) -@tvm_tirx.register_global_func("relax.run.unique") +@tvm_ffi.register_global_func("relax.run.unique") def numpy_unique( x: tvm_tirx.runtime.tensor, sorted: int, @@ -217,7 +218,7 @@ def nonzero(x: Expr, *, ty=None, loc: Location = UNKNOWN_LOC) -> Expr: return _Call("relax.nonzero", [x], loc=loc, ty=ty) # type: ignore -@tvm_tirx.register_global_func("relax.run.nonzero") +@tvm_ffi.register_global_func("relax.run.nonzero") def numpy_nonzero(x: tvm_tirx.runtime.tensor) -> tvm_tirx.runtime.tensor: np_result = np.atleast_1d(x.numpy()).nonzero() return tvm_tirx.runtime.tensor(np.stack(np_result, axis=0)) diff --git a/python/tvm/relax/op/statistical.py b/python/tvm/relax/op/statistical.py index cd06565499cd..e080eea29a81 100644 --- a/python/tvm/relax/op/statistical.py +++ b/python/tvm/relax/op/statistical.py @@ -19,7 +19,6 @@ import tvm_ffi -from tvm_tirx import DataType from tvm_tirx.ir import Attrs, PrimType from tvm_tirx.ir import Call as _Call from tvm_tirx.ir.attrs import make_node as _make_attrs @@ -291,7 +290,7 @@ class ScanopAttrs(Attrs): def cumprod( data: Expr, axis: int | None = None, - dtype: str | DataType | None = None, + dtype: str | tvm_ffi.dtype | None = None, exclusive: bool = False, *, ty=None, @@ -309,7 +308,7 @@ def cumprod( Axis along which the cumulative product is computed. The default (None) is to compute the cumprod over the flattened array. - dtype : Optional[Union[str, DataType]] + dtype : Optional[Union[str, tvm_ffi.dtype]] Type of the returned array and of the accumulator in which the elements are computed. If dtype is not specified, it defaults to the dtype of data. @@ -364,7 +363,7 @@ def cumprod( def cumsum( data: Expr, axis: int | None = None, - dtype: str | DataType | None = None, + dtype: str | tvm_ffi.dtype | None = None, exclusive: bool = False, *, ty=None, @@ -382,7 +381,7 @@ def cumsum( Axis along which the cumulative sum is computed. The default (None) is to compute the cumsum over the flattened array. - dtype : Optional[Union[str, DataType]] + dtype : Optional[Union[str, tvm_ffi.dtype]] Type of the returned array and of the accumulator in which the elements are summed. If dtype is not specified, it defaults to the dtype of data. diff --git a/python/tvm/relax/relax_to_pyfunc_converter.py b/python/tvm/relax/relax_to_pyfunc_converter.py index b1643f51d48d..e4243b56fd46 100644 --- a/python/tvm/relax/relax_to_pyfunc_converter.py +++ b/python/tvm/relax/relax_to_pyfunc_converter.py @@ -27,6 +27,7 @@ import numpy # pylint: disable=unused-import import torch import torch.nn.functional as F +import tvm_ffi import tvm import tvm_tirx @@ -635,7 +636,7 @@ def _convert_call_tir(self, call: relax.Call, args: list[Any]) -> Any: # If not found in current module, try global registry if tir_function is None: - tir_function = tvm_tirx.get_global_func(func_name) + tir_function = tvm_ffi.get_global_func(func_name) if tir_function is None: if len(converted_args) >= 2: @@ -650,7 +651,7 @@ def _convert_call_tir(self, call: relax.Call, args: list[Any]) -> Any: try: if isinstance(arg, torch.Tensor): # Convert PyTorch tensor to TVM NDArray via DLPack - tvm_arg = runtime.from_dlpack(torch.to_dlpack(arg)) + tvm_arg = tvm_ffi.from_dlpack(torch.to_dlpack(arg)) tvm_args.append(tvm_arg) else: tvm_args.append(arg) @@ -745,7 +746,7 @@ def _convert_call_dps_packed(self, call: relax.Call, args: list[Any]) -> Any: try: # Get the packed function from TVM - packed_function = tvm_tirx.get_global_func(func_name) + packed_function = tvm_ffi.get_global_func(func_name) if packed_function is None: return converted_args[0] if converted_args else torch.tensor([]) @@ -754,7 +755,7 @@ def _convert_call_dps_packed(self, call: relax.Call, args: list[Any]) -> Any: for arg in converted_args: if isinstance(arg, torch.Tensor): # Convert PyTorch tensor to TVM NDArray via DLPack - tvm_arg = runtime.from_dlpack(torch.to_dlpack(arg)) + tvm_arg = tvm_ffi.from_dlpack(torch.to_dlpack(arg)) tvm_args.append(tvm_arg) else: tvm_args.append(arg) diff --git a/python/tvm/relax/script/ir_builder/distributed/ir.py b/python/tvm/relax/script/ir_builder/distributed/ir.py index a214ed758870..6292907b2074 100644 --- a/python/tvm/relax/script/ir_builder/distributed/ir.py +++ b/python/tvm/relax/script/ir_builder/distributed/ir.py @@ -21,6 +21,7 @@ from numbers import Number import numpy as _np # type: ignore +import tvm_ffi import tvm_tirx from tvm.relax.distributed import DeviceMesh, DTensorType @@ -53,7 +54,7 @@ def const( Python scalars and NumPy values are converted to the dtype specified by ``ty.tensor_ty``. The dtype is not inferred from the Python value. """ - ty = tvm_tirx.runtime.convert(ty) + ty = tvm_ffi.convert(ty) if not isinstance(ty, DTensorType): raise TypeError("ty needs to be an instance of DTensorType. ") dtype = str(ty.tensor_ty.dtype) diff --git a/python/tvm/relax/script/ir_builder/ir.py b/python/tvm/relax/script/ir_builder/ir.py index 3cc977e31e10..72739a0888d8 100644 --- a/python/tvm/relax/script/ir_builder/ir.py +++ b/python/tvm/relax/script/ir_builder/ir.py @@ -40,7 +40,6 @@ from tvm_tirx.runtime._tensor import ( cpu, cuda, - device, ext_dev, hexagon, metal, @@ -567,7 +566,6 @@ def match_cast(value, ty, *, loc: _LocationEntry | Location = UNKNOWN_LOC): "const", "cpu", "cuda", - "device", "device_mesh", "dtype", "dummy_global_info", diff --git a/python/tvm/relax/script/ir_builder/parser_protocol.py b/python/tvm/relax/script/ir_builder/parser_protocol.py index f1badec9fa91..0dcdb9984eb4 100644 --- a/python/tvm/relax/script/ir_builder/parser_protocol.py +++ b/python/tvm/relax/script/ir_builder/parser_protocol.py @@ -28,6 +28,7 @@ from collections.abc import Callable, Sequence from typing import Any, NoReturn +import tvm_ffi import tvm_ffi as _ffi from tvm import relax as _relax @@ -36,7 +37,6 @@ from tvm.relax.utils import gen_call_tir_inputs from tvm_tirx import ir as _ir from tvm_tirx.ir.location import UnknownLoc -from tvm_tirx.runtime import Object as tvm_Object from tvm_tirx.script.ir_builder import base as _base from tvm_tirx.script.ir_builder.base import IRBuilder as _IRBuilder from tvm_tirx.script.ir_builder.parser_protocol import decl_function @@ -129,7 +129,7 @@ def func_name_(name: str) -> None: return _ffi_api.FuncName(name) -def func_attr(attrs: dict[str, tvm_Object]) -> None: +def func_attr(attrs: dict[str, tvm_ffi.Object]) -> None: """Specify the attrs of the last function frame. Parameters @@ -258,7 +258,7 @@ def bind_( if isinstance(name_loc, _base.LocationEntry): name_loc = name_loc.loc if frame_value: - if isinstance(value, _python.list | _python.tuple | _ir.Array): + if isinstance(value, _python.list | _python.tuple | tvm_ffi.Array): for index, item in enumerate(value): bind_( item, diff --git a/python/tvm/relax/script/printer.py b/python/tvm/relax/script/printer.py index b0fca7167a8c..6e19bc442c57 100644 --- a/python/tvm/relax/script/printer.py +++ b/python/tvm/relax/script/printer.py @@ -20,7 +20,6 @@ from tvm_ffi.access_path import AccessPath from tvm_tirx.ir.base import register_dialect_script_method -from tvm_tirx.runtime import Object from tvm_tirx.script.printer import PrinterConfig @@ -38,8 +37,8 @@ def _relax_script( extra_config: dict | None = None, path_to_underline: list[AccessPath] | None = None, path_to_annotate: dict[AccessPath, str] | None = None, - obj_to_underline: list[Object] | None = None, - obj_to_annotate: dict[Object, str] | None = None, + obj_to_underline: list[tvm_ffi.Object] | None = None, + obj_to_annotate: dict[tvm_ffi.Object, str] | None = None, ) -> str: merged_extra = dict(extra_config or {}) return tvm_ffi.get_global_func("script.printer.ReprPrintRelax")( diff --git a/python/tvm/relax/testing/lib_comparator.py b/python/tvm/relax/testing/lib_comparator.py index 2a25ceb5da21..e3c396bb1506 100644 --- a/python/tvm/relax/testing/lib_comparator.py +++ b/python/tvm/relax/testing/lib_comparator.py @@ -37,7 +37,7 @@ class LibCompareVMInstrument: Parameters ---------- - mod: runtime.Module + mod: tvm_ffi.Module The module of interest to be validated. device: runtime.Device diff --git a/python/tvm/relax/testing/transform.py b/python/tvm/relax/testing/transform.py index a46ef9be1990..8ececd18e96e 100644 --- a/python/tvm/relax/testing/transform.py +++ b/python/tvm/relax/testing/transform.py @@ -26,12 +26,11 @@ from tvm.relax.expr import DataflowBlock, Var from tvm_tirx.ir import Call from tvm_tirx.ir.module import IRModule -from tvm_tirx.runtime import Object def ApplyEmptyCppMutator() -> tvm_tirx.ir.transform.Pass: """Create empty cpp mutator""" - packed_func = tvm_tirx.get_global_func("relax.testing.transform.ApplyEmptyCppMutator") + packed_func = tvm_ffi.get_global_func("relax.testing.transform.ApplyEmptyCppMutator") return packed_func() @@ -42,9 +41,7 @@ def dataflow_liveness_analysis(block: DataflowBlock) -> dict[Var, tuple[int, int if "PYTEST_CURRENT_TEST" not in os.environ: logging.warning("The function dataflow_liveness_analysis is exposed for testing only.") - live_ranges = tvm_tirx.get_global_func("relax.testing.transform.DataflowLivenessAnalysis")( - block - ) # type: ignore + live_ranges = tvm_ffi.get_global_func("relax.testing.transform.DataflowLivenessAnalysis")(block) # type: ignore ret = {} for var, live_range in live_ranges.items(): ret[var] = tuple(live_range) @@ -60,7 +57,7 @@ def dataflow_alias_analysis( if "PYTEST_CURRENT_TEST" not in os.environ: logging.warning("The function dataflow_alias_analysis is exposed for testing only.") - alias_sets, tuple_map = tvm_tirx.get_global_func( + alias_sets, tuple_map = tvm_ffi.get_global_func( "relax.testing.transform.DataflowAliasAnalysis" )( block, @@ -76,7 +73,7 @@ def dataflow_alias_analysis( @tvm_ffi.register_object("relax.transform.InplaceOpportunity") -class InplaceOpportunity(Object): +class InplaceOpportunity(tvm_ffi.Object): """ Represents an opportunity to make a binding in-place. Exposed only for testing; the constructor is not exposed. @@ -102,7 +99,7 @@ def dataflow_inplace_analysis( """ if "PYTEST_CURRENT_TEST" not in os.environ: logging.warning("The function dataflow_inplace_analysis is exposed for testing only.") - index_lists = tvm_tirx.get_global_func("relax.testing.transform.DataflowInplaceAnalysis")( + index_lists = tvm_ffi.get_global_func("relax.testing.transform.DataflowInplaceAnalysis")( block, inputs, mod ) # type: ignore @@ -121,7 +118,7 @@ def dataflow_single_inplace_call( if "PYTEST_CURRENT_TEST" not in os.environ: logging.warning("The function dataflow_single_inplace_call is exposed for testing only.") - ret = tvm_tirx.get_global_func("relax.testing.transform.SingleInplaceCall")( + ret = tvm_ffi.get_global_func("relax.testing.transform.SingleInplaceCall")( mod, call, inplace_indices, diff --git a/python/tvm/relax/testing/vm.py b/python/tvm/relax/testing/vm.py index 008afc6e5013..df81dc67e4c0 100644 --- a/python/tvm/relax/testing/vm.py +++ b/python/tvm/relax/testing/vm.py @@ -20,62 +20,64 @@ from typing import Any import numpy as np # type: ignore +import tvm_ffi import tvm_tirx from tvm import relax -from tvm_tirx.runtime import Object -@tvm_tirx.register_global_func("test.vm.move") +@tvm_ffi.register_global_func("test.vm.move") def move(src): return src -@tvm_tirx.register_global_func("test.vm.add") +@tvm_ffi.register_global_func("test.vm.add") def add(a, b): ret = a.numpy() + b.numpy() return tvm_tirx.runtime.tensor(ret) -@tvm_tirx.register_global_func("test.vm.mul") +@tvm_ffi.register_global_func("test.vm.mul") def mul(a, b): ret = a.numpy() * b.numpy() return tvm_tirx.runtime.tensor(ret) -@tvm_tirx.register_global_func("test.vm.equal_zero") +@tvm_ffi.register_global_func("test.vm.equal_zero") def equal_zero(a): ret = np.all(a.numpy() == 0) return tvm_tirx.runtime.tensor(ret) -@tvm_tirx.register_global_func("test.vm.subtract_one") +@tvm_ffi.register_global_func("test.vm.subtract_one") def subtract_one(a): ret = np.subtract(a.numpy(), 1) return tvm_tirx.runtime.tensor(ret) -@tvm_tirx.register_global_func("test.vm.identity") +@tvm_ffi.register_global_func("test.vm.identity") def identity_packed(a, b): b[:] = tvm_tirx.runtime.tensor(a.numpy()) -@tvm_tirx.register_global_func("test.vm.tile") +@tvm_ffi.register_global_func("test.vm.tile") def tile_packed(a, b): b[:] = tvm_tirx.runtime.tensor(np.tile(a.numpy(), (1, 2))) -@tvm_tirx.register_global_func("test.vm.add_scalar") +@tvm_ffi.register_global_func("test.vm.add_scalar") def add_scalar(a, b): return a + b -@tvm_tirx.register_global_func("test.vm.get_device_id") +@tvm_ffi.register_global_func("test.vm.get_device_id") def get_device_id(device): return device.index -def check_saved_func(vm: relax.VirtualMachine, func_name: str, *inputs: list[Any]) -> Object: +def check_saved_func( + vm: relax.VirtualMachine, func_name: str, *inputs: list[Any] +) -> tvm_ffi.Object: # uses save_function to create a closure with the given inputs # and ensure the result is the same # (assumes the functions return tensors and that they're idempotent) @@ -87,6 +89,6 @@ def check_saved_func(vm: relax.VirtualMachine, func_name: str, *inputs: list[Any return res1 -@tvm_tirx.register_global_func("test.vm.check_if_defined") -def check_if_defined(obj: tvm_tirx.Object) -> tvm_tirx.ir.IntImm: - return tvm_tirx.runtime.convert(obj is not None) +@tvm_ffi.register_global_func("test.vm.check_if_defined") +def check_if_defined(obj: tvm_ffi.Object) -> tvm_tirx.ir.IntImm: + return tvm_ffi.convert(obj is not None) diff --git a/python/tvm/relax/training/optimizer.py b/python/tvm/relax/training/optimizer.py index f704dbad41b1..19c118303ac7 100644 --- a/python/tvm/relax/training/optimizer.py +++ b/python/tvm/relax/training/optimizer.py @@ -150,7 +150,7 @@ def _set_params_and_dtype(self, params: list[Var]) -> None: f"Optimizers only support Tensor parameters, but parameter {x.name} has " f"type {x.ty}" ) - data_type = tvm_tirx.DataType(x.ty.dtype.dtype) + data_type = tvm_ffi.dtype(x.ty.dtype.dtype) if data_type.type_code not in ( tvm_tirx.DataTypeCode.BFLOAT, tvm_tirx.DataTypeCode.FLOAT, diff --git a/python/tvm/relax/transform/ipc_allreduce_rewrite.py b/python/tvm/relax/transform/ipc_allreduce_rewrite.py index 530fcccf8633..d67d741ba1a6 100644 --- a/python/tvm/relax/transform/ipc_allreduce_rewrite.py +++ b/python/tvm/relax/transform/ipc_allreduce_rewrite.py @@ -18,6 +18,8 @@ The pass is written in Python for experiment, fast development. """ +import tvm_ffi + import tvm_tirx from tvm import relax from tvm.relax.expr import Expr, Var @@ -42,7 +44,7 @@ def __init__(self, allreduce_strategy: int) -> None: def transform_module(self, mod: IRModule, _ctx: tvm_tirx.ir.transform.PassContext) -> IRModule: """IRModule-level transformation""" - fcustom_allreduce = tvm_tirx.get_global_func( + fcustom_allreduce = tvm_ffi.get_global_func( "runtime.disco.cuda_ipc.custom_allreduce", allow_missing=True ) if fcustom_allreduce is None: diff --git a/python/tvm/relax/transform/transform.py b/python/tvm/relax/transform/transform.py index 2f634621a5bc..445c3bd94a65 100644 --- a/python/tvm/relax/transform/transform.py +++ b/python/tvm/relax/transform/transform.py @@ -32,7 +32,7 @@ import tvm_tirx.ir from tvm.relax import Expr, Type, Var from tvm.relax.dpl import DFPattern -from tvm_tirx.runtime import Object, Tensor +from tvm_tirx.runtime import Tensor from tvm_tirx.tirx import Function, IndexMap from ..expr import Var @@ -815,7 +815,7 @@ def FuseTIR() -> tvm_tirx.ir.transform.Pass: @tvm_ffi.register_object("relax.transform.PatternCheckContext") -class PatternCheckContext(Object): +class PatternCheckContext(tvm_ffi.Object): """ The input of check function `FusionPattern.check`. @@ -849,7 +849,7 @@ class PatternCheckContext(Object): @tvm_ffi.register_object("relax.transform.FusionPattern") -class FusionPattern(Object): +class FusionPattern(tvm_ffi.Object): """ The pattern used by `FuseOpsByPattern`. It's mainly DFPattern but with other information to help during the fusion pass. @@ -1418,7 +1418,7 @@ def SplitCallTIRByPattern( patterns : List[Function] The list of patterns to match. - fcodegen: Callable[[List[MatchResult]], List[Object]] + fcodegen: Callable[[List[MatchResult]], List[tvm_ffi.Object]] The function to generate the code for the matched patterns. Returns diff --git a/python/tvm/relax/vm_build.py b/python/tvm/relax/vm_build.py index 6a982e4ac833..94ac42c3f3ba 100644 --- a/python/tvm/relax/vm_build.py +++ b/python/tvm/relax/vm_build.py @@ -17,6 +17,8 @@ # pylint: disable=invalid-name, no-member """VM build logics""" +import tvm_ffi + import tvm_tirx from tvm import relax from tvm_tirx.codegen import Executable @@ -29,7 +31,7 @@ class VMExecutable(Executable): """The virtual machine executable object emitted by the VM compiler or the ExecBuilder.""" - def __init__(self, mod: tvm_tirx.runtime.Module): + def __init__(self, mod: tvm_ffi.Module): super().__init__(mod) self._stats = self.mod["stats"] self._as_text = self.mod["as_text"] @@ -90,7 +92,7 @@ def _auto_attach_system_lib_prefix( return tir_mod -def _is_device_module(mod: tvm_tirx.runtime.Module) -> bool: +def _is_device_module(mod: tvm_ffi.Module) -> bool: return mod.kind in ["cuda", "opencl", "metal", "hip", "vulkan", "webgpu"] @@ -99,7 +101,7 @@ def _vmlink( target: str | tvm_tirx.target.Target | None, tir_mod: tvm_tirx.IRModule | None = None, tir_pipeline: str | tvm_tirx.ir.transform.Pass | None = "default", - ext_libs: list[tvm_tirx.runtime.Module] | None = None, + ext_libs: list[tvm_ffi.Module] | None = None, params: dict[str, list] | None = None, *, system_lib: bool | None = None, @@ -124,7 +126,7 @@ def _vmlink( tir_mod: IRModule The input TIR IRModule to be linked together. - ext_libs: List[tvm_tirx.runtime.Module] + ext_libs: List[tvm_ffi.Module] List of compiled external modules. params: Optional[Dict[str, list]] diff --git a/python/tvm/rpc/client.py b/python/tvm/rpc/client.py index 9292f9de9459..434eca77bb12 100644 --- a/python/tvm/rpc/client.py +++ b/python/tvm/rpc/client.py @@ -50,12 +50,12 @@ def system_lib(self): Returns ------- - module : runtime.Module + module : tvm_ffi.Module The system-wide library module. See Also -------- - tvm_tirx.runtime.system_lib + tvm_ffi.system_lib """ return self.get_function("ffi.SystemLib")() @@ -88,7 +88,7 @@ def device(self, dev_type, dev_id=0): dev: RemoteDevice The session-owned device. """ - return _ffi_api.RemoteDevice(self._sess, tvm_tirx.runtime.device(dev_type, dev_id)) + return _ffi_api.RemoteDevice(self._sess, tvm_ffi.device(dev_type, dev_id)) def close(self): """Close this session; existing remote values reject further operations.""" diff --git a/python/tvm/rpc/server.py b/python/tvm/rpc/server.py index 4d1de81167b4..6919fd60411d 100644 --- a/python/tvm/rpc/server.py +++ b/python/tvm/rpc/server.py @@ -489,7 +489,7 @@ def server_init_callback(): # must import mypackage here import mypackage - tvm_tirx.register_global_func("function", mypackage.func) + tvm_ffi.register_global_func("function", mypackage.func) server = rpc.Server(host, server_init_callback=server_init_callback) """ diff --git a/python/tvm/rpc/testing.py b/python/tvm/rpc/testing.py index 00ca910aa4c5..8da5641a3761 100644 --- a/python/tvm/rpc/testing.py +++ b/python/tvm/rpc/testing.py @@ -19,43 +19,44 @@ """Testing functions for the RPC server.""" import numpy as np +import tvm_ffi import tvm_tirx # RPC test functions to be registered for unit-tests purposes -@tvm_tirx.register_global_func("rpc.test.addone") +@tvm_ffi.register_global_func("rpc.test.addone") def _addone(x): return x + 1 -@tvm_tirx.register_global_func("rpc.test.strcat") +@tvm_ffi.register_global_func("rpc.test.strcat") def _strcat(name, x): return f"{name}:{x}" -@tvm_tirx.register_global_func("rpc.test.except") +@tvm_ffi.register_global_func("rpc.test.except") def _remotethrow(name): raise ValueError(f"{name}") -@tvm_tirx.register_global_func("rpc.test.runtime_str_concat") +@tvm_ffi.register_global_func("rpc.test.runtime_str_concat") def _strcat(x, y): return x + y -@tvm_tirx.register_global_func("rpc.test.remote_tensor_func") +@tvm_ffi.register_global_func("rpc.test.remote_tensor_func") def _remote_tensor_func(y): x = np.ones((3, 4)) np.testing.assert_equal(y.numpy(), x) -@tvm_tirx.register_global_func("rpc.test.add_to_lhs") +@tvm_ffi.register_global_func("rpc.test.add_to_lhs") def _add_to_lhs(x): return lambda y: x + y -@tvm_tirx.register_global_func("rpc.test.remote_return_nd") +@tvm_ffi.register_global_func("rpc.test.remote_return_nd") def _my_module(name): # Use closure to check the ref counter correctness nd = tvm_tirx.runtime.tensor(np.zeros(10).astype("float32")) diff --git a/python/tvm/runtime/disco/session.py b/python/tvm/runtime/disco/session.py index a658ebd1653b..49f2d076a040 100644 --- a/python/tvm/runtime/disco/session.py +++ b/python/tvm/runtime/disco/session.py @@ -28,9 +28,9 @@ import numpy as np from tvm_ffi import Object, Shape, get_global_func, register_global_func, register_object +from tvm_tirx.runtime import Device from tvm_tirx.runtime._tensor import Tensor from tvm_tirx.runtime._tensor import tensor as _as_Tensor -from tvm_tirx.runtime.device import Device from . import _ffi_api, process_pool # pylint: disable=unused-import diff --git a/python/tvm/runtime/vm/vm.py b/python/tvm/runtime/vm/vm.py index 8b9443df76f5..86d3cf9d1f0b 100644 --- a/python/tvm/runtime/vm/vm.py +++ b/python/tvm/runtime/vm/vm.py @@ -24,10 +24,11 @@ from typing import TYPE_CHECKING, Any import numpy as np # type: ignore +import tvm_ffi from tvm_ffi import Function, register_global_func import tvm_tirx -from tvm_tirx.runtime import Device, Object +from tvm_tirx.runtime import Device from ..benchmark import BenchmarkResult, time_evaluator @@ -46,7 +47,7 @@ class VirtualMachine: def __init__( self, - rt_mod: "tvm_tirx.runtime.Module | Executable", + rt_mod: "tvm_ffi.Module | Executable", device: Device | list[Device], ) -> None: """ @@ -54,19 +55,19 @@ def __init__( Parameters ---------- - rt_mod: Union[tvm_tirx.runtime.Module, tvm_tirx.codegen.Executable] + rt_mod: Union[tvm_ffi.Module, tvm_tirx.codegen.Executable] Runtime module exported by the result of build. device : Union[Device, List[Device]] The device to deploy the module. """ - if not isinstance(rt_mod, tvm_tirx.runtime.Module): + if not isinstance(rt_mod, tvm_ffi.Module): from tvm_tirx.codegen import Executable # pylint: disable=import-outside-toplevel if isinstance(rt_mod, Executable): rt_mod = rt_mod.jit() else: - raise ValueError("Expect the rt_mod to be an runtime.Module") + raise ValueError("Expect the rt_mod to be an tvm_ffi.Module") self.module = rt_mod["vm_load_executable"]() self._invoke_closure = self.module["invoke_closure"] @@ -101,12 +102,12 @@ def _setup_device(self, dev: Device | list[Device]) -> None: def __getitem__(self, key: str) -> Function: return self.module[key] - def invoke_closure(self, closure: Object, *args: Any) -> Object: + def invoke_closure(self, closure: tvm_ffi.Object, *args: Any) -> tvm_ffi.Object: """Invoke a closure. Parameters ---------- - closure : Object + closure : tvm_ffi.Object The VMClosure Object. args : list[tvm_tirx.runtime.Tensor] or list[np.ndarray] @@ -114,7 +115,7 @@ def invoke_closure(self, closure: Object, *args: Any) -> Object: Returns ------- - result : Object + result : tvm_ffi.Object The output. """ return self._invoke_closure(closure, *args) @@ -178,7 +179,7 @@ def _gettype(arg): else: return "float32" - if isinstance(arg, Object): + if isinstance(arg, tvm_ffi.Object): cargs.append(arg) elif isinstance(arg, np.ndarray): nd_arr = tvm_tirx.runtime.tensor(arg, device=tvm_tirx.cpu(0)) @@ -269,7 +270,7 @@ def invoke_stateful(self, func_name: str) -> None: """ self._invoke_stateful(func_name) - def get_outputs(self, func_name: str) -> tvm_tirx.Object | tuple[Any]: + def get_outputs(self, func_name: str) -> tvm_ffi.Object | tuple[Any]: """ Get the value output by the function by the given name after a call of `invoke_stateful`. @@ -283,7 +284,7 @@ def get_outputs(self, func_name: str) -> tvm_tirx.Object | tuple[Any]: Returns ------- - ret: Union[tvm_tirx.Object, Tuple[Any]] + ret: Union[tvm_ffi.Object, Tuple[Any]] The result of the earlier call to the function via `invoke_stateful`. If the result is a tuple, it returns a list of the fields. The fields are potentially also tuples, so these can be arbitrily nested. diff --git a/python/tvm/s_tir/_tensor_intrin.py b/python/tvm/s_tir/_tensor_intrin.py index cccdd07097b5..99c9361db5fe 100644 --- a/python/tvm/s_tir/_tensor_intrin.py +++ b/python/tvm/s_tir/_tensor_intrin.py @@ -20,14 +20,13 @@ import tvm_ffi -from tvm_tirx.runtime import Object from tvm_tirx.tirx.function import Function from . import _ffi_api @tvm_ffi.register_object("s_tir.TensorIntrin") -class TensorIntrin(Object): +class TensorIntrin(tvm_ffi.Object): """A tensor intrinsic. Parameters diff --git a/python/tvm/s_tir/data_layout.py b/python/tvm/s_tir/data_layout.py index 033707ccb7c0..c5e4d51939cd 100644 --- a/python/tvm/s_tir/data_layout.py +++ b/python/tvm/s_tir/data_layout.py @@ -18,13 +18,11 @@ import tvm_ffi -from tvm_tirx.runtime import Object - from . import _ffi_api @tvm_ffi.register_object("s_tir.SLayout") -class SLayout(Object): +class SLayout(tvm_ffi.Object): """SLayout is composed of upper cases, lower cases and numbers, where upper case indicates a primal axis and the corresponding lower case with factor size indicates the subordinate axis. @@ -83,7 +81,7 @@ def factor_of(self, axis): @tvm_ffi.register_object("s_tir.SBijectiveLayout") -class SBijectiveLayout(Object): +class SBijectiveLayout(tvm_ffi.Object): """Bijective mapping for two layouts (src-layout and dst-layout). It provides shape and index conversion between each other. diff --git a/python/tvm/s_tir/dlight/base/utils.py b/python/tvm/s_tir/dlight/base/utils.py index d5b7114efb58..b10b6e865206 100644 --- a/python/tvm/s_tir/dlight/base/utils.py +++ b/python/tvm/s_tir/dlight/base/utils.py @@ -17,18 +17,19 @@ # pylint: disable=missing-docstring """Utility methods for generic GPU.""" +import tvm_ffi + import tvm_tirx from tvm import s_tir -from tvm_tirx import DataType from tvm_tirx.ir import PrimType from tvm_tirx.target import Target -def get_bytes(dtype: DataType | PrimType | str) -> int: +def get_bytes(dtype: tvm_ffi.dtype | PrimType | str) -> int: if isinstance(dtype, PrimType): dtype = dtype.dtype if isinstance(dtype, str): - dtype = DataType(dtype) + dtype = tvm_ffi.dtype(dtype) return dtype.itemsize diff --git a/python/tvm/s_tir/meta_schedule/arg_info.py b/python/tvm/s_tir/meta_schedule/arg_info.py index b711ed38185c..265cd00b8ee1 100644 --- a/python/tvm/s_tir/meta_schedule/arg_info.py +++ b/python/tvm/s_tir/meta_schedule/arg_info.py @@ -18,10 +18,10 @@ from typing import Any +import tvm_ffi from tvm_ffi import Shape, register_object from tvm_tirx.ir import IRModule -from tvm_tirx.runtime import DataType, Object from tvm_tirx.tirx import Function from . import _ffi_api @@ -29,7 +29,7 @@ @register_object("s_tir.meta_schedule.ArgInfo") -class ArgInfo(Object): +class ArgInfo(tvm_ffi.Object): """Argument information""" def as_json(self) -> Any: @@ -93,25 +93,25 @@ class TensorInfo(ArgInfo): Parameters ---------- - dtype : DataType + dtype : tvm_ffi.dtype The data type of the tensor. shape : Shape The shape of the tensor. """ - dtype: DataType + dtype: tvm_ffi.dtype shape: Shape def __init__( self, - dtype: DataType, + dtype: tvm_ffi.dtype, shape: Shape | list[int], ) -> None: """Constructor Parameters ---------- - dtype : DataType + dtype : tvm_ffi.dtype The data type of the tensor. shape : Shape The shape of the tensor. diff --git a/python/tvm/s_tir/meta_schedule/builder/__init__.py b/python/tvm/s_tir/meta_schedule/builder/__init__.py index c7ba88719947..3dedc874a18e 100644 --- a/python/tvm/s_tir/meta_schedule/builder/__init__.py +++ b/python/tvm/s_tir/meta_schedule/builder/__init__.py @@ -17,7 +17,7 @@ # under the License. """ The tvm.s_tir.meta_schedule.builder package. -Meta Schedule builders that translate IRModule to runtime.Module, +Meta Schedule builders that translate IRModule to tvm_ffi.Module, and then export """ diff --git a/python/tvm/s_tir/meta_schedule/builder/builder.py b/python/tvm/s_tir/meta_schedule/builder/builder.py index 70cfdd5754b0..ba46a5c7fa0e 100644 --- a/python/tvm/s_tir/meta_schedule/builder/builder.py +++ b/python/tvm/s_tir/meta_schedule/builder/builder.py @@ -15,11 +15,13 @@ # specific language governing permissions and limitations # under the License. # ruff: noqa: RUF012 -"""Meta Schedule builders that translate IRModule to runtime.Module, and then export""" +"""Meta Schedule builders that translate IRModule to tvm_ffi.Module, and then export""" from collections.abc import Callable from typing import Union +import tvm_ffi + # isort: off from typing import Literal @@ -27,14 +29,14 @@ from tvm_ffi import register_object from tvm_tirx.ir import IRModule -from tvm_tirx.runtime import Object, Tensor +from tvm_tirx.runtime import Tensor from tvm_tirx.target import Target from .. import _ffi_api @register_object("s_tir.meta_schedule.BuilderInput") -class BuilderInput(Object): +class BuilderInput(tvm_ffi.Object): """The builder's input. Parameters @@ -77,7 +79,7 @@ def __init__( @register_object("s_tir.meta_schedule.BuilderResult") -class BuilderResult(Object): +class BuilderResult(tvm_ffi.Object): """The builder's result. Parameters @@ -113,7 +115,7 @@ def __init__( @register_object("s_tir.meta_schedule.Builder") -class Builder(Object): +class Builder(tvm_ffi.Object): """The abstract builder interface.""" BuilderType = Union["Builder", Literal["local"]] diff --git a/python/tvm/s_tir/meta_schedule/builder/local_builder.py b/python/tvm/s_tir/meta_schedule/builder/local_builder.py index d02da1a47f2a..054f5eb0acfb 100644 --- a/python/tvm/s_tir/meta_schedule/builder/local_builder.py +++ b/python/tvm/s_tir/meta_schedule/builder/local_builder.py @@ -22,13 +22,14 @@ from collections.abc import Callable from typing import Optional, Union +import tvm_ffi from tvm_ffi import register_global_func from tvm.runtime.vm import load_param_dict, save_param_dict from tvm.s_tir.utils import derived_object from tvm_tirx.codegen import Executable from tvm_tirx.ir import IRModule -from tvm_tirx.runtime import Module, Tensor +from tvm_tirx.runtime import Tensor from tvm_tirx.support.popen_pool import MapResult, PopenPoolExecutor, StatusKind from tvm_tirx.target import Target @@ -40,9 +41,9 @@ T_BUILD = Callable[ # pylint: disable=invalid-name - [IRModule, Target, dict[str, Tensor] | None], Module + [IRModule, Target, dict[str, Tensor] | None], tvm_ffi.Module ] -T_EXPORT = Callable[[Module], str] # pylint: disable=invalid-name +T_EXPORT = Callable[[tvm_ffi.Module], str] # pylint: disable=invalid-name def _serialize_params(params: dict[str, Tensor] | None) -> bytearray | None: @@ -89,7 +90,7 @@ def default_build( mod: IRModule, target: Target, params: Optional[Dict[str, Tensor]] - ) -> Module: + ) -> tvm_ffi.Module: ... T_EXPORT : typing._GenericAlias @@ -97,7 +98,7 @@ def default_build( .. code-block:: python - def default_export(mod: Module) -> str: + def default_export(mod: tvm_ffi.Module) -> str: ... Note @@ -235,14 +236,16 @@ def _worker_func( default_export, ) # Step 1. Build the IRModule - rt_mod: Module = f_build(mod, target, _deserialize_params(params)) + rt_mod: tvm_ffi.Module = f_build(mod, target, _deserialize_params(params)) # Step 2. Export the Module artifact_path: str = f_export(rt_mod) return artifact_path @register_global_func("s_tir.meta_schedule.builder.default_build") -def default_build(mod: IRModule, target: Target, _params: dict[str, Tensor] | None) -> Module: +def default_build( + mod: IRModule, target: Target, _params: dict[str, Tensor] | None +) -> tvm_ffi.Module: """Default build function. Parameters @@ -256,7 +259,7 @@ def default_build(mod: IRModule, target: Target, _params: dict[str, Tensor] | No Returns ------- - rt_mod : Module + rt_mod : tvm_ffi.Module The built Module. """ # pylint: disable=import-outside-toplevel @@ -270,12 +273,12 @@ def default_build(mod: IRModule, target: Target, _params: dict[str, Tensor] | No @register_global_func("s_tir.meta_schedule.builder.default_export") -def default_export(mod: Module) -> str: +def default_export(mod: tvm_ffi.Module) -> str: """Default export function. Parameters ---------- - mod : Module + mod : tvm_ffi.Module The Module to be exported. Returns diff --git a/python/tvm/s_tir/meta_schedule/cost_model/cost_model.py b/python/tvm/s_tir/meta_schedule/cost_model/cost_model.py index c749a622ead0..2adb23ef8668 100644 --- a/python/tvm/s_tir/meta_schedule/cost_model/cost_model.py +++ b/python/tvm/s_tir/meta_schedule/cost_model/cost_model.py @@ -21,6 +21,8 @@ from collections.abc import Callable from typing import Union +import tvm_ffi + # isort: off from typing import Literal @@ -29,8 +31,6 @@ import numpy as np # type: ignore from tvm_ffi import register_object -from tvm_tirx.runtime import Object - from .. import _ffi_api from ..runner import RunnerResult from ..search_strategy import MeasureCandidate @@ -38,7 +38,7 @@ @register_object("s_tir.meta_schedule.CostModel") -class CostModel(Object): +class CostModel(tvm_ffi.Object): """Cost model.""" CostModelType = Union["CostModel", Literal["xgb", "mlp", "random"]] diff --git a/python/tvm/s_tir/meta_schedule/database/database.py b/python/tvm/s_tir/meta_schedule/database/database.py index d1420f22a6ca..60baf8a82c5b 100644 --- a/python/tvm/s_tir/meta_schedule/database/database.py +++ b/python/tvm/s_tir/meta_schedule/database/database.py @@ -20,6 +20,8 @@ from collections.abc import Callable from typing import Any, Optional, Union +import tvm_ffi + # isort: off from typing import Literal @@ -29,7 +31,6 @@ from tvm.s_tir.schedule import Schedule, Trace from tvm_tirx.ir.module import IRModule -from tvm_tirx.runtime import Object from tvm_tirx.target import Target from .. import _ffi_api @@ -38,7 +39,7 @@ @register_object("s_tir.meta_schedule.Workload") -class Workload(Object): +class Workload(tvm_ffi.Object): """A workload, i.e. an IRModule and its structural hash. Parameters @@ -84,7 +85,7 @@ def from_json(json_obj: Any) -> "Workload": @register_object("s_tir.meta_schedule.TuningRecord") -class TuningRecord(Object): +class TuningRecord(tvm_ffi.Object): """The class of tuning records. Parameters @@ -165,7 +166,7 @@ def from_json(json_obj: Any, workload: Workload) -> "TuningRecord": @register_object("s_tir.meta_schedule.Database") -class Database(Object): +class Database(tvm_ffi.Object): """The abstract database interface.""" DatabaseType = Union["Database", Literal["json", "memory"]] diff --git a/python/tvm/s_tir/meta_schedule/extracted_task.py b/python/tvm/s_tir/meta_schedule/extracted_task.py index fe44bbca478d..e73211e8cb78 100644 --- a/python/tvm/s_tir/meta_schedule/extracted_task.py +++ b/python/tvm/s_tir/meta_schedule/extracted_task.py @@ -16,17 +16,17 @@ # under the License. """Extracted tasks from high-level IR.""" +import tvm_ffi from tvm_ffi import register_object from tvm_tirx.ir import IRModule -from tvm_tirx.runtime import Object from tvm_tirx.target import Target from . import _ffi_api @register_object("s_tir.meta_schedule.ExtractedTask") -class ExtractedTask(Object): +class ExtractedTask(tvm_ffi.Object): """A tuning task extracted from the high-level IR Parameters diff --git a/python/tvm/s_tir/meta_schedule/feature_extractor/feature_extractor.py b/python/tvm/s_tir/meta_schedule/feature_extractor/feature_extractor.py index 942b6b18b56a..2d4750b1515b 100644 --- a/python/tvm/s_tir/meta_schedule/feature_extractor/feature_extractor.py +++ b/python/tvm/s_tir/meta_schedule/feature_extractor/feature_extractor.py @@ -20,6 +20,8 @@ from collections.abc import Callable from typing import Union +import tvm_ffi + # isort: off from typing import Literal @@ -27,7 +29,6 @@ from tvm_ffi import register_object -from tvm_tirx.runtime import Object from tvm_tirx.runtime._tensor import Tensor from .. import _ffi_api @@ -36,7 +37,7 @@ @register_object("s_tir.meta_schedule.FeatureExtractor") -class FeatureExtractor(Object): +class FeatureExtractor(tvm_ffi.Object): """Extractor for features from measure candidates for use in cost model.""" FeatureExtractorType = Union[Literal["per-store-feature"], "FeatureExtractor"] diff --git a/python/tvm/s_tir/meta_schedule/measure_callback/measure_callback.py b/python/tvm/s_tir/meta_schedule/measure_callback/measure_callback.py index ebabe6ae82b2..2e23152e6fb7 100644 --- a/python/tvm/s_tir/meta_schedule/measure_callback/measure_callback.py +++ b/python/tvm/s_tir/meta_schedule/measure_callback/measure_callback.py @@ -20,6 +20,8 @@ from collections.abc import Callable from typing import TYPE_CHECKING, Union +import tvm_ffi + # isort: off from typing import Literal @@ -27,8 +29,6 @@ from tvm_ffi import register_object -from tvm_tirx.runtime import Object - from .. import _ffi_api from ..builder import BuilderResult from ..runner import RunnerResult @@ -39,7 +39,7 @@ @register_object("s_tir.meta_schedule.MeasureCallback") -class MeasureCallback(Object): +class MeasureCallback(tvm_ffi.Object): """Rules to apply after measure results is available.""" CallbackListType = Union[list["MeasureCallback"], "MeasureCallback", Literal["default"]] diff --git a/python/tvm/s_tir/meta_schedule/mutator/mutator.py b/python/tvm/s_tir/meta_schedule/mutator/mutator.py index e506522d2d4e..c312e427fcb9 100644 --- a/python/tvm/s_tir/meta_schedule/mutator/mutator.py +++ b/python/tvm/s_tir/meta_schedule/mutator/mutator.py @@ -20,6 +20,8 @@ from collections.abc import Callable from typing import TYPE_CHECKING +import tvm_ffi + # isort: off from typing import Literal @@ -28,7 +30,6 @@ from tvm_ffi import register_object from tvm.s_tir.schedule import Trace -from tvm_tirx.runtime import Object from .. import _ffi_api @@ -36,7 +37,7 @@ from ..tune_context import TuneContext -class Mutator(Object): +class Mutator(tvm_ffi.Object): """Mutator is designed to mutate the trace to explore the design space.""" def _initialize_with_tune_context(self, context: "TuneContext") -> None: diff --git a/python/tvm/s_tir/meta_schedule/postproc/postproc.py b/python/tvm/s_tir/meta_schedule/postproc/postproc.py index f388f262b127..4e3bfea6068b 100644 --- a/python/tvm/s_tir/meta_schedule/postproc/postproc.py +++ b/python/tvm/s_tir/meta_schedule/postproc/postproc.py @@ -20,6 +20,8 @@ from collections.abc import Callable from typing import TYPE_CHECKING +import tvm_ffi + # isort: off from typing import Literal @@ -28,7 +30,6 @@ from tvm_ffi import register_object from tvm.s_tir.schedule import Schedule -from tvm_tirx.runtime import Object from .. import _ffi_api @@ -37,7 +38,7 @@ @register_object("s_tir.meta_schedule.Postproc") -class Postproc(Object): +class Postproc(tvm_ffi.Object): """Rules to apply a postprocessor to a schedule.""" def _initialize_with_tune_context(self, context: "TuneContext") -> None: diff --git a/python/tvm/s_tir/meta_schedule/profiler.py b/python/tvm/s_tir/meta_schedule/profiler.py index 20691eca6100..7ba3cbec10b2 100644 --- a/python/tvm/s_tir/meta_schedule/profiler.py +++ b/python/tvm/s_tir/meta_schedule/profiler.py @@ -20,15 +20,14 @@ from contextlib import contextmanager from typing import Optional +import tvm_ffi from tvm_ffi import register_object -from tvm_tirx.runtime import Object - from . import _ffi_api @register_object("s_tir.meta_schedule.Profiler") -class Profiler(Object): +class Profiler(tvm_ffi.Object): """Tuning time profiler.""" def __init__(self) -> None: diff --git a/python/tvm/s_tir/meta_schedule/runner/local_runner.py b/python/tvm/s_tir/meta_schedule/runner/local_runner.py index 1e5be8c74145..ca95eef494c1 100644 --- a/python/tvm/s_tir/meta_schedule/runner/local_runner.py +++ b/python/tvm/s_tir/meta_schedule/runner/local_runner.py @@ -21,9 +21,11 @@ from collections.abc import Callable from contextlib import contextmanager +import tvm_ffi + import tvm_tirx from tvm.s_tir.utils import derived_object -from tvm_tirx.runtime import Device, Module +from tvm_tirx.runtime import Device from tvm_tirx.support import cc, tar, utils from tvm_tirx.support.popen_pool import PopenPoolExecutor @@ -52,7 +54,7 @@ ] T_RUN_EVALUATOR = Callable[ # pylint: disable=invalid-name [ - Module, # The Module opened on the remote + tvm_ffi.Module, # The Module opened on the remote Device, # The device on the remote EvaluatorConfig, # The evaluator configuration list[T_ARGUMENT_LIST], # A list of argument lists @@ -156,7 +158,7 @@ def resource_handler(): rt_mod = tvm_tirx.runtime.load_module(artifact_path) # Step 2: Allocate input arguments with Profiler.timeit("LocalRunner/alloc_argument"): - device = tvm_tirx.runtime.device(device_type, 0) + device = tvm_ffi.device(device_type, 0) repeated_args: list[T_ARGUMENT_LIST] = f_alloc_argument( device, args_info, @@ -214,7 +216,7 @@ def default_alloc_argument( .. code-block:: python def default_run_evaluator( - rt_mod: Module, + rt_mod: tvm_ffi.Module, device: Device, evaluator_config: EvaluatorConfig, repeated_args: List[T_ARGUMENT_LIST], @@ -369,7 +371,7 @@ def default_alloc_argument( def default_run_evaluator( - rt_mod: Module, + rt_mod: tvm_ffi.Module, device: Device, evaluator_config: EvaluatorConfig, repeated_args: list[T_ARGUMENT_LIST], @@ -378,7 +380,7 @@ def default_run_evaluator( Parameters ---------- - rt_mod: Module + rt_mod: tvm_ffi.Module The runtime module device: Device The device to run the evaluator @@ -400,7 +402,7 @@ def default_cleanup() -> None: pass # pylint: disable=unnecessary-pass -@tvm_tirx.register_global_func("s_tir.meta_schedule.runner.get_local_runner") +@tvm_ffi.register_global_func("s_tir.meta_schedule.runner.get_local_runner") def get_local_builder() -> LocalRunner: """Get the local Runner. diff --git a/python/tvm/s_tir/meta_schedule/runner/rpc_runner.py b/python/tvm/s_tir/meta_schedule/runner/rpc_runner.py index 4795705aca13..20515f1cd81c 100644 --- a/python/tvm/s_tir/meta_schedule/runner/rpc_runner.py +++ b/python/tvm/s_tir/meta_schedule/runner/rpc_runner.py @@ -21,9 +21,11 @@ from collections.abc import Callable from contextlib import contextmanager +import tvm_ffi + from tvm.rpc import RPCSession from tvm.s_tir.utils import derived_object -from tvm_tirx.runtime import Device, Module +from tvm_tirx.runtime import Device from tvm_tirx.support.popen_pool import PopenPoolExecutor from ..logging import get_logger @@ -54,7 +56,7 @@ str, # local path to the artifact str, # remote path to the artifact ], - Module, # the Module opened on the remote + tvm_ffi.Module, # the Module opened on the remote ] T_ALLOC_ARGUMENT = Callable[ # pylint: disable=invalid-name [ @@ -68,7 +70,7 @@ T_RUN_EVALUATOR = Callable[ # pylint: disable=invalid-name [ RPCSession, # The RPC Session - Module, # The Module opened on the remote + tvm_ffi.Module, # The Module opened on the remote Device, # The device on the remote EvaluatorConfig, # The evaluator configuration list[T_ARGUMENT_LIST], # A list of argument lists @@ -178,7 +180,7 @@ def default_upload_module( session: RPCSession, local_path: str, remote_path: str, - ) -> Module: + ) -> tvm_ffi.Module: ... T_ALLOC_ARGUMENT : typing._GenericAlias @@ -201,7 +203,7 @@ def default_alloc_argument( def default_run_evaluator( session: RPCSession, - rt_mod: Module, + rt_mod: tvm_ffi.Module, device: Device, evaluator_config: EvaluatorConfig, repeated_args: List[T_ARGUMENT_LIST], @@ -390,7 +392,7 @@ def resource_handler(): with Profiler.timeit("RPCRunner/upload_module"): _, remote_path = osp.split(artifact_path) local_path: str = artifact_path - rt_mod: Module = f_upload_module(session, local_path, remote_path) + rt_mod: tvm_ffi.Module = f_upload_module(session, local_path, remote_path) # Step 3: Allocate input arguments with Profiler.timeit("RPCRunner/alloc_argument"): repeated_args: list[T_ARGUMENT_LIST] = f_alloc_argument( @@ -431,7 +433,7 @@ def default_upload_module( session: RPCSession, local_path: str, remote_path: str, -) -> Module: +) -> tvm_ffi.Module: """Default function to upload the module Parameters @@ -445,11 +447,11 @@ def default_upload_module( Returns ------- - rt_mod : Module + rt_mod : tvm_ffi.Module The runtime module """ session.upload(local_path, remote_path) - rt_mod: Module = session.load_module(remote_path) + rt_mod: tvm_ffi.Module = session.load_module(remote_path) return rt_mod @@ -488,7 +490,7 @@ def default_alloc_argument( def default_run_evaluator( session: RPCSession, # pylint: disable=unused-argument - rt_mod: Module, + rt_mod: tvm_ffi.Module, device: Device, evaluator_config: EvaluatorConfig, repeated_args: list[T_ARGUMENT_LIST], @@ -499,7 +501,7 @@ def default_run_evaluator( ---------- session: RPCSession The session to run the evaluator - rt_mod: Module + rt_mod: tvm_ffi.Module The runtime module device: Device The device to run the evaluator diff --git a/python/tvm/s_tir/meta_schedule/runner/runner.py b/python/tvm/s_tir/meta_schedule/runner/runner.py index bfaeb7c445e1..422af9f40875 100644 --- a/python/tvm/s_tir/meta_schedule/runner/runner.py +++ b/python/tvm/s_tir/meta_schedule/runner/runner.py @@ -20,6 +20,8 @@ from collections.abc import Callable from typing import Union +import tvm_ffi + # isort: off from typing import Literal @@ -27,14 +29,12 @@ from tvm_ffi import register_object -from tvm_tirx.runtime import Object - from .. import _ffi_api from ..arg_info import ArgInfo @register_object("s_tir.meta_schedule.RunnerInput") -class RunnerInput(Object): +class RunnerInput(tvm_ffi.Object): """The runner's input Parameters @@ -77,7 +77,7 @@ def __init__( @register_object("s_tir.meta_schedule.RunnerResult") -class RunnerResult(Object): +class RunnerResult(tvm_ffi.Object): """The runner's result Parameters @@ -113,7 +113,7 @@ def __init__( @register_object("s_tir.meta_schedule.RunnerFuture") -class RunnerFuture(Object): +class RunnerFuture(tvm_ffi.Object): """ A class to fetch asynchronous runner's output. This is NOT the user facing class for function overloading inheritance. @@ -169,7 +169,7 @@ def result(self) -> RunnerResult: @register_object("s_tir.meta_schedule.Runner") -class Runner(Object): +class Runner(tvm_ffi.Object): """The abstract runner interface""" RunnerType = Union["Runner", Literal["local", "rpc"]] diff --git a/python/tvm/s_tir/meta_schedule/runner/utils.py b/python/tvm/s_tir/meta_schedule/runner/utils.py index b4cc4b745599..02df6010962e 100644 --- a/python/tvm/s_tir/meta_schedule/runner/utils.py +++ b/python/tvm/s_tir/meta_schedule/runner/utils.py @@ -20,9 +20,11 @@ from collections.abc import Callable from typing import Any +import tvm_ffi + import tvm_tirx.runtime from tvm.runtime import time_evaluator -from tvm_tirx.runtime import Device, Module +from tvm_tirx.runtime import Device from .config import EvaluatorConfig @@ -88,7 +90,7 @@ def alloc_fail(*arg_info) -> None: def run_evaluator_common( - rt_mod: Module, + rt_mod: tvm_ffi.Module, device: Device, evaluator_config: EvaluatorConfig, repeated_args: list[T_ARGUMENT_LIST], @@ -97,7 +99,7 @@ def run_evaluator_common( Parameters ---------- - rt_mod: Module + rt_mod: tvm_ffi.Module The runtime module device: Device The device to run the evaluator diff --git a/python/tvm/s_tir/meta_schedule/schedule/cuda/layout_transform.py b/python/tvm/s_tir/meta_schedule/schedule/cuda/layout_transform.py index 64aa0acffeeb..4cb3e29a52c2 100644 --- a/python/tvm/s_tir/meta_schedule/schedule/cuda/layout_transform.py +++ b/python/tvm/s_tir/meta_schedule/schedule/cuda/layout_transform.py @@ -20,6 +20,8 @@ import math from collections import deque +import tvm_ffi + import tvm_tirx from tvm.s_tir import meta_schedule from tvm.s_tir.schedule import ExprRV, LoopRV, SBlockRV, Schedule @@ -499,7 +501,7 @@ def get_max_tile_size() -> int: return max_tile_size -@tvm_tirx.register_global_func("s_tir.meta_schedule.cuda.layout_transform") +@tvm_ffi.register_global_func("s_tir.meta_schedule.cuda.layout_transform") def cuda_layout_transform_schedule_rule( sch: Schedule, block: SBlockRV, testing_tile_sizes: list[int] | None = None ) -> list[Schedule]: diff --git a/python/tvm/s_tir/meta_schedule/schedule_rule/schedule_rule.py b/python/tvm/s_tir/meta_schedule/schedule_rule/schedule_rule.py index f6543cacb89f..07aa183080bc 100644 --- a/python/tvm/s_tir/meta_schedule/schedule_rule/schedule_rule.py +++ b/python/tvm/s_tir/meta_schedule/schedule_rule/schedule_rule.py @@ -23,6 +23,8 @@ from collections.abc import Callable from typing import TYPE_CHECKING +import tvm_ffi + # isort: off from typing import Literal @@ -31,7 +33,6 @@ from tvm_ffi import register_object from tvm.s_tir.schedule import SBlockRV, Schedule -from tvm_tirx.runtime import Object from .. import _ffi_api @@ -40,7 +41,7 @@ @register_object("s_tir.meta_schedule.ScheduleRule") -class ScheduleRule(Object): +class ScheduleRule(tvm_ffi.Object): """Rules to modify a block in a schedule.""" def _initialize_with_tune_context(self, context: "TuneContext") -> None: diff --git a/python/tvm/s_tir/meta_schedule/search_strategy/search_strategy.py b/python/tvm/s_tir/meta_schedule/search_strategy/search_strategy.py index 9541d2120165..5b0f3f4bb0ae 100644 --- a/python/tvm/s_tir/meta_schedule/search_strategy/search_strategy.py +++ b/python/tvm/s_tir/meta_schedule/search_strategy/search_strategy.py @@ -23,6 +23,8 @@ from collections.abc import Callable from typing import TYPE_CHECKING, Optional, Union +import tvm_ffi + # isort: off from typing import Literal @@ -30,7 +32,6 @@ from tvm_ffi import register_object from tvm.s_tir.schedule import Schedule -from tvm_tirx.runtime import Object from .. import _ffi_api from ..arg_info import ArgInfo @@ -43,7 +44,7 @@ @register_object("s_tir.meta_schedule.MeasureCandidate") -class MeasureCandidate(Object): +class MeasureCandidate(tvm_ffi.Object): """Measure candidate class. Parameters @@ -79,7 +80,7 @@ def __init__( @register_object("s_tir.meta_schedule.SearchStrategy") -class SearchStrategy(Object): +class SearchStrategy(tvm_ffi.Object): """Search strategy is the class that generates the measure candidates.""" SearchStrategyType = Union[ diff --git a/python/tvm/s_tir/meta_schedule/space_generator/space_generator.py b/python/tvm/s_tir/meta_schedule/space_generator/space_generator.py index 20afd791a2d1..daf1d0331655 100644 --- a/python/tvm/s_tir/meta_schedule/space_generator/space_generator.py +++ b/python/tvm/s_tir/meta_schedule/space_generator/space_generator.py @@ -23,6 +23,8 @@ from collections.abc import Callable from typing import TYPE_CHECKING, Union +import tvm_ffi + # isort: off from typing import Literal @@ -31,7 +33,6 @@ from tvm.s_tir.schedule import Schedule from tvm_tirx.ir import IRModule -from tvm_tirx.runtime import Object from .. import _ffi_api @@ -43,7 +44,7 @@ @register_object("s_tir.meta_schedule.SpaceGenerator") -class SpaceGenerator(Object): +class SpaceGenerator(tvm_ffi.Object): """The abstract design space generator interface.""" ScheduleFnType = ( diff --git a/python/tvm/s_tir/meta_schedule/task_scheduler/task_scheduler.py b/python/tvm/s_tir/meta_schedule/task_scheduler/task_scheduler.py index baad14d377c5..9d4cb9b7bd08 100644 --- a/python/tvm/s_tir/meta_schedule/task_scheduler/task_scheduler.py +++ b/python/tvm/s_tir/meta_schedule/task_scheduler/task_scheduler.py @@ -20,6 +20,8 @@ from collections.abc import Callable from typing import Union +import tvm_ffi + # isort: off from typing import Literal @@ -27,8 +29,6 @@ from tvm_ffi import register_object -from tvm_tirx.runtime import Object - from .. import _ffi_api from ..builder import Builder, BuilderResult from ..cost_model import CostModel @@ -43,7 +43,7 @@ @register_object("s_tir.meta_schedule.TaskRecord") -class TaskRecord(Object): +class TaskRecord(tvm_ffi.Object): """The running record of a task.""" ctx: TuneContext @@ -58,7 +58,7 @@ class TaskRecord(Object): @register_object("s_tir.meta_schedule.TaskScheduler") -class TaskScheduler(Object): +class TaskScheduler(tvm_ffi.Object): """The abstract task scheduler interface.""" tasks_: list[TaskRecord] diff --git a/python/tvm/s_tir/meta_schedule/testing/custom_builder_runner.py b/python/tvm/s_tir/meta_schedule/testing/custom_builder_runner.py index 62007c032623..1d8901609a5e 100644 --- a/python/tvm/s_tir/meta_schedule/testing/custom_builder_runner.py +++ b/python/tvm/s_tir/meta_schedule/testing/custom_builder_runner.py @@ -15,25 +15,25 @@ # specific language governing permissions and limitations # under the License. """Customized builder and runner methods""" -# pylint: disable=import-outside-toplevel +# pylint: disable=import-outside-toplevel from collections.abc import Callable import numpy as np # type: ignore +import tvm_ffi from tvm.s_tir.meta_schedule.runner import RPCConfig from tvm_tirx.codegen import Executable -from tvm_tirx.runtime import Module def run_module_via_rpc( rpc_config: RPCConfig, - lib: Module | Executable, + lib: tvm_ffi.Module | Executable, dev_type: str, args: dict[int, np.ndarray] | dict[str, np.ndarray], continuation: Callable, ): - """Execute a tvm_tirx.runtime.Module on RPC remote""" + """Execute a tvm_ffi.Module on RPC remote""" # pylint: disable=import-outside-toplevel import os import tempfile diff --git a/python/tvm/s_tir/meta_schedule/testing/tune_utils.py b/python/tvm/s_tir/meta_schedule/testing/tune_utils.py index 7112da246d22..cefbb329ccb8 100644 --- a/python/tvm/s_tir/meta_schedule/testing/tune_utils.py +++ b/python/tvm/s_tir/meta_schedule/testing/tune_utils.py @@ -19,6 +19,7 @@ from collections.abc import Callable import numpy as np # type: ignore +import tvm_ffi import tvm_tirx from tvm_tirx.runtime import Tensor @@ -81,7 +82,7 @@ def create_calculator(backend: str) -> Callable: """ def f_calculator( - rt_mod: tvm_tirx.runtime.Module, + rt_mod: tvm_ffi.Module, dev: tvm_tirx.runtime.Device, # pylint: disable=unused-argument input_data: dict[str, Tensor], ) -> list[Tensor]: @@ -89,9 +90,9 @@ def f_calculator( Parameters ---------- - rt_mod : tvm_tirx.runtime.Module + rt_mod : tvm_ffi.Module The runtime module. - dev : tvm_tirx.device + dev : tvm_ffi.device The device type to run workload. input_data : Dict[str, np.ndarray] The input data as a dictionary. diff --git a/python/tvm/s_tir/meta_schedule/testing/validate_database.py b/python/tvm/s_tir/meta_schedule/testing/validate_database.py index 029562b0b131..7e2fd08d6fee 100644 --- a/python/tvm/s_tir/meta_schedule/testing/validate_database.py +++ b/python/tvm/s_tir/meta_schedule/testing/validate_database.py @@ -449,7 +449,7 @@ def f_with_args_alloc_argument_common( ] def f_with_args_run_evaluator_common( - rt_mod: tvm_tirx.runtime.Module, + rt_mod: tvm_ffi.Module, device: tvm_tirx.runtime.Device, evaluator_config: ms.runner.EvaluatorConfig, repeated_args: list[ms.runner.rpc_runner.T_ARGUMENT_LIST], @@ -526,7 +526,7 @@ def f_with_args_alloc_argument_rpc( def f_with_args_run_evaluator_rpc( rpc_session: ms.runner.rpc_runner.RPCSession, # pylint: disable=unused-argument - rt_mod: tvm_tirx.runtime.Module, + rt_mod: tvm_ffi.Module, device: tvm_tirx.runtime.Device, evaluator_config: ms.runner.EvaluatorConfig, repeated_args: list[ms.runner.rpc_runner.T_ARGUMENT_LIST], diff --git a/python/tvm/s_tir/meta_schedule/tune_context.py b/python/tvm/s_tir/meta_schedule/tune_context.py index 2c1ad072ffa6..d8573351ca3e 100644 --- a/python/tvm/s_tir/meta_schedule/tune_context.py +++ b/python/tvm/s_tir/meta_schedule/tune_context.py @@ -19,6 +19,8 @@ from typing import TYPE_CHECKING, Optional, Union +import tvm_ffi + # isort: off from typing import Literal @@ -28,7 +30,6 @@ from tvm.s_tir import Schedule from tvm_tirx import IRModule -from tvm_tirx.runtime import Object from tvm_tirx.target import Target from tvm_tirx.tirx import Function @@ -61,7 +62,7 @@ def _normalize_mod(mod: Function | IRModule) -> IRModule: @register_object("s_tir.meta_schedule.TuneContext") -class TuneContext(Object): +class TuneContext(tvm_ffi.Object): """The tune context class is designed to contain all resources for a tuning task. Parameters diff --git a/python/tvm/s_tir/sblock_dependence_info.py b/python/tvm/s_tir/sblock_dependence_info.py index 21af539cb1fd..7dda75c711d2 100644 --- a/python/tvm/s_tir/sblock_dependence_info.py +++ b/python/tvm/s_tir/sblock_dependence_info.py @@ -17,11 +17,11 @@ """Define SBlockDependenceInfoObj that uses the SBlockScope and StmtSRef objects to store the block level dependences""" +import tvm_ffi from tvm_ffi import register_object from tvm.s_tir import SBlock from tvm_tirx.ir.module import IRModule -from tvm_tirx.runtime import Object from tvm_tirx.tirx import Function from . import _ffi_api @@ -29,7 +29,7 @@ @register_object("s_tir.SBlockDependenceInfo") -class SBlockDependenceInfo(Object): +class SBlockDependenceInfo(tvm_ffi.Object): """ SBlockDependenceInfo An object that helps build and query block level dependences using the 2 core objects diff --git a/python/tvm/s_tir/sblock_scope.py b/python/tvm/s_tir/sblock_scope.py index 13cc9de92fe6..96cd7b5db903 100644 --- a/python/tvm/s_tir/sblock_scope.py +++ b/python/tvm/s_tir/sblock_scope.py @@ -19,17 +19,17 @@ from enum import IntEnum from typing import Optional +import tvm_ffi from tvm_ffi import register_object from tvm.s_tir import SBlock from tvm_tirx.ir import For -from tvm_tirx.runtime import Object from . import _ffi_api @register_object("s_tir.StmtSRef") -class StmtSRef(Object): +class StmtSRef(tvm_ffi.Object): """An object that refers to schedulable elements in the TensorIR, aka "sref". Glossary @@ -90,7 +90,7 @@ class DepKind(IntEnum): @register_object("s_tir.Dependency") -class Dependency(Object): +class Dependency(tvm_ffi.Object): """A tuple (src, dst, kind) representing certain types of dependency. For example, (A, B, kRAW) means block B depends on block A, and the dependency kind is read-after-write, which means block B reads the result written by block A. @@ -111,7 +111,7 @@ class Dependency(Object): @register_object("s_tir.SBlockScope") -class SBlockScope(Object): +class SBlockScope(tvm_ffi.Object): """An object corresponds to each block sref in the sref tree, which tracks the producer-consumer dependency between blocks. diff --git a/python/tvm/s_tir/schedule/analysis.py b/python/tvm/s_tir/schedule/analysis.py index f2a0c3e72cdb..e99f6cf8f739 100644 --- a/python/tvm/s_tir/schedule/analysis.py +++ b/python/tvm/s_tir/schedule/analysis.py @@ -19,7 +19,6 @@ import tvm_ffi from tvm_tirx.ir import For, Var -from tvm_tirx.runtime import Object from tvm_tirx.tirx.expr import Expr from tvm_tirx.tirx.function import Function from tvm_tirx.tirx.index_map import IndexMap @@ -62,7 +61,7 @@ def suggest_index_map( @tvm_ffi.register_object("s_tir.schedule.TensorizeInfo") -class TensorizeInfo(Object): +class TensorizeInfo(tvm_ffi.Object): """Necessary information used for tensorization.""" @@ -90,7 +89,7 @@ def get_tensorize_loop_mapping( @tvm_ffi.register_object("s_tir.schedule.AutoTensorizeMappingInfo") -class AutoTensorizeMappingInfo(Object): +class AutoTensorizeMappingInfo(tvm_ffi.Object): """Necessary information used to perform transformations for tensorization.""" diff --git a/python/tvm/s_tir/schedule/instruction.py b/python/tvm/s_tir/schedule/instruction.py index dae7b8f6c89a..385373e97918 100644 --- a/python/tvm/s_tir/schedule/instruction.py +++ b/python/tvm/s_tir/schedule/instruction.py @@ -18,10 +18,9 @@ from typing import TYPE_CHECKING, Any +import tvm_ffi from tvm_ffi import register_object as _register_object -from tvm_tirx.runtime import Object - from . import _ffi_api if TYPE_CHECKING: @@ -35,7 +34,7 @@ @_register_object("s_tir.InstructionKind") -class InstructionKind(Object): +class InstructionKind(tvm_ffi.Object): """Kind of an instruction, e.g. Split, Reorder, etc. Besides the name, every kind of instruction has its own properties, including: 1) A boolean indicating if the instruction is pure, i.e. change nothing in the schedule state @@ -91,7 +90,7 @@ def get(name: str) -> "InstructionKind": @_register_object("s_tir.Instruction") -class Instruction(Object): +class Instruction(tvm_ffi.Object): """Schedule instructions each corresponds to a schedule primitive Attributes diff --git a/python/tvm/s_tir/schedule/schedule.py b/python/tvm/s_tir/schedule/schedule.py index 01286374d9be..86d5f58190ab 100644 --- a/python/tvm/s_tir/schedule/schedule.py +++ b/python/tvm/s_tir/schedule/schedule.py @@ -20,12 +20,13 @@ from collections.abc import Callable from typing import Literal +import tvm_ffi from tvm_ffi import register_error from tvm_ffi import register_object as _register_object from tvm.s_tir import SBlock from tvm_tirx.ir import Expr, FloatImm, For, GlobalVar, IntImm, IRModule, is_prim_expr -from tvm_tirx.runtime import DataTypeCode, MemoryScope, Object +from tvm_tirx.runtime import DataTypeCode, MemoryScope from tvm_tirx.tirx import Function, Var, is_tensor_var from tvm_tirx.tirx.index_map import IndexMap @@ -41,7 +42,7 @@ class ScheduleError(RuntimeError): @_register_object("s_tir.LoopRV") -class LoopRV(Object): +class LoopRV(tvm_ffi.Object): """A random variable that refers to a loop""" def __init__(self) -> None: @@ -52,7 +53,7 @@ def __init__(self) -> None: @_register_object("s_tir.SBlockRV") -class SBlockRV(Object): +class SBlockRV(tvm_ffi.Object): """A random variable that refers to a block""" def __init__(self) -> None: @@ -112,7 +113,7 @@ def _get_sblock_default_dtype(block: SBlock) -> str: @_register_object("s_tir.Schedule") -class Schedule(Object): +class Schedule(tvm_ffi.Object): """The user-facing schedule class A schedule is a set of transformations that change the order of computation but @@ -3497,7 +3498,7 @@ def two_elementwise_transformed_intermediate_tensor( ) elif not isinstance(pad_value, IndexMap): # Explicitly convert python int/float arguments to the - # tensor's type. If the default `tvm.runtime.convert` + # tensor's type. If the default `tvm_ffi.convert` # behavior is applied, these would be converted to # int32/float32, which may not match the tensor's type. if tensor_obj.ty.dtype.matches_code(DataTypeCode.INT, DataTypeCode.UINT) and isinstance( diff --git a/python/tvm/s_tir/schedule/state.py b/python/tvm/s_tir/schedule/state.py index 7279573047d6..8173f62bc9a9 100644 --- a/python/tvm/s_tir/schedule/state.py +++ b/python/tvm/s_tir/schedule/state.py @@ -20,11 +20,11 @@ from collections import namedtuple from enum import IntEnum +import tvm_ffi from tvm_ffi import register_object from tvm.s_tir import SBlock, SBlockRealize from tvm_tirx.ir import For, IRModule -from tvm_tirx.runtime import Object from tvm_tirx.tirx import Function from ..sblock_scope import SBlockScope, StmtSRef @@ -80,7 +80,7 @@ def _parse_enable_checks(enable_checks: bool) -> bool: @register_object("s_tir.ScheduleState") -class ScheduleState(Object): +class ScheduleState(tvm_ffi.Object): """The state of scheduling, which exposes a `Replace` method as the primary resort for all the scheduling primitives to manipulate the TensorIR. diff --git a/python/tvm/s_tir/schedule/trace.py b/python/tvm/s_tir/schedule/trace.py index faf02e7061bf..a8622a105a28 100644 --- a/python/tvm/s_tir/schedule/trace.py +++ b/python/tvm/s_tir/schedule/trace.py @@ -20,11 +20,11 @@ from collections.abc import Callable from typing import TYPE_CHECKING, Any +import tvm_ffi from tvm_ffi import Array, Map from tvm_ffi import register_object as _register_object from tvm_tirx.ir import FloatImm, IntImm, save_json -from tvm_tirx.runtime import Object from tvm_tirx.tirx.index_map import IndexMap from . import _ffi_api @@ -58,7 +58,7 @@ def _json_from_tvm(obj): @_register_object("s_tir.Trace") -class Trace(Object): +class Trace(tvm_ffi.Object): """An execution trace of a scheduling program. A trace has two parts: diff --git a/python/tvm/s_tir/stmt.py b/python/tvm/s_tir/stmt.py index 05bed5bcb27b..39a86fa0a388 100644 --- a/python/tvm/s_tir/stmt.py +++ b/python/tvm/s_tir/stmt.py @@ -23,13 +23,12 @@ from tvm.s_tir import IterVar from tvm_tirx.ir import Expr, Location, Scriptable, SeqStmt, Stmt, TensorRegion, Var, const from tvm_tirx.ir.location import UNKNOWN_LOC -from tvm_tirx.runtime import Object from . import _ffi_api @tvm_ffi.register_object("s_tir.MatchTensorRegion") -class MatchTensorRegion(Object, Scriptable): +class MatchTensorRegion(tvm_ffi.Object, Scriptable): """MatchTensorRegion node. Parameters @@ -82,7 +81,7 @@ class SBlock(Stmt): match_tensors: Optional[List[MatchTensorRegion]] The subregion tensor match - annotations: Optional[Mapping[str, Object]] + annotations: Optional[Mapping[str, tvm_ffi.Object]] Additional annotation hints. loc : Location, optional @@ -97,7 +96,7 @@ class SBlock(Stmt): init: SeqStmt | None alloc_tensors: list[Var] match_tensors: list[MatchTensorRegion] - annotations: Mapping[str, Object] + annotations: Mapping[str, tvm_ffi.Object] loc: Location def __init__( @@ -110,7 +109,7 @@ def __init__( init: Stmt | Sequence[Stmt] | None = None, alloc_tensors: list[Var] | None = None, match_tensors: list[MatchTensorRegion] | None = None, - annotations: Mapping[str, Object] | None = None, + annotations: Mapping[str, tvm_ffi.Object] | None = None, loc: Location = UNKNOWN_LOC, ) -> None: if alloc_tensors is None: diff --git a/python/tvm/s_tir/tensor_intrin/riscv_cpu.py b/python/tvm/s_tir/tensor_intrin/riscv_cpu.py index 5f90e09b6fe7..17280757dcf9 100644 --- a/python/tvm/s_tir/tensor_intrin/riscv_cpu.py +++ b/python/tvm/s_tir/tensor_intrin/riscv_cpu.py @@ -23,7 +23,6 @@ from tvm.s_tir import script as Ts from tvm_tirx.codegen import llvm, target_has_features -from tvm_tirx.runtime import DataType from tvm_tirx.script import tirx as T from tvm_tirx.target import Target @@ -93,21 +92,21 @@ def rvv_vec_dot_prod_desc( # LLVM only supports ELEN=32 or ELEN=64 # https://llvm.org/docs//RISCV/RISCVVectorExtension.html - d_dtype_lanes = (64 // DataType(data_dtype).bits) * lmul - w_dtype_lanes = (64 // DataType(weight_dtype).bits) * lmul + d_dtype_lanes = (64 // tvm_ffi.dtype(data_dtype).bits) * lmul + w_dtype_lanes = (64 // tvm_ffi.dtype(weight_dtype).bits) * lmul # reduction lanes narrows - o_dtype_lanes = (64 // DataType(out_dtype).bits) * lmul // n_lanes + o_dtype_lanes = (64 // tvm_ffi.dtype(out_dtype).bits) * lmul // n_lanes # data type widening case o_dtype_lanes = max(o_dtype_lanes, 2) - weight_bytes = DataType(weight_dtype).bits // 8 - output_bytes = DataType(out_dtype).bits // 8 + weight_bytes = tvm_ffi.dtype(weight_dtype).bits // 8 + output_bytes = tvm_ffi.dtype(out_dtype).bits // 8 mask_args = () if data_dtype[0] in ("i", "u") else (T.uint64(7),) wide_dtype = out_dtype - if DataType(out_dtype).bits > DataType(data_dtype).bits: + if tvm_ffi.dtype(out_dtype).bits > tvm_ffi.dtype(data_dtype).bits: wide_dtype = "".join(c for c in data_dtype if not c.isdigit()) - wide_dtype += str(DataType(data_dtype).bits * 2) + wide_dtype += str(tvm_ffi.dtype(data_dtype).bits * 2) # fmt: off @Ts.function @@ -198,18 +197,18 @@ def register_rvv_isa_intrinsics(target: Target, inventory_only=False) -> dict(): for d_dtype, w_dtype, o_dtype in zip(data_dtype, weight_dtype, output_dtype): # max elements to grouped registers - max_elems = get_max_elems(vlen, lmul=8, sew=DataType(d_dtype).bits) + max_elems = get_max_elems(vlen, lmul=8, sew=tvm_ffi.dtype(d_dtype).bits) # data widening halves available vector registers - if DataType(o_dtype).bits > DataType(d_dtype).bits: + if tvm_ffi.dtype(o_dtype).bits > tvm_ffi.dtype(d_dtype).bits: max_elems //= 2 # compute optimal LMUL for full load - lmul = max_elems // (vlen // DataType(d_dtype).bits) + lmul = max_elems // (vlen // tvm_ffi.dtype(d_dtype).bits) n_elems = max_elems while n_elems >= 4: - dt = DataType(d_dtype) - wt = DataType(w_dtype) - ot = DataType(o_dtype) + dt = tvm_ffi.dtype(d_dtype) + wt = tvm_ffi.dtype(w_dtype) + ot = tvm_ffi.dtype(o_dtype) kernel_name = "rvv_dot" kernel_name += f"_{n_elems}{dt[0]}{dt.bits}" kernel_name += f"_{n_lanes}x{n_elems}{wt[0]}{wt.bits}" diff --git a/python/tvm/s_tir/tensor_intrin/rocm.py b/python/tvm/s_tir/tensor_intrin/rocm.py index 44a9995b2143..3bdf2dad1211 100644 --- a/python/tvm/s_tir/tensor_intrin/rocm.py +++ b/python/tvm/s_tir/tensor_intrin/rocm.py @@ -26,16 +26,17 @@ extent without a warp-memory lowering pass. """ +import tvm_ffi + from tvm.s_tir import script as Ts from tvm_tirx.ir import IntImm -from tvm_tirx.runtime import convert from tvm_tirx.script import tirx as T from tvm_tirx.tirx.expr import Cast from .. import TensorIntrin from .dot_product_common import get_dp4a_intrin -lift = convert +lift = tvm_ffi.convert @Ts.function diff --git a/python/tvm/s_tir/testing/runner.py b/python/tvm/s_tir/testing/runner.py index e2384a9bfacb..f4b4474f534b 100644 --- a/python/tvm/s_tir/testing/runner.py +++ b/python/tvm/s_tir/testing/runner.py @@ -20,13 +20,15 @@ from collections.abc import Callable from typing import TYPE_CHECKING, Literal, Optional, Union +import tvm_ffi + from tvm.runtime import time_evaluator if TYPE_CHECKING: import numpy as np from tvm.s_tir.meta_schedule.runner import EvaluatorConfig, RPCConfig - from tvm_tirx.runtime import Device, Module, Tensor + from tvm_tirx.runtime import Device, Tensor # pylint: disable=import-outside-toplevel,protected-access @@ -87,18 +89,18 @@ def export_with(func): def local_run( # pylint: disable=too-many-arguments,too-many-locals - mod: "Module", + mod: "tvm_ffi.Module", device_type: str, args: list[Union["np.ndarray", "Tensor", int, float]], evaluator_config: Optional["EvaluatorConfig"] = None, - export_func: Callable[["Module", str], None] | Literal["tar", "ndk"] = "tar", + export_func: Callable[["tvm_ffi.Module", str], None] | Literal["tar", "ndk"] = "tar", output_format: str | None = None, ): """Run a TVM module on a local device. Parameters ---------- - mod : Module + mod : tvm_ffi.Module The TVM module to run. device_type : str The device type to run the module on. @@ -106,7 +108,7 @@ def local_run( # pylint: disable=too-many-arguments,too-many-locals The arguments to be fed to the module. evaluator_config : Optional[EvaluatorConfig] The evaluator configuration to use. - export_func : Union[Callable[Module, str], Literal["tar", "ndk"]] + export_func : Union[Callable[tvm_ffi.Module, str], Literal["tar", "ndk"]] The function to export the module to a file. If callable, it must be a function that takes two arguments: the module to export and the path to export to. @@ -127,7 +129,7 @@ def local_run( # pylint: disable=too-many-arguments,too-many-locals import tempfile from tvm.s_tir.meta_schedule.runner import EvaluatorConfig - from tvm_tirx.runtime import device, load_module + from tvm_tirx.runtime import load_module from tvm_tirx.support import cc, tar, utils evaluator_config = EvaluatorConfig._normalized(evaluator_config) @@ -136,7 +138,7 @@ def local_run( # pylint: disable=too-many-arguments,too-many-locals with tempfile.TemporaryDirectory() as tmp_dir: artifact_path = osp.join(tmp_dir, "tvm_tmp_mod." + output_format) export_func(mod, artifact_path) - device: Device = device(device_type, 0) + device: Device = tvm_ffi.device(device_type, 0) args = _args_to_device(args, device) if artifact_path.endswith((".o", ".tar")): @@ -164,19 +166,19 @@ def local_run( # pylint: disable=too-many-arguments,too-many-locals def rpc_run( # pylint: disable=too-many-arguments,too-many-locals - mod: "Module", + mod: "tvm_ffi.Module", device_type: str, args: list[Union["np.ndarray", "Tensor", int, float]], evaluator_config: Optional["EvaluatorConfig"] = None, rpc_config: Optional["RPCConfig"] = None, - export_func: Callable[["Module", str], None] | Literal["tar", "ndk"] = "tar", + export_func: Callable[["tvm_ffi.Module", str], None] | Literal["tar", "ndk"] = "tar", output_format: str | None = None, ): """Run a TVM module on a remote device. Parameters ---------- - mod : Module + mod : tvm_ffi.Module The TVM module to run. device_type : str The device type to run the module on. @@ -191,7 +193,7 @@ def rpc_run( # pylint: disable=too-many-arguments,too-many-locals - TVM_TRACKER_HOST - TVM_TRACKER_PORT - TVM_TRACKER_KEY - export_func : Union[Callable[Module, str], Literal["tar", "ndk"]] + export_func : Union[Callable[tvm_ffi.Module, str], Literal["tar", "ndk"]] The function to export the module to a file. If callable, it must be a function that takes two arguments: the module to export and the path to export to. diff --git a/python/tvm/te/operation.py b/python/tvm/te/operation.py index 53d6b90d7792..82e791b71479 100644 --- a/python/tvm/te/operation.py +++ b/python/tvm/te/operation.py @@ -21,6 +21,7 @@ # pylint: disable=invalid-name from numbers import Integral as _Integral +import tvm_ffi from tvm_ffi import Array import tvm @@ -29,7 +30,6 @@ import tvm_tirx.tirx._ffi_api from tvm_tirx.ir import is_prim_expr from tvm_tirx.ir.location import UNKNOWN_LOC, Location -from tvm_tirx.runtime import convert from . import _ffi_api from . import tag as _tag @@ -135,7 +135,7 @@ def compute(shape, fcompute, name="compute", tag="", attrs=None, varargs_names=N if not isinstance(body, list | tuple): body = [body] - body = convert(body) + body = tvm_ffi.convert(body) op_node = _ffi_api.ComputeOp(name, tag, attrs, dim_var, body) num = op_node.num_outputs diff --git a/python/tvm/te/reduction.py b/python/tvm/te/reduction.py index 8a7e7b3e6904..c5aaa599839c 100644 --- a/python/tvm/te/reduction.py +++ b/python/tvm/te/reduction.py @@ -29,7 +29,6 @@ from tvm_tirx.ir.location import UNKNOWN_LOC, Location from tvm_tirx.ir.prim import _ffi_api as _prim_ffi_api from tvm_tirx.ir.prim import max_value, min_value -from tvm_tirx.runtime import Object from tvm_tirx.tirx.op import _primexpr_dtype from . import _ffi_api @@ -45,7 +44,7 @@ def _iter_var_type(): @tvm_ffi.register_object("te.CommReducer") -class CommReducer(Object, Scriptable): +class CommReducer(tvm_ffi.Object, Scriptable): """Commutative reduce operator Parameters @@ -227,7 +226,7 @@ def _make_reduce(expr, axis, where=None, init=None): if init is not None: init = [init] combiner = CommReducer(lhs, rhs, result, id_elem) - if not isinstance(axis, list | tuple | tvm_tirx.ir.Array): + if not isinstance(axis, list | tuple | tvm_ffi.Array): axis = [axis] if where is None: where = tir.convert(True) diff --git a/python/tvm/te/tensor.py b/python/tvm/te/tensor.py index 2f4c896831a7..eee138d49ce3 100644 --- a/python/tvm/te/tensor.py +++ b/python/tvm/te/tensor.py @@ -21,7 +21,6 @@ from tvm_tirx.ir import Expr, OpaqueExpr, SeqStmt, const, is_prim_expr from tvm_tirx.ir.location import UNKNOWN_LOC, Location -from tvm_tirx.runtime import Object from tvm_tirx.tirx import expr as _expr from . import _ffi_api, _te_tensor_overload @@ -306,7 +305,7 @@ def name(self): @tvm_ffi.register_object("te.Operation") -class Operation(Object): +class Operation(tvm_ffi.Object): """Represent an operation that generates a tensor""" def output(self, index): diff --git a/python/tvm/topi/utils.py b/python/tvm/topi/utils.py index 4c21ddb99d2b..9b183a9cd646 100644 --- a/python/tvm/topi/utils.py +++ b/python/tvm/topi/utils.py @@ -21,6 +21,7 @@ from numbers import Integral import numpy as np +import tvm_ffi import tvm import tvm_tirx @@ -412,9 +413,7 @@ def get_shape(src_shape, src_layout, dst_layout): assert len(src_layout) == len(dst_layout), f"Incompatible layout {src_layout} vs {dst_layout}" layout_mapping = sbijective_layout(src_layout, dst_layout) - dst_indices = layout_mapping.forward_index( - tvm_tirx.runtime.convert(list(range(len(src_layout)))) - ) + dst_indices = layout_mapping.forward_index(tvm_ffi.convert(list(range(len(src_layout))))) return get_const_tuple(tuple([src_shape[i.value] for i in dst_indices])) diff --git a/src/runtime/vm/builtin.cc b/src/runtime/vm/builtin.cc index ec1bd91aa1d4..334e1d376081 100644 --- a/src/runtime/vm/builtin.cc +++ b/src/runtime/vm/builtin.cc @@ -683,7 +683,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { const auto debug_func = tvm::ffi::Function::GetGlobal(debug_func_name); TVM_FFI_CHECK(debug_func.has_value(), ValueError) << debug_func_name << " is not found. " - << "Use the decorator `@tvm.register_global_func(\"" << debug_func_name + << "Use the decorator `@tvm_ffi.register_global_func(\"" << debug_func_name << "\")` to register it."; ffi::String line_info = args[2].cast(); std::vector call_args(num_args + 1); diff --git a/tests/python/all-platform-minimal-test/test_minimal_target_codegen_llvm.py b/tests/python/all-platform-minimal-test/test_minimal_target_codegen_llvm.py index dd297f414fd2..48306276f848 100644 --- a/tests/python/all-platform-minimal-test/test_minimal_target_codegen_llvm.py +++ b/tests/python/all-platform-minimal-test/test_minimal_target_codegen_llvm.py @@ -23,6 +23,7 @@ import numpy as np import pytest +import tvm_ffi import tvm import tvm_tirx @@ -36,7 +37,7 @@ def test_llvm_add_pipeline(): """all-platform-minimal-test: Check LLVM enablement.""" nn = 128 - n = tvm_tirx.runtime.convert(nn) + n = tvm_ffi.convert(nn) A = te.placeholder((n,), name="A") B = te.placeholder((n,), name="B") AA = te.compute((n,), lambda *i: A(*i), name="A") diff --git a/tests/python/codegen/llvm/test_cross_compile.py b/tests/python/codegen/llvm/test_cross_compile.py index c1135013e831..5f43cc334bef 100644 --- a/tests/python/codegen/llvm/test_cross_compile.py +++ b/tests/python/codegen/llvm/test_cross_compile.py @@ -22,6 +22,7 @@ import numpy as np import pytest +import tvm_ffi import tvm import tvm_tirx @@ -97,7 +98,7 @@ def build_arm(): @pytest.mark.skipif( - tvm_tirx.get_global_func("rpc.LocalSession", allow_missing=True) is None, reason="need rpc" + tvm_ffi.get_global_func("rpc.LocalSession", allow_missing=True) is None, reason="need rpc" ) def test_bigendian_rpc(): """Test big endian rpc when there is a PowerPC RPC server available""" diff --git a/tests/python/contrib/test_cblas.py b/tests/python/contrib/test_cblas.py index 867f20db282f..59fc979b393e 100644 --- a/tests/python/contrib/test_cblas.py +++ b/tests/python/contrib/test_cblas.py @@ -18,6 +18,7 @@ import numpy as np import pytest +import tvm_ffi pytest.importorskip("scipy") # tvm.topi.testing imports scipy @@ -61,7 +62,7 @@ def verify(target="llvm"): if not tvm_tirx.testing.device_enabled(target): print(f"skip because {target} is not enabled...") return - if not tvm_tirx.get_global_func(lib.__name__ + ".matmul", True): + if not tvm_ffi.get_global_func(lib.__name__ + ".matmul", True): print("skip because extern function is not available") return dev = tvm_tirx.cpu(0) @@ -125,7 +126,7 @@ def test_matmul_add(): def verify_quantized_matmul_add(matrix_m, matrix_l, matrix_n, transa=False, transb=False): """Tests quantized matmul+add op""" - if not tvm_tirx.get_global_func("tvm.contrib.mkl.matmul_u8s8s32", True): + if not tvm_ffi.get_global_func("tvm.contrib.mkl.matmul_u8s8s32", True): pytest.skip("Quantized dense is supported only for MKL. TVM GPU CI uses openblas") data_dtype = "uint8" kernel_dtype = "int8" @@ -151,7 +152,7 @@ def verify(target="llvm"): if not tvm_tirx.testing.device_enabled(target): print(f"skip because {target} is not enabled...") return - if not tvm_tirx.get_global_func("tvm.contrib.mkl.matmul_u8s8s32", True): + if not tvm_ffi.get_global_func("tvm.contrib.mkl.matmul_u8s8s32", True): print("skip because extern function is not available") return dev = tvm_tirx.cpu(0) @@ -236,7 +237,7 @@ def verify(target="llvm"): if not tvm_tirx.testing.device_enabled(target): print(f"skip because {target} is not enabled...") return - if not tvm_tirx.get_global_func(lib.__name__ + ".matmul", True): + if not tvm_ffi.get_global_func(lib.__name__ + ".matmul", True): print("skip because extern function is not available") return dev = tvm_tirx.cpu(0) diff --git a/tests/python/contrib/test_cutlass_gemm.py b/tests/python/contrib/test_cutlass_gemm.py index db780ab7fcbf..dd5eca110673 100644 --- a/tests/python/contrib/test_cutlass_gemm.py +++ b/tests/python/contrib/test_cutlass_gemm.py @@ -18,6 +18,7 @@ import ml_dtypes import numpy as np import pytest +import tvm_ffi import tvm_tirx import tvm_tirx.testing @@ -36,7 +37,7 @@ def get_random_tensor(shape, dtype): def verify_group_gemm( func_name, M, N, K, num_groups, x_dtype, weight_dtype, out_dtype, use_scale, rtol, atol ): - group_gemm_func = tvm_tirx.get_global_func(func_name, allow_missing=True) + group_gemm_func = tvm_ffi.get_global_func(func_name, allow_missing=True) if group_gemm_func is None: print(f"Skipped as {func_name} is not available") return @@ -317,7 +318,7 @@ def test_fp8_e4m3_groupwise_scaled_gemm(): assert N % 128 == 0 and K % 128 == 0 # Only support N/K are multiple of 128 func_name = "cutlass.groupwise_scaled_gemm_e4m3fn_e4m3fn" - gemm_func = tvm_tirx.get_global_func(func_name, allow_missing=True) + gemm_func = tvm_ffi.get_global_func(func_name, allow_missing=True) if gemm_func is None: print(f"Skipped as {func_name} is not available") return @@ -362,7 +363,7 @@ def test_fp8_e4m3_groupwise_scaled_bmm(): assert N % 128 == 0 and K % 128 == 0 # Only support N/K are multiple of 128 func_name = "cutlass.groupwise_scaled_bmm_e4m3fn_e4m3fn" - gemm_func = tvm_tirx.get_global_func(func_name, allow_missing=True) + gemm_func = tvm_ffi.get_global_func(func_name, allow_missing=True) if gemm_func is None: print(f"Skipped as {func_name} is not available") return diff --git a/tests/python/contrib/test_dlpack.py b/tests/python/contrib/test_dlpack.py index 2486be033130..cb7066b1a69b 100644 --- a/tests/python/contrib/test_dlpack.py +++ b/tests/python/contrib/test_dlpack.py @@ -15,6 +15,7 @@ # specific language governing permissions and limitations # under the License. import numpy as np +import tvm_ffi import tvm import tvm_tirx @@ -26,20 +27,20 @@ def verify_torch_dlpack(): a = np.random.randn(1337) tvm_a = tvm_tirx.runtime.tensor(a) - np.testing.assert_equal(tvm_tirx.runtime.from_dlpack(tvm_a).numpy(), a) + np.testing.assert_equal(tvm_ffi.from_dlpack(tvm_a).numpy(), a) try: import torch import torch.utils.dlpack x = torch.rand(56, 56) - tvm_x = tvm_tirx.runtime.from_dlpack(torch.utils.dlpack.to_dlpack(x)) + tvm_x = tvm_ffi.from_dlpack(torch.utils.dlpack.to_dlpack(x)) np.testing.assert_equal(x.numpy(), tvm_x.numpy()) - y = tvm_tirx.runtime.from_dlpack(tvm_x) + y = tvm_ffi.from_dlpack(tvm_x) np.testing.assert_equal(y.numpy(), tvm_x.numpy()) np.testing.assert_equal(torch.utils.dlpack.from_dlpack(y).numpy(), tvm_x.numpy()) - n = tvm_tirx.runtime.convert(137) + n = tvm_ffi.convert(137) xx = torch.rand(137, 137) yy = torch.rand(137, 137) zz2 = torch.empty(137, 137) diff --git a/tests/python/contrib/test_example_npu.py b/tests/python/contrib/test_example_npu.py index fd0efe7bd750..123a9cb1cd5d 100644 --- a/tests/python/contrib/test_example_npu.py +++ b/tests/python/contrib/test_example_npu.py @@ -23,6 +23,7 @@ import numpy as np import pytest +import tvm_ffi import tvm_tirx import tvm_tirx.testing @@ -96,8 +97,8 @@ def main(x: R.Tensor((2, 8), "float32")) -> R.Tensor((2, 8), "float32"): # Check if the example NPU runtime is available -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 example_npu_enabled = pytest.mark.skipif( diff --git a/tests/python/contrib/test_hipblas.py b/tests/python/contrib/test_hipblas.py index 9e424bbdd34c..e47907f94169 100644 --- a/tests/python/contrib/test_hipblas.py +++ b/tests/python/contrib/test_hipblas.py @@ -17,6 +17,7 @@ # ruff: noqa: E741 import numpy as np import pytest +import tvm_ffi import tvm import tvm_tirx @@ -35,7 +36,7 @@ def verify_matmul_add(in_dtype, out_dtype, rtol=1e-5): C = hipblas.matmul(A, B, dtype=out_dtype) def verify(target="rocm"): - if not tvm_tirx.get_global_func("tvm.contrib.hipblas.matmul", True): + if not tvm_ffi.get_global_func("tvm.contrib.hipblas.matmul", True): print("skip because extern function is not available") return f = tvm.compile(te.create_function([A, B, C]), target=target) @@ -100,7 +101,7 @@ def test_matmul_add(): @pytest.mark.gpu @pytest.mark.skipif(not env.has_rocm(), reason="need rocm") def test_batch_matmul(): - if not tvm_tirx.get_global_func("tvm.contrib.hipblas.batch_matmul", True): + if not tvm_ffi.get_global_func("tvm.contrib.hipblas.batch_matmul", True): print("skip because extern function is not available") return diff --git a/tests/python/contrib/test_random.py b/tests/python/contrib/test_random.py index 8eb7a9e59448..4dd1e9413ac1 100644 --- a/tests/python/contrib/test_random.py +++ b/tests/python/contrib/test_random.py @@ -22,6 +22,7 @@ import numpy as np import pytest +import tvm_ffi import tvm import tvm_tirx @@ -40,7 +41,7 @@ def verify(target="llvm"): if not tvm_tirx.testing.device_enabled(target): print(f"skip because {target} is not enabled...") return - if not tvm_tirx.get_global_func("tvm.contrib.random.randint", True): + if not tvm_ffi.get_global_func("tvm.contrib.random.randint", True): print("skip because extern function is not available") return dev = tvm_tirx.cpu(0) @@ -65,7 +66,7 @@ def verify(target="llvm"): if not tvm_tirx.testing.device_enabled(target): print(f"skip because {target} is not enabled...") return - if not tvm_tirx.get_global_func("tvm.contrib.random.uniform", True): + if not tvm_ffi.get_global_func("tvm.contrib.random.uniform", True): print("skip because extern function is not available") return dev = tvm_tirx.cpu(0) @@ -90,7 +91,7 @@ def verify(target="llvm"): if not tvm_tirx.testing.device_enabled(target): print(f"skip because {target} is not enabled...") return - if not tvm_tirx.get_global_func("tvm.contrib.random.normal", True): + if not tvm_ffi.get_global_func("tvm.contrib.random.normal", True): print("skip because extern function is not available") return dev = tvm_tirx.cpu(0) @@ -124,16 +125,16 @@ def test_random_fill(target): """Tests random_fill function""" if not tvm_tirx.testing.device_enabled(target): pytest.skip("target is unavailable or excluded by TVM_TEST_TARGETS") - if not tvm_tirx.get_global_func("tvm.contrib.random.random_fill", True): + if not tvm_ffi.get_global_func("tvm.contrib.random.random_fill", True): pytest.skip("random_fill is not built") dev = tvm_tirx.device_from_target(target) def test_local(dev, dtype): - if not tvm_tirx.get_global_func("tvm.contrib.random.random_fill", True): + if not tvm_ffi.get_global_func("tvm.contrib.random.random_fill", True): print("skip because extern function is not available") return value = tvm_tirx.runtime.empty((512, 512), dtype, dev) - random_fill = tvm_tirx.get_global_func("tvm.contrib.random.random_fill") + random_fill = tvm_ffi.get_global_func("tvm.contrib.random.random_fill") random_fill(value) assert np.count_nonzero(value.numpy()) == 512 * 512 @@ -143,7 +144,7 @@ def test_local(dev, dtype): assert np.isfinite(np_values * np_values + np_values).any() def test_rpc(dtype): - if not tvm_tirx.get_global_func("tvm.contrib.random.random_fill", True): + if not tvm_ffi.get_global_func("tvm.contrib.random.random_fill", True): print("skip because extern function is not available") return if not tvm_tirx.testing.device_enabled("rpc") or not tvm_tirx.runtime.enabled("llvm"): @@ -195,11 +196,11 @@ def test_random_fill_mt(): def test_body(): try: num_thread_used = 1 - configure_threads = tvm_tirx.get_global_func("runtime.config_threadpool") + configure_threads = tvm_ffi.get_global_func("runtime.config_threadpool") configure_threads(1, num_thread_used) test_input = tvm_tirx.runtime.empty((10, 10)) - random_fill = tvm_tirx.get_global_func("tvm.contrib.random.random_fill_for_measure") + random_fill = tvm_ffi.get_global_func("tvm.contrib.random.random_fill_for_measure") random_fill(test_input) except: # pylint: disable=bare-except nonlocal no_exception_happened diff --git a/tests/python/contrib/test_rpc_tracker.py b/tests/python/contrib/test_rpc_tracker.py index 95a97fe776e0..0d761a159063 100644 --- a/tests/python/contrib/test_rpc_tracker.py +++ b/tests/python/contrib/test_rpc_tracker.py @@ -20,7 +20,8 @@ import logging import time -import tvm_tirx +import tvm_ffi + from tvm import rpc @@ -33,7 +34,7 @@ def check_server_drop(): # pylint: disable=import-outside-toplevel from tvm.rpc.base import TrackerCode - @tvm_tirx.register_global_func("rpc.test2.addone") + @tvm_ffi.register_global_func("rpc.test2.addone") def addone(x): return x + 1 diff --git a/tests/python/disco/test_callback.py b/tests/python/disco/test_callback.py index 967dc3f0719d..0d768b89abd3 100644 --- a/tests/python/disco/test_callback.py +++ b/tests/python/disco/test_callback.py @@ -112,7 +112,7 @@ def transform_params( temp_dir = pathlib.Path(temp_dir) # TODO(Lunderberg): Update `disco.Session.load_vm_module` to - # allow a `tvm.runtime.Module` argument. This would avoid the + # allow a `tvm_ffi.Module` argument. This would avoid the # need for a temporary file. shlib_path = temp_dir.joinpath("libtemp.so") built.export_library(shlib_path) diff --git a/tests/python/disco/test_ccl.py b/tests/python/disco/test_ccl.py index b772b45ee561..4dfdd6760a6b 100644 --- a/tests/python/disco/test_ccl.py +++ b/tests/python/disco/test_ccl.py @@ -21,6 +21,7 @@ import numpy as np import pytest +import tvm_ffi import tvm import tvm_tirx @@ -30,13 +31,12 @@ from tvm.runtime import disco as di from tvm.runtime.vm import VirtualMachine from tvm.s_tir import dlight as dl -from tvm_tirx import get_global_func if di is None: pytest.skip("disco runtime is not available", allow_module_level=True) _all_session_kinds = [di.ThreadedSession, di.ProcessSession] -_compiled_ccl = get_global_func("runtime.disco.compiled_ccl", allow_missing=True) +_compiled_ccl = tvm_ffi.get_global_func("runtime.disco.compiled_ccl", allow_missing=True) if _compiled_ccl is None: pytest.skip("Disco CCL is not enabled in this TVM build", allow_module_level=True) _ccl = [_compiled_ccl()] diff --git a/tests/python/disco/test_custom_allreduce.py b/tests/python/disco/test_custom_allreduce.py index e0c8f275e0e0..8170d0c96c9c 100644 --- a/tests/python/disco/test_custom_allreduce.py +++ b/tests/python/disco/test_custom_allreduce.py @@ -21,12 +21,10 @@ import numpy as np import pytest +import tvm_ffi from tvm_ffi import Shape -import tvm_tirx -import tvm_tirx.testing from tvm.runtime import disco -from tvm_tirx.runtime import DataType if disco is None: pytest.skip("disco runtime is not available", allow_module_level=True) @@ -48,7 +46,7 @@ class AllReduceStrategyType(enum.IntEnum): AllReduceStrategyType.AUTO, ] -_compiled_ccl = tvm_tirx.get_global_func("runtime.disco.compiled_ccl", allow_missing=True) +_compiled_ccl = tvm_ffi.get_global_func("runtime.disco.compiled_ccl", allow_missing=True) if _compiled_ccl is None: pytest.skip("Disco CCL is not enabled in this TVM build", allow_module_level=True) _ccl = [ccl for ccl in _compiled_ccl() if ccl == "nccl"] @@ -66,7 +64,7 @@ def test_allreduce(shape, ccl, strategy): dtype = "float32" falloc_ipc_storage = sess.get_global_func("runtime.disco.cuda_ipc.alloc_storage") fallreduce = sess.get_global_func("runtime.disco.cuda_ipc.custom_allreduce") - d_input = sess.call_packed(falloc_ipc_storage, Shape(shape), DataType(dtype)) + d_input = sess.call_packed(falloc_ipc_storage, Shape(shape), tvm_ffi.dtype(dtype)) array_1 = np.arange(num_elements, dtype="float32").reshape(*shape) array_2 = np.arange(start=1, stop=-(num_elements - 1), step=-1, dtype="float32").reshape(*shape) diff --git a/tests/python/disco/test_loader.py b/tests/python/disco/test_loader.py index 9ef42b65ff57..fc45b6043742 100644 --- a/tests/python/disco/test_loader.py +++ b/tests/python/disco/test_loader.py @@ -23,6 +23,7 @@ import numpy as np import pytest +import tvm_ffi from tvm_ffi import Shape, register_global_func import tvm @@ -39,7 +40,7 @@ # `runtime.disco.compiled_ccl` is registered together with the CCL runtime # functions, so its absence means the disco CCL runtime is not in this build. -_compiled_ccl = tvm_tirx.get_global_func("runtime.disco.compiled_ccl", allow_missing=True) +_compiled_ccl = tvm_ffi.get_global_func("runtime.disco.compiled_ccl", allow_missing=True) if _compiled_ccl is None or _compiled_ccl() != "nccl": pytest.skip("Disco NCCL support is not available", allow_module_level=True) diff --git a/tests/python/disco/test_nvshmem.py b/tests/python/disco/test_nvshmem.py index f9cb5be833ba..5f2e4dfecb59 100644 --- a/tests/python/disco/test_nvshmem.py +++ b/tests/python/disco/test_nvshmem.py @@ -27,6 +27,7 @@ import numpy as np import pytest +import tvm_ffi from tvm_ffi import Shape import tvm @@ -181,7 +182,7 @@ def _require_cuda_devices(num_workers): def _init_finalize(session_kind, num_workers): sess = session_kind(num_workers=num_workers) - f_init_nvshmem_uid = tvm_tirx.get_global_func("runtime.disco.nvshmem.init_nvshmem_uid") + f_init_nvshmem_uid = tvm_ffi.get_global_func("runtime.disco.nvshmem.init_nvshmem_uid") uid = f_init_nvshmem_uid() init_dfunc = sess.get_global_func("runtime.disco.nvshmem.init_nvshmem") init_dfunc(uid, num_workers, 0) @@ -194,7 +195,7 @@ def _init_finalize(session_kind, num_workers): def _empty(session_kind, num_workers): device = tvm_tirx.cuda() sess = session_kind(num_workers=num_workers) - f_init_nvshmem_uid = tvm_tirx.get_global_func("runtime.disco.nvshmem.init_nvshmem_uid") + f_init_nvshmem_uid = tvm_ffi.get_global_func("runtime.disco.nvshmem.init_nvshmem_uid") uid = f_init_nvshmem_uid() init_dfunc = sess.get_global_func("runtime.disco.nvshmem.init_nvshmem") init_dfunc(uid, num_workers, 0) @@ -226,7 +227,7 @@ def _compile(): num_workers = 2 sess = di.ProcessSession(num_workers=num_workers) - f_init_nvshmem_uid = tvm_tirx.get_global_func("runtime.disco.nvshmem.init_nvshmem_uid") + f_init_nvshmem_uid = tvm_ffi.get_global_func("runtime.disco.nvshmem.init_nvshmem_uid") uid = f_init_nvshmem_uid() init_dfunc = sess.get_global_func("runtime.disco.nvshmem.init_nvshmem") init_dfunc(uid, num_workers, 0) @@ -295,7 +296,7 @@ def _kernel_compile(compile_mode): num_workers = 2 sess = di.ProcessSession(num_workers=num_workers) - f_init_nvshmem_uid = tvm_tirx.get_global_func("runtime.disco.nvshmem.init_nvshmem_uid") + f_init_nvshmem_uid = tvm_ffi.get_global_func("runtime.disco.nvshmem.init_nvshmem_uid") uid = f_init_nvshmem_uid() init_dfunc = sess.get_global_func("runtime.disco.nvshmem.init_nvshmem") init_dfunc(uid, num_workers, 0) diff --git a/tests/python/relax/backend/adreno/utils.py b/tests/python/relax/backend/adreno/utils.py index ae0091a0ae58..402c4b7a1edc 100644 --- a/tests/python/relax/backend/adreno/utils.py +++ b/tests/python/relax/backend/adreno/utils.py @@ -20,6 +20,7 @@ import numpy as np import pytest +import tvm_ffi import tvm import tvm_tirx @@ -68,7 +69,7 @@ def __call__(self): # CLML Codegen skip_unless_adreno_clml = pytest.mark.skipif( - tvm_tirx.get_global_func("relax.is_openclml_runtime_enabled", allow_missing=True) is None, + tvm_ffi.get_global_func("relax.is_openclml_runtime_enabled", allow_missing=True) is None, reason="need adreno openclml", ) diff --git a/tests/python/relax/nvshmem/test_runtime_builtin_kv_cache_transfer.py b/tests/python/relax/nvshmem/test_runtime_builtin_kv_cache_transfer.py index 51db3bed34ed..ecb354668929 100644 --- a/tests/python/relax/nvshmem/test_runtime_builtin_kv_cache_transfer.py +++ b/tests/python/relax/nvshmem/test_runtime_builtin_kv_cache_transfer.py @@ -127,29 +127,29 @@ def set_global_func(head_dim, dtype, target): global fmerge_state, fsplit_rotary, fattention_rotary, fcopy_single_page, fcompact_copy global fnvshmem_get_uid, fnvshmem_init, fdisagg_mark_send, fdisagg_prepare_recv - fclear = tvm_tirx.get_global_func("vm.builtin.kv_state_clear") - fadd_sequence = tvm_tirx.get_global_func("vm.builtin.kv_state_add_sequence") - fremove_sequence = tvm_tirx.get_global_func("vm.builtin.kv_state_remove_sequence") - ffork_sequence = tvm_tirx.get_global_func("vm.builtin.kv_state_fork_sequence") - fenable_sliding_window_for_seq = tvm_tirx.get_global_func( + fclear = tvm_ffi.get_global_func("vm.builtin.kv_state_clear") + fadd_sequence = tvm_ffi.get_global_func("vm.builtin.kv_state_add_sequence") + fremove_sequence = tvm_ffi.get_global_func("vm.builtin.kv_state_remove_sequence") + ffork_sequence = tvm_ffi.get_global_func("vm.builtin.kv_state_fork_sequence") + fenable_sliding_window_for_seq = tvm_ffi.get_global_func( "vm.builtin.attention_kv_cache_enable_sliding_window_for_seq" ) - fpopn = tvm_tirx.get_global_func("vm.builtin.kv_state_popn") - fbegin_forward = tvm_tirx.get_global_func("vm.builtin.kv_state_begin_forward") - fend_forward = tvm_tirx.get_global_func("vm.builtin.kv_state_end_forward") - fcommit_accepted_token_tree_nodes = tvm_tirx.get_global_func( + fpopn = tvm_ffi.get_global_func("vm.builtin.kv_state_popn") + fbegin_forward = tvm_ffi.get_global_func("vm.builtin.kv_state_begin_forward") + fend_forward = tvm_ffi.get_global_func("vm.builtin.kv_state_end_forward") + fcommit_accepted_token_tree_nodes = tvm_ffi.get_global_func( "vm.builtin.attention_kv_cache_commit_accepted_token_tree_nodes" ) - fattention_with_fuse_qkv = tvm_tirx.get_global_func( + fattention_with_fuse_qkv = tvm_ffi.get_global_func( "vm.builtin.attention_kv_cache_attention_with_fused_qkv" ) - fis_empty = tvm_tirx.get_global_func("vm.builtin.attention_kv_cache_empty") - fdebug_get_kv = tvm_tirx.get_global_func("vm.builtin.attention_kv_cache_debug_get_kv") + fis_empty = tvm_ffi.get_global_func("vm.builtin.attention_kv_cache_empty") + fdebug_get_kv = tvm_ffi.get_global_func("vm.builtin.attention_kv_cache_debug_get_kv") - fnvshmem_get_uid = tvm_tirx.get_global_func("runtime.disco.nvshmem.init_nvshmem_uid") - fnvshmem_init = tvm_tirx.get_global_func("runtime.disco.nvshmem.init_nvshmem") - fdisagg_mark_send = tvm_tirx.get_global_func("vm.builtin.kv_cache_disagg_mark_send") - fdisagg_prepare_recv = tvm_tirx.get_global_func("vm.builtin.kv_cache_disagg_prepare_recv") + fnvshmem_get_uid = tvm_ffi.get_global_func("runtime.disco.nvshmem.init_nvshmem_uid") + fnvshmem_init = tvm_ffi.get_global_func("runtime.disco.nvshmem.init_nvshmem") + fdisagg_mark_send = tvm_ffi.get_global_func("vm.builtin.kv_cache_disagg_mark_send") + fdisagg_prepare_recv = tvm_ffi.get_global_func("vm.builtin.kv_cache_disagg_prepare_recv") builts = [] for tir_func in [ @@ -199,7 +199,7 @@ def set_global_func(head_dim, dtype, target): def create_kv_cache(head_dim, dtype, rope_mode, support_sliding_window): - fcreate = tvm_tirx.get_global_func("vm.builtin.paged_attention_kv_cache_create") + fcreate = tvm_ffi.get_global_func("vm.builtin.paged_attention_kv_cache_create") cache = fcreate( tvm_ffi.Shape( [ @@ -756,12 +756,12 @@ def test_paged_attention_kv_cache_transfer(kv_cache_and_config): def init_nvshmem(num_workers, pe_offset): if rank == 0: - f_init_nvshmem_uid = tvm_tirx.get_global_func("runtime.disco.nvshmem.init_nvshmem_uid") + f_init_nvshmem_uid = tvm_ffi.get_global_func("runtime.disco.nvshmem.init_nvshmem_uid") uid = f_init_nvshmem_uid() else: uid = None uid = comm.bcast(uid, root=0) - init_func = tvm_tirx.get_global_func("runtime.disco.nvshmem.init_nvshmem") + init_func = tvm_ffi.get_global_func("runtime.disco.nvshmem.init_nvshmem") init_func(uid, num_workers, pe_offset) diff --git a/tests/python/relax/nvshmem/test_runtime_builtin_kv_cache_transfer_kernel.py b/tests/python/relax/nvshmem/test_runtime_builtin_kv_cache_transfer_kernel.py index 6dbf95705831..4d11d968d961 100644 --- a/tests/python/relax/nvshmem/test_runtime_builtin_kv_cache_transfer_kernel.py +++ b/tests/python/relax/nvshmem/test_runtime_builtin_kv_cache_transfer_kernel.py @@ -18,6 +18,7 @@ import numpy as np import pytest +import tvm_ffi from tvm_ffi import Shape import tvm_tirx @@ -67,14 +68,14 @@ def test_kv_transfer_without_disco(): layer_id = 1 dev = tvm_tirx.cuda(rank) if rank == 0: - f_init_nvshmem_uid = tvm_tirx.get_global_func("runtime.disco.nvshmem.init_nvshmem_uid") + f_init_nvshmem_uid = tvm_ffi.get_global_func("runtime.disco.nvshmem.init_nvshmem_uid") uid = f_init_nvshmem_uid() else: uid = None uid = comm.bcast(uid, root=0) - init_func = tvm_tirx.get_global_func("runtime.disco.nvshmem.init_nvshmem") + init_func = tvm_ffi.get_global_func("runtime.disco.nvshmem.init_nvshmem") init_func(uid, 2, rank) - empty_func = tvm_tirx.get_global_func("runtime.disco.nvshmem.empty") + empty_func = tvm_ffi.get_global_func("runtime.disco.nvshmem.empty") pages = empty_func( Shape((num_layers, num_pages, 2, num_kv_heads, page_size, head_dim)), "float16", dev ) @@ -89,7 +90,7 @@ def test_kv_transfer_without_disco(): remote_position_map = tvm_tirx.runtime.tensor(remote_position_map_np, dev) remote_tp_group_pe_offset_np = np.array([1] * len(position_map_array), dtype=np.int32) remote_tp_group_pe_offset = tvm_tirx.runtime.tensor(remote_tp_group_pe_offset_np, dev) - transfer_func = tvm_tirx.get_global_func("nvshmem.KVTransfer") + transfer_func = tvm_ffi.get_global_func("nvshmem.KVTransfer") layer_view = pages._create_view( [num_pages, 2, num_kv_heads, page_size, head_dim], "float16", @@ -110,7 +111,7 @@ def test_kv_transfer_without_disco(): original_v = v_np[i] transferred_v = pages_np[layer_id, page_id, 1, :, offset_in_page, :] tvm_tirx.testing.assert_allclose(original_v, transferred_v) - finalize_func = tvm_tirx.get_global_func("runtime.disco.nvshmem.finalize_nvshmem") + finalize_func = tvm_ffi.get_global_func("runtime.disco.nvshmem.finalize_nvshmem") finalize_func() comm.Barrier() @@ -122,14 +123,14 @@ def test_kv_transfer_page_to_page_without_disco(): layer_id = 1 dev = tvm_tirx.cuda(rank) if rank == 0: - f_init_nvshmem_uid = tvm_tirx.get_global_func("runtime.disco.nvshmem.init_nvshmem_uid") + f_init_nvshmem_uid = tvm_ffi.get_global_func("runtime.disco.nvshmem.init_nvshmem_uid") uid = f_init_nvshmem_uid() else: uid = None uid = comm.bcast(uid, root=0) - init_func = tvm_tirx.get_global_func("runtime.disco.nvshmem.init_nvshmem") + init_func = tvm_ffi.get_global_func("runtime.disco.nvshmem.init_nvshmem") init_func(uid, 2, rank) - empty_func = tvm_tirx.get_global_func("runtime.disco.nvshmem.empty") + empty_func = tvm_ffi.get_global_func("runtime.disco.nvshmem.empty") pages = empty_func( Shape((num_layers, num_pages, 2, num_kv_heads, page_size, head_dim)), "float16", dev ) @@ -149,7 +150,7 @@ def test_kv_transfer_page_to_page_without_disco(): [1] * len(rank_0_position_map_array), dtype=np.int32 ) remote_tp_group_pe_offset = tvm_tirx.runtime.tensor(remote_tp_group_pe_offset_np, dev) - transfer_func = tvm_tirx.get_global_func("nvshmem.KVTransferPageToPage") + transfer_func = tvm_ffi.get_global_func("nvshmem.KVTransferPageToPage") layer_view = pages._create_view( [num_pages, 2, num_kv_heads, page_size, head_dim], "float16", @@ -177,7 +178,7 @@ def test_kv_transfer_page_to_page_without_disco(): rank_0_entry = pages_np[layer_id, rank_0_page_id, :, :, rank_0_offset_in_page, :] transferred_entry = new_pages_np[layer_id, page_id, :, :, offset_in_page, :] tvm_tirx.testing.assert_allclose(rank_0_entry, transferred_entry) - finalize_func = tvm_tirx.get_global_func("runtime.disco.nvshmem.finalize_nvshmem") + finalize_func = tvm_ffi.get_global_func("runtime.disco.nvshmem.finalize_nvshmem") finalize_func() comm.Barrier() @@ -188,7 +189,7 @@ def test_kv_transfer_with_disco(): comm, rank = get_comm_rank() layer_id = 1 if rank == 0: - f_init_nvshmem_uid = tvm_tirx.get_global_func("runtime.disco.nvshmem.init_nvshmem_uid") + f_init_nvshmem_uid = tvm_ffi.get_global_func("runtime.disco.nvshmem.init_nvshmem_uid") uid = f_init_nvshmem_uid() else: uid = None diff --git a/tests/python/relax/script/test_relax_script_printer.py b/tests/python/relax/script/test_relax_script_printer.py index d9dafa353644..08127d17798d 100644 --- a/tests/python/relax/script/test_relax_script_printer.py +++ b/tests/python/relax/script/test_relax_script_printer.py @@ -20,6 +20,7 @@ from __future__ import annotations import pytest +import tvm_ffi from tvm_ffi.access_path import AccessPath import tvm @@ -639,7 +640,7 @@ def func(n: T.int64) -> T.int64: def test_data_type_imm(): - obj = tvm_tirx.ir.DataTypeImm(tvm_tirx.DataType("float32")) + obj = tvm_tirx.ir.DataTypeImm(tvm_ffi.dtype("float32")) _assert_print_lines(obj, 'T.dtype("float32")') diff --git a/tests/python/relax/test_ast_printer.py b/tests/python/relax/test_ast_printer.py index 2d7107336af2..5ee48a4b8211 100644 --- a/tests/python/relax/test_ast_printer.py +++ b/tests/python/relax/test_ast_printer.py @@ -19,6 +19,7 @@ from functools import partial import numpy as np +import tvm_ffi import tvm_tirx import tvm_tirx.testing @@ -694,7 +695,7 @@ def test_string_imm(): def test_datatype_imm(): - data_type_imm = tvm_tirx.ir.DataTypeImm(tvm_tirx.DataType("int32")) + data_type_imm = tvm_tirx.ir.DataTypeImm(tvm_ffi.dtype("int32")) data_type_str = strip_whitespace(dump_ast(data_type_imm)) assert data_type_str == strip_whitespace( """ diff --git a/tests/python/relax/test_base_py_module_symbolic_shape.py b/tests/python/relax/test_base_py_module_symbolic_shape.py index 124be2b1637b..f054bc5fb8df 100644 --- a/tests/python/relax/test_base_py_module_symbolic_shape.py +++ b/tests/python/relax/test_base_py_module_symbolic_shape.py @@ -18,6 +18,7 @@ import numpy as np import pytest +import tvm_ffi import tvm_tirx from tvm import relax @@ -254,7 +255,7 @@ def test_base_py_module_call_dps_packed_symbolic(): """Test call_dps_packed with symbolic shapes.""" try: # Register a simple test function - @tvm_tirx.register_global_func("test_add_packed") + @tvm_ffi.register_global_func("test_add_packed") def test_add_packed(a, b, out): """Add two tensors element-wise.""" a_np = a.numpy() @@ -283,7 +284,7 @@ def test_base_py_module_call_dps_packed_multiple_args(): """Test call_dps_packed with multiple arguments and symbolic shapes.""" try: # Register a function that takes multiple arguments - @tvm_tirx.register_global_func("test_matmul_packed") + @tvm_ffi.register_global_func("test_matmul_packed") def test_matmul_packed(a, b, out): """Matrix multiplication.""" a_np = a.numpy() @@ -312,7 +313,7 @@ def test_base_py_module_call_dps_packed_scalar_args(): """Test call_dps_packed with scalar arguments and symbolic shapes.""" try: # Register a function that takes scalar arguments - @tvm_tirx.register_global_func("test_add_scalar_packed") + @tvm_ffi.register_global_func("test_add_scalar_packed") def test_add_scalar_packed(x, scalar, out): """Add scalar to tensor.""" x_np = x.numpy() diff --git a/tests/python/relax/test_blockbuilder_core.py b/tests/python/relax/test_blockbuilder_core.py index fc1c34aa2e88..14741faaba2d 100644 --- a/tests/python/relax/test_blockbuilder_core.py +++ b/tests/python/relax/test_blockbuilder_core.py @@ -19,6 +19,7 @@ # The test here do not depend on tvmscript to cover most basic features import pytest +import tvm_ffi import tvm import tvm.contrib.cblas @@ -37,7 +38,7 @@ @pytest.fixture(scope="module") def register_nop(): - @tvm_tirx.register_global_func("test.blockbuilder.nop") + @tvm_ffi.register_global_func("test.blockbuilder.nop") def nop(): pass diff --git a/tests/python/relax/test_codegen_cutlass.py b/tests/python/relax/test_codegen_cutlass.py index 8e4fa1dc7123..d60217605480 100644 --- a/tests/python/relax/test_codegen_cutlass.py +++ b/tests/python/relax/test_codegen_cutlass.py @@ -17,6 +17,7 @@ # ruff: noqa: E501, F841, RUF005 import numpy as np import pytest +import tvm_ffi import tvm import tvm_tirx @@ -2174,7 +2175,7 @@ def main( def _test_batched_var_len_attention( mod, seq_lens, num_head, num_kv_head, head_size, window_size=None ): - if not tvm_tirx.get_global_func("tvm.contrib.thrust.sum_scan", True): + if not tvm_ffi.get_global_func("tvm.contrib.thrust.sum_scan", True): return hidden_size = num_head * head_size diff --git a/tests/python/relax/test_codegen_dnnl.py b/tests/python/relax/test_codegen_dnnl.py index c2d694eac7a1..6380a17cff28 100644 --- a/tests/python/relax/test_codegen_dnnl.py +++ b/tests/python/relax/test_codegen_dnnl.py @@ -16,6 +16,7 @@ # under the License. import numpy as np import pytest +import tvm_ffi import tvm import tvm_tirx @@ -42,7 +43,7 @@ def main( return conv2 -has_dnnl = tvm_tirx.get_global_func("relax.ext.dnnl", True) +has_dnnl = tvm_ffi.get_global_func("relax.ext.dnnl", True) dnnl_enabled = pytest.mark.skipif( not has_dnnl, diff --git a/tests/python/relax/test_codegen_tensorrt.py b/tests/python/relax/test_codegen_tensorrt.py index 0821283d0538..8e56b10d8e49 100644 --- a/tests/python/relax/test_codegen_tensorrt.py +++ b/tests/python/relax/test_codegen_tensorrt.py @@ -16,6 +16,7 @@ # under the License. import numpy as np import pytest +import tvm_ffi import tvm import tvm_tirx @@ -45,8 +46,8 @@ def main( return out -has_tensorrt = tvm_tirx.get_global_func("relax.ext.tensorrt", True) -env_checker_runtime = tvm_tirx.get_global_func("relax.is_tensorrt_runtime_enabled", True) +has_tensorrt = tvm_ffi.get_global_func("relax.ext.tensorrt", True) +env_checker_runtime = tvm_ffi.get_global_func("relax.is_tensorrt_runtime_enabled", True) requires_tensorrt_codegen = pytest.mark.skipif( not has_tensorrt, diff --git a/tests/python/relax/test_contrib_vllm.py b/tests/python/relax/test_contrib_vllm.py index 1833b4f67900..8b9490648474 100644 --- a/tests/python/relax/test_contrib_vllm.py +++ b/tests/python/relax/test_contrib_vllm.py @@ -28,7 +28,7 @@ from tvm_tirx.script import tirx as T from tvm_tirx.testing import env -has_vllm = tvm_tirx.get_global_func("tvm.contrib.vllm.single_query_cached_kv_attention", True) +has_vllm = tvm_ffi.get_global_func("tvm.contrib.vllm.single_query_cached_kv_attention", True) vllm_enabled = pytest.mark.skipif( not has_vllm, @@ -767,10 +767,8 @@ def test_reconstruct_from_cache(): num_tokens = 8 num_blocks = 1 - reshape_and_cache_func = tvm_tirx.get_global_func("tvm.contrib.vllm.reshape_and_cache") - reconstruct_from_cache_func = tvm_tirx.get_global_func( - "tvm.contrib.vllm.reconstruct_from_cache" - ) + reshape_and_cache_func = tvm_ffi.get_global_func("tvm.contrib.vllm.reshape_and_cache") + reconstruct_from_cache_func = tvm_ffi.get_global_func("tvm.contrib.vllm.reconstruct_from_cache") def run_and_check(): dev = tvm_tirx.cuda(0) diff --git a/tests/python/relax/test_dlpack_integration.py b/tests/python/relax/test_dlpack_integration.py index 740fe025191b..a91c6deb31d7 100644 --- a/tests/python/relax/test_dlpack_integration.py +++ b/tests/python/relax/test_dlpack_integration.py @@ -26,6 +26,7 @@ import pytest import torch +import tvm_ffi import tvm_tirx import tvm_tirx.testing @@ -39,7 +40,7 @@ class TestDLPackIntegration: def test_dlpack_pytorch_to_tvm_conversion(self): pytorch_tensor = torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0], dtype=torch.float32) - tvm_tensor = tvm_tirx.runtime.from_dlpack(pytorch_tensor) + tvm_tensor = tvm_ffi.from_dlpack(pytorch_tensor) assert isinstance(tvm_tensor, tvm_tirx.runtime.Tensor) assert tvm_tensor.shape == pytorch_tensor.shape @@ -57,7 +58,7 @@ def run_and_check(): [1.0, 2.0, 3.0, 4.0, 5.0], dtype=torch.float32, device="cuda" ) - tvm_tensor = tvm_tirx.runtime.from_dlpack(pytorch_tensor) + tvm_tensor = tvm_ffi.from_dlpack(pytorch_tensor) assert isinstance(tvm_tensor, tvm_tirx.runtime.Tensor) assert tvm_tensor.shape == pytorch_tensor.shape @@ -119,7 +120,7 @@ def test_dlpack_roundtrip_conversion(self): original_tensor = torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0], dtype=torch.float32) # Convert to TVM - tvm_tensor = tvm_tirx.runtime.from_dlpack(original_tensor) + tvm_tensor = tvm_ffi.from_dlpack(original_tensor) # Convert back to PyTorch result_tensor = torch.from_dlpack(tvm_tensor) @@ -143,7 +144,7 @@ def test_dlpack_different_data_types(self): pytorch_tensor = torch.tensor([1, 2, 3], dtype=torch_dtype) # Convert to TVM - tvm_tensor = tvm_tirx.runtime.from_dlpack(pytorch_tensor) + tvm_tensor = tvm_ffi.from_dlpack(pytorch_tensor) # Convert back to PyTorch result_tensor = torch.from_dlpack(tvm_tensor) @@ -166,7 +167,7 @@ def test_dlpack_different_shapes(self): pytorch_tensor = torch.randn(shape, dtype=torch.float32) # Convert to TVM - tvm_tensor = tvm_tirx.runtime.from_dlpack(pytorch_tensor) + tvm_tensor = tvm_ffi.from_dlpack(pytorch_tensor) # Convert back to PyTorch result_tensor = torch.from_dlpack(tvm_tensor) @@ -182,7 +183,7 @@ def test_dlpack_functionality_verification(self): pytorch_tensor = torch.randn(size, dtype=torch.float32) # Test DLPack conversion - tvm_tensor_dlpack = tvm_tirx.runtime.from_dlpack(pytorch_tensor) + tvm_tensor_dlpack = tvm_ffi.from_dlpack(pytorch_tensor) # Test numpy conversion numpy_array = pytorch_tensor.detach().cpu().numpy() @@ -235,7 +236,7 @@ def test_dlpack_device_consistency(self): """Test DLPack conversion maintains device consistency.""" # Test CPU tensor cpu_tensor = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32) - cpu_tvm = tvm_tirx.runtime.from_dlpack(cpu_tensor) + cpu_tvm = tvm_ffi.from_dlpack(cpu_tensor) cpu_result = torch.from_dlpack(cpu_tvm) assert cpu_result.device.type == "cpu" @@ -250,7 +251,7 @@ def test_dlpack_memory_sharing(self): pytorch_tensor = torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0], dtype=torch.float32) # Convert to TVM - tvm_tensor = tvm_tirx.runtime.from_dlpack(pytorch_tensor) + tvm_tensor = tvm_ffi.from_dlpack(pytorch_tensor) # Modify the original tensor pytorch_tensor[0] = 10.0 @@ -269,7 +270,7 @@ def test_dlpack_batch_operations(self): pytorch_tensors = [torch.randn(5, dtype=torch.float32) for _ in range(batch_size)] # Convert all to TVM - tvm_tensors = [tvm_tirx.runtime.from_dlpack(t) for t in pytorch_tensors] + tvm_tensors = [tvm_ffi.from_dlpack(t) for t in pytorch_tensors] # Convert all back to PyTorch result_tensors = [torch.from_dlpack(t) for t in tvm_tensors] @@ -282,7 +283,7 @@ def test_dlpack_edge_cases(self): """Test DLPack conversion with edge cases.""" # Empty tensor empty_tensor = torch.tensor([], dtype=torch.float32) - empty_tvm = tvm_tirx.runtime.from_dlpack(empty_tensor) + empty_tvm = tvm_ffi.from_dlpack(empty_tensor) empty_result = torch.from_dlpack(empty_tvm) assert empty_result.shape == empty_tensor.shape @@ -290,7 +291,7 @@ def test_dlpack_edge_cases(self): # Single element tensor single_tensor = torch.tensor([42.0], dtype=torch.float32) - single_tvm = tvm_tirx.runtime.from_dlpack(single_tensor) + single_tvm = tvm_ffi.from_dlpack(single_tensor) single_result = torch.from_dlpack(single_tvm) assert single_result.shape == single_tensor.shape diff --git a/tests/python/relax/test_expr.py b/tests/python/relax/test_expr.py index ca8dfbbf4b08..3e0b30642a78 100644 --- a/tests/python/relax/test_expr.py +++ b/tests/python/relax/test_expr.py @@ -348,8 +348,8 @@ def test_string_imm(): def test_datatype_imm(): - d0 = tvm_tirx.ir.DataTypeImm(tvm_tirx.DataType("int32")) - d1 = tvm_tirx.ir.DataTypeImm(tvm_tirx.DataType("int32")) + d0 = tvm_tirx.ir.DataTypeImm(tvm_ffi.dtype("int32")) + d1 = tvm_tirx.ir.DataTypeImm(tvm_ffi.dtype("int32")) assert d0.value == "int32" _check_equal(d0, d1) _check_json_roundtrip(d0) diff --git a/tests/python/relax/test_frontend_nn_debug.py b/tests/python/relax/test_frontend_nn_debug.py index 7d08b907118d..9d64abed6d60 100644 --- a/tests/python/relax/test_frontend_nn_debug.py +++ b/tests/python/relax/test_frontend_nn_debug.py @@ -16,9 +16,8 @@ # under the License. # pylint: disable=missing-docstring import torch +import tvm_ffi -import tvm_tirx -import tvm_tirx.testing from tvm.relax.frontend import nn from tvm.relax.frontend.nn import op, spec from tvm_tirx import tirx @@ -43,7 +42,7 @@ def forward(self, x: nn.Tensor): # pylint: disable=invalid-name def test_debug_func(): - @tvm_tirx.register_global_func("testing.relax.frontend.nn.test_debug_func") + @tvm_ffi.register_global_func("testing.relax.frontend.nn.test_debug_func") def _debug( # pylint: disable=too-many-arguments lineno: str, tensor: Tensor, diff --git a/tests/python/relax/test_frontend_nn_op.py b/tests/python/relax/test_frontend_nn_op.py index d8ef0d2e227c..e8483845e478 100644 --- a/tests/python/relax/test_frontend_nn_op.py +++ b/tests/python/relax/test_frontend_nn_op.py @@ -18,6 +18,7 @@ # ruff: noqa: E501, F841 import numpy as np import pytest +import tvm_ffi import tvm import tvm_tirx @@ -909,7 +910,7 @@ def test(q: R.Tensor((1, 1, 16, 8), dtype="float32"), k: R.Tensor((64, 16, 8), d def test_empty(): - @tvm_tirx.register_global_func("test_empty_assert", override=True) + @tvm_ffi.register_global_func("test_empty_assert", override=True) def test_empty_assert(_lineo, x): assert x.shape == (10, 10) assert x.dtype == "float32" diff --git a/tests/python/relax/test_group_gemm_flashinfer.py b/tests/python/relax/test_group_gemm_flashinfer.py index 43fbd942c64a..9af22041c11d 100644 --- a/tests/python/relax/test_group_gemm_flashinfer.py +++ b/tests/python/relax/test_group_gemm_flashinfer.py @@ -23,6 +23,7 @@ import numpy as np import pytest import torch +import tvm_ffi import tvm import tvm_tirx @@ -54,7 +55,7 @@ def has_flashinfer(): def has_cutlass(): """Check if CUTLASS is available for SM90+ operations""" - if not tvm_tirx.get_global_func("device_api.cuda", True): + if not tvm_ffi.get_global_func("device_api.cuda", True): return False try: import pynvml # pylint: disable=import-outside-toplevel @@ -317,17 +318,17 @@ def generate_test_data( a_quantized.view(torch.uint8).cpu().numpy().view(fp8_dtype), device=device ) else: - a_tvm = tvm_tirx.runtime.from_dlpack(a_quantized) + a_tvm = tvm_ffi.from_dlpack(a_quantized) if dtype_b == "float8_e4m3fn": b_tvm = tvm_tirx.runtime.tensor( b_quantized.view(torch.uint8).cpu().numpy().view(fp8_dtype), device=device ) else: - b_tvm = tvm_tirx.runtime.from_dlpack(b_quantized) + b_tvm = tvm_ffi.from_dlpack(b_quantized) - scale_a_tvm = tvm_tirx.runtime.from_dlpack(scale_a) - scale_b_tvm = tvm_tirx.runtime.from_dlpack(scale_b) + scale_a_tvm = tvm_ffi.from_dlpack(scale_a) + scale_b_tvm = tvm_ffi.from_dlpack(scale_b) # Create m_indptr for grouped operation m_indptr = torch.tensor(cum_m, device=torch_device, dtype=torch.int32) diff --git a/tests/python/relax/test_nn_llm_prefill_parallel.py b/tests/python/relax/test_nn_llm_prefill_parallel.py index 7e437e1568ab..eed6d3860a1a 100644 --- a/tests/python/relax/test_nn_llm_prefill_parallel.py +++ b/tests/python/relax/test_nn_llm_prefill_parallel.py @@ -16,6 +16,7 @@ # under the License. import numpy as np import pytest +import tvm_ffi import tvm import tvm_tirx @@ -57,7 +58,7 @@ def rope(values, positions): def test_parallel_prefill(device_type, layout, rotary, causal, first_kv_length): if not tvm_tirx.testing.device_enabled(device_type): pytest.skip(f"{device_type} not enabled") - device = tvm_tirx.device(device_type) + device = tvm_ffi.device(device_type) target = tvm_tirx.target.Target.from_device(device) rng = np.random.default_rng(127) lengths = [(7, first_kv_length), (17, 513)] diff --git a/tests/python/relax/test_op_inspect.py b/tests/python/relax/test_op_inspect.py index fc2967eab7cd..9d061fcf7087 100644 --- a/tests/python/relax/test_op_inspect.py +++ b/tests/python/relax/test_op_inspect.py @@ -67,7 +67,7 @@ def main(A: R.Tensor): arg = tvm_tirx.runtime.empty([16], dtype) res = vm["main"](arg) - expected_type_code = tvm_tirx.runtime.DataType(dtype).type_code + expected_type_code = tvm_ffi.dtype(dtype).type_code assert res == expected_type_code @@ -84,7 +84,7 @@ def main(A: R.Tensor): arg = tvm_tirx.runtime.empty([16], dtype) res = vm["main"](arg) - expected_type_bits = tvm_tirx.runtime.DataType(dtype).bits + expected_type_bits = tvm_ffi.dtype(dtype).bits assert res == expected_type_bits @@ -101,7 +101,7 @@ def main(A: R.Tensor): arg = tvm_tirx.runtime.empty([16], dtype) res = vm["main"](arg) - expected_type_lanes = tvm_tirx.runtime.DataType(dtype).lanes + expected_type_lanes = tvm_ffi.dtype(dtype).lanes assert res == expected_type_lanes @@ -204,7 +204,7 @@ def main(A: R.Tensor): def test_elem_offset(elem_offset, dtype): - tvm_dtype = tvm_tirx.runtime.DataType(dtype) + tvm_dtype = tvm_ffi.dtype(dtype) backing_shape = [64, 64] view_shape = [16, 16] diff --git a/tests/python/relax/test_op_misc.py b/tests/python/relax/test_op_misc.py index 7006085344ec..06b8ccb00279 100644 --- a/tests/python/relax/test_op_misc.py +++ b/tests/python/relax/test_op_misc.py @@ -25,7 +25,7 @@ from tvm_tirx.script import tirx as T -@tvm_tirx.register_global_func("test.op.identity", override=True) +@tvm_ffi.register_global_func("test.op.identity", override=True) def identity_packed(a): return tvm_tirx.runtime.tensor(a.numpy()) diff --git a/tests/python/relax/test_relax_operators.py b/tests/python/relax/test_relax_operators.py index de07439cc653..ccde7417343e 100644 --- a/tests/python/relax/test_relax_operators.py +++ b/tests/python/relax/test_relax_operators.py @@ -338,7 +338,7 @@ def pure_copy(x: R.Tensor((3, 4), "float32")): ) return z - @tvm_tirx.register_global_func("test.inplace.add", override=True) + @tvm_ffi.register_global_func("test.inplace.add", override=True) def inplace_add(a, b): arr_a = a.numpy() arr_b = b.numpy() @@ -370,7 +370,7 @@ def inplace_add(x: R.Tensor((3, 4), "float32"), y: R.Tensor((3, 4), "float32")): assert result == tvm_arr_a assert (result.numpy() == sum).all() - @tvm_tirx.register_global_func("test.inplace.tuple_add", override=True) + @tvm_ffi.register_global_func("test.inplace.tuple_add", override=True) def inplace_tuple_add(a, b): arr_a = a.numpy() arr_b = b.numpy() @@ -379,7 +379,7 @@ def inplace_tuple_add(a, b): for j in range(len(arr_a[i])): arr_a[i][j] = arr_a[i][j] + arr_b[i][j] a.copyfrom(arr_a) - return tvm_tirx.runtime.convert([a, c]) + return tvm_ffi.convert([a, c]) @tvm_tirx.script.ir_module class CallInplaceTuple: @@ -601,8 +601,8 @@ def func(x: R.Tensor((3,), "float32")): def test_call_py_func_numpy_execution(): - register = tvm_tirx.get_global_func("vm.builtin.register_py_func") - unregister = tvm_tirx.get_global_func("vm.builtin.unregister_py_func") + register = tvm_ffi.get_global_func("vm.builtin.register_py_func") + unregister = tvm_ffi.get_global_func("vm.builtin.unregister_py_func") def numpy_pair(x): data = x.numpy() @@ -670,7 +670,7 @@ def torch_sigmoid(x): result = torch.sigmoid(x_torch) return tvm_tirx.runtime.tensor(result.numpy()) - register_func = tvm_tirx.get_global_func("vm.builtin.register_py_func") + register_func = tvm_ffi.get_global_func("vm.builtin.register_py_func") register_func("torch_relu", torch_relu) register_func("torch_sigmoid", torch_sigmoid) @@ -702,7 +702,7 @@ def multiple_calls(x: R.Tensor((2,), "float32")): expected2 = 1.0 / (1.0 + np.exp(-np.maximum(y_data, 0.0))) assert (result2.numpy() == expected2).all() - unregister_func = tvm_tirx.get_global_func("vm.builtin.unregister_py_func") + unregister_func = tvm_ffi.get_global_func("vm.builtin.unregister_py_func") unregister_func("torch_relu") unregister_func("torch_sigmoid") @@ -711,8 +711,8 @@ def test_py_func_registry_is_scoped_to_its_module(): """A module's finalizer must drop its own registrations and nothing else.""" from tvm.relax.base_py_module import BasePyModule - get_func = tvm_tirx.get_global_func("vm.builtin.get_py_func") - tvm_tirx.get_global_func("vm.builtin.register_py_func")("registry_probe", lambda x: x) + get_func = tvm_ffi.get_global_func("vm.builtin.get_py_func") + tvm_ffi.get_global_func("vm.builtin.register_py_func")("registry_probe", lambda x: x) # __new__ skips __init__'s JIT compilation; only the registration matters here. module = BasePyModule.__new__(BasePyModule) @@ -728,7 +728,7 @@ def test_py_func_registry_is_scoped_to_its_module(): assert get_func("registry_probe") is not None, "another owner's function was dropped" with pytest.raises(RuntimeError, match="not found in registry"): get_func("registry_owned") - tvm_tirx.get_global_func("vm.builtin.unregister_py_func")("registry_probe") + tvm_ffi.get_global_func("vm.builtin.unregister_py_func")("registry_probe") def test_op_to_device(): diff --git a/tests/python/relax/test_relax_to_pyfunc_converter.py b/tests/python/relax/test_relax_to_pyfunc_converter.py index bc8c08d69e5b..1a4cb544a92a 100644 --- a/tests/python/relax/test_relax_to_pyfunc_converter.py +++ b/tests/python/relax/test_relax_to_pyfunc_converter.py @@ -25,6 +25,7 @@ import pytest import torch import torch.nn.functional as F +import tvm_ffi import tvm import tvm_tirx @@ -213,7 +214,7 @@ def mock_softmax(x, axis): return x # Register the function globally - tvm_tirx.register_global_func("my_softmax", mock_softmax) + tvm_ffi.register_global_func("my_softmax", mock_softmax) class TestRelaxToPyFuncConverter: @@ -972,7 +973,7 @@ def test_packed_function_with_primvalue_args(self): def test_packed_func(x, axis): return x # Simple identity function - tvm_tirx.register_global_func("test_packed_func", test_packed_func) + tvm_ffi.register_global_func("test_packed_func", test_packed_func) @I.ir_module class PackedFuncTestModule: diff --git a/tests/python/relax/test_runtime_builtin.py b/tests/python/relax/test_runtime_builtin.py index 2f58b27741c7..b8b8901dea69 100644 --- a/tests/python/relax/test_runtime_builtin.py +++ b/tests/python/relax/test_runtime_builtin.py @@ -30,7 +30,7 @@ def test_make_shape(): MK = MakeShapeCode - make_shape = tvm_tirx.get_global_func("vm.builtin.make_shape") + make_shape = tvm_ffi.get_global_func("vm.builtin.make_shape") heap = tvm_tirx.runtime.tensor(np.arange(10).astype("int64")) s = make_shape(heap, 3, MK.USE_IMM, 10, MK.LOAD_SHAPE, 0, MK.LOAD_SHAPE, 2) @@ -39,7 +39,7 @@ def test_make_shape(): def test_match_shape(): MS = MatchShapeCode - match_shape = tvm_tirx.get_global_func("vm.builtin.match_shape") + match_shape = tvm_ffi.get_global_func("vm.builtin.match_shape") heap = tvm_tirx.runtime.tensor(np.zeros(10).astype("int64")) assert heap.numpy()[2] == 0 @@ -72,7 +72,7 @@ def test_match_shape(): def test_check_shape_info(): - check_shape_info = tvm_tirx.get_global_func("vm.builtin.check_shape_info") + check_shape_info = tvm_ffi.get_global_func("vm.builtin.check_shape_info") s = tvm_ffi.Shape([1, 2, 3]) check_shape_info(s, 3, "") @@ -88,7 +88,7 @@ def test_check_shape_info(): def test_check_tensor_info(): - check_tensor_info = tvm_tirx.get_global_func("vm.builtin.check_tensor_info") + check_tensor_info = tvm_ffi.get_global_func("vm.builtin.check_tensor_info") x = tvm_tirx.runtime.tensor(np.zeros((2, 3)).astype("int32")) check_tensor_info(x, 2, "int32", "") @@ -118,9 +118,9 @@ def test_check_tensor_info(): def test_check_tuple_info(): - check_tuple_info = tvm_tirx.get_global_func("vm.builtin.check_tuple_info") + check_tuple_info = tvm_ffi.get_global_func("vm.builtin.check_tuple_info") x = tvm_tirx.runtime.tensor(np.zeros((2, 3)).astype("int32")) - t = tvm_tirx.runtime.convert([x, x, x]) + t = tvm_ffi.convert([x, x, x]) check_tuple_info(t, 3, "") @@ -134,8 +134,8 @@ def test_check_tuple_info(): def test_check_func_info(): - check_func_info = tvm_tirx.get_global_func("vm.builtin.check_func_info") - f = tvm_tirx.runtime.convert(lambda x: x) + check_func_info = tvm_ffi.get_global_func("vm.builtin.check_func_info") + f = tvm_ffi.convert(lambda x: x) x = tvm_tirx.runtime.tensor(np.zeros((2, 3)).astype("int32")) check_func_info(f, "") @@ -146,19 +146,19 @@ def test_check_func_info(): def test_tuple_getitem(): - tuple_getitem = tvm_tirx.get_global_func("vm.builtin.tuple_getitem") + tuple_getitem = tvm_ffi.get_global_func("vm.builtin.tuple_getitem") x = tvm_tirx.runtime.tensor(np.zeros((2, 3)).astype("int32")) y = tvm_tirx.runtime.tensor(np.zeros((2, 3)).astype("int32")) - t = tvm_tirx.runtime.convert([x, y]) + t = tvm_ffi.convert([x, y]) assert tuple_getitem(t, 0) == x assert tuple_getitem(t, 1) == y def test_attention_kv_cache(): - fcreate = tvm_tirx.get_global_func("vm.builtin.attention_kv_cache_create") - fappend = tvm_tirx.get_global_func("vm.builtin.attention_kv_cache_append") - fview = tvm_tirx.get_global_func("vm.builtin.attention_kv_cache_view") + fcreate = tvm_ffi.get_global_func("vm.builtin.attention_kv_cache_create") + fappend = tvm_ffi.get_global_func("vm.builtin.attention_kv_cache_append") + fview = tvm_ffi.get_global_func("vm.builtin.attention_kv_cache_view") cache = fcreate(tvm_tirx.runtime.empty((1, 2), dtype="int32"), tvm_ffi.Shape([2, 2]), 0) num_steps = 2 @@ -172,8 +172,8 @@ def test_attention_kv_cache(): def test_tensor_cache(): - fload = tvm_tirx.get_global_func("vm.builtin.tensor_cache.load") - fget_params = tvm_tirx.get_global_func("vm.builtin.param_array_from_cache") + fload = tvm_ffi.get_global_func("vm.builtin.tensor_cache.load") + fget_params = tvm_ffi.get_global_func("vm.builtin.param_array_from_cache") param_dict = { "x_0": np.array([1, 2, 3], dtype="int32"), @@ -192,8 +192,8 @@ def test_tensor_cache(): def test_tensor_cache_update(): - fload = tvm_tirx.get_global_func("vm.builtin.tensor_cache.load") - fget_params = tvm_tirx.get_global_func("vm.builtin.param_array_from_cache") + fload = tvm_ffi.get_global_func("vm.builtin.tensor_cache.load") + fget_params = tvm_ffi.get_global_func("vm.builtin.param_array_from_cache") param_dict = { "x_0": np.array([1, 2, 3], dtype="int32"), @@ -217,9 +217,9 @@ def test_tensor_cache_update(): def test_attention_kv_cache_window_override(): - fcreate = tvm_tirx.get_global_func("vm.builtin.attention_kv_cache_create") - foverride = tvm_tirx.get_global_func("vm.builtin.attention_kv_cache_window_override") - fview = tvm_tirx.get_global_func("vm.builtin.attention_kv_cache_view") + fcreate = tvm_ffi.get_global_func("vm.builtin.attention_kv_cache_create") + foverride = tvm_ffi.get_global_func("vm.builtin.attention_kv_cache_window_override") + fview = tvm_ffi.get_global_func("vm.builtin.attention_kv_cache_view") current_pos = 4 cache = fcreate( @@ -246,9 +246,9 @@ def test_attention_kv_cache_window_override(): def test_attention_kv_cache_window_override_with_sinks(): - fcreate = tvm_tirx.get_global_func("vm.builtin.attention_kv_cache_create") - foverride = tvm_tirx.get_global_func("vm.builtin.attention_kv_cache_window_override_with_sinks") - fview = tvm_tirx.get_global_func("vm.builtin.attention_kv_cache_view") + fcreate = tvm_ffi.get_global_func("vm.builtin.attention_kv_cache_create") + foverride = tvm_ffi.get_global_func("vm.builtin.attention_kv_cache_window_override_with_sinks") + fview = tvm_ffi.get_global_func("vm.builtin.attention_kv_cache_view") num_attention_sinks = 2 has_sink = False diff --git a/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_cpu.py b/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_cpu.py index b37b2c7a1d46..f1f5e0c6d7fe 100644 --- a/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_cpu.py +++ b/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_cpu.py @@ -113,35 +113,35 @@ def set_global_func(head_dim, dtype, layer_sliding_window_size=1024): global fattn_prefill_sliding_window, fattn_decode_sliding_window global fmerge_state, fsplit_rotary, fattention_rotary, fcopy_single_page, fcompact_copy - fclear = tvm_tirx.get_global_func("vm.builtin.kv_state_clear") - fadd_sequence = tvm_tirx.get_global_func("vm.builtin.kv_state_add_sequence") - fremove_sequence = tvm_tirx.get_global_func("vm.builtin.kv_state_remove_sequence") - ffork_sequence = tvm_tirx.get_global_func("vm.builtin.kv_state_fork_sequence") - fenable_sliding_window_for_seq = tvm_tirx.get_global_func( + fclear = tvm_ffi.get_global_func("vm.builtin.kv_state_clear") + fadd_sequence = tvm_ffi.get_global_func("vm.builtin.kv_state_add_sequence") + fremove_sequence = tvm_ffi.get_global_func("vm.builtin.kv_state_remove_sequence") + ffork_sequence = tvm_ffi.get_global_func("vm.builtin.kv_state_fork_sequence") + fenable_sliding_window_for_seq = tvm_ffi.get_global_func( "vm.builtin.attention_kv_cache_enable_sliding_window_for_seq" ) - fpopn = tvm_tirx.get_global_func("vm.builtin.kv_state_popn") - fbegin_forward = tvm_tirx.get_global_func("vm.builtin.kv_state_begin_forward") - fend_forward = tvm_tirx.get_global_func("vm.builtin.kv_state_end_forward") - fcommit_accepted_token_tree_nodes = tvm_tirx.get_global_func( + fpopn = tvm_ffi.get_global_func("vm.builtin.kv_state_popn") + fbegin_forward = tvm_ffi.get_global_func("vm.builtin.kv_state_begin_forward") + fend_forward = tvm_ffi.get_global_func("vm.builtin.kv_state_end_forward") + fcommit_accepted_token_tree_nodes = tvm_ffi.get_global_func( "vm.builtin.attention_kv_cache_commit_accepted_token_tree_nodes" ) - fattention_with_fuse_qkv = tvm_tirx.get_global_func( + fattention_with_fuse_qkv = tvm_ffi.get_global_func( "vm.builtin.attention_kv_cache_attention_with_fused_qkv" ) - fattention_with_shared_kv = tvm_tirx.get_global_func( + fattention_with_shared_kv = tvm_ffi.get_global_func( "vm.builtin.attention_kv_cache_attention_with_shared_kv" ) - fis_empty = tvm_tirx.get_global_func("vm.builtin.attention_kv_cache_empty") - fdebug_get_kv = tvm_tirx.get_global_func("vm.builtin.attention_kv_cache_debug_get_kv") - fget_checkpoint_metadata = tvm_tirx.get_global_func( + fis_empty = tvm_ffi.get_global_func("vm.builtin.attention_kv_cache_empty") + fdebug_get_kv = tvm_ffi.get_global_func("vm.builtin.attention_kv_cache_debug_get_kv") + fget_checkpoint_metadata = tvm_ffi.get_global_func( "vm.builtin.attention_kv_cache_get_checkpoint_metadata" ) - fexport_page_group = tvm_tirx.get_global_func("vm.builtin.attention_kv_cache_export_page_group") - fprepare_import = tvm_tirx.get_global_func("vm.builtin.attention_kv_cache_prepare_import") - fimport_page_group = tvm_tirx.get_global_func("vm.builtin.attention_kv_cache_import_page_group") - ffinish_import = tvm_tirx.get_global_func("vm.builtin.attention_kv_cache_finish_import") - fget_sequence_length = tvm_tirx.get_global_func( + fexport_page_group = tvm_ffi.get_global_func("vm.builtin.attention_kv_cache_export_page_group") + fprepare_import = tvm_ffi.get_global_func("vm.builtin.attention_kv_cache_prepare_import") + fimport_page_group = tvm_ffi.get_global_func("vm.builtin.attention_kv_cache_import_page_group") + ffinish_import = tvm_ffi.get_global_func("vm.builtin.attention_kv_cache_finish_import") + fget_sequence_length = tvm_ffi.get_global_func( "vm.builtin.attention_kv_cache_get_sequence_length" ) @@ -226,7 +226,7 @@ def create_kv_cache( layer_sliding_window_size=None, attn_kinds=None, ): - fcreate = tvm_tirx.get_global_func("vm.builtin.paged_attention_kv_cache_create") + fcreate = tvm_ffi.get_global_func("vm.builtin.paged_attention_kv_cache_create") cache_config = [ reserved_nseq, maximum_total_seq_length, diff --git a/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_flashinfer.py b/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_flashinfer.py index 23f33e50959b..5d020d710e32 100644 --- a/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_flashinfer.py +++ b/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_flashinfer.py @@ -98,17 +98,17 @@ def set_global_func(rope_mode: RopeMode, target): global fattention_merge_state, fsplit_rotary, fcopy_single_page global ftranspose_append, fcopy_cache, fcompact_copy - fclear = tvm_tirx.get_global_func("vm.builtin.kv_state_clear") - fadd_sequence = tvm_tirx.get_global_func("vm.builtin.kv_state_add_sequence") - fremove_sequence = tvm_tirx.get_global_func("vm.builtin.kv_state_remove_sequence") - ffork_sequence = tvm_tirx.get_global_func("vm.builtin.kv_state_fork_sequence") - fpopn = tvm_tirx.get_global_func("vm.builtin.kv_state_popn") - fbegin_forward = tvm_tirx.get_global_func("vm.builtin.kv_state_begin_forward") - fend_forward = tvm_tirx.get_global_func("vm.builtin.kv_state_end_forward") - fattention_with_fuse_qkv = tvm_tirx.get_global_func( + fclear = tvm_ffi.get_global_func("vm.builtin.kv_state_clear") + fadd_sequence = tvm_ffi.get_global_func("vm.builtin.kv_state_add_sequence") + fremove_sequence = tvm_ffi.get_global_func("vm.builtin.kv_state_remove_sequence") + ffork_sequence = tvm_ffi.get_global_func("vm.builtin.kv_state_fork_sequence") + fpopn = tvm_ffi.get_global_func("vm.builtin.kv_state_popn") + fbegin_forward = tvm_ffi.get_global_func("vm.builtin.kv_state_begin_forward") + fend_forward = tvm_ffi.get_global_func("vm.builtin.kv_state_end_forward") + fattention_with_fuse_qkv = tvm_ffi.get_global_func( "vm.builtin.attention_kv_cache_attention_with_fused_qkv" ) - fdebug_get_kv = tvm_tirx.get_global_func("vm.builtin.attention_kv_cache_debug_get_kv") + fdebug_get_kv = tvm_ffi.get_global_func("vm.builtin.attention_kv_cache_debug_get_kv") flashinfer_prefill_mod = relax.backend.cuda.flashinfer.gen_flashinfer_prefill_module( dtype_q=dtype, @@ -162,7 +162,7 @@ def set_global_func(rope_mode: RopeMode, target): def create_kv_cache(rope_mode): - fcreate = tvm_tirx.get_global_func("vm.builtin.paged_attention_kv_cache_create") + fcreate = tvm_ffi.get_global_func("vm.builtin.paged_attention_kv_cache_create") support_sliding_window = 0 cache = fcreate( tvm_ffi.Shape( diff --git a/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_metadata.py b/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_metadata.py index 8217c7684e57..52b475ea51f3 100644 --- a/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_metadata.py +++ b/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_metadata.py @@ -56,8 +56,8 @@ def create_kv_cache( prefill_chunk_size_value=prefill_chunk_size, rope_ext_factors=None, ): - fcreate = tvm_tirx.get_global_func("vm.builtin.paged_attention_kv_cache_create") - dummy_func = tvm_tirx.runtime.convert(_nop) + fcreate = tvm_ffi.get_global_func("vm.builtin.paged_attention_kv_cache_create") + dummy_func = tvm_ffi.convert(_nop) return fcreate( tvm_ffi.Shape( [ @@ -99,24 +99,24 @@ def create_kv_cache( def append_tokens(kv_cache, seq_id=0, append_length=page_size + 1): - fadd_sequence = tvm_tirx.get_global_func("vm.builtin.kv_state_add_sequence") - fbegin_forward = tvm_tirx.get_global_func("vm.builtin.kv_state_begin_forward") - fend_forward = tvm_tirx.get_global_func("vm.builtin.kv_state_end_forward") + fadd_sequence = tvm_ffi.get_global_func("vm.builtin.kv_state_add_sequence") + fbegin_forward = tvm_ffi.get_global_func("vm.builtin.kv_state_begin_forward") + fend_forward = tvm_ffi.get_global_func("vm.builtin.kv_state_end_forward") fadd_sequence(kv_cache, seq_id) fbegin_forward(kv_cache, Shape([seq_id]), Shape([append_length]), None) fend_forward(kv_cache) def test_checkpoint_metadata_reports_layout_pages_and_groups(): - fget_checkpoint_metadata = tvm_tirx.get_global_func( + fget_checkpoint_metadata = tvm_ffi.get_global_func( "vm.builtin.attention_kv_cache_get_checkpoint_metadata" ) - fget_layout_hash = tvm_tirx.get_global_func("vm.builtin.attention_kv_cache_get_layout_hash") - fexport_page_group = tvm_tirx.get_global_func("vm.builtin.attention_kv_cache_export_page_group") - fprepare_import = tvm_tirx.get_global_func("vm.builtin.attention_kv_cache_prepare_import") - fimport_page_group = tvm_tirx.get_global_func("vm.builtin.attention_kv_cache_import_page_group") - ffinish_import = tvm_tirx.get_global_func("vm.builtin.attention_kv_cache_finish_import") - fget_sequence_length = tvm_tirx.get_global_func( + fget_layout_hash = tvm_ffi.get_global_func("vm.builtin.attention_kv_cache_get_layout_hash") + fexport_page_group = tvm_ffi.get_global_func("vm.builtin.attention_kv_cache_export_page_group") + fprepare_import = tvm_ffi.get_global_func("vm.builtin.attention_kv_cache_prepare_import") + fimport_page_group = tvm_ffi.get_global_func("vm.builtin.attention_kv_cache_import_page_group") + ffinish_import = tvm_ffi.get_global_func("vm.builtin.attention_kv_cache_finish_import") + fget_sequence_length = tvm_ffi.get_global_func( "vm.builtin.attention_kv_cache_get_sequence_length" ) @@ -178,7 +178,7 @@ def test_checkpoint_metadata_reports_layout_pages_and_groups(): def test_checkpoint_layout_hash_is_stable_and_layout_sensitive(): - fget_layout_hash = tvm_tirx.get_global_func("vm.builtin.attention_kv_cache_get_layout_hash") + fget_layout_hash = tvm_ffi.get_global_func("vm.builtin.attention_kv_cache_get_layout_hash") kv_cache = create_kv_cache() same_layout = create_kv_cache() @@ -205,11 +205,11 @@ def test_checkpoint_layout_hash_is_stable_and_layout_sensitive(): def test_checkpoint_metadata_rejects_unsupported_layouts_and_sequence_ids(): - fget_checkpoint_metadata = tvm_tirx.get_global_func( + fget_checkpoint_metadata = tvm_ffi.get_global_func( "vm.builtin.attention_kv_cache_get_checkpoint_metadata" ) - fget_layout_hash = tvm_tirx.get_global_func("vm.builtin.attention_kv_cache_get_layout_hash") - fexport_page_group = tvm_tirx.get_global_func("vm.builtin.attention_kv_cache_export_page_group") + fget_layout_hash = tvm_ffi.get_global_func("vm.builtin.attention_kv_cache_get_layout_hash") + fexport_page_group = tvm_ffi.get_global_func("vm.builtin.attention_kv_cache_export_page_group") sliding_cache = create_kv_cache(support_sliding_window=True) append_tokens(sliding_cache) @@ -240,8 +240,8 @@ def test_checkpoint_metadata_rejects_unsupported_layouts_and_sequence_ids(): fget_layout_hash(rope_ext_cache) tree_cache = create_kv_cache() - tvm_tirx.get_global_func("vm.builtin.kv_state_add_sequence")(tree_cache, 0) - tvm_tirx.get_global_func("vm.builtin.kv_state_begin_forward")( + tvm_ffi.get_global_func("vm.builtin.kv_state_add_sequence")(tree_cache, 0) + tvm_ffi.get_global_func("vm.builtin.kv_state_begin_forward")( tree_cache, Shape([0]), Shape([2]), Shape([-1, 0]) ) with pytest.raises(RuntimeError, match="committed token-chain state"): @@ -249,10 +249,10 @@ def test_checkpoint_metadata_rejects_unsupported_layouts_and_sequence_ids(): def test_checkpoint_export_page_group_validates_group_shape(): - fget_checkpoint_metadata = tvm_tirx.get_global_func( + fget_checkpoint_metadata = tvm_ffi.get_global_func( "vm.builtin.attention_kv_cache_get_checkpoint_metadata" ) - fexport_page_group = tvm_tirx.get_global_func("vm.builtin.attention_kv_cache_export_page_group") + fexport_page_group = tvm_ffi.get_global_func("vm.builtin.attention_kv_cache_export_page_group") kv_cache = create_kv_cache() append_tokens(kv_cache) @@ -287,10 +287,10 @@ def test_checkpoint_export_page_group_validates_group_shape(): def test_checkpoint_prepare_import_validates_metadata(): - fget_checkpoint_metadata = tvm_tirx.get_global_func( + fget_checkpoint_metadata = tvm_ffi.get_global_func( "vm.builtin.attention_kv_cache_get_checkpoint_metadata" ) - fprepare_import = tvm_tirx.get_global_func("vm.builtin.attention_kv_cache_prepare_import") + fprepare_import = tvm_ffi.get_global_func("vm.builtin.attention_kv_cache_prepare_import") kv_cache = create_kv_cache() append_tokens(kv_cache) @@ -336,11 +336,11 @@ def test_checkpoint_prepare_import_validates_metadata(): def test_checkpoint_import_page_group_validates_group_shape(): - fget_checkpoint_metadata = tvm_tirx.get_global_func( + fget_checkpoint_metadata = tvm_ffi.get_global_func( "vm.builtin.attention_kv_cache_get_checkpoint_metadata" ) - fprepare_import = tvm_tirx.get_global_func("vm.builtin.attention_kv_cache_prepare_import") - fimport_page_group = tvm_tirx.get_global_func("vm.builtin.attention_kv_cache_import_page_group") + fprepare_import = tvm_ffi.get_global_func("vm.builtin.attention_kv_cache_prepare_import") + fimport_page_group = tvm_ffi.get_global_func("vm.builtin.attention_kv_cache_import_page_group") kv_cache = create_kv_cache() append_tokens(kv_cache) @@ -379,14 +379,14 @@ def test_checkpoint_import_page_group_validates_group_shape(): def test_checkpoint_import_requires_all_groups_and_explicit_finish(): - fget_checkpoint_metadata = tvm_tirx.get_global_func( + fget_checkpoint_metadata = tvm_ffi.get_global_func( "vm.builtin.attention_kv_cache_get_checkpoint_metadata" ) - fexport_page_group = tvm_tirx.get_global_func("vm.builtin.attention_kv_cache_export_page_group") - fprepare_import = tvm_tirx.get_global_func("vm.builtin.attention_kv_cache_prepare_import") - fimport_page_group = tvm_tirx.get_global_func("vm.builtin.attention_kv_cache_import_page_group") - ffinish_import = tvm_tirx.get_global_func("vm.builtin.attention_kv_cache_finish_import") - fbegin_forward = tvm_tirx.get_global_func("vm.builtin.kv_state_begin_forward") + fexport_page_group = tvm_ffi.get_global_func("vm.builtin.attention_kv_cache_export_page_group") + fprepare_import = tvm_ffi.get_global_func("vm.builtin.attention_kv_cache_prepare_import") + fimport_page_group = tvm_ffi.get_global_func("vm.builtin.attention_kv_cache_import_page_group") + ffinish_import = tvm_ffi.get_global_func("vm.builtin.attention_kv_cache_finish_import") + fbegin_forward = tvm_ffi.get_global_func("vm.builtin.kv_state_begin_forward") source_cache = create_kv_cache() append_tokens(source_cache) diff --git a/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_mla_flashinfer.py b/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_mla_flashinfer.py index ccf530ad5739..7e0c624e46d5 100644 --- a/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_mla_flashinfer.py +++ b/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_mla_flashinfer.py @@ -95,7 +95,7 @@ def has_flashinfer(): # Register a dumb function for testing purpose. -@tvm_tirx.register_global_func("test.dumb_function", override=True) +@tvm_ffi.register_global_func("test.dumb_function", override=True) def _dumb_function(): raise RuntimeError("Dumb function isn't supposed to be accessed.") @@ -109,21 +109,21 @@ def set_global_func(dtype, target): global fattn_prefill_ragged, fattn_prefill_ragged_plan global fmerge_state, fmerge_state_additional, fcopy_single_page - fclear = tvm_tirx.get_global_func("vm.builtin.kv_state_clear") - fadd_sequence = tvm_tirx.get_global_func("vm.builtin.kv_state_add_sequence") - fremove_sequence = tvm_tirx.get_global_func("vm.builtin.kv_state_remove_sequence") - ffork_sequence = tvm_tirx.get_global_func("vm.builtin.kv_state_fork_sequence") - fpopn = tvm_tirx.get_global_func("vm.builtin.kv_state_popn") - fbegin_forward = tvm_tirx.get_global_func("vm.builtin.kv_state_begin_forward") - fend_forward = tvm_tirx.get_global_func("vm.builtin.kv_state_end_forward") - fself_attn = tvm_tirx.get_global_func("vm.builtin.attention_kv_cache_self_attention") - fcross_attn = tvm_tirx.get_global_func("vm.builtin.attention_kv_cache_cross_attention") - fappend_mla_kv = tvm_tirx.get_global_func("vm.builtin.attention_kv_cache_append_mla_kv") - fkv_merge_attn_output = tvm_tirx.get_global_func( + fclear = tvm_ffi.get_global_func("vm.builtin.kv_state_clear") + fadd_sequence = tvm_ffi.get_global_func("vm.builtin.kv_state_add_sequence") + fremove_sequence = tvm_ffi.get_global_func("vm.builtin.kv_state_remove_sequence") + ffork_sequence = tvm_ffi.get_global_func("vm.builtin.kv_state_fork_sequence") + fpopn = tvm_ffi.get_global_func("vm.builtin.kv_state_popn") + fbegin_forward = tvm_ffi.get_global_func("vm.builtin.kv_state_begin_forward") + fend_forward = tvm_ffi.get_global_func("vm.builtin.kv_state_end_forward") + fself_attn = tvm_ffi.get_global_func("vm.builtin.attention_kv_cache_self_attention") + fcross_attn = tvm_ffi.get_global_func("vm.builtin.attention_kv_cache_cross_attention") + fappend_mla_kv = tvm_ffi.get_global_func("vm.builtin.attention_kv_cache_append_mla_kv") + fkv_merge_attn_output = tvm_ffi.get_global_func( "vm.builtin.attention_kv_cache_merge_attn_output_inplace" ) - fis_empty = tvm_tirx.get_global_func("vm.builtin.attention_kv_cache_empty") - fdebug_get_kv = tvm_tirx.get_global_func("vm.builtin.attention_kv_cache_debug_get_kv_mla") + fis_empty = tvm_ffi.get_global_func("vm.builtin.attention_kv_cache_empty") + fdebug_get_kv = tvm_ffi.get_global_func("vm.builtin.attention_kv_cache_debug_get_kv_mla") flashinfer_prefill_mod = relax.backend.cuda.flashinfer.gen_flashinfer_prefill_module( dtype_q=dtype, @@ -187,8 +187,8 @@ def _initialize_weights(): def create_kv_cache(dtype): - fcreate = tvm_tirx.get_global_func("vm.builtin.paged_attention_kv_cache_create") - fdumb = tvm_tirx.get_global_func("test.dumb_function") + fcreate = tvm_ffi.get_global_func("vm.builtin.paged_attention_kv_cache_create") + fdumb = tvm_ffi.get_global_func("test.dumb_function") cache = fcreate( tvm_ffi.Shape( [ diff --git a/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_mla_tir.py b/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_mla_tir.py index ad21aa3ee4dc..8e3544c6a68f 100644 --- a/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_mla_tir.py +++ b/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_mla_tir.py @@ -80,7 +80,7 @@ # Register a dumb function for testing purpose. -@tvm_tirx.register_global_func("test.dumb_function", override=True) +@tvm_ffi.register_global_func("test.dumb_function", override=True) def _dumb_function(): raise RuntimeError("Dumb function isn't supposed to be accessed.") @@ -93,21 +93,21 @@ def set_global_func(dtype, target): global ftranspose_append, fcopy_cache, fmla_prefill, fmla_prefill_ragged global fmerge_state, fmerge_state_additional, fcopy_single_page - fclear = tvm_tirx.get_global_func("vm.builtin.kv_state_clear") - fadd_sequence = tvm_tirx.get_global_func("vm.builtin.kv_state_add_sequence") - fremove_sequence = tvm_tirx.get_global_func("vm.builtin.kv_state_remove_sequence") - ffork_sequence = tvm_tirx.get_global_func("vm.builtin.kv_state_fork_sequence") - fpopn = tvm_tirx.get_global_func("vm.builtin.kv_state_popn") - fbegin_forward = tvm_tirx.get_global_func("vm.builtin.kv_state_begin_forward") - fend_forward = tvm_tirx.get_global_func("vm.builtin.kv_state_end_forward") - fself_attn = tvm_tirx.get_global_func("vm.builtin.attention_kv_cache_self_attention") - fcross_attn = tvm_tirx.get_global_func("vm.builtin.attention_kv_cache_cross_attention") - fappend_mla_kv = tvm_tirx.get_global_func("vm.builtin.attention_kv_cache_append_mla_kv") - fkv_merge_attn_output = tvm_tirx.get_global_func( + fclear = tvm_ffi.get_global_func("vm.builtin.kv_state_clear") + fadd_sequence = tvm_ffi.get_global_func("vm.builtin.kv_state_add_sequence") + fremove_sequence = tvm_ffi.get_global_func("vm.builtin.kv_state_remove_sequence") + ffork_sequence = tvm_ffi.get_global_func("vm.builtin.kv_state_fork_sequence") + fpopn = tvm_ffi.get_global_func("vm.builtin.kv_state_popn") + fbegin_forward = tvm_ffi.get_global_func("vm.builtin.kv_state_begin_forward") + fend_forward = tvm_ffi.get_global_func("vm.builtin.kv_state_end_forward") + fself_attn = tvm_ffi.get_global_func("vm.builtin.attention_kv_cache_self_attention") + fcross_attn = tvm_ffi.get_global_func("vm.builtin.attention_kv_cache_cross_attention") + fappend_mla_kv = tvm_ffi.get_global_func("vm.builtin.attention_kv_cache_append_mla_kv") + fkv_merge_attn_output = tvm_ffi.get_global_func( "vm.builtin.attention_kv_cache_merge_attn_output_inplace" ) - fis_empty = tvm_tirx.get_global_func("vm.builtin.attention_kv_cache_empty") - fdebug_get_kv = tvm_tirx.get_global_func("vm.builtin.attention_kv_cache_debug_get_kv_mla") + fis_empty = tvm_ffi.get_global_func("vm.builtin.attention_kv_cache_empty") + fdebug_get_kv = tvm_ffi.get_global_func("vm.builtin.attention_kv_cache_debug_get_kv_mla") builts = [] for tir_func in [ @@ -147,8 +147,8 @@ def set_global_func(dtype, target): def create_kv_cache(dtype): - fcreate = tvm_tirx.get_global_func("vm.builtin.paged_attention_kv_cache_create") - fdumb = tvm_tirx.get_global_func("test.dumb_function") + fcreate = tvm_ffi.get_global_func("vm.builtin.paged_attention_kv_cache_create") + fdumb = tvm_ffi.get_global_func("test.dumb_function") cache = fcreate( tvm_ffi.Shape( [ diff --git a/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_tir.py b/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_tir.py index f898e074f897..130d67e21e3f 100644 --- a/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_tir.py +++ b/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_tir.py @@ -100,7 +100,7 @@ def test_paged_prefill_layer_sliding_window_mask(device_type): if not tvm_tirx.testing.device_enabled(device_type): pytest.skip(f"{device_type} not enabled") - dev = tvm_tirx.device(device_type) + dev = tvm_ffi.device(device_type) target = tvm_tirx.target.Target.from_device(dev) head_dim = 64 num_kv_heads = 2 @@ -168,7 +168,7 @@ def test_paged_attention_kv_cache_sliding_window_two_sequences(device_type): def run_and_check(): global device - device = tvm_tirx.device(device_type) + device = tvm_ffi.device(device_type) set_global_func(head_dim, dtype, tvm_tirx.target.Target.from_device(device)) try: cache = create_kv_cache(head_dim, dtype, RopeMode.NONE, True) @@ -203,24 +203,24 @@ def set_global_func(head_dim, dtype, target): global fattn_prefill_sliding_window, fattn_decode_sliding_window global fmerge_state, fsplit_rotary, fattention_rotary, fcopy_single_page, fcompact_copy - fclear = tvm_tirx.get_global_func("vm.builtin.kv_state_clear") - fadd_sequence = tvm_tirx.get_global_func("vm.builtin.kv_state_add_sequence") - fremove_sequence = tvm_tirx.get_global_func("vm.builtin.kv_state_remove_sequence") - ffork_sequence = tvm_tirx.get_global_func("vm.builtin.kv_state_fork_sequence") - fenable_sliding_window_for_seq = tvm_tirx.get_global_func( + fclear = tvm_ffi.get_global_func("vm.builtin.kv_state_clear") + fadd_sequence = tvm_ffi.get_global_func("vm.builtin.kv_state_add_sequence") + fremove_sequence = tvm_ffi.get_global_func("vm.builtin.kv_state_remove_sequence") + ffork_sequence = tvm_ffi.get_global_func("vm.builtin.kv_state_fork_sequence") + fenable_sliding_window_for_seq = tvm_ffi.get_global_func( "vm.builtin.attention_kv_cache_enable_sliding_window_for_seq" ) - fpopn = tvm_tirx.get_global_func("vm.builtin.kv_state_popn") - fbegin_forward = tvm_tirx.get_global_func("vm.builtin.kv_state_begin_forward") - fend_forward = tvm_tirx.get_global_func("vm.builtin.kv_state_end_forward") - fcommit_accepted_token_tree_nodes = tvm_tirx.get_global_func( + fpopn = tvm_ffi.get_global_func("vm.builtin.kv_state_popn") + fbegin_forward = tvm_ffi.get_global_func("vm.builtin.kv_state_begin_forward") + fend_forward = tvm_ffi.get_global_func("vm.builtin.kv_state_end_forward") + fcommit_accepted_token_tree_nodes = tvm_ffi.get_global_func( "vm.builtin.attention_kv_cache_commit_accepted_token_tree_nodes" ) - fattention_with_fuse_qkv = tvm_tirx.get_global_func( + fattention_with_fuse_qkv = tvm_ffi.get_global_func( "vm.builtin.attention_kv_cache_attention_with_fused_qkv" ) - fis_empty = tvm_tirx.get_global_func("vm.builtin.attention_kv_cache_empty") - fdebug_get_kv = tvm_tirx.get_global_func("vm.builtin.attention_kv_cache_debug_get_kv") + fis_empty = tvm_ffi.get_global_func("vm.builtin.attention_kv_cache_empty") + fdebug_get_kv = tvm_ffi.get_global_func("vm.builtin.attention_kv_cache_debug_get_kv") cache_key = ( str(target), @@ -292,7 +292,7 @@ def set_global_func(head_dim, dtype, target): def create_kv_cache(head_dim, dtype, rope_mode, support_sliding_window): - fcreate = tvm_tirx.get_global_func("vm.builtin.paged_attention_kv_cache_create") + fcreate = tvm_ffi.get_global_func("vm.builtin.paged_attention_kv_cache_create") cache = fcreate( tvm_ffi.Shape( [ diff --git a/tests/python/relax/test_runtime_builtin_rnn_state.py b/tests/python/relax/test_runtime_builtin_rnn_state.py index 5dc8d19954cf..75e240921f3d 100644 --- a/tests/python/relax/test_runtime_builtin_rnn_state.py +++ b/tests/python/relax/test_runtime_builtin_rnn_state.py @@ -19,6 +19,7 @@ import numpy as np import pytest +import tvm_ffi from tvm_ffi import Shape import tvm_tirx @@ -64,16 +65,16 @@ def set_global_func(): global f_begin_forward, f_end_forward, f_get, f_set, f_debug_get global f_tir_gets, f_tir_sets - f_clear = tvm_tirx.get_global_func("vm.builtin.kv_state_clear") - f_add_sequence = tvm_tirx.get_global_func("vm.builtin.kv_state_add_sequence") - f_remove_sequence = tvm_tirx.get_global_func("vm.builtin.kv_state_remove_sequence") - f_fork_sequence = tvm_tirx.get_global_func("vm.builtin.kv_state_fork_sequence") - f_popn = tvm_tirx.get_global_func("vm.builtin.kv_state_popn") - f_begin_forward = tvm_tirx.get_global_func("vm.builtin.kv_state_begin_forward") - f_end_forward = tvm_tirx.get_global_func("vm.builtin.kv_state_end_forward") - f_get = tvm_tirx.get_global_func("vm.builtin.rnn_state_get") - f_set = tvm_tirx.get_global_func("vm.builtin.rnn_state_set") - f_debug_get = tvm_tirx.get_global_func("vm.builtin.rnn_state_debug_get") + f_clear = tvm_ffi.get_global_func("vm.builtin.kv_state_clear") + f_add_sequence = tvm_ffi.get_global_func("vm.builtin.kv_state_add_sequence") + f_remove_sequence = tvm_ffi.get_global_func("vm.builtin.kv_state_remove_sequence") + f_fork_sequence = tvm_ffi.get_global_func("vm.builtin.kv_state_fork_sequence") + f_popn = tvm_ffi.get_global_func("vm.builtin.kv_state_popn") + f_begin_forward = tvm_ffi.get_global_func("vm.builtin.kv_state_begin_forward") + f_end_forward = tvm_ffi.get_global_func("vm.builtin.kv_state_end_forward") + f_get = tvm_ffi.get_global_func("vm.builtin.rnn_state_get") + f_set = tvm_ffi.get_global_func("vm.builtin.rnn_state_set") + f_debug_get = tvm_ffi.get_global_func("vm.builtin.rnn_state_debug_get") target = tvm_tirx.target.Target("cuda") @@ -95,7 +96,7 @@ def _build(tir_func): def create_rnn_state(device): - f_create = tvm_tirx.get_global_func("vm.builtin.rnn_state_create") + f_create = tvm_ffi.get_global_func("vm.builtin.rnn_state_create") init_values = [ tvm_tirx.runtime.tensor(np_zero, device=device), tvm_tirx.runtime.tensor(np_one, device=device), diff --git a/tests/python/relax/test_runtime_sampling_flashinfer.py b/tests/python/relax/test_runtime_sampling_flashinfer.py index d5d5ffc9761f..ce842d02e63f 100644 --- a/tests/python/relax/test_runtime_sampling_flashinfer.py +++ b/tests/python/relax/test_runtime_sampling_flashinfer.py @@ -18,6 +18,7 @@ import numpy as np import pytest +import tvm_ffi import tvm_tirx import tvm_tirx.testing @@ -27,7 +28,7 @@ @pytest.mark.skip(reason="Requires FlashInfer enabled and proper setup") def test_sampling(): - def load_module(name: str, static_modules: list[tvm_tirx.runtime.Module]): + def load_module(name: str, static_modules: list[tvm_ffi.Module]): assert len(static_modules) > 0 if len(static_modules) == 1: return static_modules[0] diff --git a/tests/python/relax/test_tensorrt_engine_serialization.py b/tests/python/relax/test_tensorrt_engine_serialization.py index 6769a55794f1..6ac694b67272 100644 --- a/tests/python/relax/test_tensorrt_engine_serialization.py +++ b/tests/python/relax/test_tensorrt_engine_serialization.py @@ -25,6 +25,7 @@ import numpy as np import pytest +import tvm_ffi import tvm import tvm_tirx @@ -36,8 +37,8 @@ from tvm_tirx.support import cc from tvm_tirx.testing import env -_has_codegen = tvm_tirx.get_global_func("relax.ext.tensorrt", True) is not None -_runtime_enabled = tvm_tirx.get_global_func("relax.is_tensorrt_runtime_enabled", True) +_has_codegen = tvm_ffi.get_global_func("relax.ext.tensorrt", True) is not None +_runtime_enabled = tvm_ffi.get_global_func("relax.is_tensorrt_runtime_enabled", True) _has_runtime = _runtime_enabled is not None and _runtime_enabled() pytestmark = pytest.mark.skipif(not _has_codegen, reason="TensorRT codegen is not enabled") @@ -81,7 +82,7 @@ def _legacy_bytes(module): def _module_bytes(module): """Extract a leaf's payload from the same import packer used by export_library.""" assert not module.imports - packed = tvm_tirx.get_global_func("runtime.ModulePackImportsToTensor")(module).numpy().tobytes() + packed = tvm_ffi.get_global_func("runtime.ModulePackImportsToTensor")(module).numpy().tobytes() assert struct.unpack_from("= y[1]), x[1], y[1]) return lhs, rhs - def f_identity(dtype0: tvm_tirx.DataType, dtype1: tvm_tirx.DataType): + def f_identity(dtype0: tvm_ffi.dtype, dtype1: tvm_ffi.dtype): return tvm_tirx.tirx.const(-1, dtype0), tvm.te.min_value(ty=dtype1) argmax = te.comm_reducer(f_combine, f_identity, name="argmax") diff --git a/tests/python/relax/test_transform_ipc_allreduce_rewrite.py b/tests/python/relax/test_transform_ipc_allreduce_rewrite.py index 4a071b2ee933..fcf636512971 100644 --- a/tests/python/relax/test_transform_ipc_allreduce_rewrite.py +++ b/tests/python/relax/test_transform_ipc_allreduce_rewrite.py @@ -15,6 +15,8 @@ # specific language governing permissions and limitations # under the License. +import tvm_ffi + import tvm_tirx import tvm_tirx.testing from tvm import relax @@ -68,7 +70,7 @@ def main(shape: R.Shape([m, n])): # type: ignore mod, ( Expected - if tvm_tirx.get_global_func( + if tvm_ffi.get_global_func( "runtime.disco.cuda_ipc.custom_allreduce", allow_missing=True ) is not None @@ -126,7 +128,7 @@ def main( mod, ( Expected - if tvm_tirx.get_global_func( + if tvm_ffi.get_global_func( "runtime.disco.cuda_ipc.custom_allreduce", allow_missing=True ) is not None diff --git a/tests/python/relax/test_transform_lazy_transform_params.py b/tests/python/relax/test_transform_lazy_transform_params.py index 6f2ebde8885b..23897ed507ab 100644 --- a/tests/python/relax/test_transform_lazy_transform_params.py +++ b/tests/python/relax/test_transform_lazy_transform_params.py @@ -16,6 +16,7 @@ # under the License. # ruff: noqa: F841 import numpy as np +import tvm_ffi import tvm import tvm_tirx @@ -656,11 +657,11 @@ def transform_params( transformed = {} expected = [params[0].transpose(1, 0, 2, 3), params[1]] - @tvm_tirx.register_global_func("get_item", override=True) + @tvm_ffi.register_global_func("get_item", override=True) def get_item(i): return tvm_tirx.runtime.tensor(params[i], dev) - @tvm_tirx.register_global_func("set_item", override=True) + @tvm_ffi.register_global_func("set_item", override=True) def set_item(i, value): assert i not in transformed, f"Set item called multiple times for index {i}" transformed[i] = value.numpy() diff --git a/tests/python/relax/test_vm_build.py b/tests/python/relax/test_vm_build.py index 33025c8f6553..bc814d216267 100644 --- a/tests/python/relax/test_vm_build.py +++ b/tests/python/relax/test_vm_build.py @@ -387,7 +387,7 @@ def main(x: R.Tensor((2, 3), "int32"), y: R.Tensor((2, 3), "int32")) -> R.Tensor def test_vm_emit_te_extern(): - if not tvm_tirx.get_global_func("tvm.contrib.cblas.matmul", True): + if not tvm_ffi.get_global_func("tvm.contrib.cblas.matmul", True): print("skip because extern function is not available") return bb = relax.BlockBuilder() @@ -1032,8 +1032,8 @@ def popen_check(): # Load dll, will trigger system library registration ctypes.CDLL(path_dso) # Load the system wide library - vmA = relax.VirtualMachine(tvm_tirx.runtime.system_lib("libA_"), tvm_tirx.cpu()) - vmB = relax.VirtualMachine(tvm_tirx.runtime.system_lib("libB_"), tvm_tirx.cpu()) + vmA = relax.VirtualMachine(tvm_ffi.system_lib("libA_"), tvm_tirx.cpu()) + vmB = relax.VirtualMachine(tvm_ffi.system_lib("libB_"), tvm_tirx.cpu()) retA = vmA["main"](tvm_ffi.Shape([1])) retB = vmB["main"](tvm_ffi.Shape([2])) diff --git a/tests/python/relax/test_vm_cuda_graph.py b/tests/python/relax/test_vm_cuda_graph.py index aa3aff564f07..2d10ffc18953 100644 --- a/tests/python/relax/test_vm_cuda_graph.py +++ b/tests/python/relax/test_vm_cuda_graph.py @@ -21,6 +21,7 @@ import numpy as np import pytest +import tvm_ffi import tvm import tvm_tirx @@ -182,7 +183,7 @@ def run_and_check(): cudart.cudaGetLastError.restype = ctypes.c_int cudart.cudaGetLastError() - @tvm_tirx.register_global_func( + @tvm_ffi.register_global_func( "test_vm_cuda_graph.invalid_impl_for_cudagraph", override=True ) def invalid_impl_for_cudagraph(arg_tensor): diff --git a/tests/python/relax/test_vm_execbuilder.py b/tests/python/relax/test_vm_execbuilder.py index 95b83b7f4c17..71b3272ff4b1 100644 --- a/tests/python/relax/test_vm_execbuilder.py +++ b/tests/python/relax/test_vm_execbuilder.py @@ -168,7 +168,7 @@ def test_vm_shapeof(): def test_vm_storage(): - dtype = tvm_tirx.DataType("float32") + dtype = tvm_ffi.dtype("float32") shape = (4, 6) ib = relax.ExecBuilder() with ib.function("main", num_inputs=0): diff --git a/tests/python/relax/texture/test_texture_nd.py b/tests/python/relax/texture/test_texture_nd.py index 2dfd6ba0c2cd..d8464ecba62a 100644 --- a/tests/python/relax/texture/test_texture_nd.py +++ b/tests/python/relax/texture/test_texture_nd.py @@ -21,6 +21,7 @@ import numpy as np import pytest +import tvm_ffi import tvm import tvm_tirx @@ -29,7 +30,7 @@ from tvm.relax.transform.legalize_ops import adreno as legalize_adreno from tvm.rpc import connect_tracker from tvm.s_tir import script as Ts -from tvm_tirx import DataType, IRModule, tirx +from tvm_tirx import IRModule, tirx from tvm_tirx.script import ir as I from tvm_tirx.script import tirx as T from tvm_tirx.support import ndk @@ -114,7 +115,7 @@ def postprocess_pipeline(mod: IRModule) -> IRModule: def test_texture_copy(backend, dtype, channel_size, read_width): remote = get_rpc() M, N, K = (256, 1024, 128) - lanes = channel_size // DataType(dtype).bits + lanes = channel_size // tvm_ffi.dtype(dtype).bits if read_width > lanes: return diff --git a/tests/python/runtime/rpc/test_runtime_rpc.py b/tests/python/runtime/rpc/test_runtime_rpc.py index 9da61daa389d..bd758688e289 100644 --- a/tests/python/runtime/rpc/test_runtime_rpc.py +++ b/tests/python/runtime/rpc/test_runtime_rpc.py @@ -55,7 +55,7 @@ # Windows does not support fork so we can enable Windows for testing sys.platform.startswith("win") == False and multiprocessing.get_start_method() != "fork", reason=( - "pytest + multiprocessing spawn method causes tvm_tirx.register_global_func to " + "pytest + multiprocessing spawn method causes tvm_ffi.register_global_func to " "not work on the rpc.Server." ), ) @@ -129,7 +129,7 @@ def system_library_source(native_source): @pytest.mark.skipif( - tvm_tirx.get_global_func("rpc.LocalSession", allow_missing=True) is None, reason="need rpc" + tvm_ffi.get_global_func("rpc.LocalSession", allow_missing=True) is None, reason="need rpc" ) def test_rpc_simple(): server = rpc.Server(key="x1") @@ -150,7 +150,7 @@ def check_remote(): @pytest.mark.skipif( - tvm_tirx.get_global_func("rpc.LocalSession", allow_missing=True) is None, reason="need rpc" + tvm_ffi.get_global_func("rpc.LocalSession", allow_missing=True) is None, reason="need rpc" ) def test_rpc_runtime_string(): server = rpc.Server(key="x1") @@ -166,7 +166,7 @@ def check_remote(): @pytest.mark.skipif( - tvm_tirx.get_global_func("rpc.LocalSession", allow_missing=True) is None, reason="need rpc" + tvm_ffi.get_global_func("rpc.LocalSession", allow_missing=True) is None, reason="need rpc" ) def test_rpc_array(): server = rpc.Server() @@ -184,7 +184,7 @@ def check_remote(): @pytest.mark.skipif( - tvm_tirx.get_global_func("rpc.LocalSession", allow_missing=True) is None, reason="need rpc" + tvm_ffi.get_global_func("rpc.LocalSession", allow_missing=True) is None, reason="need rpc" ) def test_rpc_large_array(): # testcase of large array creation @@ -205,7 +205,7 @@ def check_remote(): @pytest.mark.skipif("32bit" in platform.architecture()[0], reason="skipping test for i386.") @pytest.mark.skipif( - tvm_tirx.get_global_func("rpc.LocalSession", allow_missing=True) is None, reason="need rpc" + tvm_ffi.get_global_func("rpc.LocalSession", allow_missing=True) is None, reason="need rpc" ) def test_rpc_echo(): def check(remote, local_session): @@ -236,7 +236,7 @@ def check(remote, local_session): @pytest.mark.skipif( - tvm_tirx.get_global_func("rpc.LocalSession", allow_missing=True) is None, reason="need rpc" + tvm_ffi.get_global_func("rpc.LocalSession", allow_missing=True) is None, reason="need rpc" ) def test_rpc_file_exchange(): server = rpc.Server() @@ -252,7 +252,7 @@ def check_remote(): @pytest.mark.skipif( - tvm_tirx.get_global_func("rpc.LocalSession", allow_missing=True) is None, reason="need rpc" + tvm_ffi.get_global_func("rpc.LocalSession", allow_missing=True) is None, reason="need rpc" ) def test_rpc_remote_module(native_object, system_library_source, tmp_path): module = tvm_tirx.codegen.load_static_library( @@ -298,7 +298,7 @@ def check_remote(remote, name): server0.terminate() server1.terminate() - if tvm_tirx.get_global_func("rpc.CreatePipeClient", allow_missing=True) is None: + if tvm_ffi.get_global_func("rpc.CreatePipeClient", allow_missing=True) is None: return source = tmp_path / "system.cc" source.write_text(system_library_source) @@ -320,7 +320,7 @@ def check_remote(remote, name): @pytest.mark.skipif( - tvm_tirx.get_global_func("rpc.LocalSession", allow_missing=True) is None, reason="need rpc" + tvm_ffi.get_global_func("rpc.LocalSession", allow_missing=True) is None, reason="need rpc" ) def test_rpc_return_func(): server = rpc.Server(key="x1") @@ -335,7 +335,7 @@ def check_remote(): @pytest.mark.skipif( - tvm_tirx.get_global_func("rpc.LocalSession", allow_missing=True) is None, reason="need rpc" + tvm_ffi.get_global_func("rpc.LocalSession", allow_missing=True) is None, reason="need rpc" ) def test_rpc_session_constructor_args(): # start server @@ -374,7 +374,7 @@ def check_error_handling(): @pytest.mark.skipif( - tvm_tirx.get_global_func("rpc.LocalSession", allow_missing=True) is None, reason="need rpc" + tvm_ffi.get_global_func("rpc.LocalSession", allow_missing=True) is None, reason="need rpc" ) def test_rpc_return_tensor(): def run_arr_test(): @@ -397,7 +397,7 @@ def run_arr_test(): @pytest.mark.skipif( - tvm_tirx.get_global_func("rpc.LocalSession", allow_missing=True) is None, reason="need rpc" + tvm_ffi.get_global_func("rpc.LocalSession", allow_missing=True) is None, reason="need rpc" ) def test_rpc_return_remote_object(): def check(client, is_local): @@ -425,7 +425,7 @@ def check_remote(): check_remote() def check_minrpc(): - if tvm_tirx.get_global_func("rpc.CreatePipeClient", allow_missing=True) is None: + if tvm_ffi.get_global_func("rpc.CreatePipeClient", allow_missing=True) is None: return # Test minrpc server. temp = utils.tempdir() @@ -445,7 +445,7 @@ def check_minrpc(): @pytest.mark.skipif( - tvm_tirx.get_global_func("rpc.LocalSession", allow_missing=True) is None, reason="need rpc" + tvm_ffi.get_global_func("rpc.LocalSession", allow_missing=True) is None, reason="need rpc" ) def test_local_func(): client = rpc.LocalSession() @@ -464,7 +464,7 @@ def check_remote(): @pytest.mark.skipif( - tvm_tirx.get_global_func("rpc.LocalSession", allow_missing=True) is None, reason="need rpc" + tvm_ffi.get_global_func("rpc.LocalSession", allow_missing=True) is None, reason="need rpc" ) @pytest.mark.parametrize("device_key", ["test_device", "127.0.0.1:5555"]) def test_rpc_tracker_register(device_key): @@ -539,7 +539,7 @@ def _target(host, port, device_key, timeout, acquired): @pytest.mark.skipif( - tvm_tirx.get_global_func("rpc.LocalSession", allow_missing=True) is None, reason="need rpc" + tvm_ffi.get_global_func("rpc.LocalSession", allow_missing=True) is None, reason="need rpc" ) @pytest.mark.parametrize("device_key", ["test_device", "127.0.0.1:5555"]) def test_rpc_tracker_request(device_key): @@ -587,7 +587,7 @@ def test_rpc_tracker_request(device_key): @pytest.mark.skipif( - tvm_tirx.get_global_func("rpc.LocalSession", allow_missing=True) is None, reason="need rpc" + tvm_ffi.get_global_func("rpc.LocalSession", allow_missing=True) is None, reason="need rpc" ) @pytest.mark.parametrize("device_key", ["test_device", "127.0.0.1:5555"]) def test_rpc_tracker_via_proxy(device_key): @@ -631,7 +631,7 @@ def test_rpc_tracker_via_proxy(device_key): @pytest.mark.skipif( - tvm_tirx.get_global_func("rpc.LocalSession", allow_missing=True) is None, reason="need rpc" + tvm_ffi.get_global_func("rpc.LocalSession", allow_missing=True) is None, reason="need rpc" ) @pytest.mark.parametrize("with_proxy", (True, False)) def test_rpc_session_timeout_error(with_proxy): @@ -665,7 +665,7 @@ def test_rpc_session_timeout_error(with_proxy): @pytest.mark.skipif( - tvm_tirx.get_global_func("rpc.LocalSession", allow_missing=True) is None, reason="need rpc" + tvm_ffi.get_global_func("rpc.LocalSession", allow_missing=True) is None, reason="need rpc" ) @pytest.mark.parametrize("call_with_unused_argument", [True, False]) def test_remote_function_with_zero_arguments(native_library, call_with_unused_argument): diff --git a/tests/python/s_tir/base/test_s_tir_texture_scope.py b/tests/python/s_tir/base/test_s_tir_texture_scope.py index 53996902bfb4..fa0b849e1ed4 100644 --- a/tests/python/s_tir/base/test_s_tir_texture_scope.py +++ b/tests/python/s_tir/base/test_s_tir_texture_scope.py @@ -17,6 +17,7 @@ # ruff: noqa: F401 import pytest +import tvm_ffi import tvm import tvm_tirx @@ -72,7 +73,7 @@ def schedule_block(block): assert len(device_mods) == 1 (device_target, _backend_config), device_mod = next(iter(device_mods.items())) device_mod = finalize_device_passes()(device_mod) - source = tvm_tirx.get_global_func("target.build.opencl")( + source = tvm_ffi.get_global_func("target.build.opencl")( device_mod, device_target ).inspect_source() assert "__read_only image2d_array_t" in source diff --git a/tests/python/s_tir/dlight/test_primitives.py b/tests/python/s_tir/dlight/test_primitives.py index 78959930d608..00416f776d3f 100644 --- a/tests/python/s_tir/dlight/test_primitives.py +++ b/tests/python/s_tir/dlight/test_primitives.py @@ -18,6 +18,7 @@ # ruff: noqa: F841 import pytest +import tvm_ffi import tvm import tvm_tirx @@ -59,7 +60,7 @@ def main(p0: T.Tensor((), "int32"), T_stack: T.Tensor((T.int64(3),), "int32")): @pytest.mark.skipif(not env.has_cuda(), reason="need cuda") def test_normalize_function_with_scalar(): sch = tvm.s_tir.Schedule(main) - f_normalize_function = tvm_tirx.get_global_func("s_tir.schedule.NormalizeFunction") + f_normalize_function = tvm_ffi.get_global_func("s_tir.schedule.NormalizeFunction") assert f_normalize_function(sch) diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_builder.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_builder.py index 21972f99fedd..2a4682a6ae1e 100644 --- a/tests/python/s_tir/meta_schedule/test_meta_schedule_builder.py +++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_builder.py @@ -22,6 +22,7 @@ import time import pytest +import tvm_ffi from tvm_ffi import register_global_func import tvm @@ -34,7 +35,6 @@ PyBuilder, ) from tvm_tirx import script -from tvm_tirx.runtime import Module from tvm_tirx.script import tirx as T from tvm_tirx.target import Target @@ -170,7 +170,7 @@ def test_meta_schedule_error_handle_build_func(): def initializer(): @register_global_func("meta_schedule.builder.test_build") - def test_build(mod: Module, target: Target, _) -> None: # pylint: disable=unused-variable + def test_build(mod: tvm_ffi.Module, target: Target, _) -> None: # pylint: disable=unused-variable raise ValueError("Builder intended Test Error (build func).") builder = LocalBuilder(f_build="meta_schedule.builder.test_build", initializer=initializer) @@ -190,7 +190,7 @@ def test_meta_schedule_error_handle_export_func(): def initializer(): @register_global_func("meta_schedule.builder.test_export") - def test_build(mod: Module) -> str: # pylint: disable=unused-variable + def test_build(mod: tvm_ffi.Module) -> str: # pylint: disable=unused-variable raise ValueError("Builder intended Test Error (export func).") builder = LocalBuilder(f_export="meta_schedule.builder.test_export", initializer=initializer) diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_runner.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_runner.py index 55b47da64ddd..46b64c78400e 100644 --- a/tests/python/s_tir/meta_schedule/test_meta_schedule_runner.py +++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_runner.py @@ -25,6 +25,7 @@ import numpy as np import pytest +import tvm_ffi from tvm_ffi import register_global_func pytest.importorskip("tornado") # tvm.rpc.tracker (LocalRPC) requires tornado @@ -61,7 +62,7 @@ ) from tvm.s_tir.utils import derived_object from tvm_tirx.ir import FloatImm -from tvm_tirx.runtime import Device, Module +from tvm_tirx.runtime import Device from tvm_tirx.script import tirx as T from tvm_tirx.target import Target @@ -643,7 +644,7 @@ def test_alloc_argument( def test_run_evaluator( session: RPCSession, # pylint: disable=unused-argument - rt_mod: Module, + rt_mod: tvm_ffi.Module, device: Device, evaluator_config: EvaluatorConfig, repeated_args: list[Any], @@ -757,7 +758,7 @@ def test_alloc_argument( def test_run_evaluator( session: RPCSession, # pylint: disable=unused-argument - rt_mod: Module, + rt_mod: tvm_ffi.Module, device: Device, evaluator_config: EvaluatorConfig, repeated_args: list[Any], @@ -864,7 +865,7 @@ def test_alloc_argument( return repeated_args def test_run_evaluator( - rt_mod: Module, + rt_mod: tvm_ffi.Module, device: Device, evaluator_config: EvaluatorConfig, repeated_args: list[Any], diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_schedule_rule_apply_custom_rule.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_schedule_rule_apply_custom_rule.py index 4aebfa4cc523..224aa609bbbc 100644 --- a/tests/python/s_tir/meta_schedule/test_meta_schedule_schedule_rule_apply_custom_rule.py +++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_schedule_rule_apply_custom_rule.py @@ -18,6 +18,7 @@ import tempfile import pytest +import tvm_ffi import tvm import tvm_tirx @@ -46,7 +47,7 @@ def main( C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vk, vj] -@tvm_tirx.register_global_func("s_tir.meta_schedule.cpu.test_apply_custom_rule") +@tvm_ffi.register_global_func("s_tir.meta_schedule.cpu.test_apply_custom_rule") def sch_fn(sch: tvm.s_tir.Schedule, block: tvm.s_tir.SBlock) -> list[tvm.s_tir.Schedule]: raise ValueError("Intended for s_tir.meta_schedule.cpu.test_apply_custom_rule") diff --git a/tests/python/s_tir/script/test_s_tir_script_printer.py b/tests/python/s_tir/script/test_s_tir_script_printer.py index 2db7f5cdf322..958a59021104 100644 --- a/tests/python/s_tir/script/test_s_tir_script_printer.py +++ b/tests/python/s_tir/script/test_s_tir_script_printer.py @@ -20,6 +20,7 @@ import sys import pytest +import tvm_ffi from tvm_ffi.access_path import AccessPath import tvm @@ -2180,9 +2181,7 @@ def foo(): def test_str_metadata(): # This test is to check we reuse the existing metadata element for the same metadata-backed GenericConst # So metadata["ir.GenericConst"][0] will occur in the printed script for three times - str_imm = tvm_tirx.ir.GenericConst( - tvm_tirx.ir.AnyType(), tvm_tirx.runtime.convert(["aaa\nbbb\n"]) - ) + str_imm = tvm_tirx.ir.GenericConst(tvm_tirx.ir.AnyType(), tvm_ffi.convert(["aaa\nbbb\n"])) @I.ir_module class Module: diff --git a/tests/python/s_tir/test_arith_domain_touched.py b/tests/python/s_tir/test_arith_domain_touched.py index e7a874abc833..307ea22a1ffa 100644 --- a/tests/python/s_tir/test_arith_domain_touched.py +++ b/tests/python/s_tir/test_arith_domain_touched.py @@ -73,7 +73,7 @@ def test_domain_touched(): def test_domain_touched_vector(): pytest.skip("TensorRegion arithmetic in expressions not supported") - m = tvm_tirx.runtime.convert(128) + m = tvm_ffi.convert(128) @Ts.function def func(A: T.Tensor((n * m,), "float32"), B: T.Tensor((n * m,), "float32"), n: T.int32): # noqa: F821 diff --git a/tests/python/s_tir/transform/test_s_tir_transform_inject_ptx_async_copy.py b/tests/python/s_tir/transform/test_s_tir_transform_inject_ptx_async_copy.py index 84a4ff313f0c..44763ee024fc 100644 --- a/tests/python/s_tir/transform/test_s_tir_transform_inject_ptx_async_copy.py +++ b/tests/python/s_tir/transform/test_s_tir_transform_inject_ptx_async_copy.py @@ -62,7 +62,7 @@ def verify(n): def generate_global_to_shared_vectorized_copy(dtype, vector_size): num_iters = 128 // vector_size - vector_size_expr = tvm_tirx.runtime.convert(vector_size) + vector_size_expr = tvm_ffi.convert(vector_size) @Ts.function def ptx_global_to_shared_copy( @@ -329,7 +329,7 @@ def postproc_if_missing_async_support(): support_async = major >= 8 func_name = "tvm_callback_cuda_postproc" - prev_postproc = tvm_tirx.get_global_func(func_name, allow_missing=True) + prev_postproc = tvm_ffi.get_global_func(func_name, allow_missing=True) # Store the generated code prior to the post-processing. This # way, even though the generated code doesn't compile on platforms @@ -343,7 +343,7 @@ def get_original_code(): nonlocal original_code return original_code - @tvm_tirx.register_global_func(func_name, override=True) + @tvm_ffi.register_global_func(func_name, override=True) def tvm_callback_cuda_postproc(code, _): nonlocal original_code original_code = code @@ -365,7 +365,7 @@ def tvm_callback_cuda_postproc(code, _): if prev_postproc is None: tvm_ffi.registry.remove_global_func(func_name) else: - tvm_tirx.register_global_func(func_name, prev_postproc, override=True) + tvm_ffi.register_global_func(func_name, prev_postproc, override=True) @pytest.mark.gpu diff --git a/tests/python/script/test_script_printer_entry.py b/tests/python/script/test_script_printer_entry.py index 6d33e9176703..0d23cdcc9fc7 100644 --- a/tests/python/script/test_script_printer_entry.py +++ b/tests/python/script/test_script_printer_entry.py @@ -36,7 +36,7 @@ def test_script_entry_points(entry): config = PrinterConfig(extra_config={"tirx.prefix": "IR"}) assert render(tvm_tirx.ir.Range(0, 4), config) == "IR.Range(0, 4)" fallback_config = PrinterConfig(path_to_underline=[AccessPath.root().attr("missing")]) - assert render(tvm_tirx.runtime.ShapeTuple([1, 2]), fallback_config) == ( + assert render(tvm_ffi.Shape([1, 2]), fallback_config) == ( "Access path: .missing\n" "Note: No visible object for this path is rendered in TVMScript.\n\n" "Shape(1, 2)" diff --git a/tests/python/target/test_target_target.py b/tests/python/target/test_target_target.py index a9c649b5114b..48d37b0b98bd 100644 --- a/tests/python/target/test_target_target.py +++ b/tests/python/target/test_target_target.py @@ -401,7 +401,7 @@ def test_target_from_device_vulkan(input_form): def run_and_check(): dev = tvm_tirx.vulkan() target = Target.from_device("vulkan" if input_form == "string" else dev) - f_get_target_property = tvm_tirx.get_global_func("device_api.vulkan.get_target_property") + f_get_target_property = tvm_ffi.get_global_func("device_api.vulkan.get_target_property") assert target.kind.name == "vulkan" assert target.attrs["max_threads_per_block"] == dev.max_threads_per_block assert int(target.attrs["max_shared_memory_per_block"]) == dev.max_shared_memory_per_block diff --git a/tests/python/te/test_te_create_function.py b/tests/python/te/test_te_create_function.py index 5986d2a93a21..8cc09cf0788b 100644 --- a/tests/python/te/test_te_create_function.py +++ b/tests/python/te/test_te_create_function.py @@ -20,6 +20,7 @@ import numpy as np import pytest +import tvm_ffi import tvm import tvm_tirx @@ -554,7 +555,7 @@ def f_combine(x, y): rhs = tvm_tirx.tirx.Select((x[1] >= y[1]), x[1], y[1]) return lhs, rhs - def f_identity(dtype0: tvm_tirx.DataType, dtype1: tvm_tirx.DataType): + def f_identity(dtype0: tvm_ffi.dtype, dtype1: tvm_ffi.dtype): return tvm_tirx.tirx.const(-1, dtype0), tvm.te.min_value(ty=dtype1) argmax = te.comm_reducer(f_combine, f_identity, name="argmax") @@ -607,7 +608,7 @@ def f_combine(x, y): rhs = tvm_tirx.tirx.Select((x[0] >= y[0]), x[1], y[1]) return lhs, rhs - def f_identity(dtype0: tvm_tirx.DataType, dtype1: tvm_tirx.DataType): + def f_identity(dtype0: tvm_ffi.dtype, dtype1: tvm_ffi.dtype): return tvm.te.min_value(ty=dtype0), tvm_tirx.tirx.const(-1, dtype1) argmax = te.comm_reducer(f_combine, f_identity, name="argmax") diff --git a/tests/python/tirx/base/test_tirx_constructor.py b/tests/python/tirx/base/test_tirx_constructor.py index 09950c131e12..a4fd3e38c98f 100644 --- a/tests/python/tirx/base/test_tirx_constructor.py +++ b/tests/python/tirx/base/test_tirx_constructor.py @@ -89,8 +89,8 @@ def test_expr_constructor(): assert x.a == a assert x.b.same_as(b) - a = tvm_tirx.runtime.convert(tvm_tirx.tirx.Var("x", "int32") > 1) - b = tvm_tirx.runtime.convert(tvm_tirx.tirx.Var("x", "int32") == 1) + a = tvm_ffi.convert(tvm_tirx.tirx.Var("x", "int32") > 1) + b = tvm_ffi.convert(tvm_tirx.tirx.Var("x", "int32") == 1) for cls in [tvm_tirx.tirx.And, tvm_tirx.tirx.Or]: x = cls(a, b) diff --git a/tests/python/tirx/codegen/test_codegen_nvshmem.py b/tests/python/tirx/codegen/test_codegen_nvshmem.py index 85e6308cc177..5f4d6f3af3ed 100644 --- a/tests/python/tirx/codegen/test_codegen_nvshmem.py +++ b/tests/python/tirx/codegen/test_codegen_nvshmem.py @@ -21,11 +21,11 @@ import numpy as np import pytest +import tvm_ffi import tvm_tirx import tvm_tirx.testing from tvm.runtime import disco as di -from tvm_tirx.runtime import ShapeTuple from tvm_tirx.script import tirx as T from tvm_tirx.support.popen_pool import PopenWorker from tvm_tirx.testing import env @@ -49,7 +49,7 @@ def run_function(sess, function, *args): def create_nvshmem_array(sess, shape, dtype, init_data_fn=None, zero_out=True): """Create and optionally initialize an nvshmem-accessible DNDArray.""" nvshmem_empty = sess.get_global_func("runtime.disco.nvshmem.empty") - arr = nvshmem_empty(ShapeTuple(shape), dtype, None) + arr = nvshmem_empty(tvm_ffi.Shape(shape), dtype, None) if init_data_fn: for i in range(NUM_WORKERS): @@ -69,7 +69,7 @@ def test_codegen_nvshmem(): def _test_func(): ############ setup ############ sess = di.ProcessSession(num_workers=NUM_WORKERS) - f_init_nvshmem_uid = tvm_tirx.get_global_func("runtime.disco.nvshmem.init_nvshmem_uid") + f_init_nvshmem_uid = tvm_ffi.get_global_func("runtime.disco.nvshmem.init_nvshmem_uid") uid = f_init_nvshmem_uid() init_dfunc = sess.get_global_func("runtime.disco.nvshmem.init_nvshmem") init_dfunc(uid, NUM_WORKERS, 0) diff --git a/tests/python/tirx/script/test_tirx_script_ir_builder.py b/tests/python/tirx/script/test_tirx_script_ir_builder.py index c8d84e2ab9f3..3ad62e1fe92d 100644 --- a/tests/python/tirx/script/test_tirx_script_ir_builder.py +++ b/tests/python/tirx/script/test_tirx_script_ir_builder.py @@ -297,7 +297,7 @@ def test_flat_tensor_store_preserves_identity_and_emits_once(layout): def _check_stores(output): - outputs = output if isinstance(output, tuple | list | tvm_tirx.ir.Array) else [output] + outputs = output if isinstance(output, tuple | list | tvm_ffi.Array) else [output] stores = [] for tensor in outputs: @@ -627,7 +627,7 @@ def build(i, old, new): expected = [np.array([0, 2, 0, 3], dtype="int32")] inputs = [te.placeholder(array.shape, str(array.dtype)) for array in arrays] outputs = build(*inputs) - outputs = list(outputs) if isinstance(outputs, list | tuple | tvm_tirx.ir.Array) else [outputs] + outputs = list(outputs) if isinstance(outputs, list | tuple | tvm_ffi.Array) else [outputs] executable = tvm_tirx.tirx.compile(te.create_function([*inputs, *outputs]), target="llvm") actual = [ tvm_tirx.runtime.empty(tuple(int(dim) for dim in out.shape), out.dtype) for out in outputs diff --git a/tirx/python/tvm_tirx/__init__.py b/tirx/python/tvm_tirx/__init__.py index 8244906d9c9c..a6e35fd6c689 100644 --- a/tirx/python/tvm_tirx/__init__.py +++ b/tirx/python/tvm_tirx/__init__.py @@ -37,17 +37,16 @@ def _is_config_mode(): import os # ffi module must load first - from tvm_ffi import register_object, register_global_func, get_global_func + import tvm_ffi as _tvm_ffi # top-level alias from .libinfo import __version__ from .base import _RUNTIME_ONLY # tvm.runtime - from .runtime import Object - from .runtime._tensor import device, cpu, cuda, opencl, vulkan, metal + from .runtime._tensor import cpu, cuda, opencl, vulkan, metal from .runtime._tensor import vpi, rocm, ext_dev, hexagon - from .runtime import DataType, DataTypeCode + from .runtime import DataTypeCode # tvm.ir from .ir import IRModule, transform diff --git a/tirx/python/tvm_tirx/backend/cuda/cpp/asm.py b/tirx/python/tvm_tirx/backend/cuda/cpp/asm.py index 9b4e558e4837..2571abc03607 100644 --- a/tirx/python/tvm_tirx/backend/cuda/cpp/asm.py +++ b/tirx/python/tvm_tirx/backend/cuda/cpp/asm.py @@ -27,6 +27,8 @@ import re +import tvm_ffi + from tvm_tirx import ir from tvm_tirx.backend.cuda.op import func_call from tvm_tirx.runtime import MemoryScope @@ -447,8 +449,6 @@ def codegen_s_tir_cp_async_raw(*args): # PTX plain form for both .ca and .cg is just # ``cp.async..shared.global [dst], [src], cp_size;`` — three # operands, no trailing src-size / cache-policy. - from tvm_tirx import DataType - dst_ptr_in, dst_offset, src_ptr_in, src_offset, cp_size = args[:5] predicate = args[5] if len(args) == 6 else -1 cp_size_v = int(cp_size) @@ -467,7 +467,7 @@ def _elem_bytes(ptr): et = ta.element_type if not hasattr(et, "dtype"): return 1 - bits = DataType(str(et.dtype)).bits + bits = tvm_ffi.dtype(str(et.dtype)).bits assert bits % 8 == 0, f"non-byte element dtype: {et.dtype}" return bits // 8 diff --git a/tirx/python/tvm_tirx/backend/cuda/cpp/descriptors.py b/tirx/python/tvm_tirx/backend/cuda/cpp/descriptors.py index 182a261faffe..995db8a58f4e 100644 --- a/tirx/python/tvm_tirx/backend/cuda/cpp/descriptors.py +++ b/tirx/python/tvm_tirx/backend/cuda/cpp/descriptors.py @@ -563,7 +563,7 @@ def _get_tcgen05_mma_scale_vec_size(kind, scale_dtype): def _dtype_name(dtype) -> str: - """Name of a dtype given as a string, a DataType, or an object holding one.""" + """Name of a dtype given as a string, a tvm_ffi.dtype, or an object holding one.""" dtype_obj = getattr(dtype, "dtype", None) if dtype_obj is not None: return str(dtype_obj) diff --git a/tirx/python/tvm_tirx/backend/cuda/lang/alloc_pool.py b/tirx/python/tvm_tirx/backend/cuda/lang/alloc_pool.py index f47d6a125d81..d1ff8bad2eb9 100644 --- a/tirx/python/tvm_tirx/backend/cuda/lang/alloc_pool.py +++ b/tirx/python/tvm_tirx/backend/cuda/lang/alloc_pool.py @@ -21,7 +21,7 @@ import functools import operator -from tvm_tirx import DataType +import tvm_ffi # --------------------------------------------------------------------------- # ir_builder helpers — imported lazily to avoid circular deps at module level @@ -114,11 +114,11 @@ def _validate_mma_alloc_shape(shape, dtype, swizzle_mode): if not (isinstance(rows, int) and isinstance(cols, int)): return - dtype_bytes = DataType(dtype).bits // 8 + dtype_bytes = tvm_ffi.dtype(dtype).bits // 8 if dtype_bytes == 0: # Sub-byte dtype (e.g. float4); ``cols`` is already in element units, so # use a fractional check expressed via bits. - col_bits = cols * DataType(dtype).bits + col_bits = cols * tvm_ffi.dtype(dtype).bits atom_bits = _swizzle_atom_bytes(swizzle_mode) * 8 if col_bits < atom_bits or col_bits % atom_bits != 0: row_bytes = col_bits // 8 if col_bits % 8 == 0 else col_bits / 8 @@ -130,7 +130,8 @@ def _validate_mma_alloc_shape(shape, dtype, swizzle_mode): f"swizzle atom selected by {swizzle_mode.name}. " f"Use swizzle_mode=SwizzleMode.{suggestion}, or widen shape[-1] " f"to a multiple of " - f"{(atom_bits + DataType(dtype).bits - 1) // DataType(dtype).bits} elements." + f"{(atom_bits + tvm_ffi.dtype(dtype).bits - 1) // tvm_ffi.dtype(dtype).bits} " + f"elements." ) else: row_bytes = cols * dtype_bytes @@ -227,7 +228,7 @@ def alloc( # Advance in bits then round up to bytes so sub-byte dtypes (e.g. # float4_e2m1fn = 4 bits) still bump the cursor instead of leaving it # at 0 (bits // 8) and silently overlapping the next allocation. - self.offset += (_shape_product(shape) * DataType(dtype).bits + 7) // 8 + self.offset += (_shape_product(shape) * tvm_ffi.dtype(dtype).bits + 7) // 8 if self._owns_buffer: self.max_offset = max(self.max_offset, self.offset) return res diff --git a/tirx/python/tvm_tirx/backend/cuda/layout.py b/tirx/python/tvm_tirx/backend/cuda/layout.py index 20c0ee9b9f8d..f5df2a488aab 100644 --- a/tirx/python/tvm_tirx/backend/cuda/layout.py +++ b/tirx/python/tvm_tirx/backend/cuda/layout.py @@ -19,7 +19,8 @@ import copy from enum import Enum -import tvm_tirx +import tvm_ffi + from tvm_tirx.tirx.layout import Axis, ComposeLayout, Iter, Layout, R, S, TileLayout # ============================================================================ @@ -311,11 +312,10 @@ def _chk_rows(rows): layout = TileLayout(S[(m, K) : (1 @ tlane, 1 @ tcol)]) else: raise ValueError(f"tmem_mma_operand_layout: operand must be 'A' or 'D', got {operand!r}") + # PTX matrix D keeps sub-word accumulators in the low bits of a b32 cell. # A remains packed. TCol strides are expressed in tensor element units. - from tvm_tirx.runtime import DataType - - if operand == "D" and DataType(dtype).bits == 16: + if operand == "D" and tvm_ffi.dtype(dtype).bits == 16: layout = TileLayout.from_iters( shard=[ Iter(i.extent, i.stride * (2 if i.axis.name == "TCol" else 1), i.axis) @@ -339,7 +339,7 @@ class SwizzleMode(Enum): def mma_atom_layout(dtype: str, swizzle_mode: SwizzleMode | int) -> ComposeLayout: """Generate the MMA-compatible shared-memory atom layout.""" - bits = tvm_tirx.DataType(dtype).bits + bits = tvm_ffi.dtype(dtype).bits if isinstance(swizzle_mode, int): swizzle_mode = SwizzleMode(swizzle_mode) per_element = (128 // bits).bit_length() - 1 @@ -351,7 +351,7 @@ def mma_atom_layout(dtype: str, swizzle_mode: SwizzleMode | int) -> ComposeLayou def mma_atom_shape(dtype: str, swizzle_mode: SwizzleMode | int, shape: list[int] | None = None): """Generate the MMA-compatible shared-memory atom shape.""" - bits = tvm_tirx.DataType(dtype).bits + bits = tvm_ffi.dtype(dtype).bits if isinstance(swizzle_mode, int): swizzle_mode = SwizzleMode(swizzle_mode) atom_shape = { @@ -376,7 +376,7 @@ def mma_shared_layout(dtype: str, swizzle_mode: SwizzleMode | int, shape) -> Lay if swizzle_mode == SwizzleMode.SWIZZLE_NONE: # No-swizzle MMA smem is the packed-16B atom: offset e16*m + # M*e16*(k//e16) + k%e16 (e16=128/bits); plain tile if K % e16 != 0. - bits = tvm_tirx.DataType(dtype).bits + bits = tvm_ffi.dtype(dtype).bits e16 = 128 // bits m, k = int(shape[-2]), int(shape[-1]) if k % e16 == 0: diff --git a/tirx/python/tvm_tirx/backend/cuda/op.py b/tirx/python/tvm_tirx/backend/cuda/op.py index 7bbee1efe4b8..d21a8952e811 100644 --- a/tirx/python/tvm_tirx/backend/cuda/op.py +++ b/tirx/python/tvm_tirx/backend/cuda/op.py @@ -25,7 +25,7 @@ import tvm_ffi import tvm_tirx -from tvm_tirx import DataType, tirx +from tvm_tirx import tirx from tvm_tirx.ir import Attrs, Call, Op, StringImm, const from tvm_tirx.ir.location import UNKNOWN_LOC, Location from tvm_tirx.ir.op import _init_op_api, _make_op_api @@ -383,7 +383,7 @@ def ptx_cp_async_legacy( ``elem_dtype`` scales element offsets independently of the call's result ``ty``. """ - dtype = DataType(elem_dtype) + dtype = tvm_ffi.dtype(elem_dtype) elem_bits = dtype.bits * dtype.lanes dst_ptr = ptr_byte_offset(dst_ptr, dst_offset * elem_bits // 8, ty=dst_ptr.ty) src_ptr = ptr_byte_offset(src_ptr, src_offset * elem_bits // 8, ty=src_ptr.ty) diff --git a/tirx/python/tvm_tirx/backend/cuda/tile/patterns/async_copy.py b/tirx/python/tvm_tirx/backend/cuda/tile/patterns/async_copy.py index 3e2552141973..4ae853aed59d 100644 --- a/tirx/python/tvm_tirx/backend/cuda/tile/patterns/async_copy.py +++ b/tirx/python/tvm_tirx/backend/cuda/tile/patterns/async_copy.py @@ -28,9 +28,10 @@ before mapa. Layout strides are elements; chunk_bytes and src_size are bytes. """ +import tvm_ffi + from tvm_tirx.backend.cuda.ptx import PTXNamespace from tvm_tirx.ir import Bind, Evaluate, Var -from tvm_tirx.runtime import DataType from tvm_tirx.script import tirx as T from .common import Context, scalar @@ -52,7 +53,7 @@ def lower(opcode, entry, modifiers, args, params, sctx): ("dst", args[0], ("shared", "shared.dyn")), ("src", args[1], ("shared", "shared.dyn") if bulk else ("global",)), ): - bits = DataType(region.source.ty.dtype.dtype).bits + bits = tvm_ffi.dtype(region.source.ty.dtype.dtype).bits if size * 8 % bits: raise ValueError(f"{name}: chunk_bytes must contain whole storage elements") elements = size * 8 // bits diff --git a/tirx/python/tvm_tirx/backend/cuda/tile/patterns/common.py b/tirx/python/tvm_tirx/backend/cuda/tile/patterns/common.py index 272bb95092af..97797b7cda84 100644 --- a/tirx/python/tvm_tirx/backend/cuda/tile/patterns/common.py +++ b/tirx/python/tvm_tirx/backend/cuda/tile/patterns/common.py @@ -42,7 +42,7 @@ Var, ) from tvm_tirx.ir.prim import Let -from tvm_tirx.runtime import DataType, MemoryScope +from tvm_tirx.runtime import MemoryScope from tvm_tirx.script import tirx as T from tvm_tirx.sym import Analyzer from tvm_tirx.tirx.layout import TileLayout @@ -230,7 +230,7 @@ def __init__(self, ctx, name, region, atom_shape, scopes, repeat_axes, bounds): raise TypeError(f"{name}: expected a tensor region") self.ctx, self.name, self.region = ctx, name, region self.tensor, self.ty = region.source, region.source.ty - self.dtype = DataType(self.ty.dtype.dtype) + self.dtype = tvm_ffi.dtype(self.ty.dtype.dtype) self.bits = self.dtype.bits self.layout = self.ty.effective_layout() if self.layout is None: diff --git a/tirx/python/tvm_tirx/backend/cuda/tile/patterns/tma.py b/tirx/python/tvm_tirx/backend/cuda/tile/patterns/tma.py index a467872eb421..08a73682e05c 100644 --- a/tirx/python/tvm_tirx/backend/cuda/tile/patterns/tma.py +++ b/tirx/python/tvm_tirx/backend/cuda/tile/patterns/tma.py @@ -35,9 +35,11 @@ from itertools import product from math import prod +import tvm_ffi + from tvm_tirx.backend.cuda.ptx import PTXNamespace from tvm_tirx.ir import Range, TensorRegion, Tuple, Var -from tvm_tirx.runtime import DataType, MemoryScope +from tvm_tirx.runtime import MemoryScope from tvm_tirx.script import tirx as T from .common import Context, scalar @@ -62,7 +64,7 @@ def _descriptor_fields(fields, rank, box, default_dtype): force = 11 if dtype == "tf32" else -1 if force == 11: dtype = "float32" - bits = DataType(dtype).bits + bits = tvm_ffi.dtype(dtype).bits element_strides = tuple( int(scalar(x)) for x in sequence(fields.pop("element_strides", (1,) * rank)) ) @@ -98,7 +100,7 @@ def _tensor_map(ctx, region, rank, box, encoding): raise ValueError("TMA descriptor requires a global tensor region") tensor = region.source ty = tensor.ty - bits = DataType(ty.dtype.dtype).bits + bits = tvm_ffi.dtype(ty.dtype.dtype).bits packed = encoding[0] == "float4_e2m1fn" interleave = encoding[4] channels = (16 * interleave * 8 // bits) if interleave else 1 @@ -258,12 +260,12 @@ def lower(opcode, entry, modifiers, args, params, sctx): ) encoding = _descriptor_fields(params, rank, box, default_dtype) dtype, _, element_strides, swizzle, interleave, _, _ = encoding - bits = DataType(dtype).bits + bits = tvm_ffi.dtype(dtype).bits packed = dtype == "float4_e2m1fn" if packed and not load: raise ValueError("TMA 16U4_ALIGN16B supports global-to-shared loads only") if isinstance(global_, TensorRegion): - if bits != DataType(global_.source.ty.dtype.dtype).bits: + if bits != tvm_ffi.dtype(global_.source.ty.dtype.dtype).bits: raise ValueError("TMA descriptor format must match the global storage element width") tensor_map = _tensor_map(ctx, global_, rank, box, encoding) tmap = T.address_of(tensor_map) diff --git a/tirx/python/tvm_tirx/codegen/__init__.py b/tirx/python/tvm_tirx/codegen/__init__.py index 5396028ac1f0..7d1ddfa800d7 100644 --- a/tirx/python/tvm_tirx/codegen/__init__.py +++ b/tirx/python/tvm_tirx/codegen/__init__.py @@ -37,7 +37,7 @@ def build_module(mod, target): Returns ------- - module : runtime.Module + module : tvm_ffi.Module The corresponding module. """ target = Target(target) if isinstance(target, str) else target diff --git a/tirx/python/tvm_tirx/codegen/executable.py b/tirx/python/tvm_tirx/codegen/executable.py index 674b47c12761..98e10bd0ff02 100644 --- a/tirx/python/tvm_tirx/codegen/executable.py +++ b/tirx/python/tvm_tirx/codegen/executable.py @@ -59,7 +59,7 @@ def jit( The Executable returned by tvm.compile may not be directly runnable as they may contain CUDA source files and objects that are yet to be compiled and linked. - This function helps to create a runtime.Module for these cases. + This function helps to create a tvm_ffi.Module for these cases. Parameters ---------- @@ -94,7 +94,7 @@ def jit( if self._jitted_mod is not None and not force_recompile: return self._jitted_mod - # TODO(tvm-team): Update runtime.Module interface + # TODO(tvm-team): Update tvm_ffi.Module interface # to query these properties as bitmask. def _not_runnable(x): return x.kind in ("c", "static_library") @@ -190,7 +190,7 @@ def export_library( If fcompile has attribute object_format, will compile host library to that format. Otherwise, will use default format "o". - fpack_imports: function(mod: runtime.Module, is_system_lib: bool, symbol_prefix: str, + fpack_imports: function(mod: tvm_ffi.Module, is_system_lib: bool, symbol_prefix: str, workspace_dir: str) -> str Function used to pack imported modules from `mod` into a file suitable for passing to fcompile as an input file. The result can be a C source, or an .o object file, diff --git a/tirx/python/tvm_tirx/ir/__init__.py b/tirx/python/tvm_tirx/ir/__init__.py index a577eeeb6b0a..f51dae0f41d3 100644 --- a/tirx/python/tvm_tirx/ir/__init__.py +++ b/tirx/python/tvm_tirx/ir/__init__.py @@ -87,7 +87,5 @@ stmt_list, ) -from tvm_ffi import Array, Map - # Load passes after the common IR classes used by transformations are registered. from . import transform diff --git a/tirx/python/tvm_tirx/ir/attrs.py b/tirx/python/tvm_tirx/ir/attrs.py index b2136ec5404d..02fdecd563df 100644 --- a/tirx/python/tvm_tirx/ir/attrs.py +++ b/tirx/python/tvm_tirx/ir/attrs.py @@ -19,11 +19,9 @@ import tvm_ffi import tvm_ffi._ffi_api as _tvm_ffi_api -from tvm_tirx.runtime import Object - @tvm_ffi.register_object("ir.Attrs") -class Attrs(Object): +class Attrs(tvm_ffi.Object): """Attribute node, which is mainly use for defining attributes of operators. Used by function registered in python side, such as compute, schedule and alter_layout. diff --git a/tirx/python/tvm_tirx/ir/base.py b/tirx/python/tvm_tirx/ir/base.py index 07e4f4612260..ed7f76fdabd4 100644 --- a/tirx/python/tvm_tirx/ir/base.py +++ b/tirx/python/tvm_tirx/ir/base.py @@ -20,8 +20,6 @@ from tvm_ffi import register_object from tvm_ffi.serialization import from_json_graph_str, to_json_graph_str -from tvm_tirx.runtime import Object - from ..libinfo import __version__ from . import _ffi_api, json_compact @@ -44,12 +42,12 @@ def register_dialect_script_method(name, method): setattr(Scriptable, name, method) -class Node(Object): +class Node(tvm_ffi.Object): """Base class of all IR Nodes.""" @register_object("ir.EnvFunc") -class EnvFunc(Object): +class EnvFunc(tvm_ffi.Object): """Environment function. This is a global function object that can be serialized by its name. @@ -70,7 +68,7 @@ def get(name): return _ffi_api.EnvFuncGet(name) # type: ignore # pylint: disable=no-member -def load_json(json_str) -> Object: +def load_json(json_str) -> tvm_ffi.Object: """Load tvm object from json_str. Parameters @@ -80,7 +78,7 @@ def load_json(json_str) -> Object: Returns ------- - node : Object + node : tvm_ffi.Object The loaded tvm node. """ @@ -93,7 +91,7 @@ def save_json(node) -> str: Parameters ---------- - node : Object + node : tvm_ffi.Object A TVM object to be saved. Returns @@ -109,10 +107,10 @@ def assert_structural_equal(lhs, rhs, map_free_vars=False): Parameters ---------- - lhs : Object + lhs : tvm_ffi.Object The left operand. - rhs : Object + rhs : tvm_ffi.Object The left operand. map_free_vars : bool diff --git a/tirx/python/tvm_tirx/ir/expr.py b/tirx/python/tvm_tirx/ir/expr.py index 1165ccfabba9..d026dda51c71 100644 --- a/tirx/python/tvm_tirx/ir/expr.py +++ b/tirx/python/tvm_tirx/ir/expr.py @@ -27,7 +27,6 @@ import tvm_tirx from tvm_tirx.ir.location import UNKNOWN_LOC, Location -from ..runtime import Object from . import _ffi_api, _tensor_expr_overload from ._constant import const from .base import Node, Scriptable @@ -368,12 +367,12 @@ def __le__(self, other): def __eq__(self, other): if is_prim_expr(self): return _prim_expr_overload.__eq__(self, other) - return Object.__eq__(self, other) + return tvm_ffi.Object.__eq__(self, other) def __ne__(self, other): if is_prim_expr(self): return _prim_expr_overload.__ne__(self, other) - return Object.__ne__(self, other) + return tvm_ffi.Object.__ne__(self, other) def __gt__(self, other): if is_prim_expr(self): @@ -537,7 +536,7 @@ class FloatImm(Constant): value: float def __init__( - self, ty: str | tvm_tirx.DataType | PrimType, value: float, loc: Location = UNKNOWN_LOC + self, ty: str | tvm_ffi.dtype | PrimType, value: float, loc: Location = UNKNOWN_LOC ) -> None: if isinstance(ty, PrimType): ty = ty.dtype @@ -571,7 +570,7 @@ class IntImm(Constant): value: int def __init__( - self, ty: str | tvm_tirx.DataType | PrimType, value: int, loc: Location = UNKNOWN_LOC + self, ty: str | tvm_ffi.dtype | PrimType, value: int, loc: Location = UNKNOWN_LOC ) -> None: if isinstance(ty, PrimType): ty = ty.dtype @@ -618,15 +617,15 @@ class DataTypeImm(Constant): Parameters ---------- - value : str or tvm_tirx.DataType + value : str or tvm_ffi.dtype The represented data type. loc : Location, optional The source location of the literal. """ - value: tvm_tirx.DataType + value: tvm_ffi.dtype - def __init__(self, value: str | tvm_tirx.DataType, loc: Location = UNKNOWN_LOC) -> None: + def __init__(self, value: str | tvm_ffi.dtype, loc: Location = UNKNOWN_LOC) -> None: self.__init_handle_by_constructor__(_ffi_api.DataTypeImm, value, loc) @@ -798,7 +797,7 @@ def __init__( if isinstance(ty, str): ty = PtrType(PrimType("void")) if ty == "handle" else PrimType(ty) elif ty is not None: - ty = tvm_tirx.runtime.convert(ty) + ty = tvm_ffi.convert(ty) if not isinstance(ty, Type): raise TypeError("ty must be a Type or primitive dtype string") self.__init_handle_by_constructor__(_ffi_api.Var, name, ty, loc) @@ -810,7 +809,7 @@ def _lambda_type(annotation): return annotation if isinstance(annotation, str): return Var("", annotation).ty - if isinstance(annotation, tvm_tirx.DataType): + if isinstance(annotation, tvm_ffi.dtype): return tvm_tirx.ir.PrimType(annotation) if callable(annotation): value = annotation() @@ -829,7 +828,7 @@ def _lambda_result(value): elif isinstance(value, str): value = StringImm(value) else: - value = tvm_tirx.runtime.convert(value) + value = tvm_ffi.convert(value) if not isinstance(value, Expr): raise TypeError("Lambda body must be an Expr or a tuple/list of expressions") return value @@ -938,10 +937,10 @@ def from_min_extent(min_value: Expr, extent: Expr, loc: Location = UNKNOWN_LOC) """ return _ffi_api.RangeFromMinExtent(min_value, extent, loc) - def __eq__(self, other: Object) -> bool: + def __eq__(self, other: tvm_ffi.Object) -> bool: return tvm_ffi.structural_equal(self, other) - def __ne__(self, other: Object) -> bool: + def __ne__(self, other: tvm_ffi.Object) -> bool: return not self.__eq__(other) diff --git a/tirx/python/tvm_tirx/ir/function.py b/tirx/python/tvm_tirx/ir/function.py index c7814119fbdb..eb1614988205 100644 --- a/tirx/python/tvm_tirx/ir/function.py +++ b/tirx/python/tvm_tirx/ir/function.py @@ -22,7 +22,6 @@ import tvm_ffi import tvm_tirx.runtime -from tvm_tirx.runtime import Object from . import _ffi_api from .attrs import DictAttrs @@ -52,7 +51,7 @@ def with_attr(self, attr_key_or_dict, attr_value=None) -> "BaseFunc": attr_key_or_dict : Union[str, dict] The attribute key to use or a dict containing multiple key value pairs. - attr_value : Object + attr_value : tvm_ffi.Object The new attribute value. Returns @@ -70,21 +69,21 @@ def with_attr(self, attr_key_or_dict, attr_value=None) -> "BaseFunc": new_value = _ffi_api.BaseFuncWithAttr( new_value if new_value is self else new_value._move(), key, - tvm_tirx.runtime.convert(val), + tvm_ffi.convert(val), ) return new_value return _ffi_api.BaseFuncWithAttr( new_value if new_value is self else new_value._move(), attr_key_or_dict, - tvm_tirx.runtime.convert(attr_value), + tvm_ffi.convert(attr_value), ) - def with_attrs(self, attr_map: DictAttrs | dict[str, Object]) -> "BaseFunc": + def with_attrs(self, attr_map: DictAttrs | dict[str, tvm_ffi.Object]) -> "BaseFunc": """Copy the IRModule and add the given attribute map to it. Parameters ---------- - attr_map: Union[DictAttrs, Dict[str, Object]] + attr_map: Union[DictAttrs, Dict[str, tvm_ffi.Object]] The attribute map Returns ------- diff --git a/tirx/python/tvm_tirx/ir/global_info.py b/tirx/python/tvm_tirx/ir/global_info.py index 788a0d8a1149..c6991cdbf309 100644 --- a/tirx/python/tvm_tirx/ir/global_info.py +++ b/tirx/python/tvm_tirx/ir/global_info.py @@ -18,11 +18,9 @@ import tvm_ffi -from tvm_tirx.runtime import Object - @tvm_ffi.register_object("ir.GlobalInfo") -class GlobalInfo(Object): +class GlobalInfo(tvm_ffi.Object): """Base node for all global info that can appear in the IR""" def __eq__(self, other): diff --git a/tirx/python/tvm_tirx/ir/location.py b/tirx/python/tvm_tirx/ir/location.py index eeeb35d62c0e..5f3d7f962dd5 100644 --- a/tirx/python/tvm_tirx/ir/location.py +++ b/tirx/python/tvm_tirx/ir/location.py @@ -16,16 +16,16 @@ # under the License. """Source-location metadata and the shared unknown location.""" +import tvm_ffi from tvm_ffi import register_object from tvm_tirx.base import _RUNTIME_ONLY -from tvm_tirx.runtime import Object from . import _ffi_api @register_object("ir.SourceName") -class SourceName(Object): +class SourceName(tvm_ffi.Object): """An identifier for a source location. Parameters @@ -39,7 +39,7 @@ def __init__(self, name): @register_object("ir.Location") -class Location(Object): +class Location(tvm_ffi.Object): """Non-null source-location metadata attached to IR nodes. IR constructors default to the shared :data:`UNKNOWN_LOC`. diff --git a/tirx/python/tvm_tirx/ir/module.py b/tirx/python/tvm_tirx/ir/module.py index 3be29155a648..58577213609e 100644 --- a/tirx/python/tvm_tirx/ir/module.py +++ b/tirx/python/tvm_tirx/ir/module.py @@ -18,8 +18,6 @@ import tvm_ffi -from tvm_tirx.runtime import Object - from . import _ffi_api from . import expr as _expr from .attrs import DictAttrs, make_node @@ -241,7 +239,7 @@ def with_attr(self, attr_key, attr_value): attr_key : str The attribute key. - attr_value : Object + attr_value : tvm_ffi.Object The new attribute value. Returns @@ -266,11 +264,11 @@ def without_attr(self, attr_key: str) -> "IRModule": return _ffi_api.IRModuleWithoutAttr(self, attr_key) - def with_attrs(self, attr_map: DictAttrs | dict[str, Object]) -> "IRModule": + def with_attrs(self, attr_map: DictAttrs | dict[str, tvm_ffi.Object]) -> "IRModule": """Copy the IRModule and add the given attribute map to it. Parameters ---------- - attr_map: Union[DictAttrs, Dict[str, Object]] + attr_map: Union[DictAttrs, Dict[str, tvm_ffi.Object]] The attribute map Returns ------- diff --git a/tirx/python/tvm_tirx/ir/stmt.py b/tirx/python/tvm_tirx/ir/stmt.py index 2197ffd480e1..98ffef95d0d7 100644 --- a/tirx/python/tvm_tirx/ir/stmt.py +++ b/tirx/python/tvm_tirx/ir/stmt.py @@ -24,13 +24,12 @@ from tvm_tirx.ir import DictAttrs, Expr, Location, Op, Scriptable, StringImm, Var, make_node from tvm_tirx.ir.location import UNKNOWN_LOC -from tvm_tirx.runtime import Object from . import _ffi_api @tvm_ffi.register_object("ir.Stmt") -class Stmt(Object, Scriptable): +class Stmt(tvm_ffi.Object, Scriptable): """Base class of all the statements.""" @@ -146,7 +145,7 @@ class For(Stmt): body : Stmt | Sequence[Stmt] The body statement. - annotations: Optional[Mapping[str, Object]] + annotations: Optional[Mapping[str, tvm_ffi.Object]] Additional loop annotations interpreted by the consuming dialect. step : Expr @@ -161,7 +160,7 @@ class For(Stmt): extent: Expr kind: ForKind body: "SeqStmt" - annotations: Mapping[str, Object] + annotations: Mapping[str, tvm_ffi.Object] step: Expr | None loc: Location @@ -172,7 +171,7 @@ def __init__( extent: Expr, kind: ForKind, body: Stmt | Sequence[Stmt], - annotations: Mapping[str, Object] | None = None, + annotations: Mapping[str, tvm_ffi.Object] | None = None, step: Expr | None = None, loc: Location = UNKNOWN_LOC, ) -> None: diff --git a/tirx/python/tvm_tirx/ir/transform/core.py b/tirx/python/tvm_tirx/ir/transform/core.py index e85100035ce1..a013f7d0fe2a 100644 --- a/tirx/python/tvm_tirx/ir/transform/core.py +++ b/tirx/python/tvm_tirx/ir/transform/core.py @@ -22,13 +22,11 @@ import tvm_ffi -import tvm_tirx.runtime - from . import _ffi_transform_api @tvm_ffi.register_object("transform.PassInfo") -class PassInfo(tvm_tirx.runtime.Object): +class PassInfo(tvm_ffi.Object): """The class contains the meta data required by a pass. It is the container of information needed by running an optimization or analysis. This class can be extended by adding new members when more meta data is @@ -48,7 +46,7 @@ def __init__(self, opt_level, name): @tvm_ffi.register_object("transform.PassInstrument") -class PassInstrument(tvm_tirx.runtime.Object): +class PassInstrument(tvm_ffi.Object): """A pass instrument implementation. To use, a user class can either subclass from PassInstrument @@ -248,7 +246,7 @@ def create_pass_instrument(pi_cls): @tvm_ffi.register_object("transform.PassContext") -class PassContext(tvm_tirx.runtime.Object): +class PassContext(tvm_ffi.Object): """The basis where a TVM optimization/analysis runs on. Each pass context contains a number of auxiliary information that is used to help an optimization pass. Such information includes the error reporter @@ -336,7 +334,7 @@ def list_configs(): @tvm_ffi.register_object("transform.Pass") -class Pass(tvm_tirx.runtime.Object): +class Pass(tvm_ffi.Object): """The base class of all passes. All methods here are just simple wrappers that are implemented in the backend. They are defined for users to conveniently interact with the base class. diff --git a/tirx/python/tvm_tirx/ir/unique_name_supply.py b/tirx/python/tvm_tirx/ir/unique_name_supply.py index 98cd20433ed8..6b11dc4f5b21 100644 --- a/tirx/python/tvm_tirx/ir/unique_name_supply.py +++ b/tirx/python/tvm_tirx/ir/unique_name_supply.py @@ -18,13 +18,11 @@ import tvm_ffi -from tvm_tirx import Object - from . import _ffi_api @tvm_ffi.register_object("ir.UniqueNameSupply") -class UniqueNameSupply(Object): +class UniqueNameSupply(tvm_ffi.Object): """UniqueNameSupply that can be used to generate unique names. Parameters diff --git a/tirx/python/tvm_tirx/runtime/__init__.py b/tirx/python/tvm_tirx/runtime/__init__.py index aa8f098b9d56..2722d8f6c141 100644 --- a/tirx/python/tvm_tirx/runtime/__init__.py +++ b/tirx/python/tvm_tirx/runtime/__init__.py @@ -17,8 +17,8 @@ # under the License. """TVM runtime namespace.""" -from tvm_ffi import convert, Object, Module, load_module, system_lib -from tvm_ffi._dtype import dtype as DataType, DataTypeCode +from tvm_ffi import load_module +from tvm_ffi._dtype import DataTypeCode # class exposures from .device import Device @@ -26,12 +26,10 @@ from ._tensor import Tensor, tensor, empty # function exposures -from ._tensor import device, cpu, cuda, opencl, vulkan, metal -from ._tensor import vpi, rocm, ext_dev, from_dlpack +from ._tensor import cpu, cuda, opencl, vulkan, metal +from ._tensor import vpi, rocm, ext_dev from . import _ffi_api -from tvm_ffi import Shape as ShapeTuple - def enabled(target): """Whether module runtime is enabled for target diff --git a/tirx/python/tvm_tirx/runtime/_tensor.py b/tirx/python/tvm_tirx/runtime/_tensor.py index af3f79b4adf3..c7cd312f1b6a 100644 --- a/tirx/python/tvm_tirx/runtime/_tensor.py +++ b/tirx/python/tvm_tirx/runtime/_tensor.py @@ -30,26 +30,13 @@ ml_dtypes = None import tvm_ffi -from tvm_ffi import DLDeviceType, device +from tvm_ffi import DLDeviceType from tvm_tirx.runtime import Device from . import _ffi_api -def from_dlpack(ext_tensor): - """ - Convert an external tensor to an Tensor. - - Parameters - ---------- - ext_tensor : object - The external tensor to convert. - - """ - return tvm_ffi.from_dlpack(ext_tensor) - - @tvm_ffi.register_object("ffi.Tensor") class Tensor(tvm_ffi.core.Tensor): """Lightweight Tensor class of TVM runtime. @@ -360,7 +347,7 @@ def cpu(dev_id=0): dev : Device The created device """ - return device(DLDeviceType.kDLCPU, dev_id) + return tvm_ffi.device(DLDeviceType.kDLCPU, dev_id) def cuda(dev_id=0): @@ -376,7 +363,7 @@ def cuda(dev_id=0): dev : Device The created device """ - return device(DLDeviceType.kDLCUDA, dev_id) + return tvm_ffi.device(DLDeviceType.kDLCUDA, dev_id) def rocm(dev_id=0): @@ -392,7 +379,7 @@ def rocm(dev_id=0): dev : Device The created device """ - return device(DLDeviceType.kDLROCM, dev_id) + return tvm_ffi.device(DLDeviceType.kDLROCM, dev_id) def opencl(dev_id=0): @@ -408,7 +395,7 @@ def opencl(dev_id=0): dev : Device The created device """ - return device(DLDeviceType.kDLOpenCL, dev_id) + return tvm_ffi.device(DLDeviceType.kDLOpenCL, dev_id) def metal(dev_id=0): @@ -424,7 +411,7 @@ def metal(dev_id=0): dev : Device The created device """ - return device(DLDeviceType.kDLMetal, dev_id) + return tvm_ffi.device(DLDeviceType.kDLMetal, dev_id) def vpi(dev_id=0): @@ -440,7 +427,7 @@ def vpi(dev_id=0): dev : Device The created device """ - return device(DLDeviceType.kDLVPI, dev_id) + return tvm_ffi.device(DLDeviceType.kDLVPI, dev_id) def vulkan(dev_id=0): @@ -456,7 +443,7 @@ def vulkan(dev_id=0): dev : Device The created device """ - return device(DLDeviceType.kDLVulkan, dev_id) + return tvm_ffi.device(DLDeviceType.kDLVulkan, dev_id) def ext_dev(dev_id=0): @@ -477,7 +464,7 @@ def ext_dev(dev_id=0): This API is reserved for quick testing of new device by plugin device API as ext_dev. """ - return device(DLDeviceType.kDLExtDev, dev_id) + return tvm_ffi.device(DLDeviceType.kDLExtDev, dev_id) def hexagon(dev_id=0): @@ -493,7 +480,7 @@ def hexagon(dev_id=0): dev : Device The created device """ - return device(DLDeviceType.kDLHexagon, dev_id) + return tvm_ffi.device(DLDeviceType.kDLHexagon, dev_id) def webgpu(dev_id=0): @@ -509,7 +496,7 @@ def webgpu(dev_id=0): dev : Device The created device """ - return device(DLDeviceType.kDLWebGPU, dev_id) + return tvm_ffi.device(DLDeviceType.kDLWebGPU, dev_id) # Register back to FFI diff --git a/tirx/python/tvm_tirx/script/ir_builder/base.py b/tirx/python/tvm_tirx/script/ir_builder/base.py index 47b4bd756a70..ddc32545fd04 100644 --- a/tirx/python/tvm_tirx/script/ir_builder/base.py +++ b/tirx/python/tvm_tirx/script/ir_builder/base.py @@ -22,12 +22,12 @@ from inspect import signature from typing import Any, Generic, ParamSpec, TypeVar +import tvm_ffi from tvm_ffi import register_object as _register_object from tvm_ffi.dataclasses import MISSING as MISSING from tvm_tirx import ir from tvm_tirx.ir.location import UnknownLoc -from tvm_tirx.runtime import Object as _Object from . import _ffi_api @@ -99,7 +99,7 @@ def invoke(*args, **kwargs): @_register_object("script.ir_builder.IRBuilderFrame") -class IRBuilderFrame(_Object): +class IRBuilderFrame(tvm_ffi.Object): """A stack frame of the IRBuilder used to keep track of the current scope. A language variant supplies frame subclasses that retain the information @@ -132,7 +132,7 @@ def add_callback(self, callback: Callable[[], None]) -> None: @_register_object("script.ir_builder.IRBuilder") -class IRBuilder(_Object): +class IRBuilder(tvm_ffi.Object): """A shared construction context for any language variant's IR. Examples @@ -199,7 +199,7 @@ def is_in_scope() -> bool: """ return _ffi_api.IRBuilderIsInScope() # type: ignore[attr-defined] # pylint: disable=no-member - def get(self) -> _Object: + def get(self) -> tvm_ffi.Object: """Get the constructed IR.""" return _ffi_api.IRBuilderGet(self) # type: ignore[attr-defined] # pylint: disable=no-member @@ -417,7 +417,7 @@ def at(loc: LocationEntry | ir.Location | None, value: _T) -> _T: target = value.value if isinstance(value, AlreadyEmitted) else value targets = target if isinstance(target, list | tuple) else (target,) for item in targets: - if isinstance(item, _Object): + if isinstance(item, tvm_ffi.Object): IRBuilder.current()._attach_loc(item, loc) return value diff --git a/tirx/python/tvm_tirx/script/ir_builder/ir.py b/tirx/python/tvm_tirx/script/ir_builder/ir.py index 1153a8d85bfb..24ee0209d50c 100644 --- a/tirx/python/tvm_tirx/script/ir_builder/ir.py +++ b/tirx/python/tvm_tirx/script/ir_builder/ir.py @@ -18,11 +18,11 @@ from typing import NoReturn, TypeVar -from tvm_tirx import DataType +import tvm_ffi + from tvm_tirx.ir import DataTypeImm, GlobalInfo, LambdaExpr, Location, Var from tvm_tirx.ir.location import UNKNOWN_LOC from tvm_tirx.ir.type import _init_type_api -from tvm_tirx.runtime import Object as tvm_Object from . import _ffi_api from .base import IRBuilder, LocationEntry @@ -31,8 +31,8 @@ T = TypeVar("T") -def dtype(value: str | DataType) -> DataTypeImm: - """Construct a DataType-valued constant for an ordinary operation argument. +def dtype(value: str | tvm_ffi.dtype) -> DataTypeImm: + """Construct a tvm_ffi.dtype-valued constant for an ordinary operation argument. Unlike scalar constructors such as ``T.float32``, this constructs a dtype value, rather than a scalar value or variable of that dtype. @@ -117,7 +117,7 @@ def meta_var(value: T) -> T: return value -def module_attrs(attrs: dict[str, tvm_Object], allow_overwrite=False) -> None: +def module_attrs(attrs: dict[str, tvm_ffi.Object], allow_overwrite=False) -> None: """Specify the attrs of the ir_module frame. Parameters ---------- @@ -129,7 +129,7 @@ def module_attrs(attrs: dict[str, tvm_Object], allow_overwrite=False) -> None: return _ffi_api.ModuleAttrs(attrs, allow_overwrite) # type: ignore[attr-defined] # pylint: disable=no-member -def module_get_attr(attr_key: str) -> tvm_Object | None: +def module_get_attr(attr_key: str) -> tvm_ffi.Object | None: """Get the specified attr of the ir_module frame. Parameters ---------- @@ -144,7 +144,7 @@ def module_get_attr(attr_key: str) -> tvm_Object | None: def module_set_attr( - attr_key: str, attr_value: tvm_Object | None, allow_overwrite: bool = False + attr_key: str, attr_value: tvm_ffi.Object | None, allow_overwrite: bool = False ) -> None: """Set the specified attr of the ir_module frame. Parameters diff --git a/tirx/python/tvm_tirx/script/printer/doc.py b/tirx/python/tvm_tirx/script/printer/doc.py index 775dd3331125..d2c6c47060d9 100644 --- a/tirx/python/tvm_tirx/script/printer/doc.py +++ b/tirx/python/tvm_tirx/script/printer/doc.py @@ -20,17 +20,17 @@ from enum import IntEnum, unique from typing import Union +import tvm_ffi from tvm_ffi import register_object from tvm_ffi.access_path import AccessPath from tvm_tirx.ir import FloatImm, IntImm -from tvm_tirx.runtime import Object from . import _ffi_api @register_object("script.printer.Doc") -class Doc(Object): +class Doc(tvm_ffi.Object): """Base class of all Docs""" diff --git a/tirx/python/tvm_tirx/script/printer/scriptable.py b/tirx/python/tvm_tirx/script/printer/scriptable.py index a54d2ccf3540..c6823ccbadaf 100644 --- a/tirx/python/tvm_tirx/script/printer/scriptable.py +++ b/tirx/python/tvm_tirx/script/printer/scriptable.py @@ -19,17 +19,17 @@ import os from collections.abc import Sequence +import tvm_ffi from tvm_ffi import register_object from tvm_ffi.access_path import AccessPath from tvm_tirx.ir.base import register_script_printer -from tvm_tirx.runtime import Object from . import _ffi_api @register_object("script.PrinterConfig") -class PrinterConfig(Object): +class PrinterConfig(tvm_ffi.Object): """Configuration of TVMScript printer. The ``int_dtype`` and ``float_dtype`` arguments populate the string-valued @@ -64,8 +64,8 @@ def __init__( extra_config: dict | None = None, path_to_underline: list[AccessPath] | None = None, path_to_annotate: dict[AccessPath, str] | None = None, - obj_to_underline: list[Object] | None = None, - obj_to_annotate: dict[Object, str] | None = None, + obj_to_underline: list[tvm_ffi.Object] | None = None, + obj_to_annotate: dict[tvm_ffi.Object, str] | None = None, ) -> None: if num_context_lines is None: num_context_lines = -1 @@ -95,7 +95,7 @@ def __init__( ) -def _script(obj: Object, config: PrinterConfig) -> str: +def _script(obj: tvm_ffi.Object, config: PrinterConfig) -> str: return _ffi_api.Script(obj, config) # type: ignore # pylint: disable=no-member @@ -117,8 +117,8 @@ def script( extra_config: dict | None = None, path_to_underline: list[AccessPath] | None = None, path_to_annotate: dict[AccessPath, str] | None = None, - obj_to_underline: list[Object] | None = None, - obj_to_annotate: dict[Object, str] | None = None, + obj_to_underline: list[tvm_ffi.Object] | None = None, + obj_to_annotate: dict[tvm_ffi.Object, str] | None = None, ) -> str: """Print TVM IR into TVMScript text format @@ -153,9 +153,9 @@ def script( Object path to be underlined path_to_annotate : Optional[Dict[AccessPath, str]] = None Object path to be annotated - obj_to_underline : Optional[List[Object]] = None + obj_to_underline : Optional[List[tvm_ffi.Object]] = None Object to be underlined - obj_to_annotate : Optional[Dict[Object, str]] = None + obj_to_annotate : Optional[Dict[tvm_ffi.Object, str]] = None Object to be annotated Returns @@ -208,8 +208,8 @@ def show( extra_config: dict | None = None, path_to_underline: list[AccessPath] | None = None, path_to_annotate: dict[AccessPath, str] | None = None, - obj_to_underline: list[Object] | None = None, - obj_to_annotate: dict[Object, str] | None = None, + obj_to_underline: list[tvm_ffi.Object] | None = None, + obj_to_annotate: dict[tvm_ffi.Object, str] | None = None, ) -> None: """A sugar for print highlighted TVM script. @@ -265,9 +265,9 @@ def show( Object path to be underlined path_to_annotate : Optional[Dict[AccessPath, str]] = None Object path to be annotated - obj_to_underline : Optional[List[Object]] = None + obj_to_underline : Optional[List[tvm_ffi.Object]] = None Object to be underlined - obj_to_annotate : Optional[Dict[Object, str]] = None + obj_to_annotate : Optional[Dict[tvm_ffi.Object, str]] = None Object to be annotated """ @@ -308,7 +308,7 @@ def _repr(self) -> str: PrinterConfig(extra_config={"ir.comment_imports": True}), ) except Exception: - return Object.__repr__(self) + return tvm_ffi.Object.__repr__(self) # The printer supplies its explicit display capability to existing IR wrappers. diff --git a/tirx/python/tvm_tirx/support/rocm.py b/tirx/python/tvm_tirx/support/rocm.py index 962b99978778..545ea96b0a9e 100644 --- a/tirx/python/tvm_tirx/support/rocm.py +++ b/tirx/python/tvm_tirx/support/rocm.py @@ -174,7 +174,7 @@ def callback_rocm_bitcode_path(rocdl_dir=None): elif "isa_version" not in n and n not in {"irif"}: raise RuntimeError("could not find bitcode " + n) - return tvm_tirx.runtime.convert(bitcode_files) + return tvm_ffi.convert(bitcode_files) def parse_compute_version(compute_version): diff --git a/tirx/python/tvm_tirx/sym/analyzer.py b/tirx/python/tvm_tirx/sym/analyzer.py index 2925d4f25dfe..7bfb2252f819 100644 --- a/tirx/python/tvm_tirx/sym/analyzer.py +++ b/tirx/python/tvm_tirx/sym/analyzer.py @@ -23,7 +23,6 @@ import tvm_ffi from tvm_tirx import ir -from tvm_tirx.runtime import Object from tvm_tirx.sym import IntSet from . import _ffi_api @@ -53,7 +52,7 @@ class CompareResult(enum.IntEnum): @tvm_ffi.register_object("sym.ModularSet") -class ModularSet(Object): +class ModularSet(tvm_ffi.Object): """Represent range of (coeff * x + base) for x in Z""" def __init__(self, coeff, base): @@ -61,7 +60,7 @@ def __init__(self, coeff, base): @tvm_ffi.register_object("sym.ConstIntBound") -class ConstIntBound(Object): +class ConstIntBound(tvm_ffi.Object): """Represent constant integer bound Parameters @@ -105,7 +104,7 @@ def __exit__(self, ptype, value, trace): @tvm_ffi.register_object("sym.Analyzer") -class Analyzer(Object): +class Analyzer(tvm_ffi.Object): """Integer arithmetic analyzer This is a stateful analyzer class that can be used to perform diff --git a/tirx/python/tvm_tirx/sym/int_set.py b/tirx/python/tvm_tirx/sym/int_set.py index 67d28c4c439a..15ea5633347c 100644 --- a/tirx/python/tvm_tirx/sym/int_set.py +++ b/tirx/python/tvm_tirx/sym/int_set.py @@ -18,13 +18,11 @@ import tvm_ffi -from tvm_tirx.runtime import Object - from . import _ffi_api @tvm_ffi.register_object("sym.IntSet") -class IntSet(Object): +class IntSet(tvm_ffi.Object): """Represent a set of integer in one dimension.""" def is_nothing(self): diff --git a/tirx/python/tvm_tirx/sym/iter_affine_map.py b/tirx/python/tvm_tirx/sym/iter_affine_map.py index 1c8112c570e3..7865e5daa843 100644 --- a/tirx/python/tvm_tirx/sym/iter_affine_map.py +++ b/tirx/python/tvm_tirx/sym/iter_affine_map.py @@ -21,7 +21,6 @@ import tvm_ffi from tvm_tirx.ir import Expr -from tvm_tirx.runtime import Object from . import _ffi_api @@ -32,7 +31,7 @@ class IterMapExpr(Expr): @tvm_ffi.register_object("sym.IterMark") -class IterMark(Object): +class IterMark(tvm_ffi.Object): """Mark the source as an iterator in [0, extent). Parameters @@ -95,7 +94,7 @@ def __init__(self, args, base): @tvm_ffi.register_object("sym.IterMapResult") -class IterMapResult(Object): +class IterMapResult(tvm_ffi.Object): """Result of iter map detection.""" diff --git a/tirx/python/tvm_tirx/target/detect_target.py b/tirx/python/tvm_tirx/target/detect_target.py index 4d32ec8d36da..c20847fe2779 100644 --- a/tirx/python/tvm_tirx/target/detect_target.py +++ b/tirx/python/tvm_tirx/target/detect_target.py @@ -18,9 +18,10 @@ from collections.abc import Callable +import tvm_ffi from tvm_ffi import get_global_func -from ..runtime import Device, device +from ..runtime import Device from . import Target @@ -67,7 +68,7 @@ def detect_target_from_device(dev: str | Device) -> Target: The detected target. """ if isinstance(dev, str): - dev = device(dev) + dev = tvm_ffi.device(dev) device_type = Device._DEVICE_TYPE_TO_NAME[dev.dlpack_device_type()] if device_type not in SUPPORTED_DEVICE: raise ValueError( diff --git a/tirx/python/tvm_tirx/target/target.py b/tirx/python/tvm_tirx/target/target.py index 3724502d9002..b6598a3b9f56 100644 --- a/tirx/python/tvm_tirx/target/target.py +++ b/tirx/python/tvm_tirx/target/target.py @@ -20,13 +20,13 @@ from tvm_ffi import Map from tvm_ffi.core import String -from tvm_tirx.runtime import Device, Object, convert, device +from tvm_tirx.runtime import Device from . import _ffi_api @tvm_ffi.register_object("target.TargetKind") -class TargetKind(Object): +class TargetKind(tvm_ffi.Object): """Kind of a compilation target""" @property @@ -49,7 +49,7 @@ def __getattr__(self, name: str): @tvm_ffi.register_object("target.Target") -class Target(Object): +class Target(tvm_ffi.Object): """Target device information, use through TVM API. Targets can be constructed from: @@ -132,9 +132,9 @@ def __init__(self, target, host=None): same as target. """ if isinstance(target, dict | str): - target = convert(target) + target = tvm_ffi.convert(target) if isinstance(host, dict | str): - host = convert(host) + host = tvm_ffi.convert(host) if target is None or not isinstance(target, Map | String | Target | str): raise ValueError(f"target has to be a string or dictionary. instead get {type(target)}") if host is not None: @@ -259,4 +259,4 @@ def device_from_target(target, index=None): """ if not isinstance(target, Target): target = Target(target) - return device(target.get_target_device_type(), index) + return tvm_ffi.device(target.get_target_device_type(), index) diff --git a/tirx/python/tvm_tirx/testing/env.py b/tirx/python/tvm_tirx/testing/env.py index 6ce5d3cd34ec..2200d0817fc0 100644 --- a/tirx/python/tvm_tirx/testing/env.py +++ b/tirx/python/tvm_tirx/testing/env.py @@ -49,6 +49,8 @@ def test_my_cuda_kernel(): import functools import os +import tvm_ffi + import tvm_tirx __all__ = [ @@ -90,9 +92,9 @@ def test_my_cuda_kernel(): @functools.cache def _device_exists(kind: str, index: int = 0) -> bool: - """Return whether ``tvm_tirx.device(kind, index)`` is present and usable.""" + """Return whether ``tvm_ffi.device(kind, index)`` is present and usable.""" try: - return bool(tvm_tirx.device(kind, index).exist) + return bool(tvm_ffi.device(kind, index).exist) except Exception: # pylint: disable=broad-except # A missing backend / driver must skip the test, not crash collection. return False @@ -187,7 +189,7 @@ def _runtime_enabled(kind: str) -> bool: Used for kinds whose device existence does not imply the backend was compiled in -- notably ``llvm``, which maps to the always-present CPU - device, so ``tvm_tirx.device("llvm").exist`` is True even on ``USE_LLVM=OFF``. + device, so ``tvm_ffi.device("llvm").exist`` is True even on ``USE_LLVM=OFF``. """ try: return bool(tvm_tirx.runtime.enabled(kind)) @@ -286,7 +288,7 @@ def has_nccl() -> bool: """True if NCCL was built in and a CUDA device is present.""" try: return has_cuda() and ( - tvm_tirx.get_global_func("runtime.disco.nccl.init_ccl", allow_missing=True) is not None + tvm_ffi.get_global_func("runtime.disco.nccl.init_ccl", allow_missing=True) is not None ) except Exception: # pylint: disable=broad-except return False @@ -306,7 +308,7 @@ def has_nvshmem() -> bool: """ try: return has_cuda() and ( - tvm_tirx.get_global_func("runtime.disco.nvshmem.init_nvshmem_uid", allow_missing=True) + tvm_ffi.get_global_func("runtime.disco.nvshmem.init_nvshmem_uid", allow_missing=True) is not None ) except Exception: # pylint: disable=broad-except diff --git a/tirx/python/tvm_tirx/testing/utils.py b/tirx/python/tvm_tirx/testing/utils.py index 970799095b52..1ddc894ce31f 100644 --- a/tirx/python/tvm_tirx/testing/utils.py +++ b/tirx/python/tvm_tirx/testing/utils.py @@ -71,6 +71,7 @@ def test_cuda_vectorize_add(): import ml_dtypes import numpy as np import pytest +import tvm_ffi import tvm_tirx @@ -143,7 +144,7 @@ def device_enabled(target): if not env._runtime_enabled(kind): continue if kind == "cuda" and "cudnn" in tvm_tirx.target.Target(candidate).attrs.get("libs", []): - exists = tvm_tirx.get_global_func("tvm.contrib.cudnn.exists", allow_missing=True) + exists = tvm_ffi.get_global_func("tvm.contrib.cudnn.exists", allow_missing=True) if env.has_cudnn() and exists is not None and exists(): return True elif tvm_tirx.device_from_target(candidate).exist: diff --git a/tirx/python/tvm_tirx/tirx/_compile.py b/tirx/python/tvm_tirx/tirx/_compile.py index f88c114ee286..1465c409c4a6 100644 --- a/tirx/python/tvm_tirx/tirx/_compile.py +++ b/tirx/python/tvm_tirx/tirx/_compile.py @@ -18,6 +18,7 @@ # pylint: disable=invalid-name """Native TIRx compilation and target code generation.""" +import tvm_ffi from tvm_ffi import DLDeviceType import tvm_tirx @@ -129,12 +130,12 @@ def is_host_function(f): return host_mod, device_mod_dict -def _build_for_target(mod: IRModule, target: Target) -> tvm_tirx.runtime.Module: +def _build_for_target(mod: IRModule, target: Target) -> tvm_ffi.Module: """Build a runtime module from an IRModule and a Target.""" if tvm_tirx.ir.transform.PassContext.current().config.get("tirx.disable_assert", False): mod = tvm_tirx.tirx.transform.SkipAssert()(mod) build_f_name = "target.build." + target.kind.name - bf = tvm_tirx.get_global_func(build_f_name) + bf = tvm_ffi.get_global_func(build_f_name) if bf is None: raise ValueError(f"{build_f_name} is not enabled") return bf(mod, target) @@ -142,7 +143,7 @@ def _build_for_target(mod: IRModule, target: Target) -> tvm_tirx.runtime.Module: def _codegen_host_and_device( host_mod: IRModule, device_mod_dict: dict[tuple[Target, str], IRModule], target_host: Target -) -> tvm_tirx.runtime.Module: +) -> tvm_ffi.Module: """Compile each device group and attach it to the host runtime module.""" # Get the first module to get the attributes @@ -252,4 +253,4 @@ def compile( # pylint: disable=redefined-builtin # C++ compiler consumers link the generated runtime module directly. -tvm_tirx.register_global_func("tirx.compile", lambda *args, **kwargs: compile(*args, **kwargs).mod) +tvm_ffi.register_global_func("tirx.compile", lambda *args, **kwargs: compile(*args, **kwargs).mod) diff --git a/tirx/python/tvm_tirx/tirx/_tensor_view.py b/tirx/python/tvm_tirx/tirx/_tensor_view.py index a8d100d172b9..01c2ff4245a3 100644 --- a/tirx/python/tvm_tirx/tirx/_tensor_view.py +++ b/tirx/python/tvm_tirx/tirx/_tensor_view.py @@ -23,6 +23,8 @@ from numbers import Integral from typing import TYPE_CHECKING +import tvm_ffi + import tvm_tirx from tvm_tirx.runtime import MemoryScope @@ -164,9 +166,9 @@ def _infer_shape(shape): ) return shape - if len(args) == 1 and isinstance(args[0], str | tvm_tirx.DataType) and not kwargs: - cast_dtype = tvm_tirx.DataType(args[0]) - cur_dtype = tvm_tirx.DataType(buf.dtype) + if len(args) == 1 and isinstance(args[0], str | tvm_ffi.dtype) and not kwargs: + cast_dtype = tvm_ffi.dtype(args[0]) + cur_dtype = tvm_ffi.dtype(buf.dtype) storage_axis = "TCol" if _is_tmem_scope(buf.scope()) else "m" if cast_dtype.bits > cur_dtype.bits: assert cast_dtype.bits % cur_dtype.bits == 0 @@ -425,7 +427,7 @@ def _tmem_element_offset_to_column_offset(buf: Var, element_offset): if element_offset is None: return None - bits = tvm_tirx.DataType(buf.dtype).bits + bits = tvm_ffi.dtype(buf.dtype).bits bit_offset = element_offset * bits bit_offset_c = _concrete_int(bit_offset) if bit_offset_c is not None: diff --git a/tirx/python/tvm_tirx/tirx/function.py b/tirx/python/tvm_tirx/tirx/function.py index daa9b20d3056..724d5c99572e 100644 --- a/tirx/python/tvm_tirx/tirx/function.py +++ b/tirx/python/tvm_tirx/tirx/function.py @@ -25,7 +25,6 @@ import tvm_tirx.runtime from tvm_tirx.ir import BaseFunc, Scriptable, SeqStmt, Stmt, Var from tvm_tirx.ir.location import UNKNOWN_LOC, Location -from tvm_tirx.runtime import Object from . import _ffi_api from .expr import Expr @@ -67,7 +66,7 @@ def __init__( ret_type = tvm_tirx.ir.Type.missing() param_list = [] for x in params: - x = tvm_tirx.runtime.convert(x) if not isinstance(x, Object) else x + x = tvm_ffi.convert(x) if not isinstance(x, tvm_ffi.Object) else x if not isinstance(x, Var): raise TypeError("params can only contain Var or Var") param_list.append(x) diff --git a/tirx/python/tvm_tirx/tirx/index_map.py b/tirx/python/tvm_tirx/tirx/index_map.py index 0ddc1755202c..9c7323b8136c 100644 --- a/tirx/python/tvm_tirx/tirx/index_map.py +++ b/tirx/python/tvm_tirx/tirx/index_map.py @@ -26,14 +26,13 @@ import tvm_tirx from tvm_tirx.ir import Expr, Range, Var -from tvm_tirx.runtime import Object from tvm_tirx.runtime._tensor import Tensor from . import _ffi_api @tvm_ffi.register_object("tirx.IndexMap") -class IndexMap(Object): +class IndexMap(tvm_ffi.Object): """A mapping from multi-dimensional indices to another set of multi-dimensional indices Parameters diff --git a/tirx/python/tvm_tirx/tirx/layout.py b/tirx/python/tvm_tirx/tirx/layout.py index 3a20a07e6f75..057ead1e3160 100644 --- a/tirx/python/tvm_tirx/tirx/layout.py +++ b/tirx/python/tvm_tirx/tirx/layout.py @@ -24,7 +24,6 @@ from tvm_ffi.dataclasses import Enum import tvm_tirx -from tvm_tirx.runtime import Object from tvm_tirx.tirx.expr import Expr from . import _ffi_api @@ -32,7 +31,7 @@ @tvm_ffi.register_object("tirx.Layout") -class Layout(Object): +class Layout(tvm_ffi.Object): def verify_well_formed(self) -> bool: """Verify if the layout is well-formed. @@ -504,7 +503,7 @@ def _to_offset_expr(x: _OffsetExprLike) -> _OffsetExpr: @tvm_ffi.register_object("tirx.Iter") -class Iter(Object): +class Iter(tvm_ffi.Object): """A memory layout that tiles data across devices.""" extent: Expr diff --git a/tirx/python/tvm_tirx/tirx/op/__init__.py b/tirx/python/tvm_tirx/tirx/op/__init__.py index c65dac08a9dc..15f6e662b75a 100644 --- a/tirx/python/tvm_tirx/tirx/op/__init__.py +++ b/tirx/python/tvm_tirx/tirx/op/__init__.py @@ -662,7 +662,7 @@ def stack_make_dltensor( call : Expr The call expression. """ - if isinstance(arr_dtype, str | tvm_tirx.DataType | tvm_tirx.ir.PrimType): + if isinstance(arr_dtype, str | tvm_ffi.dtype | tvm_tirx.ir.PrimType): arr_dtype = const(0, dtype=arr_dtype) return call_intrin( @@ -1373,7 +1373,7 @@ def get_vscale_expr(dtype: str | tvm_ffi.dtype, min_size: int = 128) -> Expr: Parameters ---------- - dtype : Union[str, tvm_ffi.DataType] + dtype : Union[str, tvm_ffi.dtype] Element data type. min_size : int The minimum size of the scalable vector in bits. diff --git a/tirx/python/tvm_tirx/tirx/script/ir_builder/external_kernel.py b/tirx/python/tvm_tirx/tirx/script/ir_builder/external_kernel.py index f5c53c4994c0..0a4e1490080e 100644 --- a/tirx/python/tvm_tirx/tirx/script/ir_builder/external_kernel.py +++ b/tirx/python/tvm_tirx/tirx/script/ir_builder/external_kernel.py @@ -28,7 +28,6 @@ from tvm_tirx import libinfo, tirx from tvm_tirx.backend.config import backend_config_json, merge_backend_configs from tvm_tirx.ir import Expr, PtrType, is_prim_expr -from tvm_tirx.runtime import Module class BaseKernel: # pylint: disable=too-few-public-methods @@ -36,7 +35,7 @@ class BaseKernel: # pylint: disable=too-few-public-methods def compile_to_device_module( self, launch_args, *args, **kwargs - ) -> tuple[str, Module, list[Any]]: + ) -> tuple[str, tvm_ffi.Module, list[Any]]: """Compile the kernel to a device module.""" raise NotImplementedError() @@ -92,7 +91,7 @@ def _create_cuda_module( Returns ------- - kernel_module : Module + kernel_module : tvm_ffi.Module The CUDA module. """ tvm_metadata = self._format_tvm_module_metadata( @@ -124,7 +123,7 @@ def compile_to_device_module( # pylint: disable=arguments-differ grid: list[list[int | tirx.Expr]], *args: list[Any], **kwargs: dict[str, Any], - ) -> tuple[str, Module, list[Any]]: + ) -> tuple[str, tvm_ffi.Module, list[Any]]: """Compile the kernel to a device module.""" kernel_name = kwargs["kernel_name"] assert len(grid) == 2, ( @@ -234,7 +233,7 @@ def call_kernel( ) # Attach the kernel module to the current IRModule - external_mods: list[Module] = module_get_attr("external_mods") or [] + external_mods: list[tvm_ffi.Module] = module_get_attr("external_mods") or [] kernel_exists = any([mod.implements_function(kernel_name) for mod in external_mods]) if kernel_exists: logging.debug("Kernel %s already exists in the IRModule", kernel_name) diff --git a/tirx/python/tvm_tirx/tirx/script/ir_builder/ir.py b/tirx/python/tvm_tirx/tirx/script/ir_builder/ir.py index 2affe6d25cb5..a79838854e88 100644 --- a/tirx/python/tvm_tirx/tirx/script/ir_builder/ir.py +++ b/tirx/python/tvm_tirx/tirx/script/ir_builder/ir.py @@ -23,8 +23,10 @@ from collections.abc import Callable from typing import TYPE_CHECKING, Any, TypeVar +import tvm_ffi + import tvm_tirx -from tvm_tirx import DataType, ir +from tvm_tirx import ir from tvm_tirx import ir as _ir from tvm_tirx import tirx as tir @@ -91,7 +93,7 @@ def _get_elem_offset(elem_offset, byte_offset, dtype: str): return elem_offset if byte_offset is None: return None - return byte_offset * 8 // (DataType(_normalize_prim_type(dtype).dtype).bits) + return byte_offset * 8 // (tvm_ffi.dtype(_normalize_prim_type(dtype).dtype).bits) _meta_construction_state = threading.local() @@ -243,7 +245,9 @@ def normalize(value): args = [] if placement is not None: if not isinstance(placement, ir.Tuple): - placement = placement if isinstance(placement, list | tuple | ir.Array) else [placement] + placement = ( + placement if isinstance(placement, list | tuple | tvm_ffi.Array) else [placement] + ) placement = ir.Tuple([normalize(value) for value in placement]) args.append(placement) return ir.Call( diff --git a/tirx/python/tvm_tirx/tirx/script/ir_builder/parser_protocol.py b/tirx/python/tvm_tirx/tirx/script/ir_builder/parser_protocol.py index 8ef324f80200..7f297df198e0 100644 --- a/tirx/python/tvm_tirx/tirx/script/ir_builder/parser_protocol.py +++ b/tirx/python/tvm_tirx/tirx/script/ir_builder/parser_protocol.py @@ -28,6 +28,8 @@ from collections.abc import Sequence from typing import Any +import tvm_ffi + import tvm_tirx # isort: off @@ -257,7 +259,7 @@ def bind_( """ name_loc = loc if name_loc is None else name_loc if frame_value: - if isinstance(value, _python.list | _python.tuple | _ir.Array): + if isinstance(value, _python.list | _python.tuple | tvm_ffi.Array): for index, item in enumerate(value): bind_( item, @@ -385,7 +387,7 @@ def emit_(value: Any, *, loc: _Loc = None) -> None: or getattr(type(value), "_is_meta_class", False) ): return - if isinstance(value, list | tuple | _ir.Array): + if isinstance(value, list | tuple | tvm_ffi.Array): for item in value: emit_(item, loc=loc) return @@ -450,7 +452,7 @@ def tensor_store( """ from tvm_tirx.sym import Analyzer # pylint: disable=import-outside-toplevel - if not isinstance(indices, list | tuple | _ir.Array): + if not isinstance(indices, list | tuple | tvm_ffi.Array): indices = [indices] expr_indices = [] diff --git a/tirx/python/tvm_tirx/tirx/script/ir_builder/triton.py b/tirx/python/tvm_tirx/tirx/script/ir_builder/triton.py index 33ae1d87f4f2..d7dbfe6ae72a 100644 --- a/tirx/python/tvm_tirx/tirx/script/ir_builder/triton.py +++ b/tirx/python/tvm_tirx/tirx/script/ir_builder/triton.py @@ -20,13 +20,13 @@ from typing import Any import triton +import tvm_ffi from packaging import version from triton.runtime.jit import type_canonicalisation_dict import tvm_tirx from tvm_tirx import tirx from tvm_tirx.ir import PrimType, PtrType, is_prim_expr -from tvm_tirx.runtime import Module from tvm_tirx.sym.utils import get_const_int from .external_kernel import BaseKernel @@ -55,7 +55,7 @@ def compile_to_device_module( launch_args: list[int | tirx.Expr], *args: list[Any], **kwargs: dict[str, Any], - ) -> tuple[str, Module, list[Any]]: + ) -> tuple[str, tvm_ffi.Module, list[Any]]: """Compile the kernel to a device module. Parameters diff --git a/tirx/python/tvm_tirx/tirx/tensor.py b/tirx/python/tvm_tirx/tirx/tensor.py index f497854fdd1f..642636dc066a 100644 --- a/tirx/python/tvm_tirx/tirx/tensor.py +++ b/tirx/python/tvm_tirx/tirx/tensor.py @@ -100,7 +100,7 @@ def __init__( if elem_offset is not None and byte_offset is not None: raise ValueError("elem_offset and byte_offset cannot be set at the same time") if byte_offset is not None: - elem_offset = byte_offset * 8 // tvm_tirx.DataType(dtype.dtype).bits + elem_offset = byte_offset * 8 // tvm_ffi.dtype(dtype.dtype).bits if data_alignment is None: data_alignment = align elif align not in (0, data_alignment): @@ -220,7 +220,7 @@ def view(self, expr, *args, **kwargs) -> "Var": Supported signatures are ``view(*shape, layout=None)``, where shape can contain ``-1`` to indicate that the dimension size is auto-inferred, and - ``view(dtype: Union[str, tvm_tirx.DataType])``. + ``view(dtype: Union[str, tvm_ffi.dtype])``. Returns ------- @@ -284,7 +284,7 @@ def _dtype(self, expr): def _byte_offset(self, expr): _check_tensor_property_receiver(self, expr) - return self.elem_offset * tvm_tirx.DataType(self.dtype).bits // 8 + return self.elem_offset * tvm_ffi.dtype(self.dtype).bits // 8 def _sub(self, expr): _check_tensor_property_receiver(self, expr) diff --git a/tirx/python/tvm_tirx/tirx/tile/dispatch.py b/tirx/python/tvm_tirx/tirx/tile/dispatch.py index 8da1f65e873a..48e476d70022 100644 --- a/tirx/python/tvm_tirx/tirx/tile/dispatch.py +++ b/tirx/python/tvm_tirx/tirx/tile/dispatch.py @@ -18,12 +18,13 @@ Mirrors the C++ header ``include/tvm/tirx/tile/context.h``. """ -# pylint: disable=no-member +import tvm_ffi + +# pylint: disable=no-member from tvm_ffi import register_object from tvm_tirx.ir import Expr, Range, Scriptable, Stmt, Var -from tvm_tirx.runtime import Object from tvm_tirx.target import Target from .. import _ffi_api @@ -31,7 +32,7 @@ @register_object("tirx.DispatchContext") -class DispatchContext(Object, Scriptable): +class DispatchContext(tvm_ffi.Object, Scriptable): """DispatchContext node. Parameters @@ -48,10 +49,10 @@ class DispatchContext(Object, Scriptable): var_range_map : Dict[Var, Range] A map from loop variables to their ranges. - callbacks : Dict[str, Object] + callbacks : Dict[str, tvm_ffi.Object] The callbacks of the dispatch context. - shared_state : Dict[str, Object] + shared_state : Dict[str, tvm_ffi.Object] Shared state persisting across dispatch calls within a single lowering pass. """ @@ -60,8 +61,8 @@ class DispatchContext(Object, Scriptable): launch_params: dict[str, tuple[Var, Expr]] var_range_map: dict[Var, Range] alloc_only: bool - callbacks: dict[str, Object] - shared_state: dict[str, Object] + callbacks: dict[str, tvm_ffi.Object] + shared_state: dict[str, tvm_ffi.Object] inter: dict[str, list] intra: dict[str, list] scope_kind: str @@ -78,8 +79,8 @@ def __init__( launch_params: dict[str, tuple[Var, Expr]], var_range_map: dict[Var, Range], alloc_only: bool = False, - callbacks: dict[str, Object] = {}, - shared_state: dict[str, Object] = {}, + callbacks: dict[str, tvm_ffi.Object] = {}, + shared_state: dict[str, tvm_ffi.Object] = {}, inter: dict[str, list] | None = None, intra: dict[str, list] | None = None, scope_kind: str = "", @@ -147,7 +148,7 @@ def add_post_tensor_def_stmt(self, tensor: Var, stmt: Stmt) -> None: """ _ffi_api.DispatchContextAddPostTensorDefStmt(self, tensor, stmt) # pylint: disable=no-member - def cache_get(self, key: str) -> Object | None: + def cache_get(self, key: str) -> tvm_ffi.Object | None: """Look up a cached value by key. Parameters @@ -157,19 +158,19 @@ def cache_get(self, key: str) -> Object | None: Returns ------- - Optional[Object] + Optional[tvm_ffi.Object] The cached value, or None on miss. """ return _ffi_api.DispatchContextSharedStateGet(self, key) - def cache_set(self, key: str, value: Object) -> None: + def cache_set(self, key: str, value: tvm_ffi.Object) -> None: """Store a value in the cross-dispatch cache. Parameters ---------- key : str Cache key (built by the caller from construction parameters). - value : Object + value : tvm_ffi.Object The object to cache (e.g. a Var or Var). """ _ffi_api.DispatchContextSharedStateSet(self, key, value) diff --git a/tirx/tests/python/backend/cuda/test_module_export.py b/tirx/tests/python/backend/cuda/test_module_export.py index f356ac39f181..9f4b75db9cfe 100644 --- a/tirx/tests/python/backend/cuda/test_module_export.py +++ b/tirx/tests/python/backend/cuda/test_module_export.py @@ -19,6 +19,7 @@ import numpy as np import pytest +import tvm_ffi import tvm_tirx import tvm_tirx.support.nvcc @@ -94,10 +95,10 @@ def run_and_check(): a_np = np.random.uniform(size=12).astype("float32") a_nd = tvm_tirx.runtime.tensor(a_np, dev) b_nd = tvm_tirx.runtime.tensor(a_np, dev) - syslibA = tvm_tirx.runtime.system_lib("modA_") - syslibB = tvm_tirx.runtime.system_lib("modB_") + syslibA = tvm_ffi.system_lib("modA_") + syslibB = tvm_ffi.system_lib("modB_") # reload same lib twice - syslibA = tvm_tirx.runtime.system_lib("modA_") + syslibA = tvm_ffi.system_lib("modA_") syslibA["my_inplace_update"](a_nd) syslibB["my_inplace_update"](b_nd) np.testing.assert_equal(a_nd.numpy(), a_np + 1) diff --git a/tirx/tests/python/backend/hexagon/codegen/test_basic.py b/tirx/tests/python/backend/hexagon/codegen/test_basic.py index 7e300f2e71c9..b46b83427bd8 100644 --- a/tirx/tests/python/backend/hexagon/codegen/test_basic.py +++ b/tirx/tests/python/backend/hexagon/codegen/test_basic.py @@ -19,6 +19,7 @@ from pathlib import Path import pytest +import tvm_ffi from tvm_ffi import register_global_func import tvm_tirx @@ -31,7 +32,7 @@ @pytest.fixture(autouse=True) def register_linker(): name = "tvm.contrib.hexagon.link_shared" - original_linker = tvm_tirx.get_global_func(name) + original_linker = tvm_ffi.get_global_func(name) def link_shared(so_name, _objs, _extra_args=None): # Codegen retains the output bytes, but these tests only inspect the diff --git a/tirx/tests/python/backend/metal/codegen/test_basic.py b/tirx/tests/python/backend/metal/codegen/test_basic.py index 6a3a01358d81..60dac4c89e1c 100644 --- a/tirx/tests/python/backend/metal/codegen/test_basic.py +++ b/tirx/tests/python/backend/metal/codegen/test_basic.py @@ -203,9 +203,9 @@ def func(A: T.Tensor((16), "float32"), B: T.Tensor((16), "float32"), x: T.float3 with T.launch_thread("threadIdx.x", 16) as i: B[i] = A[i] + x - previous_callback = tvm_tirx.get_global_func("tvm_callback_metal_compile", allow_missing=True) + previous_callback = tvm_ffi.get_global_func("tvm_callback_metal_compile", allow_missing=True) - @tvm_tirx.register_global_func("tvm_callback_metal_compile", override=True) + @tvm_ffi.register_global_func("tvm_callback_metal_compile", override=True) def compile_metal(src, target): return xcode.compile_metal(src) @@ -220,7 +220,7 @@ def compile_metal(src, target): if previous_callback is None: tvm_ffi.registry.remove_global_func("tvm_callback_metal_compile") else: - tvm_tirx.register_global_func( + tvm_ffi.register_global_func( "tvm_callback_metal_compile", previous_callback, override=True ) @@ -250,7 +250,7 @@ def kernel(): C_frag.data, 0, A_frag.data, 0, B_frag.data, 0, C_frag.data, 0 ) - metal_codegen = tvm_tirx.get_global_func("target.build.metal") + metal_codegen = tvm_ffi.get_global_func("target.build.metal") module = metal_codegen(Module, tvm_tirx.target.Target("metal")) source = module.inspect_source() @@ -261,7 +261,7 @@ def kernel(): def _build_metal(mod): - build = tvm_tirx.get_global_func("target.build.metal") + build = tvm_ffi.get_global_func("target.build.metal") return build(mod, tvm_tirx.target.Target("metal")) @@ -525,7 +525,7 @@ def kernel(): typed_buffer[0] = T.float16(1) void_buffer[0] = T.float16(2) - metal_codegen = tvm_tirx.get_global_func("target.build.metal") + metal_codegen = tvm_ffi.get_global_func("target.build.metal") module = metal_codegen(Module, tvm_tirx.target.Target("metal")) source = module.inspect_source() diff --git a/tirx/tests/python/backend/metal/test_compile_callback.py b/tirx/tests/python/backend/metal/test_compile_callback.py index dff1253a64a9..9fc04ceec019 100644 --- a/tirx/tests/python/backend/metal/test_compile_callback.py +++ b/tirx/tests/python/backend/metal/test_compile_callback.py @@ -47,7 +47,7 @@ def inspect_callback(src, target): seen["src"] = src return (src, "metal") - tvm_tirx.register_global_func("tvm_callback_metal_compile", inspect_callback, override=True) + tvm_ffi.register_global_func("tvm_callback_metal_compile", inspect_callback, override=True) try: f = tvm_tirx.tirx.compile(Module, target="metal") dev = tvm_tirx.metal() @@ -96,7 +96,7 @@ def mixed_callback(src, target): # Second kernel declares textual MSL, contradicting the metallib above. return (src, "metal") - tvm_tirx.register_global_func("tvm_callback_metal_compile", mixed_callback, override=True) + tvm_ffi.register_global_func("tvm_callback_metal_compile", mixed_callback, override=True) try: with pytest.raises(Exception, match="inconsistent formats"): tvm_tirx.tirx.compile(Module, target="metal") diff --git a/tirx/tests/python/backend/vulkan/codegen/test_basic.py b/tirx/tests/python/backend/vulkan/codegen/test_basic.py index e06ac3192695..be20d5e9ff44 100644 --- a/tirx/tests/python/backend/vulkan/codegen/test_basic.py +++ b/tirx/tests/python/backend/vulkan/codegen/test_basic.py @@ -20,6 +20,7 @@ import numpy as np import pytest +import tvm_ffi import tvm_tirx import tvm_tirx.testing @@ -527,7 +528,7 @@ def kernel(): T.evaluate(A_buf[0]) target = tvm_tirx.target.Target("vulkan") - vulkan_codegen = tvm_tirx.get_global_func("target.build.vulkan") + vulkan_codegen = tvm_ffi.get_global_func("target.build.vulkan") vulkan_codegen(AllocationBacked, target) @I.ir_module diff --git a/tirx/tests/python/backend/webgpu/codegen/test_basic.py b/tirx/tests/python/backend/webgpu/codegen/test_basic.py index 20bc5b7595f6..174fa2097b8d 100644 --- a/tirx/tests/python/backend/webgpu/codegen/test_basic.py +++ b/tirx/tests/python/backend/webgpu/codegen/test_basic.py @@ -18,6 +18,7 @@ import re import pytest +import tvm_ffi import tvm_tirx import tvm_tirx.testing @@ -43,7 +44,7 @@ def main(A: T.Tensor((8,), "float32"), B: T.Tensor((8,), "float32")): def _build_webgpu(mod, target="webgpu"): - build = tvm_tirx.get_global_func("target.build.webgpu") + build = tvm_ffi.get_global_func("target.build.webgpu") return build(mod, tvm_tirx.target.Target(target)) diff --git a/tirx/tests/python/codegen/gpu_common/test_allreduce.py b/tirx/tests/python/codegen/gpu_common/test_allreduce.py index 5a26cd1412da..14bce0a2b48a 100644 --- a/tirx/tests/python/codegen/gpu_common/test_allreduce.py +++ b/tirx/tests/python/codegen/gpu_common/test_allreduce.py @@ -110,11 +110,11 @@ def run_and_check(): @pytest.fixture def optional_metal_compile_callback(define_metal_compile_callback): name = "tvm_callback_metal_compile" - cached = tvm_tirx.get_global_func(name, allow_missing=True) + cached = tvm_ffi.get_global_func(name, allow_missing=True) if define_metal_compile_callback: - @tvm_tirx.register_global_func(name, override=True) + @tvm_ffi.register_global_func(name, override=True) def compile_metal(src, target): from tvm_tirx.support.xcode import ( compile_metal, # pylint: disable=import-outside-toplevel @@ -128,7 +128,7 @@ def compile_metal(src, target): if cached is None: tvm_ffi.registry.remove_global_func(name) else: - tvm_tirx.register_global_func(name, cached, override=True) + tvm_ffi.register_global_func(name, cached, override=True) @pytest.mark.gpu diff --git a/tirx/tests/python/codegen/gpu_common/test_device.py b/tirx/tests/python/codegen/gpu_common/test_device.py index 2ad2fb2722e9..a1f9982ea917 100644 --- a/tirx/tests/python/codegen/gpu_common/test_device.py +++ b/tirx/tests/python/codegen/gpu_common/test_device.py @@ -16,6 +16,7 @@ # under the License. import numpy as np import pytest +import tvm_ffi import tvm_tirx import tvm_tirx.testing @@ -46,7 +47,7 @@ def check_target(target): f = tvm_tirx.tirx.compile(Module, target=target) def run_and_check(): - dev = tvm_tirx.device(target_kind, 0) + dev = tvm_ffi.device(target_kind, 0) a = tvm_tirx.runtime.empty((12,), dtype="uint64", device=dev) f(a) assert a.numpy()[0] == value + 3 @@ -89,7 +90,7 @@ def check_target(device, host): n = 1027 def run_and_check(): - dev = tvm_tirx.device(device, 0) + dev = tvm_ffi.device(device, 0) a = tvm_tirx.runtime.tensor(np.random.uniform(size=n).astype("float32"), dev) b = tvm_tirx.runtime.tensor(np.random.uniform(size=()).astype("float32"), dev) d = tvm_tirx.runtime.tensor(np.zeros(n, dtype="float32"), dev) diff --git a/tirx/tests/python/codegen/gpu_common/test_intrinsics.py b/tirx/tests/python/codegen/gpu_common/test_intrinsics.py index 11cc3b6c9ce6..1726c5380aad 100644 --- a/tirx/tests/python/codegen/gpu_common/test_intrinsics.py +++ b/tirx/tests/python/codegen/gpu_common/test_intrinsics.py @@ -18,6 +18,7 @@ import numpy as np import pytest +import tvm_ffi import tvm_tirx import tvm_tirx.testing @@ -96,7 +97,7 @@ def check_device(device): f.export_library(path_dso, fcompile=cc.cross_compiler("g++")) def run_and_check(): - dev = tvm_tirx.device(device, 0) + dev = tvm_ffi.device(device, 0) def popen_check(): import tvm_tirx @@ -123,7 +124,7 @@ def check_c(device): f = tvm_tirx.tirx.compile(main, target=tvm_tirx.target.Target(device, host="c")) def run_and_check(): - dev = tvm_tirx.device(device, 0) + dev = tvm_ffi.device(device, 0) a = tvm_tirx.runtime.tensor(np.random.uniform(size=1024).astype("float32"), dev) b = tvm_tirx.runtime.tensor(np.zeros(1024, dtype="float32"), dev) f["main"](a, b) diff --git a/tirx/tests/python/codegen/test_common.py b/tirx/tests/python/codegen/test_common.py index 104138f92e43..f031cebc5810 100644 --- a/tirx/tests/python/codegen/test_common.py +++ b/tirx/tests/python/codegen/test_common.py @@ -17,6 +17,7 @@ import numpy as np import pytest +import tvm_ffi import tvm_tirx import tvm_tirx.testing @@ -153,7 +154,7 @@ def kernel(A_ptr: T.handle("float32", "global")): B[0] = T.float32(1) mod = tvm_tirx.IRModule({"kernel": kernel}) - build = tvm_tirx.get_global_func(f"target.build.{target}") + build = tvm_ffi.get_global_func(f"target.build.{target}") source = build(mod, tvm_tirx.target.Target(target)).inspect_source() assert f"{qualifier}float* B" in source assert f"(({qualifier}float*)" in source diff --git a/tirx/tests/python/codegen/test_extern.py b/tirx/tests/python/codegen/test_extern.py index b65f8534823b..44d68508e015 100644 --- a/tirx/tests/python/codegen/test_extern.py +++ b/tirx/tests/python/codegen/test_extern.py @@ -17,6 +17,7 @@ # ruff: noqa: F841 import numpy as np import pytest +import tvm_ffi import tvm_tirx import tvm_tirx.testing @@ -84,7 +85,7 @@ class Module: def main(A: T.Tensor((1024,), "float32"), C: T.Tensor((1024,), "float32")): T.evaluate(T.call_packed("my_extern_array_func1", A, C)) - @tvm_tirx.register_global_func + @tvm_ffi.register_global_func def my_extern_array_func1(aa, bb): aa.copyto(bb) diff --git a/tirx/tests/python/codegen/test_static_init.py b/tirx/tests/python/codegen/test_static_init.py index ad88e1e99c36..8900e2f4713c 100644 --- a/tirx/tests/python/codegen/test_static_init.py +++ b/tirx/tests/python/codegen/test_static_init.py @@ -17,6 +17,7 @@ import ctypes import numpy as np +import tvm_ffi import tvm_tirx from tvm_tirx.script import ir as I @@ -24,7 +25,7 @@ def test_static_init(): - @tvm_tirx.register_global_func("test_static_callback") + @tvm_ffi.register_global_func("test_static_callback") def test_cb(sh, A): assert isinstance(sh, ctypes.c_void_p) return sh diff --git a/tirx/tests/python/ir/test_attrs.py b/tirx/tests/python/ir/test_attrs.py index 10ffba20957b..c39a9c4e51dc 100644 --- a/tirx/tests/python/ir/test_attrs.py +++ b/tirx/tests/python/ir/test_attrs.py @@ -54,8 +54,8 @@ def test_attrs_equal(): dattr2 = tvm_tirx.ir.make_node("ir.DictAttrs", x=1, y=None) tvm_tirx.ir.assert_structural_equal(dattr0, dattr1) assert not tvm_ffi.structural_equal(dattr0, dattr2) - assert not tvm_ffi.structural_equal({"x": 1}, tvm_tirx.runtime.convert(1)) - assert not tvm_ffi.structural_equal([1, 2], tvm_tirx.runtime.convert(1)) + assert not tvm_ffi.structural_equal({"x": 1}, tvm_ffi.convert(1)) + assert not tvm_ffi.structural_equal([1, 2], tvm_ffi.convert(1)) def test_assert_structural_equal_reports_mismatch(): diff --git a/tirx/tests/python/ir/test_json.py b/tirx/tests/python/ir/test_json.py index 776005f69d5d..df5aec125002 100644 --- a/tirx/tests/python/ir/test_json.py +++ b/tirx/tests/python/ir/test_json.py @@ -169,7 +169,7 @@ def test_json_containers(): assert type(values[3]) is bool a = tvm_tirx.ir.Var("a", tvm_tirx.ir.PrimType("int32")) b = tvm_tirx.ir.Var("b", tvm_tirx.ir.PrimType("int32")) - objects = tvm_tirx.runtime.convert({a: 2, b: 3}) + objects = tvm_ffi.convert({a: 2, b: 3}) assert a in objects and b in dict(objects.items()) assert len(objects) == 2 assert tvm_tirx.ir.prim.Add(a, tvm_tirx.ir.IntImm("int32", 1)) not in objects diff --git a/tirx/tests/python/ir/test_type.py b/tirx/tests/python/ir/test_type.py index 5445401f7174..1e5b87497ee2 100644 --- a/tirx/tests/python/ir/test_type.py +++ b/tirx/tests/python/ir/test_type.py @@ -19,6 +19,7 @@ import uuid import pytest +import tvm_ffi import tvm_tirx from tvm_tirx.runtime import MemoryScope @@ -46,7 +47,7 @@ def test_prim_type(): def test_func_type(): - arg_types = tvm_tirx.runtime.convert([]) + arg_types = tvm_ffi.convert([]) ret_type = tvm_tirx.ir.PrimType("float32") tf = tvm_tirx.ir.FuncType(arg_types, ret_type) assert tf.arg_types == arg_types @@ -61,7 +62,7 @@ def test_func_type(): def test_tuple_type(): tf = tvm_tirx.ir.FuncType([], tvm_tirx.ir.TupleType([])) tt = tvm_tirx.ir.PrimType("float32") - fields = tvm_tirx.runtime.convert([tf, tt]) + fields = tvm_ffi.convert([tf, tt]) tup_ty = tvm_tirx.ir.TupleType(fields) assert tup_ty.loc.same_as(tvm_tirx.ir.UnknownLoc()) diff --git a/tirx/tests/python/runtime/test_device_api.py b/tirx/tests/python/runtime/test_device_api.py index 715b0bb9fba5..01a977ea783c 100644 --- a/tirx/tests/python/runtime/test_device_api.py +++ b/tirx/tests/python/runtime/test_device_api.py @@ -37,7 +37,7 @@ def test_check_if_device_exists(): cmd = [ sys.executable, "-c", - "import tvm_tirx; tvm_tirx.device('cuda').exist", + "import tvm_ffi; import tvm_tirx; tvm_ffi.device('cuda').exist", ] subprocess.check_call( cmd, diff --git a/tirx/tests/python/runtime/test_memory_scope.py b/tirx/tests/python/runtime/test_memory_scope.py index c2b1f51129e2..042001e41a20 100644 --- a/tirx/tests/python/runtime/test_memory_scope.py +++ b/tirx/tests/python/runtime/test_memory_scope.py @@ -95,6 +95,7 @@ def test_runtime_only_registration(): sys.executable, "-c", """ +import tvm_ffi import tvm_tirx from tvm_tirx.runtime import MemoryScope assert tvm_tirx.base._RUNTIME_ONLY @@ -102,7 +103,7 @@ def test_runtime_only_registration(): assert MemoryScope(name).name == name for name, allocator in (("nvshmem", "runtime.disco.nvshmem.empty"), ("ipc_memory", "runtime.disco.cuda_ipc.empty")): - if tvm_tirx.get_global_func(allocator, allow_missing=True) is not None: + if tvm_ffi.get_global_func(allocator, allow_missing=True) is not None: assert MemoryScope(name).name == name else: try: diff --git a/tirx/tests/python/runtime/test_module.py b/tirx/tests/python/runtime/test_module.py index dbcc77d73aea..9d7b80a50907 100644 --- a/tirx/tests/python/runtime/test_module.py +++ b/tirx/tests/python/runtime/test_module.py @@ -47,9 +47,8 @@ def system_library_source(native_source): def test_module_properties(source_module, native_library): - assert tvm_tirx.runtime.Module is tvm_ffi.Module + assert isinstance(source_module, tvm_ffi.Module) assert tvm_tirx.runtime.load_module is tvm_ffi.load_module - assert tvm_tirx.runtime.system_lib is tvm_ffi.system_lib assert not hasattr(source_module, "export_library") assert not hasattr(source_module, "time_evaluator") assert source_module.is_binary_serializable() @@ -106,9 +105,10 @@ def test_system_library(system_library_source, native_compile_options, tmp_path) import ctypes import sys import numpy as np +import tvm_ffi import tvm_tirx handle = ctypes.CDLL(sys.argv[1]) -module = tvm_tirx.runtime.system_lib() +module = tvm_ffi.system_lib() assert module["answer"]() == 42 a = tvm_tirx.runtime.tensor(np.arange(10, dtype="float32")) b = tvm_tirx.runtime.empty((10,), "float32") diff --git a/tirx/tests/python/runtime/test_tensor.py b/tirx/tests/python/runtime/test_tensor.py index d564d2743306..e1897d8c1305 100644 --- a/tirx/tests/python/runtime/test_tensor.py +++ b/tirx/tests/python/runtime/test_tensor.py @@ -18,6 +18,7 @@ import numpy as np import pytest +import tvm_ffi import tvm_tirx import tvm_tirx.testing @@ -74,7 +75,7 @@ def test_subview_must_be_in_range_of_immediate_parent(): def test_from_dlpack_shape_one(): torch = pytest.importorskip("torch") values = torch.randn(1, 16) - tensor = tvm_tirx.runtime.from_dlpack(torch.utils.dlpack.to_dlpack(values)) + tensor = tvm_ffi.from_dlpack(torch.utils.dlpack.to_dlpack(values)) np.testing.assert_equal(tensor.numpy(), values.numpy()) view = tensor._create_view((2, 8)) np.testing.assert_equal(view.numpy(), values.numpy().reshape(2, 8)) @@ -123,7 +124,7 @@ def test_tensor_json_roundtrip(container): def test_tensor_container_identity(): tensor = tvm_tirx.runtime.tensor([1, 2, 3]) - values = tvm_tirx.runtime.convert([tensor, tensor]) + values = tvm_ffi.convert([tensor, tensor]) assert values[0].same_as(tensor) assert values[1].same_as(tensor) assert isinstance(values[0], tvm_tirx.runtime.Tensor) diff --git a/tirx/tests/python/script/test_basic_usage.py b/tirx/tests/python/script/test_basic_usage.py index a7718074b426..54183f0d33e3 100644 --- a/tirx/tests/python/script/test_basic_usage.py +++ b/tirx/tests/python/script/test_basic_usage.py @@ -22,13 +22,14 @@ import linecache import pytest +import tvm_ffi # Script-local bindings are observed through the constructed IR. # ruff: noqa: F841 from minilang import Value import tvm_tirx -from tvm_tirx import DataType, ir +from tvm_tirx import ir from tvm_tirx.ir import prim from tvm_tirx.ir._prim_expr_overload import EqualOp from tvm_tirx.script import ir as I @@ -579,7 +580,7 @@ def test_datatype_value_roundtrip(language): spelling = 'T.dtype("float32")' source = f"@M.function\ndef main():\n M.record({spelling})\n" value = entry.parse(source, extra_vars={"M": language.M, "T": tvm_tirx.script.tirx}).body[0][1] - ir.assert_structural_equal(value, ir.DataTypeImm(DataType("float32"))) + ir.assert_structural_equal(value, ir.DataTypeImm(tvm_ffi.dtype("float32"))) printed = ir.IRModule(attrs={"dtype": value}).script() assert f'T.module_attrs({{"dtype": {spelling}}})' in printed diff --git a/tirx/tests/python/sym/test_analyzer_object.py b/tirx/tests/python/sym/test_analyzer_object.py index b3c239f81a78..9d10a76bd0bd 100644 --- a/tirx/tests/python/sym/test_analyzer_object.py +++ b/tirx/tests/python/sym/test_analyzer_object.py @@ -16,11 +16,11 @@ # under the License. import pytest +import tvm_ffi import tvm_tirx import tvm_tirx.testing from tvm_tirx import tirx -from tvm_tirx.runtime import Object from tvm_tirx.sym.analyzer import CompareResult @@ -28,7 +28,7 @@ def test_analyzer_is_ffi_object_with_persistent_state(): analyzer = tvm_tirx.sym.Analyzer() x = tirx.Var("x", "int64") - assert isinstance(analyzer, Object) + assert isinstance(analyzer, tvm_ffi.Object) analyzer.bind(x, tvm_tirx.ir.Range(0, 8)) assert analyzer.const_int_bound_is_bound(x) diff --git a/tirx/tests/python/target/test_llvm_features_info.py b/tirx/tests/python/target/test_llvm_features_info.py index df2b9643d078..13326dfa914d 100644 --- a/tirx/tests/python/target/test_llvm_features_info.py +++ b/tirx/tests/python/target/test_llvm_features_info.py @@ -15,6 +15,7 @@ # specific language governing permissions and limitations # under the License. import pytest +import tvm_ffi import tvm_tirx from tvm_tirx import codegen @@ -34,19 +35,17 @@ def test_llvm_targets(capfd): assert len(codegen.llvm.get_system_triple()) assert len(codegen.llvm.get_system_x86_vendor()) # check ffi vs python - assert ( - codegen.llvm.get_system_cpu() == tvm_tirx.get_global_func("codegen.llvm_get_system_cpu")() - ) + assert codegen.llvm.get_system_cpu() == tvm_ffi.get_global_func("codegen.llvm_get_system_cpu")() assert ( codegen.llvm.get_system_triple() - == tvm_tirx.get_global_func("codegen.llvm_get_system_triple")() + == tvm_ffi.get_global_func("codegen.llvm_get_system_triple")() ) assert ( codegen.llvm.get_system_x86_vendor() - == tvm_tirx.get_global_func("codegen.llvm_get_system_x86_vendor")() + == tvm_ffi.get_global_func("codegen.llvm_get_system_x86_vendor")() ) assert str(codegen.llvm.get_targets()) == str( - tvm_tirx.get_global_func("codegen.llvm_get_targets")() + tvm_ffi.get_global_func("codegen.llvm_get_targets")() ) tvm_tirx.codegen.llvm.get_cpu_features( @@ -99,7 +98,7 @@ def test_target_features(min_llvm_version, llvm_target, cpu_arch, cpu_features, assert cpu_arch in codegen.llvm.get_cpu_archlist(target) # check ffi vs python assert str(codegen.llvm.get_cpu_archlist(target)) == str( - tvm_tirx.get_global_func("codegen.llvm_get_cpu_archlist")(target) + tvm_ffi.get_global_func("codegen.llvm_get_cpu_archlist")(target) ) ## @@ -114,5 +113,5 @@ def test_target_features(min_llvm_version, llvm_target, cpu_arch, cpu_features, # check ffi vs python for feat in cpu_features: assert str(codegen.llvm.cpu_has_features(feat, target)) == str( - tvm_tirx.get_global_func("codegen.llvm_cpu_has_feature")(feat, target) + tvm_ffi.get_global_func("codegen.llvm_cpu_has_feature")(feat, target) ) diff --git a/tirx/tests/python/target/test_x86_features.py b/tirx/tests/python/target/test_x86_features.py index 59da3b2015a2..c56d96a48707 100644 --- a/tirx/tests/python/target/test_x86_features.py +++ b/tirx/tests/python/target/test_x86_features.py @@ -16,8 +16,8 @@ # under the License. # ruff: noqa: F841 import pytest +import tvm_ffi -import tvm_tirx from tvm_tirx import codegen from tvm_tirx.codegen import target_has_features from tvm_tirx.target import Target @@ -205,7 +205,7 @@ def test_x86_target_features(min_llvm_version, tvm_target, x86_feature, is_suppo # check for feature via the ffi llvm api (no explicit target, no context target) try: assert ( - tvm_tirx.get_global_func("codegen.target_has_feature")(x86_feature, None) + tvm_ffi.get_global_func("codegen.target_has_feature")(x86_feature, None) == is_supported ) assert False @@ -229,7 +229,7 @@ def test_x86_target_features(min_llvm_version, tvm_target, x86_feature, is_suppo if isinstance(x86_feature, str): # check for feature via the ffi llvm api (with explicit target, no context target) assert ( - tvm_tirx.get_global_func("codegen.target_has_feature")(x86_feature, Target(tvm_target)) + tvm_ffi.get_global_func("codegen.target_has_feature")(x86_feature, Target(tvm_target)) == is_supported ) @@ -246,7 +246,7 @@ def test_x86_target_features(min_llvm_version, tvm_target, x86_feature, is_suppo # check for feature via the ffi llvm api (current context target) ( sum( - tvm_tirx.get_global_func("codegen.target_has_feature")(feat, None) + tvm_ffi.get_global_func("codegen.target_has_feature")(feat, None) for feat in x86_feature ) > 0 diff --git a/tirx/tests/python/tirx/base/test_index_map.py b/tirx/tests/python/tirx/base/test_index_map.py index 134f9d245312..7136e12f82d6 100644 --- a/tirx/tests/python/tirx/base/test_index_map.py +++ b/tirx/tests/python/tirx/base/test_index_map.py @@ -134,14 +134,14 @@ def test_nonbijective_inverse_gives_error(): inverse=lambda i, j: [4 * i + j], pre_shape=[16], post_shape=[T.int32(4), T.int32(4)], - padding=lambda i, j: tvm_tirx.runtime.convert(False), + padding=lambda i, j: tvm_ffi.convert(False), ), "right_padding": dict( forward=lambda i: [i // 4, i % 4], inverse=lambda i, j: [4 * i + j], pre_shape=[15], post_shape=[T.int32(4), T.int32(4)], - padding=lambda i, j: tvm_tirx.tirx.And(i == 3, tvm_tirx.runtime.convert(3) == j), + padding=lambda i, j: tvm_tirx.tirx.And(i == 3, tvm_ffi.convert(3) == j), ), "left_padding": dict( forward=lambda i: [(i + 1) // 4, (i + 1) % 4], @@ -157,7 +157,7 @@ def test_nonbijective_inverse_gives_error(): post_shape=[T.int32(4), T.int32(4)], padding=lambda i, j: tvm_tirx.tirx.Or( tvm_tirx.tirx.And(i == 0, j < 1), - tvm_tirx.tirx.And(i == 3, tvm_tirx.runtime.convert(3) == j), + tvm_tirx.tirx.And(i == 3, tvm_ffi.convert(3) == j), ), ), "dynamic_size": dict( @@ -186,7 +186,7 @@ def test_nonbijective_inverse_gives_error(): padding=lambda i_outer, j_outer, i_inner, j_inner: tvm_tirx.tirx.Or( tvm_tirx.tirx.Or( tvm_tirx.tirx.And(i_outer == 0, i_inner < 1), - tvm_tirx.tirx.And(i_outer == 3, tvm_tirx.runtime.convert(3) == i_inner), + tvm_tirx.tirx.And(i_outer == 3, tvm_ffi.convert(3) == i_inner), ), tvm_tirx.tirx.Or( tvm_tirx.tirx.And(j_outer == 0, j_inner < 5), @@ -227,7 +227,7 @@ def test_nonbijective_inverse_gives_error(): inverse=lambda i, j: [i * 4 + j], pre_shape=[3], post_shape=[T.int32(1), T.int32(4)], - padding=lambda i, j: tvm_tirx.runtime.convert(3) == j, + padding=lambda i, j: tvm_ffi.convert(3) == j, ), } ) diff --git a/tirx/tests/python/tirx/base/test_nodes.py b/tirx/tests/python/tirx/base/test_nodes.py index 1ce2d99cd641..a0840dde38b7 100644 --- a/tirx/tests/python/tirx/base/test_nodes.py +++ b/tirx/tests/python/tirx/base/test_nodes.py @@ -16,6 +16,7 @@ # under the License. import numpy as np import pytest +import tvm_ffi import tvm_tirx from tvm_tirx import ir @@ -113,7 +114,7 @@ def test_cast(): def test_attr(): - a = tvm_tirx.runtime.convert(1) + a = tvm_ffi.convert(1) assert a == 1 try: a.no_field @@ -400,15 +401,15 @@ def test_tensor_load_rejects_two_scalable_types(): def test_intimm_cond(): - x = tvm_tirx.runtime.convert(1) - y = tvm_tirx.runtime.convert(1) + x = tvm_ffi.convert(1) + y = tvm_ffi.convert(1) s = {x} assert y in s assert x == y assert x < 20 assert not (x >= 20) assert x < 10 and y < 10 - assert not tvm_tirx.runtime.convert(x != 1) + assert not tvm_ffi.convert(x != 1) assert x == 1 diff --git a/tirx/tests/python/tirx/base/test_scalable_datatype.py b/tirx/tests/python/tirx/base/test_scalable_datatype.py index a5a11a97593d..071800afd7e3 100644 --- a/tirx/tests/python/tirx/base/test_scalable_datatype.py +++ b/tirx/tests/python/tirx/base/test_scalable_datatype.py @@ -16,6 +16,7 @@ # under the License. import pytest +import tvm_ffi import tvm_tirx from tvm_tirx import tirx @@ -28,7 +29,7 @@ def test_create_scalable_data_type_python_api(): - dtype = tvm_tirx.DataType("float32xvscalex4") + dtype = tvm_ffi.dtype("float32xvscalex4") assert str(dtype) == "float32xvscalex4" @@ -63,7 +64,7 @@ def my_func(): def test_invalid_data_type(): with pytest.raises(ValueError): - tvm_tirx.DataType("float32x4xvscale") + tvm_ffi.dtype("float32x4xvscale") if __name__ == "__main__": diff --git a/tirx/tests/python/tirx/base/test_tensor_type.py b/tirx/tests/python/tirx/base/test_tensor_type.py index 8e47538640fa..4973f775931c 100644 --- a/tirx/tests/python/tirx/base/test_tensor_type.py +++ b/tirx/tests/python/tirx/base/test_tensor_type.py @@ -18,6 +18,7 @@ import uuid import pytest +import tvm_ffi import tvm_tirx import tvm_tirx.testing @@ -107,7 +108,7 @@ def test_tensor_type_and_projection(dtype, scope): assert tuple(tensor.ty.shape) == (m, n) assert tuple(tensor.shape) == (m, n) assert tensor.ty.dtype == tvm_tirx.ir.PrimType(dtype) - assert tensor.dtype == tvm_tirx.DataType(dtype) + assert tensor.dtype == tvm_ffi.dtype(dtype) assert tensor.data.ty == tvm_tirx.ir.PtrType(tvm_tirx.ir.PrimType(dtype), scope) assert tensor.data.op.name == "tirx.tensor_data_ptr" assert tensor.data.args[0].same_as(tensor) diff --git a/tirx/tests/python/tirx/base/test_texture_scope.py b/tirx/tests/python/tirx/base/test_texture_scope.py index acc658911b8a..a15a4ebcb05b 100644 --- a/tirx/tests/python/tirx/base/test_texture_scope.py +++ b/tirx/tests/python/tirx/base/test_texture_scope.py @@ -14,6 +14,8 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. +import tvm_ffi + import tvm_tirx import tvm_tirx.testing from tvm_tirx.script import tirx as T @@ -45,7 +47,7 @@ def texture_kernel( target = tvm_tirx.target.Target({"kind": "opencl", "keys": ["adreno"]}) mod = tvm_tirx.IRModule.from_expr(texture_kernel.with_attr("target", target)) - source = tvm_tirx.get_global_func("target.build.opencl")(mod, target).inspect_source() + source = tvm_ffi.get_global_func("target.build.opencl")(mod, target).inspect_source() assert "__read_only image2d_array_t" in source assert "__write_only image2d_array_t" in source assert "READ_IMAGEF" in source diff --git a/tirx/tests/python/tirx/codegen/test_blackwell.py b/tirx/tests/python/tirx/codegen/test_blackwell.py index 73dff1d15774..dbf8cb69af30 100644 --- a/tirx/tests/python/tirx/codegen/test_blackwell.py +++ b/tirx/tests/python/tirx/codegen/test_blackwell.py @@ -403,7 +403,7 @@ def run_and_check(): @pytest.mark.skipif(not env.has_cuda_compute(10), reason="need cuda compute >= 10.0") def test_tcgen05_cp_ld_roundtrip(): dtype = "float32" - dtype_bits = tvm_tirx.DataType(dtype).bits + dtype_bits = tvm_ffi.dtype(dtype).bits HEIGHT = 128 WIDTH = 64 N_COLS = 512 diff --git a/tirx/tests/python/tirx/codegen/test_cuda.py b/tirx/tests/python/tirx/codegen/test_cuda.py index 1bb95ed34ea8..629880d3d708 100644 --- a/tirx/tests/python/tirx/codegen/test_cuda.py +++ b/tirx/tests/python/tirx/codegen/test_cuda.py @@ -20,6 +20,7 @@ import numpy as np import pytest +import tvm_ffi import tvm_tirx import tvm_tirx.testing @@ -188,7 +189,7 @@ def test_legacy_tensorcore_ops_require_downstream_codegen(op_name): ) target = tvm_tirx.target.Target({"kind": "cuda", "arch": "sm_80"}) with pytest.raises(RuntimeError, match="Unresolved call.*" + re.escape(call.op.name)): - tvm_tirx.get_global_func("target.build.cuda")(tvm_tirx.IRModule.from_expr(func), target) + tvm_ffi.get_global_func("target.build.cuda")(tvm_tirx.IRModule.from_expr(func), target) @pytest.mark.parametrize("target_kind", ["cuda", "opencl"]) @@ -203,7 +204,7 @@ def main(data: T.handle("float32")): target = tvm_tirx.target.Target(target_kind) with pytest.raises(RuntimeError, match="warp storage is unsupported"): - tvm_tirx.get_global_func("target.build." + target_kind)( + tvm_ffi.get_global_func("target.build." + target_kind)( tvm_tirx.IRModule.from_expr(main), target ) @@ -226,7 +227,7 @@ def main(data: T.handle("float32"), output: T.handle): monkeypatch.setenv("TVM_COMPILE_FORCE_FALLBACK", "1") target = tvm_tirx.target.Target({"kind": "cuda", "arch": "sm_80"}) - source = tvm_tirx.get_global_func("target.build.cuda")( + source = tvm_ffi.get_global_func("target.build.cuda")( tvm_tirx.IRModule.from_expr(main), target ).inspect_source() assert "float* data, void* output" in source @@ -267,7 +268,7 @@ def test_vector_pointer_preserves_packed_offset(monkeypatch): lowered = tvm_tirx.tirx.transform.LowerIntrin()(tvm_tirx.IRModule.from_expr(func)) monkeypatch.setenv("TVM_COMPILE_FORCE_FALLBACK", "1") - source = tvm_tirx.get_global_func("target.build.cuda")(lowered, target).inspect_source() + source = tvm_ffi.get_global_func("target.build.cuda")(lowered, target).inspect_source() call = next(line.strip() for line in source.splitlines() if line.strip().startswith("consume(")) assert "make_int4" not in call diff --git a/tirx/tests/python/tirx/codegen/test_hopper.py b/tirx/tests/python/tirx/codegen/test_hopper.py index 63f9ff48cebc..34426a662588 100644 --- a/tirx/tests/python/tirx/codegen/test_hopper.py +++ b/tirx/tests/python/tirx/codegen/test_hopper.py @@ -19,6 +19,7 @@ import numpy as np import pytest +import tvm_ffi import tvm_tirx import tvm_tirx.testing @@ -427,7 +428,7 @@ def test_cp_async_bulk_tensor_global_to_shared_unicast(dtype, inputs): import ml_dtypes def get_ir(shape, tma_args): - t_dtype = tvm_tirx.DataType(dtype) + t_dtype = tvm_ffi.dtype(dtype) total_bytes = math.prod(shape) * t_dtype.bits // 8 coord = [0 for _ in shape] tma_args_copy = tma_args.copy() @@ -618,7 +619,7 @@ def test_tensormap_encode_tiled_runtime_validation(shape, dtype, encode_args, er @pytest.mark.skipif(not env.has_cuda_compute(9), reason="need cuda compute >= 9.0") def test_cp_async_bulk_tensor_global_to_shared_swizzle(swizzle, dtype): def get_ir(swizzle, dtype): - dtype = tvm_tirx.DataType(dtype) + dtype = tvm_ffi.dtype(dtype) elem_bytes = dtype.bits // 8 shape = [16, 64] @@ -699,7 +700,7 @@ def main( A_np = [i for i in range(total_elems)] A_np = np.array(A_np).astype(dtype) B_np = np.zeros((total_elems,)).astype(dtype) - dtype = tvm_tirx.DataType(dtype) + dtype = tvm_ffi.dtype(dtype) layout = T.ComposeLayout( int(math.log2(128 // dtype.bits)), swizzle, @@ -1046,8 +1047,8 @@ def get_ir( ): coordA = [0 for _ in shapeA] coordB = [0 for _ in shapeB] - A_bytes = tvm_tirx.DataType(in_dtype).bits // 8 * math.prod(shapeA) - B_bytes = tvm_tirx.DataType(in_dtype).bits // 8 * math.prod(shapeB) + A_bytes = tvm_ffi.dtype(in_dtype).bits // 8 * math.prod(shapeA) + B_bytes = tvm_ffi.dtype(in_dtype).bits // 8 * math.prod(shapeB) C_elems = math.prod(shapeC) // 128 @@ -1163,7 +1164,7 @@ def main( transA = transB = True swizzleA = swizzleB = 3 - t_in_dtype = tvm_tirx.DataType(in_dtype) + t_in_dtype = tvm_ffi.dtype(in_dtype) elem_bytes = t_in_dtype.bits // 8 target = tvm_tirx.target.Target("cuda") @@ -1222,7 +1223,7 @@ def get_ir( shapeA, shapeB, shapeC, B_tma_args, in_dtype, in_dtype_bits, out_dtype, B_encode_args ): coordB = [0 for _ in shapeB] - B_bytes = tvm_tirx.DataType(in_dtype).bits // 8 * math.prod(shapeB) + B_bytes = tvm_ffi.dtype(in_dtype).bits // 8 * math.prod(shapeB) A_elems = math.prod(shapeA) // 128 C_elems = math.prod(shapeC) // 128 @@ -1350,7 +1351,7 @@ def main( transB = True swizzleB = 3 - t_in_dtype = tvm_tirx.DataType(in_dtype) + t_in_dtype = tvm_ffi.dtype(in_dtype) elem_bytes = t_in_dtype.bits // 8 target = tvm_tirx.target.Target("cuda") diff --git a/tirx/tests/python/tirx/layout_fixtures.py b/tirx/tests/python/tirx/layout_fixtures.py index e58b68b02dea..fdbb8446b814 100644 --- a/tirx/tests/python/tirx/layout_fixtures.py +++ b/tirx/tests/python/tirx/layout_fixtures.py @@ -16,7 +16,8 @@ # under the License. """Operand layouts shared by CUDA kernel tests; not part of the public API.""" -import tvm_tirx +import tvm_ffi + from tvm_tirx.backend.cuda.axis import Axis as CUDAAxis from tvm_tirx.tirx.layout import Iter, R, S, TileLayout @@ -28,7 +29,7 @@ def tcgen05_register_layout(instr_shape, tensor_shape, dtype): """Build the logical register image exercised by TCGEN05 readback tests.""" - bits = tvm_tirx.runtime.DataType(dtype).bits + bits = tvm_ffi.dtype(dtype).bits rows, cols = tensor_shape if instr_shape == "32x32b" and rows == 64: diff --git a/tirx/tests/python/tirx/script/test_printer_invisible_path_info.py b/tirx/tests/python/tirx/script/test_printer_invisible_path_info.py index 51ca34485cd1..bf4eb803348a 100644 --- a/tirx/tests/python/tirx/script/test_printer_invisible_path_info.py +++ b/tirx/tests/python/tirx/script/test_printer_invisible_path_info.py @@ -17,6 +17,7 @@ """TIRx script printer invisible path info.""" +import tvm_ffi from tvm_ffi.access_path import AccessPath import tvm_tirx @@ -58,7 +59,7 @@ def test_render_invisible_path_info_exact_visible_path_omits_note(): def test_render_invisible_path_info_without_visible_path(): result = _script( - tvm_tirx.runtime.ShapeTuple([1, 2]), + tvm_ffi.Shape([1, 2]), PrinterConfig(path_to_underline=[AccessPath.root().attr("missing")]), ) assert result == ( diff --git a/tirx/tests/python/tirx/test_cuda_backend_config.py b/tirx/tests/python/tirx/test_cuda_backend_config.py index 0407ff581a33..0963e9311f78 100644 --- a/tirx/tests/python/tirx/test_cuda_backend_config.py +++ b/tirx/tests/python/tirx/test_cuda_backend_config.py @@ -567,7 +567,7 @@ def main( def test_legacy_payloads(): import struct - load = tvm_tirx.get_global_func("ffi.Module.load_from_bytes.cuda", allow_missing=True) + load = tvm_ffi.get_global_func("ffi.Module.load_from_bytes.cuda", allow_missing=True) if load is None: pytest.skip("CUDA runtime module loader is not linked") diff --git a/tirx/tests/python/tirx/tile/cuda/test_cp_async_bulk.py b/tirx/tests/python/tirx/tile/cuda/test_cp_async_bulk.py index 654eccaf9cda..571faf8e534b 100644 --- a/tirx/tests/python/tirx/tile/cuda/test_cp_async_bulk.py +++ b/tirx/tests/python/tirx/tile/cuda/test_cp_async_bulk.py @@ -46,7 +46,7 @@ def _make_dsmem_dispatch_call(shape, dtype, src_layout, dst_layout, chunks=1): # Each test explicitly chooses its number of bulk chunks. Keep both # original physical maps; flattening only changes the logical shape. canonical = (math.prod(shape),) if chunks == 1 else shape - chunk_bytes = math.prod(shape) * tvm_tirx.DataType(dtype).bits // 8 // chunks + chunk_bytes = math.prod(shape) * tvm_ffi.dtype(dtype).bits // 8 // chunks src = tvm_tirx.tirx.decl_tensor(canonical, dtype, "A", scope="shared.dyn", layout=src_layout) dst = tvm_tirx.tirx.decl_tensor(canonical, dtype, "B", scope="shared.dyn", layout=dst_layout) op_call = T.cuda.tile.apply( @@ -154,7 +154,7 @@ def test_dsmem(shape, dtype, src_spec, dst_spec, expected): CLUSTER_N = 2 n_elements = functools.reduce(lambda a, b: a * b, shape, 1) - copy_bytes = n_elements * tvm_tirx.DataType(dtype).bits // 8 + copy_bytes = n_elements * tvm_ffi.dtype(dtype).bits // 8 canonical = (n_elements,) if expected == 1 else shape src_phys = _layout_physical_elements(src_layout) dst_phys = _layout_physical_elements(dst_layout) diff --git a/tirx/tests/python/tirx/tile/cuda/test_cp_async_bulk_tensor.py b/tirx/tests/python/tirx/tile/cuda/test_cp_async_bulk_tensor.py index 937f0bd67c21..7c4429985839 100644 --- a/tirx/tests/python/tirx/tile/cuda/test_cp_async_bulk_tensor.py +++ b/tirx/tests/python/tirx/tile/cuda/test_cp_async_bulk_tensor.py @@ -842,8 +842,8 @@ def _original_addresses(kwargs): def _recipe_views(kwargs, recipe): """Declare the caller's explicit zero-copy aliases on the original storage.""" dtype = kwargs.get("dtype", "float16") - bits = tvm_tirx.DataType(recipe.dtype).bits - original_bits = tvm_tirx.DataType(dtype).bits + bits = tvm_ffi.dtype(recipe.dtype).bits + original_bits = tvm_ffi.dtype(dtype).bits ga, sa = _original_addresses(kwargs) shape = tuple(reversed(recipe.dims)) strides = tuple(reversed((1, *(s * 8 // bits for s in recipe.strides)))) diff --git a/tirx/tests/python/tirx/tile/cuda/test_ld_st_gmem_smem.py b/tirx/tests/python/tirx/tile/cuda/test_ld_st_gmem_smem.py index 12218563ce10..695fe6c1ff23 100644 --- a/tirx/tests/python/tirx/tile/cuda/test_ld_st_gmem_smem.py +++ b/tirx/tests/python/tirx/tile/cuda/test_ld_st_gmem_smem.py @@ -25,6 +25,7 @@ import numpy as np import pytest +import tvm_ffi import tvm_tirx import tvm_tirx.testing @@ -243,7 +244,7 @@ def _explicit_staging_kernel(g_shape, s_shape, g_layout, s_layout, start, word_d rows, cols = s_shape lanes = min(rows, 32) reg_layout = TileLayout(S[s_shape : (1 @ tx, 1)]) - bits = tvm_tirx.DataType(word_dtype).bits + bits = tvm_ffi.dtype(word_dtype).bits suffix = f"v{words}.b{bits}" if words > 1 else f"b{bits}" load, store = "ld.global." + suffix, "st.shared." + suffix # The raw argument includes gaps in the original global layout. diff --git a/tirx/tests/python/tirx/tile/cuda/test_ld_st_registers.py b/tirx/tests/python/tirx/tile/cuda/test_ld_st_registers.py index f9e1df7e97f2..dc06902ae0fb 100644 --- a/tirx/tests/python/tirx/tile/cuda/test_ld_st_registers.py +++ b/tirx/tests/python/tirx/tile/cuda/test_ld_st_registers.py @@ -50,7 +50,7 @@ def _build_roundtrip_kernel(scope, n_threads, k, dtype, non_r_scope): shape = (n_threads, k) r_layout = _r_layout(scope, shape) # The caller chooses a native width from this test's fixed row size. - words = min(4, k * tvm_tirx.DataType(dtype).bits // 32) + words = min(4, k * tvm_ffi.dtype(dtype).bits // 32) load = f"ld.{non_r_scope}.v{words}.u32" store = f"st.{non_r_scope}.v{words}.u32" if non_r_scope == "shared": @@ -273,7 +273,7 @@ def run_and_check(): def test_copy_g2l_l2g_vec_load(task, dtype): g_shape, l_shape, g_region, thread_cnt, layoutA, layoutB, layoutLocal = task - words = 2 if tvm_tirx.DataType(dtype).bits == 8 else 4 + words = 2 if tvm_ffi.dtype(dtype).bits == 8 else 4 load, store = f"ld.global.v{words}.b32", f"st.global.v{words}.b32" r_gmem = tuple(slice(g_region[i][0], g_region[i][1]) for i in range(len(g_shape))) diff --git a/tirx/tests/python/tirx/tile/cuda/test_pattern_math.py b/tirx/tests/python/tirx/tile/cuda/test_pattern_math.py index fa309958b8bc..c43219598ebd 100644 --- a/tirx/tests/python/tirx/tile/cuda/test_pattern_math.py +++ b/tirx/tests/python/tirx/tile/cuda/test_pattern_math.py @@ -19,6 +19,7 @@ import numpy as np import pytest +import tvm_ffi import tvm_tirx from tvm_tirx.backend.cuda.axis import laneid, warpid @@ -75,7 +76,7 @@ def kernel(A: T.Tensor((2, 32, 7, 1), "float32"), Output: T.Tensor((2, 32, 7, 1) ) @pytest.mark.parametrize("opcode,kind", [("add", "binary"), ("mul", "binary"), ("fma", "ternary")]) def test_native_packed(dtype, suffix, opcode, kind): - bits = tvm_tirx.DataType(dtype).bits * 2 + bits = tvm_ffi.dtype(dtype).bits * 2 storage = "uint" + str(bits) instruction = f"{opcode}.rn.{suffix}" diff --git a/tirx/tests/python/tirx/tile/cuda/test_pattern_mma.py b/tirx/tests/python/tirx/tile/cuda/test_pattern_mma.py index 74fe7724f0e4..3474176a3285 100644 --- a/tirx/tests/python/tirx/tile/cuda/test_pattern_mma.py +++ b/tirx/tests/python/tirx/tile/cuda/test_pattern_mma.py @@ -19,6 +19,7 @@ import numpy as np import pytest +import tvm_ffi import tvm_tirx from tvm_tirx.backend.cuda.axis import laneid @@ -101,8 +102,8 @@ def test_tcgen05_mma(m, ws, dtype, k, kind, acc_dtype, major, swizzle=3): if major == "mn" and kind == "tf32" and swizzle != 4: pytest.skip("TF32 MN-major uses the separately tested 32B-atomicity descriptor") n, nk = 64, 3 - packed = 32 // tvm_tirx.DataType(acc_dtype).bits - bits = tvm_tirx.DataType(dtype).bits + packed = 32 // tvm_ffi.dtype(acc_dtype).bits + bits = tvm_ffi.dtype(dtype).bits row_bytes = 128 if swizzle == 4 else 16 << swizzle swizzle_len = 2 if swizzle == 4 else swizzle atom_len = 2 if swizzle == 4 else 3 diff --git a/tirx/tests/python/tirx/tile/cuda/test_tcgen05_cp.py b/tirx/tests/python/tirx/tile/cuda/test_tcgen05_cp.py index 2c1d05818dcd..2513fb5cc000 100644 --- a/tirx/tests/python/tirx/tile/cuda/test_tcgen05_cp.py +++ b/tirx/tests/python/tirx/tile/cuda/test_tcgen05_cp.py @@ -26,6 +26,7 @@ import numpy as np import pytest +import tvm_ffi import tvm_tirx import tvm_tirx.testing @@ -97,7 +98,7 @@ def _emit_cp(dst, src, t_region, s_region, shape, opcode): operand's physical layout. Each region retains its original row/column start. """ atom_rows, atom_bits = _shape_dims(shape) - bits = tvm_tirx.DataType(src.dtype).bits + bits = tvm_ffi.dtype(src.dtype).bits words = atom_bits // 32 s_row, s_col = s_region t_row, t_col = t_region @@ -266,7 +267,7 @@ def _build_case( instruction_shape=None, ): """Assemble (kernel, host-check closure) for one cp configuration.""" - bits = tvm_tirx.runtime.DataType(dtype).bits + bits = tvm_ffi.dtype(dtype).bits rows, atom_bits = _shape_dims(shape) epa = atom_bits // bits # elems per lane per cp instruction C = epa * n_mid # copied cols (elems) diff --git a/tirx/tests/python/tirx/tile/cuda/test_tcgen05_ld_st.py b/tirx/tests/python/tirx/tile/cuda/test_tcgen05_ld_st.py index 0b221613c1cb..9430c1ce754d 100644 --- a/tirx/tests/python/tirx/tile/cuda/test_tcgen05_ld_st.py +++ b/tirx/tests/python/tirx/tile/cuda/test_tcgen05_ld_st.py @@ -40,6 +40,7 @@ import numpy as np import pytest +import tvm_ffi import tvm_tirx import tvm_tirx.testing @@ -61,7 +62,7 @@ def _window(region): def _words(region): tensor = _window(region) - return tensor if tvm_tirx.DataType(tensor.dtype).bits == 32 else tensor.view("uint32") + return tensor if tvm_ffi.dtype(tensor.dtype).bits == 32 else tensor.view("uint32") def _atom_words(region, shape, rows): @@ -270,7 +271,7 @@ def _run_roundtrip_16b( frag_rows_override=None, tmem_datapath: str = "D", ): - bits = tvm_tirx.runtime.DataType(dtype).bits + bits = tvm_ffi.dtype(dtype).bits assert bits == 16 elem_per_32b = 2 K_cols_fp32 = _COL_FACTOR_FP32[shape] * rep @@ -913,7 +914,7 @@ def _run_load_test(shape: str, rep: int, dtype: str): """Stage A into TMEM via .32x32b, then read it back as the fragment via ..x with an explicit register layout, and compare each thread's registers against the expected layout-derived value.""" - bits = tvm_tirx.runtime.DataType(dtype).bits + bits = tvm_ffi.dtype(dtype).bits elem_per_32b = 32 // bits # Per-warp fp32 col span x number of warps in one warpgroup covers the # fragment column footprint. The TMEM allocation is sized for the same @@ -1105,7 +1106,7 @@ def test_tcgen05_st_16xnb_store(shape, rep, dtype): """ if rep not in _SHAPE_REPS[shape]: pytest.skip(f"rep {rep} not valid for {shape}") - bits = tvm_tirx.runtime.DataType(dtype).bits + bits = tvm_ffi.dtype(dtype).bits elem_per_32b = 32 // bits K_cols_fp32 = _COL_FACTOR_FP32[shape] * rep K_cols_elem = K_cols_fp32 * elem_per_32b @@ -1646,7 +1647,7 @@ def next_power_of_2(x): return 1 return 1 << (x - 1).bit_length() - bits = tvm_tirx.runtime.DataType(dtype).bits + bits = tvm_ffi.dtype(dtype).bits if 128 % bits != 0 or 32 % bits != 0: pytest.skip(f"dtype {dtype} is not supported") @@ -1771,7 +1772,7 @@ def next_power_of_2(x): return 1 return 1 << (x - 1).bit_length() - bits = tvm_tirx.runtime.DataType(dtype).bits + bits = tvm_ffi.dtype(dtype).bits if 128 % bits != 0 or 32 % bits != 0: pytest.skip(f"dtype {dtype} is not supported") @@ -1894,7 +1895,7 @@ def next_power_of_2(x): return 1 return 1 << (x - 1).bit_length() - bits = tvm_tirx.runtime.DataType(dtype).bits + bits = tvm_ffi.dtype(dtype).bits if 128 % bits != 0 or 32 % bits != 0: pytest.skip(f"dtype {dtype} is not supported") diff --git a/tirx/tests/python/tirx/tile/cuda/test_tcgen05_mma.py b/tirx/tests/python/tirx/tile/cuda/test_tcgen05_mma.py index e5161ca06a2d..0355aac865ed 100644 --- a/tirx/tests/python/tirx/tile/cuda/test_tcgen05_mma.py +++ b/tirx/tests/python/tirx/tile/cuda/test_tcgen05_mma.py @@ -22,6 +22,7 @@ import numpy as np import pytest +import tvm_ffi import tvm_tirx import tvm_tirx.testing @@ -75,9 +76,9 @@ def test_gemm_tcgen05_cta_group_1(task): assert C_shape[0] == 128 assert C_region[0] == (0, 128) assert len(C_shape) == 2 - A_elem_bytes = tvm_tirx.runtime.DataType(A_dtype).bits // 8 - B_elem_bytes = tvm_tirx.runtime.DataType(B_dtype).bits // 8 - C_elem_bytes = tvm_tirx.runtime.DataType(C_dtype).bits // 8 + A_elem_bytes = tvm_ffi.dtype(A_dtype).bits // 8 + B_elem_bytes = tvm_ffi.dtype(B_dtype).bits // 8 + C_elem_bytes = tvm_ffi.dtype(C_dtype).bits // 8 C_elem_32b = 4 // C_elem_bytes cols_alloc = max(32, next_power_of_2(C_shape[1] // C_elem_32b)) A_layout = mma_shared_layout(A_dtype, A_swizzle_mode, A_shape) @@ -466,9 +467,9 @@ def test_gemm_tcgen05_cta_group_2(task): assert C_shape[0] == 256 assert C_region[0] == (0, 128) assert len(C_shape) == 2 - A_elem_bytes = tvm_tirx.runtime.DataType(A_dtype).bits // 8 - B_elem_bytes = tvm_tirx.runtime.DataType(B_dtype).bits // 8 - C_elem_bytes = tvm_tirx.runtime.DataType(C_dtype).bits // 8 + A_elem_bytes = tvm_ffi.dtype(A_dtype).bits // 8 + B_elem_bytes = tvm_ffi.dtype(B_dtype).bits // 8 + C_elem_bytes = tvm_ffi.dtype(C_dtype).bits // 8 C_elem_32b = 4 // C_elem_bytes cols_alloc = max(32, next_power_of_2(C_shape[1] // C_elem_32b)) @@ -693,9 +694,9 @@ def test_gemm_tcgen05_cta_group_2_layout_b(): B_shape = (N_half, K) # per CTA: N_logical // cta_group C_shape = (M_per_cta * 2, N_logical) # global output - A_elem_bytes = tvm_tirx.runtime.DataType(A_dtype).bits // 8 - B_elem_bytes = tvm_tirx.runtime.DataType(B_dtype).bits // 8 - C_elem_32b = 4 // (tvm_tirx.runtime.DataType(C_dtype).bits // 8) + A_elem_bytes = tvm_ffi.dtype(A_dtype).bits // 8 + B_elem_bytes = tvm_ffi.dtype(B_dtype).bits // 8 + C_elem_32b = 4 // (tvm_ffi.dtype(C_dtype).bits // 8) cols_alloc = max(32, next_power_of_2(N_half // C_elem_32b)) A_layout = mma_shared_layout(A_dtype, swizzle_mode, A_shape) @@ -1124,9 +1125,9 @@ def test_gemm_block_scaled_fp8_cta_group_1(task): assert C_region[0] == (0, 128) assert len(C_shape) == 2 - A_elem_bytes = max(1, tvm_tirx.runtime.DataType(A_dtype).bits // 8) - B_elem_bytes = max(1, tvm_tirx.runtime.DataType(B_dtype).bits // 8) - C_elem_bytes = tvm_tirx.runtime.DataType(C_dtype).bits // 8 + A_elem_bytes = max(1, tvm_ffi.dtype(A_dtype).bits // 8) + B_elem_bytes = max(1, tvm_ffi.dtype(B_dtype).bits // 8) + C_elem_bytes = tvm_ffi.dtype(C_dtype).bits // 8 C_elem_32b = 4 // C_elem_bytes cols_alloc = max(32, next_power_of_2(C_shape[1] // C_elem_32b)) @@ -1147,7 +1148,7 @@ def test_gemm_block_scaled_fp8_cta_group_1(task): sf_mma_k = 1 # fp8: 1 scale factor per MMA iteration sfa_layout = scale_factor_tmem_layout(M, SF_K=sf_mma_k * 1, sf_per_mma=sf_mma_k) sfb_layout = scale_factor_tmem_layout(N, SF_K=sf_mma_k * 1, sf_per_mma=sf_mma_k) - sf_epc = 32 // tvm_tirx.runtime.DataType(SF_dtype).bits + sf_epc = 32 // tvm_ffi.dtype(SF_dtype).bits SFA_TMEM_SPACING = (int(sfa_layout.span("TCol")) + sf_epc - 1) // sf_epc SFA_TMEM_START = width SFB_TMEM_START = SFA_TMEM_START + SFA_TMEM_SPACING @@ -1445,9 +1446,9 @@ def test_gemm_block_scaled_fp8_cta_group_2(task): assert C_region[0] == (0, 128) assert len(C_shape) == 2 - A_elem_bytes = max(1, tvm_tirx.runtime.DataType(A_dtype).bits // 8) - B_elem_bytes = max(1, tvm_tirx.runtime.DataType(B_dtype).bits // 8) - C_elem_bytes = tvm_tirx.runtime.DataType(C_dtype).bits // 8 + A_elem_bytes = max(1, tvm_ffi.dtype(A_dtype).bits // 8) + B_elem_bytes = max(1, tvm_ffi.dtype(B_dtype).bits // 8) + C_elem_bytes = tvm_ffi.dtype(C_dtype).bits // 8 C_elem_32b = 4 // C_elem_bytes cols_alloc = max(32, next_power_of_2(C_shape[1] // C_elem_32b)) @@ -1469,7 +1470,7 @@ def test_gemm_block_scaled_fp8_cta_group_2(task): sf_mma_k = 1 # fp8: 1 scale factor per MMA iteration sf_layout = scale_factor_tmem_layout(128, SF_K=sf_mma_k * 1, sf_per_mma=sf_mma_k) - sf_epc = 32 // tvm_tirx.runtime.DataType(SF_dtype).bits + sf_epc = 32 // tvm_ffi.dtype(SF_dtype).bits SF_TMEM_SPACING = (int(sf_layout.span("TCol")) + sf_epc - 1) // sf_epc N_cols = C_region[1][1] - C_region[1][0] SFA_TMEM_START = N_cols @@ -1829,7 +1830,7 @@ def test_gemm_block_scaled_nvfp4_cta_group_1(): A_fp4_shape = (M, K) B_fp4_shape = (N, K) - C_elem_bytes = tvm_tirx.runtime.DataType(C_dtype).bits // 8 + C_elem_bytes = tvm_ffi.dtype(C_dtype).bits // 8 C_elem_32b = 4 // C_elem_bytes cols_alloc = max(32, next_power_of_2(C_shape[1] // C_elem_32b)) @@ -1843,7 +1844,7 @@ def test_gemm_block_scaled_nvfp4_cta_group_1(): sf_mma_k = 4 # nvfp4: 4 scale factors per MMA iteration (MMA_K=64, SF_VEC=16) sfa_layout = scale_factor_tmem_layout(M, SF_K=sf_mma_k * 1, sf_per_mma=sf_mma_k) sfb_layout = scale_factor_tmem_layout(N, SF_K=sf_mma_k * 1, sf_per_mma=sf_mma_k) - sf_epc = 32 // tvm_tirx.runtime.DataType(SF_dtype).bits + sf_epc = 32 // tvm_ffi.dtype(SF_dtype).bits SFA_TMEM_SPACING = (int(sfa_layout.span("TCol")) + sf_epc - 1) // sf_epc SFA_TMEM_START = width SFB_TMEM_START = SFA_TMEM_START + SFA_TMEM_SPACING @@ -2152,7 +2153,7 @@ def test_gemm_block_scaled_nvfp4_cta_group_2(): A_packed_shape = (M_total, K // 2) # (256, 128) B_packed_shape = (N_total, K // 2) # (64, 128) - C_elem_bytes = tvm_tirx.runtime.DataType(C_dtype).bits // 8 + C_elem_bytes = tvm_ffi.dtype(C_dtype).bits // 8 C_elem_32b = 4 // C_elem_bytes cols_alloc = max(32, next_power_of_2(C_shape[1] // C_elem_32b)) @@ -2166,7 +2167,7 @@ def test_gemm_block_scaled_nvfp4_cta_group_2(): sf_mma_k = 4 # nvfp4: 4 scale factors per MMA iteration sfa_layout = scale_factor_tmem_layout(M_per_cta, SF_K=sf_mma_k * 1, sf_per_mma=sf_mma_k) sfb_layout = scale_factor_tmem_layout(N_total, SF_K=sf_mma_k * 1, sf_per_mma=sf_mma_k) - sf_epc = 32 // tvm_tirx.runtime.DataType(SF_dtype).bits + sf_epc = 32 // tvm_ffi.dtype(SF_dtype).bits SFA_TMEM_SPACING = (int(sfa_layout.span("TCol")) + sf_epc - 1) // sf_epc (int(sfb_layout.span("TCol")) + sf_epc - 1) // sf_epc SFA_TMEM_START = width @@ -2517,9 +2518,9 @@ def test_gemm_block_scaled_fp8_sf_id(): A_shape = (M, K) B_shape = (N, K) - A_elem_bytes = max(1, tvm_tirx.runtime.DataType(A_dtype).bits // 8) - B_elem_bytes = max(1, tvm_tirx.runtime.DataType(B_dtype).bits // 8) - C_elem_bytes = tvm_tirx.runtime.DataType(C_dtype).bits // 8 + A_elem_bytes = max(1, tvm_ffi.dtype(A_dtype).bits // 8) + B_elem_bytes = max(1, tvm_ffi.dtype(B_dtype).bits // 8) + C_elem_bytes = tvm_ffi.dtype(C_dtype).bits // 8 C_elem_32b = 4 // C_elem_bytes cols_alloc = max(32, next_power_of_2(C_shape[1] // C_elem_32b)) @@ -2535,7 +2536,7 @@ def test_gemm_block_scaled_fp8_sf_id(): num_ki = K // MMA_K # 4: distinct SF positions per call sfa_layout = scale_factor_tmem_layout(M, SF_K=sf_mma_k * num_ki, sf_per_mma=sf_mma_k) sfb_layout = scale_factor_tmem_layout(N, SF_K=sf_mma_k * num_ki, sf_per_mma=sf_mma_k) - sf_epc = 32 // tvm_tirx.runtime.DataType(SF_dtype).bits + sf_epc = 32 // tvm_ffi.dtype(SF_dtype).bits SFA_TMEM_SPACING = (int(sfa_layout.span("TCol")) + sf_epc - 1) // sf_epc SFA_TMEM_START = N SFB_TMEM_START = SFA_TMEM_START + SFA_TMEM_SPACING @@ -3044,11 +3045,11 @@ def test_gemm_tcgen05_arbitrary_tiles(task, copy_perm, a_major, b_major): B_ref_perm = B_spec[5] if len(B_spec) > 5 else None M = C_region[0][1] - C_region[0][0] N = C_region[1][1] - C_region[1][0] - C_elem_bytes = tvm_tirx.runtime.DataType(C_dtype).bits // 8 + C_elem_bytes = tvm_ffi.dtype(C_dtype).bits // 8 C_elem_32b = 4 // C_elem_bytes cols_alloc = max(32, next_power_of_2(C_shape[1] // C_elem_32b)) - A_elem_bytes = tvm_tirx.runtime.DataType(A_dtype).bits // 8 - B_elem_bytes = tvm_tirx.runtime.DataType(B_dtype).bits // 8 + A_elem_bytes = tvm_ffi.dtype(A_dtype).bits // 8 + B_elem_bytes = tvm_ffi.dtype(B_dtype).bits // 8 # Accept either swizzle mode (int) or pre-built layout A_layout = ( A_swizzle_mode @@ -3455,11 +3456,11 @@ def _run_dense_gemm( A_swizzle, B_swizzle = 3, 3 A_layout = mma_shared_layout(A_dtype, A_swizzle, A_shape) B_layout = mma_shared_layout(B_dtype, B_swizzle, B_shape) - C_elem_32b = 4 // (tvm_tirx.runtime.DataType(C_dtype).bits // 8) + C_elem_32b = 4 // (tvm_ffi.dtype(C_dtype).bits // 8) cols_alloc = max(32, next_power_of_2(N // C_elem_32b)) total_bytes = functools.reduce(operator.mul, A_shape, 1) * ( - tvm_tirx.runtime.DataType(A_dtype).bits // 8 - ) + functools.reduce(operator.mul, B_shape, 1) * (tvm_tirx.runtime.DataType(B_dtype).bits // 8) + tvm_ffi.dtype(A_dtype).bits // 8 + ) + functools.reduce(operator.mul, B_shape, 1) * (tvm_ffi.dtype(B_dtype).bits // 8) instruction_k = 8 if is_AB_tf32 else 32 kind = "tf32" if is_AB_tf32 else "f8f6f4" atom_width = 32 if is_AB_tf32 else 128 diff --git a/tirx/tests/python/tirx/tile/cuda/test_tcgen05_mma_codegen.py b/tirx/tests/python/tirx/tile/cuda/test_tcgen05_mma_codegen.py index 034e7dc8812b..8a3b61687709 100644 --- a/tirx/tests/python/tirx/tile/cuda/test_tcgen05_mma_codegen.py +++ b/tirx/tests/python/tirx/tile/cuda/test_tcgen05_mma_codegen.py @@ -20,6 +20,7 @@ import re import pytest +import tvm_ffi import tvm_tirx from tirx.layout_fixtures import scale_factor_tmem_layout @@ -386,7 +387,7 @@ def test_gemm_tcgen05_explicit_mma_m_is_cluster_total_for_cta_group_2(rows, inst def _packed_fixture(a_form="packed", dtype="bfloat16", M=64, cta=2): """The old column-major alias is explicitly unfolded to packed 16B lines.""" N, K = 256, 64 - e16 = 128 // tvm_tirx.DataType(dtype).bits + e16 = 128 // tvm_ffi.dtype(dtype).bits a_layout = ( TileLayout(S[(M, K) : (1, M)]) if a_form == "column_major" diff --git a/tirx/tests/python/tirx/transform/test_lower_tvm_builtin.py b/tirx/tests/python/tirx/transform/test_lower_tvm_builtin.py index e11d35c31fce..2a0b34ba8f9e 100644 --- a/tirx/tests/python/tirx/transform/test_lower_tvm_builtin.py +++ b/tirx/tests/python/tirx/transform/test_lower_tvm_builtin.py @@ -17,6 +17,7 @@ import numpy as np import pytest +import tvm_ffi import tvm_tirx import tvm_tirx.testing @@ -25,7 +26,7 @@ from tvm_tirx.testing import env -@tvm_tirx.register_global_func("tvm.test_matmul") +@tvm_ffi.register_global_func("tvm.test_matmul") def my_matmul(a, b, c): c.copyfrom(np.dot(a.numpy(), b.numpy()))