[WebNN EP] Support Einsum op

[peishenyan]

Adds support for einsum via WebNN matmul, transpose, reshape, reducesum, identity and element-wise binary ops.

[Honry]

Thanks @peishenyan, some comments.

[peishenyan]

Thanks @Honry , comments addressed.
@fdwr , @guschmue , PTAL, thanks!

[peishenyan]

Hi, all. Last week, I implemented Einsum based on the DML execution provider's implementation, which only support a few conditions. However, Einsum('bhwc,hkc->bhwk', A, B) cannot be handled in ``Segment Anything Encoder''.

Now, according to Einsum Formula in ML Operator Formulas, I implement pair-wise operand processing for handling two inputs and kept the previously implemented single operand processing (Identity, Transpose and ReduceSum.)

[peishenyan]

Since WebNN does not support Triangular op so far and Diagonal depends on triangular, WebNN will fallback when the Einsum op is used for Diagonal operation.

One more thing, in Segment Anything Encoder model, the output shape of Einsum is not defined, which is the input of some Unsqueeze operation, while WebNN does not support dynamic shape. We need additional work to solve this problem in another work branch.

[peishenyan]

I will leave for about one week for my research paper submission. Responses may be delayed.

[Honry]

Thanks @peishenyan!

@guschmue, @fdwr, PTAL thanks!

One more thing, in Segment Anything Encoder model, the output shape of Einsum is not defined, which is the input of some Unsqueeze operation, while WebNN does not support dynamic shape. We need additional work to solve this problem in another work branch.

@guschmue, @fdwr, the output shape of Einsum is not provided in the ONNX graph even it is a static shape model. You know its output shape could only be calculated via its equation expression and input shape.

So my question is, is that possible to calculate and add its output shape info to the graph during ORT Web graph optimization?

[fdwr]

So my question is, is that possible to calculate and add its output shape info to the graph during ORT Web graph optimization?

Are you saying that ORT is not doing shape inference properly for EinSum? The DML EP (which also needs constant shapes before DML CompileGraph is called) registers its own shape inference function for EinSum here

onnxruntime/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorHelper.cpp

Lines 1632 to 1641 in 1e69b61

	// Generate output dimensions from corresponding input tensor labels.
	// e.g. Given ij,jk->ij with [2,3] and [3,5], the output is [2,5].
	std::vector<uint32_t> outputDimensions;
	auto outputLabelIndices = m_components.back().GetLabels(m_labelIndices);
	for (auto labelIndex : outputLabelIndices)
	{
	outputDimensions.push_back(labelSizes[labelIndex]);
	}
	return { EdgeShapes(outputDimensions) };

, which (if you've already parsed the components) isn't too long/complicated.

Since WebNN does not support Triangular op so far and Diagonal depends on triangular, WebNN will fallback when the Einsum op is used for Diagonal operation.

👍 Lisha is working on it (and I drew some whiteboard diagrams on decomposing it for DML feature level 4). It should be available soon.

... my research paper submission

Thanks for creating this - hope your research paper goes well.

[Honry]

Are you saying that ORT is not doing shape inference properly for EinSum?

Right, ORT doesn't do shape inference for EinSum, we couldn't get the input shape of EinSum node's following nodes from the graph, e.g. in following snapshot, we couldn't get the Usqueeze node's input shape. Thus this op will be fallback to wasm as we do dynamic input shape check at https://github.com/microsoft/onnxruntime/blob/main/onnxruntime/core/providers/webnn/builders/helper.cc#L73.

The DML EP (which also needs constant shapes before DML CompileGraph is called) registers its own shape inference function for EinSum here

Interesting, how and when does DML EP register the shape info to the graph?

[fdwr]

Thanks Peishen. I'm still trying to figure out PairwiseOperandProcess (some more comments/visuals/examples could help, especially future generations), but I reviewed the rest of it.

[peishenyan]

Hi @fdwr , I have addressed most of the comments, but the code is not ready for final review now.
As @Honry said before, the output shape of Einsum is not provided in the ONNX graph. We found that we could use onnx shape inference to achieve almost all output shapes except the output shape of Einsum. Since there was only Rank Inference function for Einsum instead of Shape Inference, we implemented Shape Infernce for Einusm in onnx/onnx#6010 to address this problem, which is still waiting for review and merge.
Furthermore, I am working on supporting parsing ellipsis in equation and diagnoal now~

[peishenyan]

Hi @Honry @fdwr , I have addressed all the comments and we decide to implement the extended features (diagonal and ellipsis parsing) in the future, which will not be included in this PR. This PR is ready for review. PTAL. Thanks

[tianleiwu]

/azp run Windows ARM64 QNN CI Pipeline,Windows x64 QNN CI Pipeline,Windows CPU CI Pipeline,Windows GPU CI Pipeline,Windows GPU TensorRT CI Pipeline,ONNX Runtime Web CI Pipeline,Linux CPU CI Pipeline,Linux CPU Minimal Build E2E CI Pipeline,Linux GPU CI Pipeline,Linux GPU TensorRT CI Pipeline

[tianleiwu]

/azp run Linux OpenVINO CI Pipeline,Linux QNN CI Pipeline,MacOS CI Pipeline,orttraining-amd-gpu-ci-pipeline,orttraining-linux-ci-pipeline,orttraining-linux-gpu-ci-pipeline,orttraining-ortmodule-distributed,onnxruntime-binary-size-checks-ci-pipeline,Big Models

[azure-pipelines]

Azure Pipelines successfully started running 9 pipeline(s).

[azure-pipelines]

Azure Pipelines successfully started running 10 pipeline(s).

[tianleiwu]

Any test can cover the op in WebNN EP?

[Honry]

Any test can cover the op in WebNN EP?

@tianleiwu, good point!

@peishenyan has ever tested the following unit tests, but which are all fp64 data type that currently is not supported by WebNN. Is that possible to change the data type to fp32?
https://github.com/microsoft/onnxruntime/blob/main/js/web/test/suite-test-list.jsonc#L1731-L1735

[fdwr]

TY Peishen. Some more thoughts. I'll make a few trivial edits directly (little typos), and for the renaming, we can go ahead and submit as-is if you can follow up in a small CR to make some of the local variable names cleaner.

[fdwr]

/azp run Windows ARM64 QNN CI Pipeline,Windows x64 QNN CI Pipeline,Windows CPU CI Pipeline,Windows GPU CI Pipeline,Windows GPU TensorRT CI Pipeline,ONNX Runtime Web CI Pipeline,Linux CPU CI Pipeline,Linux CPU Minimal Build E2E CI Pipeline,Linux GPU CI Pipeline,Linux GPU TensorRT CI Pipeline

[fdwr]

/azp run Linux OpenVINO CI Pipeline,Linux QNN CI Pipeline,MacOS CI Pipeline,orttraining-amd-gpu-ci-pipeline,orttraining-linux-ci-pipeline,orttraining-linux-gpu-ci-pipeline,orttraining-ortmodule-distributed,onnxruntime-binary-size-checks-ci-pipeline,Big Models

[azure-pipelines]

Azure Pipelines successfully started running 10 pipeline(s).

[azure-pipelines]

Azure Pipelines successfully started running 9 pipeline(s).

[fdwr]

/azp run Linux Android Emulator QNN CI Pipeline

[azure-pipelines]

Azure Pipelines successfully started running 1 pipeline(s).

[fdwr]

/azp run Windows ARM64 QNN CI Pipeline,Windows x64 QNN CI Pipeline,Windows CPU CI Pipeline,Windows GPU CI Pipeline,Windows GPU TensorRT CI Pipeline,ONNX Runtime Web CI Pipeline,Linux CPU CI Pipeline,Linux CPU Minimal Build E2E CI Pipeline,Linux GPU CI Pipeline,Linux GPU TensorRT CI Pipeline

[azure-pipelines]

Azure Pipelines successfully started running 8 pipeline(s).

[azure-pipelines]

Azure Pipelines successfully started running 10 pipeline(s).

[peishenyan]

Hi @fdwr, do you have any other comment? Could we merge this branch?

[fdwr]

👍

[fdwr]

/azp run ONNX Runtime Web CI Pipeline,Windows GPU CI Pipeline,Linux Android Emulator QNN CI Pipeline

[fdwr]

/azp run Linux CPU CI Pipeline,Linux CPU Minimal Build E2E CI Pipeline,Linux GPU CI Pipeline,Linux GPU TensorRT CI Pipeline,Linux OpenVINO CI Pipeline,Linux QNN CI Pipeline,MacOS CI Pipeline,Windows ARM64 QNN CI Pipeline,Windows CPU CI Pipeline

[fdwr]

/azp run Windows GPU CUDA CI Pipeline,Windows GPU DML CI Pipeline,Windows GPU Doc Gen CI Pipeline

[azure-pipelines]

Azure Pipelines successfully started running 2 pipeline(s).

[fdwr]

/azp run Windows GPU TensorRT CI Pipeline,onnxruntime-binary-size-checks-ci-pipeline,orttraining-linux-ci-pipeline,orttraining-linux-gpu-ci-pipeline,orttraining-ortmodule-distributed,Windows x64 QNN CI Pipeline,Big Models

[azure-pipelines]

Azure Pipelines successfully started running 3 pipeline(s).

[azure-pipelines]

Azure Pipelines successfully started running 4 pipeline(s).

[azure-pipelines]

Azure Pipelines successfully started running 9 pipeline(s).

[peishenyan]

I have updated the code to fix the above comment, please re-trigger the CI tests. Thanks!

The following is the result of failed"Windows GPU DML CI pipeline, but it does not appear to be related to this PR.

1: [ RUN      ] NhwcMaxPoolContribOpTest.MaxPool2D
1: 2024-11-12 06:36:00.9277083 [W:onnxruntime:Default, base_tester.cc:830 onnxruntime::test::BaseTester::ExecuteModelForEps] registered execution providers DmlExecutionProvider  were unable to run the model.
1: 2024-11-12 06:36:01.1277950 [W:onnxruntime:Default, base_tester.cc:830 onnxruntime::test::BaseTester::ExecuteModelForEps] registered execution providers DmlExecutionProvider  were unable to run the model.
1: 2024-11-12 06:36:01.3399722 [W:onnxruntime:Default, base_tester.cc:830 onnxruntime::test::BaseTester::ExecuteModelForEps] registered execution providers DmlExecutionProvider  were unable to run the model.
1: 2024-11-12 06:36:01.5518211 [W:onnxruntime:Default, base_tester.cc:830 onnxruntime::test::BaseTester::ExecuteModelForEps] registered execution providers DmlExecutionProvider  were unable to run the model.
1: D:\a\_work\1\s\onnxruntime\core\providers\dml\dml_provider_factory.cc(567)\onnxruntime_test_all.exe!00007FF62D76D73A: (caller: 00007FF62D76DAE5) Exception(1) tid(1658) 8007000E Not enough memory resources are available to complete this operation.
1: 
1: Provider:DmlExecutionProvider
1: unknown file: error: C++ exception with description "D:\a\_work\1\s\onnxruntime\core\providers\dml\dml_provider_factory.cc(567)\onnxruntime_test_all.exe!00007FF62D76D73A: (caller: 00007FF62D76DAE5) Exception(1) tid(1658) 8007000E Not enough memory resources are available to complete this operation.
1: " thrown in the test body.
1: 
1: [  FAILED  ] NhwcMaxPoolContribOpTest.MaxPool2D (4516 ms)

[fdwr]

The following is the result of failed"Windows GPU DML CI pipeline, but it does not appear to be related to this PR.

Hmm, yeah, definitely unrelated.

[fdwr]

/azp run ONNX Runtime Web CI Pipeline,Windows GPU CI Pipeline,Linux Android Emulator QNN CI Pipeline

[fdwr]

/azp run Linux CPU CI Pipeline,Linux CPU Minimal Build E2E CI Pipeline,Linux GPU CI Pipeline,Linux GPU TensorRT CI Pipeline,Linux OpenVINO CI Pipeline,Linux QNN CI Pipeline,MacOS CI Pipeline,Windows ARM64 QNN CI Pipeline,Windows CPU CI Pipeline

[fdwr]

/azp run Windows GPU CUDA CI Pipeline,Windows GPU DML CI Pipeline,Windows GPU Doc Gen CI Pipeline

[fdwr]

/azp run Windows GPU TensorRT CI Pipeline,onnxruntime-binary-size-checks-ci-pipeline,orttraining-linux-ci-pipeline,orttraining-linux-gpu-ci-pipeline,orttraining-ortmodule-distributed,Windows x64 QNN CI Pipeline,Big Models

[azure-pipelines]

Azure Pipelines successfully started running 2 pipeline(s).

[azure-pipelines]

Azure Pipelines successfully started running 3 pipeline(s).

[azure-pipelines]

Azure Pipelines successfully started running 4 pipeline(s).

[azure-pipelines]

Azure Pipelines successfully started running 9 pipeline(s).

[fdwr]

Rerunning failed CI's.

[fdwr]

I retried the "Linux GPU CI Pipeline", but I still see a failure 🤔.

[fdwr]

Can you remerge with master to include this? #22828