init

Author: metascroy

backends/apple/coreml/partition/coreml_partitioner.py

@@ -28,12 +28,20 @@
 
 class OperatorsSupportedForCoreMLBackend(OperatorSupportBase):
     def __init__(
-        self, skip_ops_for_coreml_delegation: Optional[List[str]] = None
+        self,
+        skip_ops_for_coreml_delegation: Optional[List[str]] = None,
+        lower_full_graph: bool = False,
     ) -> None:
         if skip_ops_for_coreml_delegation is None:
             skip_ops_for_coreml_delegation = []
         super().__init__()
         self.skip_ops_for_coreml_delegation = skip_ops_for_coreml_delegation
+        self.lower_full_graph = lower_full_graph
+        if self.lower_full_graph:
+            assert (
+                len(self.skip_ops_for_coreml_delegation or []) == 0
+            ), "Cannot have skip_ops_for_coreml_delegation when lower_full_graph is True"
+        self._logged_skips = set()
 
     def is_node_supported(self, submodules, node: torch.fx.Node) -> bool:
         # get_attr node can always be supported on any backend
@@ -44,14 +52,74 @@ def is_node_supported(self, submodules, node: torch.fx.Node) -> bool:
             # skip ops if specified by user
             node_target_name = getattr(node.target, "__name__", "").lower()
             if node_target_name in (self.skip_ops_for_coreml_delegation or []):
+                skip_str = (
+                    "Skipping op for CoreML delegation because it is in skip_ops_for_coreml_delegation: "
+                    + node_target_name
+                )
+                if skip_str not in self._logged_skips:
+                    logging.info(skip_str)
+                    self._logged_skips.add(skip_str)
+                assert (
+                    not self.lower_full_graph
+                ), "Cannot have skip_ops_for_coreml_delegation when lower_full_graph is True"
+                return False
+
+            # 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
+            if not self.lower_full_graph and any(
+                isinstance(arg, torch.fx.Node)
+                and isinstance(
+                    arg.meta.get("val", None),
+                    (torch.SymInt, torch.SymBool, torch.SymFloat),
+                )
+                for arg in node.args
+            ):
+                skip_str = (
+                    "Skipping op for CoreML delegation because it contains symbolic args: "
+                    + node_target_name
+                )
+                if skip_str not in self._logged_skips:
+                    logging.info(skip_str)
+                    self._logged_skips.add(skip_str)
+                assert not self.lower_full_graph
                 return False
+
             # query coremltools to see if node is supported
-            return ct.converters.mil.frontend.torch.is_torch_fx_node_supported(node)
+            is_supported = ct.converters.mil.frontend.torch.is_torch_fx_node_supported(
+                node
+            )
+            if not is_supported:
+                if self.lower_full_graph:
+                    raise NotImplementedError(
+                        f"""CoreML does not support the op {node_target_name}, but you have set lower_full_graph=True in the CoreMLPartitioner.
+
+Please set lower_full_graph=False in the CoreMLPartitioner to allow running unsupported ops outside of CoreML.  Note that setting lower_full_graph=False may affect performance of CoreML and the available features.
+As an alternative to setting lower_full_graph=False, you can try rewriting your model to avoid using this op.
+
+Also consider filing an issue with Apple's coremltools repo to request support for the op: https://github.com/apple/coremltools/issues
+Do not file an issue with ExecuTorch for op support.
+"""
+                    )
+                skip_str = (
+                    "Skipping op for CoreML delegation because it is not supported by CoreML: "
+                    + node_target_name
+                )
+                if skip_str not in self._logged_skips:
+                    logging.info(skip_str)
+                    self._logged_skips.add(skip_str)
+            return is_supported
         # cowardly refuse to support all other types of node:
         # 1. placeholder / output nodes should not be tagged
         #    reference: https://github.com/pytorch/executorch/pull/1398
         # 2. call_module / call_method should have been replaced with call_function?
         else:
+            skip_str = (
+                "Skipping op for CoreML delegation because it is not get_attr or call_function: "
+                + node.op
+            )
+            if skip_str not in self._logged_skips:
+                logging.info(skip_str)
+                self._logged_skips.add(skip_str)
             return False
 
 
@@ -62,6 +130,8 @@ def __init__(
         skip_ops_for_coreml_delegation: Optional[List[str]] = None,
         compile_specs: Optional[List[CompileSpec]] = None,
         take_over_mutable_buffer: Optional[bool] = True,
+        lower_full_graph: bool = False,
+        tag_constant_data: bool = True,
     ) -> None:
         if skip_ops_for_coreml_delegation is None:
             skip_ops_for_coreml_delegation = []
@@ -71,6 +141,8 @@ def __init__(
             compile_specs=compile_specs if compile_specs is not None else [],
         )
         self.take_over_mutable_buffer = take_over_mutable_buffer
+        self.lower_full_graph = lower_full_graph
+        self.tag_constant_data = tag_constant_data
 
     def partition(self, exported_program: ExportedProgram) -> PartitionResult:
         # Run the CapabilityBasedPartitioner to return the largest possible
@@ -80,7 +152,9 @@ def partition(self, exported_program: ExportedProgram) -> PartitionResult:
 
         capability_partitioner = CapabilityBasedPartitioner(
             exported_program.graph_module,
-            OperatorsSupportedForCoreMLBackend(self.skip_ops_for_coreml_delegation),
+            OperatorsSupportedForCoreMLBackend(
+                self.skip_ops_for_coreml_delegation, self.lower_full_graph
+            ),
             allows_single_node_partition=True,
         )
         partition_list = capability_partitioner.propose_partitions()
@@ -90,7 +164,8 @@ def partition(self, exported_program: ExportedProgram) -> PartitionResult:
                 node.meta["delegation_tag"] = tag
                 partition_tags[tag] = self.delegation_spec
 
-        tag_constant_data(exported_program)
+        if self.tag_constant_data:
+            tag_constant_data(exported_program)
         if self.take_over_mutable_buffer:
             logger.info(
                 "Core ML partitioner will take over torch mutable buffer as Core ML state, "
@@ -109,7 +184,9 @@ def ops_to_not_decompose(
         self, ep: ExportedProgram
     ) -> Tuple[List[torch._ops.OpOverload], Optional[Callable[[torch.fx.Node], bool]]]:
         do_not_decompose = []
-        op_support = OperatorsSupportedForCoreMLBackend()
+        op_support = OperatorsSupportedForCoreMLBackend(
+            self.skip_ops_for_coreml_delegation, self.lower_full_graph
+        )
         _logged_warnings = set()
 
         # CoreML prevents certain ops (like triu) from lowering to CoreML when put in the ExecuTorch op namespace

backends/apple/coreml/test/test_coreml_partitioner.py

@@ -2,6 +2,7 @@
 #
 # Please refer to the license found in the LICENSE file in the root directory of the source tree.
 
+import copy
 import unittest
 
 import coremltools as ct
@@ -16,6 +17,24 @@
 from executorch.exir.backend.utils import format_delegated_graph
 
 
+@torch.library.custom_op("unsupported::linear", mutates_args=())
+def _(
+    x: torch.Tensor,
+    w: torch.Tensor,
+    b: torch.Tensor,
+) -> torch.Tensor:
+    return torch.ops.aten.linear.default(x, w, b)
+
+
+@torch.library.register_fake("unsupported::linear")
+def _(
+    x: torch.Tensor,
+    w: torch.Tensor,
+    b: torch.Tensor,
+) -> torch.Tensor:
+    return torch.ops.aten.linear.default(x, w, b)
+
+
 class TestCoreMLPartitioner(unittest.TestCase):
     edge_compile_config = executorch.exir.EdgeCompileConfig()
 
@@ -200,10 +219,99 @@ def forward(self, q, k_val, input_pos):
             "getitem",
         ]
 
+    def test_lower_full_graph(self):
+        class Model(torch.nn.Module):
+            def forward(self, a, x, b):
+                out = torch.ops.aten.linear.default(a, x, b)
+                out2 = torch.ops.unsupported.linear.default(out, x, b)
+                return out2
+
+        model = Model()
+        model.eval()
+
+        example_inputs = (torch.randn(2, 2), torch.randn(2, 2), torch.randn(2, 2))
+        exir_program_aten = torch.export.export(model, example_inputs, strict=True)
+        edge_program_manager = executorch.exir.to_edge(exir_program_aten)
+        edge_program_manager2 = copy.deepcopy(edge_program_manager)
+
+        delegated_program_manager = edge_program_manager.to_backend(CoreMLPartitioner())
+
+        print(delegated_program_manager.exported_program())
+
+        for node in delegated_program_manager.exported_program().graph.nodes:
+            if node.op == "call_function":
+                assert node.target.__name__ in [
+                    "unsupported.linear.default",
+                    "executorch_call_delegate",
+                    "getitem",
+                ], node.target.__name__
+
+        with self.assertRaises(NotImplementedError):
+            edge_program_manager2.to_backend(CoreMLPartitioner(lower_full_graph=True))
+
+    def test_symint_arg(self):
+        class Model(torch.nn.Module):
+            def forward(self, x, w, b, y):
+                val = y.item()
+                out = torch.ops.unsupported.linear.default(x, w, b + val) + val
+                out2 = torch.ops.aten.linear.default(out, w, b) + val
+                return out2
+
+        model = Model()
+        model.eval()
+        example_inputs = (
+            torch.randn(2, 2),
+            torch.randn(2, 2),
+            torch.randn(2, 2),
+            torch.tensor(2),
+        )
+        exir_program_aten = torch.export.export(model, example_inputs)
+
+        edge_program_manager = executorch.exir.to_edge(exir_program_aten)
+
+        delegated_program_manager = edge_program_manager.to_backend(CoreMLPartitioner())
+
+        # This op has symbolic args
+        assert (
+            "torch.ops.aten.scalar_tensor.default"
+            in delegated_program_manager.exported_program().graph_module.code
+        )
+
+    def test_tag_constant_data_false(self):
+        class Model(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.linear = torch.nn.Linear(2, 2)
+
+            def forward(self, x):
+                return self.linear(x)
+
+        model = Model()
+        model.eval()
+        example_inputs = (torch.randn(2, 2),)
+        exir_program_aten = torch.export.export(model, example_inputs)
+
+        edge_program_manager = executorch.exir.to_edge_transform_and_lower(
+            exir_program_aten,
+            partitioner=[CoreMLPartitioner(tag_constant_data=False)],
+        )
+        for node in edge_program_manager.exported_program().graph.nodes:
+            if (
+                node.op == "call_function"
+                and node.target.__name__ == "executorch_call_delegate"
+            ):
+                break
+
+        # lowered_module_0, x, p_linear_weight, p_linear_bias
+        assert len(node.args) == 4
+
 
 if __name__ == "__main__":
     test_runner = TestCoreMLPartitioner()
     test_runner.test_add_sub_skip_mm()
     test_runner.test_vit_skip_conv()
     test_runner.test_ops_to_not_decompose()
     test_runner.test_buffer()
+    test_runner.test_lower_full_graph()
+    test_runner.test_symint_arg()
+    test_runner.test_tag_constant_data_false()