From 043ef5c95fad9a19a5563103cd7929633468ab4f Mon Sep 17 00:00:00 2001 From: Wanming Lin Date: Tue, 11 Jun 2024 23:27:14 +0800 Subject: [PATCH] [WebNN EP] Support latest WebNN softmax op (#20827) Latest WebNN softmax supports N-D input and axis parameter. --- js/web/docs/webnn-operators.md | 2 +- .../webnn/builders/impl/softmax_op_builder.cc | 82 ++----------------- 2 files changed, 6 insertions(+), 78 deletions(-) diff --git a/js/web/docs/webnn-operators.md b/js/web/docs/webnn-operators.md index 966c93a85a..7c2f10a034 100644 --- a/js/web/docs/webnn-operators.md +++ b/js/web/docs/webnn-operators.md @@ -80,7 +80,7 @@ operators and the supported opset domain/versions in **WebNN EP** by ONNX Runtim | Softsign | ai.onnx(7+) | softsign | ✗ | ✓ | | | Sin | ai.onnx(7+) | sin | ✗ | ✓ | | | Slice | ai.onnx(7-9, 10, 11-12, 13+) | slice | ✓ | ✓ | Input 'starts', 'ends', 'axes', and 'steps' if present must be a constant, only supports 'steps' value 1 | -| Softmax | ai.onnx(7-10, 11-12, 13+) | softmax | ✓ | ✓ | Only supports input rank >= 2 | +| Softmax | ai.onnx(7-10, 11-12, 13+) | softmax | ✓ | ✓ | | | Split | ai.onnx(7-10, 11-12, 13-17, 18+) | split | ✓ | ✓ | Input 'split' if present should be a constant | | Sqrt | ai.onnx(7-12, 13+) | sqrt | ✓ | ✓ | | | Squeeze | ai.onnx(7-10, 11-12, 13-20, 21+) | reshape | ✓ | ✓ | Input 'axes' if present should be a constant | diff --git a/onnxruntime/core/providers/webnn/builders/impl/softmax_op_builder.cc b/onnxruntime/core/providers/webnn/builders/impl/softmax_op_builder.cc index 283badb98a..95c1dbd518 100644 --- a/onnxruntime/core/providers/webnn/builders/impl/softmax_op_builder.cc +++ b/onnxruntime/core/providers/webnn/builders/impl/softmax_op_builder.cc @@ -33,82 +33,16 @@ Status SoftmaxOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder, const logging::Logger& logger) const { const auto& input_defs = node.InputDefs(); emscripten::val input = model_builder.GetOperand(node.InputDefs()[0]->Name()); - emscripten::val output = emscripten::val::object(); std::vector input_shape; ORT_RETURN_IF_NOT(GetShape(*input_defs[0], input_shape, logger), "Cannot get shape"); const auto input_size = input_shape.size(); + NodeAttrHelper helper(node); - if (node.SinceVersion() < 13) { - int32_t axis = helper.Get("axis", 1); - axis = static_cast(HandleNegativeAxis(axis, input_size)); - // Coerce the input into a 2-dimensional tensor with dimensions [a_0 * ... * a_{k-1}, a_k * ... * a_{n-1}]. - if (input_size != 2) { - int32_t first_dim = static_cast(std::reduce(input_shape.begin(), input_shape.begin() + axis, - 1, std::multiplies())); - int32_t second_dim = static_cast(std::reduce(input_shape.begin() + axis, input_shape.end(), - 1, std::multiplies())); - emscripten::val new_shape = emscripten::val::array(std::vector{first_dim, second_dim}); - input = model_builder.GetBuilder().call("reshape", input, new_shape); - } + const int32_t default_axis = node.SinceVersion() < 13 ? 1 : -1; + int32_t axis = helper.Get("axis", default_axis); + axis = static_cast(HandleNegativeAxis(axis, input_size)); - output = model_builder.GetBuilder().call("softmax", input); - - // Reshape output to the same shape of input. - if (input_size != 2) { - emscripten::val new_shape = emscripten::val::array(); - for (size_t i = 0; i < input_size; i++) { - new_shape.call("push", static_cast(input_shape[i])); - } - output = model_builder.GetBuilder().call("reshape", output, new_shape); - } - } else { - int32_t axis = helper.Get("axis", -1); - axis = static_cast(HandleNegativeAxis(axis, input_size)); - // Wraparound for transpose the target axis to the last. - // WebNN compute the softmax values of the 2-D input tensor along axis 1. - // https://www.w3.org/TR/webnn/#api-mlgraphbuilder-softmax-method - if (axis != static_cast(input_shape.size() - 1)) { - emscripten::val options = emscripten::val::object(); - std::vector permutation(input_shape.size()); - std::iota(permutation.begin(), permutation.end(), 0); - std::rotate(permutation.begin() + axis, permutation.begin() + axis + 1, permutation.end()); - options.set("permutation", emscripten::val::array(permutation)); - input = model_builder.GetBuilder().call("transpose", input, options); - } - // Wraparound for reshape input tensor to 2-D. - if (input_shape.size() != 2) { - uint32_t first_dim = static_cast(std::reduce(input_shape.begin(), input_shape.begin() + axis, - 1, std::multiplies())); - first_dim *= static_cast(std::reduce(input_shape.begin() + axis + 1, input_shape.end(), - 1, std::multiplies())); - uint32_t second_dim = static_cast(input_shape[axis]); - emscripten::val new_shape = emscripten::val::array(std::vector{first_dim, second_dim}); - input = model_builder.GetBuilder().call("reshape", input, new_shape); - } - - output = model_builder.GetBuilder().call("softmax", input); - - // Restore from 2-D to the original shape. - if (input_shape.size() != 2) { - std::vector new_shape; - std::transform(input_shape.begin(), input_shape.begin() + axis, std::back_inserter(new_shape), - [](int64_t dim) -> uint32_t { return static_cast(dim); }); - std::transform(input_shape.begin() + axis + 1, input_shape.end(), std::back_inserter(new_shape), - [](int64_t dim) -> uint32_t { return static_cast(dim); }); - new_shape.push_back(static_cast(input_shape[axis])); - output = model_builder.GetBuilder().call("reshape", - output, emscripten::val::array(new_shape)); - } - // Restore the corresponding axis back to the initial position from the last position. - if (axis != static_cast(input_shape.size() - 1)) { - emscripten::val options = emscripten::val::object(); - std::vector permutation(input_shape.size()); - std::iota(permutation.begin(), permutation.end(), 0); - std::rotate(permutation.rbegin(), permutation.rbegin() + 1, permutation.rend() - axis); - options.set("permutation", emscripten::val::array(permutation)); - output = model_builder.GetBuilder().call("transpose", output, options); - } - } + emscripten::val output = model_builder.GetBuilder().call("softmax", input, axis); model_builder.AddOperand(node.OutputDefs()[0]->Name(), std::move(output)); return Status::OK(); } @@ -123,12 +57,6 @@ bool SoftmaxOpBuilder::IsOpSupportedImpl(const InitializedTensorSet& /* initiali std::vector input_shape; if (!GetShape(*input_defs[0], input_shape, logger)) return false; - const auto input_size = input_shape.size(); - if (input_size < 2) { - LOGS(logger, VERBOSE) << "SoftMax only support input size >= 2d shape, input is " - << input_size << "d shape"; - return false; - } return true; }