[WebNN EP] Support Einsum op #19558

peishenyan · 2024-02-19T07:49:48Z

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

Honry

Thanks @peishenyan, some comments.

onnxruntime/core/providers/webnn/builders/impl/einsum_op_builder.cc

peishenyan · 2024-02-19T09:01:12Z

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

peishenyan · 2024-02-26T12:09:27Z

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 · 2024-02-26T12:28:14Z

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 · 2024-02-26T12:33:04Z

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

Honry · 2024-02-27T00:50:32Z

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?

onnxruntime/core/providers/webnn/builders/impl/einsum_op_builder.cc

fdwr · 2024-02-27T08:51:04Z

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 · 2024-02-27T15:04:21Z

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?

onnxruntime/core/providers/webnn/builders/impl/einsum_op_builder.cc

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.

onnxruntime/core/providers/webnn/builders/impl/einsum_op_builder.cc

peishenyan · 2024-03-12T09:10:24Z

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~

onnxruntime/core/providers/webnn/builders/impl/einsum_op_builder.cc

peishenyan · 2024-03-26T12:05:08Z

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 · 2024-04-05T13:17:13Z

/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 · 2024-04-05T13:17:14Z

/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 · 2024-04-05T13:17:53Z

Azure Pipelines successfully started running 9 pipeline(s).

azure-pipelines · 2024-04-05T13:17:53Z

Azure Pipelines successfully started running 10 pipeline(s).

tianleiwu · 2024-04-06T02:51:55Z

Any test can cover the op in WebNN EP?

Honry · 2024-04-06T14:06:18Z

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.

onnxruntime/core/providers/webnn/builders/impl/einsum_op_builder.cc

fdwr · 2024-06-21T22:00:46Z

onnxruntime/core/providers/webnn/builders/impl/einsum_op_builder.cc

+
+  std::vector<uint32_t> sequence_o(output_indices.size());
+  std::iota(sequence_o.begin(), sequence_o.end(), 0);
+  if (v != sequence_o) {


if (output_permutation != sequence_o) {

We should use a clearer name than just "v". Does output_permutation make sense here? With the exception of really common cases (like i for loop counters and x,y,z for coordinates), single letter variable names are best avoided. Similarly s and t are unclear what they stand for?

onnxruntime/core/providers/webnn/builders/impl/einsum_op_builder.cc

fdwr · 2024-06-21T22:13:59Z

/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 · 2024-06-21T22:14:08Z

/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 · 2024-06-21T22:14:32Z

Azure Pipelines successfully started running 10 pipeline(s).

azure-pipelines · 2024-06-21T22:14:47Z

Azure Pipelines successfully started running 9 pipeline(s).

fdwr · 2024-06-21T22:15:13Z

/azp run Linux Android Emulator QNN CI Pipeline

azure-pipelines · 2024-06-21T22:15:23Z

Azure Pipelines successfully started running 1 pipeline(s).

fdwr · 2024-06-25T02:03:00Z

/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 · 2024-06-25T02:03:05Z

/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 · 2024-06-25T02:03:36Z

Azure Pipelines successfully started running 10 pipeline(s).

azure-pipelines · 2024-06-25T02:03:38Z

Azure Pipelines successfully started running 9 pipeline(s).

fdwr · 2024-06-25T02:04:04Z

/azp run Linux Android Emulator QNN CI Pipeline

azure-pipelines · 2024-06-25T02:04:12Z

Azure Pipelines successfully started running 1 pipeline(s).

peishenyan · 2024-06-25T02:30:58Z

Hi @fdwr, long time no see. Really appreciate your help. I will address the comments today~
BTW, I remember a few months ago when I ran ORT WebNN EP without onnx shape inference, it aborted because the output could not be registered (no output shape). I'm not sure if this branch can be merged before onnx 1.17 is ready, and I will confirm this today.

peishenyan · 2024-06-25T16:40:47Z

Hi all, I have addressed the comments in the latest commit.
However, without onnx v1.17 the ORT with WebNN EP will crash because the output of einsum op is not registered correctly. This branch should not be merged until onnx v1.17 is ready. Thanks.

fdwr · 2024-06-25T20:56:19Z

This branch should not be merged until onnx v1.17 is ready.

Thanks for confirmation. Ok, will defer merging. So then for any models that need the shape inference logic (like the Segment Anything encoder https://github.com/microsoft/webnn-developer-preview/blob/main/demos/segment-anything/index.js#L739), we'll need a custom build of ORT and copy the .wasm files onto the server.

fdwr · 2024-10-03T04:08:51Z

Relevant! #21897

Address comments register einsum op add logs for einsum Fully implement Einsum op for WebNN EP Address comments address comments fix bugs fix bugs add some comments and fix einsum_type error lint and apply the fixes address comments add TODO and fix identity case remove useless code fix permutation error Apply suggestions from code review Typos and minor declaration. address comments add datatype check and modify IsOpSupported check add diagonal and fix some bugs

fix bugs

Honry · 2024-10-11T01:28:24Z

#22376 this is landed!

@fdwr, Peishen has resolved the conflict and applied latest OpSupportLimits feature, pls. take another look, thanks!

tianleiwu · 2024-10-11T04:15:22Z

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

tianleiwu · 2024-10-11T04:15:23Z

/azp run Linux GPU CI Pipeline,Linux GPU TensorRT CI Pipeline,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

tianleiwu · 2024-10-11T04:15:24Z

/azp run Big Models,Linux Android Emulator QNN CI Pipeline,Android CI Pipeline,iOS CI Pipeline,ONNX Runtime React Native CI Pipeline

azure-pipelines · 2024-10-11T04:15:48Z

Azure Pipelines successfully started running 5 pipeline(s).

azure-pipelines · 2024-10-11T04:15:55Z

Azure Pipelines successfully started running 8 pipeline(s).

azure-pipelines · 2024-10-11T04:15:57Z

Azure Pipelines successfully started running 10 pipeline(s).

peishenyan · 2024-11-07T05:53:44Z

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

Honry reviewed Feb 19, 2024

View reviewed changes

Honry reviewed Feb 27, 2024

View reviewed changes

onnxruntime/core/providers/webnn/builders/impl/einsum_op_builder.cc Outdated Show resolved Hide resolved

guschmue added the ep:WebNN WebNN execution provider label Feb 27, 2024

Honry reviewed Feb 29, 2024

View reviewed changes

onnxruntime/core/providers/webnn/builders/impl/einsum_op_builder.cc Outdated Show resolved Hide resolved

fdwr reviewed Mar 5, 2024

View reviewed changes

github-advanced-security bot found potential problems Mar 13, 2024

View reviewed changes

onnxruntime/core/providers/webnn/builders/impl/einsum_op_builder.cc Fixed Show fixed Hide fixed

fdwr reviewed Mar 14, 2024

View reviewed changes

onnxruntime/core/providers/webnn/builders/impl/einsum_op_builder.cc Outdated Show resolved Hide resolved

fdwr reviewed Jun 21, 2024

View reviewed changes

peishenyan force-pushed the webnn_einsum_op branch 2 times, most recently from 525893c to c5ad5c3 Compare August 2, 2024 08:11

fdwr mentioned this pull request Aug 15, 2024

Support for transformers webmachinelearning/webnn#375

Open

fdwr mentioned this pull request Oct 10, 2024

Pick changes from onnx/onnx#6010 to support EinSum shape inference #22376

Merged

peishenyan added 3 commits October 10, 2024 15:30

add einsum options for webnn and update docs

152f2d8

fix bugs

Resolve conflicts and use latest OpSupportLimits feature

af7d5f7

peishenyan force-pushed the webnn_einsum_op branch from c5ad5c3 to af7d5f7 Compare October 10, 2024 13:55

Merge branch 'microsoft:main' into webnn_einsum_op

644c741

[WebNN EP] Support Einsum op #19558

Are you sure you want to change the base?

[WebNN EP] Support Einsum op #19558

Conversation

peishenyan commented Feb 19, 2024

Honry left a comment

Choose a reason for hiding this comment

peishenyan commented Feb 19, 2024

peishenyan commented Feb 26, 2024 • edited Loading

peishenyan commented Feb 26, 2024

peishenyan commented Feb 26, 2024

Honry commented Feb 27, 2024

fdwr commented Feb 27, 2024

Honry commented Feb 27, 2024

fdwr left a comment

Choose a reason for hiding this comment

peishenyan commented Mar 12, 2024

peishenyan commented Mar 26, 2024

tianleiwu commented Apr 5, 2024

tianleiwu commented Apr 5, 2024

azure-pipelines bot commented Apr 5, 2024

azure-pipelines bot commented Apr 5, 2024

tianleiwu commented Apr 6, 2024

Honry commented Apr 6, 2024

fdwr left a comment

Choose a reason for hiding this comment

fdwr Jun 21, 2024 • edited Loading

Choose a reason for hiding this comment

fdwr commented Jun 21, 2024

fdwr commented Jun 21, 2024

azure-pipelines bot commented Jun 21, 2024

azure-pipelines bot commented Jun 21, 2024

fdwr commented Jun 21, 2024

azure-pipelines bot commented Jun 21, 2024

fdwr commented Jun 25, 2024

fdwr commented Jun 25, 2024

azure-pipelines bot commented Jun 25, 2024

azure-pipelines bot commented Jun 25, 2024

fdwr commented Jun 25, 2024

azure-pipelines bot commented Jun 25, 2024

peishenyan commented Jun 25, 2024

peishenyan commented Jun 25, 2024

fdwr commented Jun 25, 2024 • edited Loading

fdwr commented Oct 3, 2024

Honry commented Oct 11, 2024

tianleiwu commented Oct 11, 2024

tianleiwu commented Oct 11, 2024

tianleiwu commented Oct 11, 2024

azure-pipelines bot commented Oct 11, 2024

azure-pipelines bot commented Oct 11, 2024

azure-pipelines bot commented Oct 11, 2024

peishenyan commented Nov 7, 2024

peishenyan commented Feb 26, 2024 •

edited

Loading

fdwr Jun 21, 2024 •

edited

Loading

fdwr commented Jun 25, 2024 •

edited

Loading