Skip to content

Commit ddd4b1e

Browse files
wejoncywejoncy
committed
address comments
1 parent 4c710c7 commit ddd4b1e

File tree

7 files changed

+104
-97
lines changed

7 files changed

+104
-97
lines changed

‎onnxruntime/core/providers/coreml/builders/impl/binary_op_builder.cc‎

Lines changed: 40 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,44 @@ bool CheckIfBothInputShapesMatch(const Node& node, const logging::Logger& logger
5656
}
5757
} // namespace
5858

59+
// Add variadic inputs to the model builder
60+
// in onnx spec, some node allows variadic inputs, such as max(x, y, z, ...)
61+
// while in coreml, maximum op only allows two inputs maximum(x, y)
62+
// the conversion is doing the following:
63+
// max(x, y, z, ...) -> max(max(x, y), z, ...)
64+
#if defined(COREML_ENABLE_MLPROGRAM)
65+
static void AddVariadicInputsToMB(std::unique_ptr<CoreML::Specification::MILSpec::Operation>* op, ModelBuilder& model_builder, const Node& node,
66+
const logging::Logger& logger) {
67+
using namespace CoreML::Specification::MILSpec;
68+
const auto& input_defs(node.InputDefs());
69+
std::string_view layer_input_name_x = model_builder.GetUniqueName(node, "variadic");
70+
auto input_dtype = input_defs[0]->TypeAsProto()->tensor_type().elem_type();
71+
const int32_t elem_type = static_cast<int32_t>(input_dtype);
72+
std::vector<int64_t> x0_shape, x1_shape;
73+
GetShape(*input_defs[0], x0_shape, logger);
74+
GetShape(*input_defs[1], x1_shape, logger);
75+
if (x0_shape.size() < x1_shape.size()) {
76+
std::swap(x0_shape, x1_shape);
77+
}
78+
std::fill(x0_shape.begin(), x0_shape.end(), -1);
79+
std::unique_ptr<Operation> op_prev = std::move(*op);
80+
for (size_t i = 2; i < input_defs.size(); i++) {
81+
AddIntermediateOperationOutput(*op_prev, layer_input_name_x, elem_type, x0_shape);
82+
std::unique_ptr<Operation> op_cur = model_builder.CreateOperation(node, op_prev->type());
83+
AddOperationInput(*op_cur, "x", layer_input_name_x);
84+
AddOperationInput(*op_cur, "y", input_defs[i]->Name());
85+
model_builder.AddOperation(std::move(op_prev));
86+
op_prev = std::move(op_cur);
87+
layer_input_name_x = model_builder.GetUniqueName(node, "variadic");
88+
GetShape(*input_defs[i], x1_shape, logger);
89+
if (x0_shape.size() < x1_shape.size()) {
90+
x0_shape.resize(x1_shape.size(), -1);
91+
}
92+
}
93+
*op = std::move(op_prev);
94+
}
95+
#endif
96+
5997
Status BinaryOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder, const Node& node,
6098
const logging::Logger& logger) const {
6199
const auto& op_type(node.OpType());
@@ -90,30 +128,8 @@ Status BinaryOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder, const
90128
AddOperationInput(*op, "x", input_defs[0]->Name());
91129
AddOperationInput(*op, "y", input_defs[1]->Name());
92130
if (input_defs.size() > 2) {
93-
std::string_view layer_input_name_x = model_builder.GetUniqueName(node, "variadic");
94-
auto input_dtype = input_defs[0]->TypeAsProto()->tensor_type().elem_type();
95-
const int32_t elem_type = static_cast<int32_t>(input_dtype);
96-
std::vector<int64_t> x0_shape, x1_shape;
97-
GetShape(*input_defs[0], x0_shape, logger);
98-
GetShape(*input_defs[1], x1_shape, logger);
99-
for (size_t i = 0; i < x0_shape.size(); i++) {
100-
x0_shape[i] = std::max(x0_shape[i], x1_shape[i]);
101-
}
102-
std::unique_ptr<Operation> op_prev = std::move(op);
103-
for (size_t i = 2; i < input_defs.size(); i++) {
104-
AddIntermediateOperationOutput(*op_prev, layer_input_name_x, elem_type, x0_shape);
105-
std::unique_ptr<Operation> op_cur = model_builder.CreateOperation(node, coreml_op_type);
106-
AddOperationInput(*op_cur, "x", layer_input_name_x);
107-
AddOperationInput(*op_cur, "y", input_defs[i]->Name());
108-
model_builder.AddOperation(std::move(op_prev));
109-
op_prev = std::move(op_cur);
110-
layer_input_name_x = model_builder.GetUniqueName(node, "variadic");
111-
GetShape(*input_defs[i], x1_shape, logger);
112-
for (size_t i = 0; i < x0_shape.size(); i++) {
113-
x0_shape[i] = std::max(x0_shape[i], x1_shape[i]);
114-
}
115-
}
116-
op = std::move(op_prev);
131+
// "max" node may have variadic inputs
132+
AddVariadicInputsToMB(&op, model_builder, node, logger);
117133
}
118134
AddOperationOutput(*op, *node.OutputDefs()[0]);
119135
model_builder.AddOperation(std::move(op));

‎onnxruntime/core/providers/coreml/builders/impl/reduction_op_builder.cc‎

Lines changed: 10 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -80,6 +80,10 @@ Status ReductionOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder, co
8080
coreml_op_type = "reduce_mean";
8181
} else if (op_type == "ReduceMax") {
8282
coreml_op_type = "reduce_max";
83+
} else if (op_type == "ReduceMin") {
84+
coreml_op_type = "reduce_min";
85+
} else if (op_type == "ReduceProd") {
86+
coreml_op_type = "reduce_prod";
8387
} else {
8488
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
8589
"ReductionOpBuilder::AddToModelBuilderImpl, unexpected op: ", op_type);
@@ -120,7 +124,10 @@ Status ReductionOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder, co
120124
bool ReductionOpBuilder::IsOpSupportedImpl(const Node& node, const OpBuilderInputParams& input_params,
121125
const logging::Logger& logger) const {
122126
const auto& input_defs = node.InputDefs();
123-
127+
if (!input_params.create_mlprogram &&
128+
(node.OpType() == "ReduceMax" || node.OpType() == "ReduceMin" || node.OpType() == "ReduceProd")) {
129+
return false;
130+
}
124131
NodeAttrHelper helper(node);
125132

126133
// noop_with_empty_axes defaults to false and is only available in newer opsets where 'axes' is an optional input
@@ -140,12 +147,10 @@ bool ReductionOpBuilder::IsOpSupportedImpl(const Node& node, const OpBuilderInpu
140147
empty_axes = axes->int64_data_size() == 0;
141148
}
142149
if (empty_axes && noop_with_empty_axes && !input_params.create_mlprogram) {
143-
LOGS(logger, VERBOSE) << "CoreML doesn't support noop on empty axes for reduction layers";
144-
return false;
145-
}
146-
if (!input_params.create_mlprogram && node.OpType() == "ReduceMax") {
150+
LOGS(logger, VERBOSE) << "NeuralNetwork doesn't support noop on empty axes for reduction layers";
147151
return false;
148152
}
153+
149154
return true;
150155
}
151156

‎onnxruntime/core/providers/coreml/builders/impl/shape_op_builder.cc‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -74,9 +74,6 @@ Status ShapeOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder, const
7474
bool ShapeOpBuilder::IsOpSupportedImpl(const Node& node, const OpBuilderInputParams& input_params,
7575
const logging::Logger& logger) const {
7676
const auto* tensor_shape = node.InputDefs()[0]->Shape();
77-
if (tensor_shape == nullptr) {
78-
return false;
79-
}
8077

8178
NodeAttrHelper node_attr_helper{node};
8279
if (!input_params.create_mlprogram) {
@@ -96,6 +93,10 @@ bool ShapeOpBuilder::IsOpSupportedImpl(const Node& node, const OpBuilderInputPar
9693
int64_t start = HandleNegativeAxis(node_attr_helper.Get("start", 0), tensor_shape->dim_size());
9794
size = size - start;
9895
if (size == 0) {
96+
LOGS(logger, VERBOSE) << "Shape does not support slicing when size is 0";
97+
return false;
98+
} else if (size != tensor_shape->dim_size() && tensor_shape == nullptr) {
99+
LOGS(logger, VERBOSE) << "Shape does not support slicing when tensor_shape is not available";
99100
return false;
100101
}
101102
}

‎onnxruntime/core/providers/coreml/builders/impl/softmax_op_builder.cc‎

Lines changed: 22 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -35,49 +35,43 @@ Status SoftmaxOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder,
3535
NodeAttrHelper helper(node);
3636
int32_t axis_default_value = (node.SinceVersion() < 13) ? 1 : -1;
3737
const auto axis = helper.Get("axis", axis_default_value);
38-
auto axis_nonnegative = HandleNegativeAxis(axis, data_shape.size());
38+
auto axis_nonnegative = HandleNegativeAxis(axis, data_shape.size());
3939

4040
#if defined(COREML_ENABLE_MLPROGRAM)
41+
// CoreML's softmax match onnx's softmax behavior since opset 13.
42+
// For opset < 13, we need to reshape to 2D and set axis to -1 to simulate onnx softmax behavior.
43+
// [B,D,...](onnx softmax opset 12, axis=1)->[B,D*...](CoreML softmax, axis=-1)->[B,D,...](reshape back)
4144
if (model_builder.CreateMLProgram()) {
4245
using namespace CoreML::Specification::MILSpec;
4346
auto input_dtype = node.InputDefs()[0]->TypeAsProto()->tensor_type().elem_type();
4447
const int32_t elem_type = static_cast<int32_t>(input_dtype);
4548

4649
std::string_view layer_input_name_x = node.InputDefs()[0]->Name();
47-
std::vector<int64_t> shape1(2, 1);
48-
if (node.SinceVersion() < 13 && axis_nonnegative != static_cast<int64_t>(data_shape.size()) - 1) {
50+
const bool need_reshape = node.SinceVersion() < 13 && axis_nonnegative != static_cast<int64_t>(data_shape.size()) - 1;
51+
std::vector<int64_t> target_shape;
52+
if (need_reshape) {
4953
// reshape to 2D to simulate onnx softmax behavior
5054
auto reshape1 = model_builder.CreateOperation(node, "reshape", "pre");
51-
for (size_t i = 0, shape1_i = 0; i < data_shape.size(); i++) {
52-
if (static_cast<int64_t>(i) == axis_nonnegative) {
53-
shape1_i++;
54-
}
55-
shape1[shape1_i] *= data_shape[i];
56-
}
57-
if (axis_nonnegative == 0) {
58-
shape1[0] = shape1[1];
59-
shape1.pop_back();
60-
}
61-
axis_nonnegative = -1;
55+
TensorShape input_shape(data_shape);
56+
target_shape.push_back(input_shape.SizeToDimension(axis_nonnegative ));
57+
target_shape.push_back(input_shape.SizeFromDimension(axis_nonnegative ));
58+
axis_nonnegative = 1;
6259
AddOperationInput(*reshape1, "x", layer_input_name_x);
63-
AddOperationInput(*reshape1, "shape", model_builder.AddConstant(reshape1->type(), "shape1", shape1));
60+
AddOperationInput(*reshape1, "shape", model_builder.AddConstant(reshape1->type(), "shape1", target_shape));
6461
layer_input_name_x = model_builder.GetUniqueName(node, "ln_reshape1_");
65-
AddIntermediateOperationOutput(*reshape1, layer_input_name_x, elem_type, shape1);
62+
AddIntermediateOperationOutput(*reshape1, layer_input_name_x, elem_type, target_shape);
6663
model_builder.AddOperation(std::move(reshape1));
6764
}
6865
std::unique_ptr<Operation> op = model_builder.CreateOperation(node, "softmax");
6966
AddOperationInput(*op, "x", layer_input_name_x);
70-
AddOperationInput(*op, "axis", model_builder.AddScalarConstant(op->type(), "axis", axis_nonnegative));
71-
std::string_view ln_output_name;
72-
if (node.SinceVersion() < 13 && axis_nonnegative != static_cast<int64_t>(data_shape.size()) - 1) {
73-
ln_output_name = model_builder.GetUniqueName(node, "ln_reshape1_");
74-
AddIntermediateOperationOutput(*op, ln_output_name, elem_type, shape1);
75-
} else {
67+
AddOperationInput(*op, "axis", model_builder.AddScalarConstant(op->type(), "axis", axis_nonnegative ));
68+
if (!need_reshape) {
7669
AddOperationOutput(*op, *node.OutputDefs()[0]);
77-
}
78-
model_builder.AddOperation(std::move(op));
79-
80-
if (node.SinceVersion() < 13 && axis_nonnegative != static_cast<int64_t>(data_shape.size()) - 1) {
70+
model_builder.AddOperation(std::move(op));
71+
} else {
72+
std::string_view ln_output_name = model_builder.GetUniqueName(node, "ln_reshape1_");
73+
AddIntermediateOperationOutput(*op, ln_output_name, elem_type, target_shape);
74+
model_builder.AddOperation(std::move(op));
8175
auto reshape2 = model_builder.CreateOperation(node, "reshape", "post");
8276
AddOperationInput(*reshape2, "x", ln_output_name);
8377
AddOperationInput(*reshape2, "shape", model_builder.AddConstant(reshape2->type(), "shape2", data_shape));
@@ -97,8 +91,8 @@ Status SoftmaxOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder,
9791
// note: if opsets < 13, onnx Softmax coerces the input shape to be 2D based on axis.
9892
// we need to manually reshape to 2D and apply SoftmaxND to axis -1 to achieve equivalent results for CoreML.
9993
TensorShape input_shape(data_shape);
100-
const auto size_to_dimension = input_shape.SizeToDimension(axis_nonnegative);
101-
const auto size_from_dimension = input_shape.SizeFromDimension(axis_nonnegative);
94+
const auto size_to_dimension = input_shape.SizeToDimension(axis_nonnegative );
95+
const auto size_from_dimension = input_shape.SizeFromDimension(axis_nonnegative );
10296

10397
TensorShapeVector target_shape;
10498
target_shape.push_back(size_to_dimension);

‎onnxruntime/core/providers/coreml/builders/impl/squeeze_op_builder.cc‎

Lines changed: 17 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
#include "core/providers/coreml/shape_utils.h"
1212
#include "core/providers/shared/utils/utils.h"
1313
#include "core/optimizer/initializer.h"
14+
#include "core/providers/cpu/tensor/unsqueeze.h"
1415

1516
namespace onnxruntime {
1617
namespace coreml {
@@ -27,7 +28,7 @@ class SqueezeOpBuilder : public BaseOpBuilder {
2728
};
2829

2930
namespace {
30-
void GetAxes(ModelBuilder& model_builder, const Node& node, std::vector<int64_t>& axes) {
31+
void GetAxes(ModelBuilder& model_builder, const Node& node, TensorShapeVector& axes) {
3132
// Squeeze opset 13 use input as axes
3233
if (node.SinceVersion() > 12) {
3334
// If axes is not provided, return an empty axes as default to squeeze all
@@ -41,7 +42,8 @@ void GetAxes(ModelBuilder& model_builder, const Node& node, std::vector<int64_t>
4142
}
4243
} else {
4344
NodeAttrHelper helper(node);
44-
axes = helper.Get("axes", std::vector<int64_t>());
45+
auto axes_attr = helper.Get("axes", std::vector<int64_t>());
46+
axes.assign(axes_attr.begin(), axes_attr.end());
4547
}
4648
}
4749
} // namespace
@@ -58,7 +60,7 @@ Status SqueezeOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder,
5860
std::unique_ptr<COREML_SPEC::NeuralNetworkLayer> layer = model_builder.CreateNNLayer(node);
5961
const auto& input_defs(node.InputDefs());
6062
auto* coreml_squeeze = layer->mutable_squeeze();
61-
std::vector<int64_t> axes;
63+
TensorShapeVector axes;
6264
GetAxes(model_builder, node, axes);
6365
std::vector<int64_t> input_shape;
6466
GetShape(*input_defs[0], input_shape, logger);
@@ -72,24 +74,12 @@ Status SqueezeOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder,
7274

7375
if (coreml_op_type == "squeeze") {
7476
if (!axes.empty()) {
75-
AddOperationInput(*op, "axes", model_builder.AddConstant(op->type(), "axes", axes));
77+
// coreml squeeze op does support negative axes
78+
AddOperationInput(*op, "axes", model_builder.AddConstant(op->type(), "axes", AsSpan(axes)));
7679
}
7780
} else {
78-
for (auto& axis : axes) {
79-
axis = HandleNegativeAxis(axis, input_shape.size() + axes.size());
80-
}
81-
std::vector<int64_t> new_shape(axes.size() + input_shape.size(), 1);
82-
std::sort(axes.begin(), axes.end());
83-
// For example: Given an input tensor (data) of shape [3, 4, 5],
84-
// then Unsqueeze(data, axes=[0, 4]) outputs a tensor (expanded) containing same data as data but with shape [1, 3, 4, 5, 1].
85-
for (size_t i = 0, ori_i = 0, axes_i = 0; i < new_shape.size(); i++) {
86-
if ((axes_i >= axes.size() || static_cast<int64_t>(i) != axes[axes_i]) && input_shape.size() >= ori_i) {
87-
new_shape[i] = input_shape[ori_i++];
88-
} else {
89-
axes_i++;
90-
}
91-
}
92-
AddOperationInput(*op, "shape", model_builder.AddConstant(op->type(), "shape", new_shape));
81+
TensorShapeVector output_shape = UnsqueezeBase::ComputeOutputShape(TensorShape(input_shape), axes);
82+
AddOperationInput(*op, "shape", model_builder.AddConstant(op->type(), "shape", AsSpan(output_shape)));
9383
}
9484
AddOperationOutput(*op, *node.OutputDefs()[0]);
9585
model_builder.AddOperation(std::move(op));
@@ -118,7 +108,7 @@ bool SqueezeOpBuilder::IsOpSupportedImpl(const Node& node, const OpBuilderInputP
118108
if (node.SinceVersion() > 12 && input_defs.size() > 1) {
119109
const auto& axes_name = input_defs[1]->Name();
120110
if (!input_params.graph_viewer.GetConstantInitializer(axes_name)) {
121-
LOGS(logger, VERBOSE) << "Input axes of Squeeze must be known";
111+
LOGS(logger, VERBOSE) << "Input axes must be known";
122112
return false;
123113
}
124114
}
@@ -127,19 +117,19 @@ bool SqueezeOpBuilder::IsOpSupportedImpl(const Node& node, const OpBuilderInputP
127117
if (!input_params.create_mlprogram) {
128118
return false;
129119
}
130-
int64_t rank = -1;
120+
121+
int64_t num_of_new_dims = 0;
131122
if (node.SinceVersion() > 12) {
132-
const auto& axes_tensor = *input_params.graph_viewer.GetConstantInitializer(input_defs[1]->Name());
133-
Initializer unpacked_tensor(axes_tensor);
134-
rank = unpacked_tensor.size();
123+
num_of_new_dims = node.InputDefs()[1]->Shape()->dim(0).dim_value();
135124
} else {
136125
NodeAttrHelper helper(node);
137126
auto axes = helper.Get("axes", std::vector<int64_t>());
138-
rank = static_cast<int64_t>(axes.size());
127+
num_of_new_dims = static_cast<int64_t>(axes.size());
139128
}
129+
140130
std::vector<int64_t> input_shape;
141-
if (!GetShape(*input_defs[0], input_shape, logger) || input_shape.size() + rank > 5) {
142-
LOGS(logger, VERBOSE) << "Unsqueeze with rank > 5 is not supported";
131+
if (!GetShape(*input_defs[0], input_shape, logger) || input_shape.size() + num_of_new_dims > 5) {
132+
LOGS(logger, VERBOSE) << "Unsqueeze with num_of_new_dims > 5 is not supported";
143133
return false;
144134
}
145135
}

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,9 @@ static OpBuilderRegistrations CreateOpBuilderRegistrations() {
4545

4646
// Reduction ops
4747
CreateReductionOpBuilder("ReduceMean", op_registrations);
48+
CreateReductionOpBuilder("ReduceMin", op_registrations);
4849
CreateReductionOpBuilder("ReduceMax", op_registrations);
50+
CreateReductionOpBuilder("ReduceProd", op_registrations);
4951
CreateReductionOpBuilder("ReduceSum", op_registrations);
5052

5153
// Normalization ops

‎onnxruntime/core/providers/cpu/tensor/unsqueeze.h‎

Lines changed: 9 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -20,16 +20,7 @@ class UnsqueezeBase {
2020
};
2121

2222
Status PrepareCompute(OpKernelContext* context, Prepare& p) const;
23-
24-
protected:
25-
UnsqueezeBase(const OpKernelInfo& info) {
26-
size_t num_inputs = info.GetInputCount();
27-
if (num_inputs == 1) { // axes must be a valid attribute
28-
ORT_ENFORCE(info.GetAttrs("axes", axes_).IsOK(), "Missing/Invalid 'axes' attribute value");
29-
}
30-
}
31-
32-
static TensorShapeVector ComputeOutputShape(
23+
static TensorShapeVector ComputeOutputShape(
3324
const TensorShape& input_shape,
3425
const TensorShapeVector& axes) {
3526
TensorShapeVector output_shape;
@@ -59,6 +50,14 @@ class UnsqueezeBase {
5950
return output_shape;
6051
}
6152

53+
protected:
54+
UnsqueezeBase(const OpKernelInfo& info) {
55+
size_t num_inputs = info.GetInputCount();
56+
if (num_inputs == 1) { // axes must be a valid attribute
57+
ORT_ENFORCE(info.GetAttrs("axes", axes_).IsOK(), "Missing/Invalid 'axes' attribute value");
58+
}
59+
}
60+
6261
TensorShapeVector axes_;
6362
};
6463

0 commit comments

Comments
 (0)