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 1a764d7993 [Tests] Reduce runtime of slow Python tests (#20006)
1a764d7993 is described below
commit 1a764d799316c7d73123750a0a3f7ce40c7d5cba
Author: Shushi Hong <[email protected]>
AuthorDate: Wed Jul 15 18:07:20 2026 -0400
[Tests] Reduce runtime of slow Python tests (#20006)
This PR reduces the runtime of several slow Python test groups:
- Parameterize LLVM division and CUDA vectorized-cast cases so
pytest-xdist can schedule them independently.
- Replace exhaustive ONNX execution with structural importer checks plus
representative numerical cases, and avoid registering unsupported
backend cases.
- Reuse compiled paged-attention kernels across compatible test cases.
Targeted measurements showed:
- LLVM division: 25.34s → 19.87s
- CUDA vectorized casts: 142.85s → 108.40s
- Paged-attention CPU: 316.32s → 192.50s
- ONNX Conv: 25.84s → 1.81s
- ONNX Reduce: 20.66s → 4.04s
This PR also fixes a latent CUDA Graph cleanup bug that could leave
`cudaErrorStreamCaptureInvalidated` in the worker thread and cause
unrelated subsequent GPU tests to fail.
---
src/runtime/vm/cuda/cuda_graph_builtin.cc | 9 +-
tests/python/codegen/test_target_codegen_cuda.py | 143 +++---
tests/python/codegen/test_target_codegen_llvm.py | 107 +++--
tests/python/relax/test_frontend_onnx.py | 478 +++++++++++----------
tests/python/relax/test_frontend_onnx_backend.py | 41 +-
...runtime_builtin_paged_attention_kv_cache_cpu.py | 75 ++--
...runtime_builtin_paged_attention_kv_cache_tir.py | 83 ++--
tests/python/relax/test_vm_cuda_graph.py | 16 +
8 files changed, 563 insertions(+), 389 deletions(-)
diff --git a/src/runtime/vm/cuda/cuda_graph_builtin.cc
b/src/runtime/vm/cuda/cuda_graph_builtin.cc
index b51b287d81..75876dba97 100644
--- a/src/runtime/vm/cuda/cuda_graph_builtin.cc
+++ b/src/runtime/vm/cuda/cuda_graph_builtin.cc
@@ -123,7 +123,14 @@ class CUDACaptureStream {
TVM_FFI_CHECK_CUDA_ERROR(cudaStreamBeginCapture(capture_stream_,
cudaStreamCaptureModeGlobal));
}
~CUDACaptureStream() noexcept(false) {
- cudaStreamEndCapture(capture_stream_, output_graph_);
+ cudaError_t capture_error = cudaStreamEndCapture(capture_stream_,
output_graph_);
+ if (capture_error != cudaSuccess) {
+ // The capture may have been invalidated by the exception that is
+ // currently unwinding the stack. Do not throw a second exception from
+ // this destructor, but clear CUDA's thread-local error so that it is not
+ // reported by an unrelated CUDA call later in the same host thread.
+ cudaGetLastError();
+ }
TVM_FFI_CHECK_SAFE_CALL(TVMFFIEnvSetStream(kDLCUDA, device_id_,
prev_default_stream_, nullptr));
}
diff --git a/tests/python/codegen/test_target_codegen_cuda.py
b/tests/python/codegen/test_target_codegen_cuda.py
index a820cf5f67..8bb32f68a4 100644
--- a/tests/python/codegen/test_target_codegen_cuda.py
+++ b/tests/python/codegen/test_target_codegen_cuda.py
@@ -607,76 +607,93 @@ def test_cuda_floormod_with_vectorization():
tvm.testing.run_with_gpu_lock(run_and_check)
+_VECTORIZED_CAST_TYPES_4 = [
+ "float16",
+ "float32",
+ "int8",
+ "uint8",
+ "int16",
+ "uint16",
+ "int32",
+ "uint32",
+ "float64",
+ "int64",
+ "uint64",
+]
+_VECTORIZED_CAST_TYPES_8 = [
+ "float16",
+ "float32",
+ "int8",
+ "uint8",
+ "int16",
+ "uint16",
+ "int32",
+ "uint32",
+]
+
+
+def _skip_vectorized_cast(t0, t1):
+ if t0 == t1:
+ return True
+ # CUDA does support cast between {u}int8 and fp16.
+ skip_set = {"float16", "uint8", "int8"}
+ return t0 in skip_set and t1 in skip_set
+
+
+_VECTORIZED_CAST_CASES = [
+ (t0, t1, 4)
+ for t0 in _VECTORIZED_CAST_TYPES_4
+ for t1 in _VECTORIZED_CAST_TYPES_4
+ if not _skip_vectorized_cast(t0, t1)
+] + [
+ (t0, t1, 8)
+ for t0 in _VECTORIZED_CAST_TYPES_8
+ for t1 in _VECTORIZED_CAST_TYPES_8
+ if not _skip_vectorized_cast(t0, t1)
+]
+_VECTORIZED_CAST_CASES += [("int8", "uint8", 16), ("uint8", "int8", 16)]
+
+
@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
-def test_vectorized_casts():
- def check(t0, t1, factor):
- if (t0 == "float16" or t1 == "float16") and not
have_fp16(tvm.cuda(0).compute_version):
- print("Skip because gpu does not have fp16 support")
- return
[email protected]("t0,t1,factor", _VECTORIZED_CAST_CASES)
+def test_vectorized_casts(t0, t1, factor):
+ if (t0 == "float16" or t1 == "float16") and not
have_fp16(tvm.cuda(0).compute_version):
+ print("Skip because gpu does not have fp16 support")
+ return
- n = 128
- num_thread = n // factor
+ n = 128
+ num_thread = n // factor
- @I.ir_module(s_tir=True)
- class Module:
- @T.prim_func(s_tir=True)
- def main(A: T.Buffer((n,), t0), B: T.Buffer((n,), t1), C:
T.Buffer((n,), t0)):
- T.func_attr({"tirx.noalias": True})
- for i_0 in T.thread_binding(num_thread, thread="threadIdx.x"):
- for i_1 in T.vectorized(factor):
- with T.sblock("C"):
- v_i = T.axis.spatial(n, i_0 * factor + i_1)
- T.reads(A[v_i], B[v_i])
- T.writes(C[v_i])
- C[v_i] = A[v_i] + T.Cast(t0, B[v_i])
-
- func = tvm.compile(Module, target="cuda")
+ @I.ir_module(s_tir=True)
+ class Module:
+ @T.prim_func(s_tir=True)
+ def main(A: T.Buffer((n,), t0), B: T.Buffer((n,), t1), C:
T.Buffer((n,), t0)):
+ T.func_attr({"tirx.noalias": True})
+ for i_0 in T.thread_binding(num_thread, thread="threadIdx.x"):
+ for i_1 in T.vectorized(factor):
+ with T.sblock("C"):
+ v_i = T.axis.spatial(n, i_0 * factor + i_1)
+ T.reads(A[v_i], B[v_i])
+ T.writes(C[v_i])
+ C[v_i] = A[v_i] + T.Cast(t0, B[v_i])
- # correctness
- def run_and_check():
- dev = tvm.cuda(0)
- low, high = (0, 20) if t0.startswith("u") or t1.startswith("u")
else (-10, 10)
- a_np = np.random.randint(low, high, size=n).astype(t0)
- b_np = np.random.randint(low, high, size=n).astype(t1)
- c_np = (a_np + b_np).astype(t0)
- a_nd = tvm.runtime.tensor(a_np, dev)
- b_nd = tvm.runtime.tensor(b_np, dev)
- c_nd = tvm.runtime.tensor(np.zeros(c_np.shape, dtype=c_np.dtype),
dev)
- func(a_nd, b_nd, c_nd)
- tvm.testing.assert_allclose(c_nd.numpy(), c_np, rtol=1e-3)
+ func = tvm.compile(Module, target="cuda")
- tvm.testing.run_with_gpu_lock(run_and_check)
+ # correctness
+ def run_and_check():
+ dev = tvm.cuda(0)
+ low, high = (0, 20) if t0.startswith("u") or t1.startswith("u") else
(-10, 10)
+ a_np = np.random.randint(low, high, size=n).astype(t0)
+ b_np = np.random.randint(low, high, size=n).astype(t1)
+ c_np = (a_np + b_np).astype(t0)
+ a_nd = tvm.runtime.tensor(a_np, dev)
+ b_nd = tvm.runtime.tensor(b_np, dev)
+ c_nd = tvm.runtime.tensor(np.zeros(c_np.shape, dtype=c_np.dtype), dev)
+ func(a_nd, b_nd, c_nd)
+ tvm.testing.assert_allclose(c_nd.numpy(), c_np, rtol=1e-3)
- def skip(t0, t1):
- if t0 == t1:
- return True
- # CUDA does support cast between {u}int8 and fp16.
- skip_set = {"float16", "uint8", "int8"}
- if t0 in skip_set and t1 in skip_set:
- return True
- return False
-
- types_4 = [
- "float16",
- "float32",
- "int8",
- "uint8",
- "int16",
- "uint16",
- "int32",
- "uint32",
- "float64",
- "int64",
- "uint64",
- ]
- types_8 = ["float16", "float32", "int8", "uint8", "int16", "uint16",
"int32", "uint32"]
- for t0, t1 in [(x, y) for x in types_4 for y in types_4 if not skip(x, y)]:
- check(t0, t1, 4)
- for t0, t1 in [(x, y) for x in types_8 for y in types_8 if not skip(x, y)]:
- check(t0, t1, 8)
- check("int8", "uint8", 16)
- check("uint8", "int8", 16)
+ tvm.testing.run_with_gpu_lock(run_and_check)
def sched(compute_fn, dtype, n=128):
diff --git a/tests/python/codegen/test_target_codegen_llvm.py
b/tests/python/codegen/test_target_codegen_llvm.py
index 4f67f1803c..9de6b73bb1 100644
--- a/tests/python/codegen/test_target_codegen_llvm.py
+++ b/tests/python/codegen/test_target_codegen_llvm.py
@@ -558,8 +558,74 @@ def test_alignment():
assert has_call_to_assume()
+def _llvm_div_cases():
+ dividend_ranges = [
+ (-12, -12),
+ (-11, -1),
+ (-11, 0),
+ (0, 0),
+ (12, 12),
+ (1, 11),
+ (0, 11),
+ (-11, 11),
+ ]
+ divisor_ranges = [
+ (-11, -1),
+ (-11, 1),
+ (-4, -4),
+ (-2, -2),
+ (1, 11),
+ (0, 11),
+ (4, 4),
+ (2, 2),
+ (-11, 11),
+ ]
+
+ for start, end in dividend_ranges:
+ for dstart, dend in divisor_ranges:
+ for dtype in ["int32", "int8"]:
+ for floor_div in [False, True]:
+ yield pytest.param(
+ start,
+ end,
+ dstart,
+ dend,
+ dtype,
+ floor_div,
+ id=f"{dtype}-{'floor' if floor_div else
'trunc'}-{start}:{end}-{dstart}:{dend}",
+ )
+ if start >= 0 and dstart >= 0:
+ for floor_div in [False, True]:
+ yield pytest.param(
+ start,
+ end,
+ dstart,
+ dend,
+ "uint32",
+ floor_div,
+ id=f"uint32-{'floor' if floor_div else
'trunc'}-{start}:{end}-{dstart}:{dend}",
+ )
+
+ for dstart, dend in [(0, 11), (1, 11), (2, 2), (4, 4)]:
+ for start, end in [(123, 133), (0, 255)]:
+ for floor_div in [False, True]:
+ yield pytest.param(
+ start,
+ end,
+ dstart,
+ dend,
+ "uint8",
+ floor_div,
+ id=f"uint8-{'floor' if floor_div else
'trunc'}-{start}:{end}-{dstart}:{dend}",
+ )
+
+
@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
-def test_llvm_div():
[email protected](
+ "start,end,dstart,dend,dtype,floor_div",
+ list(_llvm_div_cases()),
+)
+def test_llvm_div(start, end, dstart, dend, dtype, floor_div):
"""Check that the semantics of div and mod is correct"""
def check(start, end, dstart, dend, dtype, floor_div=False):
@@ -658,44 +724,7 @@ def test_llvm_div():
f"but should be {mref}"
)
- # Try different ranges to cover different cases
- for start, end in [
- (-12, -12),
- (-11, -1),
- (-11, 0),
- (0, 0),
- (12, 12),
- (1, 11),
- (0, 11),
- (-11, 11),
- ]:
- for dstart, dend in [
- (-11, -1),
- (-11, 1),
- (-4, -4),
- (-2, -2),
- (1, 11),
- (0, 11),
- (4, 4),
- (2, 2),
- (-11, 11),
- ]:
- if end < start or dend < dstart or (dend == 0 and dstart == 0) or
dend == 0:
- continue
- check(start, end, dstart, dend, "int32", floor_div=False)
- check(start, end, dstart, dend, "int32", floor_div=True)
- check(start, end, dstart, dend, "int8", floor_div=False)
- check(start, end, dstart, dend, "int8", floor_div=True)
- if start >= 0 and dstart >= 0:
- check(start, end, dstart, dend, "uint32", floor_div=False)
- check(start, end, dstart, dend, "uint32", floor_div=True)
-
- # Additional tests for uint8
- for dstart, dend in [(0, 11), (1, 11), (2, 2), (4, 4)]:
- check(123, 133, dstart, dend, "uint8", floor_div=False)
- check(123, 133, dstart, dend, "uint8", floor_div=True)
- check(0, 255, dstart, dend, "uint8", floor_div=False)
- check(0, 255, dstart, dend, "uint8", floor_div=True)
+ check(start, end, dstart, dend, dtype, floor_div)
@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
diff --git a/tests/python/relax/test_frontend_onnx.py
b/tests/python/relax/test_frontend_onnx.py
index 84c3260bf6..7269594104 100644
--- a/tests/python/relax/test_frontend_onnx.py
+++ b/tests/python/relax/test_frontend_onnx.py
@@ -3494,84 +3494,129 @@ def test_shrink():
tvm.ir.assert_structural_equal(tvm_model, ExpectedCustom)
[email protected]("stride", [1, 2])
[email protected]("dilation", [1, 2])
+def _make_conv_model(input_shape, weight_shape, stride, dilation, pad, bias,
auto_pad):
+ nd = len(weight_shape) - 2
+ groups = input_shape[1] // weight_shape[1]
+ node_attrs = {
+ "strides": [stride] * nd,
+ "dilations": [dilation] * nd,
+ "group": groups,
+ }
+ if auto_pad == "VALID":
+ output_shape = [input_shape[0], weight_shape[0]] + [
+ (input_shape[i] - dilation * (weight_shape[i] - 1) - 1) // stride
+ 1
+ for i in range(2, len(input_shape))
+ ]
+ node_attrs["auto_pad"] = auto_pad
+ elif auto_pad in ("SAME_UPPER", "SAME_LOWER"):
+ output_shape = [input_shape[0], weight_shape[0]] + [
+ (input_shape[i] + stride - 1) // stride for i in range(2,
len(input_shape))
+ ]
+ node_attrs["auto_pad"] = auto_pad
+ else:
+ output_shape = [input_shape[0], weight_shape[0]] + [
+ (input_shape[i] + 2 * pad - dilation * (weight_shape[i] - 1) - 1)
// stride + 1
+ for i in range(2, len(input_shape))
+ ]
+ node_attrs["pads"] = [pad] * nd * 2
+
+ conv_node = helper.make_node(
+ "Conv",
+ inputs=["x", "w"] + (["b"] if bias else []),
+ outputs=["y"],
+ **node_attrs,
+ )
+ graph = helper.make_graph(
+ [conv_node],
+ "conv_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,
[output_shape[1]])]
+ if bias
+ else []
+ ),
+ outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT,
output_shape)],
+ )
+ model = helper.make_model(
+ graph,
+ producer_name="conv_test",
+ opset_imports=[helper.make_opsetid("", 14)],
+ )
+ return model, output_shape, groups
+
+
+CONV_IMPORT_CONFIGS = [
+ *(("VALID", stride, dilation, 0) for stride in [1, 2] for dilation in [1,
2]),
+ *((auto_pad, stride, 1, 0) for auto_pad in ["SAME_UPPER", "SAME_LOWER"]
for stride in [1, 2]),
+ *(
+ ("NOTSET", stride, dilation, pad)
+ for stride in [1, 2]
+ for dilation in [1, 2]
+ for pad in [0, 2]
+ ),
+]
+
+
+def _verify_conv_import(auto_pad, stride, dilation, pad, bias, nd, groups):
+ input_shape = [1, 4] + [8] * nd
+ weight_shape = [4, 4 // groups] + [3] * nd
+ model, output_shape, expected_groups = _make_conv_model(
+ input_shape, weight_shape, stride, dilation, pad, bias, auto_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"
+ 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 int(conv_call.attrs.groups) == expected_groups
+ assert ("relax.add" in collect_relax_call_ops(func)) == bias
+
+ expected_padding = (pad,) * (nd * 2) if auto_pad == "NOTSET" else (0,) *
(nd * 2)
+ assert tuple(int(value) for value in conv_call.attrs.padding) ==
expected_padding
+
+
@pytest.mark.parametrize("bias", [True, False])
[email protected]("pad", [0, 2])
[email protected]("auto_pad", ["SAME_UPPER", "SAME_LOWER", "VALID"])
-def test_conv(stride: int, dilation: int, pad: int, bias: bool, auto_pad: str):
- def _verify_conv(input_shape, weight_shape):
- nd = len(weight_shape) - 2
- if auto_pad == "VALID":
- output_shape = [input_shape[0], weight_shape[0]] + [
- (input_shape[i] - dilation * (weight_shape[i] - 1) - 1) //
stride + 1
- for i in range(2, len(input_shape))
- ]
- bias_shape = [output_shape[1]]
- conv_node = helper.make_node(
- "Conv",
- inputs=["x", "w"] + (["b"] if bias else []),
- outputs=["y"],
- strides=[stride] * nd,
- dilations=[dilation] * nd,
- auto_pad=auto_pad,
- group=input_shape[1] // weight_shape[1],
- )
- elif auto_pad in ("SAME_UPPER", "SAME_LOWER"):
- if dilation == 2:
- # auto_pad = "SAME" and dilation = 2 is not supported in ONNX
- return
- output_shape = [input_shape[0], weight_shape[0]] + [
- (input_shape[i] + stride - 1) // stride for i in range(2,
len(input_shape))
- ]
- bias_shape = [output_shape[1]]
- conv_node = helper.make_node(
- "Conv",
- inputs=["x", "w"] + (["b"] if bias else []),
- outputs=["y"],
- strides=[stride] * nd,
- dilations=[dilation] * nd,
- auto_pad=auto_pad,
- group=input_shape[1] // weight_shape[1],
- )
- else:
- output_shape = [input_shape[0], weight_shape[0]] + [
- (input_shape[i] + 2 * pad - dilation * (weight_shape[i] - 1) -
1) // stride + 1
- for i in range(2, len(input_shape))
- ]
- bias_shape = [output_shape[1]]
- conv_node = helper.make_node(
- "Conv",
- inputs=["x", "w"] + (["b"] if bias else []),
- outputs=["y"],
- strides=[stride] * nd,
- dilations=[dilation] * nd,
- pads=[pad] * nd * 2,
- group=input_shape[1] // weight_shape[1],
- )
- graph = helper.make_graph(
- [conv_node],
- "conv_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)],
- )
[email protected]("nd", [1, 2, 3])
[email protected]("groups", [1, 2])
+def test_conv_import(bias, nd, groups):
+ for auto_pad, stride, dilation, pad in CONV_IMPORT_CONFIGS:
+ _verify_conv_import(auto_pad, stride, dilation, pad, bias, nd, groups)
- model = helper.make_model(graph, producer_name="conv_test")
- check_correctness(model, atol=1e-4)
- # Conv1D
- _verify_conv([3, 4, 32], [4, 4, 3])
- _verify_conv([3, 4, 32], [2, 4, 3]) # group=2
- # Conv2D
- _verify_conv([3, 4, 32, 32], [4, 4, 3, 3])
- _verify_conv([3, 4, 32, 32], [2, 4, 3, 3]) # group=2
- # Conv3D
- _verify_conv([3, 4, 32, 32, 32], [4, 4, 3, 3, 3])
- _verify_conv([3, 4, 32, 32, 32], [2, 4, 3, 3, 3]) # group=2
[email protected](
+ "nd, groups, auto_pad, stride, dilation, pad, bias",
+ [
+ (1, 1, "VALID", 1, 2, 0, False),
+ (1, 2, "NOTSET", 2, 1, 2, True),
+ (2, 1, "SAME_UPPER", 2, 1, 0, True),
+ (2, 2, "SAME_LOWER", 2, 1, 0, False),
+ (3, 2, "VALID", 2, 1, 0, True),
+ (3, 1, "NOTSET", 1, 2, 2, False),
+ ],
+)
+def test_conv_numerical(nd, groups, auto_pad, stride, dilation, pad, bias):
+ input_shape = [1, 4] + [8] * nd
+ weight_shape = [4, 4 // groups] + [3] * nd
+ model, _, _ = _make_conv_model(input_shape, weight_shape, stride,
dilation, pad, bias, auto_pad)
+ check_correctness(model, opset=14, atol=1e-4)
@pytest.mark.parametrize("stride", [2])
@@ -5144,64 +5189,97 @@ def create_composite_reduce_test_parameters_axes_attr():
return output
[email protected]("func, dynamic, opset",
create_reduce_test_parameters_axes_attr())
-def test_all_reduce_funcs_axes_attr(func, dynamic, opset):
- def verify_reduce_func(func, data, axis, keepdims):
- inshape = data.shape
- outshape = np.sum(data, axis=axis, keepdims=keepdims == 1).shape
+def _verify_reduce_numerical(
+ func: str,
+ opset: int,
+ *,
+ axes_as_input: bool,
+ axes,
+ noop_with_empty_axes: bool = False,
+ dynamic: bool = False,
+ keepdims: bool = False,
+):
+ input_shape = [3, 3, 3]
+ node_inputs = ["x"]
+ initializers = []
+ node_attrs = {"keepdims": keepdims}
- if axis:
- node = onnx.helper.make_node(
- func, inputs=["x"], outputs=["y"], axes=axis, keepdims=keepdims
+ if axes_as_input:
+ node_attrs["noop_with_empty_axes"] = noop_with_empty_axes
+ if axes is not None:
+ axes_np = np.asarray(axes, dtype=np.int64)
+ initializers.append(
+ helper.make_tensor(
+ name="reduce_axes",
+ data_type=TensorProto.INT64,
+ dims=axes_np.shape,
+ vals=axes_np,
+ )
)
- else:
- node = onnx.helper.make_node(func, inputs=["x"], outputs=["y"],
keepdims=keepdims)
-
- if dynamic:
- in_list = ["?" for _ in range(len(inshape))]
- out_list = ["?" for _ in range(len(outshape))]
- else:
- in_list = list(inshape)
- out_list = list(outshape)
- graph = helper.make_graph(
- [node],
- "reduce_test",
- inputs=[helper.make_tensor_value_info("x", TensorProto.FLOAT,
in_list)],
- outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT,
out_list)],
- )
+ node_inputs.append("reduce_axes")
+ elif axes:
+ node_attrs["axes"] = axes
- model = helper.make_model(graph, producer_name="reduce_test")
- inputs_dict = {"x": data}
- # Reduction ops accumulate arithmetic errors, so we use a higher
tolerance.
- check_correctness(model, inputs_dict, opset=opset, rtol=1e-4,
atol=1e-4)
+ if axes_as_input and noop_with_empty_axes and not axes:
+ output_shape = input_shape
+ else:
+ axis = None if not axes else tuple(axes)
+ output_shape = list(np.sum(np.empty(input_shape), axis=axis,
keepdims=keepdims).shape)
- for keepdims in [True, False]:
- verify_reduce_func(
- func, np.random.randn(3, 2, 2).astype(np.float32), axis=None,
keepdims=keepdims
- )
+ graph_input_shape = ["?"] * len(input_shape) if dynamic else input_shape
+ graph_output_shape = ["?"] * len(output_shape) if dynamic else output_shape
- verify_reduce_func(
- func, np.random.randn(3, 2, 3).astype(np.float32), axis=None,
keepdims=keepdims
- )
+ node = helper.make_node(func, inputs=node_inputs, outputs=["y"],
**node_attrs)
+ graph = helper.make_graph(
+ [node],
+ "reduce_numerical_test",
+ inputs=[helper.make_tensor_value_info("x", TensorProto.FLOAT,
graph_input_shape)],
+ initializer=initializers,
+ outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT,
graph_output_shape)],
+ )
+ model = helper.make_model(
+ graph,
+ producer_name="reduce_numerical_test",
+ opset_imports=[helper.make_opsetid("", opset)],
+ )
+ inputs = {"x": np.random.randn(*input_shape).astype(np.float32)}
+ check_correctness(model, inputs, opset=opset, rtol=1e-4, atol=1e-4)
- verify_reduce_func(
- func, np.random.randn(3, 3, 3).astype(np.float32), axis=(1,),
keepdims=keepdims
- )
- verify_reduce_func(
- func, np.random.randn(3, 3, 3, 1).astype(np.float32), axis=(1, 2),
keepdims=keepdims
- )
[email protected]("func, dynamic, opset",
create_reduce_test_parameters_axes_attr())
+def test_all_reduce_funcs_axes_attr(func, dynamic, opset):
+ for keepdims in [True, False]:
+ for input_shape, axes in REDUCE_AXES_ATTR_TEST_CASES:
+ expected = _make_reduce_expected_ir(func, input_shape, axes,
False, keepdims, dynamic)
+ verify_composite_reduce_axes_attr_ir(
+ func, input_shape, axes, keepdims, dynamic, opset, expected
+ )
- verify_reduce_func(
- func, np.random.randn(3, 3, 3, 1).astype(np.float32), axis=(1,),
keepdims=keepdims
- )
- verify_reduce_func(
- func, np.random.randn(1, 3, 4, 1).astype(np.float32), axis=(1,),
keepdims=keepdims
- )
[email protected](
+ "func, opset, dynamic, keepdims",
+ [
+ ("ReduceMax", 11, False, False),
+ ("ReduceMean", 11, False, False),
+ ("ReduceMean", 13, True, True),
+ ("ReduceMin", 11, False, False),
+ ("ReduceProd", 11, False, False),
+ ("ReduceProd", 13, False, False),
+ ("ReduceSum", 11, False, False),
+ ],
+)
+def test_reduce_funcs_axes_attr_numerical(func, opset, dynamic, keepdims):
+ _verify_reduce_numerical(
+ func,
+ opset,
+ axes_as_input=False,
+ axes=[1],
+ dynamic=dynamic,
+ keepdims=keepdims,
+ )
-def _make_composite_reduce_expected_ir(
+def _make_reduce_expected_ir(
func: str,
input_shape: list[int],
axes,
@@ -5229,8 +5307,19 @@ def _make_composite_reduce_expected_ir(
parser_vars["axes_shape"] = axes_shape
params.append(' reduce_axes: R.Tensor(axes_shape,
dtype="int64")')
+ basic_reduce_op = {
+ "ReduceMax": R.max,
+ "ReduceMean": R.mean,
+ "ReduceMin": R.min,
+ "ReduceProd": R.prod,
+ "ReduceSum": R.sum,
+ }.get(func)
+
if noop_with_empty_axes and not axes:
body = [" gv = x"]
+ elif basic_reduce_op is not None:
+ parser_vars["reduce_op"] = basic_reduce_op
+ body = [" gv = reduce_op(x, axis=axis, keepdims=keepdims)"]
elif func == "ReduceSumSquare":
body = [
" lv = R.multiply(x, x)",
@@ -5296,7 +5385,7 @@ def test_composite_reduce_funcs_axes_attr_ir():
for keepdims in [True, False]:
for dynamic in [True, False]:
for input_shape, axes in REDUCE_AXES_ATTR_TEST_CASES:
- expected = _make_composite_reduce_expected_ir(
+ expected = _make_reduce_expected_ir(
func, input_shape, axes, False, keepdims, dynamic
)
for opset in [13, 11]:
@@ -5391,107 +5480,52 @@ def verify_composite_reduce_axes_input_ir(
@pytest.mark.parametrize("func, dynamic, opset",
create_reduce_test_parameters_axes_input())
def test_all_reduce_funcs_axes_input(func, dynamic, opset):
- def verify_reduce_func(func, data, axes, keepdims,
noop_with_empty_axes=False):
- inshape = data.shape
- inputs = ["x"]
- initializers = []
-
- # Optional `axes` input
- if axes is not None:
- axes_name = "reduce_axes"
- axes_np = np.asarray(axes, dtype=np.int64)
- axes_init = helper.make_tensor(
- name=axes_name,
- data_type=TensorProto.INT64,
- dims=axes_np.shape,
- vals=axes_np,
- )
- initializers.append(axes_init)
- inputs.append(axes_name)
-
- # Determine input and output shapes
- if not axes and not noop_with_empty_axes:
- outshape = np.sum(data, axis=None, keepdims=keepdims).shape
- elif not axes and noop_with_empty_axes:
- outshape = inshape
- else:
- outshape = np.sum(data, axis=axes, keepdims=keepdims).shape
-
- if dynamic:
- in_list = ["?"] * len(inshape)
- out_list = ["?"] * len(outshape)
- else:
- in_list = list(inshape)
- out_list = list(outshape)
-
- # Make a model node
- node = helper.make_node(
- func,
- inputs=inputs,
- outputs=["y"],
- keepdims=keepdims,
- noop_with_empty_axes=noop_with_empty_axes,
- )
-
- # Make a model graph and a model
- graph = helper.make_graph(
- [node],
- "reduce18_test",
- inputs=[helper.make_tensor_value_info("x", TensorProto.FLOAT,
in_list)],
- initializer=initializers,
- outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT,
out_list)],
- )
- model = helper.make_model(graph, producer_name="reduce18_test")
-
- inputs_dict = {"x": data}
- check_correctness(model, inputs_dict, opset=opset, rtol=1e-4,
atol=1e-4)
-
- # Verify
for keepdims in [True, False]:
- # no `axes` input && `noop_with_empty_axes` = 0 -> reduce over all
dimensions.
- verify_reduce_func(
- func,
- np.random.randn(3, 2, 2).astype(np.float32),
- axes=[],
- keepdims=keepdims,
- noop_with_empty_axes=False,
- )
-
- # no `axes` input && `noop_with_empty_axes` = 0 -> reduce over all
dimensions.
- verify_reduce_func(
- func,
- np.random.randn(3, 2, 2).astype(np.float32),
- axes=None,
- keepdims=keepdims,
- noop_with_empty_axes=False,
- )
-
- # no `axes` input && `noop_with_empty_axes` = 1 -> return the input
unchanged.
- verify_reduce_func(
- func,
- np.random.randn(4, 3).astype(np.float32),
- axes=[],
- keepdims=keepdims,
- noop_with_empty_axes=True,
- )
+ for input_shape, axes, noop_with_empty_axes in
REDUCE_AXES_INPUT_TEST_CASES:
+ expected = _make_reduce_expected_ir(
+ func,
+ input_shape,
+ axes,
+ noop_with_empty_axes,
+ keepdims,
+ dynamic,
+ axes_as_input=True,
+ )
+ verify_composite_reduce_axes_input_ir(
+ func,
+ input_shape,
+ axes,
+ noop_with_empty_axes,
+ keepdims,
+ dynamic,
+ opset,
+ expected,
+ )
- # no `axes` input && `noop_with_empty_axes` = 1 -> return the input
unchanged.
- # (onnxruntime bug) Runtime error on the onnxruntime part
- # verify_reduce_func(
- # func,
- # np.random.randn(4, 3).astype(np.float32),
- # axes=None,
- # keepdims=keepdims,
- # noop_with_empty_axes=True,
- # )
- # `axes` provided -> reduce over specified axes.
- verify_reduce_func(
- func,
- np.random.randn(3, 3, 3, 1).astype(np.float32),
- axes=(1, 2),
- keepdims=keepdims,
- )
[email protected](
+ "func, opset, axes, noop_with_empty_axes, dynamic, keepdims",
+ [
+ ("ReduceMax", 18, [1], False, False, False),
+ ("ReduceMean", 18, [1], False, True, True),
+ ("ReduceMin", 18, [1], False, False, False),
+ ("ReduceProd", 18, [1], False, False, False),
+ ("ReduceSum", 13, [1], False, False, False),
+ ("ReduceSum", 13, [], True, False, False),
+ ],
+)
+def test_reduce_funcs_axes_input_numerical(
+ func, opset, axes, noop_with_empty_axes, dynamic, keepdims
+):
+ _verify_reduce_numerical(
+ func,
+ opset,
+ axes_as_input=True,
+ axes=axes,
+ noop_with_empty_axes=noop_with_empty_axes,
+ dynamic=dynamic,
+ keepdims=keepdims,
+ )
def test_composite_reduce_funcs_axes_input_ir():
@@ -5499,7 +5533,7 @@ def test_composite_reduce_funcs_axes_input_ir():
for keepdims in [True, False]:
for dynamic in [True, False]:
for input_shape, axes, noop_with_empty_axes in
REDUCE_AXES_INPUT_TEST_CASES:
- expected = _make_composite_reduce_expected_ir(
+ expected = _make_reduce_expected_ir(
func,
input_shape,
axes,
diff --git a/tests/python/relax/test_frontend_onnx_backend.py
b/tests/python/relax/test_frontend_onnx_backend.py
index 3a80473cb1..1b7526aa1f 100644
--- a/tests/python/relax/test_frontend_onnx_backend.py
+++ b/tests/python/relax/test_frontend_onnx_backend.py
@@ -29,6 +29,8 @@ semantic verification.
"""
+import re
+
import numpy as np
import pytest
@@ -119,12 +121,6 @@ class TVMRelaxBackend(Backend):
return device == "CPU"
-# ---------------------------------------------------------------------------
-# Test registration
-# ---------------------------------------------------------------------------
-
-backend_test = onnx.backend.test.BackendTest(TVMRelaxBackend, __name__)
-
# Operators where ALL ONNX node tests pass on the Relax importer.
# Each prefix covers the base test and all its variants
# (e.g. test_add, test_add_bcast, test_add_uint8).
@@ -206,7 +202,36 @@ _INCLUDE_OPS = [
"xor",
]
-for _op in _INCLUDE_OPS:
- backend_test.include(rf"^test_{_op}(?:_.*)?(?:_cpu|_cuda)$")
+# ``BackendTest.include`` registers every ONNX backend test and marks the
+# non-matches as skipped. Filter at the runner's registration hook instead so
+# CI only generates allowlisted cases for devices the backend supports.
+_INCLUDE_PATTERN = re.compile(
+ rf"^test_(?:{'|'.join(map(re.escape,
_INCLUDE_OPS))})(?:_.*)?(?:_cpu|_cuda)$"
+)
+
+
+class _AllowlistedBackendTest(onnx.backend.test.BackendTest):
+ """Register only allowlisted, supported variants with the ONNX test
runner."""
+
+ def _add_test(
+ self,
+ category,
+ test_name,
+ test_func,
+ report_item,
+ devices=("CPU", "CUDA"),
+ **kwargs,
+ ):
+ devices = [
+ device
+ for device in devices
+ if self.backend.supports_device(device)
+ and _INCLUDE_PATTERN.search(f"{test_name}_{device.lower()}")
+ ]
+ if devices:
+ super()._add_test(category, test_name, test_func, report_item,
devices, **kwargs)
+
+
+backend_test = _AllowlistedBackendTest(TVMRelaxBackend, __name__)
globals().update(backend_test.test_cases)
diff --git
a/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_cpu.py
b/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_cpu.py
index 57f4061d96..dfd0135178 100644
--- a/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_cpu.py
+++ b/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_cpu.py
@@ -17,6 +17,7 @@
# ruff: noqa: E501, E741
import enum
import itertools
+import json
import numpy as np
import pytest
@@ -88,6 +89,8 @@ fattention_rotary = None
fcopy_single_page = None
fcompact_copy = None
+_COMPILED_KERNEL_CACHE = {}
+
def set_global_func(head_dim, dtype):
global fclear, fadd_sequence, fremove_sequence, ffork_sequence,
fenable_sliding_window_for_seq
@@ -121,33 +124,51 @@ def set_global_func(head_dim, dtype):
fdebug_get_kv =
tvm.get_global_func("vm.builtin.attention_kv_cache_debug_get_kv")
target = tvm.target.Target.from_device(device)
- builts = []
- for tir_func in [
- _kv_cache_transpose_append(num_kv_heads, head_dim, dtype),
- _kv_cache_debug_get_kv(num_layers, num_kv_heads, head_dim, dtype),
- _attention_prefill_cpu(num_kv_heads, num_qo_heads, head_dim, dtype,
False, rope_scaling),
- _attention_decode_cpu(num_kv_heads, num_qo_heads, head_dim, dtype,
False, rope_scaling),
- _attention_prefill_cpu(num_kv_heads, num_qo_heads, head_dim, dtype,
True, rope_scaling),
- _attention_decode_cpu(num_kv_heads, num_qo_heads, head_dim, dtype,
True, rope_scaling),
- _attention_prefill_ragged_cpu(
- num_kv_heads, num_qo_heads, head_dim, head_dim, dtype, rope_scaling
- ),
- tree_attn_cpu(num_kv_heads, num_qo_heads, head_dim, dtype,
rope_scaling),
- tree_attn_with_paged_kv_cache_cpu(
- num_kv_heads, num_qo_heads, head_dim, dtype, rope_scaling
- ),
- _merge_state_inplace_cpu(dtype),
- llama_rope_with_position_map(
- rope_theta, rope_scale, head_dim, num_qo_heads, num_kv_heads,
dtype, rope_scaling
- ),
- _copy_single_page_cpu(num_kv_heads, page_size, head_dim, dtype),
- _compact_kv_copy_cpu(num_kv_heads, head_dim, dtype),
- ]:
- mod = tvm.IRModule({"main": tir_func})
- with target:
- mod = dl.ApplyDefaultSchedule(dl.gpu.Fallback())(mod)
- f = tvm.tirx.build(mod["main"], target=target)
- builts.append(f.main)
+ cache_key = (
+ str(target),
+ num_layers,
+ num_qo_heads,
+ num_kv_heads,
+ head_dim,
+ page_size,
+ dtype,
+ rope_scale,
+ rope_theta,
+ json.dumps(rope_scaling, sort_keys=True),
+ )
+ builts = _COMPILED_KERNEL_CACHE.get(cache_key)
+ if builts is None:
+ builts = []
+ for tir_func in [
+ _kv_cache_transpose_append(num_kv_heads, head_dim, dtype),
+ _kv_cache_debug_get_kv(num_layers, num_kv_heads, head_dim, dtype),
+ _attention_prefill_cpu(
+ num_kv_heads, num_qo_heads, head_dim, dtype, False,
rope_scaling
+ ),
+ _attention_decode_cpu(num_kv_heads, num_qo_heads, head_dim, dtype,
False, rope_scaling),
+ _attention_prefill_cpu(num_kv_heads, num_qo_heads, head_dim,
dtype, True, rope_scaling),
+ _attention_decode_cpu(num_kv_heads, num_qo_heads, head_dim, dtype,
True, rope_scaling),
+ _attention_prefill_ragged_cpu(
+ num_kv_heads, num_qo_heads, head_dim, head_dim, dtype,
rope_scaling
+ ),
+ tree_attn_cpu(num_kv_heads, num_qo_heads, head_dim, dtype,
rope_scaling),
+ tree_attn_with_paged_kv_cache_cpu(
+ num_kv_heads, num_qo_heads, head_dim, dtype, rope_scaling
+ ),
+ _merge_state_inplace_cpu(dtype),
+ llama_rope_with_position_map(
+ rope_theta, rope_scale, head_dim, num_qo_heads, num_kv_heads,
dtype, rope_scaling
+ ),
+ _copy_single_page_cpu(num_kv_heads, page_size, head_dim, dtype),
+ _compact_kv_copy_cpu(num_kv_heads, head_dim, dtype),
+ ]:
+ mod = tvm.IRModule({"main": tir_func})
+ with target:
+ mod = dl.ApplyDefaultSchedule(dl.gpu.Fallback())(mod)
+ f = tvm.tirx.build(mod["main"], target=target)
+ builts.append(f.main)
+ builts = tuple(builts)
+ _COMPILED_KERNEL_CACHE[cache_key] = builts
(
ftranspose_append,
diff --git
a/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_tir.py
b/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_tir.py
index c48ddeb7df..fe7fc0ff07 100644
--- a/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_tir.py
+++ b/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_tir.py
@@ -17,6 +17,7 @@
# ruff: noqa: E501, E741
import functools
import itertools
+import json
import pytest
import torch
@@ -87,6 +88,8 @@ fattention_rotary = None
fcopy_single_page = None
fcompact_copy = None
+_COMPILED_KERNEL_CACHE = {}
+
def set_global_func(head_dim, dtype, target):
global fclear, fadd_sequence, fremove_sequence, ffork_sequence,
fenable_sliding_window_for_seq
@@ -119,35 +122,57 @@ def set_global_func(head_dim, dtype, target):
fis_empty = tvm.get_global_func("vm.builtin.attention_kv_cache_empty")
fdebug_get_kv =
tvm.get_global_func("vm.builtin.attention_kv_cache_debug_get_kv")
- builts = []
- for tir_func in [
- _kv_cache_transpose_append(num_kv_heads, head_dim, dtype),
- _kv_cache_debug_get_kv(num_layers, num_kv_heads, head_dim, dtype),
- _attention_prefill(
- num_kv_heads, num_qo_heads, head_dim, dtype, False, rope_scaling,
target
- ),
- _attention_decode(num_kv_heads, num_qo_heads, head_dim, dtype, False,
rope_scaling, target),
- _attention_prefill(num_kv_heads, num_qo_heads, head_dim, dtype, True,
rope_scaling, target),
- _attention_decode(num_kv_heads, num_qo_heads, head_dim, dtype, True,
rope_scaling, target),
- _attention_prefill_ragged(
- num_kv_heads, num_qo_heads, head_dim, head_dim, dtype,
rope_scaling, target
- ),
- tree_attn(num_kv_heads, num_qo_heads, head_dim, dtype, rope_scaling,
target),
- tree_attn_with_paged_kv_cache(
- num_kv_heads, num_qo_heads, head_dim, dtype, rope_scaling, target
- ),
- _merge_state_inplace(num_qo_heads, head_dim, dtype, target),
- llama_rope_with_position_map(
- rope_theta, rope_scale, head_dim, num_qo_heads, num_kv_heads,
dtype, rope_scaling
- ),
- _copy_single_page(num_kv_heads, page_size, head_dim, dtype, target),
- _compact_kv_copy(num_kv_heads, head_dim, dtype, target),
- ]:
- mod = tvm.IRModule({"main": tir_func})
- with target:
- mod = dl.ApplyDefaultSchedule(dl.gpu.Fallback())(mod)
- f = tvm.tirx.build(mod["main"], target=target)
- builts.append(f.main)
+ cache_key = (
+ str(target),
+ num_layers,
+ num_qo_heads,
+ num_kv_heads,
+ head_dim,
+ page_size,
+ dtype,
+ rope_scale,
+ rope_theta,
+ json.dumps(rope_scaling, sort_keys=True),
+ )
+ builts = _COMPILED_KERNEL_CACHE.get(cache_key)
+ if builts is None:
+ builts = []
+ for tir_func in [
+ _kv_cache_transpose_append(num_kv_heads, head_dim, dtype),
+ _kv_cache_debug_get_kv(num_layers, num_kv_heads, head_dim, dtype),
+ _attention_prefill(
+ num_kv_heads, num_qo_heads, head_dim, dtype, False,
rope_scaling, target
+ ),
+ _attention_decode(
+ num_kv_heads, num_qo_heads, head_dim, dtype, False,
rope_scaling, target
+ ),
+ _attention_prefill(
+ num_kv_heads, num_qo_heads, head_dim, dtype, True,
rope_scaling, target
+ ),
+ _attention_decode(
+ num_kv_heads, num_qo_heads, head_dim, dtype, True,
rope_scaling, target
+ ),
+ _attention_prefill_ragged(
+ num_kv_heads, num_qo_heads, head_dim, head_dim, dtype,
rope_scaling, target
+ ),
+ tree_attn(num_kv_heads, num_qo_heads, head_dim, dtype,
rope_scaling, target),
+ tree_attn_with_paged_kv_cache(
+ num_kv_heads, num_qo_heads, head_dim, dtype, rope_scaling,
target
+ ),
+ _merge_state_inplace(num_qo_heads, head_dim, dtype, target),
+ llama_rope_with_position_map(
+ rope_theta, rope_scale, head_dim, num_qo_heads, num_kv_heads,
dtype, rope_scaling
+ ),
+ _copy_single_page(num_kv_heads, page_size, head_dim, dtype,
target),
+ _compact_kv_copy(num_kv_heads, head_dim, dtype, target),
+ ]:
+ mod = tvm.IRModule({"main": tir_func})
+ with target:
+ mod = dl.ApplyDefaultSchedule(dl.gpu.Fallback())(mod)
+ f = tvm.tirx.build(mod["main"], target=target)
+ builts.append(f.main)
+ builts = tuple(builts)
+ _COMPILED_KERNEL_CACHE[cache_key] = builts
(
ftranspose_append,
diff --git a/tests/python/relax/test_vm_cuda_graph.py
b/tests/python/relax/test_vm_cuda_graph.py
index 663578227f..a842f035dd 100644
--- a/tests/python/relax/test_vm_cuda_graph.py
+++ b/tests/python/relax/test_vm_cuda_graph.py
@@ -16,6 +16,9 @@
# under the License.
# ruff: noqa: E501
+import ctypes
+import ctypes.util
+
import numpy as np
import pytest
@@ -170,6 +173,12 @@ def test_capture_error_is_recoverable():
def run_and_check():
dev = tvm.cuda()
+ cudart_path = ctypes.util.find_library("cudart")
+ assert cudart_path is not None, "Unable to locate the CUDA runtime
library"
+ cudart = ctypes.CDLL(cudart_path)
+ cudart.cudaGetLastError.argtypes = []
+ cudart.cudaGetLastError.restype = ctypes.c_int
+ cudart.cudaGetLastError()
@tvm.register_global_func("test_vm_cuda_graph.invalid_impl_for_cudagraph",
override=True)
def invalid_impl_for_cudagraph(arg_tensor):
@@ -186,6 +195,13 @@ def test_capture_error_is_recoverable():
with pytest.raises(RuntimeError):
vm["main"](arg)
+ # cudaGetLastError is host-thread-local, so query it in the same
+ # callback that triggered the invalid capture.
+ cuda_error = cudart.cudaGetLastError()
+ assert cuda_error == 0, (
+ f"CUDA error state was not cleared after failed graph capture:
{cuda_error}"
+ )
+
tvm.testing.run_with_gpu_lock(run_and_check)
