Core ML: run ops that read uint8 and int8 inputs on the CPU, refuse ops coremltools gets wrong, and fix fills, tanh GELU, RMSNorm and repeated outputs

backends/apple/coreml/compiler/coreml_preprocess.py

@@ -17,6 +17,9 @@
 
 import coremltools as ct
 import coremltools.optimize as cto
+from coremltools.converters.mil import Builder as mb
+from coremltools.converters.mil.mil.passes.graph_pass import AbstractGraphPass
+from coremltools.converters.mil.mil.passes.pass_registry import register_pass
 from executorch.backends.apple.coreml import executorchcoreml
 from executorch.backends.apple.coreml.compiler.enumerated_shape_utils import (
     _get_ct_inputs,
@@ -67,6 +70,33 @@ def _gathers_sharing_a_weight(mlmodel: ct.models.MLModel) -> List[str]:
     return shared
 
 
+@register_pass(namespace="executorch", name="copy_reused_outputs")
+class _CopyReusedOutputs(AbstractGraphPass):
+    """
+    Gives each output of the program a value of its own.
+
+    coremltools converts ops such as clone into nothing, so a partition can
+    return one of its own inputs, or one value twice. Core ML then hands back
+    such an output wrong, such as all zeros or the bytes of another output.
+    """
+
+    def apply(self, prog):
+        for function in prog.functions.values():
+            seen = set(function.inputs.values())
+            outputs = []
+            for output in function.outputs:
+                if output in seen:
+                    with function:
+                        copy = mb.identity(x=output)
+                    # coremltools wants every op to carry the source program's
+                    # scopes, which set_outputs, unlike a var replacement, skips.
+                    function._copy_scope_info(output, copy)
+                    output = copy
+                seen.add(output)
+                outputs.append(output)
+            function.set_outputs(outputs)
+
+
 class COMPILE_SPEC_KEYS(Enum):
     COMPUTE_UNITS = "compute_units"
     MODEL_TYPE = "model_type"
@@ -565,6 +595,7 @@ def _convert_to_mlmodel(
         )
         compute_units = CoreMLBackend.compute_unit_from_compile_specs(compile_specs)
         pass_pipeline = CoreMLBackend.pass_pipeline_from_compile_specs(compile_specs)
+        pass_pipeline.append_pass("executorch::copy_reused_outputs")
         enumerated_shapes = CoreMLBackend.enumerated_shapes_from_compile_specs(
             compile_specs
         )

backends/apple/coreml/compiler/torch_ops.py

@@ -11,10 +11,13 @@
 import numpy as np
 import torch as _torch
 from coremltools import _logger
+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
 from coremltools.converters.mil.frontend.torch.ops import (
     _get_inputs,
     _get_kwinputs,
+    _make_fill_op,
     noop,
     NUM_TO_NUMPY_DTYPE,
     NUM_TO_TORCH_DTYPE,
@@ -49,6 +52,113 @@ def split_copy(context, node):
     split(context, node)
 
 
+# https://github.com/apple/coremltools/pull/2863
+# coremltools reads approximate only from the positional inputs, but torch.export
+# passes it as a keyword, so gelu(approximate="tanh") became the exact gelu.
+@register_torch_op(override=True)
+def gelu(context, node):
+    inputs = _get_inputs(context, node, expected=[1, 2])
+    if len(inputs) == 2:
+        approximate = inputs[1].val
+    else:
+        # Read the keyword as is: _get_kwinputs would bind "tanh" to a node named tanh.
+        approximate = (node.kwinputs or {}).get("approximate", ["none"])[0]
+    mode = {"none": "EXACT", "tanh": "TANH_APPROXIMATION"}[approximate]
+    context.add(mb.gelu(x=inputs[0], mode=mode, name=node.name))
+
+
+def _torch_dtype(node):
+    """The dtype torch gives the output of node, if node carries it, as a node of
+    an ExportedProgram does."""
+    val = (node.meta or {}).get("val")
+    return val.dtype if isinstance(val, _torch.Tensor) else None
+
+
+def _core_ml_type(node):
+    """The numpy type Core ML uses for the dtype torch gives the output of node."""
+    dtype = _torch_dtype(node)
+    if dtype == _torch.bool:
+        return np.bool_
+    if dtype == _torch.float16:
+        return np.float16
+    return np.float32 if dtype.is_floating_point else np.int32
+
+
+def _add_fill(context, node, shape, value):
+    context.add(_make_fill_op(shape, _core_ml_type(node)(value), node.name))
+
+
+# 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. Not fixed in coremltools yet, except for
+# zeros_like in https://github.com/apple/coremltools/pull/2865. A TorchScript
+# node, which carries no dtype, and a fill value known only at run time keep
+# the coremltools converter.
+@register_torch_op(override=True)
+def full(context, node):
+    inputs = _get_inputs(context, node, min_expected=2)
+    if _torch_dtype(node) is None or inputs[1].val is None:
+        return _ct_ops.full(context, node)
+    _add_fill(context, node, inputs[0], inputs[1].val)
+
+
+@register_torch_op(override=True)
+def new_full(context, node):
+    inputs = _get_inputs(context, node, min_expected=3)
+    if _torch_dtype(node) is None or inputs[2].val is None:
+        return _ct_ops.new_full(context, node)
+    _add_fill(context, node, inputs[1], inputs[2].val)
+
+
+@register_torch_op(torch_alias=["fill"], override=True)
+def full_like(context, node):
+    inputs = _get_inputs(context, node, min_expected=2)
+    if _torch_dtype(node) is None or inputs[1].val is None:
+        return _ct_ops.full_like(context, node)
+    _add_fill(context, node, mb.shape(x=inputs[0]), inputs[1].val)
+
+
+# Lowering turns empty, new_empty and empty_like into _empty_dim_order, which
+# coremltools converts with its dtype, so these aliases matter only when dim
+# order ops are skipped.
+@register_torch_op(torch_alias=["empty", "empty.memory_format"], override=True)
+def zeros(context, node):
+    if _torch_dtype(node) is None:
+        return _ct_ops.zeros(context, node)
+    _add_fill(context, node, _get_inputs(context, node, min_expected=1)[0], 0)
+
+
+@register_torch_op(torch_alias=["new_empty"], override=True)
+def new_zeros(context, node):
+    if _torch_dtype(node) is None:
+        return _ct_ops.new_zeros(context, node)
+    _add_fill(context, node, _get_inputs(context, node, min_expected=2)[1], 0)
+
+
+@register_torch_op(torch_alias=["empty_like"], override=True)
+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)
+
+
+@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)
+
+
+@register_torch_op(override=True)
+def scalar_tensor(context, node):
+    x = _get_inputs(context, node, expected=[1, 5])[0]
+    if _torch_dtype(node) is None or x.val is None:
+        return _ct_ops.scalar_tensor(context, node)
+    context.add(mb.const(val=_core_ml_type(node)(x.val), name=node.name))
+
+
 def is_fbcode():
     return not hasattr(_torch.version, "git_version")
 
@@ -58,6 +168,7 @@ def is_fbcode():
         _empty_dim_order,
         _to_dim_order_copy,
     )
+    from coremltools.converters.mil.frontend.torch.ops import rms_norm as _rms_norm
 
     # This is a temporary hack to register the alias "dim_order_ops._to_dim_order_copy",
     # which was missed by coremltools
@@ -71,6 +182,18 @@ def _to_dim_order_copy_TMP_EXECUTORCH_ALIAS_HACK(context, node):
     def _empty_dim_order_TMP_EXECUTORCH_ALIAS_HACK(context, node):
         _empty_dim_order(context, node)
 
+    # coremltools 9.0 requires the eps that torch.nn.RMSNorm(eps=None) leaves out.
+    # Torch then uses the float32 machine epsilon, also for float16 inputs, while
+    # https://github.com/apple/coremltools/pull/2743 would use 1e-5, so keep this
+    # override after that change lands.
+    @register_torch_op(override=True)
+    def rms_norm(context, node):
+        inputs = list(node.inputs) + [None] * (4 - len(node.inputs))
+        if inputs[3] is None:
+            inputs[3] = float(np.finfo(np.float32).eps)
+        node.inputs = inputs
+        _rms_norm(context, node)
+
 else:
     # TODO: remove this case when fbcode updates to coremltools 9.0
     @register_torch_op(