Do not delegate CoreML reflect/replicate pads beyond the last two dims

backends/apple/coreml/partition/coreml_partitioner.py

@@ -83,6 +83,41 @@ def _is_arg_min_max_over_flattened_input(node: torch.fx.Node) -> bool:
     return dim is None
 
 
+_PAD_TARGETS = (
+    torch.ops.aten.pad.default,
+    exir_ops.edge.aten.pad.default,
+)
+
+_PAD_3D_TARGETS = (
+    torch.ops.aten.reflection_pad3d.default,
+    torch.ops.aten.replication_pad3d.default,
+    exir_ops.edge.aten.reflection_pad3d.default,
+    exir_ops.edge.aten.replication_pad3d.default,
+)
+
+
+def _is_pad_unsupported_by_coreml(node: torch.fx.Node) -> bool:
+    """Non-constant padding that Core ML cannot load.
+
+    Core ML only supports ``reflect`` / ``replicate`` padding on the last two
+    dimensions; anything further out (e.g. ``ReflectionPad3d``, or padding
+    just the depth dimension of a 5D input) lowers, then fails to load with
+    "Padding for more than two dimensions only supports `constant` mode".
+    Rejecting the op lets it decompose into ops Core ML does run.  See
+    pytorch/executorch#11708 and pytorch/executorch#11709.
+    """
+    if node.target in _PAD_3D_TARGETS:
+        return True
+    if node.target not in _PAD_TARGETS:
+        return False
+    mode = node.args[2] if len(node.args) >= 3 else node.kwargs.get("mode", "constant")
+    if mode not in ("reflect", "replicate"):
+        return False
+    # ``pad`` lists (before, after) pairs starting from the last dimension.
+    pad = node.args[1]
+    return any(p != 0 for p in pad[4:])
+
+
 def _is_view_op(op: torch._ops.OpOverload) -> bool:
     schema = op._schema
     if len(schema.arguments) == 0:
@@ -159,6 +194,13 @@ def should_override_support(self, node) -> bool:
             )
             return True
 
+        if _is_pad_unsupported_by_coreml(node):
+            self.log_once(
+                "Non-constant padding outside the last two dimensions is not "
+                "supported by CoreML.  Overriding op support."
+            )
+            return True
+
         # TODO: enable this after bugs in ExecuTorch's partitioner are fixed
         # # If lower_full_graph=False, do not partition nodes with symbolic args because it can result in symbolic args
         # # in the placeholders due to partitioning, which CoreML does not support

backends/apple/coreml/test/test_coreml_partitioner.py

@@ -433,6 +433,84 @@ def forward(self, x):
         self.assertIn("executorch_call_delegate", op_names)
         self.assertNotIn("aten.argmax.default", op_names)
 
+    def test_nonconstant_pad_beyond_last_two_dims_is_decomposed(self):
+        """
+        Regression test for https://github.com/pytorch/executorch/issues/11708
+        and https://github.com/pytorch/executorch/issues/11709.
+
+        Core ML only supports reflect / replicate padding on the last two
+        dimensions.  Handing it a wider pad lowers fine but the model fails to
+        load ("Padding for more than two dimensions only supports `constant`
+        mode").  The partitioner must not keep such a pad whole, so it
+        decomposes into ops Core ML can run (index / arange / clamp) and the
+        model stays delegated; 2D reflect / replicate padding is delegated
+        intact as ``pad``.
+        """
+
+        class Pad(torch.nn.Module):
+            def __init__(self, pad, mode):
+                super().__init__()
+                self.pad = pad
+                self.mode = mode
+
+            def forward(self, x):
+                return torch.nn.functional.pad(x, self.pad, mode=self.mode)
+
+        def lower(model, x):
+            ep = torch.export.export(model, (x,), strict=True)
+            edge = executorch.exir.to_edge_transform_and_lower(
+                ep, partitioner=[CoreMLPartitioner()]
+            )
+            op_names = [
+                n.target.__name__
+                for n in edge.exported_program().graph.nodes
+                if n.op == "call_function"
+            ]
+            delegated = format_delegated_graph(edge.exported_program().graph_module)
+            if _TEST_RUNTIME:
+                et_prog = edge.to_executorch()
+                program = Runtime.get().load_program(et_prog.buffer)
+                et_outputs = program.load_method("forward").execute((x,))[0]
+                self.assertTrue(
+                    torch.allclose(et_outputs, model(x), atol=1e-02, rtol=1e-02),
+                    model,
+                )
+            return ep, op_names, delegated
+
+        x5 = torch.randn(1, 3, 6, 6, 6)
+        unsupported = [
+            (torch.nn.ReflectionPad3d(2), x5),
+            (torch.nn.ReplicationPad3d(2), x5),
+            # Pads a single dimension, but not one of the last two.
+            (Pad((0, 0, 0, 0, 1, 1), "reflect"), x5),
+        ]
+        for model, x in unsupported:
+            model = model.eval()
+            ep, op_names, delegated = lower(model, x)
+            not_decomposed, _ = CoreMLPartitioner().ops_to_not_decompose(ep)
+            self.assertNotIn(torch.ops.aten.pad.default, not_decomposed, model)
+            # The pad is decomposed, not kept whole: no pad op is left in the
+            # graph or handed to Core ML, and its decomposition is delegated.
+            self.assertIn("executorch_call_delegate", op_names, model)
+            self.assertNotIn("aten.pad.default", op_names, model)
+            self.assertNotIn("aten.pad.default", delegated, model)
+            self.assertIn("aten.index.Tensor", delegated, model)
+
+        x4 = torch.randn(1, 3, 6, 6)
+        supported = [
+            Pad((1, 1, 1, 1), "reflect"),
+            Pad((1, 1, 1, 1), "replicate"),
+        ]
+        for model in supported:
+            model = model.eval()
+            ep, op_names, delegated = lower(model, x4)
+            not_decomposed, _ = CoreMLPartitioner().ops_to_not_decompose(ep)
+            self.assertIn(torch.ops.aten.pad.default, not_decomposed, model)
+            # The pad itself is handed to Core ML.
+            self.assertIn("executorch_call_delegate", op_names, model)
+            self.assertNotIn("aten.pad.default", op_names, model)
+            self.assertIn("aten.pad.default", delegated, model)
+
     def test_tied_embedding_is_quantized_with_its_linear(self):
         """A weight that is both an embedding table and an output projection lowers.
 
@@ -572,6 +650,7 @@ def forward(self, x):
     test_runner.test_take_over_constant_data_false()
     test_runner.test_aten_rand_default_falls_back_to_portable()
     test_runner.test_aten_randn_is_still_delegated()
+    test_runner.test_nonconstant_pad_beyond_last_two_dims_is_decomposed()
     test_runner.test_tied_embedding_is_quantized_with_its_linear()
     test_runner.test_deprecation_warning_for_to_backend_workflow()
     test_runner.test_no_warning_for_to_edge_transform_and_lower_workflow()