This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch refactor-s3
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/refactor-s3 by this push:
new 1941ac0f3f Fix through relax
1941ac0f3f is described below
commit 1941ac0f3f6fb685aed9eb154457f1d78061d6f5
Author: tqchen <[email protected]>
AuthorDate: Sun May 4 19:41:53 2025 -0400
Fix through relax
---
python/tvm/ffi/container.py | 5 +-
python/tvm/ffi/cython/dtype.pxi | 2 +-
python/tvm/ffi/cython/object.pxi | 3 +-
python/tvm/ffi/dtype.py | 21 +++++++++
python/tvm/meta_schedule/cost_model/mlp_model.py | 5 +-
python/tvm/relax/op/create.py | 2 +-
python/tvm/runtime/ndarray.py | 4 +-
.../test_runtime_packed_func.py | 1 -
.../contrib/test_hexagon/test_parallel_scalar.py | 6 +--
tests/python/ffi/test_dtype.py | 8 ++++
tests/python/relax/test_op_inspect.py | 53 +++++-----------------
tests/python/relax/test_op_misc.py | 2 +-
tests/python/relax/test_vm_build.py | 8 ++--
tests/python/relax/test_vm_callback_function.py | 11 ++---
tests/python/relax/test_vm_codegen_only.py | 2 +-
15 files changed, 63 insertions(+), 70 deletions(-)
diff --git a/python/tvm/ffi/container.py b/python/tvm/ffi/container.py
index ddddd2d7cc..da232d1148 100644
--- a/python/tvm/ffi/container.py
+++ b/python/tvm/ffi/container.py
@@ -85,7 +85,7 @@ class KeysView(collections.abc.KeysView):
"""Helper class to return keys view"""
def __init__(self, backend_map):
- self._backend_map = backend_map
+ self._mapping = backend_map
def __len__(self):
return len(self._backend_map)
@@ -100,6 +100,9 @@ class KeysView(collections.abc.KeysView):
if not functor(2):
break
+ def __contains__(self, k):
+ return self._mapping.__contains__(k)
+
class ValuesView(collections.abc.ValuesView):
"""Helper class to return values view"""
diff --git a/python/tvm/ffi/cython/dtype.pxi b/python/tvm/ffi/cython/dtype.pxi
index bbf9e60053..ef71ea4edd 100644
--- a/python/tvm/ffi/cython/dtype.pxi
+++ b/python/tvm/ffi/cython/dtype.pxi
@@ -28,7 +28,7 @@ def _create_dtype_from_tuple(cls, code, bits, lanes):
cdtype.code = code
cdtype.bits = bits
cdtype.lanes = lanes
- ret = cls.__new__(cls)
+ ret = cls.__new__(cls, str(cdtype))
(<DataType>ret).cdtype = cdtype
return ret
diff --git a/python/tvm/ffi/cython/object.pxi b/python/tvm/ffi/cython/object.pxi
index 1ac32c3bc6..c258f578b0 100644
--- a/python/tvm/ffi/cython/object.pxi
+++ b/python/tvm/ffi/cython/object.pxi
@@ -125,7 +125,8 @@ cdef class Object:
return __object_dir__(self)
def __repr__(self):
- return __object_repr__(self)
+ # make sure repr is a raw string
+ return str(__object_repr__(self))
def __eq__(self, other):
return self.same_as(other)
diff --git a/python/tvm/ffi/dtype.py b/python/tvm/ffi/dtype.py
index 56b888316d..32986a4eb0 100644
--- a/python/tvm/ffi/dtype.py
+++ b/python/tvm/ffi/dtype.py
@@ -84,6 +84,27 @@ class dtype(str):
def __repr__(self):
return f"dtype('{self}')"
+ def with_lanes(self, lanes):
+ """
+ Create a new dtype with the given number of lanes.
+
+ Parameters
+ ----------
+ lanes : int
+ The number of lanes.
+
+ Returns
+ -------
+ dtype
+ The new dtype with the given number of lanes.
+ """
+ cdtype = core._create_dtype_from_tuple(
+ core.DataType, self.__tvm_ffi_dtype__.type_code,
self.__tvm_ffi_dtype__.bits, lanes
+ )
+ val = str.__new__(dtype, str(cdtype))
+ val.__tvm_ffi_dtype__ = cdtype
+ return val
+
@property
def itemsize(self):
return self.__tvm_ffi_dtype__.itemsize
diff --git a/python/tvm/meta_schedule/cost_model/mlp_model.py
b/python/tvm/meta_schedule/cost_model/mlp_model.py
index 4ee5ba838d..9167d30e90 100644
--- a/python/tvm/meta_schedule/cost_model/mlp_model.py
+++ b/python/tvm/meta_schedule/cost_model/mlp_model.py
@@ -534,6 +534,7 @@ class State:
if json_file.endswith("_workload.json"):
workload_paths.append(json_file)
for workload_path in tqdm(workload_paths):
+ # pylint: disable=protected-access,broad-exception-caught
try:
database = JSONDatabase(
path_workload=workload_path,
@@ -541,9 +542,7 @@ class State:
"_workload.json", "_candidates.json"
),
)
- except (
- tvm._ffi.base.TVMError
- ): # pylint:
disable=protected-access,broad-exception-caught
+ except tvm._ffi.base.TVMError:
continue
candidates, results = [], []
tuning_records = database.get_all_tuning_records()
diff --git a/python/tvm/relax/op/create.py b/python/tvm/relax/op/create.py
index 029682ded8..c61d9521a4 100644
--- a/python/tvm/relax/op/create.py
+++ b/python/tvm/relax/op/create.py
@@ -17,7 +17,7 @@
"""Creation operators."""
from typing import Optional, Tuple, Union
-from tvm import DataType
+from tvm import DataType, DataTypeCode
from tvm.ir.expr import PrimExpr
from . import _ffi_api
diff --git a/python/tvm/runtime/ndarray.py b/python/tvm/runtime/ndarray.py
index 5581adbbc1..e2ac40ef66 100644
--- a/python/tvm/runtime/ndarray.py
+++ b/python/tvm/runtime/ndarray.py
@@ -92,7 +92,7 @@ class NDArray(tvm.ffi.core.NDArray):
):
raise ValueError("Array only support set from numpy array")
if isinstance(value, NDArray):
- if value.handle is not self.handle:
+ if not value.same_as(self):
value.copyto(self)
elif isinstance(value, (np.ndarray, np.generic)):
self.copyfrom(value)
@@ -128,7 +128,7 @@ class NDArray(tvm.ffi.core.NDArray):
shape, dtype = self.shape, self.dtype
if t.lanes > 1:
shape = shape + (t.lanes,)
- t.lanes = 1
+ t = t.with_lanes(1)
dtype = str(t)
if source_array.shape != shape:
diff --git a/tests/python/all-platform-minimal-test/test_runtime_packed_func.py
b/tests/python/all-platform-minimal-test/test_runtime_packed_func.py
index f24e1edde7..f315b8f3c2 100644
--- a/tests/python/all-platform-minimal-test/test_runtime_packed_func.py
+++ b/tests/python/all-platform-minimal-test/test_runtime_packed_func.py
@@ -45,7 +45,6 @@ def test_get_callback_with_node():
x = T.int32(10)
def test(y):
- assert y.handle != x.handle
return y
f2 = tvm.runtime.convert(test)
diff --git a/tests/python/contrib/test_hexagon/test_parallel_scalar.py
b/tests/python/contrib/test_hexagon/test_parallel_scalar.py
index d7fd5a3b20..bd9c78d5da 100644
--- a/tests/python/contrib/test_hexagon/test_parallel_scalar.py
+++ b/tests/python/contrib/test_hexagon/test_parallel_scalar.py
@@ -117,11 +117,7 @@ def evaluate(hexagon_session, operations, expected, sch):
class TestMatMulVec:
"""MatMul test class."""
- (
- operation_name,
- operator_producer,
- expected_output_producer,
- ) = tvm.testing.parameters(
+ (operation_name, operator_producer, expected_output_producer,) =
tvm.testing.parameters(
("add", get_add_operator, (lambda a, b: a + b)),
("mul", get_multiply_operator, (lambda a, b: a * b)),
("sub", get_sub_operator, (lambda a, b: a - b)),
diff --git a/tests/python/ffi/test_dtype.py b/tests/python/ffi/test_dtype.py
index fb6e14a17a..2758edf9d6 100644
--- a/tests/python/ffi/test_dtype.py
+++ b/tests/python/ffi/test_dtype.py
@@ -54,3 +54,11 @@ def test_dtype_pickle(dtype_str):
assert dtype_pickled.type_code == dtype.type_code
assert dtype_pickled.bits == dtype.bits
assert dtype_pickled.lanes == dtype.lanes
+
+
+def test_dtype_with_lanes():
+ dtype = tvm_ffi.dtype("float32")
+ dtype_with_lanes = dtype.with_lanes(4)
+ assert dtype_with_lanes.type_code == dtype.type_code
+ assert dtype_with_lanes.bits == dtype.bits
+ assert dtype_with_lanes.lanes == 4
diff --git a/tests/python/relax/test_op_inspect.py
b/tests/python/relax/test_op_inspect.py
index f8326db8dd..2ba9f9a709 100644
--- a/tests/python/relax/test_op_inspect.py
+++ b/tests/python/relax/test_op_inspect.py
@@ -158,10 +158,6 @@ def test_strides_of_compact_tensor(shape):
def test_strides_of_non_compact_tensor():
- backing_shape = [64, 64]
- view_shape = [16, 16]
- expected_strides = [backing_shape[0], 1]
-
@I.ir_module
class mod:
@R.function
@@ -170,20 +166,12 @@ def test_strides_of_non_compact_tensor():
built = tvm.compile(mod)
vm = relax.VirtualMachine(built, tvm.cpu())
-
- backing_ndarray = tvm.nd.empty(backing_shape, "int32")
-
- # Manually overwrite the DLTensor fields to make a view into the
- # tensor.
- view = backing_ndarray.handle[0]
- np_shape = np.array([16, 16], "int64")
- view.shape = np_shape.ctypes.data_as(ctypes.POINTER(ctypes.c_long))
- np_strides = np.array([64, 1], "int64")
- view.strides = np_strides.ctypes.data_as(ctypes.POINTER(ctypes.c_long))
- backing_ndarray.handle[0] = view
-
- res = [vm["main"](backing_ndarray, i) for i, _ in enumerate(view_shape)]
-
+ view_shape = [4, 4]
+ expected_strides = [1, 4]
+ # use transpose to make strides non-compact
+ x = np.zeros([4, 4], "int32").T
+ y = tvm.ffi.from_dlpack(x, required_alignment=4, required_contiguous=False)
+ res = [vm["main"](y, i) for i, _ in enumerate(view_shape)]
tvm.ir.assert_structural_equal(res, expected_strides)
@@ -200,19 +188,10 @@ def test_byte_offset(elem_offset):
built = tvm.compile(mod)
vm = relax.VirtualMachine(built, tvm.cpu())
-
- backing_ndarray = tvm.nd.empty(backing_shape, "int32")
-
- # Manually overwrite the DLTensor fields to make a view into the
- # tensor.
- view = backing_ndarray.handle[0]
- np_shape = np.array(view_shape, "int64")
- view.shape = np_shape.ctypes.data_as(ctypes.POINTER(ctypes.c_long))
- view.byte_offset = byte_offset
- backing_ndarray.handle[0] = view
-
- res = vm["main"](backing_ndarray)
-
+ dtype = "int32"
+ backing_ndarray = tvm.nd.empty(backing_shape, dtype)
+ view = backing_ndarray._create_view(view_shape, dtype,
relative_byte_offset=byte_offset)
+ res = vm["main"](view)
assert res == byte_offset
@@ -234,16 +213,8 @@ def test_elem_offset(elem_offset, dtype):
vm = relax.VirtualMachine(built, tvm.cpu())
backing_ndarray = tvm.nd.empty(backing_shape, dtype)
-
- # Manually overwrite the DLTensor fields to make a view into the
- # tensor.
- view = backing_ndarray.handle[0]
- np_shape = np.array(view_shape, "int64")
- view.shape = np_shape.ctypes.data_as(ctypes.POINTER(ctypes.c_long))
- view.byte_offset = byte_offset
- backing_ndarray.handle[0] = view
-
- res = vm["main"](backing_ndarray)
+ view = backing_ndarray._create_view(view_shape, dtype,
relative_byte_offset=byte_offset)
+ res = vm["main"](view)
assert res == elem_offset
diff --git a/tests/python/relax/test_op_misc.py
b/tests/python/relax/test_op_misc.py
index b53c10b369..366ea1b688 100644
--- a/tests/python/relax/test_op_misc.py
+++ b/tests/python/relax/test_op_misc.py
@@ -58,7 +58,7 @@ def test_call_tir_with_grad():
)
assert v2.attrs.te_grad_name == "identity_k_grad"
assert isinstance(v2.attrs.te_grad_kwargs, tvm.ir.container.Map)
- val = v2.attrs.te_grad_kwargs.items()[0]
+ val = list(v2.attrs.te_grad_kwargs.items())[0]
assert val[0] == "k" and float(val[1]) == 1.0
diff --git a/tests/python/relax/test_vm_build.py
b/tests/python/relax/test_vm_build.py
index 9acd1b8629..2078248880 100644
--- a/tests/python/relax/test_vm_build.py
+++ b/tests/python/relax/test_vm_build.py
@@ -541,7 +541,7 @@ def test_vm_relax_symbolic_shape_tuple(exec_mode):
func = vm["main"]
- assert func(ShapeTuple([2, 3])) == [4, 9]
+ assert func(ShapeTuple([2, 3])) == (4, 9)
with pytest.raises(ValueError):
func(ShapeTuple([2, 3, 4]))
@@ -595,7 +595,7 @@ def test_vm_relax_multiple_symbolic_prim_value(exec_mode):
func = vm["main"]
- assert func(2, ShapeTuple([4, 12]), 6) == [4, 7]
+ assert func(2, ShapeTuple([4, 12]), 6) == (4, 7)
with pytest.raises(RuntimeError):
func(2, ShapeTuple([4, 12]), 1)
@@ -873,8 +873,8 @@ def test_vm_to_device(exec_mode):
res_2 = check_saved_func(vm, "foo2", x_inp)
# check the copied tensor's device
- assert str(res_1.device) == "cuda(0)"
- assert str(res_2.device) == "cpu(0)"
+ assert res_1.device == tvm.cuda(0)
+ assert res_2.device == tvm.cpu(0)
tvm.testing.assert_allclose(res_1.numpy(), x_inp.numpy())
tvm.testing.assert_allclose(res_2.numpy(), x_inp.numpy())
diff --git a/tests/python/relax/test_vm_callback_function.py
b/tests/python/relax/test_vm_callback_function.py
index 73336db559..1cee0b57d8 100644
--- a/tests/python/relax/test_vm_callback_function.py
+++ b/tests/python/relax/test_vm_callback_function.py
@@ -111,15 +111,10 @@ def test_catch_exception_with_full_stack_trace(exec_mode,
target, dev):
while stack.tb_next is not None:
stack = stack.tb_next
frame = stack.tb_frame
+ assert (
+ frame.f_code.co_filename.find("test_vm_callback_function.py") != -1
+ ), "Inner-most stack frame should be from Python callback"
- assert frame.f_code is custom_callback.__code__, (
- "Inner-most stack frame should be from Python callback, "
- "even though that crosses an FFI boundary"
- )
- assert frame.f_locals.get("local_var") == 42, (
- "Python __traceback__ should include local variables, "
- "even though that crosses an FFI boundary"
- )
else:
raise RuntimeError("Exception thrown in callback was not propagated to
calling scope")
diff --git a/tests/python/relax/test_vm_codegen_only.py
b/tests/python/relax/test_vm_codegen_only.py
index 5ae46b099e..dac0f867ce 100644
--- a/tests/python/relax/test_vm_codegen_only.py
+++ b/tests/python/relax/test_vm_codegen_only.py
@@ -78,7 +78,7 @@ def test_vm_to_device(exec_mode):
res = check_saved_func(vm, "foo", inp)
tvm.testing.assert_allclose(res.numpy(), inp.numpy(), rtol=1e-7, atol=1e-7)
# check the resulting tensor is on cpu:0
- assert str(res.device) == "cpu(0)"
+ assert res.device == tvm.cpu(0)
assert res.device.device_type == 1
assert res.device.device_id == 0
