This is an automated email from the ASF dual-hosted git repository. tqchen pushed a commit to branch tirx-standalone-root in repository https://gitbox.apache.org/repos/asf/tvm.git
commit b7a4e89c1871855dd40ed580671d4a68073ddcab Author: Tianqi Chen <[email protected]> AuthorDate: Fri Oct 16 22:37:41 2026 -0400 [REFACTOR][Python] Use TVM-FFI APIs directly (#20769) Use canonical TVM-FFI APIs directly instead of maintaining duplicate TIRx root, runtime and IR exports. Migrate Python callers, examples and references to `tvm_ffi`, while preserving TVM-specific runtime behavior, device convenience constructors and IR diagnostics. --- tirx/python/tvm_tirx/__init__.py | 7 ++- tirx/python/tvm_tirx/backend/cuda/cpp/asm.py | 6 +-- .../tvm_tirx/backend/cuda/cpp/descriptors.py | 2 +- .../tvm_tirx/backend/cuda/lang/alloc_pool.py | 11 ++-- tirx/python/tvm_tirx/backend/cuda/layout.py | 14 ++--- tirx/python/tvm_tirx/backend/cuda/op.py | 4 +- .../backend/cuda/tile/patterns/async_copy.py | 5 +- .../tvm_tirx/backend/cuda/tile/patterns/common.py | 4 +- .../tvm_tirx/backend/cuda/tile/patterns/tma.py | 12 +++-- tirx/python/tvm_tirx/codegen/__init__.py | 2 +- tirx/python/tvm_tirx/codegen/executable.py | 6 +-- tirx/python/tvm_tirx/ir/__init__.py | 2 - tirx/python/tvm_tirx/ir/attrs.py | 4 +- tirx/python/tvm_tirx/ir/base.py | 16 +++--- tirx/python/tvm_tirx/ir/expr.py | 25 +++++---- tirx/python/tvm_tirx/ir/function.py | 11 ++-- tirx/python/tvm_tirx/ir/global_info.py | 4 +- tirx/python/tvm_tirx/ir/location.py | 6 +-- tirx/python/tvm_tirx/ir/module.py | 8 ++- tirx/python/tvm_tirx/ir/stmt.py | 9 ++-- tirx/python/tvm_tirx/ir/transform/core.py | 10 ++-- tirx/python/tvm_tirx/ir/unique_name_supply.py | 4 +- tirx/python/tvm_tirx/runtime/__init__.py | 10 ++-- tirx/python/tvm_tirx/runtime/_tensor.py | 35 ++++-------- tirx/python/tvm_tirx/script/ir_builder/base.py | 10 ++-- tirx/python/tvm_tirx/script/ir_builder/ir.py | 14 ++--- tirx/python/tvm_tirx/script/printer/doc.py | 4 +- tirx/python/tvm_tirx/script/printer/scriptable.py | 28 +++++----- tirx/python/tvm_tirx/support/rocm.py | 2 +- tirx/python/tvm_tirx/sym/analyzer.py | 7 ++- tirx/python/tvm_tirx/sym/int_set.py | 4 +- tirx/python/tvm_tirx/sym/iter_affine_map.py | 5 +- tirx/python/tvm_tirx/target/detect_target.py | 5 +- tirx/python/tvm_tirx/target/target.py | 12 ++--- tirx/python/tvm_tirx/testing/env.py | 12 +++-- tirx/python/tvm_tirx/testing/utils.py | 3 +- tirx/python/tvm_tirx/tirx/_compile.py | 9 ++-- tirx/python/tvm_tirx/tirx/_tensor_view.py | 10 ++-- tirx/python/tvm_tirx/tirx/function.py | 3 +- tirx/python/tvm_tirx/tirx/index_map.py | 3 +- tirx/python/tvm_tirx/tirx/layout.py | 5 +- tirx/python/tvm_tirx/tirx/op/__init__.py | 4 +- .../tirx/script/ir_builder/external_kernel.py | 9 ++-- tirx/python/tvm_tirx/tirx/script/ir_builder/ir.py | 10 ++-- .../tirx/script/ir_builder/parser_protocol.py | 8 +-- .../tvm_tirx/tirx/script/ir_builder/triton.py | 4 +- tirx/python/tvm_tirx/tirx/tensor.py | 6 +-- tirx/python/tvm_tirx/tirx/tile/dispatch.py | 27 +++++----- .../python/backend/cuda/test_module_export.py | 7 +-- .../python/backend/hexagon/codegen/test_basic.py | 3 +- .../python/backend/metal/codegen/test_basic.py | 12 ++--- .../python/backend/metal/test_compile_callback.py | 4 +- .../python/backend/vulkan/codegen/test_basic.py | 3 +- .../python/backend/webgpu/codegen/test_basic.py | 3 +- .../python/codegen/gpu_common/test_allreduce.py | 6 +-- .../tests/python/codegen/gpu_common/test_device.py | 5 +- .../python/codegen/gpu_common/test_intrinsics.py | 5 +- tirx/tests/python/codegen/test_common.py | 3 +- tirx/tests/python/codegen/test_extern.py | 3 +- tirx/tests/python/codegen/test_static_init.py | 3 +- tirx/tests/python/ir/test_attrs.py | 4 +- tirx/tests/python/ir/test_json.py | 2 +- tirx/tests/python/ir/test_type.py | 5 +- tirx/tests/python/runtime/test_device_api.py | 2 +- tirx/tests/python/runtime/test_memory_scope.py | 3 +- tirx/tests/python/runtime/test_module.py | 6 +-- tirx/tests/python/runtime/test_tensor.py | 5 +- tirx/tests/python/script/test_basic_usage.py | 5 +- tirx/tests/python/sym/test_analyzer_object.py | 4 +- .../tests/python/target/test_llvm_features_info.py | 15 +++--- tirx/tests/python/target/test_x86_features.py | 8 +-- tirx/tests/python/tirx/base/test_index_map.py | 10 ++-- tirx/tests/python/tirx/base/test_nodes.py | 9 ++-- .../python/tirx/base/test_scalable_datatype.py | 5 +- tirx/tests/python/tirx/base/test_tensor_type.py | 3 +- tirx/tests/python/tirx/base/test_texture_scope.py | 4 +- tirx/tests/python/tirx/codegen/test_blackwell.py | 2 +- tirx/tests/python/tirx/codegen/test_cuda.py | 9 ++-- tirx/tests/python/tirx/codegen/test_hopper.py | 17 +++--- tirx/tests/python/tirx/layout_fixtures.py | 5 +- .../script/test_printer_invisible_path_info.py | 3 +- tirx/tests/python/tirx/test_cuda_backend_config.py | 2 +- .../python/tirx/tile/cuda/test_cp_async_bulk.py | 4 +- .../tirx/tile/cuda/test_cp_async_bulk_tensor.py | 4 +- .../python/tirx/tile/cuda/test_ld_st_gmem_smem.py | 3 +- .../python/tirx/tile/cuda/test_ld_st_registers.py | 4 +- .../python/tirx/tile/cuda/test_pattern_math.py | 3 +- .../python/tirx/tile/cuda/test_pattern_mma.py | 5 +- .../tests/python/tirx/tile/cuda/test_tcgen05_cp.py | 5 +- .../python/tirx/tile/cuda/test_tcgen05_ld_st.py | 15 +++--- .../python/tirx/tile/cuda/test_tcgen05_mma.py | 63 +++++++++++----------- .../tirx/tile/cuda/test_tcgen05_mma_codegen.py | 3 +- .../tirx/transform/test_lower_tvm_builtin.py | 3 +- 93 files changed, 362 insertions(+), 356 deletions(-) diff --git a/tirx/python/tvm_tirx/__init__.py b/tirx/python/tvm_tirx/__init__.py index 8244906d9c..a6e35fd6c6 100644 --- a/tirx/python/tvm_tirx/__init__.py +++ b/tirx/python/tvm_tirx/__init__.py @@ -37,17 +37,16 @@ if TYPE_CHECKING or not _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 9b4e558e48..2571abc036 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 @@ the PTX text, or the offset-scaling ``cp.async`` form the legacy 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.<v>.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 codegen_s_tir_cp_async_raw(*args): 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 182a261faf..995db8a58f 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 @@ device_intrinsic( 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 f47d6a125d..d1ff8bad2e 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 @@ from __future__ import annotations 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 @@ class SMEMPool: # 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 20c0ee9b9f..f5df2a488a 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 tmem_mma_operand_layout( 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 7bbee1efe4..d21a8952e8 100644 --- a/tirx/python/tvm_tirx/backend/cuda/op.py +++ b/tirx/python/tvm_tirx/backend/cuda/op.py @@ -25,7 +25,7 @@ from functools import cache 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 3e25521419..4ae853aed5 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 @@ cache_policy. DSMEM instead requires mbar and remote_cta_id, both interpreted 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 272bb95092..97797b7cda 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 @@ from tvm_tirx.ir import ( 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 @@ class Region: 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 a467872eb4..08a73682e0 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 @@ opaque TensorMap and packed/interleaved representations. 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 5396028ac1..7d1ddfa800 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 674b47c127..98e10bd0ff 100644 --- a/tirx/python/tvm_tirx/codegen/executable.py +++ b/tirx/python/tvm_tirx/codegen/executable.py @@ -59,7 +59,7 @@ class Executable: 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 @@ class Executable: 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 @@ class Executable: 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 a577eeeb6b..f51dae0f41 100644 --- a/tirx/python/tvm_tirx/ir/__init__.py +++ b/tirx/python/tvm_tirx/ir/__init__.py @@ -87,7 +87,5 @@ from .stmt import ( 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 b2136ec540..02fdecd563 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 07e4f46122..ed7f76fdab 100644 --- a/tirx/python/tvm_tirx/ir/base.py +++ b/tirx/python/tvm_tirx/ir/base.py @@ -20,8 +20,6 @@ import tvm_ffi 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 @@ class EnvFunc(Object): 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 1165ccfabb..d026dda51c 100644 --- a/tirx/python/tvm_tirx/ir/expr.py +++ b/tirx/python/tvm_tirx/ir/expr.py @@ -27,7 +27,6 @@ import tvm_ffi 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 @@ class ExprOperand: 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 @@ class Var(_CallableExprWithOp): 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 @@ class Range(Node, Scriptable): """ 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 c7814119fb..eb16149882 100644 --- a/tirx/python/tvm_tirx/ir/function.py +++ b/tirx/python/tvm_tirx/ir/function.py @@ -22,7 +22,6 @@ from enum import IntEnum 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 @@ class BaseFunc(_CallableExprWithOp): 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 @@ class BaseFunc(_CallableExprWithOp): 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 788a0d8a11..c6991cdbf3 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 eeeb35d62c..5f3d7f962d 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 @@ class SourceName(Object): @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 3be29155a6..5857721360 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 @@ class IRModule(Node, Scriptable): attr_key : str The attribute key. - attr_value : Object + attr_value : tvm_ffi.Object The new attribute value. Returns @@ -266,11 +264,11 @@ class IRModule(Node, Scriptable): 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 2197ffd480..98ffef95d0 100644 --- a/tirx/python/tvm_tirx/ir/stmt.py +++ b/tirx/python/tvm_tirx/ir/stmt.py @@ -24,13 +24,12 @@ import tvm_ffi 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 @@ class For(Stmt): 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 e85100035c..a013f7d0fe 100644 --- a/tirx/python/tvm_tirx/ir/transform/core.py +++ b/tirx/python/tvm_tirx/ir/transform/core.py @@ -22,13 +22,11 @@ import inspect 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 @@ class PassInfo(tvm_tirx.runtime.Object): @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 pass_instrument(pi_cls=None): @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 @@ class PassContext(tvm_tirx.runtime.Object): @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 98cd20433e..6b11dc4f5b 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 aa8f098b9d..2722d8f6c1 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 .memory_scope import MemoryScope 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 af3f79b4ad..c7cd312f1b 100644 --- a/tirx/python/tvm_tirx/runtime/_tensor.py +++ b/tirx/python/tvm_tirx/runtime/_tensor.py @@ -30,26 +30,13 @@ except ImportError: 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 47b4bd756a..ddc32545fd 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 functools import wraps 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 resolve_global_info_args( @_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 @@ class IRBuilderFrame(_Object): @_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 @@ class IRBuilder(_Object): """ 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 1153a8d85b..24ee0209d5 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 @@ from .frame import IRModuleFrame 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 775dd33311..d2c6c47060 100644 --- a/tirx/python/tvm_tirx/script/printer/doc.py +++ b/tirx/python/tvm_tirx/script/printer/doc.py @@ -20,17 +20,17 @@ from collections.abc import Sequence 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 a54d2ccf35..c6823ccbad 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 @@ class PrinterConfig(Object): 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 @@ class PrinterConfig(Object): ) -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 @@ class _ScriptableMethods: 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 @@ class _ScriptableMethods: 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 @@ class _ScriptableMethods: 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 @@ class _ScriptableMethods: 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 962b999787..545ea96b0a 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 2925d4f25d..7bfb2252f8 100644 --- a/tirx/python/tvm_tirx/sym/analyzer.py +++ b/tirx/python/tvm_tirx/sym/analyzer.py @@ -23,7 +23,6 @@ from contextlib import contextmanager 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 @@ class ModularSet(Object): @tvm_ffi.register_object("sym.ConstIntBound") -class ConstIntBound(Object): +class ConstIntBound(tvm_ffi.Object): """Represent constant integer bound Parameters @@ -105,7 +104,7 @@ class ConstraintScope: @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 67d28c4c43..15ea563334 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 1c8112c570..7865e5daa8 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 @@ from enum import IntEnum 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 @@ class IterSumExpr(IterMapExpr): @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 4d32ec8d36..c20847fe27 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 3724502d90..b6598a3b9f 100644 --- a/tirx/python/tvm_tirx/target/target.py +++ b/tirx/python/tvm_tirx/target/target.py @@ -20,13 +20,13 @@ import tvm_ffi 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 @@ class TargetFeatures: @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 @@ class Target(Object): 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 6ce5d3cd34..2200d0817f 100644 --- a/tirx/python/tvm_tirx/testing/env.py +++ b/tirx/python/tvm_tirx/testing/env.py @@ -49,6 +49,8 @@ Three kinds of probe live here: import functools import os +import tvm_ffi + import tvm_tirx __all__ = [ @@ -90,9 +92,9 @@ __all__ = [ @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 970799095b..1ddc894ce3 100644 --- a/tirx/python/tvm_tirx/testing/utils.py +++ b/tirx/python/tvm_tirx/testing/utils.py @@ -71,6 +71,7 @@ import time 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 f88c114ee2..1465c409c4 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 _split_host_device_modules( 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 a8d100d172..01c2ff4245 100644 --- a/tirx/python/tvm_tirx/tirx/_tensor_view.py +++ b/tirx/python/tvm_tirx/tirx/_tensor_view.py @@ -23,6 +23,8 @@ import re 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 view(buf: Var, *args, **kwargs) -> Var: ) 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 daa9b20d30..724d5c9957 100644 --- a/tirx/python/tvm_tirx/tirx/function.py +++ b/tirx/python/tvm_tirx/tirx/function.py @@ -25,7 +25,6 @@ import tvm_tirx 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 @@ class Function(BaseFunc, Scriptable): 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 0ddc175520..9c7323b813 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_ffi 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 3a20a07e6f..057ead1e31 100644 --- a/tirx/python/tvm_tirx/tirx/layout.py +++ b/tirx/python/tvm_tirx/tirx/layout.py @@ -24,7 +24,6 @@ import tvm_ffi 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 @@ from .thread_scope import ThreadScope @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 c65dac08a9..15f6e662b7 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 f5c53c4994..0a4e149008 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 __version__ as tvm_version 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 @@ class BaseKernel: # pylint: disable=too-few-public-methods Returns ------- - kernel_module : Module + kernel_module : tvm_ffi.Module The CUDA module. """ tvm_metadata = self._format_tvm_module_metadata( @@ -124,7 +123,7 @@ class SourceKernel(BaseKernel): # pylint: disable=too-few-public-methods 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 2affe6d25c..a79838854e 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 @@ import threading 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 alloc_tensor( 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 8ef324f802..7f297df198 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 @@ import builtins as _python 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 33ae1d87f4..d7dbfe6ae7 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 @@ class TritonKernel(BaseKernel): 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 aab69a485e..43c476ec30 100644 --- a/tirx/python/tvm_tirx/tirx/tensor.py +++ b/tirx/python/tvm_tirx/tirx/tensor.py @@ -100,7 +100,7 @@ class TensorType(Type): 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 @@ class TensorType(Type): 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 @@ class TensorType(Type): 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 8da1f65e87..48e476d700 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 @@ from ..thread_scope import ThreadScope @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 @@ class DispatchContext(Object, Scriptable): 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 @@ class DispatchContext(Object, Scriptable): """ _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 @@ class DispatchContext(Object, Scriptable): 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 f356ac39f1..9f4b75db9c 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 ctypes import numpy as np import pytest +import tvm_ffi import tvm_tirx import tvm_tirx.support.nvcc @@ -94,10 +95,10 @@ def test_cuda_multi_lib(): 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 7e300f2e71..b46b83427b 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 @@ import re from pathlib import Path import pytest +import tvm_ffi from tvm_ffi import register_global_func import tvm_tirx @@ -31,7 +32,7 @@ from tvm_tirx.testing import env @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 6a3a01358d..60dac4c89e 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 test_func_with_trailing_pod_params(): 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 test_func_with_trailing_pod_params(): 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 test_codegen_simdgroup_buffer_data(): 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 test_codegen_simdgroup_buffer_data(): 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 test_codegen_pointer_byte_offsets_preserve_storage_scope(): 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 dff1253a64..9fc04ceec0 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 test_metal_compile_callback_source_passthrough(): 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 test_metal_compile_callback_mixed_formats_rejected(): # 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 e06ac31926..be20d5e9ff 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 re import numpy as np import pytest +import tvm_ffi import tvm_tirx import tvm_tirx.testing @@ -527,7 +528,7 @@ def test_codegen_decl_buffer(): 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 20bc5b7595..174fa2097b 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 test_codegen_buffer_access_modes(): 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 5a26cd1412..14bce0a2b4 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 @@ define_metal_compile_callback = tvm_tirx.testing.parameter(True, False) @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 optional_metal_compile_callback(define_metal_compile_callback): 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 2ad2fb2722..a1f9982ea9 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 test_large_uint_imm(): 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 test_add_pipeline(): 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 11cc3b6c9c..1726c5380a 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 @@ from functools import partial import numpy as np import pytest +import tvm_ffi import tvm_tirx import tvm_tirx.testing @@ -96,7 +97,7 @@ def test_device_module_dump(): 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 test_device_module_dump(): 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 104138f92e..f031cebc58 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 test_decl_buffer_offset_preserves_storage_scope(target, qualifier): 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 b65f853482..44d68508e0 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 @@ def test_pack_buffer_simple(): 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 ad88e1e99c..8900e2f471 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 @@ from tvm_tirx.script import tirx as T 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 10ffba2095..c39a9c4e51 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 776005f69d..df5aec1250 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 5445401f71..1e5b87497e 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 715b0bb9fb..01a977ea78 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 c2b1f51129..042001e41a 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 @@ for name in ("global", "shared", "shared.dyn", "warp", "local"): 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 dbcc77d73a..9d7b80a509 100644 --- a/tirx/tests/python/runtime/test_module.py +++ b/tirx/tests/python/runtime/test_module.py @@ -47,9 +47,8 @@ static RegisterSystemLibrary register_system_library; 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 d564d27433..e1897d8c13 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 a7718074b4..54183f0d33 100644 --- a/tirx/tests/python/script/test_basic_usage.py +++ b/tirx/tests/python/script/test_basic_usage.py @@ -22,13 +22,14 @@ import inspect 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 b3c239f81a..9d10a76bd0 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 df2b9643d0..13326dfa91 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 59da3b2015..c56d96a487 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 134f9d2453..7136e12f82 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 @@ padding_test_case = tvm_tirx.testing.parameter( 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 @@ padding_test_case = tvm_tirx.testing.parameter( 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 @@ padding_test_case = tvm_tirx.testing.parameter( 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 @@ padding_test_case = tvm_tirx.testing.parameter( 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 1ce2d99cd6..a0840dde38 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 a5a11a9759..071800afd7 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 @@ Tests for scalable data types. 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 test_tvm_script_create_scalable_tir_intrin(): 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 8e47538640..4973f77593 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 acc658911b..a15a4ebcb0 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 test_texture_scope(): 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 73dff1d157..dbf8cb69af 100644 --- a/tirx/tests/python/tirx/codegen/test_blackwell.py +++ b/tirx/tests/python/tirx/codegen/test_blackwell.py @@ -403,7 +403,7 @@ def test_tcgen05_ld_st_roundtrip(): @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 1bb95ed34e..629880d3d7 100644 --- a/tirx/tests/python/tirx/codegen/test_cuda.py +++ b/tirx/tests/python/tirx/codegen/test_cuda.py @@ -20,6 +20,7 @@ import re 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 test_warp_storage_requires_explicit_local_layout(target_kind): 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 test_typed_buffer_bindings(monkeypatch): 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 63f9ff48ce..34426a6625 100644 --- a/tirx/tests/python/tirx/codegen/test_hopper.py +++ b/tirx/tests/python/tirx/codegen/test_hopper.py @@ -19,6 +19,7 @@ import math 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 test_cp_async_bulk_tensor_global_to_shared_swizzle(swizzle, dtype): 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 test_wgmma_ss_nt(): ): 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 test_wgmma_ss_nt(): 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 test_wgmma_rs_nt(): 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 test_wgmma_rs_nt(): 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 e58b68b02d..fdbb8446b8 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 @@ _TCGEN05_COL_FACTOR_FP32 = {"32x32b": 1, "16x64b": 2, "16x128b": 4, "16x256b": 8 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 51ca34485c..bf4eb80334 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 0407ff581a..0963e9311f 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 @@ assert 'tvm_tirx' not in sys.modules 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 654eccaf9c..571faf8e53 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 937f0bd67c..7c44299858 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 12218563ce..695fe6c1ff 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 math 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 f9e1df7e97..dc06902ae0 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 test_reg_roundtrip_gapped_permuted_storage(): 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 fa309958b8..c43219598e 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 test_ewise_composite(operation): ) @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 74fe7724f0..3474176a32 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 2c1d05818d..2513fb5cc0 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 itertools 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 0b221613c1..9430c1ce75 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 @@ that the M=64 fragment's row positions hold the expected register data. 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 .<shape>.x<rep> 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 test_copy_tmem2reg_async(dtype, width_32b): 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 test_copy_tmem2reg(dtype, width_32b, offset_32b): 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 test_copy_tmem2reg_sliced_local(dtype, width_32b, local_offset_32b): 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 e5161ca06a..0355aac865 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 operator 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 034e7dc881..8a3b616877 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 e11d35c31f..2a0b34ba8f 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.script import tirx as T 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()))
