Allow checking whether operators/types are supported for a backend before creating a graph

[huningxin]

This feedback is from @RafaelCintron (Thanks!) when reviewing Chromium CL-4828993 review.

In the current spec and implementation, if any operators/types combination are not supported by a backend, for example DirectML backend doesn't support the dilations of average pooling operator, the errors will be thrown when user code calls MLGraphBuilder.build(). Rafael mentioned this might be too late because user code may already download all the weights.

We should add new WebNN APIs to help developers determine which of these operators/type combinations will fail BEFORE they attempt to create a graph. This will help them avoid large model downloads with operators/types the backend does not support.

[wacky6]

🤔 Or perhaps WebNN needs a weight loading step?

This is essentially the pattern I see with popular models (e.g. Real-ESRGAN, Stable diffusion). The graph describes the architecture and doesn't include weights.

If we want to fail before building the grpah, I'm afraid the operand methods (e.g. conv2d) needs to be async (to allow querying the backend), or make every feature feature-detectable (think of sending a feature matrix to the renderer).

[huningxin]

@wacky6

Or perhaps WebNN needs a weight loading step?

This is a great idea.

According to the current WebNN spec, the WebGPU context has MLCommandEncoder.initializeGraph(graph) method for preparing and optimizing constant input data before subsequent graph executions.

We probably can reuse this method for default context and extend it for weights uploading. For example, MLContext.initializeGraph(graph, weights), it would allow to upload the weights to constant operands after building the graph.

If the constants/weights uploading could be decoupled from graph building, developers can build the graph topology/architecture first before downloading the weights. If there are any unsupported error, they may handle the error or stop there. Only when the graph building succeeds, they proceed to download the weights from network and upload to device by MLContext.initializeGraph().

[wacky6]

Weight loading would substantially change how graphBuilder.constant() works. This will be a non-trivial change to implementation and usage.

We need to match constant nodes with weight buffers. Probably do-able if we ask clients to retain references to constant MLOperand and pass pair<constNode, arrayBuffer> to loadWeights().

The overall usage will be:

// Build graph
const x = graphBuilder.input()
let constNode = graphBuilder.constant(shape)
const y = graphBuilder.add(x, constNode)
graphBuilder.constant()

// Compile the graph on backend
// Will reject if any operators in the graph isn't supported.
const graph = await graphBuilder.build(y)

// Download weight, match them with array buffer
const weightNodes = new Map()
weightNodes.set(constNode, arrayBuffer)   // Note the `constNode` reference here

// Load weight and make graph computable.
away graph.loadWeight(weightNodes, [arrayBuffer1])

// Compute
await MLContext.compute(graph, {...inputs})

I'd recommend including this issue in v2. Diffuser / LLM are both large.

Stable diffusion XL already hit 7GB (fp16). Maybe its size will double again this year :)

[huningxin]

We need to match constant nodes with weight buffers. Probably do-able if we ask clients to retain references to constant MLOperand and pass pair<constNode, arrayBuffer> to loadWeights().

Alternatively, we may reuse MLNamedArrayBufferViews for weights uploading. As for common model formats, e.g., ONNX or TFLite, a weight usually has a name which can be used to build a constant operand when building the graph and to map a buffer when uploading the weights.

The overall usage will be:

The usage looks good. I slightly modified the pseudo code by using MLNamedArrayBufferViews:

// Build graph
const x = graphBuilder.input('x', shape)
let constNode = graphBuilder.constant('constNode', shape)
const y = graphBuilder.add(x, constNode)

// Compile the graph on backend
// Will reject if any operators in the graph isn't supported.
const graph = await graphBuilder.build({y})

// Download weight, match them with array buffer by name
const weights = {'constNode': arrayBuffer}

// Load weight and make graph computable.
await graph.loadWeight(weights)

// Compute
await MLContext.compute(graph, {...inputs})

[reillyeon]

I had this same thought when reviewing Chromium CL 5232604. In that example there are a number of input shapes which aren't supported by XNNPACK. Generating these errors at build time rather than when the MLGraphBuilder.matmul() function is called will make it all but impossible to debug why the build failed.

In the XNNPACK case it is easy enough to do these checks earlier and in the renderer process, but I can see how for backends with a more elaborate build step this becomes much more difficult. Perhaps a build error API could return a reference to the problematic MLOperand so that debug tools can visualize the graph and highlight the problem. This would also aid in frameworks performing feature detection, though as @RafaelCintron pointed out in the original comment, the earlier the incompatibility can be detected the less wasted effort there is.

[philloooo]

what's the expected behavior for the framework when an operator/type is not supported?

If we just want to fail quick, then make weight loading separate is sufficient. But if we want to allow frameworks to adjust the graph partition/ deploy a different model that works for the current webnn backend, then we do need to be able to feature detect in some way to surface the particular unsupported features.

[philloooo]

Hi! Given we are having more and more data type support differences across platforms (e.g. #653, #675, #283 (comment)), I'd like to make more progress on this issue.

What do folks think about this structure that starts with just exposing the data type support level?

context.opSupportLimits():
{
    "input": {
        "dataType": [
            "float16",
            "float32"
        ]
    },
    "constant": {
        "dataType": [
            "float16",
            "float32",
            "int32"
        ]
    },
    "gather": {
        "input": {
            "dataType": [
                "float16",
                "float32"
            ]
        },
        "indices": {
            "dataType": [
                "uint32",
                "int32",
                "int64"
            ]
        }
    },
    "softmax": {
        "input": {
            "dataType": [
                "float16",
                "float32"
            ]
        }
    },
   ...
}

Probing supported feature limits should be tied to a MLContext as it’s particular to the underlying devices.
MLContext.opSupportLimits() returns supported types for all ops with all operand parameters.

In this proposal “dataType” is in a nested field to make it more extensible, so later we can specify other limits like {“rankRange”: [0, 4]}.

This will be generated using the intersection between the WebNN supported types(for chromium, currently defined in graph_validation_utils) and underlying platform’s supported types. Since MLContext creation is async, this information can be gathered at MLContext creation time by querying the backend.

• For input and constant, since they are data operands themselves, the value is an object containing all the limits, currently it’s only the dataType limit.
• For operator methods, the value is an object with all the operand parameters.
• For operator methods with overloads(activation functions), the value is a superset of all the operand parameters. Which for activations are: e.g. {“softmax”: {“input”: {“dataType”: [“float32”, “float16”]}}}

This is also useful even if a context supports all types declared in WebNN spec. As it provides a programmable way for frameworks to check the supported types for all ops.

thougths? @huningxin @fdwr
@Honry interested in your thoughts on whether this makes sense for ONNX WebNN EP.