From 1e3cd86d800eea180eaea61d9f72109e1d4197e4 Mon Sep 17 00:00:00 2001 From: shiyi Date: Sat, 28 Sep 2024 05:23:08 +0800 Subject: [PATCH] [WebNN EP] Support LSTM op (#20293) --- js/web/docs/webnn-operators.md | 1 + js/web/test/suite-test-list.jsonc | 6 +- .../core/providers/webnn/builders/helper.h | 1 + .../webnn/builders/impl/lstm_op_builder.cc | 278 ++++++++++++++++++ .../webnn/builders/op_builder_factory.cc | 4 + .../webnn/builders/op_builder_factory.h | 1 + 6 files changed, 288 insertions(+), 3 deletions(-) create mode 100644 onnxruntime/core/providers/webnn/builders/impl/lstm_op_builder.cc diff --git a/js/web/docs/webnn-operators.md b/js/web/docs/webnn-operators.md index 6fd4f9af20..6c50f37527 100644 --- a/js/web/docs/webnn-operators.md +++ b/js/web/docs/webnn-operators.md @@ -53,6 +53,7 @@ operators and the supported opset domain/versions in **WebNN EP** by ONNX Runtim | LessOrEqual | ai.onnx(12-15, 16+) | lesserOrEqual | ✓ | ✓ | | | Log | ai.onnx(7-12, 13+) | log | ✓ | ✓ | | | LpPool | ai.onnx(7-10, 11-17, 18+) | l2Pool2d | ✗ | ✓ | Only supports 4-D input, 2-D 'kernel_shape', 'p' value is 2 | +| LSTM | ai.onnx(7-13, 14-21, 22+) | lstm | ✓ | ✓ | Only supports 'layout' == 0, 'input_forget' == 0. 'clip' is not supported. The activation functions in 'activations' must be one of 'Relu', 'Tanh', 'Sigmoid'. Forward and backward activations must be the same if bidirectional. 'sequence_lens' if present should be constant with values equal to the first dimension length of input 'X' | | MatMul | ai.onnx(7-8, 9-12, 13+) | matmul | ✓ | ✓ | | | Max | ai.onnx(7, 8-11, 12, 13+) | max | ✓ | ✓ | | | MaxPool | ai.onnx(7, 8-9, 10, 11, 12+) | maxPool2d | ✓ | ✓ | Only supports 4-D input, 2-D 'kernel_shape', 'storage_order' != 1, one output | diff --git a/js/web/test/suite-test-list.jsonc b/js/web/test/suite-test-list.jsonc index 5c1e2e27a6..ae708467be 100644 --- a/js/web/test/suite-test-list.jsonc +++ b/js/web/test/suite-test-list.jsonc @@ -1912,9 +1912,9 @@ // "test_lrn_default", // "test_lrn", // // "test_lstm_batchwise", - // // "test_lstm_defaults", - // // "test_lstm_with_initial_bias", - // // "test_lstm_with_peepholes", + "test_lstm_defaults", + "test_lstm_with_initial_bias", + "test_lstm_with_peepholes", "test_matmul_2d", "test_matmul_3d", "test_matmul_4d", diff --git a/onnxruntime/core/providers/webnn/builders/helper.h b/onnxruntime/core/providers/webnn/builders/helper.h index dd4a8acc66..7ba1d18fa1 100644 --- a/onnxruntime/core/providers/webnn/builders/helper.h +++ b/onnxruntime/core/providers/webnn/builders/helper.h @@ -195,6 +195,7 @@ static const InlinedHashMap op_map = { {"LessOrEqual", "lesserOrEqual"}, {"Log", "log"}, {"LpPool", "l2Pool2d"}, + {"LSTM", "lstm"}, {"MatMul", "matmul"}, {"MatMulInteger", "matmulInteger"}, {"Max", "max"}, diff --git a/onnxruntime/core/providers/webnn/builders/impl/lstm_op_builder.cc b/onnxruntime/core/providers/webnn/builders/impl/lstm_op_builder.cc new file mode 100644 index 0000000000..6213b039fb --- /dev/null +++ b/onnxruntime/core/providers/webnn/builders/impl/lstm_op_builder.cc @@ -0,0 +1,278 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Copyright (c) Intel Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/providers/common.h" +#include "core/providers/shared/utils/utils.h" +#include "core/providers/webnn/builders/helper.h" +#include "core/providers/webnn/builders/model_builder.h" +#include "core/providers/webnn/builders/op_builder_factory.h" + +#include "base_op_builder.h" + +namespace onnxruntime::webnn { + +class LstmOpBuilder : public BaseOpBuilder { + // Add operator related. + public: + void AddInitializersToSkip(ModelBuilder& model_builder, const Node& node) const override; + + private: + Status AddToModelBuilderImpl(ModelBuilder& model_builder, const Node& node, + const logging::Logger& logger) const override ORT_MUST_USE_RESULT; + + // Operator support related. + private: + bool IsOpSupportedImpl(const InitializedTensorSet& initializers, const Node& node, + const WebnnDeviceType /*device_type*/, const logging::Logger& logger) const override; + bool HasSupportedInputsImpl(const Node& node, const emscripten::val& wnn_limits, + const logging::Logger& logger) const override; + bool HasSupportedOutputsImpl(const Node& node, const emscripten::val& wnn_limits, + const logging::Logger& logger) const override; +}; + +void LstmOpBuilder::AddInitializersToSkip(ModelBuilder& model_builder, const Node& node) const { + if (node.InputDefs().size() > 4 && node.InputDefs()[4]->Exists()) { + model_builder.AddInitializerToSkip(node.InputDefs()[4]->Name()); // sequence_lens + model_builder.AddInputToSkip(node.InputDefs()[4]->Name()); + } +} + +Status LstmOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder, const Node& node, + const logging::Logger& logger) const { + NodeAttrHelper helper(node); + uint32_t hidden_size = helper.Get("hidden_size", 1); + + const auto& input_defs = node.InputDefs(); + std::vector input_shape; + ORT_RETURN_IF_NOT(GetShape(*input_defs[0], input_shape, logger), "Cannot get input's shape"); + uint32_t steps = static_cast(input_shape[0]); + + emscripten::val input = model_builder.GetOperand(input_defs[0]->Name()); + emscripten::val weight = model_builder.GetOperand(input_defs[1]->Name()); + emscripten::val recurrent_weight = model_builder.GetOperand(input_defs[2]->Name()); + + emscripten::val options = emscripten::val::object(); + options.set("label", node.Name()); + options.set("layout", emscripten::val("iofg")); + + if (input_defs.size() > 3 && input_defs[3]->Exists()) { + emscripten::val bias = model_builder.GetOperand(input_defs[3]->Name()); + emscripten::val split_options = emscripten::val::object(); + split_options.set("axis", 1); + split_options.set("label", node.Name() + "_split"); + // Split it to bias and recurrentBias. + emscripten::val splitted_biases = + model_builder.GetBuilder().call("split", bias, /*splits*/ 2, split_options); + options.set("bias", splitted_biases[0]); + options.set("recurrentBias", splitted_biases[1]); + } + if (input_defs.size() > 5 && input_defs[5]->Exists()) { + options.set("initialHiddenState", model_builder.GetOperand(input_defs[5]->Name())); + } + if (input_defs.size() > 6 && input_defs[6]->Exists()) { + options.set("initialCellState", model_builder.GetOperand(input_defs[6]->Name())); + } + if (input_defs.size() > 7 && input_defs[7]->Exists()) { + options.set("peepholeWeight", model_builder.GetOperand(input_defs[7]->Name())); + } + + std::string direction = helper.Get("direction", "forward"); + if (direction == "forward") { + options.set("direction", emscripten::val("forward")); + } else if (direction == "reverse") { + options.set("direction", emscripten::val("backward")); + } else if (direction == "bidirectional") { + options.set("direction", emscripten::val("both")); + } + + const auto& output_defs = node.OutputDefs(); + bool has_Y = output_defs.size() > 0 && output_defs[0]->Exists(); + bool has_Y_h = output_defs.size() > 1 && output_defs[1]->Exists(); + bool has_Y_c = output_defs.size() > 2 && output_defs[2]->Exists(); + options.set("returnSequence", has_Y); + + if (helper.HasAttr("activations")) { + const auto activations = helper.Get("activations", std::vector{"Sigmoid", "Tanh", "Tanh"}); + emscripten::val opt_activations = emscripten::val::array(); + for (size_t i = 0; i < 3; ++i) { + const std::string& activation = activations[i]; + if (activation == "Relu") { + opt_activations.call("push", emscripten::val("relu")); + } else if (activation == "Sigmoid") { + opt_activations.call("push", emscripten::val("sigmoid")); + } else if (activation == "Tanh") { + opt_activations.call("push", emscripten::val("tanh")); + } + } + + options.set("activations", opt_activations); + } + + emscripten::val outputs = model_builder.GetBuilder().call("lstm", input, weight, recurrent_weight, + steps, hidden_size, options); + + if (has_Y) { + model_builder.AddOperand(output_defs[0]->Name(), outputs[2]); + } + if (has_Y_h) { + model_builder.AddOperand(output_defs[1]->Name(), outputs[0]); + } + if (has_Y_c) { + model_builder.AddOperand(output_defs[2]->Name(), outputs[1]); + } + + return Status::OK(); +} + +bool LstmOpBuilder::IsOpSupportedImpl(const InitializedTensorSet& initializers, const Node& node, + const WebnnDeviceType /*device_type*/, const logging::Logger& logger) const { + const auto& input_defs = node.InputDefs(); + if (input_defs.size() < 3) { + LOGS(logger, ERROR) << "LSTM: input size must be greater than or equal to 3"; + return false; + } + + std::vector input_shape; + if (!GetShape(*input_defs[0], input_shape, logger)) { + LOGS(logger, ERROR) << "Cannot get input's shape"; + return false; + } + int32_t steps = static_cast(input_shape[0]); + + if (input_defs.size() > 4 && input_defs[4]->Exists()) { + if (!Contains(initializers, input_defs[4]->Name())) { + LOGS(logger, ERROR) << "LSTM: sequence_lens must be constant"; + return false; + } + + const auto& sequence_lens_tensor = *initializers.at(input_defs[4]->Name()); + std::vector sequence_lens; + if (!ReadIntArrayFrom1DTensor(sequence_lens_tensor, sequence_lens, logger)) { + LOGS(logger, ERROR) << "Cannot read sequence lens tensor"; + return false; + } + if (std::any_of(sequence_lens.begin(), sequence_lens.end(), + [steps](int32_t lens) -> bool { return steps != lens; })) { + LOGS(logger, ERROR) << "LSTM: every sequence length must be equal to input shape[0]"; + return false; + } + } + + NodeAttrHelper helper(node); + if (helper.HasAttr("activations")) { + const auto activations = helper.Get("activations", std::vector{"Sigmoid", "Tanh", "Tanh"}); + + if (activations.size() >= 6) { + if (activations[0] != activations[3] || activations[1] != activations[4] || activations[2] != activations[5]) { + LOGS(logger, ERROR) << "LSTM: forward and backward activations must be the same"; + return false; + } + } + + const InlinedHashSet supported_activations = {"Relu", "Tanh", "Sigmoid"}; + if (std::any_of(activations.begin(), activations.end(), + [&supported_activations](const std::string& activation) -> bool { + return !supported_activations.contains(activation); + })) { + LOGS(logger, ERROR) << "LSTM: activations must be one of Relu, Tanh, Sigmoid"; + return false; + } + } + + if (helper.Get("clip", std::numeric_limits::max()) != std::numeric_limits::max()) { + LOGS(logger, ERROR) << "LSTM: clip is not supported"; + return false; + } + + if (helper.Get("input_forget", 0) != 0) { + LOGS(logger, ERROR) << "LSTM: input_forget == 1 is not supported"; + return false; + } + + if (helper.Get("layout", 0) != 0) { + LOGS(logger, ERROR) << "LSTM: batchwise (layout == 1) is not supported"; + return false; + } + + return true; +} + +bool LstmOpBuilder::HasSupportedInputsImpl(const Node& node, const emscripten::val& wnn_limits, + const logging::Logger& logger) const { + const auto& input_defs = node.InputDefs(); + const auto& op_type = node.OpType(); + int32_t input0_type = 0; // input data type + int32_t input1_type = 0; // weight data type + int32_t input2_type = 0; // recurrentWeight data type + int32_t input3_type = 0; // bias data type + // input4 sequence_lens is skipped. + int32_t input5_type = 0; // initialHiddenState data type + int32_t input6_type = 0; // initialCellState data type + int32_t input7_type = 0; // peepholeWeight data type + bool has_input3 = input_defs.size() > 3 && input_defs[3]->Exists(); + bool has_input5 = input_defs.size() > 5 && input_defs[5]->Exists(); + bool has_input6 = input_defs.size() > 6 && input_defs[6]->Exists(); + bool has_input7 = input_defs.size() > 7 && input_defs[7]->Exists(); + + if (!GetType(*input_defs[0], input0_type, logger) || + !GetType(*input_defs[1], input1_type, logger) || + !GetType(*input_defs[2], input2_type, logger) || + (has_input3 && !GetType(*input_defs[3], input3_type, logger)) || + (has_input5 && !GetType(*input_defs[5], input5_type, logger)) || + (has_input6 && !GetType(*input_defs[6], input6_type, logger)) || + (has_input7 && !GetType(*input_defs[7], input7_type, logger))) { + return false; + } + + InlinedVector input_types = {input0_type, input1_type, input2_type}; + if (has_input3) { + input_types.push_back(input3_type); + } + if (has_input5) { + input_types.push_back(input5_type); + } + if (has_input6) { + input_types.push_back(input6_type); + } + if (has_input7) { + input_types.push_back(input7_type); + } + if (!AreInputDataTypesSame(op_type, input_types, logger)) { + return false; + } + + return IsDataTypeSupportedByOp(op_type, input0_type, wnn_limits, "input", "X", logger); +} + +bool LstmOpBuilder::HasSupportedOutputsImpl(const Node& node, + const emscripten::val& wnn_limits, + const logging::Logger& logger) const { + const auto& output_defs = node.OutputDefs(); + const auto& op_type = node.OpType(); + int32_t Y_type = 0; + int32_t Y_h_type = 0; + int32_t Y_c_type = 0; + bool has_Y = output_defs.size() > 0 && output_defs[0]->Exists(); + bool has_Y_h = output_defs.size() > 1 && output_defs[1]->Exists(); + bool has_Y_c = output_defs.size() > 2 && output_defs[2]->Exists(); + + if (has_Y && GetType(*output_defs[0], Y_type, logger)) { + return IsDataTypeSupportedByOp(op_type, Y_type, wnn_limits, "outputs", "Y", logger); + } + if (has_Y_h && GetType(*output_defs[1], Y_h_type, logger)) { + return IsDataTypeSupportedByOp(op_type, Y_h_type, wnn_limits, "outputs", "Y_h", logger); + } + if (has_Y_c && GetType(*output_defs[2], Y_c_type, logger)) { + return IsDataTypeSupportedByOp(op_type, Y_c_type, wnn_limits, "outputs", "Y_c", logger); + } + + return false; +} + +void CreateLstmOpBuilder(const std::string& op_type, OpBuilderRegistrations& op_registrations) { + op_registrations.builders.push_back(std::make_unique()); + op_registrations.op_builder_map.emplace(op_type, op_registrations.builders.back().get()); +} + +} // namespace onnxruntime::webnn diff --git a/onnxruntime/core/providers/webnn/builders/op_builder_factory.cc b/onnxruntime/core/providers/webnn/builders/op_builder_factory.cc index 93a2b232a7..9df09af01b 100644 --- a/onnxruntime/core/providers/webnn/builders/op_builder_factory.cc +++ b/onnxruntime/core/providers/webnn/builders/op_builder_factory.cc @@ -121,6 +121,10 @@ static OpBuilderRegistrations CreateOpBuilderRegistrations() { CreateLogicalOpBuilder("Not", op_registrations); } + { // LSTM + CreateLstmOpBuilder("LSTM", op_registrations); + } + { // Max/Min CreateMaxMinOpBuilder("Max", op_registrations); CreateMaxMinOpBuilder("Min", op_registrations); diff --git a/onnxruntime/core/providers/webnn/builders/op_builder_factory.h b/onnxruntime/core/providers/webnn/builders/op_builder_factory.h index 61fe6d936e..398dfc2d3f 100644 --- a/onnxruntime/core/providers/webnn/builders/op_builder_factory.h +++ b/onnxruntime/core/providers/webnn/builders/op_builder_factory.h @@ -34,6 +34,7 @@ void CreateGatherOpBuilder(const std::string& op_type, OpBuilderRegistrations& o void CreateGemmOpBuilder(const std::string& op_type, OpBuilderRegistrations& op_registrations); void CreateGruOpBuilder(const std::string& op_type, OpBuilderRegistrations& op_registrations); void CreateLogicalOpBuilder(const std::string& op_type, OpBuilderRegistrations& op_registrations); +void CreateLstmOpBuilder(const std::string& op_type, OpBuilderRegistrations& op_registrations); void CreateMaxMinOpBuilder(const std::string& op_type, OpBuilderRegistrations& op_registrations); void CreateNormalizationOpBuilder(const std::string& op_type, OpBuilderRegistrations& op_registrations); void CreatePadOpBuilder(const std::string& op_type, OpBuilderRegistrations& op_registrations);