diff --git a/onnxruntime/core/providers/webnn/builders/helper.cc b/onnxruntime/core/providers/webnn/builders/helper.cc index 5d1794ed0a..44e6953db4 100644 --- a/onnxruntime/core/providers/webnn/builders/helper.cc +++ b/onnxruntime/core/providers/webnn/builders/helper.cc @@ -136,21 +136,33 @@ bool IsSupportedDataType(const int32_t data_type, supported_data_types.end(); } -bool IsValidMultidirectionalBroadcast(std::vector& shape_a, - std::vector& shape_b, - const logging::Logger& logger) { +bool GetBidirectionalBroadcastShape(std::vector& shape_a, + std::vector& shape_b, + std::vector& 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; } diff --git a/onnxruntime/core/providers/webnn/builders/helper.h b/onnxruntime/core/providers/webnn/builders/helper.h index 4ee3f891f9..496f886e5a 100644 --- a/onnxruntime/core/providers/webnn/builders/helper.h +++ b/onnxruntime/core/providers/webnn/builders/helper.h @@ -277,9 +277,9 @@ static const std::unordered_set webnn_supp bool IsSupportedDataType(const int32_t data_type, const std::unordered_set& supported_data_types); -bool IsValidMultidirectionalBroadcast(std::vector& shape_a, - std::vector& shape_b, - const logging::Logger& logger); +bool GetBidirectionalBroadcastShape(std::vector& shape_a, + std::vector& shape_b, + std::vector& output_shape); bool SetWebnnDataType(emscripten::val& desc, const int32_t data_type); diff --git a/onnxruntime/core/providers/webnn/builders/impl/expand_op_builder.cc b/onnxruntime/core/providers/webnn/builders/impl/expand_op_builder.cc index 9d4de45fda..9c75c00fa9 100644 --- a/onnxruntime/core/providers/webnn/builders/impl/expand_op_builder.cc +++ b/onnxruntime/core/providers/webnn/builders/impl/expand_op_builder.cc @@ -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 new_shape; + std::vector 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 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 output_shape; + ORT_RETURN_IF_NOT(GetBidirectionalBroadcastShape(input_shape, new_shape, output_shape), "Cannot get output shape."); + emscripten::val output = model_builder.GetBuilder().call("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 output_shape; + if (!GetBidirectionalBroadcastShape(input_shape, new_shape, output_shape)) { LOGS(logger, VERBOSE) << "The input cannot expand to shape " << GetShapeString(new_shape); return false; }