Support flexible input sizes

[huningxin]

There are models that may require flexible input sizes that are determined at inference time.

Examples include some vision models that may need to work on multiple resolutions, e.g., MODNet (a model for real-time portrait matting), the model accepts input in shape of [batch_size,3,height,width].

Transformers may work on arbitrary input lengths. When using KV cache, they may also need to increase the KV cache length by 1 per each inference.

For speech recognition model, e.g., Whisper, the encoder needs to work on arbitrary input lengths, its input shape is [batch_size, feature_size, encoder_sequence_length]. The decoder (with past) needs to increase the KV cache length by 1 per each inference. i.e. the shape of input named "past_key_values.0.decoder.key" is [batch_size,6,past_decoder_sequence_length,64], the shape of corresponding output named "present.0.decoder.key" is [batch_size,6,past_decoder_sequence_length + 1,64].

For other language models, e.g., Qwen2.5-0.5B-Instruct, the model takes input in shape [batch_size,sequence_length]. Similarly, it also increases KV cache length at inference time, i.e. the shape of input "past_key_values.0.key" is [batch_size,2,past_sequence_length,64], the corresponding output "present.0.key" shape is [batch_size,2,past_sequence_length + 1,64]

Lack of the support for flexible input sizes increases the complexity of using WebNN for those models. For example, apps need to modify the model and fix the input size when compiling the WebNN graph. At the inference time, apps need to resize the image input, or pad the language input to maximum length before passing the input to WebNN graph.

Flexible input sizes are already supported by native frameworks. For example ONNX allows model inputs having dynamic dimension. At inference time, ONNX Runtime allows apps to specify arbitrary values for those dynamic dimensions. CoreML also supports flexible input shape and allows to set multiple size options or size ranges.

/cc @fdwr @xenova

[fdwr]

A few considerations:

• API role - what is the role of WebNN in the stack, and how high-level should it be? There's an analog in graphics API's for retained mode vs immediate mode, more like SVG/WPF vs OpenGL/D3D. WebNN is an API rather than a model file format (where a model file naturally needs these dynamic dimensions to be reused with different input shapes), and so it's been the onus of the caller to resolve these questions before it even reaches WebNN. It's not really a limitation of WebNN itself, but rather in the intermediate layers that translate from higher level representations into the WebNN API. You can imagine a simple layer between a TF/ONNX model that stamps out the explicit shapes, creates a few static WebNN graphs of different sizes, ideally sharing the constants between all instances. So from the API point-of-view, nothing is really missing, but that said, there are other considerations to weigh...
• API impact - a number of operators will be impacted, including input (as the tensor MLOperandDescriptor.shape is potentially incomplete, not just a list of integers), shape (because currently any constructed operator can tell you it's shape, which may now need to return an optional list or list of optionals rather than a concrete list of integers), reshape (because you might need unspecified free dimensions or need operations like flatten/squeeze/unsqueeze), pooling with output size rounding, and possibly others like convTranspose.
• Execution speed of such models - will it improve speed or hinder speed (vs the fixed input size approach offline)? Lazily rebuilding the graph multiple times on demand (either externally or internally) for changing inputs will likely cost more than executing multiple static iterations, but if you could cache multiple instantiations, it could be as-fast/faster.
• Memory sharing - memory sharing of constants isn't really made easier/harder with dynamic shapes because one can always create the constants via the constant(MLTensor tensor) overload so they are shared by the context across multiple built graphs.
• Latency of stages - if the input shapes are not known until dispatch, then suddenly dispatch becomes heavyweight, as memory planning would potentially need to be redone each time inputs change sizes, but dispatch really should be a light go-do-thing-at-light-speed anymore but rather than redoing shape inference and memory planning too.
• Graph construction overhead - how much runtime is spent in just the graph construction phase (the JS calls to create the object)? Allowing a single graph (with dynamic shapes) to be bound to different input shapes to instantiate different instances might save JS call time.
• Internal complexity - If multiple input shapes are passed to dispatch over time, is the previous graph invalidated, or do we keep the last few shapes in case we see the same shape again (lest we execute the whole graph loop again in the LLM)? Shouldn't this be user controlled?

For the above reasons (I'd like dispatch to still be light and for the user to have control over which instantiated shapes to preserve), then if we add dynamic shapes, I'd like an intermediate step WebNN between build and dispatch where full shape computation and memory planning can occur:

1. build graph (may be dynamic or static)
2. ⭐ produce finalized graph with known bound input shapes
3. dispatch graph of static shapes

e.g.

...
const context = await navigator.ml.createContext();
const builder = new MLGraphBuilder(context);
const constant = builder.constant(descriptor, new Float32Array(4).fill(0.2));
...
const potentiallyDynamicShapeGraph = await builder.build({'C': C});

const [inputTensorA, inputTensorB, outputTensorC] = await ...;
context.writeTensor(inputTensorA, new Float32Array(4).fill(1.0));
...

const inputs = {'A': inputTensorA, 'B': inputTensorB};
const outputs = {'C': outputTensorC};

// Finalize shapes                                        ⭐
const knownShapeGraph = await potentiallyDynamicShapeGraph.computeShapes(inputs, outputs);
// knownShapeGraph is just an MLGraph where all inputs/outputs are resolved concretely.
// If all input shapes were already known, then the original MLGraph is already such a graph that can be
// passed directly to dispatch, and calling computeShapes (or some other name) just returns
// the same graph back. Any knownShapeGraph can be held onto by the caller for later iterations,
// and multiple known-shape-graphs of differing input shapes can be held for reuse.

context.dispatch(knownShapeGraph, inputs, outputs);
// Dispatch requires the MLGraph to have well defined shapes on its outputs, otherwise it's an error.

[huningxin]

@fdwr

• API role - ... and so it's been the onus of the caller to resolve these questions before it even reaches WebNN.

ONNX Runtime may pass dynamic shape tensors to its execution providers. WebNN EP currently has to fall operators with dynamic inputs back to Wasm EP.

Also /cc @guschmue

[fdwr]

WebNN EP currently has to fall operators with dynamic inputs back to Wasm EP.

Note the DML EP executes models that support dynamic input shapes (not as efficiently as if you pass the expected shapes up front, but it still runs) by creating the operator on demand or delaying construction of the graph until shape information is known, and also by setting tensor affinity for shapes to CPU-side to avoid any GPU/NPU readback (see here and here and here).

[mtavenrath]

To add some more context to the request, there are three different types of dynamic shapes usually used in neural network models:

• Completely Unknown Sizes: The dimensions of a tensor are entirely undefined before inference.
• Symbolic Sizes: Multiple tensor shape dimensions use the same symbol, which is assigned a concrete value before an inference run begins.
• Tensor-Derived Sizes: A tensor's dimensions are determined by data from another tensor. These dimensions can change during an inference run or be based on the output of a previous run.

More information about these types can be found in the TensorRT Documentation.

Should WebNN support all three types of dynamic shapes?

Also, dynamic shapes can be either bound or unbound. A bound dynamic shape has a specified upper limit (and sometimes a lower limit) for its dimensions.

Since system memory is finite and reallocating memory during each inference run is inefficient, I'd highly recommended that WebNN supports only bound dynamic shapes. This approach allows the backend to fail gracefully during the validation and compilation stages if it lacks sufficient resources or if the underlying hardware cannot handle the specified upper-bound dimensions.

[anssiko]

Related resolution from Kobe F2F:

RESOLUTION: study more backends and do prototyping before more formally specifying solution