gemini-code-assist[bot] commented on code in PR #19923:
URL: https://github.com/apache/tvm/pull/19923#discussion_r3524633813
##########
tests/python/relax/test_pipeline.py:
##########
@@ -149,3 +150,79 @@ 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():
+ """The CUDA legalization pipeline thread-binds a legalized elementwise
kernel.
+
+ Device-free (no GPU required): runs the CUDA `legalize_passes`, which end
+ right after DLight scheduling, so the only TIR PrimFunc left is the `power`
+ kernel itself (no later host-side shape helpers to confuse the check). The
+ kernel must carry a GPU thread binding, otherwise `VerifyMemory` would
reject
+ it during a real build.
+ """
+ target = tvm.target.Target(
+ "cuda -keys=cuda,gpu -arch=sm_86 -max_num_threads=1024
-thread_warp_size=32"
+ )
+ with target:
+ seq = tvm.transform.Sequential(relax.pipeline.legalize_passes(target))
+ mod = seq(PowerModule)
+
+ 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 legalization"
+ for func in prim_funcs:
+ assert _has_thread_binding(func), (
+ "power PrimFunc left without a GPU thread binding (VerifyMemory
would fall)"
+ )
Review Comment:
Asserting that **all** `PrimFunc`s in the module have thread bindings is
fragile. While currently there is only one `PrimFunc` (the legalized `power`
kernel), future updates to the Relax legalization or lowering passes might
introduce host-side helper `PrimFunc`s (e.g., for shape calculations or
metadata) which do not require thread bindings. This would cause the test to
fail unexpectedly.
Additionally, there is a minor typo in the assertion message (`VerifyMemory
would fall` -> `VerifyMemory would fail`).
We can make this more robust by filtering the `PrimFunc`s to only check
those whose names contain `"power"`.
```suggestion
prim_funcs = [
(g_var, func)
for g_var, func in mod.functions_items()
if isinstance(func, tvm.tirx.PrimFunc) and "power" in g_var.name_hint
]
assert prim_funcs, "expected at least one power TIR PrimFunc after
legalization"
for _, func in prim_funcs:
assert _has_thread_binding(func), (
"power PrimFunc left without a GPU thread binding (VerifyMemory
would fail)"
)
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
