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 5128fc0655 [Relax] Keep a call_tir's out_ty consistent with its
arguments when canonicalizing (#20494)
5128fc0655 is described below
commit 5128fc065591cc4b6e2a0fc9b51718611cc4c0ea
Author: Akaash Parthasarathy <[email protected]>
AuthorDate: Sat Oct 3 16:56:27 2026 -0700
[Relax] Keep a call_tir's out_ty consistent with its arguments when
canonicalizing (#20494)
`CanonicalizeTIRVariables` substitutes the known value of a symbolic
variable into type annotations, but keeps a `MatchCast` whose variable
is used at runtime. The argument of a later `call_tir` then still
carries the variable while the call's `out_ty` had it replaced, and the
`call_tir` validator from #20480 rejects the pair:
```
The `out_ty` argument for R.call_tir must be compatible with the PrimFunc.
However, the PrimFunc's signature implies that the output should be
R.Tensor((1, new_s1, new_s2, 2, 2, 1024), dtype="float16"), but the `out_ty`
argument was R.Tensor((1, 12, 12, 2, 2, 1024), dtype="float16")
```
The same happens when two `MatchCast`s rename one dimension, because a
value is substituted one level at a time. MLC LLM hits both while
compiling Phi-3.5-vision.
1. In the canonicalizer, visit a `call_tir`'s arguments first and derive
the type they imply. Keep each output of the canonical `out_ty` that
still follows from the arguments. For a tensor output that does not,
take the implied shape and keep its dtype and `vdevice`, which the
PrimFunc signature does not carry.
2. Expose `InferCallTIROutputTypeFromArguments` in `op_common.h` for
that
3. Add tests for a kept `MatchCast`, including an output whose shape the
arguments do not determine, and for chained `MatchCast`s
---
src/relax/op/op.cc | 2 +-
src/relax/op/op_common.h | 7 ++
src/relax/transform/canonicalize_bindings.cc | 65 ++++++++++++
.../relax/test_transform_canonicalize_bindings.py | 114 +++++++++++++++++++++
4 files changed, 187 insertions(+), 1 deletion(-)
diff --git a/src/relax/op/op.cc b/src/relax/op/op.cc
index 9c15ce1dd4..29b268250d 100644
--- a/src/relax/op/op.cc
+++ b/src/relax/op/op.cc
@@ -309,7 +309,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
* \return The `arg_ty`, if it can be inferred from the arguments.
* Otherwise, std::nullopt.
*/
-static ffi::Optional<Type> InferCallTIROutputTypeFromArguments(
+ffi::Optional<Type> InferCallTIROutputTypeFromArguments(
Type func_ty, Type arg_ty, ffi::Optional<ffi::Array<int64_t>>
opt_inplace_indices) {
auto opt_callee_ty = func_ty.as<FuncType>();
TVM_FFI_CHECK(opt_callee_ty, TypeError)
diff --git a/src/relax/op/op_common.h b/src/relax/op/op_common.h
index 7b30b84038..8d25435fe8 100644
--- a/src/relax/op/op_common.h
+++ b/src/relax/op/op_common.h
@@ -238,6 +238,13 @@ inline Type InferTypeUnary(const Call& call, const
BlockBuilder&, FType f_comput
return InferTypeUnary<require_float_dtype>(call, f_compute_out_dtype);
}
+/*!
+ * \brief Derive the output type of a call_tir from the callee signature and
the
+ * argument types, or nullopt when it cannot be derived.
+ */
+ffi::Optional<Type> InferCallTIROutputTypeFromArguments(
+ Type func_ty, Type arg_ty, ffi::Optional<ffi::Array<int64_t>>
opt_inplace_indices);
+
/*!
* \brief Infer the type by returning the type of the input argument.
* \param call The context Call to the operator.
diff --git a/src/relax/transform/canonicalize_bindings.cc
b/src/relax/transform/canonicalize_bindings.cc
index 1c7b831193..a282f745a2 100644
--- a/src/relax/transform/canonicalize_bindings.cc
+++ b/src/relax/transform/canonicalize_bindings.cc
@@ -28,6 +28,7 @@
#include <tvm/ffi/extra/structural_mutate.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/relax/analysis.h>
+#include <tvm/relax/attrs/op.h>
#include <tvm/relax/expr.h>
#include <tvm/relax/expr_functor.h>
#include <tvm/relax/transform.h>
@@ -35,6 +36,8 @@
#include <tvm/relax/utils.h>
#include <tvm/tirx/stmt_functor.h>
+#include "../op/op_common.h"
+
namespace tvm {
namespace relax {
using namespace tvm::prim;
@@ -52,6 +55,68 @@ class SymbolicVarCanonicalizer : public ExprMutator {
return CanonicalizeShapeValue(expr);
}
+ Expr VisitExpr_(const CallNode* op) final {
+ static const Op& call_tir_op = Op::Get("relax.call_tir");
+ static const Op& call_tir_inplace_op = Op::Get("relax.call_tir_inplace");
+ static const Op& call_tir_with_grad_op =
Op::Get("relax.call_tir_with_grad");
+ bool is_call_tir = op->op.same_as(call_tir_op) ||
op->op.same_as(call_tir_inplace_op) ||
+ op->op.same_as(call_tir_with_grad_op);
+ if (!is_call_tir || op->args.size() != 2 || op->ty_args.size() != 1) {
+ return ExprMutator::VisitExpr_(op);
+ }
+ // The out_ty of a call_tir is checked against the argument types.
Canonicalizing the
+ // two independently can leave them mentioning different variables for the
same value,
+ // so the parts of the canonical out_ty that no longer follow from the
arguments are
+ // replaced by what the arguments imply.
+ Expr new_op = this->VisitExpr(op->op);
+ ffi::Array<Expr> new_args =
+ op->args.Map([this](const Expr& arg) { return this->VisitExpr(arg); });
+ Type out_ty = this->VisitExprDepTypeField(op->ty_args[0]);
+ ffi::Optional<ffi::Array<int64_t>> inplace_indices;
+ if (const auto* attrs = op->attrs.as<CallTIRInplaceAttrs>()) {
+ inplace_indices = attrs->inplace_indices;
+ }
+ auto implied = InferCallTIROutputTypeFromArguments(GetType(new_args[0]),
GetType(new_args[1]),
+ inplace_indices);
+ if (implied.has_value()) {
+ out_ty = ReconcileOutType(implied.value(), out_ty);
+ }
+ bool unchanged =
+ new_op.same_as(op->op) && new_args.same_as(op->args) &&
out_ty.same_as(op->ty_args[0]);
+ if (unchanged) {
+ return ffi::GetRef<Expr>(op);
+ }
+ return Call(Type::Missing(), new_op, new_args, op->attrs, {out_ty},
op->span);
+ }
+
+ /*!
+ * \brief Keep each output of out_ty that follows from the implied type. For
a tensor that
+ * does not, take the implied shape and keep the dtype and vdevice, since
the PrimFunc
+ * signature carries no vdevice and erases the shape of an output with a
dimension the
+ * arguments do not determine.
+ */
+ static Type ReconcileOutType(const Type& implied, const Type& out_ty) {
+ if (IsBaseOf(implied, out_ty)) {
+ return out_ty;
+ }
+ const auto* implied_tuple = implied.as<TupleTypeNode>();
+ const auto* out_tuple = out_ty.as<TupleTypeNode>();
+ if (implied_tuple && out_tuple && implied_tuple->fields.size() ==
out_tuple->fields.size()) {
+ ffi::Array<Type> fields;
+ for (size_t i = 0; i < out_tuple->fields.size(); ++i) {
+ fields.push_back(ReconcileOutType(implied_tuple->fields[i],
out_tuple->fields[i]));
+ }
+ return TupleType(fields, out_tuple->span);
+ }
+ const auto* implied_tensor = implied.as<TensorTypeNode>();
+ const auto* out_tensor = out_ty.as<TensorTypeNode>();
+ if (implied_tensor && out_tensor && implied_tensor->shape.has_value()) {
+ return TensorType(implied_tensor->shape.value(), out_tensor->dtype,
out_tensor->vdevice,
+ out_tensor->span);
+ }
+ return implied;
+ }
+
Expr VisitExpr_(const ShapeExprNode* op) final {
if (!canonicalize_shape_values_) return ffi::GetRef<Expr>(op);
ffi::Array<PrimExpr> values =
diff --git a/tests/python/relax/test_transform_canonicalize_bindings.py
b/tests/python/relax/test_transform_canonicalize_bindings.py
index 1b10229337..1a4c1d1b53 100644
--- a/tests/python/relax/test_transform_canonicalize_bindings.py
+++ b/tests/python/relax/test_transform_canonicalize_bindings.py
@@ -26,6 +26,7 @@ from tvm.ir.base import assert_structural_equal
from tvm.relax.transform.transform import CanonicalizeBindings
from tvm.script import ir as I
from tvm.script import relax as R
+from tvm.script import s_tir as Ts
from tvm.script import tirx as T
@@ -1394,3 +1395,116 @@ def test_trace_partial_tuple_through_round_trip():
if __name__ == "__main__":
tvm.testing.main()
+
+
+def test_call_tir_out_ty_follows_arguments_when_match_cast_is_kept():
+ """A call_tir's out_ty must match what its arguments imply.
+
+ The match_cast defining `n` stays because `n` is passed to a kernel at
+ runtime, so the argument of the later call_tir keeps `n` in its shape.
+ Canonicalizing the call's out_ty on its own would replace `n` with 12
+ and leave the two disagreeing. Only that output is rewritten: the second
+ output's shape cannot be derived from the arguments, and neither output's
+ vdevice is in the PrimFunc signature, so both are kept.
+ """
+ n = T.dynamic("n", "int64")
+ m = T.dynamic("m", "int64")
+ k = T.dynamic("k", "int64")
+
+ @I.ir_module
+ class Before:
+ I.module_global_infos({"vdevice": [R.vdevice("llvm", 0),
R.vdevice("llvm", 1)]})
+
+ @Ts.prim_func(private=True)
+ def transpose(
+ x: T.Buffer((1, n, 2, m, 2, 8), "float16"),
+ y: T.Buffer((1, n, m, 2, 2, 8), "float16"),
+ z: T.Buffer((k,), "float16"),
+ ):
+ for i0, i1, i2, i3, i4, i5 in T.grid(1, n, m, 2, 2, 8):
+ with Ts.sblock("b"):
+ v0, v1, v2, v3, v4, v5 = Ts.axis.remap("SSSSSS", [i0, i1,
i2, i3, i4, i5])
+ y[v0, v1, v2, v3, v4, v5] = x[v0, v1, v3, v2, v4, v5]
+
+ @Ts.prim_func(private=True)
+ def add_scalar(x: T.Buffer((1, 8), "float16"), s: T.int64, y:
T.Buffer((1, 8), "float16")):
+ for i in T.serial(8):
+ with Ts.sblock("b"):
+ vi = Ts.axis.spatial(8, i)
+ y[0, vi] = x[0, vi] + T.Cast("float16", s)
+
+ @R.function
+ def main(x: R.Tensor((1, 12, 2, 12, 2, 8), "float16"), w: R.Tensor((1,
8), "float16")):
+ cls = Before
+ with R.dataflow():
+ lv = R.match_cast(x, R.Tensor((1, n, 2, m, 2, 8), "float16"))
+ p = R.call_tir(
+ cls.transpose,
+ (lv,),
+ out_ty=[
+ R.Tensor((1, n, m, 2, 2, 8), "float16", "llvm:1"),
+ R.Tensor((16,), "float16", "llvm:1"),
+ ],
+ )
+ y = R.call_tir(cls.add_scalar, (w, n), out_ty=R.Tensor((1, 8),
"float16"))
+ gv = (p, y)
+ R.output(gv)
+ return gv
+
+ after = relax.transform.CanonicalizeBindings()(Before)
+ lv_binding, p_binding = after["main"].body.blocks[0].bindings[:2]
+ assert isinstance(lv_binding, relax.MatchCast)
+ assert p_binding.value.args[1].fields[0].same_as(lv_binding.var)
+ transposed, extra = p_binding.value.ty_args[0].fields
+ # The out_ty still names the variables the argument carries.
+ assert transposed.shape[1].same_as(n)
+ assert transposed.shape[2].same_as(m)
+ assert [int(dim) for dim in extra.shape] == [16]
+ expected_vdevice =
Before["main"].body.blocks[0].bindings[1].value.ty_args[0].fields[0].vdevice
+ for output in (transposed, extra):
+ tvm.ir.assert_structural_equal(output.vdevice, expected_vdevice)
+ relax.analysis.well_formed(after)
+
+
+def test_call_tir_out_ty_follows_arguments_through_chained_match_casts():
+ """Two match_casts rename the same dimension before a call_tir."""
+ n = T.dynamic("n", "int64")
+ a = T.dynamic("a", "int64")
+ b = T.dynamic("b", "int64")
+
+ @I.ir_module
+ class Before:
+ @Ts.prim_func(private=True)
+ def copy(x: T.Buffer((n, 8), "float16"), y: T.Buffer((n, 8),
"float16")):
+ for i, j in T.grid(n, 8):
+ with Ts.sblock("b"):
+ vi, vj = Ts.axis.remap("SS", [i, j])
+ y[vi, vj] = x[vi, vj]
+
+ @Ts.prim_func(private=True)
+ def add_scalar(x: T.Buffer((1, 8), "float16"), s: T.int64, y:
T.Buffer((1, 8), "float16")):
+ for i in T.serial(8):
+ with Ts.sblock("b"):
+ vi = Ts.axis.spatial(8, i)
+ y[0, vi] = x[0, vi] + T.Cast("float16", s)
+
+ @R.function
+ def main(x: R.Tensor((n, 8), "float16"), w: R.Tensor((1, 8),
"float16")):
+ cls = Before
+ with R.dataflow():
+ lv1 = R.match_cast(x, R.Tensor((a, 8), "float16"))
+ lv2 = R.match_cast(lv1, R.Tensor((b, 8), "float16"))
+ c = R.call_tir(cls.copy, (lv2,), out_ty=R.Tensor((b, 8),
"float16"))
+ y = R.call_tir(cls.add_scalar, (w, n), out_ty=R.Tensor((1, 8),
"float16"))
+ gv = (c, y)
+ R.output(gv)
+ return gv
+
+ after = relax.transform.CanonicalizeBindings()(Before)
+ for binding in after["main"].body.blocks[0].bindings:
+ value = binding.value
+ if isinstance(value, relax.Call) and value.op ==
tvm.ir.Op.get("relax.call_tir"):
+ # Rebuilding the call re-runs the validator, so this passes only
when the
+ # out_ty agrees with the argument types.
+ relax.Call(value.op, value.args, value.attrs, value.ty_args)
+ relax.analysis.well_formed(after)
