up

Author: metascroy

backends/apple/coreml/partition/coreml_partitioner.py

@@ -41,7 +41,12 @@ def __init__(
             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()
+        self._logged_msgs = set()
+
+    def log_once(self, msg: str) -> None:
+        if msg not in self._logged_msgs:
+            logging.info(msg)
+            self._logged_msgs.add(msg)
 
     def is_node_supported(self, submodules, node: torch.fx.Node) -> bool:
         # get_attr node can always be supported on any backend
@@ -52,37 +57,32 @@ 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 = (
+                self.log_once(
                     "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
+            # TODO: enable this after bugs in to_edge_transform_and_lower 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
+            # 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
+            # ):
+            #     self.log_once(
+            #         "Skipping op for CoreML delegation because it contains symbolic args: "
+            #         + node_target_name
+            #     )
+            #     assert not self.lower_full_graph
+            #     return False
 
             # query coremltools to see if node is supported
             is_supported = ct.converters.mil.frontend.torch.is_torch_fx_node_supported(
@@ -100,26 +100,20 @@ def is_node_supported(self, submodules, node: torch.fx.Node) -> bool:
 Do not file an issue with ExecuTorch for op support.
 """
                     )
-                skip_str = (
+                self.log_once(
                     "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 = (
+            self.log_once(
                 "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
 
 
@@ -131,7 +125,7 @@ def __init__(
         compile_specs: Optional[List[CompileSpec]] = None,
         take_over_mutable_buffer: Optional[bool] = True,
         lower_full_graph: bool = False,
-        tag_constant_data: bool = True,
+        take_over_constant_data: bool = True,
     ) -> None:
         if skip_ops_for_coreml_delegation is None:
             skip_ops_for_coreml_delegation = []
@@ -142,7 +136,7 @@ def __init__(
         )
         self.take_over_mutable_buffer = take_over_mutable_buffer
         self.lower_full_graph = lower_full_graph
-        self.tag_constant_data = tag_constant_data
+        self.take_over_constant_data = take_over_constant_data
 
     def partition(self, exported_program: ExportedProgram) -> PartitionResult:
         # Run the CapabilityBasedPartitioner to return the largest possible
@@ -164,7 +158,7 @@ def partition(self, exported_program: ExportedProgram) -> PartitionResult:
                 node.meta["delegation_tag"] = tag
                 partition_tags[tag] = self.delegation_spec
 
-        if self.tag_constant_data:
+        if self.take_over_constant_data:
             tag_constant_data(exported_program)
         if self.take_over_mutable_buffer:
             logger.info(
@@ -215,5 +209,4 @@ def ops_to_not_decompose(
                     if warn_str not in _logged_warnings:
                         logger.warning(warn_str)
                         _logged_warnings.add(warn_str)
-
         return do_not_decompose, None

backends/apple/coreml/test/test_coreml_partitioner.py

@@ -3,6 +3,7 @@
 # Please refer to the license found in the LICENSE file in the root directory of the source tree.
 
 import copy
+import sys
 import unittest
 
 import coremltools as ct
@@ -15,6 +16,7 @@
 from executorch.backends.apple.coreml.compiler import CoreMLBackend
 from executorch.backends.apple.coreml.partition import CoreMLPartitioner
 from executorch.exir.backend.utils import format_delegated_graph
+from executorch.runtime import Runtime
 
 
 @torch.library.custom_op("unsupported::linear", mutates_args=())
@@ -35,6 +37,10 @@ def _(
     return torch.ops.aten.linear.default(x, w, b)
 
 
+_TEST_RUNTIME = sys.platform == "darwin"
+_TEST_RUNTIME = False  # Disable until segfault fixed: https://github.com/pytorch/executorch/issues/12408
+
+
 class TestCoreMLPartitioner(unittest.TestCase):
     edge_compile_config = executorch.exir.EdgeCompileConfig()
 
@@ -236,8 +242,6 @@ def forward(self, a, x, b):
 
         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 [
@@ -249,51 +253,62 @@ def forward(self, a, x, b):
         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):
+    # def test_symint_arg(self):
+    #     class Model(torch.nn.Module):
+    #         def forward(self, x, w, b, y):
+    #             val = y.item()
+    #             torch._check(val >= 0)
+    #             torch._check(val < 2)
+    #             out = torch.ops.aten.linear.default(x, w, b)
+    #             out2 = out.relu()[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(skip_ops_for_coreml_delegation=["aten.scalar_tensor.default"]))
+
+    #     # This op has symbolic args
+    #     assert (
+    #         "torch.ops.aten._assert_scalar.default"
+    #         in delegated_program_manager.exported_program().graph_module.code
+    #     )
+
+    #     if _TEST_RUNTIME:
+    #         et_prog = delegated_program_manager.to_executorch()
+    #         runtime = Runtime.get()
+    #         program = runtime.load_program(et_prog.buffer)
+    #         method = program.load_method("forward")
+    #         et_outputs = method.execute(*example_inputs)[0]
+    #         eager_outputs = model(*example_inputs)
+    #         self.assertTrue(torch.allclose(et_outputs, eager_outputs, atol=1e-02, rtol=1e-02))
+
+    def test_take_over_constant_data_false(self):
         class Model(torch.nn.Module):
             def __init__(self):
                 super().__init__()
-                self.linear = torch.nn.Linear(2, 2)
+                self.linear = torch.nn.Linear(50, 100)
 
             def forward(self, x):
                 return self.linear(x)
 
         model = Model()
         model.eval()
-        example_inputs = (torch.randn(2, 2),)
+        example_inputs = (torch.randn(2, 50),)
         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)],
+            partitioner=[CoreMLPartitioner(take_over_constant_data=False)],
         )
         for node in edge_program_manager.exported_program().graph.nodes:
             if (
@@ -305,6 +320,20 @@ def forward(self, x):
         # lowered_module_0, x, p_linear_weight, p_linear_bias
         assert len(node.args) == 4
 
+        if _TEST_RUNTIME:
+            et_prog = edge_program_manager.to_executorch()
+            runtime = Runtime.get()
+            program = runtime.load_program(et_prog.buffer)
+            method = program.load_method("forward")
+            et_outputs = method.execute(*example_inputs)[0]
+            eager_outputs = model(*example_inputs)
+            self.assertTrue(
+                torch.allclose(et_outputs, eager_outputs, atol=1e-02, rtol=1e-02)
+            )
+
+            with open("/tmp/et_model.pte", "wb") as file:
+                et_prog.write_to_file(file)
+
 
 if __name__ == "__main__":
     test_runner = TestCoreMLPartitioner()
@@ -313,5 +342,5 @@ def forward(self, x):
     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()
+    # test_runner.test_symint_arg()
+    test_runner.test_take_over_constant_data_false()