Core ML: run ops on the CPU when Core ML cannot take or return their values (#24317)

## What is wrong today

coremltools cannot take a uint8 or int8 tensor as a Core ML input, such
as camera pixels or int8 tensors in a quantized model. Lowering fails
with `KeyError: torch.uint8`. An input sized by an expression such as
`s0 + 1` fails too, unless the program declares it.

Some values cannot leave Core ML either: the delegate has no int8 data
type and returns only tensors, so neither an int8 tensor nor a symbolic
value, such as a size, can be an output.

`ops_to_not_decompose` can also keep an op whole that then runs on the
CPU, where the portable kernels may lack it, such as the matmul of a
bfloat16 linear.

`zeros_like` and `fill` put a constant as big as their input into the
model when its shape is known, and since #24424 `ones_like` and
`full_like` do too.

## What this change does

- The partitioner leaves on the CPU each op that reads such a tensor, or
a symbolic value computed on the CPU, from outside its partition, and
each op that makes an int8 tensor or a symbolic value used outside its
partition. Weights, constants and state that Core ML takes over are not
inputs. The search repeats as ops move out, and partitions left reading
only sizes, or nothing below iOS 18, move out too.
- State stays in the program when no op reads it, or when a CPU op or a
second partition also uses it or the model returns its new value. An op
that reads such uint8 or int8 state, or makes its new int8 value, moves
to the CPU.
- With `lower_full_graph=True`, a Core ML part that would return an int8
tensor now fails while lowering, with the reason, not at run time.
- `ops_to_not_decompose` stops keeping an op whole when the partitioner
moves some call of it to the CPU, or when Core ML rejects a call that
decomposing makes. It decomposes the program as lowering will only when
that may change the answer.
- From iOS 16, those four fills use Core ML's `fill_like` instead.
- The docs list the values Core ML cannot take or return, and point
camera apps to the `ImageProcessor` in `extension/image`.

## Testing

- `test_coreml_partitioner.py`, `test_torch_ops.py` and
`test_coreml_recipes.py` pass on Linux and macOS, where most new tests
also compare outputs with eager PyTorch.
- Most new tests fail on main.

Author: shoumikhin

backends/apple/coreml/compiler/torch_ops.py

@@ -13,7 +13,7 @@
 
 import numpy as np
 import torch as _torch
-from coremltools import _logger
+from coremltools import _logger, target
 from coremltools.converters.mil import Builder as mb
 from coremltools.converters.mil.frontend import _utils
 from coremltools.converters.mil.frontend.torch import ops as _ct_ops
@@ -31,6 +31,9 @@
     register_torch_op,
 )
 from coremltools.converters.mil.mil import types
+from coremltools.converters.mil.mil.block import (
+    is_current_opset_version_compatible_with,
+)
 from coremltools.converters.mil.mil.types.symbolic import is_symbolic
 from executorch.exir.dim_order_utils import get_memory_format
 
@@ -92,6 +95,16 @@ def _add_fill(context, node, shape, value):
     context.add(_make_fill_op(shape, _core_ml_type(node)(value), node.name))
 
 
+def _add_fill_like(context, node, x, value):
+    # A fill of a shape known at conversion becomes a constant as big as x in
+    # the model, which fill_like avoids from iOS 16.
+    if is_current_opset_version_compatible_with(target.iOS16):
+        value = _core_ml_type(node)(value)
+        context.add(mb.fill_like(ref_tensor=x, value=value, name=node.name))
+    else:
+        _add_fill(context, node, mb.shape(x=x), value)
+
+
 # coremltools gives these fills the type of their fill value or of their input,
 # or makes them float, instead of the dtype torch gives them, so an int or bool
 # fill could come out as float. https://github.com/apple/coremltools/pull/2865
@@ -119,7 +132,7 @@ def _fill_like(context, node, convert):
     inputs = _get_inputs(context, node, min_expected=2)
     if _torch_dtype(node) is None or inputs[1].val is None:
         return convert(context, node)
-    _add_fill(context, node, mb.shape(x=inputs[0]), inputs[1].val)
+    _add_fill_like(context, node, inputs[0], inputs[1].val)
 
 
 @register_torch_op(override=True)
@@ -153,15 +166,15 @@ def zeros_like(context, node):
     if _torch_dtype(node) is None:
         return _ct_ops.zeros_like(context, node)
     x = _get_inputs(context, node, min_expected=1)[0]
-    _add_fill(context, node, mb.shape(x=x), 0)
+    _add_fill_like(context, node, x, 0)
 
 
 @register_torch_op(override=True)
 def ones_like(context, node):
     if _torch_dtype(node) is None:
         return _ct_ops.ones_like(context, node)
     x = _get_inputs(context, node, min_expected=1)[0]
-    _add_fill(context, node, mb.shape(x=x), 1)
+    _add_fill_like(context, node, x, 1)
 
 
 @register_torch_op(override=True)