Add dynamic shape explainer

[miaobin]

This is the initial draft to summarize the discussion on dynamic shapes #883 .

The explainer covers named/unnamed dynamic dimensions, deferred (dispatch-time) shape validation, computeShapes() API, and a new family of shape-as-data (*Dynamic) operators.

Open questions and considered alternatives are called out explicitly in the explainer and feedback on this doc would be very welcome.

[miaobin]

A POC based on the current explainer is WIP. And it has been validated against some real-world Transformer, LLM and image generation models.
With credits to @Honry for the ORT WebNN EP POC and assistance with the explainer, @huningxin and @fdwr for the initial review of this explainer.

[fdwr]

👍 I have some thoughts, but it's 95% 👌.

[w3cbot]

anssiko marked as non substantive for IPR from ash-nazg.

[anssiko]

I cleared the automatic IPR check for this PR, since explainer documents are consider non-substantive from the W3C Patent Policy point of view.

Details

This IPR check is to ensure normative portions of the specification come from organizations who participate the WG. Furthermore, in this case, I can attest the authors of this PR are affiliated with Intel, and as such, any normative portions are reusable in the spec PR, as appropriate.

[anssiko]

@reillyeon you had good questions comments on our last call for this. Do you have some other Googlers in mind who should review this explainer PR?

When this PR lands, the team will start landing the implementation in smaller chunks for further validation of this approach.

[reillyeon]

Overall I think this approach is reasonable. My main concern is that by making shapes dynamic there are more opportunities for memory safety issues in the implementation.

[anssiko]

With two approvals, editors are welcome to merge this explainer PR at will.

Thank you @miaobin for this contribution and Reilly, Dwayne, everyone for your review.

[fdwr]

👍 Cool.

[miaobin]

@huningxin The latest Patchset of the dynamic shape POC runs dynamic shape inference and validation in Blink at dispatch time and also makes MLGraph.computeShapes() synchronous. PTAL, thanks.

[theHamsta]

the post-compilation on dispatch might be finger printable (below discussion suggest computeShapes is the natural place for post-compilation and shape specialization). Also, the implementation hitting an very unfortunate shape might have a finger printable cost that allows guesses about the used HW and implementation.

[theHamsta]

for those operators, symbolic shape inference would be inherently more difficult as the symbolic expression propagation would no longer be just shapes of tensors but also allowed to be actual data of the tensors.

The consequence of such operators is also that computeShape might be that it is possible to construct networks where computeShape is not simpler by any amount than actual dispatch:

• you pass you data as shapes
• you transform shape into data via the shape operator
• you run inference on the data (convolutions, matrix multiply)
• you transform data back into shape via those dynamic variants

So, your WebNN implementation is required to run actual inference during computeShapes and therefore also needs to have access to the weights. You could do that with the actual backend framework, but this framework will do its own kind of validation and inference behavior. How can a separate WebNN validation and shape inference logic from the validation and shape inference logic of my backend? I also know that many implementation will not be able to perform arbitrary meta on shape-to-data-and-back. Some stuff will work other things not. What does WebNN actually require implementation to work?

Another open question here is then for computeShape or shape inference in general, is it expected that the runtime on CPU or the accelerator runs the shape computation? For ONNX runtime, it is something weird in-between. TensorRT performs some analysis and tries to separate data paths where some tensors are shape tensors which ought to be interpreted by the runtime and others are GPU tensors which will be calculated by the runtime. Others, are kind of shape tensors, but they depend on inputs (computeShapes does not allow to depend on input data, but it still allows to transform input shapes into data) or are truely dispatch dependent (NonZero). Such tensors can just be interpreted by an accelerator by actually performing the calculation on a worst case maximum of memory and then output the calculated output side. A runtime has to classify those cases and it will work for most constructed network. However, there exist constructible networks that will just not work with most runtime. Which of those are required by WebNN to work?

When both computeShape and dispatch may use network constants where should I put them? On the accelerator or the runtime? Also this has to be determined by some kind of analysis. Some of those constant are required both for shape inference and inference. Would it make sense to separate those data paths and have one type MLOperand for data evaluated during dispatch and another MLShapeOperand which is only interpreted on CPU by the runtime? I can go from input shapes into calculations of MLShapeOperand which I interpret, but I have separate them for the data that is calculated during dispatch and on the accelerator. I would prefer to be able to reason tensor shapes. MLShapeOperand would be interpreted during computeShapes (or the validation of dispatch), MLOperand is interpreted from the accelerator during dispatch.

In RustNN, we currently construct the graph and translate it to the semantics of the backend framework. We preferably put the constant to the accelerator and the accelerator framework. We don't want to keep a copy of the constant for the RustNN runtime. If it comes now to dispatch or computeShape, we can either rely on semantics of the implementation framework to do computeShape/dispatch which might not agree with WebNN semantics (or whether it agrees is determined on build) or I duplicate all inference logic into something that just does shape inference conforming to WebNN.

[reillyeon]

I had a similar concern that computing shapes could effectively require executing the model but I convinced myself that the intent here is that the compute graph for shapes is separate from the compute graph for tensors. However the reuse of the MLOperand type makes that hard to see. We probably need to make this explicit by either adding an isShape property to an operand (forbidding non-shape operators from being used to compute shapes) or by using a completely different operand type.

[miaobin]

What does WebNN actually require implementation to work?

That was a real gap. The explainer said "no blessed list of operators" without saying what must therefore work, which reads as unbounded. It conflated two questions, and I've split them:

What a graph may express is not restricted by operator identity. Restricting it that way is what @fdwr objected to earlier, and I still think that's right. What is restricted is the chain's root: shape() outputs and build-time constants only. That restriction is the load-bearing one. It is why a shape chain cannot become the model.

ONNX answers the same question with a deliberately narrow closed set — the 11 operators carrying a data-propagation function (Add, Cast, Concat, Gather, Mul, Shape, Size, Slice, Squeeze, Sub, Unsqueeze) while ORT's symbolic_shape_infer.py covers a much wider set. A required minimum for WebNN belongs between the two. Naming it is work this proposal still owes, and I've written that down rather than leaving it implied; for calibration, the prototype today is a strict superset of the ONNX set.

Is it expected that the runtime on CPU or the accelerator runs the shape computation?

On CPU. A shape chain is integer work on vectors of at most rank length, and an implementation that cannot evaluate one fails cleanly rather than dispatching it.

Which of those are required by WebNN to work?

No classification is needed because the boundary is very clear: a chain bottoms out either at shape() and constants, or at a tensor's data. And that route isn't merely an implementation choice for us: there is no worst case to allocate (the model is deliberately unbounded, so bounded dimensions are a precondition), dispatch() takes caller-allocated outputs so the produced size would need a way back, and computeShapes() could not answer for such an output at all. I've expanded the Tensor-Derived open question to record that.

When both computeShape and dispatch may use network constants where should I put them?

That is a good question. Similar to the approach you described for RustNN, WebNN currently uploads constants to the service side as soon as they are created. Indeed, to perform validation based on actual input shapes, we need to retain certain constants related to shape inference. And the POC additionally caps each collected constant at a shape vector's worth of elements (16, covering Pad's 2 * max_rank). So you keep a handful of ≤16-element vectors, not anything model-sized, and the constants that are needed for both paths cost about a hundred bytes of duplication rather than a second copy of the model.
About the compute logic what we did is much smaller than "duplicate all inference logic". It is an integer interpreter over roughly twenty operators on vectors of at most rank length — about 880 lines in our implementation.

About MLShapeOperand

I added a section titled "Making the shape subgraph explicit" to the "Open Questions" section to facilitate further discussion.

[theHamsta]

That is a good question. Similar to the approach you described for RustNN, WebNN currently uploads constants to the service side as soon as they are created. Indeed, to perform validation based on actual input shapes, we need to retain certain constants related to shape inference.

Ah, so the intent is that computeShape is performed by the browser WebNN runtime on a certain sub-section of the graph and the rest of the graph is then performed by some accelerator framework. So, an implementation would try to split the graph into those two parts. Both for the shape graph and for the accelerator graph I need an implementation of WebNN to run all (or a fraction) the operators.

On CPU. A shape chain is integer work on vectors of at most rank length, and an implementation that cannot evaluate one fails cleanly rather than dispatching it.

Is there any guidance on how complete such a shape chain should be in comparison to a full WebNN implementation? Or would a typical approach be to just run it with the CPU backend of your WebNN implementation? Or just an ad-hoc interpreter for what you think that networks would typically use?

[theHamsta]

TRT has a function to get the output shapes given certain input shapes. ONNX runtime can only try to evaluate symbolic inference if it has been performed before on the ONNX, for the usual case, inference is required to know the output shapes as they also allow operators like NonZero where only execution determines the final output shape.

For the shape interpreter, I could create a new hand-written implementation of it or run ONNX runtime inference only on the shape graph part of the overall graph.

[miaobin]

To be precise about what the prototype does: we identify the shape subgraph by walking backward from the shape operands of the *Dynamic operators. That walk bottoms out at shape() outputs and build-time constants — it cannot bottom out anywhere else, because the interpreter has no access to tensor data (an input
operand simply evaluates to "unknown"). That root restriction is what keeps the shape subgraph a small subgraph rather than the whole model.

A standalone interpreter evaluates that subgraph on the CPU. What it implements in the POC is the set (Add, Cast, Concat, Gather, Mul, Shape, Size, Slice, Squeeze, Sub, Unsqueeze): arithmetic and structural transforms over small 1-D vectors, plus enough floating-point for cases like floor(dim * scale). It grew from what the models we tested actually needed rather than from a designed list; [symbolic_shape_infer.py](https://github.com/microsoft/onnxruntime/blob/main/onnxruntime/python/tools/symbolic_shape_infer.py#L130) is the wider reference we point at when asking where a required floor should sit.

Worth stressing: this does not partition the graph or move any work off the accelerator. The graph handed to the backend is unchanged — shape() and the *Dynamic operators are emitted as ordinary nodes, and the backend executes the whole graph however it sees fit. The CPU evaluation is a side computation, used for validation and for answering computeShapes().

Two things make it necessary:

1. Allocation ordering. dispatch() takes caller-allocated output MLTensors, and requires each one's descriptor to equal the graph's output descriptor — same data type, same shape, dimension by dimension. An MLTensor's shape is fixed at creation, and dispatch() returns undefined synchronously, so it has no way to hand back tensors it allocated itself. The caller — an application, or a framework such as ORT Web — therefore has to know the output shapes for the given input shapes before it calls. Without computeShapes() it would have to reimplement WebNN's shape inference just to size those allocations.

2. Validation has to happen before work reaches the backend. The renderer is untrusted, so the privileged service re-runs the same inference independently rather than taking the renderer's word for it. And dispatch() has no error-reporting channel: a failure that reaches the backend cannot be reported per-dispatch, only escalated to context loss — in our implementation, a session-run failure currently tears down every WebNN context in the GPU process. So "can these shapes be resolved, and do they match what the caller allocated?" has to be answered up front, on the CPU, before anything is submitted.

[theHamsta]

It grew from what the models we tested actually needed rather than from a designed list;

Could an actual design for WebNN have a similar list of limited operators that might grow over time? It would give devs and implementations some orientation on what is expected to work and allow to make certain assumptions on shape

That subgraphs for shape inference for dynamic operators always originate at shape nodes or other dynamic operators and that their outputs are always the shape inputs of dynamic operators, is something I didn't realize during the read of the document but only when thinking about it over the weekend. It simplifies what dynamic shapes can do. It falls naturally out of this document that computeShape only depends on input shapes. Could it be mentioned in this document?

[miaobin]

Could an actual design for WebNN have a similar list of limited operators that might grow over time?

Yes, that is the intent, and I'll make it explicit. A required minimum set, normative, that later revisions can grow the same way the operator set itself grows. A floor rather than a ceiling: an implementation may resolve more, and one that cannot resolve a chain must say so cleanly rather than guess. That is also what keeps it compatible with @fdwr's earlier point: the list constrains what an implementation must be able to resolve, not what a graph may express.

Could it be mentioned in this document?

Yes, the invariant is currently one sentence near the end of that section, and only one of its two ends is stated. I'll state both, up front:

• every chain ends at the shape argument of a *Dynamic operator — that is what makes it a shape chain, and the only way a computed value can influence a shape;
• every chain begins at a shape() output or a build-time constant — those constants being exactly the ones your shapeConstant proposal would have the caller declare;

Therefore the interpreter never reads tensor data, and computeShapes() is a pure function of the input shapes.

[theHamsta]

Regarding the separations of constants: we could separate MLGraphBuilder.constant from MLGraphBuilder.shape_constant

• computeShape would only be allowed to read from shape_constant but not from a regular constant
• regular inference would be allowed to read both

This would allow to immediately send .constant to the accelerator without needing to wait for the user to construct the full graph and be able to analyze it.

To me it seems that actually computeShape never needs to read the contents of constants unless they are input to one of the rank changing or dynamic operator variants. Those operators basically transform the input chain to their shape argument from the dispatch part of the graph to the computeShape part of the computation.

[theHamsta]

Please feel free to close my comments! If you feel that there is nothing actionable. Most of the comments are mainly to clarify our understanding of this proposal. Please excuse if this generates any kind of noise for you!

[miaobin]

Not noise at all! I really appreciate the thorough read. A number of your comments led directly to changes in the explainer, and your comments are genuinely useful for pressure-testing whether the explanation actually holds up. Thanks for taking the time!

[miaobin]

Hi @theHamsta, thank you for the depth of this review; it changed the document in several places. I'd like to check whether the current state is one you'd be comfortable approving.