diff --git a/backends/apple/coreml/compiler/coreml_preprocess.py b/backends/apple/coreml/compiler/coreml_preprocess.py index 931d1552606..8afb334577d 100644 --- a/backends/apple/coreml/compiler/coreml_preprocess.py +++ b/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 ) diff --git a/backends/apple/coreml/compiler/torch_ops.py b/backends/apple/coreml/compiler/torch_ops.py index 7f6abe980e1..f1f87d17e1a 100644 --- a/backends/apple/coreml/compiler/torch_ops.py +++ b/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( diff --git a/backends/apple/coreml/partition/coreml_partitioner.py b/backends/apple/coreml/partition/coreml_partitioner.py index 4453818f956..6d8b4f9e104 100644 --- a/backends/apple/coreml/partition/coreml_partitioner.py +++ b/backends/apple/coreml/partition/coreml_partitioner.py @@ -4,7 +4,9 @@ import inspect import logging -from typing import Callable, List, Optional, Tuple +import operator +import weakref +from typing import Any, Callable, Collection, List, Optional, Set, Tuple import coremltools as ct @@ -25,9 +27,21 @@ ) from executorch.exir.backend.utils import tag_constant_data, tag_mutated_buffer from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.passes.dim_order_ops_registry import DimOrderOpsMap +from executorch.exir.passes.replace_aten_with_edge_pass import aten_to_edge +from executorch.exir.tracer import _default_decomposition_table +from torch._export.passes.replace_view_ops_with_view_copy_ops_pass import ( + get_view_copy_of_view_op, +) from torch.export.exported_program import ExportedProgram -from torch.fx.passes.infra.partitioner import CapabilityBasedPartitioner -from torch.fx.passes.operator_support import OperatorSupportBase +from torch.export.graph_signature import OutputKind +from torch.fx.experimental.symbolic_shapes import statically_known_true, sym_eq +from torch.fx.passes.infra.partitioner import CapabilityBasedPartitioner, Partition +from torch.fx.passes.operator_support import ( + chain, + create_op_support, + OperatorSupportBase, +) logger = logging.getLogger(__name__) logger.setLevel(get_coreml_log_level(default_level=logging.INFO)) @@ -154,6 +168,286 @@ def _reads_or_produces_bfloat16(node: torch.fx.Node) -> bool: return any(isinstance(v, torch.Tensor) and v.dtype == torch.bfloat16 for v in vals) +_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, +) + + +def _is_min_max_of_bool(node: torch.fx.Node) -> bool: + """coremltools 9.0 cannot compute a min or max over a bool tensor. + + It raises ``ValueError`` for one, such as the all-true check that + torch.distributions exports, instead of reporting the op as unsupported. + Integer tensors work. + """ + if node.target not in _MIN_MAX_TARGETS: + return False + val = node.args[0].meta.get("val") + 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, +) + + +def _is_index_put_coremltools_cannot_convert(node: torch.fx.Node, ios18: bool) -> bool: + """coremltools 9.0 converts only some forms of ``index_put``. + + 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, a ``v`` that broadcasts to that part, and + a non-bool ``y`` below iOS 18. 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. + + This check cannot see index values, so a negative index next to a slice + stays in Core ML: it fails, or gives wrong values when it comes from an + input. + """ + if node.target not in _INDEX_PUT_TARGETS: + return False + indices = [i if i is None else i.meta["val"] for i in node.args[1]] + given = {dim: index for dim, index in enumerate(indices) if index is not None} + if len(given) == len(indices): + return 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 + ) + accumulate = ( + node.args[3] if len(node.args) > 3 else node.kwargs.get("accumulate", False) + ) + y = node.args[0].meta["val"] + if ( + accumulate + or (y.dtype == torch.bool and not ios18) + or any( + index.dtype == torch.bool + or index.dim() > 1 + or not statically_known_true(index.numel() == 1) + for index in given.values() + ) + ): + return True + selected = [1 if dim in given else size for dim, size in enumerate(y.shape)] + values = node.args[2].meta["val"].shape + reorders = len(given) > 1 or any(index.dim() == 0 for index in given.values()) + after = len(selected) - 1 - max(given) + return ( + any(isinstance(size, torch.SymInt) for size in selected) + or len(values) > len(selected) + or not all( + statically_known_true(v == 1) or statically_known_true(v == s) + for v, s in zip(reversed(values), reversed(selected)) + ) + or ( + reorders + and not all( + statically_known_true(v == 1) + for v in values[: max(0, len(values) - after)] + ) + ) + ) + + +def _makes_tensor( + node: torch.fx.Node, dtypes: Optional[Collection[torch.dtype]] = None +) -> bool: + """Whether node makes a tensor, of one of dtypes if given.""" + val = node.meta.get("val") + vals = val if isinstance(val, (list, tuple)) else [val] + return any( + isinstance(v, torch.Tensor) and (dtypes is None or v.dtype in dtypes) + for v in vals + ) + + +def _cannot_be_coreml_input(node: torch.fx.Node) -> bool: + """Whether node makes a uint8 or int8 tensor, which Core ML cannot take as a + program input. + + coremltools types each program input through ``TORCH_DTYPE_TO_MIL_DTYPE`` + and raises ``KeyError`` for a dtype missing from it instead of reporting the + op as unsupported. Only uint8 and int8 move to the CPU: the portable kernels + handle them, while other missing dtypes, such as uint16, would only fail + later, at run time. + """ + return _makes_tensor(node, (torch.uint8, torch.int8)) + + +def _find_readers_of_unsupported_inputs( + partitions: List[Partition], + kept_inside: Collection[str], + left_out: Collection[torch.fx.Node], +) -> Set[torch.fx.Node]: + """The nodes that read a tensor Core ML cannot take as an input from outside + their partition, or a symbolic value computed by an op left out of Core ML. + + A tensor made inside the partition, or a weight, constant or state the + partitioner moves into Core ML, is not an input. + """ + return { + node + for partition in partitions + for node in partition.nodes + if any( + arg not in partition.nodes + and arg.name not in kept_inside + and ( + _cannot_be_coreml_input(arg) + or ( + arg in left_out + and isinstance( + arg.meta.get("val"), + (torch.SymInt, torch.SymFloat, torch.SymBool), + ) + ) + ) + for arg in node.all_input_nodes + ) + } + + +def _find_makers_of_int8_outputs(partitions: List[Partition]) -> Set[torch.fx.Node]: + """The nodes that make an int8 tensor read outside their partition. + + Such a tensor becomes a Core ML output, and the Core ML runtime cannot + return int8. An op with several outputs hands each one out through a + getitem, which always stays with the op, so the op itself is returned. + """ + return { + node.args[0] if node.target is operator.getitem else node + for partition in partitions + for node in partition.nodes + if _makes_tensor(node, (torch.int8,)) + and any(user not in partition.nodes for user in node.users) + } + + +def _find_nodes_of_partitions_without_tensor_inputs( + partitions: List[Partition], + kept_inside: Collection[str], + allow_partitions_without_inputs: bool, +) -> Set[torch.fx.Node]: + """The nodes of the partitions that read no tensor from outside, except the + weights, constants and state the partitioner moves into Core ML. + + Core ML cannot take a symbolic size as an input, and below iOS 18 it + rejects a model with no inputs, so such a partition cannot be delegated. + """ + nodes = set() + for partition in partitions: + inputs = [ + arg + for node in partition.nodes + for arg in node.all_input_nodes + if arg not in partition.nodes and arg.name not in kept_inside + ] + reads_tensor = any(_makes_tensor(arg) for arg in inputs) + if not reads_tensor and (inputs or not allow_partitions_without_inputs): + nodes.update(partition.nodes) + return nodes + + +def _propose_partitions( + graph_module: torch.fx.GraphModule, + op_support: OperatorSupportBase, + kept_inside: Collection[str], + allow_partitions_without_inputs: bool, +) -> Tuple[List[Partition], Set[torch.fx.Node]]: + """Propose Core ML partitions, leaving out every op that reads a tensor + Core ML cannot take as an input from outside its partition or makes one it + cannot return, and then every partition this leaves without a tensor input. + + Returns the partitions and the ops left out. + """ + left_out: Set[torch.fx.Node] = set() + support = chain(create_op_support(lambda _, node: node not in left_out), op_support) + # Each op moved out of Core ML can hand such a tensor to the next op, so + # repeat until no partition reads one. + while True: + partitions = CapabilityBasedPartitioner( + graph_module, support, allows_single_node_partition=True + ).propose_partitions() + new = _find_readers_of_unsupported_inputs( + partitions, kept_inside, left_out + ) | _find_makers_of_int8_outputs(partitions) + if new or left_out: + # Moving those ops out can leave a partition that reads no tensor, + # such as the index math of a gather. + new |= _find_nodes_of_partitions_without_tensor_inputs( + partitions, kept_inside, allow_partitions_without_inputs + ) + # Stop once a round leaves out nothing new, so the loop always ends. + new -= left_out + if not new: + return partitions, left_out + left_out |= new + + +def _reads_from_a_list(user: torch.fx.Node, node: torch.fx.Node) -> bool: + return any( + isinstance(arg, (list, tuple)) and any(item is node for item in arg) + for arg in (*user.args, *user.kwargs.values()) + ) + + +def _untag_constants_shared_through_lists(ep: ExportedProgram) -> None: + """Untags each constant that ops in more than one place, CoreML partitions + or the CPU, read, when one of them reads it from a list, such as the + indices of an ``index_put``. + + to_backend gives each place its own copy of a shared constant, but it only + swaps in the copy where an op takes the constant directly, so a constant + read from a list fails to lower. An untagged constant is passed to each + CoreML partition as an input instead. + """ + for node in ep.graph.nodes: + if node.op != "placeholder" or "delegation_tag" not in node.meta: + continue + places = {user.meta.get("delegation_tag") for user in node.users} + if len(places) > 1 and any( + _reads_from_a_list(user, node) for user in node.users + ): + del node.meta["delegation_tag"] + + +def _edge_op_names(target: Any) -> Set[str]: + """The lowercase names target can have once lowered to the edge dialect, + which partition() matches skip names against: an ATen op gains its + namespace, a view op becomes its copy, and some ops become dim order ops. + """ + if not isinstance(target, torch._ops.OpOverload): + return {getattr(target, "__name__", "").lower()} + edge = aten_to_edge(get_view_copy_of_view_op(target._schema) or target) + return {edge.__name__.lower(), DimOrderOpsMap.get(edge, edge).__name__.lower()} + + def _is_view_op(op: torch._ops.OpOverload) -> bool: schema = op._schema if len(schema.arguments) == 0: @@ -168,12 +462,14 @@ def __init__( skip_ops_for_coreml_delegation: Optional[List[str]] = None, lower_full_graph: bool = False, log: bool = False, + ios18: bool = False, ) -> 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.ios18 = ios18 self._logged_msgs = set() self._log = log @@ -195,7 +491,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,20 +506,14 @@ 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, @@ -229,25 +521,22 @@ def should_override_support(self, node) -> bool: 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} over a bool tensor is not supported by CoreML." + + if _is_index_put_coremltools_cannot_convert(node, self.ios18): + return f"{name} with these indices 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 +555,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 +573,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. @@ -344,6 +640,7 @@ def __init__( self.lower_full_graph = lower_full_graph self.take_over_constant_data = take_over_constant_data self._logged_msgs = set() + self._ops_to_not_decompose_cache = None if self.lower_full_graph: assert ( @@ -393,24 +690,79 @@ def partition(self, exported_program: ExportedProgram) -> PartitionResult: logger.info("CoreMLPartitioner::partition") partition_tags = {} - capability_partitioner = CapabilityBasedPartitioner( + kept_inside = self._placeholders_kept_inside(exported_program) + partition_list, left_out = _propose_partitions( exported_program.graph_module, _OperatorsSupportedForCoreMLBackend( self.skip_ops_for_coreml_delegation, self.lower_full_graph, log=True, + ios18=self._targets_ios18(), ), - allows_single_node_partition=True, + kept_inside, + self._targets_ios18(), ) - partition_list = capability_partitioner.propose_partitions() + if left_out: + inputs = sorted( + { + f"{arg.name} ({arg.meta['val'].dtype})" + for node in left_out + for arg in node.all_input_nodes + if arg not in left_out + and arg.name not in kept_inside + and isinstance(arg.meta.get("val"), torch.Tensor) + and _cannot_be_coreml_input(arg) + } + ) + outputs = sorted( + node.name + for node in left_out + if _makes_tensor(node, (torch.int8,)) + and any(user not in left_out for user in node.users) + ) + reasons = [] + if inputs: + reasons.append( + "reads a uint8 or int8 tensor, which CoreML cannot take as an " + "input, from outside a CoreML partition: " + ", ".join(inputs) + ) + if outputs: + reasons.append( + "makes an int8 tensor, which CoreML cannot return, for use " + "outside of CoreML: " + ", ".join(outputs) + ) + msg = ( + f"The model {', and '.join(reasons)}." + if reasons + else "The model uses a uint8 or int8 tensor CoreML cannot handle." + ) + if self.lower_full_graph: + raise NotImplementedError( + f"{msg} Use a dtype CoreML supports for these tensors, or set " + "lower_full_graph=False in the CoreMLPartitioner to run the ops " + "that use them outside of CoreML." + ) + self.log_once( + f"{msg} Running these ops outside of CoreML: " + + ", ".join( + sorted( + { + getattr(node.target, "__name__", str(node.target)) + for node in left_out + } + ) + ) + ) for partition in partition_list: for node in partition.nodes: tag = f"tag{partition.id}" node.meta["delegation_tag"] = tag partition_tags[tag] = self.delegation_spec + self._check_state_is_used_in_one_place(exported_program) if self.take_over_constant_data: tag_constant_data(exported_program) + _untag_constants_shared_through_lists(exported_program) if self.take_over_mutable_buffer: logger.info( "Core ML partitioner will take over torch mutable buffer as Core ML state, " @@ -427,17 +779,94 @@ def partition(self, exported_program: ExportedProgram) -> PartitionResult: def log_once(self, msg: str) -> None: if msg not in self._logged_msgs: - logging.info(msg) + logger.info(msg) self._logged_msgs.add(msg) + def _check_state_is_used_in_one_place(self, ep: ExportedProgram) -> None: + """Raises if ops in more than one place, CoreML partitions or the CPU, + read or update a buffer CoreML takes over as state, since each place + would then work on its own copy of it. + """ + signature = ep.graph_signature + nodes = {node.name: node for node in ep.graph.nodes} + updates = { + spec.target: nodes[spec.arg.name] + for spec in signature.output_specs + if spec.kind == OutputKind.BUFFER_MUTATION + } + for name in self._state_kept_inside(ep): + users = [user for user in nodes[name].users if user.op != "output"] + # tag_mutated_buffer leaves a buffer no partition reads outside of CoreML. + if all(user.meta.get("delegation_tag") is None for user in users): + continue + users.append(updates[signature.inputs_to_buffers[name]]) + tags = {user: user.meta.get("delegation_tag") for user in users} + if len(set(tags.values())) > 1: + hint = ( + "" + if self.lower_full_graph + else " Set take_over_mutable_buffer=False in the " + "CoreMLPartitioner to keep the state outside of CoreML." + ) + places = ", ".join( + f"{user.name} ({'CoreML ' + tag if tag else 'CPU'})" + for user, tag in tags.items() + ) + raise NotImplementedError( + f"CoreML takes over {name} as state, but it is used in more " + f"than one place: {places}. Each place would use its own " + f"copy.{hint}" + ) + + def _targets_ios18(self) -> bool: + """Whether the deployment target is iOS 18 or later, which Core ML needs + for a partition with no inputs.""" + target = CoreMLBackend.min_deployment_target_from_compile_specs( + self.delegation_spec.compile_specs + ) + return target is not None and target >= ct.target.iOS18 + + def _placeholders_kept_inside(self, ep: ExportedProgram) -> Set[str]: + """The weights, constants and state this partitioner moves into Core ML.""" + signature = ep.graph_signature + mutated = set(signature.buffers_to_mutate.values()) + kept = self._state_kept_inside(ep) + if self.take_over_constant_data: + kept |= set(signature.inputs_to_parameters) + kept |= set(signature.inputs_to_lifted_tensor_constants) + kept |= { + name + for name, buffer in signature.inputs_to_buffers.items() + if buffer not in mutated + } + return kept + + def _state_kept_inside(self, ep: ExportedProgram) -> Set[str]: + """The mutable buffers this partitioner moves into Core ML as state.""" + if not self.take_over_mutable_buffer: + return set() + signature = ep.graph_signature + mutated = set(signature.buffers_to_mutate.values()) + return { + name + for name, buffer in signature.inputs_to_buffers.items() + if buffer in mutated + } + def ops_to_not_decompose( self, ep: ExportedProgram ) -> Tuple[List[torch._ops.OpOverload], Optional[Callable[[torch.fx.Node], bool]]]: + # to_edge_transform_and_lower asks twice in a row for the same program, + # and the dry run in _drop_ops_left_outside can decompose all of it. + cached = self._ops_to_not_decompose_cache + if cached is not None and cached[0]() is ep: + return list(cached[1]), None do_not_decompose = [] op_support = _OperatorsSupportedForCoreMLBackend( self.skip_ops_for_coreml_delegation, self.lower_full_graph, log=False, + ios18=self._targets_ios18(), ) # CoreML prevents certain ops (like triu) from lowering to CoreML when put in the ExecuTorch op namespace @@ -452,21 +881,24 @@ def ops_to_not_decompose( torch.ops.aten.repeat_interleave.self_int, torch.ops.aten.repeat_interleave.self_Tensor, ] + + def is_supported(submodules, node: torch.fx.Node) -> bool: + try: + return op_support.is_node_supported(submodules, 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 False + rejected = set() for node in ep.graph.nodes: if node.op == "call_function" and isinstance( node.target, torch._ops.OpOverload ): - try: - 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}" - ) - supported = False - if not supported: + if not is_supported(None, node): rejected.add(node.target) elif node.target not in do_not_decompose_blocklist and not _is_view_op( node.target @@ -474,4 +906,80 @@ def ops_to_not_decompose( 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 + keep = [op for op in do_not_decompose if op not in rejected] + keep = self._drop_ops_left_outside(ep, keep, rejected, is_supported) + # A weak reference, so the partitioner does not keep the program alive. + self._ops_to_not_decompose_cache = (weakref.ref(ep), keep) + return list(keep), None + + def _drop_ops_left_outside( + self, + ep: ExportedProgram, + keep: List[torch._ops.OpOverload], + rejected: Set[torch._ops.OpOverload], + is_supported: Callable[[Any, torch.fx.Node], bool], + ) -> List[torch._ops.OpOverload]: + """Drops every kept op that partition() moves out of CoreML at some call + for a uint8 or int8 tensor, or that CoreML rejects once the program is + decomposed, since the portable kernels may have no kernel for it. An op + skipped by name stays kept, as before. + + partition() sees the program once lowering decomposes the ops that are + not kept. That can make new calls of a kept op that CoreML rejects, such + as the matmul of a bfloat16 linear, or change which ops read a uint8 or + int8 tensor from outside of CoreML. So this decomposes the program the + same way and runs the same search, again after each drop that changes + the decomposition. It does not see transform passes, or the ops an + earlier partitioner takes. + """ + kept_inside = self._placeholders_kept_inside(ep) + skip = set(self.skip_ops_for_coreml_delegation) + not_kept = set(rejected) + decompositions = _default_decomposition_table() + while keep: + table = dict(decompositions) + for op in keep: + table.pop(op, None) + # Decomposing only views and blocklisted ops changes nothing here. + decomposed = any( + node.target in table and node.target in not_kept + for node in ep.graph.nodes + ) + graph_module = ( + ep.run_decompositions(table).graph_module + if decomposed + else ep.graph_module + ) + left_out = set() + if any( + node.name not in kept_inside and _cannot_be_coreml_input(node) + for node in graph_module.graph.nodes + ): + _, left_out = _propose_partitions( + graph_module, + create_op_support( + lambda submodules, node: is_supported(submodules, node) + and not (skip and _edge_op_names(node.target) & skip) + ), + kept_inside, + self._targets_ios18(), + ) + kept = set(keep) + late = { + node.target + for node in graph_module.graph.nodes + if node.target in kept + and (node in left_out or (decomposed and not is_supported(None, node))) + } + if not late: + break + for op in late: + self.log_once( + f"Not keeping {op} whole, because some call of it runs outside " + "of CoreML." + ) + not_kept |= late + keep = [op for op in keep if op not in late] + if not any(op in decompositions for op in late): + break + return keep diff --git a/backends/apple/coreml/test/test_coreml_partitioner.py b/backends/apple/coreml/test/test_coreml_partitioner.py index 0129eb1d33f..729c88a4907 100644 --- a/backends/apple/coreml/test/test_coreml_partitioner.py +++ b/backends/apple/coreml/test/test_coreml_partitioner.py @@ -14,6 +14,7 @@ import torch import torchvision +from coremltools.converters.mil.mil import types from executorch.backends.apple.coreml.compiler import CoreMLBackend from executorch.backends.apple.coreml.partition import CoreMLPartitioner @@ -521,6 +522,732 @@ def forward(self, q): out = program.load_method("forward").execute((x,))[0] self.assertTrue(torch.allclose(out, model(x), atol=5e-2, rtol=5e-2)) + def test_op_made_by_decomposing_a_rejected_op_is_not_kept_whole(self): + """ + A float32 matmul keeps aten.matmul whole. A bfloat16 linear is + rejected and decomposes into another matmul, which Core ML rejects too. + That matmul must be decomposed as well, since the portable kernels + have no matmul, or the program fails to load. + """ + + class Model(torch.nn.Module): + def __init__(self): + super().__init__() + self.proj = torch.nn.Linear(8, 8, bias=False).to(torch.bfloat16) + self.register_buffer("inv_freq", torch.rand(1, 4, 1)) + + def forward(self, x, pos): + freqs = self.inv_freq @ pos[:, None, :].float() + return self.proj(x).float(), freqs.cos() + + model = Model().eval() + inputs = (torch.rand(1, 6, 8, dtype=torch.bfloat16), torch.arange(6)[None]) + ep = torch.export.export(model, inputs, 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.matmul.default", op_names) + if _TEST_RUNTIME: + program = Runtime.get().load_program(edge.to_executorch().buffer) + outputs = program.load_method("forward").execute(inputs) + for out, ref in zip(outputs, model(*inputs)): + self.assertTrue(torch.allclose(out, ref, atol=1e-2)) + + def test_min_max_of_bool_is_skipped(self): + """ + coremltools raises on a min or max over a bool tensor, such as the + all-true check torch.distributions exports. The partitioner must leave + those ops to the portable backend and still delegate the rest. + """ + + class Model(torch.nn.Module): + def forward(self, x): + ok = (x.tril() == x).view(1, -1) + flags = ok.min(-1)[0].to(x.dtype) + ok.amax(-1).to(x.dtype) + both = torch.maximum(ok, ok.flip(-1)).view(x.shape).to(x.dtype) + return x * 2 + flags + both + + model = Model().eval() + x = torch.rand(14, 14) + 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.min.dim", op_names) + self.assertIn("aten.amax.default", op_names) + self.assertIn("aten.maximum.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=1e-2)) + + def test_ops_coremltools_gets_wrong(self): + """ + coremltools makes a filled int or bool tensor as float in some cases, + such as ``full`` without a dtype or ``zeros`` with one. Index math + with such a tensor then fails to lower, or loses integers above 2048 + to float16 precision, so the backend gives fills the dtype torch gives + them. coremltools also fails some forms of ``index_put``, such as + ``y[:, idx] = v`` with two entries in ``idx``, or writes them to the + wrong place, so the partitioner refuses those, and they run on the + portable backend. With lower_full_graph=True, lowering fails and says + why. + """ + + class Positions(torch.nn.Module): + def __init__(self): + super().__init__() + self.emb = torch.nn.Embedding(16, 4) + + def forward(self, x, mask): + pos = torch.full((1, 16), fill_value=0) + pos[mask] = torch.arange(16)[None][mask] + return x + self.emb(pos) + + class Counts(torch.nn.Module): + def forward(self, ids): + ones = torch.ones_like(ids) + ids.new_zeros((2, 3)) + twos = ids.new_full((2, 3), 2.0) + torch.empty_like(ids).fill_(0.0) + zeros = torch.zeros((2, 3), dtype=torch.long) + torch.zeros_like(ids) + count = torch.scalar_tensor(1001.0, dtype=torch.long) + return (ids + ones + twos + zeros) * count + + class Flags(torch.nn.Module): + def forward(self, mask): + return mask & torch.full_like(mask, 1) + + # Turns a bool mask into a float one with zeros_like(mask, dtype=float). + class Attention(torch.nn.MultiheadAttention): + def forward(self, x, mask): + return super().forward(x, x, x, attn_mask=mask, need_weights=False)[0] + + class Scatter(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("cols", torch.tensor([0, 2])) + self.register_buffer("col", torch.tensor([1])) + self.register_buffer("last", torch.tensor(3)) + self.register_buffer("rows", torch.tensor([[1], [2]])) + + def forward(self, x, v): + y = x.clone() + y[:, self.cols] = v + y[self.rows, self.cols] = v[:2] + torch.ops.aten.index_put_(y, [None, self.col], v[:, :1], True) + torch.ops.aten.index_put_(y, [None, self.last], v[:, 0]) + mask = y > 0.5 + mask[:, self.col] = True + return y * 2 + mask + + class Scatter3d(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("last", torch.tensor(2)) + self.register_buffer("mid", torch.tensor(1)) + + def forward(self, x, v): + y = x.clone() + torch.ops.aten.index_put_(y, [None, None, self.last], v) + # Values that vary only after the index stay in Core ML. + torch.ops.aten.index_put_(y, [None, self.mid], v[0] * torch.arange(3.0)) + return y * 2 + + mask = torch.arange(16)[None] % 3 != 0 + ios18 = CoreMLBackend.generate_compile_specs( + minimum_deployment_target=ct.target.iOS18 + ) + cases = [ + (Positions(), (torch.rand(1, 16, 4), mask), "aten.full.default", 0, None), + (Counts(), (torch.arange(6).reshape(2, 3),), "aten.zeros.default", 0, None), + (Flags(), (mask,), "aten.full_like.default", 0, None), + ( + Attention(16, 2, batch_first=True), + (torch.randn(1, 6, 16), torch.ones(6, 6, dtype=torch.bool).triu(1)), + "aten.zeros_like.default", + 0, + None, + ), + ( + Scatter(), + (torch.rand(3, 4), torch.rand(3, 2)), + "aten.index_put.default", + 5, + None, + ), + # From iOS 18 the bool write stays in Core ML, and reads its index + # from the same buffer as the accumulate write on the CPU. + ( + Scatter(), + (torch.rand(3, 4), torch.rand(3, 2)), + "aten.index_put.default", + 4, + ios18, + ), + ( + Scatter3d(), + (torch.rand(2, 2, 3), torch.rand(2, 1)), + "aten.index_put.default", + 1, + None, + ), + ] + for model, inputs, cpu_op, cpu_calls, compile_specs in cases: + with self.subTest( + model=type(model).__name__, ios18=compile_specs is not None + ): + ep = torch.export.export(model.eval(), inputs, strict=True) + edge = executorch.exir.to_edge_transform_and_lower( + ep, partitioner=[CoreMLPartitioner(compile_specs=compile_specs)] + ) + 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.assertEqual(op_names.count(cpu_op), cpu_calls) + if _TEST_RUNTIME: + program = Runtime.get().load_program(edge.to_executorch().buffer) + out = program.load_method("forward").execute(inputs)[0] + ref = model(*inputs) + self.assertEqual(out.dtype, ref.dtype) + self.assertTrue(torch.allclose(out, ref, atol=1e-2)) + counts = torch.export.export(Counts(), cases[1][1], strict=True) + program = ct.convert(counts.run_decompositions({}), convert_to="milinternal") + self.assertEqual(program.functions["main"].outputs[0].dtype, types.int32) + with self.assertRaisesRegex(NotImplementedError, "with these indices"): + executorch.exir.to_edge_transform_and_lower( + torch.export.export(Scatter3d(), cases[-1][1], strict=True), + partitioner=[CoreMLPartitioner(lower_full_graph=True)], + ) + + def test_int8_made_for_the_cpu_runs_on_the_cpu(self): + """ + The Core ML runtime cannot return an int8 tensor, so an op that makes + one that is read outside Core ML, here the model output, runs on the + CPU. So does an op that hands one out through a getitem, which always + stays with the op. + """ + + class Model(torch.nn.Module): + def forward(self, x): + return (x * 2).to(torch.int8) + + class Split(torch.nn.Module): + def forward(self, x): + return (x * 2).to(torch.int8).split(2, dim=1)[0] + + # Small integers, which Core ML's float16 math keeps exact. + x = torch.randint(-50, 50, (1, 4, 4, 4)).float() + for model in (Model().eval(), Split().eval()): + with self.subTest(model=type(model).__name__): + 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("dim_order_ops._to_dim_order_copy.default", op_names) + buffer = edge.to_executorch().buffer + if _TEST_RUNTIME: + program = Runtime.get().load_program(buffer) + out = program.load_method("forward").execute((x,))[0] + self.assertTrue(torch.equal(out, model(x))) + with self.assertRaisesRegex(NotImplementedError, "makes an int8"): + executorch.exir.to_edge_transform_and_lower( + ep, partitioner=[CoreMLPartitioner(lower_full_graph=True)] + ) + + def test_uint8_camera_frame_is_read_outside_the_delegate(self): + """ + Core ML cannot take a uint8 or int8 input, so the ops that read such a + camera frame stay on the CPU up to its conversion to float, and the + float ops after it are still delegated. + """ + + class Model(torch.nn.Module): + def forward(self, image): + return torch.relu(image.permute(2, 0, 1).float() / 127.0 - 0.25) + + model = Model().eval() + for dtype in (torch.uint8, torch.int8): + with self.subTest(dtype=dtype): + info = torch.iinfo(dtype) + x = torch.randint(info.min, info.max + 1, (8, 8, 3), dtype=dtype) + 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("aten.permute_copy.default", op_names) + self.assertIn("dim_order_ops._to_dim_order_copy.default", op_names) + self.assertIn("executorch_call_delegate", op_names) + for name in ("aten.div.Tensor", "aten.sub.Tensor", "aten.relu.default"): + self.assertNotIn(name, op_names) + buffer = edge.to_executorch().buffer + if _TEST_RUNTIME: + program = Runtime.get().load_program(buffer) + out = program.load_method("forward").execute((x,))[0] + self.assertTrue(torch.allclose(out, model(x), atol=1e-2)) + with self.assertRaisesRegex( + NotImplementedError, "image \\(torch\\.u?int8\\)" + ): + executorch.exir.to_edge_transform_and_lower( + ep, partitioner=[CoreMLPartitioner(lower_full_graph=True)] + ) + + def test_uint8_input_read_by_ops_with_and_without_cpu_kernels(self): + """ + An op that reads a uint8 input runs on the CPU, also when it reads it + through an op with several outputs. An op that is not core ATen, such + as pad, max pooling, roll or zeros_like, may have no CPU kernel, so it + is decomposed. Below iOS 18 the index math of a gather also runs on + the CPU, since it would leave Core ML a partition with no inputs. + """ + + class Model(torch.nn.Module): + def __init__(self, op): + super().__init__() + self.op = op + + def forward(self, image): + return self.op(image).float() / 255.0 + + # Each op, and the op that has to run on the CPU in its place. + ops = { + "pad": ( + lambda x: torch.nn.functional.pad(x, (1, 1, 1, 1)), + "aten.constant_pad_nd.default", + ), + "max_pool2d": ( + lambda x: torch.nn.functional.max_pool2d(x, 2), + "aten.max_pool2d_with_indices.default", + ), + "reflect_pad": ( + lambda x: torch.nn.functional.pad(x, (1, 1, 1, 1), mode="reflect"), + "aten.arange.start_step", + ), + "roll": (lambda x: torch.roll(x, 1, 2), "aten.arange.start_step"), + "zeros_like": ( + lambda x: torch.where(x > 100, x, torch.zeros_like(x)), + "aten.full_like.default", + ), + "channel_swap": (lambda x: x[:, [2, 1, 0]], "aten.index.Tensor"), + "masked_fill": (lambda x: x.masked_fill(x < 10, 0), "aten.where.self"), + "max_then_pad": ( + lambda x: torch.nn.functional.pad(x.max(dim=1).values, (1, 1)), + "aten.constant_pad_nd.default", + ), + "split_then_pad": ( + lambda x: torch.nn.functional.pad(x.split(1, dim=1)[0], (1, 1)), + "aten.constant_pad_nd.default", + ), + } + x = torch.randint(0, 256, (1, 3, 8, 8), dtype=torch.uint8) + ios18 = CoreMLBackend.generate_compile_specs( + minimum_deployment_target=ct.target.iOS18 + ) + # From iOS 18 index math with no inputs stays in Core ML instead. + cases = [(name, op, cpu_op, None) for name, (op, cpu_op) in ops.items()] + cases += [ + (name, op, cpu_op, ios18) + for name, (op, cpu_op) in ops.items() + if cpu_op == "aten.arange.start_step" + ] + for name, op, cpu_op, compile_specs in cases: + with self.subTest(op=name, ios18=compile_specs is not None): + model = Model(op).eval() + ep = torch.export.export(model, (x,), strict=True) + edge = executorch.exir.to_edge_transform_and_lower( + ep, partitioner=[CoreMLPartitioner(compile_specs=compile_specs)] + ) + op_names = [ + n.target.__name__ + for n in edge.exported_program().graph.nodes + if n.op == "call_function" + ] + if compile_specs is None: + self.assertIn(cpu_op, op_names) + else: + self.assertNotIn(cpu_op, op_names) + self.assertIn("executorch_call_delegate", op_names) + buffer = edge.to_executorch().buffer + if _TEST_RUNTIME: + program = Runtime.get().load_program(buffer) + out = program.load_method("forward").execute((x,))[0] + self.assertTrue(torch.allclose(out, model(x), atol=1e-2)) + + def test_uint8_input_with_dynamic_shape(self): + """ + With a dynamic height, the index math of a decomposed roll reads only + a symbolic size, which Core ML cannot take as an input, so it runs on + the CPU too. So do a reshape to a size read from the uint8 input, and + a scale computed from it. + """ + + class Model(torch.nn.Module): + def forward(self, image): + rolled = torch.roll(image, 1, 2).float() / 255.0 + scaled = rolled * (image.shape[2] / 8) + return scaled.reshape(1, 3 * image.shape[2], 8) + + model = Model().eval() + x = torch.randint(0, 256, (1, 3, 8, 8), dtype=torch.uint8) + ep = torch.export.export( + model, + (x,), + dynamic_shapes=({2: torch.export.Dim("height", min=4, max=64)},), + strict=True, + ) + ios18 = CoreMLBackend.generate_compile_specs( + minimum_deployment_target=ct.target.iOS18 + ) + # Also from iOS 18, where only a partition with no inputs at all may + # stay in Core ML. + for compile_specs in (None, ios18): + with self.subTest(ios18=compile_specs is not None): + edge = executorch.exir.to_edge_transform_and_lower( + ep, partitioner=[CoreMLPartitioner(compile_specs=compile_specs)] + ) + op_names = [ + n.target.__name__ + for n in edge.exported_program().graph.nodes + if n.op == "call_function" + ] + for name in ("aten.arange.start_step", "aten.view_copy.default"): + self.assertIn(name, op_names) + self.assertIn("executorch_call_delegate", op_names) + buffer = edge.to_executorch().buffer + if _TEST_RUNTIME: + y = torch.randint(0, 256, (1, 3, 12, 8), dtype=torch.uint8) + program = Runtime.get().load_program(buffer) + out = program.load_method("forward").execute((y,))[0] + self.assertTrue(torch.allclose(out, model(y), atol=1e-2)) + + def test_uint8_input_with_state(self): + """ + A CPU op that reads a uint8 input can separate a read of state from its + update, or make the update itself, so two places would each use their + own copy of the state. + Lowering refuses that, at any deployment target, unless the state stays + outside of Core ML. State that only a partition with no inputs uses + moves to the CPU with it below iOS 18, and state that only CPU ops read + stays there, even when Core ML makes its new value. + """ + + class Split(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("state", torch.zeros(1, 3, 2, 2)) + + def forward(self, image): + y = (image + (self.state * 2).to(torch.uint8)).float() + self.state.add_(y) + return y + 1 + + class CopyIn(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("state", torch.zeros(1, 3, 2, 2)) + + def forward(self, image): + y = self.state * 2 + 1 + self.state.copy_(image.float()) + return y + + class Counter(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("count", torch.zeros(1)) + + def forward(self, image): + self.count.add_(1) + return (image + self.count.to(torch.uint8)).float() + + class ReadOnCpu(CopyIn): + def forward(self, image): + y = torch.asinh(self.state) + self.state.copy_(image.float() + 1) + return y + + x = torch.randint(0, 3, (1, 3, 2, 2), dtype=torch.uint8) + ios18 = CoreMLBackend.generate_compile_specs( + minimum_deployment_target=ct.target.iOS18 + ) + for model_class, compile_specs in ( + (Split, None), + (Split, ios18), + (CopyIn, ios18), + ): + with self.subTest( + model=model_class.__name__, ios18=compile_specs is not None + ): + ep = torch.export.export(model_class().eval(), (x,), strict=True) + with self.assertRaisesRegex( + NotImplementedError, "b_state as state, but .* place: \\w+ \\(" + ): + executorch.exir.to_edge_transform_and_lower( + ep, partitioner=[CoreMLPartitioner(compile_specs=compile_specs)] + ) + lowered = [ + (Split, CoreMLPartitioner(take_over_mutable_buffer=False)), + (Counter, CoreMLPartitioner()), + (ReadOnCpu, CoreMLPartitioner()), + ] + for model_class, partitioner in lowered: + with self.subTest(model=model_class.__name__): + ep = torch.export.export(model_class().eval(), (x,), strict=True) + edge = executorch.exir.to_edge_transform_and_lower( + ep, partitioner=[partitioner] + ) + buffer = edge.to_executorch().buffer + if _TEST_RUNTIME: + method = Runtime.get().load_program(buffer).load_method("forward") + eager = model_class().eval() + for _ in range(2): + out = method.execute((x,))[0] + self.assertTrue(torch.allclose(out, eager(x), atol=1e-2)) + + def test_uint8_input_size_that_slices_state(self): + """ + A slice of state to a size read from a uint8 input has to run on the + CPU. When Core ML uses that state too, the slice would read a stale + copy, so lowering refuses it unless the state stays outside of Core ML. + When only CPU ops use it, the state stays on the CPU. + """ + + class Model(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("hist", torch.zeros(64)) + + def forward(self, image, y): + out = y * self.hist + self.hist[: image.shape[2]].sum() + self.hist.add_(1) + return out + + class SliceUpdate(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("hist", torch.zeros(64)) + + def forward(self, image, y): + self.hist[: image.shape[2]] += 1 + return y * 2 + self.hist + + inputs = ( + torch.randint(0, 256, (1, 3, 8, 8), dtype=torch.uint8), + torch.randn(64), + ) + dynamic_shapes = ({2: torch.export.Dim("height", min=2, max=64)}, None) + compile_specs = CoreMLBackend.generate_compile_specs( + minimum_deployment_target=ct.target.iOS18 + ) + ep = torch.export.export( + Model().eval(), inputs, dynamic_shapes=dynamic_shapes, strict=True + ) + with self.assertRaisesRegex(NotImplementedError, "take_over_mutable_buffer"): + executorch.exir.to_edge_transform_and_lower( + ep, partitioner=[CoreMLPartitioner(compile_specs=compile_specs)] + ) + lowered = [ + ( + Model, + CoreMLPartitioner( + compile_specs=compile_specs, take_over_mutable_buffer=False + ), + ), + (SliceUpdate, CoreMLPartitioner(compile_specs=compile_specs)), + ] + for model_class, partitioner in lowered: + with self.subTest(model=model_class.__name__): + ep = torch.export.export( + model_class().eval(), + inputs, + dynamic_shapes=dynamic_shapes, + strict=True, + ) + edge = executorch.exir.to_edge_transform_and_lower( + ep, partitioner=[partitioner] + ) + buffer = edge.to_executorch().buffer + if _TEST_RUNTIME: + method = Runtime.get().load_program(buffer).load_method("forward") + eager = model_class().eval() + for _ in range(2): + out = method.execute(inputs)[0] + self.assertTrue(torch.allclose(out, eager(*inputs), atol=1e-2)) + + def test_uint8_buffer_not_taken_over_is_read_outside_the_delegate(self): + """ + A uint8 buffer that Core ML does not take over is a program input, so + the op that reads it runs on the CPU. A uint8 buffer Core ML takes + over, the default, still fails to lower: coremltools cannot cast a + uint8 constant. + """ + + class Model(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer( + "lut", torch.randint(0, 256, (8,), dtype=torch.uint8) + ) + + def forward(self, x): + return x * (self.lut.float() / 255.0) + + model = Model().eval() + x = torch.randn(2, 8) + ep = torch.export.export(model, (x,), strict=True) + with self.assertRaisesRegex(ValueError, "tensor\\[8,uint8\\]"): + executorch.exir.to_edge_transform_and_lower( + ep, partitioner=[CoreMLPartitioner()] + ) + edge = executorch.exir.to_edge_transform_and_lower( + ep, partitioner=[CoreMLPartitioner(take_over_constant_data=False)] + ) + op_names = [ + n.target.__name__ + for n in edge.exported_program().graph.nodes + if n.op == "call_function" + ] + self.assertIn("dim_order_ops._to_dim_order_copy.default", op_names) + self.assertIn("executorch_call_delegate", op_names) + buffer = edge.to_executorch().buffer + if _TEST_RUNTIME: + program = Runtime.get().load_program(buffer) + out = program.load_method("forward").execute((x,))[0] + self.assertTrue(torch.allclose(out, model(x), atol=1e-2)) + + def test_uint8_made_inside_the_delegate_keeps_ops_whole(self): + """ + A uint8 tensor made inside the delegate is not an input, so the ops + that read it stay in the delegate, whole: a pad, and a rot90 that only + Core ML's decomposition supports followed by a roll. + """ + + class Pad(torch.nn.Module): + def forward(self, x): + image = torch.nn.functional.pad((x * 255).to(torch.uint8), (1, 1)) + return torch.nn.functional.pad( + image.float() / 255.0, (1, 1), mode="reflect" + ) + + class RotateRoll(torch.nn.Module): + def forward(self, x): + image = torch.rot90((x * 255).to(torch.uint8), 1, (2, 3)) + return torch.roll(image, 1, 2).float() / 255.0 + + for model, x in ( + (Pad(), torch.rand(1, 3, 8)), + (RotateRoll(), torch.rand(1, 3, 8, 8)), + ): + with self.subTest(model=type(model).__name__): + ep = torch.export.export(model.eval(), (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.assertEqual(op_names, {"executorch_call_delegate", "getitem"}) + buffer = edge.to_executorch().buffer + if _TEST_RUNTIME: + program = Runtime.get().load_program(buffer) + out = program.load_method("forward").execute((x,))[0] + self.assertTrue(torch.allclose(out, model(x), atol=1e-2)) + + def test_ops_after_a_cpu_op_on_a_tensor_made_inside(self): + """ + A uint8 tensor made in Core ML and passed through an op that stays on + the CPU, such as diagonal, a skipped flip, view or cast, or an add that + reads it from another partition past acosh, sends the pad after it to + the CPU too. The pad is decomposed there, since the portable kernels have no + pad, and so is a pad that makes an int8 model output. + """ + + class DiagonalPad(torch.nn.Module): + def forward(self, x): + d = torch.diagonal((x * 255).to(torch.uint8), 0, 2, 3) + return torch.nn.functional.pad(d, (1, 1)).float() / 255.0 + + class FlipPad(torch.nn.Module): + def forward(self, x): + d = (x * 255).to(torch.uint8).flip(-1) + return torch.nn.functional.pad(d, (1, 1)).float() / 255.0 + + class ViewPad(torch.nn.Module): + def forward(self, x): + d = (x * 255).to(torch.uint8).view(1, 3, 8) + return torch.nn.functional.pad(d, (1, 1)).float() / 255.0 + + class AcoshPad(torch.nn.Module): + def forward(self, x): + d = (x * 255).to(torch.uint8) + d = d + torch.acosh(d.float() + 2).to(torch.uint8) + return torch.nn.functional.pad(d, (1, 1)).float() / 255.0 + + class PadInt8(torch.nn.Module): + def forward(self, x): + return torch.nn.functional.pad((x * 2).to(torch.int8), (1, 1)) + + for model, x, skip in ( + (DiagonalPad(), torch.rand(1, 3, 8, 8), None), + (FlipPad(), torch.rand(1, 3, 8), ["aten.flip.default"]), + (ViewPad(), torch.rand(1, 3, 2, 4), ["aten.view_copy.default"]), + ( + FlipPad(), + torch.rand(1, 3, 8), + ["dim_order_ops._to_dim_order_copy.default"], + ), + (AcoshPad(), torch.rand(1, 3, 8), None), + (PadInt8(), torch.randint(-25, 25, (1, 3, 8)).float(), None), + ): + with self.subTest(model=type(model).__name__, skip=skip): + ep = torch.export.export(model.eval(), (x,), strict=True) + edge = executorch.exir.to_edge_transform_and_lower( + ep, + partitioner=[ + CoreMLPartitioner(skip_ops_for_coreml_delegation=skip) + ], + ) + op_names = [ + n.target.__name__ + for n in edge.exported_program().graph.nodes + if n.op == "call_function" + ] + self.assertIn("aten.constant_pad_nd.default", op_names) + self.assertIn("executorch_call_delegate", op_names) + buffer = edge.to_executorch().buffer + if _TEST_RUNTIME: + program = Runtime.get().load_program(buffer) + out = program.load_method("forward").execute((x,))[0] + self.assertTrue( + torch.allclose(out.float(), model(x).float(), atol=1e-2) + ) + def test_tied_embedding_is_quantized_with_its_linear(self): """A weight that is both an embedding table and an output projection lowers. diff --git a/backends/apple/coreml/test/test_torch_ops.py b/backends/apple/coreml/test/test_torch_ops.py index 10c3f01a585..a7c33410cd9 100644 --- a/backends/apple/coreml/test/test_torch_ops.py +++ b/backends/apple/coreml/test/test_torch_ops.py @@ -9,10 +9,13 @@ import executorch.exir import torch +from coremltools.converters.mil.mil import types +from coremltools.converters.mil.mil.passes.pass_registry import PASS_REGISTRY from executorch.backends.apple.coreml.compiler import CoreMLBackend from executorch.backends.apple.coreml.partition import CoreMLPartitioner from executorch.backends.apple.coreml.test.test_coreml_utils import ( + is_fbcode, IS_VALID_TEST_RUNTIME, ) from executorch.exir.backend.utils import format_delegated_graph @@ -40,17 +43,23 @@ def _get_test_model(self): example_inputs = (torch.LongTensor([0]),) return model, example_inputs - def _compare_outputs(self, executorch_program, eager_program, example_inputs): + def _compare_outputs( + self, executorch_program, eager_program, example_inputs, atol=1e-2, rtol=1e-2 + ): if not IS_VALID_TEST_RUNTIME: return runtime = Runtime.get() program = runtime.load_program(executorch_program.buffer) method = program.load_method("forward") - et_outputs = method.execute(example_inputs)[0] + et_outputs = method.execute(example_inputs) eager_outputs = eager_program(*example_inputs) - self.assertTrue( - torch.allclose(et_outputs, eager_outputs, atol=1e-02, rtol=1e-02) - ) + if isinstance(eager_outputs, torch.Tensor): + eager_outputs = (eager_outputs,) + self.assertEqual(len(et_outputs), len(eager_outputs)) + for et_output, eager_output in zip(et_outputs, eager_outputs): + self.assertTrue( + torch.allclose(et_output, eager_output, atol=atol, rtol=rtol) + ) def test_dequantize_affine_b4w_embedding(self): model, example_inputs = self._get_test_model() @@ -296,14 +305,27 @@ def test_dequantize_codebook_embedding_per_grouped_row(self): self._compare_outputs(et_prog, model, example_inputs) def test__clone_dim_order_contiguous(self): + """ + The clones convert to nothing, so the Core ML program returns its own + input, and one value twice, which Core ML used to return wrong. + """ + class Model(torch.nn.Module): def forward(self, x): - return torch.ops.dim_order_ops._clone_dim_order( - x, dim_order=[0, 1, 2, 3] + y = x * 2 + return ( + torch.ops.dim_order_ops._clone_dim_order(x, dim_order=[0, 1, 2, 3]), + y, + torch.ops.dim_order_ops._clone_dim_order(y, dim_order=[0, 1, 2, 3]), ) model, example_inputs = Model(), (torch.randn(1, 3, 8, 8),) ep = torch.export.export(model, example_inputs) + program = ct.convert(ep.run_decompositions({}), convert_to="milinternal") + PASS_REGISTRY["executorch::copy_reused_outputs"](program) + function = program.functions["main"] + self.assertEqual(len(set(function.outputs)), len(function.outputs)) + self.assertFalse(set(function.outputs) & set(function.inputs.values())) delegated_program = executorch.exir.to_edge_transform_and_lower( ep, partitioner=[self._coreml_partitioner()], @@ -317,6 +339,89 @@ def forward(self, x): et_prog = delegated_program.to_executorch() self._compare_outputs(et_prog, model, example_inputs) + def test_gelu_tanh(self): + """ + torch.export passes approximate="tanh" to gelu as a keyword, and + coremltools reads it only from the positional inputs. If the keyword + is dropped, the exact gelu runs, which differs by up to 4.7e-4. + """ + model = torch.nn.GELU(approximate="tanh").eval() + example_inputs = (torch.linspace(-8, 8, 4096),) + compile_specs = CoreMLBackend.generate_compile_specs( + compute_precision=ct.precision.FLOAT32 + ) + ep = torch.export.export(model, example_inputs) + program = ct.convert(ep.run_decompositions({}), convert_to="milinternal") + self.assertEqual( + [op.mode.val for op in program.find_ops(op_type="gelu")], + ["TANH_APPROXIMATION"], + ) + delegated_program = executorch.exir.to_edge_transform_and_lower( + ep, partitioner=[CoreMLPartitioner(compile_specs=compile_specs)] + ) + for node in delegated_program.exported_program().graph.nodes: + if node.op == "call_function": + self.assertIn( + node.target.__name__, ["executorch_call_delegate", "getitem"] + ) + et_prog = delegated_program.to_executorch() + self._compare_outputs(et_prog, model, example_inputs, atol=1e-5, rtol=1e-5) + + def test_float16_fill(self): + """A float16 fill stays float16, so the ops that read it stay float16.""" + + class Model(torch.nn.Module): + def forward(self, x): + return x * torch.ones_like(x) + torch.zeros_like(x) + + ep = torch.export.export(Model(), (torch.rand(2, 3, dtype=torch.float16),)) + program = ct.convert( + ep.run_decompositions({}), + convert_to="milinternal", + compute_precision=ct.precision.FLOAT32, + minimum_deployment_target=ct.target.iOS16, + ) + self.assertEqual(program.functions["main"].outputs[0].dtype, types.fp16) + + @unittest.skipIf(is_fbcode(), "The coremltools in fbcode has no rms_norm.") + def test_rms_norm_without_eps(self): + """ + torch.nn.RMSNorm(eps=None) leaves eps out of rms_norm, which + coremltools requires. Torch then uses the float32 machine epsilon, + which the output for a one-hot row of 65536 values depends on. + """ + x = torch.zeros(1, 65536) + x[0, 0] = 1 + compile_specs = CoreMLBackend.generate_compile_specs( + compute_precision=ct.precision.FLOAT32 + ) + for model in ( + torch.nn.RMSNorm(65536), + torch.nn.RMSNorm(65536, elementwise_affine=False), + ): + with self.subTest(elementwise_affine=model.elementwise_affine): + ep = torch.export.export(model.eval(), (x,)) + program = ct.convert( + ep.run_decompositions({}), + convert_to="milinternal", + compute_precision=ct.precision.FLOAT32, + ) + self.assertEqual( + [op.y.val for op in program.find_ops(op_type="add")], + [torch.finfo(torch.float32).eps], + ) + delegated_program = executorch.exir.to_edge_transform_and_lower( + ep, partitioner=[CoreMLPartitioner(compile_specs=compile_specs)] + ) + for node in delegated_program.exported_program().graph.nodes: + if node.op == "call_function": + self.assertIn( + node.target.__name__, + ["executorch_call_delegate", "getitem"], + ) + et_prog = delegated_program.to_executorch() + self._compare_outputs(et_prog, model, (x,), rtol=1e-3) + def test_dequantize_affine_below_ios18_raises_with_hint(self): """ Regression test for https://github.com/pytorch/executorch/issues/13122. diff --git a/docs/source/backends/coreml/coreml-partitioner.md b/docs/source/backends/coreml/coreml-partitioner.md index fa7a1de8277..cf009e78f32 100644 --- a/docs/source/backends/coreml/coreml-partitioner.md +++ b/docs/source/backends/coreml/coreml-partitioner.md @@ -5,8 +5,8 @@ The Core ML partitioner API allows for configuration of the model delegation to - `skip_ops_for_coreml_delegation`: Allows you to skip ops for delegation by Core ML. By default, all ops that Core ML supports will be delegated. See [here](https://github.com/pytorch/executorch/blob/14ff52ff89a89c074fc6c14d3f01683677783dcd/backends/apple/coreml/test/test_coreml_partitioner.py#L42) for an example of skipping an op for delegation. - `compile_specs`: A list of `CompileSpec`s for the Core ML backend. These control low-level details of Core ML delegation, such as the compute unit (CPU, GPU, ANE), the iOS deployment target, and the compute precision (FP16, FP32). These are discussed more below. -- `take_over_mutable_buffer`: A boolean that indicates whether PyTorch mutable buffers in stateful models should be converted to [Core ML `MLState`](https://developer.apple.com/documentation/coreml/mlstate). If set to `False`, mutable buffers in the PyTorch graph are converted to graph inputs and outputs to the Core ML lowered module under the hood. Generally, setting `take_over_mutable_buffer` to true will result in better performance, but using `MLState` requires iOS >= 18.0, macOS >= 15.0, and Xcode >= 16.0. -- `take_over_constant_data`: A boolean that indicates whether PyTorch constant data like model weights should be consumed by the Core ML delegate. If set to False, constant data is passed to the Core ML delegate as inputs. By default, take_over_constant_data=True. +- `take_over_mutable_buffer`: A boolean that indicates whether PyTorch mutable buffers in stateful models should be converted to [Core ML `MLState`](https://developer.apple.com/documentation/coreml/mlstate). If set to `False`, mutable buffers in the PyTorch graph are converted to graph inputs and outputs to the Core ML lowered module under the hood. Generally, setting `take_over_mutable_buffer` to true will result in better performance, but using `MLState` requires iOS >= 18.0, macOS >= 15.0, and Xcode >= 16.0. If a buffer that Core ML takes over is also used on the CPU or in another Core ML partition, lowering fails and names the ops, since each place would use its own copy; set `take_over_mutable_buffer` to False for such a model. +- `take_over_constant_data`: A boolean that indicates whether PyTorch constant data like model weights should be consumed by the Core ML delegate. If set to False, constant data is passed to the Core ML delegate as inputs. By default, take_over_constant_data=True. A uint8 buffer that the model casts, such as a lookup table, fails to lower with the default settings, since coremltools cannot cast a uint8 constant. Storing such a table as float16 or int32 avoids that. take_over_constant_data=False also works, since the cast then runs on the CPU, but it passes every weight to Core ML as an input, and runs the dequantize of int8 or uint8 weights on the CPU. - `lower_full_graph`: A boolean that indicates whether the entire graph must be lowered to Core ML. If set to True and Core ML does not support an op, an error is raised during lowering. If set to False and Core ML does not support an op, the op is executed on the CPU by ExecuTorch. Although setting `lower_full_graph`=False can allow a model to lower where it would otherwise fail, it can introduce performance overhead in the model when there are unsupported ops. You will see warnings about unsupported ops during lowering if there are any. By default, `lower_full_graph`=False.