[Relax] Legalize shape_to_tensor to device kernel

[guan404ming]

Why

Fixes #19925. relax.shape_to_tensor had no legalization, so it always lowered to the host packed func relax.run.shape_to_tensor, producing a CPU tensor regardless of the target device.

How

• Register a legalization that emits the shape values as a call_tir TE kernel, passing symbolic dims via tir_vars.
• Fall back to the packed func when the shape values are unknown (ShapeStructInfo without values).

[gemini-code-assist]

Warning

Gemini encountered an error creating the review. You can try again by commenting /gemini review.

[guan404ming]

/gemini review

[gemini-code-assist]

Warning

Gemini encountered an error creating the review. You can try again by commenting /gemini review.

[tlopex]

I just tested the symbolic case and it fails in the full build pipeline.

This legalization produces a call_tir with no tensor inputs and non-empty tir_vars, e.g. [m, n]. FoldConstant treats the empty tensor argument tuple as constant and attempts to execute the generated PrimFunc without
passing m and n, resulting in:

TypeError: Expected 3 arguments when calling:
  tir_function(output, m, n)

A minimal reproducer is:

def main(x: R.Tensor(("m", "n"), "float32")):
    return R.shape_to_tensor(R.shape_of(x))

relax.build(mod, target="cuda")

The current symbolic and mixed tests only run LegalizeOps, so they do not exercise this failure. Could we make FoldConstant skip call_tir calls with non-empty symbolic tir_vars (or otherwise handle those scalar arguments),
and add an end-to-end build test?

[guan404ming]

Thanks, reproduced exactly as described. The llvm pipeline doesn't run FoldConstant, which is why the existing tests missed it.

Fixed by bailing out when call_tir carries tir_vars (args.size() > 2). Added test_call_tir_with_tir_vars_not_folded. I kept it as a FoldConstant unit test rather than an end-to-end relax.build(..., target="cuda") test so it doesn't need a CUDA CI node, happy to add the build test if you'd prefer.

[tlopex]

LGTM

[guan404ming]

Thanks!