From f34ed212bea1e7ee9a0b19f9dbf177344a41debf Mon Sep 17 00:00:00 2001 From: "Guan-Ming (Wesley) Chiu" <105915352+guan404ming@users.noreply.github.com> Date: Mon, 6 Jul 2026 19:54:37 +0800 Subject: [PATCH] [Relax] Legalize shape_to_tensor to device kernel --- .../relax/transform/legalize_ops/create.py | 27 +++- ..._transform_legalize_ops_create_datatype.py | 119 +++++++++++++++++- 2 files changed, 144 insertions(+), 2 deletions(-) diff --git a/python/tvm/relax/transform/legalize_ops/create.py b/python/tvm/relax/transform/legalize_ops/create.py index 00383f8326a8..f4514e496e31 100644 --- a/python/tvm/relax/transform/legalize_ops/create.py +++ b/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 diff --git a/tests/python/relax/test_transform_legalize_ops_create_datatype.py b/tests/python/relax/test_transform_legalize_ops_create_datatype.py index dda27def445d..255da3671db9 100644 --- a/tests/python/relax/test_transform_legalize_ops_create_datatype.py +++ b/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