From 3a4af382faade1f1e0a7129658aa5ae5a05df381 Mon Sep 17 00:00:00 2001
From: john-rocky <rockyshikoku@gmail.com>
Date: Wed, 13 May 2026 03:22:34 +0900
Subject: [PATCH] [CoreML EP] Support Gather with scalar constant 'indices' on
 the MLProgram path

Fixes #28180. GatherOpBuilder rejected ONNX Gather with rank-0 'indices',
forcing a CPU partition split. PyTorch's exporter emits this routinely - e.g.
GFPGAN 1024x1024 has 16 per-layer style-code Gathers with a scalar constant
index, splitting the CoreML subgraph in two.

Building a small MIL program (mb.gather with a rank-0 mb.const indices)
end-to-end against iOS15 confirms MIL gather accepts rank-0 indices and drops
the gathered axis - matching ONNX semantics:

  main[CoreML3](%x: (1, 16, 512, fp32)) {
    %gather_0: (1, 512, fp32) = gather(x=%x, indices=3, axis=1)
  }

The original rejection was driven by ORT-side boundary handling, not by MIL:
RegisterModelInputOutput rewrites any rank-0 boundary tensor to {1}
(MLMultiArray has no rank-0), and on the NN path RegisterInitializers does
the same for rank-0 initializers (LoadConstantND requires rank >= 1).
Constant initializers on the MLProgram path go through
OnnxTensorToCoreMLTensor which preserves rank, so MIL gather can consume
them directly.

Therefore: allow scalar 'indices' on the MLProgram path only when it is a
constant initializer. NN path and non-initializer scalar 'indices' remain
rejected. AddToModelBuilderImpl is unchanged.

Existing CPU OpTester cases cover the shape and now stay on CoreML EP:
- GatherOpTest.Gather_axis0_scalar_indices
- GatherOpTest.Gather_axis1_scalar_indices
---
 .../coreml/builders/impl/gather_op_builder.cc | 35 ++++++++++++++-----
 1 file changed, 27 insertions(+), 8 deletions(-)

diff --git a/onnxruntime/core/providers/coreml/builders/impl/gather_op_builder.cc b/onnxruntime/core/providers/coreml/builders/impl/gather_op_builder.cc
index 8b58f5dc6c927..add16b569b2a0 100644
--- a/onnxruntime/core/providers/coreml/builders/impl/gather_op_builder.cc
+++ b/onnxruntime/core/providers/coreml/builders/impl/gather_op_builder.cc
@@ -74,27 +74,46 @@ bool GatherOpBuilder::HasSupportedInputsImpl(const Node& node, const OpBuilderIn
   return true;
 }
 
-bool GatherOpBuilder::IsOpSupportedImpl(const Node& node, const OpBuilderInputParams& /*input_params*/,
+bool GatherOpBuilder::IsOpSupportedImpl(const Node& node, const OpBuilderInputParams& input_params,
                                         const logging::Logger& logger) const {
+  const auto& input_defs = node.InputDefs();
   std::vector<int64_t> data_shape, indices_shape;
-  if (!GetShape(*node.InputDefs()[0], data_shape, logger)) {
+  if (!GetShape(*input_defs[0], data_shape, logger)) {
     LOGS(logger, VERBOSE) << "Failed to get 'data' shape";
     return false;
   }
 
-  if (!GetShape(*node.InputDefs()[1], indices_shape, logger)) {
+  if (!GetShape(*input_defs[1], indices_shape, logger)) {
     LOGS(logger, VERBOSE) << "Failed to get 'indices' shape";
     return false;
   }
 
-  // Don't allow scalar 'indices' input.
-  // We convert scalar inputs to tensors with shape [1] before providing them to CoreML.
-  // This modification changes the shape of the Gather output.
+  // Scalar (rank-0) 'indices' input.
+  //
+  // MIL `gather` accepts rank-0 indices and produces the correct ONNX output
+  // (the gathered axis is dropped). However, the CoreML EP reshapes any rank-0
+  // *graph-boundary* tensor to {1} when it crosses into the CoreML subgraph
+  // (see ModelBuilder::RegisterModelInputOutput) and a rank-0 input cannot be
+  // represented as an MLMultiArray. So we only allow scalar indices when they
+  // are a constant initializer: those flow through OnnxTensorToCoreMLTensor
+  // with rank preserved and MIL gather can consume them directly.
+  //
+  // On the NeuralNetwork path scalar initializers are also reshaped to {1}
+  // (ModelBuilder::RegisterInitializers, LoadConstantND requires rank >= 1),
+  // so the gather output shape ends up wrong there. Keep rejecting that case.
   if (indices_shape.empty()) {
-    LOGS(logger, VERBOSE) << "Gather does not support scalar 'indices'";
-    return false;
+    if (!input_params.create_mlprogram) {
+      LOGS(logger, VERBOSE) << "Gather does not support scalar 'indices' on the NeuralNetwork path";
+      return false;
+    }
+    if (input_params.graph_viewer.GetConstantInitializer(input_defs[1]->Name()) == nullptr) {
+      LOGS(logger, VERBOSE) << "Gather with scalar 'indices' is only supported when 'indices' is a constant initializer";
+      return false;
+    }
   }
 
+  // ONNX Gather output rank = data_rank + indices_rank - 1.
+  // For scalar indices (rank 0) this is data_rank - 1, which is what MIL also produces.
   if (data_shape.size() + indices_shape.size() - 1 > 5) {
     LOGS(logger, VERBOSE) << "Gather does not support output with rank greater than 5";
     return false;