[WebNN EP] Fixed bug in Expand implementation (#21163)

ONNX's Expand supports bidirectionally broadcast, while WebNN's expand
op only supports unidirectionally broadcast. Thus we should calculate
the output shape for 'newShape' input of WebNN's expand op.
This commit is contained in:
Wanming Lin 2024-06-27 23:09:13 +08:00 committed by GitHub
parent a1bbfeb306
commit b49788e68b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 31 additions and 21 deletions

View file

@ -136,21 +136,33 @@ bool IsSupportedDataType(const int32_t data_type,
supported_data_types.end();
}
bool IsValidMultidirectionalBroadcast(std::vector<int64_t>& shape_a,
std::vector<int64_t>& shape_b,
const logging::Logger& logger) {
bool GetBidirectionalBroadcastShape(std::vector<int64_t>& shape_a,
std::vector<int64_t>& shape_b,
std::vector<int64_t>& output_shape) {
size_t size_a = shape_a.size();
size_t size_b = shape_b.size();
size_t smaller_size = std::min(size_a, size_b);
for (size_t i = 0; i < smaller_size; i++) {
size_t larger_size = std::max(size_a, size_b);
output_shape.resize(larger_size);
for (size_t i = 0; i < larger_size; i++) {
// right alignment
size_t axis_a = size_a - i - 1;
size_t axis_b = size_b - i - 1;
// Broadcastable tensors must either have each dimension the same size or equal to one.
if (shape_a[axis_a] != shape_b[axis_b] && shape_a[axis_a] != 1 && shape_b[axis_b] != 1) {
return false;
if (i < smaller_size) {
// Broadcastable tensors must either have each dimension the same size or equal to one.
if (shape_a[axis_a] != shape_b[axis_b] && shape_a[axis_a] != 1 && shape_b[axis_b] != 1) {
return false;
}
output_shape[larger_size - i - 1] = std::max(shape_a[axis_a], shape_b[axis_b]);
} else {
// For the remaining dimensions in the larger tensor, copy the dimension size directly to the output shape.
output_shape[larger_size - i - 1] = (size_a > size_b) ? shape_a[axis_a] : shape_b[axis_b];
}
}
return true;
}

View file

@ -277,9 +277,9 @@ static const std::unordered_set<ONNX_NAMESPACE::TensorProto_DataType> webnn_supp
bool IsSupportedDataType(const int32_t data_type,
const std::unordered_set<ONNX_NAMESPACE::TensorProto_DataType>& supported_data_types);
bool IsValidMultidirectionalBroadcast(std::vector<int64_t>& shape_a,
std::vector<int64_t>& shape_b,
const logging::Logger& logger);
bool GetBidirectionalBroadcastShape(std::vector<int64_t>& shape_a,
std::vector<int64_t>& shape_b,
std::vector<int64_t>& output_shape);
bool SetWebnnDataType(emscripten::val& desc, const int32_t data_type);

View file

@ -44,18 +44,19 @@ Status ExpandOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder,
const auto& input_defs = node.InputDefs();
const auto& initializers(model_builder.GetInitializerTensors());
const auto& shape_tensor = *initializers.at(input_defs[1]->Name());
std::vector<int32_t> new_shape;
std::vector<int64_t> new_shape;
ORT_RETURN_IF_NOT(ReadIntArrayFrom1DTensor(shape_tensor, new_shape, logger), "Cannot get shape.");
emscripten::val input = model_builder.GetOperand(input_defs[0]->Name());
std::vector<int64_t> input_shape;
ORT_RETURN_IF_NOT(GetShape(*input_defs[0], input_shape, logger), "Cannot get input's shape.");
if (new_shape.size() < input_shape.size()) {
// Enlarge new shape to input.rank, right aligned with leading ones
new_shape.insert(new_shape.begin(), input_shape.size() - new_shape.size(), 1);
}
std::vector<int64_t> output_shape;
ORT_RETURN_IF_NOT(GetBidirectionalBroadcastShape(input_shape, new_shape, output_shape), "Cannot get output shape.");
emscripten::val output =
model_builder.GetBuilder().call<emscripten::val>("expand",
input, emscripten::val::array(new_shape));
input,
emscripten::val::array(GetVecUint32FromVecInt64(output_shape)));
model_builder.AddOperand(node.OutputDefs()[0]->Name(), std::move(output));
return Status::OK();
}
@ -95,11 +96,8 @@ bool ExpandOpBuilder::IsOpSupportedImpl(const InitializedTensorSet& initializers
return false;
}
if (new_shape.size() > input_shape.size()) {
LOGS(logger, VERBOSE) << "The size of shape must be less than or equal to the rank of input.";
}
if (!IsValidMultidirectionalBroadcast(input_shape, new_shape, logger)) {
std::vector<int64_t> output_shape;
if (!GetBidirectionalBroadcastShape(input_shape, new_shape, output_shape)) {
LOGS(logger, VERBOSE) << "The input cannot expand to shape " << GetShapeString(new_shape);
return false;
}