From 9c6481fa2def0a0f78a0cf11acae9b5732d7c828 Mon Sep 17 00:00:00 2001 From: Wanming Lin Date: Tue, 4 Jun 2024 05:12:11 +0800 Subject: [PATCH] [WebNN EP] Enable ArgMax and ArgMin for CPU backend (#20865) WebNN TFLite backend supports ArgMax and ArgMin, but only supports 'select_last_index' value is 0. --- js/web/docs/webnn-operators.md | 4 ++-- onnxruntime/core/providers/webnn/builders/helper.h | 4 ++-- .../webnn/builders/impl/argmax_min_op_builder.cc | 13 +++++++++++-- 3 files changed, 15 insertions(+), 6 deletions(-) diff --git a/js/web/docs/webnn-operators.md b/js/web/docs/webnn-operators.md index 7ec3eda8ce..1df40b71a0 100644 --- a/js/web/docs/webnn-operators.md +++ b/js/web/docs/webnn-operators.md @@ -13,8 +13,8 @@ operators and the supported opset domain/versions in **WebNN EP** by ONNX Runtim |:------:|:------:|:------:|:-:|:-:|:------| | Abs | ai.onnx(7-12, 13+) | abs | ✓ | ✓ | | | Add | ai.onnx(7-12, 13, 14+) | add | ✓ | ✓ | | -| ArgMax | ai.onnx(7-10, 11, 12, 13+) | argMax | ✗ | ✓ | | -| ArgMin | ai.onnx(7-10, 11, 12, 13+) | argMin | ✗ | ✓ | | +| ArgMax | ai.onnx(7-10, 11, 12, 13+) | argMax | ✓ | ✓ | WebNN CPU backend only supports 'select_last_index' value is 0 | +| ArgMin | ai.onnx(7-10, 11, 12, 13+) | argMin | ✓ | ✓ | WebNN CPU backend only supports 'select_last_index' value is 0 | | AveragePool | ai.onnx(7-9, 10, 11, 12-18, 19+) | averagePool2d | ✓ | ✓ | Only supports 4-D input, 2-D 'kernel_shape', 'count_include_pad' value is 0 | | BatchNormalization | ai.onnx(7-8, 9-13, 14, 15+) | batchNormalization | ✗ | ✓ | Only supports 'training_mode' value is 0, one output | | Cast | ai.onnx(7-8, 9-12, 13-18, 19-20, 21+) | cast | ✗ | ✓ | | diff --git a/onnxruntime/core/providers/webnn/builders/helper.h b/onnxruntime/core/providers/webnn/builders/helper.h index 486f7f69be..7c84a14ba7 100644 --- a/onnxruntime/core/providers/webnn/builders/helper.h +++ b/onnxruntime/core/providers/webnn/builders/helper.h @@ -157,8 +157,8 @@ std::vector> GetSupportedNodes(const GraphViewer& graph_v static const InlinedHashMap op_map = { {"Abs", {"abs", true}}, {"Add", {"add", true}}, - {"ArgMax", {"argMax", false}}, - {"ArgMin", {"argMin", false}}, + {"ArgMax", {"argMax", true}}, + {"ArgMin", {"argMin", true}}, {"AveragePool", {"averagePool2d", true}}, {"BatchNormalization", {"batchNormalization", false}}, {"Cast", {"cast", false}}, diff --git a/onnxruntime/core/providers/webnn/builders/impl/argmax_min_op_builder.cc b/onnxruntime/core/providers/webnn/builders/impl/argmax_min_op_builder.cc index 7926311f3c..f8b77b6350 100644 --- a/onnxruntime/core/providers/webnn/builders/impl/argmax_min_op_builder.cc +++ b/onnxruntime/core/providers/webnn/builders/impl/argmax_min_op_builder.cc @@ -21,7 +21,7 @@ class ArgMaxMinOpBuilder : public BaseOpBuilder { // Operator support related. bool IsOpSupportedImpl(const InitializedTensorSet& initializers, const Node& node, - WebnnDeviceType /* device_type */, const logging::Logger& logger) const override; + WebnnDeviceType device_type, const logging::Logger& logger) const override; bool HasSupportedInputsImpl(const Node& node, const WebnnDeviceType device_type, const logging::Logger& logger) const override; }; @@ -68,7 +68,7 @@ Status ArgMaxMinOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder, // Operator support related. bool ArgMaxMinOpBuilder::IsOpSupportedImpl(const InitializedTensorSet& /* initializers */, const Node& node, - WebnnDeviceType /* device_type */, + WebnnDeviceType device_type, const logging::Logger& logger) const { const auto& input_defs = node.InputDefs(); @@ -76,6 +76,15 @@ bool ArgMaxMinOpBuilder::IsOpSupportedImpl(const InitializedTensorSet& /* initia if (!GetShape(*input_defs[0], input_shape, logger)) return false; + // WebNN CPU backend only supports select_last_index = 0. + if (device_type == WebnnDeviceType::CPU) { + NodeAttrHelper helper(node); + const auto select_last_index = helper.Get("select_last_index", 0); + if (select_last_index) { + LOGS(logger, VERBOSE) << "ArgMax/ArgMin with select_last_index = 1 is not supported on WebNN CPU backend."; + return false; + } + } return true; }