Fixes to_edge_transform_and_lower when unsupported ops are asked for preservation (#8776)

Summary:
If a partitioner requests to_edge_transform_and_lower keep mutable / aliasing ops (e.g., transpose, view, permute, etc), lowering with ExecuTorch fails because those ops cannot be functionalized when wrapped in the EDGE_DO_NOT_DECOMP namespace as custom ops.

This PR filters out unsupported ops that backends request for preservation.


Differential Revision: D70333876

Pulled By: metascroy

Author: metascroy

backends/apple/coreml/partition/coreml_partitioner.py

@@ -111,10 +111,16 @@ def ops_to_not_decompose(
         do_not_decompose = []
         op_support = OperatorsSupportedForCoreMLBackend()
         for node in ep.graph.nodes:
-            if (
-                node.op == "call_function"
-                and isinstance(node.target, torch._ops.OpOverload)
-                and op_support.is_node_supported(None, node)
+            if node.op == "call_function" and isinstance(
+                node.target, torch._ops.OpOverload
             ):
-                do_not_decompose.append(node.target)
+                try:
+                    if op_support.is_node_supported(None, node):
+                        do_not_decompose.append(node.target)
+                except Exception as e:
+                    # CoreML's op_support.is_node_supported will sometimes throw
+                    # for unsupported ops, rather than returning False
+                    logger.warning(
+                        f"Encountered exception when checking node support: {e}"
+                    )
         return do_not_decompose, None

backends/apple/coreml/test/test_coreml_partitioner.py

@@ -82,11 +82,28 @@ def test_vit_skip_conv(self):
 
     def test_ops_to_not_decompose(self):
         class Model(torch.nn.Module):
+            def __init__(self) -> None:
+                super().__init__()
+
+                buffer = torch.ones(1)
+                self.register_buffer("buffer", buffer)
+
             def forward(self, q, k, v, mask):
-                return torch.ops.aten.scaled_dot_product_attention.default(
+                out = torch.ops.aten.scaled_dot_product_attention.default(
                     q, k, v, attn_mask=mask
                 )
 
+                # Add non-functional and alias ops
+                # These will be removed by ExecuTorch in non-decomposition
+                # table because they cannot be functionalized
+                out = out.transpose(1, 2)
+                out = out.view(1, -1)
+                out = out.permute(0, 1)
+                out = out.add_(self.buffer)
+                out = torch.ops.aten.view_copy.default(out, (-1,))
+                out = out.select(0, 0)
+                return out
+
         model = Model()
         model.eval()
 
@@ -107,6 +124,9 @@ def forward(self, q, k, v, mask):
         edge_program_manager = executorch.exir.to_edge_transform_and_lower(
             ep, partitioner=[coreml_partitioner]
         )
+        print(
+            format_delegated_graph(edge_program_manager.exported_program().graph_module)
+        )
         self.assertTrue(
             "executorch.exir.dialects.edge._ops.aten.scaled_dot_product_attention.default"
             in format_delegated_graph(

exir/program/_program.py

@@ -26,6 +26,7 @@
 from executorch.exir.emit._emitter import _DelegateDebugIdentifierMap
 from executorch.exir.error import ExportError
 from executorch.exir.graph_module import get_control_flow_submodules
+from executorch.exir.operator.convert import _pybind_schema_to_native_schema
 from executorch.exir.pass_base import PassBase
 from executorch.exir.pass_manager import PassType
 from executorch.exir.passes import (
@@ -836,6 +837,9 @@ def _replace_aten_ops_with_transformed_ops(
         ops_set_to_not_decompose, check_op_support = partitioner.ops_to_not_decompose(
             program
         )
+        ops_set_to_not_decompose = _remove_invalid_ops_for_not_decompose(
+            ops_set_to_not_decompose
+        )
 
         for op_aten in ops_set_to_not_decompose:
             _register_no_decomp_op(op_aten)
@@ -965,6 +969,37 @@ def _sanity_check_graph_for_non_decomp_ops(
                     logging.warning(warning_str)
 
 
+def _remove_invalid_ops_for_not_decompose(
+    ops_to_not_decompose: List[torch._ops.OpOverload],
+) -> List[torch._ops.OpOverload]:
+    def keep(op):
+        schema = op._schema
+        native_schema = _pybind_schema_to_native_schema(schema)
+        if native_schema.is_mutable:
+            logging.warn(f"Op {op} was requested for preservation by partitioner.  This request is ignored because it is mutable.")
+            return False
+
+        if native_schema.aliased_return_names() != [None]:
+            logging.warn(f"Op {op} was requested for preservation by partitioner.  This request is ignored because it aliases output.")
+            return False
+
+        # Explicit block list of ops that don't work if asked for
+        # preservation
+        if op in [
+            # Hits infinte recursion error when op is in 
+            # EDGE_DO_NOT_DECOMP namespace 
+            torch.ops.aten._to_copy.default,
+            # type scalar->tensor type promotion does not work on op 
+            # in EDGE_DO_NOT_DECOMP namespace  
+            torch.ops.aten.mul.Tensor,
+            ]:
+            logging.warn(f"Op {op} was requested for preservation by partitioner.  This request is ignored because it is in a blocklist.")
+            return False
+        return True
+
+    return list(filter(keep, ops_to_not_decompose))
+
+
 def _gen_edge_manager_for_partitioners(
     partitioner: Dict[str, List[Partitioner]],
     aten_programs: Dict[str, ExportedProgram],
@@ -992,6 +1027,9 @@ def _gen_edge_manager_for_partitioners(
             all_ops_no_decomp = set()
             for curr_partitioner in partitioner.get(name, []):
                 curr_ops_no_decomp, _ = curr_partitioner.ops_to_not_decompose(program)
+                curr_ops_no_decomp = _remove_invalid_ops_for_not_decompose(
+                    curr_ops_no_decomp
+                )
                 all_ops_no_decomp |= set(curr_ops_no_decomp)
 
             table = _default_decomposition_table()
@@ -1113,6 +1151,7 @@ def to_edge_transform_and_lower(
             curr_op_set, check_op_support = curr_partitioner.ops_to_not_decompose(
                 program
             )
+            curr_op_set = _remove_invalid_ops_for_not_decompose(curr_op_set)
             ops_set_to_not_decompose = ops_set_to_not_decompose.union(curr_op_set)
             _sanity_check_graph_for_non_decomp_ops(
                 name,