This is an automated email from the ASF dual-hosted git repository.
tlopex 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 99e1c2cda8 [Relax][Test] Cover default GPU pipeline scheduling for
R.power/elementwise kernels (#19923)
99e1c2cda8 is described below
commit 99e1c2cda8d5e5215f55bef51581e124d173f8ba
Author: Neo Chien <[email protected]>
AuthorDate: Wed Aug 5 03:41:32 2026 +0800
[Relax][Test] Cover default GPU pipeline scheduling for R.power/elementwise
kernels (#19923)
Hi Commiters,
PR #19384 fixed routing in `vm_build.py:249-265` but no regression test
was added. This PR adds regression coverage for GPU scheduling of
legalized elementwise ops for issue #19548.
Closes the test coverage gap:
1. Device-Free Test: Verifies CUDA pipeline produces thread-bound
kernels
2. GPU-Test: end-to-end compile + execute, verify numerical correctness
---------
Co-authored-by: cchung100m <[email protected]>
Co-authored-by: gemini-code-assist[bot]
<176961590+gemini-code-assist[bot]@users.noreply.github.com>
---
tests/python/relax/test_pipeline.py | 87 +++++++++++++++++++++++++++++++++++++
1 file changed, 87 insertions(+)
diff --git a/tests/python/relax/test_pipeline.py
b/tests/python/relax/test_pipeline.py
index f0b1c4291a..ae7a889940 100644
--- a/tests/python/relax/test_pipeline.py
+++ b/tests/python/relax/test_pipeline.py
@@ -23,6 +23,7 @@ import tvm.testing
from tvm import relax
from tvm.script import relax as R
from tvm.script import tirx as T
+from tvm.testing import env
def test_pipeline_compile():
@@ -149,3 +150,89 @@ def test_non_gpu_target_raises_error(target_name,
pipeline_func):
target = tvm.target.Target(target_name)
with pytest.raises(ValueError, match="not yet supported"):
pipeline_func(target)
+
+
+# An elementwise binary op with a scalar constant operand. `R.power(x, const)`
+# legalizes to a single elementwise TIR PrimFunc, which the default GPU
pipeline
+# must schedule (bind to GPU threads). Without a thread binding the kernel
+# access memory from the host and `VerifyMemory` rejects it at build time
+# ("... is directly accessed by the host memory ... Did you forget to bind?").
[email protected]_module
+class PowerModule:
+ @R.function
+ def main(x: R.Tensor((1, 2, 1, 1), dtype="float32")) -> R.Tensor((1, 2, 1,
1), dtype="float32"):
+ with R.dataflow():
+ y: R.Tensor((1, 2, 1, 1), dtype="float32") = R.power(x,
R.const(2.0, "float32"))
+ R.output(y)
+ return y
+
+
+def _has_thread_binding(func: tvm.tirx.PrimFunc) -> bool:
+ """Whether the PrimFunc body contains a GPU thread-binding loop."""
+ found = False
+
+ def _visit(node):
+ nonlocal found
+ if isinstance(node, tvm.tirx.For) and node.kind ==
tvm.tirx.ForKind.THREAD_BINDING:
+ found = True
+
+ tvm.tirx.stmt_functor.post_order_visit(func.body, _visit)
+ return found
+
+
+def test_default_cuda_pipeline_schedules_power():
+ """Exercise CUDA pipeline selection and verify compute kernel is
thread-bound.
+
+ Device-free test (no GPU required) that exercises the backend-level
pipeline
+ selection path (get_default_pipeline) for CUDA targets and verifies that
+ legalized elementwise compute kernels receive GPU thread bindings.
+
+ Note: The full pipeline selection routing in vm_build.py (Layer 1: "gpu" in
+ target.keys decision) is covered by the GPU-gated end-to-end test
+ (test_power_cuda_build_and_run). This test focuses on Layer 2:
backend-specific
+ pipeline composition (DLight scheduling application).
+ """
+ target = tvm.target.Target({"kind": "cuda", "arch": "sm_86"})
+
+ # Verify: Target must be recognized as GPU (Layer 1 prerequisite)
+ assert "gpu" in target.keys, "CUDA target must be marked as GPU for proper
pipeline selection"
+
+ # Exercise: Get the backend-specific default pipeline for this target
+ pipeline = relax.pipeline.get_default_pipeline(target)
+ mod = pipeline(PowerModule)
+
+ # Verify: At least one compute kernel has thread binding.
+ # (Using any() avoids false-fail from host-side shape helper kernels
+ # that may be generated by finalize_passes but don't need GPU binding)
+ prim_funcs = [func for _, func in mod.functions.items() if
isinstance(func, tvm.tirx.PrimFunc)]
+ assert prim_funcs, "expected at least one TIR PrimFunc after pipeline"
+
+ has_bound_kernel = any(_has_thread_binding(f) for f in prim_funcs)
+ assert has_bound_kernel, (
+ "At least one compute kernel must have GPU thread binding; "
+ "shape helper kernels may not be bound"
+ )
+
+
[email protected]
[email protected](not env.has_cuda(), reason="need cuda")
+def test_power_cuda_build_and_run():
+ """End-to-End build and run of an elementwise `R.power` kernel on CUDA.
+
+ Compiles through `tvm.compile`, which for a GPU target selects the
+ target-specific default pipeline (with DLight scheduling), then executes
the
+ kernel and checks the result.
+ """
+ dev = tvm.cuda(0)
+ target = tvm.target.Target.from_device(dev)
+
+ ex = tvm.compile(PowerModule, target=target)
+ vm = relax.VirtualMachine(ex, dev)
+
+ x_np = np.random.rand(1, 2, 1, 1).astype(np.float32)
+ out = vm["main"](tvm.runtime.tensor(x_np, dev))
+ tvm.testing.assert_allclose(out.numpy(), x_np**2, rtol=1e-6, atol=1e-6)
+
+
+if __name__ == "__main__":
+ tvm.testing.main()
