[Relax] Legalize shape_to_tensor to device kernel

python/tvm/relax/backend/dispatch_sort_scan.py

@@ -86,9 +86,7 @@ def visit_call_(self, call: relax.Call) -> relax.Expr:
                 if self.is_gpu_target(tgt):
                     te_func = topi.gpu.searchsorted
             out_dtype = "int32" if call.attrs.out_int32 else "int64"
-            return self.builder_.call_te(
-                te_func, boundaries, input_tensor, right, out_dtype
-            )
+            return self.builder_.call_te(te_func, boundaries, input_tensor, right, out_dtype)
         if call.op.name == "relax.sort":
             tgt = self._get_target(call.ty)
             te_func = topi.sort

python/tvm/relax/transform/legalize_ops/create.py

@@ -21,7 +21,7 @@
 import numpy as np
 
 import tvm
-from tvm import tirx, topi
+from tvm import te, tirx, topi
 from tvm.ir import Call
 
 from ...block_builder import BlockBuilder
@@ -130,6 +130,31 @@ def is_const_scalar(x: tirx.Expr):
         return bb.call_te(topi.arange, start, end, step, dtype)
 
 
+@register_legalize("relax.shape_to_tensor")
+def _shape_to_tensor(bb: BlockBuilder, call: Call) -> Expr:
+    shape = call.args[0]
+    values = shape.values if isinstance(shape, ShapeExpr) else shape.ty.values
+    if values is None:
+        return call
+    values = list(values)
+    n = len(values)
+    symbolic = [v for v in values if not isinstance(v, tirx.IntImm)]
+
+    def te_shape_to_tensor(*sym):
+        sym = list(sym)
+        resolved = [v if isinstance(v, tirx.IntImm) else sym.pop(0) for v in values]
+
+        def fcompute(i):
+            result = tirx.const(0, "int64")
+            for idx in range(n - 1, -1, -1):
+                result = tirx.if_then_else(i == idx, tirx.Cast("int64", resolved[idx]), result)
+            return result
+
+        return te.compute((n,), fcompute, name="shape_to_tensor")
+
+    return bb.call_te(te_shape_to_tensor, *symbolic, primfunc_name_hint="shape_to_tensor")
+
+
 @register_legalize("relax.hamming_window")
 def _hamming_window(bb: BlockBuilder, call: Call) -> Expr:
     assert len(call.args) == 4

src/relax/transform/fold_constant.cc

@@ -280,6 +280,9 @@ class ConstantFolder : public ExprMutator {
 
     if (!func || !arr_args) return {};
 
+    // tir_vars are passed as extra scalar arguments to the PrimFunc, which we cannot supply here.
+    if (call->args.size() > 2) return {};
+
     // Handle tuple output: ty_args[0] is a TupleType.
     if (const auto* tuple_ty = call->ty_args[0].as<TupleTypeNode>()) {
       return ConstEvaluateCallTIRTuple(func.value(), arr_args.value(), tuple_ty);

tests/python/relax/test_frontend_onnx.py

@@ -6563,7 +6563,9 @@ def verify_pad(input_shape, pads, expected, mode="constant", value=0.0, opset=14
             if axes is not None:
                 axes = np.array(axes, dtype=np.int64)
                 node_inputs = ["input", "pads", "", "axes"]
-                initializer.append(helper.make_tensor("axes", TensorProto.INT64, (len(axes),), axes))
+                initializer.append(
+                    helper.make_tensor("axes", TensorProto.INT64, (len(axes),), axes)
+                )
 
             node = helper.make_node("Pad", inputs=node_inputs, outputs=["output"], mode=mode)
             graph = helper.make_graph(
@@ -6628,7 +6630,9 @@ def verify_pad(input_shape, pads, expected, mode="constant", value=0.0, opset=14
         verify_pad(
             input_shape,
             pads,
-            _make_pad_expected_ir(input_shape, pads, mode=mode, value=value, opset=opset, axes=axes),
+            _make_pad_expected_ir(
+                input_shape, pads, mode=mode, value=value, opset=opset, axes=axes
+            ),
             mode,
             value,
             opset,

tests/python/relax/test_transform_fold_constant.py

@@ -585,5 +585,30 @@ def expected(c1: R.Tensor((2048,), "float32")):
     tvm.ir.assert_structural_equal(after, expected)
 
 
+def test_call_tir_with_tir_vars_not_folded():
+    """call_tir with symbolic tir_vars cannot be const-evaluated."""
+
+    @tvm.script.ir_module
+    class Module:
+        @T.prim_func(private=True, s_tir=True)
+        def shape_to_tensor(out: T.Buffer((T.int64(1),), "int64"), m: T.int64):
+            for i in range(T.int64(1)):
+                with T.sblock("out"):
+                    vi = T.axis.remap("S", [i])
+                    out[vi] = m
+
+        @R.function
+        def main(x: R.Tensor(("m",), "float32")):
+            m = T.int64()
+            cls = Module
+            gv = relax.call_tir(
+                cls.shape_to_tensor, R.tuple(), R.Tensor((1,), "int64"), tir_vars=R.shape([m])
+            )
+            return gv
+
+    after = relax.transform.FoldConstant()(Module)
+    tvm.ir.assert_structural_equal(after, Module)
+
+
 if __name__ == "__main__":
     tvm.testing.main()

tests/python/relax/test_transform_legalize_ops_create_datatype.py

@@ -14,7 +14,7 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-# ruff: noqa: E501
+# ruff: noqa: E501, F841
 
 import tvm
 import tvm.testing
@@ -617,6 +617,123 @@ def arange(var_T_arange: T.handle, n: T.int64):
     tvm.ir.assert_structural_equal(mod, Expected)
 
 
+def test_shape_to_tensor():
+    # fmt: off
+    @tvm.script.ir_module
+    class ShapeToTensor:
+        @R.function
+        def main(x: R.Tensor((2, 3, 4), "float32")):
+            gv = R.shape_to_tensor(R.shape_of(x))
+            return gv
+
+    @tvm.script.ir_module
+    class Expected:
+        @R.function
+        def main(x: R.Tensor((2, 3, 4), "float32")) -> R.Tensor((3,), "int64"):
+            cls = Expected
+            gv: R.Shape([2, 3, 4]) = R.shape_of(x)
+            gv_1 = R.call_tir(cls.shape_to_tensor, R.tuple(), out_ty=R.Tensor((3,), dtype="int64"))
+            return gv_1
+
+        @T.prim_func(private=True, s_tir=True)
+        def shape_to_tensor(shape_to_tensor: T.Buffer((T.int64(3),), "int64")):
+            T.func_attr({"tirx.noalias": True})
+            for i in range(T.int64(3)):
+                with T.sblock("shape_to_tensor"):
+                    v_i = T.axis.spatial(T.int64(3), i)
+                    shape_to_tensor[v_i] = T.if_then_else(v_i == T.int64(0), T.int64(2), T.if_then_else(v_i == T.int64(1), T.int64(3), T.if_then_else(v_i == T.int64(2), T.int64(4), T.int64(0))))
+    # fmt: on
+
+    mod = LegalizeOps()(ShapeToTensor)
+    tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_shape_to_tensor_symbolic():
+    # fmt: off
+    @tvm.script.ir_module
+    class ShapeToTensor:
+        @R.function
+        def main(x: R.Tensor(("m", "n"), "float32")):
+            gv = R.shape_to_tensor(R.shape_of(x))
+            return gv
+
+    @tvm.script.ir_module
+    class Expected:
+        @R.function
+        def main(x: R.Tensor(("m", "n"), "float32")) -> R.Tensor((2,), "int64"):
+            m = T.int64()
+            n = T.int64()
+            cls = Expected
+            gv: R.Shape([m, n]) = R.shape_of(x)
+            gv_1 = R.call_tir(cls.shape_to_tensor, R.tuple(), out_ty=R.Tensor((2,), dtype="int64"), tir_vars=R.shape([m, n]))
+            return gv_1
+
+        @T.prim_func(private=True, s_tir=True)
+        def shape_to_tensor(shape_to_tensor: T.Buffer((T.int64(2),), "int64"), m: T.int64, n: T.int64):
+            T.func_attr({"tirx.noalias": True})
+            for i in range(T.int64(2)):
+                with T.sblock("shape_to_tensor"):
+                    v_i = T.axis.spatial(T.int64(2), i)
+                    shape_to_tensor[v_i] = T.if_then_else(v_i == T.int64(0), m, T.if_then_else(v_i == T.int64(1), n, T.int64(0)))
+    # fmt: on
+
+    mod = LegalizeOps()(ShapeToTensor)
+    tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_shape_to_tensor_mixed():
+    # fmt: off
+    @tvm.script.ir_module
+    class ShapeToTensor:
+        @R.function
+        def main(x: R.Tensor(("m", 3), "float32")):
+            gv = R.shape_to_tensor(R.shape_of(x))
+            return gv
+
+    @tvm.script.ir_module
+    class Expected:
+        @R.function
+        def main(x: R.Tensor(("m", 3), "float32")) -> R.Tensor((2,), "int64"):
+            m = T.int64()
+            cls = Expected
+            gv: R.Shape([m, 3]) = R.shape_of(x)
+            gv_1 = R.call_tir(cls.shape_to_tensor, R.tuple(), out_ty=R.Tensor((2,), dtype="int64"), tir_vars=R.shape([m]))
+            return gv_1
+
+        @T.prim_func(private=True, s_tir=True)
+        def shape_to_tensor(shape_to_tensor: T.Buffer((T.int64(2),), "int64"), m: T.int64):
+            T.func_attr({"tirx.noalias": True})
+            for i in range(T.int64(2)):
+                with T.sblock("shape_to_tensor"):
+                    v_i = T.axis.spatial(T.int64(2), i)
+                    shape_to_tensor[v_i] = T.if_then_else(v_i == T.int64(0), m, T.if_then_else(v_i == T.int64(1), T.int64(3), T.int64(0)))
+    # fmt: on
+
+    mod = LegalizeOps()(ShapeToTensor)
+    tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_shape_to_tensor_unknown_values():
+    @tvm.script.ir_module
+    class ShapeToTensor:
+        @R.function
+        def main(s: R.Shape(ndim=2)):
+            gv = R.shape_to_tensor(s)
+            return gv
+
+    @tvm.script.ir_module
+    class Expected:
+        @R.function
+        def main(s: R.Shape(ndim=2)) -> R.Tensor((2,), "int64"):
+            gv: R.Tensor((2,), dtype="int64") = R.call_pure_packed(
+                "relax.run.shape_to_tensor", s, ty_args=(R.Tensor((2,), dtype="int64"),)
+            )
+            return gv
+
+    mod = LegalizeOps()(ShapeToTensor)
+    tvm.ir.assert_structural_equal(mod, Expected)
+
+
 def test_tril():
     # fmt: off
     @tvm.script.ir_module