diff --git a/backends/apple/coreml/partition/coreml_partitioner.py b/backends/apple/coreml/partition/coreml_partitioner.py index b4e550d8479..51da61b38f3 100644 --- a/backends/apple/coreml/partition/coreml_partitioner.py +++ b/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 diff --git a/backends/apple/coreml/test/test_coreml_partitioner.py b/backends/apple/coreml/test/test_coreml_partitioner.py index 78dd114eeda..0cd84ee9102 100644 --- a/backends/apple/coreml/test/test_coreml_partitioner.py +++ b/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()