up

Author: metascroy

backends/apple/coreml/partition/coreml_partitioner.py

@@ -66,7 +66,7 @@ def is_node_supported(self, submodules, node: torch.fx.Node) -> bool:
                 ), "Cannot have skip_ops_for_coreml_delegation when lower_full_graph is True"
                 return False
 
-            # TODO: enable this after bugs in to_edge_transform_and_lower are fixed
+            # 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
             # if not self.lower_full_graph and any(
@@ -137,6 +137,7 @@ def __init__(
         self.take_over_mutable_buffer = take_over_mutable_buffer
         self.lower_full_graph = lower_full_graph
         self.take_over_constant_data = take_over_constant_data
+        self._logged_msgs = set()
 
     def partition(self, exported_program: ExportedProgram) -> PartitionResult:
         # Run the CapabilityBasedPartitioner to return the largest possible
@@ -174,14 +175,18 @@ def partition(self, exported_program: ExportedProgram) -> PartitionResult:
             tagged_exported_program=exported_program, partition_tags=partition_tags
         )
 
+    def log_once(self, msg: str) -> None:
+        if msg not in self._logged_msgs:
+            logging.info(msg)
+            self._logged_msgs.add(msg)
+
     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(
             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
         # TODO: upstream fixes, but pending ET consuming a new published version of coremltools with the
@@ -205,8 +210,7 @@ def ops_to_not_decompose(
                 except Exception as e:
                     # CoreML's op_support.is_node_supported will sometimes throw
                     # for unsupported ops, rather than returning False
-                    warn_str = f"Encountered exception when checking node support: {e}"
-                    if warn_str not in _logged_warnings:
-                        logger.warning(warn_str)
-                        _logged_warnings.add(warn_str)
+                    self.log_once(
+                        f"Encountered exception when checking node support, treating node as unsupported: {e}"
+                    )
         return do_not_decompose, None

backends/apple/coreml/test/test_coreml_partitioner.py

@@ -253,6 +253,7 @@ def forward(self, a, x, b):
         with self.assertRaises(NotImplementedError):
             edge_program_manager2.to_backend(CoreMLPartitioner(lower_full_graph=True))
 
+    # TODO: enable this after bugs are fixed in ExecuTorch's partitioner
     # def test_symint_arg(self):
     #     class Model(torch.nn.Module):
     #         def forward(self, x, w, b, y):
@@ -331,9 +332,6 @@ def forward(self, x):
                 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()