[WebNN EP] Support MatMul 1D (#19862)

### Description
Support MatMul 1D inputs by combining Reshape and ReduceMean.



### Motivation and Context
ONNX MatMul can support 1D inputs, which is disabled in
`IsOpSupportedImpl`.
This commit is contained in:
zesongw 2024-03-20 23:32:57 +08:00 committed by GitHub
parent 6ff31e06d5
commit 7e18cb4c35
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -42,6 +42,26 @@ Status GemmOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder, const N
if (!GetShape(*input_defs[a_idx], a_shape, logger)) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Can not get shape of A.");
}
std::vector<int64_t> b_shape;
if (!GetShape(*input_defs[b_idx], b_shape, logger)) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Can not get shape of B.");
}
// If the first argument is 1-D, it is promoted to a matrix by prepending a 1 to its dimensions.
bool extended_a_shape = false;
if (a_shape.size() == 1) {
extended_a_shape = true;
a_shape.insert(a_shape.begin(), 1);
a = model_builder.GetBuilder().call<emscripten::val>("reshape", a,
emscripten::val::array(GetVecUint32FromVecInt64(a_shape)));
}
// If the second argument is 1-D, it is promoted to a matrix by appending a 1 to its dimensions.
bool extended_b_shape = false;
if (b_shape.size() == 1) {
extended_b_shape = true;
b_shape.push_back(1);
b = model_builder.GetBuilder().call<emscripten::val>("reshape", b,
emscripten::val::array(GetVecUint32FromVecInt64(b_shape)));
}
// The inputs of MatMul must be at least 3D for WebNN CPU backend. Use GEMM for 2D case.
// TODO: Remove this workaround when it is fixed in Chromium.
if (model_builder.GetWebnnDeviceType() == WebnnDeviceType::CPU && a_shape.size() == 2) {
@ -49,6 +69,27 @@ Status GemmOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder, const N
} else {
output = model_builder.GetBuilder().call<emscripten::val>("matmul", a, b);
}
// If the inputs are both 1D reduce the output to a scalar.
if (extended_a_shape && extended_b_shape) {
output = model_builder.GetBuilder().call<emscripten::val>("reshape", output, emscripten::val::array());
}
// After matrix multiplication the prepended 1 is removed.
else if (extended_a_shape) {
std::vector<uint32_t> new_shape;
for (size_t i = 0; i < b_shape.size() - 2; i++) {
new_shape.push_back(narrow<uint32_t>(b_shape[i]));
}
new_shape.push_back(narrow<uint32_t>(b_shape.back()));
output = model_builder.GetBuilder().call<emscripten::val>("reshape", output, emscripten::val::array(new_shape));
}
// After matrix multiplication the appended 1 is removed.
else if (extended_b_shape) {
std::vector<uint32_t> new_shape;
for (size_t i = 0; i < a_shape.size() - 1; i++) {
new_shape.push_back(narrow<uint32_t>(a_shape[i]));
}
output = model_builder.GetBuilder().call<emscripten::val>("reshape", output, emscripten::val::array(new_shape));
}
} else if (op_type == "MatMulInteger") {
emscripten::val a_zero_point = emscripten::val::null();
emscripten::val b_zero_point = emscripten::val::null();
@ -152,10 +193,10 @@ bool GemmOpBuilder::IsOpSupportedImpl(const InitializedTensorSet& initializers,
}
if (op_type == "MatMul") {
if (a_shape.size() < 2 || b_shape.size() < 2) {
LOGS(logger, VERBOSE) << "Inputs of MatMul must be at least 2D";
return false;
}
// If the first argument is 1-D, it is promoted to a matrix by prepending a 1 to its dimensions.
// If the second argument is 1-D, it is promoted to a matrix by appending a 1 to its dimensions.
if (a_shape.size() == 1) a_shape.insert(a_shape.begin(), 1);
if (b_shape.size() == 1) b_shape.push_back(1);
// WebNN CPU backend has two more constraints.
// https://source.chromium.org/chromium/chromium/src/+/main:third_party/blink/renderer/modules/ml/webnn/ml_graph_xnnpack.cc;l=1177