Skip to content

[Relax] Keep a call_tir's out_ty consistent with its arguments when canonicalizing - #20494

Open
akaashrp wants to merge 1 commit into
apache:mainfrom
akaashrp:fix/canonicalize-runtime-symbolic-vars
Open

akaashrp wants to merge 1 commit into
apache:mainfrom
akaashrp:fix/canonicalize-runtime-symbolic-vars

Conversation

@akaashrp

Copy link
Copy Markdown
Contributor

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 MatchCasts 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 the canonical out_ty when it still follows from the arguments, otherwise use the implied type
  2. Expose InferCallTIROutputTypeFromArguments in op_common.h for that
  3. Add tests for a kept MatchCast and for chained MatchCasts

…anonicalizing

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 keeps the variable
while the call's out_ty had it replaced, and the call_tir validator added
in apache#20480 rejects the pair. The same happens when two MatchCasts rename
one dimension, since a value is substituted one level at a time.

Canonicalize the arguments first and derive the type they imply. Use the
canonical out_ty when it still follows from the arguments, otherwise the
implied type. Expose InferCallTIROutputTypeFromArguments for that.

MLC LLM hit this compiling Phi-3.5-vision after the pin moved past apache#20480.
@akaashrp
akaashrp requested a review from tlopex September 30, 2026 04:27

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant