diff --git a/python/tvm/relax/backend/dispatch_sort_scan.py b/python/tvm/relax/backend/dispatch_sort_scan.py index 929e2e278138..76951718c497 100644 --- a/python/tvm/relax/backend/dispatch_sort_scan.py +++ b/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 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/src/relax/transform/fold_constant.cc b/src/relax/transform/fold_constant.cc index 274ad7b0a50c..3930cf33a38f 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()) { return ConstEvaluateCallTIRTuple(func.value(), arr_args.value(), tuple_ty); diff --git a/tests/python/relax/test_frontend_onnx.py b/tests/python/relax/test_frontend_onnx.py index f96984e42375..a076f01c4af2 100644 --- a/tests/python/relax/test_frontend_onnx.py +++ b/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, diff --git a/tests/python/relax/test_transform_fold_constant.py b/tests/python/relax/test_transform_fold_constant.py index 0867a9d63058..f6720131c0e6 100644 --- a/tests/python/relax/test_transform_fold_constant.py +++ b/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() 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