microsoft
[WebNN EP] Support Einsum op
Adds support for einsum via WebNN matmul, transpose, reshape, reducesum, identity and element-wise binary ops.
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.)
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.
I will leave for about one week for my research paper submission. Responses may be delayed.
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?
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 https://github.com/microsoft/onnxruntime/blob/1e69b612382205b0588f08d2b808b12e32a50a51/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorHelper.cpp#L1632-L1641, 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.
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?
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 https://github.com/onnx/onnx/pull/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~
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
/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
/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[bot] avatar](https://avatars.githubusercontent.com/in/9426?v=4 "Published by a pub.dev verified publisher"){kind=small}
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
/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
/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
/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
/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
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.
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.
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.