Skip to content

[Relax] Fix FoldConstant crash on data-dependent reshape with unknown shape - #20488

Open
cchung100m wants to merge 3 commits into
apache:mainfrom
cchung100m:issue-20260
Open

cchung100m wants to merge 3 commits into
apache:mainfrom
cchung100m:issue-20260

Conversation

@cchung100m

@cchung100m cchung100m commented Sep 29, 2026 •

Copy link
Copy Markdown
Contributor

Hi Committers,

This PR fixes issues #20260. Any suggestions would be appreciated if you are available.

Root Cause

FoldConstant has a speculative step that calls an op's registered FLegalize function to check if it would fold into an all-constant call_tir. Unlike the real LegalizeOps pass, this speculative probe never checks whether the call's argument/return shapes are actually known before invoking FLegalize.

When reshape's target shape comes from an unfolded tensor_to_shape call (i.e., a genuinely non-constant tensor, so the shape values are unknown), this lets the unresolved shape reach reshape's legalizer (reshape_call_te), which asserts - with no message - that the shape has already resolved to a ShapeExpr. That's what actually crashes.

Note: the original issue attributes the crash to fold_constant.cc's shape_to_tensor special case, but that code path is dead for this repro (shape_to_tensor has had a real FLegalize since #19957) - it was already stale when the issue was filed. This PR fixes the crash reachable today (the tensor_to_shape -> reshape case). The issue's other example (shape_to_tensor fed directly into reshape) is unrelated to this crash path. It is now rejected earlier, at graph-construction time, by unrelated, more recent type-inference changes - not by this PR.

Solutions

  • Extract the "can this op be legalized right now" check (previously a private lambda inside LegalizeOps) into a shared CanLegalizedCall helper in src/relax/transform/utils.h, reused by both LegalizeOps and FoldConstant.
  • FoldConstant now checks CanLegalizeCall before its speculative legalize-and-fold probe; when shapes aren't statically known, it skips folding and restores the call's original type (previously reset to Type::Missing(), which could let a later stage infer an ill-formed type referencing a block-scoped DataflowVar).
  • Hardened reshape/broadcast_to/collapse_sum_*'s shared FLegalize (reshape_call_te) to raise a descriptive ValueError instead of a bare assert, as defense in depth.

@cchung100m
cchung100m force-pushed the issue-20260 branch 6 times, most recently from 2953448 to 56ad79a Compare September 29, 2026 12:38
@cchung100m
cchung100m marked this pull request as ready for review September 30, 2026 02:53
@cchung100m

Copy link
Copy Markdown
Contributor Author

cc @tlopex @mshr-h

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