[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 <rachguo@rachguos-Mac-mini.local>
Co-authored-by: rachguo <rachguo@rachguos-Mini.attlocal.net>

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<const NodeArg*> 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;
 }
 

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();
 }
 

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>& 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<COREML_SPEC::NeuralNetworkLayer> layer) {
   auto* neural_network = coreml_model_->mutable_neuralnetwork();
   neural_network->mutable_layers()->AddAllocated(layer.release());

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

@@ -45,6 +45,7 @@ class ModelBuilder {
 
   std::unique_ptr<CoreML::Specification::Model> coreml_model_;
   std::unordered_set<std::string> scalar_outputs_;
+  std::unordered_set<std::string> int64_outputs_;
   std::unordered_map<std::string, OnnxTensorInfo> input_output_info_;
 
   std::unordered_set<std::string> 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);
 };
 

onnxruntime/core/providers/coreml/coreml_execution_provider.cc

@@ -191,6 +191,11 @@ common::Status CoreMLExecutionProvider::Compile(const std::vector<FusedNodeAndGr
           if (model->IsScalarOutput(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<FusedNodeAndGr
             case ONNX_NAMESPACE::TensorProto_DataType_INT32:
               output_buffer = ort.GetTensorMutableData<int32_t>(output_tensor);
               break;
+            case ONNX_NAMESPACE::TensorProto_DataType_INT64:
+              output_buffer = ort.GetTensorMutableData<int64_t>(output_tensor);
+              break;
             default:
               return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL,
                                      "Unsupported type: ", output_type, " for output: ", output_name);

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> execution_;
   std::unordered_set<std::string> scalar_outputs_;
+  std::unordered_set<std::string> int64_outputs_;
 
   std::vector<std::string> inputs_;
   std::vector<std::string> outputs_;
@@ -66,6 +69,10 @@ class Model {
   void SetScalarOutputs(std::unordered_set<std::string>&& scalar_outputs) {
     scalar_outputs_ = std::move(scalar_outputs);
   }
+
+  void SetInt64Outputs(std::unordered_set<std::string>&& int64_outputs) {
+    int64_outputs_ = std::move(int64_outputs);
+  }
 };
 
 }  // namespace coreml

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<int32_t*>(model_output_data);
+          int64_t* output_tensor_buffer_prime = static_cast<int64_t*>(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);
 }