From 187743726b29e1c994c965ff2bde4bc36e8f98d2 Mon Sep 17 00:00:00 2001 From: Rachel Guo <35738743+YUNQIUGUO@users.noreply.github.com> Date: Fri, 9 Jul 2021 09:08:05 -0700 Subject: [PATCH] [CoreML EP] Add Int32<->Int64 handling around coreml ep (#8183) * initial int32-int64 type handling * initial * clean and fix UT error * modify code comments * address partial pr comments * minor update * address pr comments Co-authored-by: rachguo Co-authored-by: rachguo --- .../coreml/builders/impl/argmax_op_builder.cc | 77 ++++++++++--------- .../coreml/builders/impl/cast_op_builder.cc | 6 +- .../coreml/builders/model_builder.cc | 14 ++++ .../providers/coreml/builders/model_builder.h | 4 + .../coreml/coreml_execution_provider.cc | 8 ++ .../core/providers/coreml/model/model.h | 7 ++ .../core/providers/coreml/model/model.mm | 28 ++++++- 7 files changed, 101 insertions(+), 43 deletions(-) diff --git a/onnxruntime/core/providers/coreml/builders/impl/argmax_op_builder.cc b/onnxruntime/core/providers/coreml/builders/impl/argmax_op_builder.cc index aedd338697..e5b0d015e3 100644 --- a/onnxruntime/core/providers/coreml/builders/impl/argmax_op_builder.cc +++ b/onnxruntime/core/providers/coreml/builders/impl/argmax_op_builder.cc @@ -32,20 +32,33 @@ Status ArgMaxOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder, NodeAttrHelper helper(node); const auto axis = helper.Get("axis", 0); - const auto keep_dims = helper.Get("keep_dims", 1); - const bool removedim = keep_dims != 1; + const auto keepdims = helper.Get("keepdims", 1); + const bool removedim = keepdims != 1; auto* coreml_argmax = layer->mutable_argmax(); coreml_argmax->set_axis(axis); coreml_argmax->set_removedim(removedim); - // Get ArgMax's next node(Cast)'s outputdefs - auto it = node.OutputEdgesBegin(); - const auto* succ_node(graph_viewer.GetNode(it->GetNode().Index())); + // There are two cases here: + // 1. Special Case (ArgMax-Cast(from int64 to int32)), we fuse the Argmax's output/Cast's input + // (We still have this special case here because CoreML model does not have Cast) + // 2. Otherwise, we add Argmax layer normally + if (node.GetOutputEdgesCount() == 1) { + auto it = node.OutputEdgesBegin(); + const auto* succ_node(graph_viewer.GetNode(it->GetNode().Index())); + // If Argmax's successive node is a Cast from int64 to int32 output + // The 'cast to' type is checked in operater supported related, omit the check here + if (succ_node->OpType() == "Cast") { + // Skip the cast's input/argmax's output + *layer->mutable_input()->Add() = node.InputDefs()[0]->Name(); + *layer->mutable_output()->Add() = succ_node->OutputDefs()[0]->Name(); + model_builder.AddLayer(std::move(layer)); + return Status::OK(); + } + } - // Skip the cast's input/argmax's output *layer->mutable_input()->Add() = node.InputDefs()[0]->Name(); - *layer->mutable_output()->Add() = succ_node->OutputDefs()[0]->Name(); + *layer->mutable_output()->Add() = node.OutputDefs()[0]->Name(); model_builder.AddLayer(std::move(layer)); return Status::OK(); @@ -53,44 +66,34 @@ Status ArgMaxOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder, // Operator support related -bool ArgMaxOpBuilder::IsOpSupportedImpl(const Node& node, const OpBuilderInputParams& input_params, +bool ArgMaxOpBuilder::IsOpSupportedImpl(const Node& node, const OpBuilderInputParams& /*input_params*/, const logging::Logger& logger) const { - // Check if Argmax's output is the graph output - const auto& graph_output_list = input_params.graph_viewer.GetOutputs(); - std::unordered_set graph_outputs(graph_output_list.cbegin(), graph_output_list.cend()); - const auto& output_defs = node.OutputDefs(); - for (const auto* output_def : output_defs) { - if (graph_outputs.count(output_def) != 0) { - LOGS(logger, VERBOSE) << "ArgMax not supported when it produces a graph output"; - return false; - } - } - - // Case where argmax has multiple succeeding nodes is not supported - if (node.GetOutputEdgesCount() > 1) { - LOGS(logger, VERBOSE) << "Multiple nodes consuming ArgMax's output"; - return false; - } - - const auto& succ_node = node.OutputEdgesBegin()->GetNode(); - - /*We're only handling the case: an ArgMax op followed by a Cast to int32 type right now so as to fuse - the int64 output (not a supported output type by CoreML model) of ArgMax.*/ - if (succ_node.OpType() != "Cast") { - LOGS(logger, VERBOSE) << "ArgMax not supported when next node is not [Cast]" - << "Current next node: [" << succ_node.OpType() - << "]"; - return false; - } - // Attribute `select_last_index` of ArgMax op is not supported NodeAttrHelper helper(node); const auto select_last_index = helper.Get("select_last_index", 0); if (select_last_index != 0) { - LOGS(logger, VERBOSE) << "selected_last_index for ArgMax is not supported"; + LOGS(logger, VERBOSE) << "select_last_index for ArgMax is not supported"; return false; } + // If there are multiple downstream nodes and cast (toint32) is one of them + // not supported, exit here + // Otherwise, for general multiple downstream nodes, supported + if (node.GetOutputEdgesCount() > 1) { + for (auto it = node.OutputEdgesBegin(), end = node.OutputEdgesEnd(); it != end; ++it) { + const auto& op_type = it->GetNode().OpType(); + if (op_type == "Cast") { + // Check if the output type of cast node is int32 + NodeAttrHelper helper(it->GetNode()); + const auto cast_to_type = helper.Get("to", ONNX_NAMESPACE::TensorProto::UNDEFINED); + if (cast_to_type == ONNX_NAMESPACE::TensorProto::INT32) { + LOGS(logger, VERBOSE) << "Argmax has both cast and other downstream nodes."; + return false; + } + } + } + } + return true; } diff --git a/onnxruntime/core/providers/coreml/builders/impl/cast_op_builder.cc b/onnxruntime/core/providers/coreml/builders/impl/cast_op_builder.cc index 8117d3edee..d4d46a22c8 100644 --- a/onnxruntime/core/providers/coreml/builders/impl/cast_op_builder.cc +++ b/onnxruntime/core/providers/coreml/builders/impl/cast_op_builder.cc @@ -28,9 +28,9 @@ class CastOpBuilder : public BaseOpBuilder { Status CastOpBuilder::AddToModelBuilderImpl(ModelBuilder& /* model_builder */, const Node& /* node */, const logging::Logger& /* logger */) const { - // Right now we're only handling an ArgMax op followed by a Cast to int32 type. - // This can fuse the ArgMax's int64 output type which is not supported in CoreML model. - // And that ArgMax fused with the cast node produces an int32 output, so we're skipping adding the Cast node here. + // This is a special handling case for ArgMax Op, where argmax is followed by a cast to int32 type. + // The ArgMax is fused with the Cast node and produces an int32 output. + // Cast node is not provided in CoreML model, so we're skipping adding the Cast node here. return Status::OK(); } diff --git a/onnxruntime/core/providers/coreml/builders/model_builder.cc b/onnxruntime/core/providers/coreml/builders/model_builder.cc index 843c5602d1..2c75c5aeaf 100644 --- a/onnxruntime/core/providers/coreml/builders/model_builder.cc +++ b/onnxruntime/core/providers/coreml/builders/model_builder.cc @@ -154,6 +154,15 @@ Status ModelBuilder::RegisterModelInputOutput(const NodeArg& node_arg, bool is_i case ONNX_NAMESPACE::TensorProto_DataType_INT32: multi_array->set_datatype(COREML_SPEC::ArrayFeatureType::INT32); break; + case ONNX_NAMESPACE::TensorProto_DataType_INT64: + if (!is_input) { + // If we have an int64 output type, since COREML_SPEC:ArrayFeatureType does not support INT64 + // we assign it to be INT32 here + multi_array->set_datatype(COREML_SPEC::ArrayFeatureType::INT32); + // Record the output names and we need to change them back to Int64 when CoreML EP returns these values to ORT + AddInt64Output(name); + } + break; default: { // TODO: support other type return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, @@ -203,6 +212,7 @@ Status ModelBuilder::Compile(std::unique_ptr& model, const std::string& p ORT_RETURN_IF_ERROR(SaveCoreMLModel(path)); model.reset(new Model(path, logger_, coreml_flags_)); model->SetScalarOutputs(std::move(scalar_outputs_)); + model->SetInt64Outputs(std::move(int64_outputs_)); model->SetInputOutputInfo(std::move(input_output_info_)); return model->LoadModel(); } @@ -225,6 +235,10 @@ void ModelBuilder::AddScalarOutput(const std::string& output_name) { scalar_outputs_.insert(output_name); } +void ModelBuilder::AddInt64Output(const std::string& output_name) { + int64_outputs_.insert(output_name); +} + void ModelBuilder::AddLayer(std::unique_ptr layer) { auto* neural_network = coreml_model_->mutable_neuralnetwork(); neural_network->mutable_layers()->AddAllocated(layer.release()); diff --git a/onnxruntime/core/providers/coreml/builders/model_builder.h b/onnxruntime/core/providers/coreml/builders/model_builder.h index 8df6e20484..a7e2d2e467 100644 --- a/onnxruntime/core/providers/coreml/builders/model_builder.h +++ b/onnxruntime/core/providers/coreml/builders/model_builder.h @@ -45,6 +45,7 @@ class ModelBuilder { std::unique_ptr coreml_model_; std::unordered_set scalar_outputs_; + std::unordered_set int64_outputs_; std::unordered_map input_output_info_; std::unordered_set skipped_initializers_; @@ -70,6 +71,9 @@ class ModelBuilder { // Record the onnx scalar output names void AddScalarOutput(const std::string& output_name); + // Record the onnx int64 type output names + void AddInt64Output(const std::string& output_name); + static const IOpBuilder* GetOpBuilder(const Node& node); }; diff --git a/onnxruntime/core/providers/coreml/coreml_execution_provider.cc b/onnxruntime/core/providers/coreml/coreml_execution_provider.cc index d4a64839a3..c7b71235df 100644 --- a/onnxruntime/core/providers/coreml/coreml_execution_provider.cc +++ b/onnxruntime/core/providers/coreml/coreml_execution_provider.cc @@ -191,6 +191,11 @@ common::Status CoreMLExecutionProvider::Compile(const std::vectorIsScalarOutput(output_name)) output_shape.clear(); + // Since CoreML EP only accepts int32 output type and onnx requires int64 output, + // We are going to set the model output (from int32) ->int64 + if (model->IsInt64Output(output_name)) + output_type = ONNX_NAMESPACE::TensorProto_DataType_INT64; + auto* output_tensor = ort.KernelContext_GetOutput(context, i, output_shape.data(), output_shape.size()); @@ -202,6 +207,9 @@ common::Status CoreMLExecutionProvider::Compile(const std::vector(output_tensor); break; + case ONNX_NAMESPACE::TensorProto_DataType_INT64: + output_buffer = ort.GetTensorMutableData(output_tensor); + break; default: return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Unsupported type: ", output_type, " for output: ", output_name); diff --git a/onnxruntime/core/providers/coreml/model/model.h b/onnxruntime/core/providers/coreml/model/model.h index e5bb246679..3a1f032196 100644 --- a/onnxruntime/core/providers/coreml/model/model.h +++ b/onnxruntime/core/providers/coreml/model/model.h @@ -33,6 +33,8 @@ class Model { bool IsScalarOutput(const std::string& output_name) const; + bool IsInt64Output(const std::string& output_name) const; + // Mutex for exclusive lock to this model object OrtMutex& GetMutex() { return mutex_; } @@ -48,6 +50,7 @@ class Model { private: std::unique_ptr execution_; std::unordered_set scalar_outputs_; + std::unordered_set int64_outputs_; std::vector inputs_; std::vector outputs_; @@ -66,6 +69,10 @@ class Model { void SetScalarOutputs(std::unordered_set&& scalar_outputs) { scalar_outputs_ = std::move(scalar_outputs); } + + void SetInt64Outputs(std::unordered_set&& int64_outputs) { + int64_outputs_ = std::move(int64_outputs); + } }; } // namespace coreml diff --git a/onnxruntime/core/providers/coreml/model/model.mm b/onnxruntime/core/providers/coreml/model/model.mm index 43f15078f2..ce5da01948 100644 --- a/onnxruntime/core/providers/coreml/model/model.mm +++ b/onnxruntime/core/providers/coreml/model/model.mm @@ -237,6 +237,8 @@ [output_name cStringUsingEncoding:NSUTF8StringEncoding]); } + auto model_output_type = data.dataType; + auto& output_tensor = output.second; size_t num_elements = accumulate(output_tensor.tensor_info.shape.begin(), @@ -249,16 +251,32 @@ switch (type) { case ONNX_NAMESPACE::TensorProto_DataType_FLOAT: output_data_byte_size = num_elements * sizeof(float); + memcpy(output_tensor.buffer, model_output_data, output_data_byte_size); break; case ONNX_NAMESPACE::TensorProto_DataType_INT32: output_data_byte_size = num_elements * sizeof(int32_t); + memcpy(output_tensor.buffer, model_output_data, output_data_byte_size); + break; + // For this case, since Coreml Spec only uses int32 for model output while onnx provides + // int64 for model output data type. We are doing a type casting (int32 -> int64) here + // when copying the model to ORT + case ONNX_NAMESPACE::TensorProto_DataType_INT64: + output_data_byte_size = num_elements * sizeof(int64_t); + if (model_output_type == MLMultiArrayDataTypeInt32) { + int32_t* model_output_data_prime = static_cast(model_output_data); + int64_t* output_tensor_buffer_prime = static_cast(output_tensor.buffer); + for (size_t i = 0; i < num_elements; i++) { + output_tensor_buffer_prime[i] = model_output_data_prime[i]; + } + } + ORT_RETURN_IF_NOT(model_output_type == MLMultiArrayDataTypeInt32, + "Coreml model_output_type is not MLMultiArrayDataTypeInt32 for the case") break; default: return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, - "Output data type is not float/int32, actual type: ", + "Output data type is not supported, actual type: ", type); - } - memcpy(output_tensor.buffer, model_output_data, output_data_byte_size); + } } return onnxruntime::common::Status::OK(); @@ -335,6 +353,10 @@ bool Model::IsScalarOutput(const std::string& output_name) const { return Contains(scalar_outputs_, output_name); } +bool Model::IsInt64Output(const std::string& output_name) const { + return Contains(int64_outputs_, output_name); +} + const OnnxTensorInfo& Model::GetInputOutputInfo(const std::string& name) const { return input_output_info_.at(name); }