[WebNN EP] Add contraints for MatMul (#19713)

### Description
Add constraints to MatMul:
- The input must be at least 2D.
- CPU backend: The input rank must be the same.
- CPU backend: The input shape except for the last two axis must be the
same.



### Motivation and Context
Prevent regression for some models.
This commit is contained in:
zesongw 2024-03-02 08:55:50 +08:00 committed by GitHub
parent a0521f899e
commit de3158e78d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -91,44 +91,33 @@ Status GemmOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder, const N
bool GemmOpBuilder::IsOpSupportedImpl(const InitializedTensorSet& initializers,
const Node& node,
const WebnnDeviceType /* device_type */,
const WebnnDeviceType device_type,
const logging::Logger& logger) const {
(void)initializers;
const auto& op_type = node.OpType();
const auto& input_defs(node.InputDefs());
const size_t a_idx = 0, b_idx = 1, c_idx = 2; // A*B+C
std::vector<int64_t> a_shape;
if (!GetShape(*input_defs[a_idx], a_shape, logger))
return false;
if (Product(a_shape) == 0) {
LOGS(logger, VERBOSE) << "A must be non-empty";
return false;
}
std::vector<int64_t> b_shape;
if (!GetShape(*input_defs[b_idx], b_shape, logger))
return false;
if (Product(b_shape) == 0) {
LOGS(logger, VERBOSE) << "B must be non-empty";
return false;
}
if (op_type == "Gemm") {
std::vector<int64_t> a_shape;
{
if (!GetShape(*input_defs[a_idx], a_shape, logger))
return false;
if (a_shape.size() != 2) {
LOGS(logger, VERBOSE) << "A must be 2D";
return false;
}
if (Product(a_shape) == 0) {
LOGS(logger, VERBOSE) << "A must be non-empty";
return false;
}
}
std::vector<int64_t> b_shape;
{
if (!GetShape(*input_defs[b_idx], b_shape, logger))
return false;
if (b_shape.size() != 2) {
LOGS(logger, VERBOSE) << "B must be 2D";
return false;
}
if (Product(b_shape) == 0) {
LOGS(logger, VERBOSE) << "B must be non-empty";
return false;
}
if (a_shape.size() != 2 || b_shape.size() != 2) {
LOGS(logger, VERBOSE) << "A and B must be 2D for Gemm";
return false;
}
// C of Gemm.
@ -162,6 +151,30 @@ 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;
}
// 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
// TODO: Remove this workaround when Chromium enables broadcast for MatMul on WebNN CPU backend.
if (device_type == WebnnDeviceType::CPU) {
if (a_shape.size() != b_shape.size()) {
LOGS(logger, VERBOSE) << "The rank of two inputs for WebNN CPU backend MatMul must be the same.";
return false;
}
for (size_t i = 0; i < a_shape.size() - 2; i++) {
if (a_shape[i] != b_shape[i]) {
LOGS(logger, VERBOSE) << "WebNN CPU backend can't support broadcasting for MatMul.";
return false;
}
}
}
}
return true;
}