[CoreML EP] Add support for PRelu (#11474)

onnxruntime/core/providers/coreml/builders/impl/activation_op_builder.cc

@@ -2,18 +2,24 @@
 // Licensed under the MIT License.
 
 #ifdef __APPLE__
+#include "core/framework/tensorprotoutils.h"
+#include "core/providers/coreml/builders/impl/builder_utils.h"
 #include "core/providers/coreml/builders/model_builder.h"
 #endif
+#include "core/providers/common.h"
+#include "core/providers/coreml/builders/helper.h"
+#include "core/providers/coreml/builders/impl/base_op_builder.h"
 #include "core/providers/coreml/builders/op_builder_factory.h"
 
-#include "base_op_builder.h"
-
 namespace onnxruntime {
 namespace coreml {
 
 class ActivationOpBuilder : public BaseOpBuilder {
   // Add operator related
 #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 ORT_MUST_USE_RESULT;
@@ -21,15 +27,62 @@ class ActivationOpBuilder : public BaseOpBuilder {
 
   // Operator support related
  private:
+  bool IsOpSupportedImpl(const Node& node, const OpBuilderInputParams& input_params,
+                         const logging::Logger& logger) const override;
   int GetMinSupportedOpSet(const Node& node) const override;
 };
 
 // Add operator related
 
 #ifdef __APPLE__
+void ActivationOpBuilder::AddInitializersToSkip(ModelBuilder& model_builder, const Node& node) const {
+  const auto& op_type = node.OpType();
+  const auto& input_defs = node.InputDefs();
+  if (op_type == "PRelu") {
+    // skip slope as it's already embedded as a weight in the coreml layer
+    model_builder.AddInitializerToSkip(input_defs[1]->Name());
+  }
+}
+
+namespace {
+Status AddPReluWeight(ModelBuilder& model_builder, const Node& node,
+                      const logging::Logger& logger,
+                      COREML_SPEC::ActivationPReLU& prelu) {
+  // add slope initializer as alpha weight
+  const auto& slope_tensor = *model_builder.GetInitializerTensors().at(node.InputDefs()[1]->Name());
+  const auto slope_tensor_num_elements = gsl::narrow<size_t>(Product(slope_tensor.dims()));
+  if (slope_tensor_num_elements != 1) {
+    ORT_RETURN_IF_ERROR(CreateCoreMLWeight(*prelu.mutable_alpha(), slope_tensor));
+  } else {
+    // TODO: CoreML crashes with single element slope, hence this special case. Remove when fixed.
+    // https://github.com/apple/coremltools/issues/1488
+
+    // "broadcast" single value by creating a CoreML weight with num_channels copies of it
+    ORT_RETURN_IF_NOT(slope_tensor.data_type() == ONNX_NAMESPACE::TensorProto_DataType_FLOAT,
+                      "slope initializer has unsupported data type: ", slope_tensor.data_type());
+
+    std::vector<int64_t> x_shape;
+    ORT_RETURN_IF_NOT(GetShape(*node.InputDefs()[0], x_shape, logger), "Failed to get shape of X.");
+
+    // assume X has 3 or 4 dimensions, that was checked in IsPReluOpSupported()
+    const auto num_channels = x_shape[x_shape.size() - 3];
+
+    std::vector<uint8_t> unpacked_tensor;
+    ORT_RETURN_IF_ERROR(onnxruntime::utils::UnpackInitializerData(slope_tensor, unpacked_tensor));
+    float value;
+    std::memcpy(&value, unpacked_tensor.data(), sizeof(value));
+
+    auto& weight_values = *prelu.mutable_alpha()->mutable_floatvalue();
+    weight_values.Clear();
+    weight_values.Resize(num_channels, value);
+  }
+  return Status::OK();
+}
+}  // namespace
+
 Status ActivationOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder,
                                                   const Node& node,
-                                                  const logging::Logger& /* logger */) const {
+                                                  const logging::Logger& logger) const {
   std::unique_ptr<COREML_SPEC::NeuralNetworkLayer> layer = CreateNNLayer(model_builder, node);
 
   const auto& op_type(node.OpType());
@@ -39,6 +92,9 @@ Status ActivationOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder,
     layer->mutable_activation()->mutable_tanh();
   } else if (op_type == "Relu") {
     layer->mutable_activation()->mutable_relu();
+  } else if (op_type == "PRelu") {
+    auto* prelu = layer->mutable_activation()->mutable_prelu();
+    ORT_RETURN_IF_ERROR(AddPReluWeight(model_builder, node, logger, *prelu));
   } else {
     return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
                            "ActivationOpBuilder::AddToModelBuilderImpl, unknown op: ", op_type);
@@ -54,6 +110,68 @@ Status ActivationOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder,
 
 // Operator support related
 
+namespace {
+// assumes that node.OpType() == "PRelu"
+bool IsPReluOpSupported(const Node& node, const OpBuilderInputParams& input_params,
+                        const logging::Logger& logger) {
+  const auto& input_defs = node.InputDefs();
+
+  // X input rank must be 3 or 4
+  std::vector<int64_t> x_shape;
+  if (!GetShape(*input_defs[0], x_shape, logger)) {
+    return false;
+  }
+
+  const auto x_rank = x_shape.size();
+  if (x_rank == 3 || x_rank == 4) {
+    LOGS(logger, VERBOSE) << "PRelu 'X' input must have 3 or 4 dimensions, it has " << x_rank << " dimensions";
+    return false;
+  }
+
+  // slope input must be a constant initializer
+  if (!input_params.graph_viewer.IsConstantInitializer(input_defs[1]->Name(), true)) {
+    LOGS(logger, VERBOSE) << "PRelu 'slope' input must be a constant initializer tensor";
+    return false;
+  }
+
+  // slope must either:
+  // - have shape [C, 1, 1]
+  // - have 1 element
+  {
+    std::vector<int64_t> slope_shape;
+    if (!GetShape(*input_defs[1], slope_shape, logger)) {
+      return false;
+    }
+    const bool has_per_channel_slopes =
+        (slope_shape.size() == 3 && std::all_of(slope_shape.begin() + 1, slope_shape.end(),
+                                                [](int64_t dim) { return dim == 1; }));
+    const bool has_single_slope = Product(slope_shape) == 1;
+    if (!has_per_channel_slopes && !has_single_slope) {
+      LOGS(logger, VERBOSE) << "PRelu 'slope' input must either have shape [C, 1, 1] or have a single value";
+      return false;
+    }
+
+    if (has_single_slope && x_shape[x_rank - 3] == 1) {
+      // TODO: CoreML crashes with single element slope, hence this special case. Remove when fixed.
+      // https://github.com/apple/coremltools/issues/1488
+      LOGS(logger, VERBOSE) << "PRelu single 'slope' value in CoreML weight is not supported";
+      return false;
+    }
+  }
+
+  return true;
+}
+}  // namespace
+
+bool ActivationOpBuilder::IsOpSupportedImpl(const Node& node, const OpBuilderInputParams& input_params,
+                                            const logging::Logger& logger) const {
+  const auto& op_type = node.OpType();
+  if (op_type == "PRelu") {
+    return IsPReluOpSupported(node, input_params, logger);
+  }
+  return true;
+}
+
 int ActivationOpBuilder::GetMinSupportedOpSet(const Node& /* node */) const {
   // All ops opset 5- uses consumed_inputs attribute which is not supported for now
   return 6;
@@ -68,6 +186,7 @@ void CreateActivationOpBuilder(const std::string& op_type, OpBuilderRegistration
           "Sigmoid",
           "Tanh",
           "Relu",
+          "PRelu",
       };
 
   op_registrations.builders.push_back(std::make_unique<ActivationOpBuilder>());

onnxruntime/core/providers/coreml/builders/impl/base_op_builder.cc

@@ -21,10 +21,12 @@ bool HasExternalInitializer(const InitializedTensorSet& initializers, const Node
                             const logging::Logger& logger) {
   for (const auto* node_arg : node.InputDefs()) {
     const auto& input_name(node_arg->Name());
-    if (!Contains(initializers, input_name))
+    const auto initializer_it = initializers.find(input_name);
+    if (initializer_it == initializers.end()) {
       continue;
+    }
 
-    const auto& tensor = *initializers.at(input_name);
+    const auto& tensor = *initializer_it->second;
     if (tensor.has_data_location() &&
         tensor.data_location() == ONNX_NAMESPACE::TensorProto_DataLocation_EXTERNAL) {
       LOGS(logger, VERBOSE) << "Initializer [" << input_name
@@ -140,4 +142,4 @@ bool BaseOpBuilder::HasSupportedOpSet(const Node& node,
 }
 
 }  // namespace coreml
-}  // namespace onnxruntime
+}  // namespace onnxruntime

onnxruntime/core/providers/coreml/builders/impl/builder_utils.cc

@@ -111,4 +111,4 @@ common::Status CreateCoreMLWeight(CoreML::Specification::WeightParams& weight,
 }  // namespace coreml
 }  // namespace onnxruntime
 
-#endif
+#endif

onnxruntime/core/providers/coreml/builders/op_builder_factory.cc

@@ -23,6 +23,7 @@ static OpBuilderRegistrations CreateOpBuilderRegistrations() {
     CreateActivationOpBuilder("Sigmoid", op_registrations);
     CreateActivationOpBuilder("Tanh", op_registrations);
     CreateActivationOpBuilder("Relu", op_registrations);
+    CreateActivationOpBuilder("PRelu", op_registrations);
   }
 
   {  // Transpose

onnxruntime/core/providers/coreml/builders/op_builder_factory.h

@@ -14,8 +14,6 @@ struct OpBuilderRegistrations {
 };
 
 // Get the lookup table with IOpBuilder delegates for different onnx operators
-// Note, the lookup table should have same number of entries as the result of CreateOpSupportCheckers()
-// in op_support_checker.h
 const std::unordered_map<std::string, const IOpBuilder*>& GetOpBuilders();
 
 void CreateBinaryOpBuilder(const std::string& op_type, OpBuilderRegistrations& op_registrations);

onnxruntime/test/providers/cpu/activation/activation_op_test.cc

@@ -4,6 +4,7 @@
 #include "activation_op_test.h"
 #include "core/providers/cpu/activation/activations.h"
 #include "test/common/cuda_op_test_utils.h"
+#include "test/common/tensor_op_test_utils.h"
 
 namespace onnxruntime {
 namespace test {
@@ -377,6 +378,7 @@ TEST_F(ActivationOpTest, Selu_GH10726) {
                           [](float x) { return x <= 0 ? gamma * (alpha * exp(x) - alpha) : gamma * x; },
                           {{"alpha", alpha}, {"gamma", gamma}});
 }
+
 TEST_F(ActivationOpTest, PRelu) {
   OpTester test("PRelu");
 
@@ -396,24 +398,31 @@ TEST_F(ActivationOpTest, PRelu) {
 }
 
 TEST_F(ActivationOpTest, PRelu_SingleSlope) {
-  OpTester test("PRelu");
+  auto test = [](bool slope_is_initializer) {
+    SCOPED_TRACE(MakeString("slope_is_initializer: ", slope_is_initializer));
 
-  auto formula = [](float x, float slope) { return x < 0 ? slope * x : x; };
+    OpTester test("PRelu");
 
-  auto inputs = {1.0f, -4.0f, 0.0f, -9.0f};
-  auto slope = 1.5f;
-  std::vector<float> outputs;
-  for (auto& input : inputs)
-    outputs.push_back(formula(input, slope));
+    auto formula = [](float x, float slope) { return x < 0 ? slope * x : x; };
 
-  std::vector<int64_t> dims{2, 2};
-  test.AddInput<float>("X", dims, inputs);
-  test.AddInput<float>("slope", {}, {slope});
-  test.AddOutput<float>("Y", dims, outputs);
-  test.Run();
+    auto inputs = {1.0f, 2.0f, -4.0f, 3.0f, 0.0f, 5.0f, -9.0f, 8.0f};
+    auto slope = 1.5f;
+    std::vector<float> outputs;
+    for (auto& input : inputs)
+      outputs.push_back(formula(input, slope));
+
+    std::vector<int64_t> dims{2, 2, 2};
+    test.AddInput<float>("X", dims, inputs);
+    test.AddInput<float>("slope", {}, {slope}, slope_is_initializer);
+    test.AddOutput<float>("Y", dims, outputs);
+    test.Run();
+  };
+
+  test(true /* slope_is_initializer */);
+  test(false /* slope_is_initializer */);
 }
 
-TEST_F(ActivationOpTest, PRelu_MultiChannel) {
+TEST_F(ActivationOpTest, PRelu_MultiChannel3D) {
   OpTester test("PRelu");
 
   auto formula = [](float x, float slope) { return x < 0 ? slope * x : x; };
@@ -435,6 +444,40 @@ TEST_F(ActivationOpTest, PRelu_MultiChannel) {
   test.Run();
 }
 
+TEST_F(ActivationOpTest, PRelu_MultiChannel4D) {
+  RandomValueGenerator random{2345};
+
+  auto test = [&](bool slope_is_initializer,
+                  int64_t n, int64_t c, int64_t h, int64_t w) {
+    SCOPED_TRACE(MakeString("slope_is_initializer: ", slope_is_initializer,
+                            ", n: ", n, ", c: ", c, ", h: ", h, ", w: ", w));
+
+    OpTester test("PRelu");
+
+    auto formula = [](float x, float slope) { return x < 0 ? slope * x : x; };
+
+    const std::vector<int64_t> x_dims{n, c, h, w};
+    const std::vector<int64_t> slope_dims{c, 1, 1};
+    std::vector<float> inputs = random.Uniform<float>(x_dims, -16.0f, 16.0f);
+    std::vector<float> slopes = random.Uniform<float>(slope_dims, -1.0f, 1.0f);
+    std::vector<float> outputs;
+    for (unsigned i = 0; i < inputs.size(); i++) {
+      outputs.push_back(formula(inputs[i], slopes[i / (h * w) % c]));
+    }
+
+    test.AddInput<float>("X", x_dims, inputs);
+    test.AddInput<float>("slope", slope_dims, slopes, slope_is_initializer);
+    test.AddOutput<float>("Y", x_dims, outputs);
+    test.Run();
+  };
+
+  test(true /* slope_is_initializer */, 5, 4, 3, 2);
+  test(false, 5, 4, 3, 2);
+
+  test(true, 3, 1, 1, 1);
+  test(false, 3, 1, 1, 1);
+}
+
 TEST_F(ActivationOpTest, Softplus) {
   TestActivationOp<float>("Softplus",
                           input_values,