diff --git a/tests/python/relax/test_pipeline.py b/tests/python/relax/test_pipeline.py index f0b1c4291a30..ae7a889940f1 100644 --- a/tests/python/relax/test_pipeline.py +++ b/tests/python/relax/test_pipeline.py @@ -23,6 +23,7 @@ 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?"). +@tvm.script.ir_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" + ) + + +@pytest.mark.gpu +@pytest.mark.skipif(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()