Refuse ops coremltools cannot convert in the Core ML partitioner, and say why

backends/apple/coreml/partition/coreml_partitioner.py

@@ -26,6 +26,7 @@
 from executorch.exir.backend.utils import tag_constant_data, tag_mutated_buffer
 from executorch.exir.dialects._ops import ops as exir_ops
 from torch.export.exported_program import ExportedProgram
+from torch.fx.experimental.symbolic_shapes import statically_known_true, sym_eq
 from torch.fx.passes.infra.partitioner import CapabilityBasedPartitioner
 from torch.fx.passes.operator_support import OperatorSupportBase
 
@@ -162,18 +163,141 @@ def _is_view_op(op: torch._ops.OpOverload) -> bool:
     return (alias_info is not None) and (not alias_info.is_write)
 
 
+_MIN_MAX_TARGETS = (
+    torch.ops.aten.min.default,
+    torch.ops.aten.min.dim,
+    torch.ops.aten.min.other,
+    torch.ops.aten.max.default,
+    torch.ops.aten.max.dim,
+    torch.ops.aten.max.other,
+    torch.ops.aten.amin.default,
+    torch.ops.aten.amax.default,
+    torch.ops.aten.minimum.default,
+    torch.ops.aten.maximum.default,
+    exir_ops.edge.aten.min.default,
+    exir_ops.edge.aten.min.dim,
+    exir_ops.edge.aten.min.other,
+    exir_ops.edge.aten.max.default,
+    exir_ops.edge.aten.max.dim,
+    exir_ops.edge.aten.max.other,
+    exir_ops.edge.aten.amin.default,
+    exir_ops.edge.aten.amax.default,
+    exir_ops.edge.aten.minimum.default,
+    exir_ops.edge.aten.maximum.default,
+)
+
+
+# No upstream issue yet. For example, converting x.amax() of a bool x raises
+# 'Op "amax" (op_type: reduce_max) Input x="x" expects tensor or scalar of
+# dtype from type domain ['fp16', 'fp32', 'int32'] but got tensor[3,2,bool]'.
+def _is_min_max_of_bool(node: torch.fx.Node) -> bool:
+    """coremltools 9.0 cannot compute a min or max of bool tensors.
+
+    It raises ``ValueError`` when every input is bool, so that the result is
+    bool, such as for the all-true check that torch.distributions exports,
+    instead of reporting the op as unsupported. A bool next to a float or int
+    tensor works, and so do int tensors.
+    """
+    if node.target not in _MIN_MAX_TARGETS:
+        return False
+    val = node.meta.get("val")
+    if isinstance(val, (list, tuple)):
+        val = val[0]
+    return isinstance(val, torch.Tensor) and val.dtype == torch.bool
+
+
+_INDEX_PUT_TARGETS = (
+    torch.ops.aten.index_put.default,
+    exir_ops.edge.aten.index_put.default,
+)
+
+
+# https://github.com/apple/coremltools/issues/2141 covers a bool index, which
+# coremltools 9.0 converts only on its own. No upstream issue covers the other
+# forms yet. For example, y[:, idx] = v with idx = torch.tensor([0, 2]) raises
+# "Input squeeze_0 has rank 1 != other inputs rank 0".
+def _index_put_refusal(
+    node: torch.fx.Node, minimum_deployment_target: Optional[ct.target]
+) -> Optional[str]:
+    """The form of an ``index_put`` that coremltools 9.0 cannot convert, such
+    as "with accumulate=True next to a slice", or None for any other node.
+
+    With a whole slice, such as ``y[:, idx] = v``, it writes ``v`` into the
+    part of ``y`` the indices select, keeping each indexed dim as size 1. That
+    needs ``accumulate=False``, int indices of shape ``()`` or ``(1,)``, a
+    static size in every other dim, and a non-bool ``y`` below iOS 18. For a
+    size known only at run time, it raises "Cannot add const". With a 0-d
+    index or several indices, torch lays out that part differently, so ``v``
+    may only vary in the dims after the last index. Without a whole slice,
+    several indices need the same shape and no bool mask. It raises or writes
+    the wrong place for any other form, such as ``idx`` with two entries,
+    instead of reporting the op as unsupported.
+    """
+    if node.target not in _INDEX_PUT_TARGETS:
+        return None
+    y, values = node.args[0].meta.get("val"), node.args[2].meta.get("val")
+    indices = [i if i is None else i.meta.get("val") for i in node.args[1]]
+    given = {dim: index for dim, index in enumerate(indices) if index is not None}
+    if not given or not all(
+        isinstance(t, torch.Tensor) for t in (y, values, *given.values())
+    ):
+        return None
+    if len(given) == len(indices):
+        if len(indices) > 1 and any(
+            index.dtype == torch.bool
+            or not statically_known_true(sym_eq(index.shape, indices[0].shape))
+            for index in indices
+        ):
+            return "with indices that differ in shape or include a bool mask"
+        return None
+    accumulate = (
+        node.args[3] if len(node.args) > 3 else node.kwargs.get("accumulate", False)
+    )
+    if accumulate:
+        return "with accumulate=True next to a slice"
+    if any(
+        index.dtype == torch.bool
+        or index.dim() > 1
+        or not statically_known_true(index.numel() == 1)
+        for index in given.values()
+    ):
+        return (
+            "with an index other than an int tensor of shape () or (1,) next to a slice"
+        )
+    selected = [1 if dim in given else size for dim, size in enumerate(y.shape)]
+    if any(isinstance(size, torch.SymInt) for size in selected):
+        return "with a size known only at run time next to a slice"
+    reorders = len(given) > 1 or any(index.dim() == 0 for index in given.values())
+    after = len(selected) - 1 - max(given)
+    if reorders and not all(
+        statically_known_true(v == 1)
+        for v in values.shape[: max(0, values.dim() - after)]
+    ):
+        return (
+            "with a 0-d index or several indices, and values that vary before "
+            "the last index"
+        )
+    if y.dtype == torch.bool and (
+        minimum_deployment_target is None or minimum_deployment_target < ct.target.iOS18
+    ):
+        return "on a bool tensor with a minimum_deployment_target below iOS 18"
+    return None
+
+
 class _OperatorsSupportedForCoreMLBackend(OperatorSupportBase):
     def __init__(
         self,
         skip_ops_for_coreml_delegation: Optional[List[str]] = None,
         lower_full_graph: bool = False,
         log: bool = False,
+        minimum_deployment_target: Optional[ct.target] = None,
     ) -> None:
         if skip_ops_for_coreml_delegation is None:
             skip_ops_for_coreml_delegation = []
         super().__init__()
         self.skip_ops_for_coreml_delegation = skip_ops_for_coreml_delegation
         self.lower_full_graph = lower_full_graph
+        self.minimum_deployment_target = minimum_deployment_target
         self._logged_msgs = set()
         self._log = log
 
@@ -195,7 +319,9 @@ def should_skip_op_for_delegation(self, node_target_name: str) -> bool:
             return True
         return False
 
-    def should_override_support(self, node) -> bool:
+    def override_reason(self, node) -> Optional[str]:
+        """Why CoreML cannot take node, even when coremltools converts its op."""
+        name = getattr(node.target, "__name__", node.target)
         # https://github.com/apple/coremltools/issues/2573
         if (
             node.target
@@ -208,46 +334,38 @@ def should_override_support(self, node) -> bool:
             and "alpha" in node.kwargs
             and node.kwargs["alpha"] != 1
         ):
-            self.log_once(
-                "torch.ops.aten.{sub, add}.Tensor with alpha != 1 is not supported by CoreML.  Overriding support."
-            )
-            return True
+            return "torch.ops.aten.{sub, add}.Tensor with alpha != 1 is not supported by CoreML."
 
         # Ops that are unsupported by CoreML purely on the basis of their
         # target — no per-arg conditions to check.  Grouped by upstream issue
         # so the comment trail still points at the underlying coremltools /
         # executorch bug for each entry.
         if node.target in _UNSUPPORTED_OP_TARGETS:
-            self.log_once(
-                f"{node.target} is not supported by CoreML.  Overriding op support."
-            )
-            return True
+            return f"{name} is not supported by CoreML."
 
         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
+            return f"{name} on a tensor above rank 5 is not supported by CoreML."
 
         if _is_arg_min_max_over_flattened_input(node):
-            self.log_once(
+            return (
                 "torch.ops.aten.{argmax, argmin}.default with dim=None is "
-                "not supported by CoreML.  Overriding op support."
+                "not supported by CoreML."
             )
-            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
+            return f"{name} on a bfloat16 tensor is not supported by CoreML."
+
+        if _is_min_max_of_bool(node):
+            return f"{name} on bool tensors is not supported by CoreML."
+
+        refusal = _index_put_refusal(node, self.minimum_deployment_target)
+        if refusal is not None:
+            return f"{name} {refusal} is not supported by CoreML."
 
         # 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
@@ -266,7 +384,7 @@ def should_override_support(self, node) -> bool:
         #     )
         #     return True
 
-        return False
+        return None
 
     def is_node_supported(self, submodules, node: torch.fx.Node) -> bool:
         # get_attr node can always be supported on any backend
@@ -284,10 +402,17 @@ def is_node_supported(self, submodules, node: torch.fx.Node) -> bool:
             is_supported = ct.converters.mil.frontend.torch.is_torch_fx_node_supported(
                 node
             )
-            if self.should_override_support(node):
+            reason = self.override_reason(node)
+            if reason is not None:
+                self.log_once(f"{reason}  Overriding op support.")
                 is_supported = False
 
             if not is_supported:
+                if self.lower_full_graph and reason is not None:
+                    raise NotImplementedError(
+                        f"{reason}  Set lower_full_graph=False in the "
+                        "CoreMLPartitioner to run it outside of CoreML."
+                    )
                 if self.lower_full_graph:
                     raise NotImplementedError(
                         f"""CoreML does not support the op {node_target_name}, but you have set lower_full_graph=True in the CoreMLPartitioner.
@@ -343,6 +468,11 @@ 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._minimum_deployment_target = (
+            CoreMLBackend.min_deployment_target_from_compile_specs(
+                self.delegation_spec.compile_specs
+            )
+        )
         self._logged_msgs = set()
 
         if self.lower_full_graph:
@@ -399,6 +529,7 @@ def partition(self, exported_program: ExportedProgram) -> PartitionResult:
                 self.skip_ops_for_coreml_delegation,
                 self.lower_full_graph,
                 log=True,
+                minimum_deployment_target=self._minimum_deployment_target,
             ),
             allows_single_node_partition=True,
         )
@@ -438,6 +569,7 @@ def ops_to_not_decompose(
             self.skip_ops_for_coreml_delegation,
             self.lower_full_graph,
             log=False,
+            minimum_deployment_target=self._minimum_deployment_target,
         )
 
         # CoreML prevents certain ops (like triu) from lowering to CoreML when put in the ExecuTorch op namespace