Skip to content

Recover Conv/ConvTranspose rank from weight when input shape is unknown#29149

Open
fanchenkong1 wants to merge 1 commit into
microsoft:mainfrom
fanchenkong1:rank-recover
Open

Recover Conv/ConvTranspose rank from weight when input shape is unknown#29149
fanchenkong1 wants to merge 1 commit into
microsoft:mainfrom
fanchenkong1:rank-recover

Conversation

@fanchenkong1

Copy link
Copy Markdown
Contributor

Recover Conv/ConvTranspose rank from weight when input shape is unknown, enabling layout transformation to NHWC for more nodes.

Description

The layout transformer skips converting a node to NHWC when input[0] has no inferred shape.

For Conv and ConvTranspose operators, the data input (input[0]) and the weight (input[1]) always share the same rank. When the input rank is unknown, recover it from the weight.

Performance Impact

Measured on Kokoro-82M-v1.0-ONNX text-to-speech model (onnx-community/Kokoro-82M-v1.0-ONNX) with WebGPU ep,

Platform Latency reduction Speedup
Intel Wildcat Lake −32.0% 1.47×
Intel Panther Lake −20.0% 1.25×

This change yields a 1.2–1.5× speedup on the Kokoro-82M text-to-speech model.

The layout transformer skips converting a node to NHWC when input[0] has
no inferred shape. But the NCHW<->NHWC permutation depends only on rank.

For Conv/ConvTranspose the data input and weight share the same rank, so
when input[0]'s rank is unknown, recover it from the weight at input[1].
@fanchenkong1 fanchenkong1 marked this pull request as ready for review June 18, 2026 04:44
@fanchenkong1

Copy link
Copy Markdown
Contributor Author

@qjia7 @guschmue This change is ready for review. PTAL, thanks!

@qjia7 qjia7 left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Correctness

The change is sound. Per ONNX spec, Conv and ConvTranspose require W (input[1])
at the same rank as X (input[0]) — [M, C/group, k1..kn] vs [N, C, d1..dn].
Falling back to the weight's rank when the data input's rank is unknown is safe.
Downstream (ChannelFirstToLastPerm / ChannelLastToFirstPerm) only needs the
rank, not the full shape. FusedConv is covered via the existing op_type
normalization to "Conv".

Simplicity

The defensive guard node->Inputs().size() > 1 && !node->Inputs()[1].empty() is
redundant. Per ONNX spec, W is a mandatory input for both Conv and
ConvTranspose — a node missing it would already be malformed and rejected
upstream. The empty-string convention only applies to optional inputs (like
B).

Suggest simplifying to:

if (!input_rank.has_value() && (op_type == "Conv" || op_type == "ConvTranspose")) {
  input_rank = api_graph->GetValueInfo(node->Inputs()[1])->ShapeRank();
}

Using ShapeRank() over Shape()->size() is the right API choice.

Security

No new attack surface. Reads an existing graph value-info; no allocation, no
unchecked arithmetic.

Testing

A unit test that constructs a Conv with unknown input[0] rank but known weight
rank, runs the layout transformer, and asserts the transpose is inserted would
lock the behavior in.

Verdict

Approve. One optional simplification (drop the redundant guard) and an optional
test.

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.

2 participants