Repository navigation
[CoreML] Fall back for ML Program MatMul with a dynamic K on only one input - #34195
Open
Vijay Kapse (vijay-kapse) wants to merge 1 commit into
Open
Vijay Kapse (vijay-kapse) wants to merge 1 commit into
Vijay Kapse (vijay-kapse) wants to merge 1 commit into
Conversation
… input The CoreML EP declares dynamic input dims with a default size of 1, and CoreML runs shape inference on those default shapes when it loads the model. For a MatMul where K is static on one input and dynamic on the other, the default shapes give mismatched K values and the model fails to load with "shapes of x and y are not broadcastable". Leave those MatMuls to the CPU EP. A static K of 1 matches the default, and a constant 2D B (a linear layer) still loads, so both stay on CoreML. Fixes microsoft#34194
|
Azure Pipelines: There may be pipelines that require an authorized user to comment /azp run to run. |
Contributor
There was a problem hiding this comment.
🟡 Changes recommended
The MatMul test repeats the static-A/dynamic-B K configuration, leaving the intended dynamic-A/static-B fallback with a nonconstant B untested.
1 open issue
What changed in this PR
This PR addresses #34194 by routing MatMul operations that can fail CoreML ML Program loading to the CPU execution provider.
Changes:
- Rejects MatMul when K is dynamic on only one input, except when the static K is 1 or B is a constant 2D tensor.
- Adds tests for CPU fallback and supported CoreML cases, checking provider assignment and outputs against CPU results.
- Leaves Gemm and NeuralNetwork-format behavior unchanged.
| File | Description |
|---|---|
onnxruntime/test/providers/coreml/coreml_basic_test.cc |
Adds a dynamic-K test helper and fallback and supported-case tests. |
onnxruntime/core/providers/coreml/builders/impl/gemm_op_builder.cc |
Adds the ML Program MatMul support check and preserves supported exceptions. |
Review effort: Balanced
| } | ||
|
|
||
| TEST(CoreMLExecutionProviderTest, MatMulDynamicKOnAStaticKOnB) { | ||
| RunMatMulDynamicKTest({-1, -1, 96}, {-1, -1, -1}, {1, 5, 96}, {1, 96, 8}, false, ExpectedEPNodeAssignment::None); |
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
In ML Program format, the CoreML EP now leaves a MatMul on the CPU EP when its K dimension is static on one input and dynamic on the other. Two exceptions stay on CoreML: a static K of 1, and a constant 2D B (a linear layer), because both still load. The check is in
GemmOpBuilder::IsOpSupportedImpl.Motivation and Context
Fixes #34194.
ModelBuilder::RegisterModelInputOutputdeclares every dynamic input dim with a default size of 1, and CoreML runs shape inference on the default shapes when it loads the model. ForA[b, 2, t, 96] @ B[p, q, r, s]the defaults are[1, 2, 1, 96] @ [1, 1, 1, 1], so the load fails and session creation fails with:This is the relative-position attention MatMul in piper (VITS) models.
I checked which shapes fail with the 1.31.0 wheel (ML Program, single MatMul, output compared against the CPU EP):
[b,2,t,96][p,q,r,s](input)[b,2,t,k][p,q,96,s](input)[m,k]/[m,96][96,n]/[k,n](input)[b,t,96][b2,k,n](input)[b,2,t,k][1,1,96,9](constant)[b,t,k][1,96,8](constant)[b,2,t,k][p,q,r,s](K dynamic on both)[m,k][1,n](static K = 1)[m,k],[b,t,k],[b,2,t,k][96,8](constant 2D)The new check rejects exactly the rows that fail and keeps the ones that work on CoreML, so no MatMul that loads today moves to the CPU. For a constant 2D B, CoreML's E5 compiler still logs the same shape-inference error, but CoreML falls back to another engine and the results match the CPU EP, so I left that case alone. Gemm is unchanged: its B is always a constant 2D tensor.
Testing
CoreMLExecutionProviderTest.MatMulStaticKOnADynamicKOnBandMatMulDynamicKOnAStaticKOnBcheck that the MatMul falls back and the output matches the CPU EP. Both fail onmainwith the load error above and pass with this change.MatMulDynamicKSupportedCaseschecks that the three supported cases stay on CoreML; it passes both before and after.onnxruntime_provider_test(macOS arm64, Release,--use_coreml).en_US-amy-medium.onnxwith ML Program: with this change and [CoreML] Saturate out-of-range int64 constants to int32 instead of throwing #34193 together, session creation now succeeds. The first run then stops on a separate limitation: aNonZeroin the duration predictor returns an empty tensor, which the CoreML EP rejects as a partition input. That's not part of this change.