[Relax][Test] Cover default GPU pipeline scheduling for R.power/elementwise kernels

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()