up

Author: metascroy

backends/apple/coreml/test/test_coreml_partitioner.py

@@ -85,9 +85,6 @@ 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):
                 out = torch.ops.aten.scaled_dot_product_attention.default(
                     q, k, v, attn_mask=mask
@@ -99,7 +96,10 @@ def forward(self, q, k, v, mask):
                 out = out.transpose(1, 2)
                 out = out.view(1, -1)
                 out = out.permute(0, 1)
-                out = out.add_(self.buffer)
+                out = out.add_(1.0)
+                out = out.mul_(2.0)
+                out = out.div_(3.0)
+                out = out.sub_(4.0)
                 out = torch.ops.aten.view_copy.default(out, (-1,))
                 out = out.select(0, 0)
                 return out

exir/program/_program.py

@@ -976,8 +976,33 @@ 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,
+            # scalar to tensor type promotion does not work on ops
+            # in EDGE_DO_NOT_DECOMP namespace
+            torch.ops.aten.mul.Tensor,
+            torch.ops.aten.add.Tensor,
+            torch.ops.aten.sub.Tensor,
+            torch.ops.aten.div.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