From f316bc57c473f657bb9365a87bdd4c88f8735e3d Mon Sep 17 00:00:00 2001 From: Shukant Pal Date: Wed, 24 May 2023 04:16:59 -0400 Subject: [PATCH] [CoreML EP] Implement Unary & Reduce operators (#15532) ### Description This change is a follow-up to #15327. It adds Unary operators (Sqrt, Reciprocal) and Reduce operators (ReduceSum, ReduceMean). I've tried to follow existing patterns in the code :-) ### Motivation and Context This reduces fragmentation across EPs when using CoreML on macOS, thereby speeding up execution. --------- Co-authored-by: Edward Chen <18449977+edgchen1@users.noreply.github.com> --- .../builders/impl/reduction_op_builder.cc | 127 ++++++++++++++++++ .../coreml/builders/impl/unary_op_builder.cc | 58 ++++++++ .../coreml/builders/op_builder_factory.cc | 11 ++ .../coreml/builders/op_builder_factory.h | 2 + .../cpu/reduction/reduction_ops_test.cc | 38 +++++- .../github/apple/coreml_supported_ops.md | 5 +- 6 files changed, 239 insertions(+), 2 deletions(-) create mode 100644 onnxruntime/core/providers/coreml/builders/impl/reduction_op_builder.cc create mode 100644 onnxruntime/core/providers/coreml/builders/impl/unary_op_builder.cc diff --git a/onnxruntime/core/providers/coreml/builders/impl/reduction_op_builder.cc b/onnxruntime/core/providers/coreml/builders/impl/reduction_op_builder.cc new file mode 100644 index 0000000000..6a2014e795 --- /dev/null +++ b/onnxruntime/core/providers/coreml/builders/impl/reduction_op_builder.cc @@ -0,0 +1,127 @@ +// Copyright (c) Shukant Pal. +// Licensed under the MIT License. + +#include "core/providers/common.h" +#include "core/providers/shared/utils/utils.h" + +#ifdef __APPLE__ +#include "core/providers/coreml/builders/model_builder.h" +#endif +#include "core/providers/coreml/builders/helper.h" +#include "core/providers/coreml/builders/op_builder_factory.h" +#include "core/optimizer/initializer.h" + +#include "base_op_builder.h" + +namespace onnxruntime { +namespace coreml { + +class ReductionOpBuilder : public BaseOpBuilder { +#ifdef __APPLE__ + 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; +#endif + private: + bool IsOpSupportedImpl(const Node& node, const OpBuilderInputParams& input_params, + const logging::Logger& logger) const override; +}; + +#ifdef __APPLE__ +namespace { +template +void AddReductionParams(T* params, const std::vector& axes, bool keepdims, bool noop_with_empty_axes) { + params->set_keepdims(keepdims); + + for (auto& axis : axes) + params->add_axes(axis); + + if (axes.size() == 0 && !noop_with_empty_axes) + params->set_reduceall(true); +} +} // namespace + +void ReductionOpBuilder::AddInitializersToSkip(ModelBuilder& model_builder, const Node& node) const { + const auto& input_defs(node.InputDefs()); + + // We have already embedded the axes into the CoreML layer. + // No need to copy them later to reduce memory consumption. + if (input_defs.size() > 1) + model_builder.AddInitializerToSkip(input_defs[1]->Name()); +} + +Status ReductionOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder, const Node& node, + const logging::Logger& /* logger */) const { + const auto& op_type(node.OpType()); + const auto& input_defs(node.InputDefs()); + const auto& initializers(model_builder.GetInitializerTensors()); + + std::vector axes; + + NodeAttrHelper helper(node); + if (input_defs.size() > 1 && input_defs[1]->Exists()) { + auto& axes_tensor = *initializers.at(input_defs[1]->Name()); + Initializer axes_initializer(axes_tensor); + int64_t* data = axes_initializer.data(); + int64_t size = axes_initializer.size(); + + axes = std::vector(data, data + size); + } else if (helper.HasAttr("axes")) { + axes = helper.Get("axes", std::vector{}); + } + + const bool keepdims = helper.Get("keepdims", 1) != 0; + const bool noop_with_empty_axes = helper.Get("noop_with_empty_axes", 0) != 0; + + std::unique_ptr layer = CreateNNLayer(model_builder, node); + + if (op_type == "ReduceSum") { + AddReductionParams(layer->mutable_reducesum(), axes, keepdims, noop_with_empty_axes); + } else if (op_type == "ReduceMean") { + AddReductionParams(layer->mutable_reducemean(), axes, keepdims, noop_with_empty_axes); + } else { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "ReductionOpBuilder::AddToModelBuilderImpl, unknown op: ", op_type); + } + + *layer->mutable_input()->Add() = node.InputDefs()[0]->Name(); + *layer->mutable_output()->Add() = node.OutputDefs()[0]->Name(); + + model_builder.AddLayer(std::move(layer)); + return Status::OK(); +} +#endif + +bool ReductionOpBuilder::IsOpSupportedImpl(const Node& node, const OpBuilderInputParams& input_params, + const logging::Logger& logger) const { + const auto& input_defs = node.InputDefs(); + + if (input_defs.size() > 1 && input_defs[1]->Exists()) { + const auto& axes_name = input_defs[1]->Name(); + const auto& initializers = input_params.graph_viewer.GetAllInitializedTensors(); + if (!Contains(initializers, axes_name)) { + LOGS(logger, VERBOSE) << "Axes of reduction must be a constant initializer"; + return false; + } + + NodeAttrHelper helper(node); + + if (initializers.at(axes_name)->int64_data_size() == 0 && helper.Get("noop_with_empty_axes", 0) != 0) { + LOGS(logger, VERBOSE) << "CoreML doesn't support noop on empty axes for reduction layers" << std::endl; + return false; + } + } + + return true; +} + +void CreateReductionOpBuilder(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 coreml +} // namespace onnxruntime \ No newline at end of file diff --git a/onnxruntime/core/providers/coreml/builders/impl/unary_op_builder.cc b/onnxruntime/core/providers/coreml/builders/impl/unary_op_builder.cc new file mode 100644 index 0000000000..660755b43c --- /dev/null +++ b/onnxruntime/core/providers/coreml/builders/impl/unary_op_builder.cc @@ -0,0 +1,58 @@ +// Copyright (c) Shukant Pal. +// Licensed under the MIT License. + +#include "core/providers/common.h" + +#ifdef __APPLE__ +#include "core/providers/coreml/builders/model_builder.h" +#endif +#include "core/providers/coreml/builders/helper.h" +#include "core/providers/coreml/builders/op_builder_factory.h" + +#include "base_op_builder.h" + +namespace onnxruntime { +namespace coreml { + +class UnaryOpBuilder : public BaseOpBuilder { + private: +#ifdef __APPLE__ + Status AddToModelBuilderImpl(ModelBuilder& model_builder, const Node& node, + const logging::Logger& logger) const override; +#endif +}; + +#ifdef __APPLE__ +Status UnaryOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder, const Node& node, + const logging::Logger& /* logger */) const { + const auto& op_type(node.OpType()); + const auto& input_defs(node.InputDefs()); + + std::unique_ptr layer = CreateNNLayer(model_builder, node); + + if (op_type == "Sqrt") { + layer->mutable_unary()->set_type(COREML_SPEC::UnaryFunctionLayerParams::SQRT); + } else if (op_type == "Reciprocal") { + layer->mutable_unary()->set_type(COREML_SPEC::UnaryFunctionLayerParams::INVERSE); + } else { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "UnaryOpBuilder::AddToModelBuilderImpl, unknown op: ", op_type); + } + + *layer->mutable_input()->Add() = input_defs[0]->Name(); + *layer->mutable_output()->Add() = node.OutputDefs()[0]->Name(); + + model_builder.AddLayer(std::move(layer)); + return Status::OK(); +} +#endif + +// Operator support related + +void CreateUnaryOpBuilder(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 coreml +} // namespace onnxruntime \ No newline at end of file diff --git a/onnxruntime/core/providers/coreml/builders/op_builder_factory.cc b/onnxruntime/core/providers/coreml/builders/op_builder_factory.cc index 317d0116bd..0ba83cb51d 100644 --- a/onnxruntime/core/providers/coreml/builders/op_builder_factory.cc +++ b/onnxruntime/core/providers/coreml/builders/op_builder_factory.cc @@ -99,6 +99,17 @@ static OpBuilderRegistrations CreateOpBuilderRegistrations() { CreatePadOpBuilder("Pad", op_registrations); } + { // Unary + CreateUnaryOpBuilder("Sqrt", op_registrations); + CreateUnaryOpBuilder("Reciprocal", op_registrations); + } + + { // Reduction + // ReduceMean is used in layer normalization which seems to be problematic in Python tests. + CreateReductionOpBuilder("ReduceMean", op_registrations); + CreateReductionOpBuilder("ReduceSum", op_registrations); + } + return op_registrations; } diff --git a/onnxruntime/core/providers/coreml/builders/op_builder_factory.h b/onnxruntime/core/providers/coreml/builders/op_builder_factory.h index c129a5c38f..63d1637ddb 100644 --- a/onnxruntime/core/providers/coreml/builders/op_builder_factory.h +++ b/onnxruntime/core/providers/coreml/builders/op_builder_factory.h @@ -35,6 +35,8 @@ void CreateCastOpBuilder(const std::string& op_type, OpBuilderRegistrations& op_ void CreateFlattenOpBuilder(const std::string& op_type, OpBuilderRegistrations& op_registrations); void CreateLRNOpBuilder(const std::string& op_type, OpBuilderRegistrations& op_registrations); void CreatePadOpBuilder(const std::string& op_type, OpBuilderRegistrations& op_registrations); +void CreateUnaryOpBuilder(const std::string& op_type, OpBuilderRegistrations& op_registrations); +void CreateReductionOpBuilder(const std::string& op_type, OpBuilderRegistrations& op_registrations); } // namespace coreml } // namespace onnxruntime diff --git a/onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc b/onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc index d8841b4f5d..afc70b1a5c 100644 --- a/onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc +++ b/onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc @@ -1547,6 +1547,40 @@ TEST(ReductionOpTest, ReduceMean_int32) { test.Run(); } +TEST(ReductionOpTest, ReduceMean_axes_input) { + OpTester test("ReduceMean", 18, onnxruntime::kOnnxDomain); + test.AddAttribute("keepdims", (int64_t)1); + test.AddInput("data", {3, 2, 2}, + {1, 2, + 3, 4, + + 5, 6, + 7, 8, + + 9, 10, + 11, 12}); + test.AddInput("axes", {2}, std::vector{0, 2}, true); + test.AddOutput("reduced", {1, 2, 1}, {5.5, 7.5}); + + // TODO: DNNL, TensorRT, and OpenVINO dont support "axes" input in opset 18, re-enable after + test.Run(OpTester::ExpectResult::kExpectSuccess, "", + {kDnnlExecutionProvider, kTensorrtExecutionProvider, kOpenVINOExecutionProvider, kDmlExecutionProvider}); +} + +TEST(ReductionOpTest, ReduceMean_do_not_keepdims_axes_input_initializer) { + OpTester test("ReduceMean", 18, onnxruntime::kOnnxDomain); + test.AddAttribute("keepdims", (int64_t)0); + test.AddInput("data", {1, 2, 2}, + {1.0f, 2.0f, + 3.0f, 4.0f}); + test.AddInput("axes", {1}, std::vector{1}, true); + test.AddOutput("reduced", {1, 2}, {2.0f, 3.0f}); + + // TODO: DNNL, TensorRT, and OpenVINO dont support "axes" input in opset 18, re-enable after + test.Run(OpTester::ExpectResult::kExpectSuccess, "", + {kDnnlExecutionProvider, kTensorrtExecutionProvider, kOpenVINOExecutionProvider, kDmlExecutionProvider}); +} + TEST(ReductionOpTest, ReduceMean0DTensor) { OpTester test("ReduceMean"); test.AddInput("data", {}, {2}); @@ -4838,7 +4872,9 @@ TEST(ReductionOpTest, ReduceSum_RK_parallel) { } } test.AddOutput("reduced", {32}, expected); - test.Run(); + + // CoreML does not provide 1e-5 precision here (it's off by 1e-4) + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kCoreMLExecutionProvider}); } TEST(ReductionOpTest, ReduceSum_RK_keepdims) { diff --git a/tools/ci_build/github/apple/coreml_supported_ops.md b/tools/ci_build/github/apple/coreml_supported_ops.md index cfd0901bc2..daa607727d 100644 --- a/tools/ci_build/github/apple/coreml_supported_ops.md +++ b/tools/ci_build/github/apple/coreml_supported_ops.md @@ -25,11 +25,14 @@ Keep in sync with doco generated from /docs/execution-providers/CoreML-Execution |ai.onnx:Pad|Only constant mode and last two dim padding is supported.
Input pads and constant_value should be constant.
If provided, axes should be constant.| |ai.onnx:Pow|Only supports cases when both inputs are fp32.| |ai.onnx:PRelu|Input slope should be constant.
Input slope should either have shape [C, 1, 1] or have 1 element.| +|ai.onnx:Reciprocal|| +|ai.onnx.ReduceSum|| |ai.onnx:Relu|| |ai.onnx:Reshape|| |ai.onnx:Resize|| |ai.onnx:Sigmoid|| |ai.onnx:Squeeze|| +|ai.onnx:Sqrt|| |ai.onnx:Sub|| |ai.onnx:Tanh|| -|ai.onnx:Transpose|| +|ai.onnx:Transpose|| \ No newline at end of file