Fix crashes, wrong outputs and lowering failures in several backends and in exir

backends/aoti/aoti_backend.py

@@ -7,6 +7,7 @@
 import contextlib
 import hashlib
 import os
+import re
 import typing
 from abc import ABC, abstractmethod
 from enum import Enum
@@ -168,7 +169,8 @@ def release_moved_tensors(
     def collect_unsupported_fallback_kernels(cls, missing_fallback_kernels: Set[str]):
         """
         Context manager to collect unsupported fallback kernels during compilation.
-        Monitors both extern kernel calls and runtime lookup.
+        Monitors extern kernel calls, runtime lookup, and the scatter and
+        index_put fallbacks, which call their kernels directly.
         """
         supported_kernels = cls.get_supported_fallback_kernels()
 
@@ -178,6 +180,10 @@ def collect_unsupported_fallback_kernels(cls, missing_fallback_kernels: Set[str]
         original_generate_fallback_kernel_with_runtime_lookup_aot = (
             CppWrapperCpu.generate_fallback_kernel_with_runtime_lookup_aot
         )
+        original_generate_scatter_fallback = CppWrapperCpu._generate_scatter_fallback
+        original_generate_index_put_fallback = (
+            CppWrapperCpu._generate_index_put_fallback
+        )
 
         def generate_c_shim_extern_kernel_call_and_collect_unsupported_kernels(
             self, kernel: str, *args: Any, **kwargs: Any
@@ -200,10 +206,36 @@ def generate_fallback_kernel_with_runtime_lookup_aot_and_collect_unsupported_ker
                 self, op_overload, *args, **kwargs
             )
 
+        def collect_kernels_it_writes(original: Any) -> Any:
+            # These fallbacks write their shim call directly, so read the
+            # kernel names from the lines they write.
+            def generate_and_collect_unsupported_kernels(
+                self, *args: Any, **kwargs: Any
+            ) -> None:
+                lines: List[str] = []
+                writeline = self.writeline
+                self.writeline = lambda line: (lines.append(str(line)), writeline(line))
+                try:
+                    return original(self, *args, **kwargs)
+                finally:
+                    self.writeline = writeline
+                    for line in lines:
+                        for kernel in re.findall(r"\b(aoti_torch_\w+)\(", line):
+                            if kernel not in supported_kernels:
+                                missing_fallback_kernels.add(kernel)
+
+            return generate_and_collect_unsupported_kernels
+
         CppWrapperCpu.generate_c_shim_extern_kernel_call = (
             generate_c_shim_extern_kernel_call_and_collect_unsupported_kernels
         )
         CppWrapperCpu.generate_fallback_kernel_with_runtime_lookup_aot = generate_fallback_kernel_with_runtime_lookup_aot_and_collect_unsupported_kernels
+        CppWrapperCpu._generate_scatter_fallback = collect_kernels_it_writes(
+            original_generate_scatter_fallback
+        )
+        CppWrapperCpu._generate_index_put_fallback = collect_kernels_it_writes(
+            original_generate_index_put_fallback
+        )
 
         try:
             yield
@@ -214,6 +246,12 @@ def generate_fallback_kernel_with_runtime_lookup_aot_and_collect_unsupported_ker
             CppWrapperCpu.generate_fallback_kernel_with_runtime_lookup_aot = (
                 original_generate_fallback_kernel_with_runtime_lookup_aot
             )
+            CppWrapperCpu._generate_scatter_fallback = (
+                original_generate_scatter_fallback
+            )
+            CppWrapperCpu._generate_index_put_fallback = (
+                original_generate_index_put_fallback
+            )
 
     @classmethod
     def preprocess(

backends/aoti/passes/replace_view_copy_with_view.py

@@ -31,6 +31,12 @@
     # natively by inductor, so map it here like the other view-copy ops.
     torch.ops.aten.split_copy.Tensor: torch.ops.aten.split.Tensor,
     ops.edge.aten.split_copy.Tensor: ops.edge.aten.split.Tensor,
+    # Edge turns ``detach`` into ``detach_copy``, which has no c-shim either.
+    # Without autograd it is ``alias_copy``, which Inductor can decompose,
+    # unlike ``detach_copy``. A view would let two outputs share one handle,
+    # and ``clone`` fails Edge verification in dim order mode.
+    torch.ops.aten.detach_copy.default: torch.ops.aten.alias_copy.default,
+    ops.edge.aten.detach_copy.default: ops.edge.aten.alias_copy.default,
 }
 
 
@@ -40,19 +46,26 @@ class ReplaceViewCopyWithViewPass(ExportPass):
     def call(self, graph_module: fx.GraphModule) -> PassResult:
         graph_changed = False
 
-        for node in graph_module.graph.nodes:
-            if node.op != "call_function" or node.target not in _VIEW_TARGETS:
+        # Control flow branches, such as those of torch.cond, are submodules.
+        for module in graph_module.modules():
+            if not isinstance(module, fx.GraphModule):
                 continue
 
-            if self._has_blocking_user(node, node.users.keys()):
-                continue
+            module_changed = False
+            for node in module.graph.nodes:
+                if node.op != "call_function" or node.target not in _VIEW_TARGETS:
+                    continue
+
+                if self._has_blocking_user(node, node.users.keys()):
+                    continue
 
-            node.target = _VIEW_TARGETS[node.target]
-            graph_changed = True
+                node.target = _VIEW_TARGETS[node.target]
+                module_changed = True
 
-        if graph_changed:
-            graph_module.graph.lint()
-            graph_module.recompile()
+            if module_changed:
+                module.graph.lint()
+                module.recompile()
+                graph_changed = True
 
         return PassResult(graph_module, graph_changed)
 

backends/apple/coreml/partition/coreml_partitioner.py

@@ -63,6 +63,65 @@
 )
 
 
+# Ops that coremltools can carry through a tensor above rank 5, by rewriting a
+# reshape -> transpose -> reshape chain into rank 4 steps.
+_HIGH_RANK_LAYOUT_TARGETS = frozenset(
+    [
+        exir_ops.edge.aten.view_copy.default,
+        exir_ops.edge.aten.permute_copy.default,
+        exir_ops.edge.aten.clone.default,
+        exir_ops.edge.dim_order_ops._clone_dim_order.default,
+    ]
+)
+
+
+def _produces_tensor_above_rank_5(node: torch.fx.Node) -> bool:
+    val = node.meta.get("val")
+    vals = list(val) if isinstance(val, (list, tuple)) else [val]
+    return any(isinstance(v, torch.Tensor) and v.dim() > 5 for v in vals)
+
+
+def _touches_unsupported_high_rank_tensor(
+    node: torch.fx.Node, is_skipped: Callable[[torch.fx.Node], bool]
+) -> bool:
+    """Core ML tensors have at most rank 5.
+
+    coremltools 9.0 raises "Core ML only supports tensors with rank <= 5" for
+    any tensor above rank 5 that it cannot rewrite away. It only rewrites a
+    static, unbranched layout chain, so such a tensor is fine only when every
+    node connected to it through tensors above rank 5 is a delegated view,
+    permute or clone, and every node producing such a tensor has a static shape
+    and a single user.
+    """
+    seen, todo = set(), [node]
+    while todo:
+        n = todo.pop()
+        if n in seen:
+            continue
+        seen.add(n)
+        high_rank_inputs = [
+            arg for arg in n.all_input_nodes if _produces_tensor_above_rank_5(arg)
+        ]
+        high_rank_output = _produces_tensor_above_rank_5(n)
+        if not high_rank_inputs and not high_rank_output:
+            continue
+        if (
+            n.op != "call_function"
+            or n.target not in _HIGH_RANK_LAYOUT_TARGETS
+            or is_skipped(n)
+        ):
+            return True
+        if high_rank_output and (
+            len(n.users) != 1
+            or any(isinstance(d, torch.SymInt) for d in n.meta["val"].shape)
+        ):
+            return True
+        todo += high_rank_inputs
+        if high_rank_output:
+            todo += list(n.users)
+    return False
+
+
 _ARG_MIN_MAX_TARGETS = (
     torch.ops.aten.argmax.default,
     torch.ops.aten.argmin.default,
@@ -83,6 +142,18 @@ def _is_arg_min_max_over_flattened_input(node: torch.fx.Node) -> bool:
     return dim is None
 
 
+def _reads_or_produces_bfloat16(node: torch.fx.Node) -> bool:
+    """coremltools 9.0 has no bfloat16 type.
+
+    It raises ``KeyError: torch.bfloat16`` on any op that reads or produces a
+    bfloat16 tensor, instead of reporting the op as unsupported.
+    """
+    val = node.meta.get("val")
+    vals = list(val) if isinstance(val, (list, tuple)) else [val]
+    vals += [arg.meta.get("val") for arg in node.all_input_nodes]
+    return any(isinstance(v, torch.Tensor) and v.dtype == torch.bfloat16 for v in vals)
+
+
 def _is_view_op(op: torch._ops.OpOverload) -> bool:
     schema = op._schema
     if len(schema.arguments) == 0:
@@ -152,13 +223,32 @@ def should_override_support(self, node) -> bool:
             )
             return True
 
+        if _touches_unsupported_high_rank_tensor(
+            node,
+            lambda n: self.should_skip_op_for_delegation(
+                getattr(n.target, "__name__", "").lower()
+            ),
+        ):
+            self.log_once(
+                f"{node.target} on a tensor above rank 5 is not supported by "
+                "CoreML.  Overriding op support."
+            )
+            return True
+
         if _is_arg_min_max_over_flattened_input(node):
             self.log_once(
                 "torch.ops.aten.{argmax, argmin}.default with dim=None is "
                 "not supported by CoreML.  Overriding op support."
             )
             return True
 
+        if _reads_or_produces_bfloat16(node):
+            self.log_once(
+                f"{node.target} on a bfloat16 tensor is not supported by CoreML.  "
+                "Overriding op support."
+            )
+            return True
+
         # 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
@@ -362,21 +452,26 @@ def ops_to_not_decompose(
             torch.ops.aten.repeat_interleave.self_int,
             torch.ops.aten.repeat_interleave.self_Tensor,
         ]
+        rejected = set()
         for node in ep.graph.nodes:
             if node.op == "call_function" and isinstance(
                 node.target, torch._ops.OpOverload
             ):
                 try:
-                    if (
-                        op_support.is_node_supported(None, node)
-                        and node.target not in do_not_decompose_blocklist
-                        and not _is_view_op(node.target)
-                    ):
-                        do_not_decompose.append(node.target)
+                    supported = op_support.is_node_supported(None, node)
                 except Exception as e:
                     # CoreML's op_support.is_node_supported will sometimes throw
                     # for unsupported ops, rather than returning False
                     self.log_once(
                         f"Encountered exception when checking node support, treating node as unsupported: {e}"
                     )
-        return do_not_decompose, None
+                    supported = False
+                if not supported:
+                    rejected.add(node.target)
+                elif node.target not in do_not_decompose_blocklist and not _is_view_op(
+                    node.target
+                ):
+                    do_not_decompose.append(node.target)
+        # Without a per-node filter, a preserved op stays whole at every call, so
+        # one rejected call would be left with no kernel.
+        return [op for op in do_not_decompose if op not in rejected], None

backends/apple/coreml/test/test_coreml_partitioner.py

@@ -386,6 +386,61 @@ def forward(self, x):
         self.assertIn("executorch_call_delegate", op_names)
         self.assertNotIn("aten.randn.default", op_names)
 
+    def test_ops_above_rank_5_are_skipped(self):
+        """
+        Core ML tensors have at most rank 5.  coremltools only rewrites an
+        unbranched view, permute, view chain through a higher rank (a window
+        partition), and raises on any other op above rank 5.  The partitioner
+        must leave those ops to the portable backend and keep such a layout
+        chain delegated.
+        """
+
+        class Model(torch.nn.Module):
+            def forward(self, x):
+                windows = x.view(2, 2, 2, 2, 2, 8).permute(0, 1, 3, 2, 4, 5)
+                windows = windows.reshape(-1, 4, 8).relu()
+                branched = x.view(2, 2, 2, 2, 2, 8).permute(0, 1, 3, 2, 4, 5)
+                branched = (
+                    branched.reshape(-1, 4, 8).sum() + branched.reshape(32, 8).sum()
+                )
+                stacked = torch.stack([x.unsqueeze(0), x.unsqueeze(0) * 2])
+                return windows.sum() + branched + stacked.sum(0).relu().sum()
+
+        model = Model().eval()
+        x = torch.randn(2, 4, 4, 8)
+        ep = torch.export.export(model, (x,), strict=True)
+        edge = executorch.exir.to_edge_transform_and_lower(
+            ep, partitioner=[CoreMLPartitioner()]
+        )
+        op_names = [
+            n.target.__name__
+            for n in edge.exported_program().graph.nodes
+            if n.op == "call_function"
+        ]
+        self.assertIn("executorch_call_delegate", op_names)
+        self.assertIn("aten.sum.dim_IntList", op_names)
+        self.assertEqual(op_names.count("aten.permute_copy.default"), 1)
+        if _TEST_RUNTIME:
+            program = Runtime.get().load_program(edge.to_executorch().buffer)
+            out = program.load_method("forward").execute((x,))[0]
+            self.assertTrue(torch.allclose(out, model(x), atol=5e-2, rtol=5e-2))
+
+        # A skipped permute leaves the chain's views with no rank 6 rewrite.
+        edge = executorch.exir.to_edge_transform_and_lower(
+            ep,
+            partitioner=[
+                CoreMLPartitioner(
+                    skip_ops_for_coreml_delegation=["aten.permute_copy.default"]
+                )
+            ],
+        )
+        op_names = [
+            n.target.__name__
+            for n in edge.exported_program().graph.nodes
+            if n.op == "call_function"
+        ]
+        self.assertEqual(op_names.count("aten.permute_copy.default"), 2)
+
     def test_argmax_argmin_dim_none_is_skipped(self):
         """
         Regression test for https://github.com/pytorch/executorch/issues/11715.
@@ -433,6 +488,39 @@ def forward(self, x):
         self.assertIn("executorch_call_delegate", op_names)
         self.assertNotIn("aten.argmax.default", op_names)
 
+    def test_bfloat16_ops_are_skipped(self):
+        """
+        coremltools has no bfloat16 type and raises on ops that use it.  The
+        partitioner must leave those ops to the portable backend and still
+        delegate the rest, also when the same op has a float32 call that Core
+        ML takes (as under bfloat16 autocast).
+        """
+
+        class Model(torch.nn.Module):
+            def forward(self, q):
+                a = torch.nn.functional.scaled_dot_product_attention(q, q, q)
+                b = q.to(torch.bfloat16)
+                c = torch.nn.functional.scaled_dot_product_attention(b, b, b)
+                return a + c.float()
+
+        model = Model().eval()
+        x = torch.randn(1, 2, 4, 8)
+        ep = torch.export.export(model, (x,), strict=True)
+        edge = executorch.exir.to_edge_transform_and_lower(
+            ep, partitioner=[CoreMLPartitioner()]
+        )
+        op_names = [
+            n.target.__name__
+            for n in edge.exported_program().graph.nodes
+            if n.op == "call_function"
+        ]
+        self.assertIn("executorch_call_delegate", op_names)
+        self.assertNotIn("aten.scaled_dot_product_attention.default", op_names)
+        if _TEST_RUNTIME:
+            program = Runtime.get().load_program(edge.to_executorch().buffer)
+            out = program.load_method("forward").execute((x,))[0]
+            self.assertTrue(torch.allclose(out, model(x), atol=5e-2, rtol=5e-2))
+
     def test_tied_embedding_is_quantized_with_its_linear(self):
         """A weight that is both an embedding table and an output projection lowers.