This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new a82b34dc9f [Tests] Reduce redundant ONNX and PyTorch integration tests
(#20026)
a82b34dc9f is described below
commit a82b34dc9f545b95b4645e9741379157dcf5b1ce
Author: Shushi Hong <[email protected]>
AuthorDate: Sat Jul 18 17:16:31 2026 -0400
[Tests] Reduce redundant ONNX and PyTorch integration tests (#20026)
This PR reduces repeated Relax frontend and integration test work while
preserving distinct coverage. It reworks ONNX ConvTranspose tests into
direct importer checks plus 11 numerical cases covering all ranks,
asymmetric padding, grouping, bias, dilation, and output padding.
Targeted runtime improves from 17.97s to 3.55s.
- removes duplicate ONNX Pow, unused dynamic Squeeze parameterizations,
and irrelevant Resize ROI value permutations.
- consolidates overlapping PyTorch integration tests while preserving
symbolic shapes, TIR, I.pyfunc, and packed-function coverage.
- removes the redundant BasePyModule aggregate suite, moves its unique
output-only call_tir case into the DLPack test, and removes a DLPack
test that swallowed all exceptions.
---
tests/python/relax/test_base_py_module.py | 219 -----------------
tests/python/relax/test_dlpack_integration.py | 32 +--
tests/python/relax/test_frontend_onnx.py | 255 +++++++++++---------
tests/python/relax/test_pytorch_integration.py | 311 ++-----------------------
tests/python/relax/test_tvmscript_pyfunc.py | 1 +
5 files changed, 186 insertions(+), 632 deletions(-)
diff --git a/tests/python/relax/test_base_py_module.py
b/tests/python/relax/test_base_py_module.py
deleted file mode 100644
index bdbeabf84b..0000000000
--- a/tests/python/relax/test_base_py_module.py
+++ /dev/null
@@ -1,219 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-# ruff: noqa: F401
-"""
-Test BasePyModule core functionality.
-
-This test verifies:
-1. BasePyModule instantiation and basic methods
-2. TIR function compilation and execution
-3. Python function integration
-4. DLPack conversion between PyTorch and TVM
-"""
-
-import numpy as np
-import pytest
-import torch
-
-import tvm
-from tvm import relax, tirx
-from tvm.relax import BasePyModule
-from tvm.script import relax as R
-from tvm.script import tirx as T
-
-
-class TestBasePyModule:
- """Test BasePyModule core functionality."""
-
- def test_base_py_module_instantiation(self):
- @T.prim_func(s_tir=True)
- def simple_func(A: T.Buffer((10,), "float32"), B: T.Buffer((10,),
"float32")):
- for i in T.grid(10):
- B[i] = A[i] * 2.0
-
- ir_mod = tvm.IRModule({"simple_func": simple_func})
- device = tvm.cpu(0)
- py_mod = BasePyModule(ir_mod, device)
-
- assert isinstance(py_mod, BasePyModule)
- assert hasattr(py_mod, "call_tir")
- assert hasattr(py_mod, "call_dps_packed")
- assert hasattr(py_mod, "compiled_tir_funcs")
-
- def test_base_py_module_instantiation_gpu(self):
- @T.prim_func(s_tir=True)
- def simple_func(A: T.Buffer((10,), "float32"), B: T.Buffer((10,),
"float32")):
- for i in T.grid(10):
- B[i] = A[i] * 2.0
-
- ir_mod = tvm.IRModule({"simple_func": simple_func})
-
- if tvm.cuda().exist:
-
- def run_and_check():
- device = tvm.cuda(0)
- py_mod = BasePyModule(ir_mod, device)
-
- assert isinstance(py_mod, BasePyModule)
- assert hasattr(py_mod, "call_tir")
- assert hasattr(py_mod, "call_dps_packed")
- assert hasattr(py_mod, "compiled_tir_funcs")
- # Check if target contains "cuda" instead of exact match
- assert "cuda" in str(py_mod.target)
-
- tvm.testing.run_with_gpu_lock(run_and_check)
- else:
- pytest.skip("CUDA not available")
-
- def test_tir_function_compilation(self):
- @T.prim_func(s_tir=True)
- def add_func(
- A: T.Buffer((5,), "float32"), B: T.Buffer((5,), "float32"), C:
T.Buffer((5,), "float32")
- ):
- for i in T.grid(5):
- C[i] = A[i] + B[i]
-
- ir_mod = tvm.IRModule({"add_func": add_func})
- device = tvm.cpu(0)
- py_mod = BasePyModule(ir_mod, device)
-
- assert "add_func" in py_mod.tir_func_names
- assert "add_func" in py_mod.compiled_tir_funcs
-
- def test_call_tir_with_pytorch_tensors(self):
- @T.prim_func(s_tir=True)
- def scale_func(A: T.Buffer((4,), "float32"), B: T.Buffer((4,),
"float32")):
- for i in T.grid(4):
- B[i] = A[i] * T.float32(2.5)
-
- ir_mod = tvm.IRModule({"scale_func": scale_func})
- device = tvm.cpu(0)
- py_mod = BasePyModule(ir_mod, device)
-
- input_tensor = torch.tensor([1.0, 2.0, 3.0, 4.0], dtype=torch.float32)
- scale_value = 2.5
-
- result = py_mod.call_tir(scale_func, [input_tensor], R.Tensor((4,),
"float32"))
-
- assert isinstance(result, torch.Tensor)
- assert result.shape == (4,)
- expected = input_tensor * scale_value
- assert torch.allclose(result, expected, atol=1e-5)
-
- def test_call_tir_with_pytorch_tensors_gpu(self):
- if tvm.cuda().exist:
-
- def run_and_check():
- # Create a simple IRModule without TIR functions for GPU
testing
- ir_mod = tvm.IRModule({})
- device = tvm.cuda(0)
- py_mod = BasePyModule(ir_mod, device)
-
- # Test basic GPU functionality without TIR compilation issues
- assert isinstance(py_mod, BasePyModule)
- assert hasattr(py_mod, "call_tir")
- assert hasattr(py_mod, "call_dps_packed")
- assert "cuda" in str(py_mod.target)
-
- # Test that we can create GPU tensors and they work
- input_tensor = torch.tensor(
- [1.0, 2.0, 3.0, 4.0], dtype=torch.float32, device="cuda"
- )
- assert input_tensor.device.type == "cuda"
- assert input_tensor.shape == (4,)
-
- tvm.testing.run_with_gpu_lock(run_and_check)
- else:
- pytest.skip("CUDA not available")
-
- def test_dlpack_conversion_pytorch_to_tvm(self):
- @T.prim_func(s_tir=True)
- def identity_func(A: T.Buffer((3,), "float32"), B: T.Buffer((3,),
"float32")):
- for i in T.grid(3):
- B[i] = A[i]
-
- ir_mod = tvm.IRModule({"identity_func": identity_func})
- device = tvm.cpu(0)
- py_mod = BasePyModule(ir_mod, device)
-
- input_tensor = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32)
-
- result = py_mod.call_tir(identity_func, [input_tensor], R.Tensor((3,),
"float32"))
-
- assert isinstance(result, torch.Tensor)
- assert torch.allclose(result, input_tensor, atol=1e-5)
-
- def test_dlpack_conversion_tvm_to_pytorch(self):
- @T.prim_func(s_tir=True)
- def constant_func(B: T.Buffer((2,), "float32")):
- for i in T.grid(2):
- B[i] = T.float32(5.0)
-
- ir_mod = tvm.IRModule({"constant_func": constant_func})
- device = tvm.cpu(0)
- py_mod = BasePyModule(ir_mod, device)
-
- result = py_mod.call_tir(constant_func, [], R.Tensor((2,), "float32"))
-
- assert isinstance(result, torch.Tensor)
- assert result.shape == (2,)
- expected = torch.tensor([5.0, 5.0], dtype=torch.float32)
- assert torch.allclose(result, expected, atol=1e-5)
-
- def test_add_python_function(self):
- ir_mod = tvm.IRModule({})
- device = tvm.cpu(0)
- py_mod = BasePyModule(ir_mod, device)
-
- def custom_activation(x):
- return torch.tanh(x)
-
- py_mod.add_python_function("custom_activation", custom_activation)
-
- assert hasattr(py_mod, "custom_activation")
- assert "custom_activation" in py_mod.pyfuncs
-
- input_tensor = torch.tensor([1.0, -1.0, 0.0], dtype=torch.float32)
- result = py_mod.custom_activation(input_tensor)
-
- assert isinstance(result, torch.Tensor)
- expected = torch.tanh(input_tensor)
- assert torch.allclose(result, expected, atol=1e-5)
-
- def test_call_dps_packed_with_python_function(self):
- ir_mod = tvm.IRModule({})
- device = tvm.cpu(0)
- py_mod = BasePyModule(ir_mod, device)
-
- def my_softmax(tensor, dim):
- return torch.softmax(tensor, dim=dim)
-
- py_mod.add_python_function("my_softmax", my_softmax)
-
- input_tensor = torch.tensor([[1.0, 2.0], [3.0, 4.0]],
dtype=torch.float32)
-
- result = py_mod.call_dps_packed(
- "my_softmax", [input_tensor, 1], R.Tensor((2, 2), "float32")
- )
-
- assert isinstance(result, torch.Tensor)
- expected = torch.softmax(input_tensor, dim=1)
- assert torch.allclose(result, expected, atol=1e-5)
-
-
-if __name__ == "__main__":
- tvm.testing.main()
diff --git a/tests/python/relax/test_dlpack_integration.py
b/tests/python/relax/test_dlpack_integration.py
index 9023b4fab6..734ba7b0c6 100644
--- a/tests/python/relax/test_dlpack_integration.py
+++ b/tests/python/relax/test_dlpack_integration.py
@@ -14,7 +14,6 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-# ruff: noqa: F401, F811, F841
"""
Test DLPack integration between PyTorch and TVM.
@@ -23,15 +22,13 @@ This test verifies:
2. DLPack conversion from TVM to PyTorch
3. Data integrity preservation during conversion
4. Functionality equivalence between DLPack and numpy fallback
-5. Error handling for unsupported data types
"""
-import numpy as np
import pytest
import torch
import tvm
-from tvm import relax, tirx
+import tvm.testing
from tvm.relax import BasePyModule
from tvm.script import relax as R
from tvm.script import tirx as T
@@ -200,21 +197,6 @@ class TestDLPackIntegration:
assert result_dlpack.shape == pytorch_tensor.shape
assert result_dlpack.dtype == pytorch_tensor.dtype
- def test_dlpack_error_handling(self):
- """Test DLPack error handling for unsupported operations."""
- # Test with non-contiguous tensor
- pytorch_tensor = torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0],
dtype=torch.float32)
- non_contiguous = pytorch_tensor[::2] # Create non-contiguous view
-
- # This should work (PyTorch handles non-contiguous tensors)
- try:
- tvm_tensor = tvm.runtime.from_dlpack(non_contiguous)
- result_tensor = torch.from_dlpack(tvm_tensor)
- assert torch.allclose(non_contiguous, result_tensor, atol=1e-5)
- except Exception as e:
- # If it fails, that's also acceptable
- pass
-
def test_dlpack_with_base_py_module(self):
"""Test DLPack conversion within BasePyModule context."""
@@ -224,7 +206,12 @@ class TestDLPackIntegration:
for i in T.grid(3):
B[i] = A[i]
- ir_mod = tvm.IRModule({"identity_func": identity_func})
+ @T.prim_func(s_tir=True)
+ def constant_func(B: T.Buffer((2,), "float32")):
+ for i in T.grid(2):
+ B[i] = T.float32(5.0)
+
+ ir_mod = tvm.IRModule({"identity_func": identity_func,
"constant_func": constant_func})
device = tvm.cpu(0)
py_mod = BasePyModule(ir_mod, device)
@@ -238,6 +225,11 @@ class TestDLPackIntegration:
assert isinstance(result, torch.Tensor)
assert torch.allclose(result, input_tensor, atol=1e-5)
+ # Preserve the output-only call_tir path, where there are no input
tensors
+ # from which to infer or transfer the output.
+ result = py_mod.call_tir(constant_func, [], R.Tensor((2,), "float32"))
+ torch.testing.assert_close(result, torch.full((2,), 5.0))
+
def test_dlpack_device_consistency(self):
"""Test DLPack conversion maintains device consistency."""
# Test CPU tensor
diff --git a/tests/python/relax/test_frontend_onnx.py
b/tests/python/relax/test_frontend_onnx.py
index 6b449624b2..7e3da83df0 100644
--- a/tests/python/relax/test_frontend_onnx.py
+++ b/tests/python/relax/test_frontend_onnx.py
@@ -3549,91 +3549,155 @@ def test_conv_numerical(nd, groups, auto_pad, stride,
dilation, pad, bias):
check_correctness(model, opset=14, atol=1e-4)
[email protected]("stride", [2])
[email protected]("dilation", [1])
[email protected]("bias", [True, False])
[email protected]("pad", [0, 2])
[email protected]("output_pad", [0, 1])
-def test_conv_transpose(stride: int, dilation: int, pad: int, bias: bool,
output_pad: int):
- def _verify_conv_transpose(input_shape, weight_shape):
- nd = len(weight_shape) - 2
- output_shape = [input_shape[0], weight_shape[0]] + [
- (input_shape[i] - 1) * stride
- - 2 * pad
- + dilation * (weight_shape[i] - 1)
- + output_pad
- + 1
- for i in range(2, len(input_shape))
+def _make_conv_transpose_model(
+ nd, groups, auto_pad, stride, dilation, pad, bias, output_pad,
spatial_extent=8
+):
+ input_shape = [1, 4] + [spatial_extent] * nd
+ weight_shape = [4, 4 // groups] + [3] * nd
+ output_channels = weight_shape[1] * groups
+ effective_kernel = dilation * (weight_shape[2] - 1) + 1
+
+ node_attrs = {
+ "strides": [stride] * nd,
+ "dilations": [dilation] * nd,
+ "output_padding": [output_pad] * nd,
+ "group": groups,
+ }
+ if auto_pad == "NOTSET":
+ node_attrs["pads"] = [pad] * nd * 2
+ output_spatial = [
+ (spatial_extent - 1) * stride - 2 * pad + effective_kernel +
output_pad
+ ] * nd
+ elif auto_pad == "VALID":
+ node_attrs["auto_pad"] = auto_pad
+ node_attrs["kernel_shape"] = weight_shape[2:]
+ output_spatial = [(spatial_extent - 1) * stride + effective_kernel +
output_pad] * nd
+ else:
+ node_attrs["auto_pad"] = auto_pad
+ node_attrs["kernel_shape"] = weight_shape[2:]
+ output_spatial = [spatial_extent * stride] * nd
+
+ output_shape = [input_shape[0], output_channels, *output_spatial]
+ conv_node = helper.make_node(
+ "ConvTranspose",
+ inputs=["x", "w"] + (["b"] if bias else []),
+ outputs=["y"],
+ **node_attrs,
+ )
+ graph = helper.make_graph(
+ [conv_node],
+ "conv_transpose_test",
+ inputs=[
+ helper.make_tensor_value_info("x", TensorProto.FLOAT, input_shape),
+ helper.make_tensor_value_info("w", TensorProto.FLOAT,
weight_shape),
]
- bias_shape = [output_shape[1]]
- conv_node = helper.make_node(
- "ConvTranspose",
- inputs=["x", "w"] + (["b"] if bias else []),
- outputs=["y"],
- strides=[stride] * nd,
- dilations=[dilation] * nd,
- pads=[pad] * nd * 2,
- output_padding=[output_pad] * nd,
- group=input_shape[1] // weight_shape[1],
- )
- graph = helper.make_graph(
- [conv_node],
- "conv_transpose_test",
- inputs=[
- helper.make_tensor_value_info("x", TensorProto.FLOAT,
input_shape),
- helper.make_tensor_value_info("w", TensorProto.FLOAT,
weight_shape),
- ]
- + ([helper.make_tensor_value_info("b", TensorProto.FLOAT,
bias_shape)] if bias else []),
- outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT,
output_shape)],
- )
+ + (
+ [helper.make_tensor_value_info("b", TensorProto.FLOAT,
[output_channels])]
+ if bias
+ else []
+ ),
+ outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT,
output_shape)],
+ )
+ model = helper.make_model(
+ graph,
+ producer_name="conv_transpose_test",
+ opset_imports=[helper.make_opsetid("", 14)],
+ )
+ return model, output_shape
- model = helper.make_model(graph, producer_name="conv_transpose_test")
- check_correctness(model, atol=1e-4)
-
- # ConvTranspose1D
- _verify_conv_transpose([3, 4, 32], [4, 4, 3])
- _verify_conv_transpose([3, 4, 32], [4, 2, 3]) # group=2
- # ConvTranspose2D
- _verify_conv_transpose([3, 4, 32, 32], [4, 4, 3, 3])
- _verify_conv_transpose([3, 4, 32, 32], [4, 2, 3, 3]) # group=2
- # ConvTranspose3D
- _verify_conv_transpose([3, 4, 12, 12, 12], [4, 4, 3, 3, 3])
- _verify_conv_transpose([3, 4, 12, 12, 12], [4, 2, 3, 3, 3]) # group=2
-
-
[email protected]("auto_pad", ["SAME_UPPER", "SAME_LOWER", "VALID"])
[email protected]("stride", [1, 2])
-def test_conv_transpose_auto_pad(auto_pad: str, stride: int):
- def _verify(input_shape, weight_shape):
- nd = len(weight_shape) - 2
- conv_node = helper.make_node(
- "ConvTranspose",
- inputs=["x", "w"],
- outputs=["y"],
- kernel_shape=weight_shape[2:],
- strides=[stride] * nd,
- auto_pad=auto_pad,
- )
- graph = helper.make_graph(
- [conv_node],
- "conv_transpose_auto_pad_test",
- inputs=[
- helper.make_tensor_value_info("x", TensorProto.FLOAT,
input_shape),
- helper.make_tensor_value_info("w", TensorProto.FLOAT,
weight_shape),
- ],
- outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT,
None)],
- )
- model = helper.make_model(graph,
producer_name="conv_transpose_auto_pad_test")
- check_correctness(model, atol=1e-4)
- # ConvTranspose1D / 2D / 3D
- _verify([1, 1, 8], [1, 1, 3])
- _verify([1, 1, 8, 8], [1, 1, 3, 3])
- _verify([1, 1, 4, 4, 4], [1, 1, 3, 3, 3])
+CONV_TRANSPOSE_IMPORT_CONFIGS = [
+ ("NOTSET", 1, 1, 0, 0),
+ ("NOTSET", 2, 2, 2, 1),
+ ("SAME_UPPER", 1, 1, 0, 0),
+ ("SAME_UPPER", 2, 1, 0, 0),
+ ("SAME_LOWER", 2, 1, 0, 0),
+ ("VALID", 1, 2, 0, 0),
+ ("VALID", 2, 1, 0, 1),
+]
+
+
+def _verify_conv_transpose_import(nd, groups, auto_pad, stride, dilation, pad,
bias, output_pad):
+ model, output_shape = _make_conv_transpose_model(
+ nd, groups, auto_pad, stride, dilation, pad, bias, output_pad
+ )
+ tvm_model = from_onnx(model, opset=14, keep_params_in_input=True)
+ func = tvm_model["main"]
+
+ conv_op_name = f"relax.nn.conv{nd}d_transpose"
+ conv_calls = []
+
+ def visit(expr):
+ if (
+ isinstance(expr, relax.Call)
+ and isinstance(expr.op, tvm.ir.Op)
+ and expr.op.name == conv_op_name
+ ):
+ conv_calls.append(expr)
+
+ relax.analysis.post_order_visit(func.body, visit)
+ assert len(conv_calls) == 1
+ conv_call = conv_calls[0]
+ assert tuple(int(value) for value in func.ret_ty.shape.values) ==
tuple(output_shape)
+ assert tuple(int(value) for value in conv_call.attrs.strides) == (stride,)
* nd
+ assert tuple(int(value) for value in conv_call.attrs.dilation) ==
(dilation,) * nd
+ assert tuple(int(value) for value in conv_call.attrs.output_padding) ==
(output_pad,) * nd
+ assert int(conv_call.attrs.groups) == groups
+ assert ("relax.add" in collect_relax_call_ops(func)) == bias
+
+ if auto_pad == "NOTSET":
+ expected_padding = (pad,) * (nd * 2)
+ elif auto_pad == "VALID":
+ expected_padding = (0,) * (nd * 2)
+ else:
+ total_pad = max((3 - 1) * dilation + 1 + output_pad - stride, 0)
+ pad_begin = total_pad // 2 if auto_pad == "SAME_UPPER" else total_pad
- total_pad // 2
+ expected_padding = (pad_begin,) * nd + (total_pad - pad_begin,) * nd
+ assert tuple(int(value) for value in conv_call.attrs.padding) ==
expected_padding
-def test_pow():
- verify_binary("Pow", [32, 32], [32, 32], [32, 32])
[email protected]("bias", [True, False])
[email protected]("nd", [1, 2, 3])
[email protected]("groups", [1, 2])
+def test_conv_transpose_import(bias, nd, groups):
+ for auto_pad, stride, dilation, pad, output_pad in
CONV_TRANSPOSE_IMPORT_CONFIGS:
+ _verify_conv_transpose_import(nd, groups, auto_pad, stride, dilation,
pad, bias, output_pad)
+
+
[email protected](
+ "nd, groups, auto_pad, stride, dilation, pad, bias, output_pad",
+ [
+ # Broad attribute coverage, including dilation and output padding.
+ (1, 1, "NOTSET", 2, 2, 2, False, 1),
+ (2, 2, "VALID", 1, 2, 0, True, 0),
+ (3, 1, "NOTSET", 1, 1, 0, True, 0),
+ (3, 2, "VALID", 2, 1, 0, False, 1),
+ # Each rank uses a distinct Relax op and legalizer. Exercise both
+ # directions of asymmetric SAME padding numerically for every rank.
+ (1, 2, "SAME_UPPER", 2, 1, 0, True, 0),
+ (1, 1, "SAME_LOWER", 2, 1, 0, False, 0),
+ (2, 2, "SAME_UPPER", 2, 1, 0, True, 0),
+ (2, 1, "SAME_LOWER", 2, 1, 0, False, 0),
+ (3, 1, "SAME_UPPER", 2, 1, 0, False, 0),
+ (3, 2, "SAME_LOWER", 2, 1, 0, True, 0),
+ # Preserve the 2-D output_padding regression through LLVM execution.
+ (2, 1, "NOTSET", 2, 1, 2, True, 1),
+ ],
+)
+def test_conv_transpose_numerical(nd, groups, auto_pad, stride, dilation, pad,
bias, output_pad):
+ spatial_extent = 4 if nd == 3 else 8
+ model, _ = _make_conv_transpose_model(
+ nd,
+ groups,
+ auto_pad,
+ stride,
+ dilation,
+ pad,
+ bias,
+ output_pad,
+ spatial_extent,
+ )
+ check_correctness(model, opset=14, atol=1e-4)
@pytest.mark.parametrize("reverse", [True, False])
@@ -3874,10 +3938,8 @@ def test_squeeze_constant():
verify_squeeze_constant(None, ExpectedSqueezeConstantAll)
[email protected]("axis", [[0]])
[email protected]("A", [8, 16, 32])
[email protected]("B", [8, 16, 32])
-def test_dynamic_squeeze(axis, A, B):
+def test_dynamic_squeeze():
+ axis = [0]
squeeze_node = helper.make_node("Squeeze", ["x", "axes"], ["y"])
shape = [1, "A", "B"]
@@ -7525,29 +7587,14 @@ def test_tile_dynamic_repeats():
def _generate_roi_cases():
- # Base case when with_roi is False
- roi_list = [
+ return [
pytest.param(False, None, False, id="no_roi"),
+ pytest.param(True, [], True, id="empty_roi_constant"),
+ pytest.param(True, [], False, id="empty_roi_initializer"),
+ pytest.param(True, [0.1, 0.2, 0.9, 0.8], True,
id="spatial_roi_constant"),
+ pytest.param(True, [0.1, 0.2, 0.9, 0.8], False,
id="spatial_roi_initializer"),
]
- # Valid when with_roi is True and with_constant is True/False
- roi_cases = [
- [],
- [0.0, 0.0, 0.0, 0.0],
- [0.0, 0.0, 1.0, 1.0],
- [0.1, 0.1, 0.9, 0.9],
- [0.2, 0.2, 0.8, 0.8],
- [0.3, 0.3, 0.7, 0.7],
- [0.4, 0.4, 0.6, 0.6],
- [0.5, 0.5, 0.5, 0.5],
- [0.1, 0.2, 0.9, 0.8],
- ]
- for roi in roi_cases:
- roi_list.append(pytest.param(True, roi, True,
id=f"roi_{'_'.join(str(x) for x in roi)}"))
- roi_list.append(pytest.param(True, roi, False,
id=f"roi_{'_'.join(str(x) for x in roi)}"))
-
- return roi_list
-
@pytest.mark.parametrize("with_roi, roi_list, with_constant",
_generate_roi_cases())
def test_resize(with_roi, roi_list, with_constant):
diff --git a/tests/python/relax/test_pytorch_integration.py
b/tests/python/relax/test_pytorch_integration.py
index b643b3e147..2be8fe3f60 100644
--- a/tests/python/relax/test_pytorch_integration.py
+++ b/tests/python/relax/test_pytorch_integration.py
@@ -14,7 +14,6 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-# ruff: noqa: F401
"""
Test PyTorch integration with TVM Relax.
@@ -23,16 +22,14 @@ This test verifies:
2. Cross-function calls between Python, TIR, and Relax functions
3. Dynamic Python function addition and execution
4. End-to-end pipeline testing
-5. Error handling and edge cases
+5. Missing packed-function error handling
"""
-import numpy as np
import pytest
import torch
import torch.nn.functional as F
import tvm
-from tvm import relax, tirx
from tvm.relax import BasePyModule
from tvm.script import ir as I
from tvm.script import relax as R
@@ -87,305 +84,41 @@ class PyTorchIntegrationModule(BasePyModule):
class TestPyTorchIntegration:
- def test_module_creation_and_instantiation(self):
- module = PyTorchIntegrationModule
-
- assert hasattr(module, "__call__"), "Module should be callable"
-
- device = tvm.cpu(0)
- instance = module(device)
-
- assert isinstance(instance, BasePyModule), "Instance should be
BasePyModule"
-
- required_methods = ["main", "call_tir", "call_dps_packed"]
- for method in required_methods:
- assert hasattr(instance, method), f"Instance should have method:
{method}"
-
- def test_module_creation_and_instantiation_gpu(self):
- module = PyTorchIntegrationModule
-
- if tvm.cuda().exist:
-
- def run_and_check():
- assert hasattr(module, "__call__"), "Module should be callable"
-
- device = tvm.cuda(0)
- instance = module(device)
-
- assert isinstance(instance, BasePyModule), "Instance should be
BasePyModule"
- required_methods = ["main", "call_tir", "call_dps_packed"]
- for method in required_methods:
- assert hasattr(instance, method), f"Instance should have
method: {method}"
- assert "cuda" in str(instance.target)
-
- tvm.testing.run_with_gpu_lock(run_and_check)
- else:
- pytest.skip("CUDA not available")
-
- def test_python_function_execution(self):
- """Test that Python functions execute correctly."""
- module = PyTorchIntegrationModule
- device = tvm.cpu(0)
- instance = module(device)
-
- # Test my_identity_func
- input_tensor = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32)
- result = instance.my_identity_func(input_tensor)
-
- assert isinstance(result, torch.Tensor)
- assert torch.allclose(result, input_tensor, atol=1e-5)
-
- def test_tir_function_execution(self):
- """Test that TIR functions execute correctly."""
- module = PyTorchIntegrationModule
- device = tvm.cpu(0)
- instance = module(device)
-
- # Test matmul function
- n = 3
- x = torch.randn(n, 16, dtype=torch.float32)
- w = torch.randn(16, 20, dtype=torch.float32)
-
- result = instance.call_tir(instance.matmul, [x, w], R.Tensor((n, 20),
"float32"))
-
- assert isinstance(result, torch.Tensor)
- assert result.shape == (n, 20)
-
- # Verify result with PyTorch matmul
- expected = torch.matmul(x, w)
- assert torch.allclose(result, expected, atol=1e-3)
-
- def test_dynamic_python_function_addition(self):
- """Test adding Python functions dynamically."""
- module = PyTorchIntegrationModule
- device = tvm.cpu(0)
- instance = module(device)
-
- # Define a custom function
- def custom_activation(x):
- return torch.sigmoid(x)
-
- # Add the function
- instance.add_python_function("custom_activation", custom_activation)
-
- # Verify function is added
- assert hasattr(instance, "custom_activation")
- assert "custom_activation" in instance.pyfuncs
-
- # Test function execution
- input_tensor = torch.tensor([1.0, -1.0, 0.0], dtype=torch.float32)
- result = instance.custom_activation(input_tensor)
-
- assert isinstance(result, torch.Tensor)
- expected = torch.sigmoid(input_tensor)
- assert torch.allclose(result, expected, atol=1e-5)
-
- def test_call_dps_packed_with_dynamic_function(self):
- """Test call_dps_packed with dynamically added function."""
- module = PyTorchIntegrationModule
- device = tvm.cpu(0)
- instance = module(device)
-
- # Define my_softmax function
- def my_softmax(tensor, dim):
- """Custom softmax function for testing call_dps_packed."""
- # Convert TVM Tensor to PyTorch tensor if needed
- if hasattr(tensor, "numpy"):
- tensor = torch.from_numpy(tensor.numpy())
- return F.softmax(tensor, dim=dim)
-
- # Add the function
- instance.my_softmax = my_softmax
-
- # Test call_dps_packed
- input_tensor = torch.tensor([[1.0, 2.0], [3.0, 4.0]],
dtype=torch.float32)
-
- result = instance.call_dps_packed(
- "my_softmax", [input_tensor, 1], R.Tensor((2, 2), "float32")
- )
-
- assert isinstance(result, torch.Tensor)
- expected = F.softmax(input_tensor, dim=1)
- assert torch.allclose(result, expected, atol=1e-5)
-
def test_end_to_end_pipeline(self):
- module = PyTorchIntegrationModule
- device = tvm.cpu(0)
- instance = module(device)
-
- def my_softmax(tensor, dim):
- if hasattr(tensor, "numpy"):
- tensor = torch.from_numpy(tensor.numpy())
- return F.softmax(tensor, dim=dim)
-
- instance.my_softmax = my_softmax
-
- n = 5
- x = torch.randn(n, 16, dtype=torch.float32)
- w = torch.randn(16, 20, dtype=torch.float32)
-
- result = instance.main(x, w)
+ instance = PyTorchIntegrationModule(tvm.cpu(0))
- assert isinstance(result, torch.Tensor)
- assert result.shape == (n, 20)
- assert result.dtype == torch.float32
-
- def test_end_to_end_pipeline_gpu(self):
- module = PyTorchIntegrationModule
-
- if tvm.cuda().exist:
-
- def run_and_check():
- device = tvm.cuda(0)
- instance = module(device)
-
- # Test basic GPU functionality without complex TIR operations
- assert isinstance(instance, BasePyModule)
- assert "cuda" in str(instance.target)
-
- # Test that we can create and work with GPU tensors
- n = 5
- x = torch.randn(n, 16, dtype=torch.float32, device="cuda")
- w = torch.randn(16, 20, dtype=torch.float32, device="cuda")
-
- assert x.device.type == "cuda"
- assert w.device.type == "cuda"
- assert x.shape == (n, 16)
- assert w.shape == (16, 20)
-
- # Test basic PyTorch operations on GPU
- result = torch.matmul(x, w)
- assert isinstance(result, torch.Tensor)
- assert result.shape == (n, 20)
- assert result.dtype == torch.float32
- assert result.device.type == "cuda"
-
- tvm.testing.run_with_gpu_lock(run_and_check)
- else:
- pytest.skip("CUDA not available")
-
- def test_cross_function_data_flow(self):
- """Test data flow between different function types."""
- module = PyTorchIntegrationModule
- device = tvm.cpu(0)
- instance = module(device)
-
- # Add required functions
def my_softmax(tensor, dim):
- if hasattr(tensor, "numpy"):
- tensor = torch.from_numpy(tensor.numpy())
return F.softmax(tensor, dim=dim)
- instance.my_softmax = my_softmax
+ instance.add_python_function("my_softmax", my_softmax)
+ assert "my_softmax" in instance.pyfuncs
- # Create test data
- n = 4
- x = torch.randn(n, 16, dtype=torch.float32)
+ torch.manual_seed(0)
w = torch.randn(16, 20, dtype=torch.float32)
- # Execute step by step to verify data flow
- # Step 1: TIR matmul
- lv = instance.call_tir(instance.matmul, [x, w], R.Tensor((n, 20),
"float32"))
- assert isinstance(lv, torch.Tensor)
- assert lv.shape == (n, 20)
-
- # Step 2: ReLU
- lv1 = F.relu(lv)
- assert isinstance(lv1, torch.Tensor)
- assert lv1.shape == (n, 20)
-
- # Step 3: Softmax via call_dps_packed
- lv2 = instance.call_dps_packed("my_softmax", [lv1, 1], R.Tensor((n,
20), "float32"))
- assert isinstance(lv2, torch.Tensor)
- assert lv2.shape == (n, 20)
-
- # Step 4: Identity function
- lv3 = instance.my_identity_func(lv2)
- assert isinstance(lv3, torch.Tensor)
- assert lv3.shape == (n, 20)
-
- # Verify final result matches expected
- expected = F.softmax(F.relu(torch.matmul(x, w)), dim=1)
- assert torch.allclose(lv3, expected, atol=1e-3)
-
- def test_error_handling(self):
- """Test error handling for various edge cases."""
- module = PyTorchIntegrationModule
- device = tvm.cpu(0)
- instance = module(device)
-
- # Test with missing function
- with pytest.raises(Exception):
- instance.call_dps_packed(
- "non_existent_function", [torch.tensor([1.0])], R.Tensor((1,),
"float32")
- )
-
- # Test with wrong tensor shapes
- x = torch.randn(3, 16, dtype=torch.float32)
- w = torch.randn(15, 20, dtype=torch.float32) # Wrong shape
-
- with pytest.raises(Exception):
- instance.call_tir(instance.matmul, [x, w], R.Tensor((3, 20),
"float32"))
-
- def test_tensor_type_preservation(self):
- module = PyTorchIntegrationModule
- device = tvm.cpu(0)
- instance = module(device)
-
- def my_softmax(tensor, dim):
- if hasattr(tensor, "numpy"):
- tensor = torch.from_numpy(tensor.numpy())
- return F.softmax(tensor, dim=dim)
-
- instance.my_softmax = my_softmax
-
- # Test with float32 data type (TIR function is hardcoded for float32)
- test_dtype = torch.float32
- n = 3
- x = torch.randn(n, 16, dtype=test_dtype)
- w = torch.randn(16, 20, dtype=test_dtype)
-
- result = instance.main(x, w)
-
- # Verify type preservation
- assert result.dtype == test_dtype
- assert isinstance(result, torch.Tensor)
- assert result.shape == (n, 20)
- assert result.dtype == torch.float32
-
- def test_batch_processing(self):
- """Test processing multiple inputs in batch."""
- module = PyTorchIntegrationModule
- device = tvm.cpu(0)
- instance = module(device)
-
- # Add required functions
- def my_softmax(tensor, dim):
- if hasattr(tensor, "numpy"):
- tensor = torch.from_numpy(tensor.numpy())
- return F.softmax(tensor, dim=dim)
-
- instance.my_softmax = my_softmax
-
- # Process multiple inputs
- batch_size = 5
- results = []
-
- for i in range(batch_size):
- n = 3 + i # Varying batch sizes
+ # Reuse the compiled module with two different symbolic extents.
+ for n in (3, 5):
x = torch.randn(n, 16, dtype=torch.float32)
- w = torch.randn(16, 20, dtype=torch.float32)
-
result = instance.main(x, w)
- results.append(result)
+ expected = F.softmax(F.relu(torch.matmul(x, w)), dim=1)
assert isinstance(result, torch.Tensor)
assert result.shape == (n, 20)
+ assert result.dtype == torch.float32
+ torch.testing.assert_close(result, expected, rtol=1e-5, atol=1e-3)
- # Verify all results are valid
- assert len(results) == batch_size
- for result in results:
- assert isinstance(result, torch.Tensor)
+ def test_missing_packed_function(self):
+ instance = BasePyModule(tvm.IRModule({}), tvm.cpu(0))
+
+ with pytest.raises(
+ ValueError,
+ match="Function 'non_existent_function' not found as a global
function",
+ ):
+ instance.call_dps_packed(
+ "non_existent_function",
+ [torch.tensor([1.0])],
+ R.Tensor((1,), "float32"),
+ )
if __name__ == "__main__":
diff --git a/tests/python/relax/test_tvmscript_pyfunc.py
b/tests/python/relax/test_tvmscript_pyfunc.py
index 950c0ccea4..8e3bb6211c 100644
--- a/tests/python/relax/test_tvmscript_pyfunc.py
+++ b/tests/python/relax/test_tvmscript_pyfunc.py
@@ -209,6 +209,7 @@ class TestTVMScriptPyFunc:
assert isinstance(instance, BasePyModule), "Instance should be
BasePyModule"
assert hasattr(instance, "pyfuncs"), "Instance should have
pyfuncs"
+ assert "cuda" in str(instance.target)
x = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32,
device="cuda")
result = instance.pytorch_processor(x)
