This is an automated email from the ASF dual-hosted git repository.
tlopex pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 6cd73cd1bf [Relax] Legalize shape_to_tensor to device kernel (#19957)
6cd73cd1bf is described below
commit 6cd73cd1bf1c4b08a07d44dbadd95fdb086e3b53
Author: Guan-Ming Chiu <[email protected]>
AuthorDate: Sun Jul 12 10:39:09 2026 +0800
[Relax] Legalize shape_to_tensor to device kernel (#19957)
## 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).
---
python/tvm/relax/transform/legalize_ops/create.py | 27 ++++-
src/relax/transform/fold_constant.cc | 3 +
tests/python/relax/test_transform_fold_constant.py | 25 +++++
.../test_transform_legalize_ops_create_datatype.py | 119 ++++++++++++++++++++-
4 files changed, 172 insertions(+), 2 deletions(-)
diff --git a/python/tvm/relax/transform/legalize_ops/create.py
b/python/tvm/relax/transform/legalize_ops/create.py
index 00383f8326..f4514e496e 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 _arange(bb: BlockBuilder, call: Call) -> 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/src/relax/transform/fold_constant.cc
b/src/relax/transform/fold_constant.cc
index 274ad7b0a5..3930cf33a3 100644
--- a/src/relax/transform/fold_constant.cc
+++ b/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);
diff --git a/tests/python/relax/test_transform_fold_constant.py
b/tests/python/relax/test_transform_fold_constant.py
index 0867a9d630..f6720131c0 100644
--- a/tests/python/relax/test_transform_fold_constant.py
+++ b/tests/python/relax/test_transform_fold_constant.py
@@ -585,5 +585,30 @@ def test_fold_large_op_with_tensor_input():
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()
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 dda27def44..255da3671d 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 test_arange_symbolic():
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
