diff --git a/backends/apple/coreml/compiler/torch_ops.py b/backends/apple/coreml/compiler/torch_ops.py index ed9cf0d5e18..3a914f7acbc 100644 --- a/backends/apple/coreml/compiler/torch_ops.py +++ b/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,7 +166,7 @@ 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) @@ -161,7 +174,7 @@ 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) diff --git a/backends/apple/coreml/partition/coreml_partitioner.py b/backends/apple/coreml/partition/coreml_partitioner.py index 4453818f956..db85afb3b9f 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, Dict, 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 torch.export.exported_program import ExportedProgram -from torch.fx.passes.infra.partitioner import CapabilityBasedPartitioner -from torch.fx.passes.operator_support import OperatorSupportBase +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, OutputKind +from torch.fx.experimental.proxy_tensor import make_fx +from torch.fx.experimental.symbolic_shapes import is_concrete_int +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,464 @@ 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) +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 _makes_symbolic_value(node: torch.fx.Node) -> bool: + """Whether node makes a symbolic value, such as a size known only at run + time.""" + return isinstance( + node.meta.get("val"), (torch.SymInt, torch.SymFloat, torch.SymBool) + ) + + +def _dynamic_size_names(ep: ExportedProgram) -> Set[str]: + """The names of the dynamic sizes ep has range constraints for, which are + the sizes coremltools can give a Core ML input besides numbers.""" + return {str(symbol) for symbol in ep.range_constraints} + + +# https://github.com/pytorch/executorch/issues/11686 +def _cannot_be_coreml_input(node: torch.fx.Node, sizes: Collection[str]) -> bool: + """Whether node makes a tensor Core ML cannot take as a program input: a + uint8 or int8 tensor, or one with a size that is neither a number nor in + sizes, such as s0 + 1. + + 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. It looks each dynamic size up by name in the range + constraints of the program, and raises ``ValueError`` for one it cannot + find. + """ + val = node.meta.get("val") + vals = val if isinstance(val, (list, tuple)) else [val] + return any( + isinstance(v, torch.Tensor) + and ( + v.dtype in (torch.uint8, torch.int8) + or any( + not is_concrete_int(size) and str(size) not in sizes for size in v.shape + ) + ) + for v in vals + ) + + +def _cannot_be_read_from_outside( + arg: torch.fx.Node, + kept_inside: Collection[str], + sizes: Collection[str], + outside_coreml: Callable[[torch.fx.Node], bool], +) -> bool: + """Whether a Core ML op cannot read arg from outside its partition: a + tensor Core ML cannot take as an input, or a symbolic value computed + outside of Core ML, which it cannot take either. A weight, constant or + state the partitioner moves into Core ML is not an input. + """ + return arg.name not in kept_inside and ( + _cannot_be_coreml_input(arg, sizes) + or (_makes_symbolic_value(arg) and outside_coreml(arg)) + ) + + +def _find_readers_of_unsupported_inputs( + partitions: List[Partition], kept_inside: Collection[str], sizes: Collection[str] +) -> 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 outside of Core ML. + """ + inside = {node for partition in partitions for node in partition.nodes} + return { + node + for partition in partitions + for node in partition.nodes + if any( + arg not in partition.nodes + and _cannot_be_read_from_outside( + arg, kept_inside, sizes, lambda arg: arg not in inside + ) + for arg in node.all_input_nodes + ) + } + + +def _cannot_be_coreml_output(node: torch.fx.Node) -> bool: + """Whether node makes a value the Core ML delegate cannot return. + + The delegate has no int8 data type to hand an int8 tensor back to + ExecuTorch, and it returns only tensors, so a symbolic value cannot be an + output either. + """ + return _makes_tensor(node, (torch.int8,)) or _makes_symbolic_value(node) + + +def _find_makers_of_unsupported_outputs( + partitions: List[Partition], state_writes: Collection[torch.fx.Node] +) -> Set[torch.fx.Node]: + """The nodes that make a value the Core ML delegate cannot return, for use + outside their partition. A write of state that Core ML keeps is no output. + + 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 _cannot_be_coreml_output(node) + and any( + user not in partition.nodes + and not (user.op == "output" and node in state_writes) + for user in node.users + ) + } + + +def _find_unsupported_inputs_of_left_out_ops( + left_out: Set[torch.fx.Node], + inside: Set[torch.fx.Node], + kept_inside: Collection[str], + sizes: Collection[str], +) -> Set[torch.fx.Node]: + """The tensors Core ML cannot take as an input that explain why the ops in + left_out run outside of Core ML.""" + read: Set[torch.fx.Node] = set() + for node in left_out: + args = [ + arg + for arg in node.all_input_nodes + if arg.name not in kept_inside + and isinstance(arg.meta.get("val"), torch.Tensor) + and _cannot_be_coreml_input(arg, sizes) + ] + # A value the delegate cannot return, made by an op left out of Core + # ML, is no reason for its readers: that op may be out only for them. + outside = [ + arg + for arg in args + if arg not in inside + and not (arg in left_out and _cannot_be_coreml_output(arg)) + ] + if outside: + read.update(arg for arg in outside if arg not in left_out) + # A tensor that a Core ML partition makes explains the op only when no + # tensor from outside does, and the op is not left out for making a + # value the delegate cannot return. + elif not ( + _cannot_be_coreml_output(node) + and any(user not in left_out for user in node.users) + ): + read.update(arg for arg in args if arg in inside) + return read + + +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 _find_ops_left_out_of_any_partition( + graph_module: torch.fx.GraphModule, + op_support: OperatorSupportBase, + kept_inside: Collection[str], + sizes: Collection[str], + state_writes: Collection[torch.fx.Node], +) -> Set[torch.fx.Node]: + """The ops _propose_partitions leaves out however it partitions the rest: + each op that reads a tensor Core ML cannot take as an input, or a symbolic + value, from an op outside of Core ML, and each op that makes an int8 tensor + or a symbolic value an op outside of Core ML uses, other than to write + state Core ML may keep. + """ + left_out: Set[torch.fx.Node] = set() + supported = { + node + for node in graph_module.graph.nodes + if op_support.is_node_supported(None, node) + } + + def outside(node: torch.fx.Node) -> bool: + return node in left_out or node not in supported + + # Leaving an op out can make an op before or after it read or make such a + # tensor across the edge of Core ML, so repeat until nothing changes. + changed = True + while changed: + changed = False + for node in graph_module.graph.nodes: + if outside(node): + continue + if any( + outside(arg) + and _cannot_be_read_from_outside(arg, kept_inside, sizes, outside) + for arg in node.all_input_nodes + ): + left_out.add(node) + changed = True + elif _cannot_be_coreml_output(node) and any( + outside(user) and not (user.op == "output" and node in state_writes) + for user in node.users + ): + maker = node.args[0] if node.target is operator.getitem else node + changed |= maker not in left_out + left_out.add(maker) + return left_out + + +def _computed_from_kept_inside( + graph_module: torch.fx.GraphModule, kept_inside: Collection[str] +) -> Set[torch.fx.Node]: + """The nodes that read no tensor but the weights, constants and state the + partitioner moves into Core ML, or nodes computed from them.""" + nodes: Set[torch.fx.Node] = set() + for node in graph_module.graph.nodes: + if node.op == "placeholder": + if node.name in kept_inside: + nodes.add(node) + elif all(arg in nodes for arg in node.all_input_nodes if _makes_tensor(arg)): + nodes.add(node) + return nodes + + +def _find_ops_partitioning_may_leave_out( + graph_module: torch.fx.GraphModule, + kept_inside: Collection[str], + sizes: Collection[str], + is_supported: Callable[[torch.fx.Node], bool], + states: Collection[str], +) -> Set[torch.fx.Node]: + """The ops _propose_partitions may leave out of Core ML. + + Those that read a tensor Core ML cannot take as an input, make a value the + delegate cannot return, or read a symbolic value from an op that may be + outside of Core ML set it off. State can stay in the program, so reading it + can be reading an input. Once it leaves out one, a partition left without + tensor inputs moves out too, and such a partition holds only ops that + compute from the weights, constants and state the partitioner moves into + Core ML. + """ + weights_and_constants = set(kept_inside).difference(states) + set_off: Set[torch.fx.Node] = set() + for node in graph_module.graph.nodes: + if node.op == "placeholder": + continue + if _cannot_be_coreml_output(node) or any( + _cannot_be_read_from_outside( + arg, + weights_and_constants, + sizes, + lambda arg: arg in set_off or not is_supported(arg), + ) + for arg in node.all_input_nodes + ): + set_off.add(node) + if not set_off: + return set() + return set_off | _computed_from_kept_inside(graph_module, kept_inside) + + +def _states_left_in_program( + partitions: List[Partition], + state_writers: Dict[torch.fx.Node, List[torch.fx.Node]], +) -> Set[str]: + """The state tag_mutated_buffer leaves in the program for these partitions: + state that an op outside one partition reads or writes, also through its + new value. + """ + partition_of = { + node: partition.id for partition in partitions for node in partition.nodes + } + left = set() + for state, writers in state_writers.items(): + users = [*state.users, *writers] + users += [ + user for writer in writers for user in writer.users if user.op != "output" + ] + ids = {partition_of.get(user) for user in users} + if None in ids or len(ids) > 1: + left.add(state.name) + return left + + +def _propose_partitions( + graph_module: torch.fx.GraphModule, + op_support: OperatorSupportBase, + kept_inside: Collection[str], + sizes: Collection[str], + allow_partitions_without_inputs: bool, + state_writers: Dict[torch.fx.Node, List[torch.fx.Node]], +) -> Tuple[List[Partition], Set[torch.fx.Node]]: + """Propose Core ML partitions, leaving out every op that reads a value + Core ML cannot take as an input from outside its partition or makes one + the delegate cannot return for use outside, and then every partition this + leaves without a tensor input. State the partitions leave in the program + is an input. + + Returns the partitions and the ops left out. + """ + # Searching for partitions is slow on a large graph, so first leave out + # the ops that do not depend on how the search goes. Most graphs then need + # a single search. + left_out = _find_ops_left_out_of_any_partition( + graph_module, + op_support, + kept_inside, + sizes, + {writer for writers in state_writers.values() for writer in writers}, + ) + 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() + left_in_program = _states_left_in_program(partitions, state_writers) + kept = set(kept_inside).difference(left_in_program) + state_writes = { + writer + for state, writers in state_writers.items() + if state.name not in left_in_program + for writer in writers + } + new = _find_readers_of_unsupported_inputs( + partitions, kept, sizes + ) | _find_makers_of_unsupported_outputs(partitions, state_writes) + 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. So once any op moves out, + # every such partition that Core ML cannot take moves out too. If + # none does, they stay: below iOS 18, lowering then fails on a + # partition with no inputs and asks for iOS 18. + # _find_ops_partitioning_may_leave_out relies on this. + new |= _find_nodes_of_partitions_without_tensor_inputs( + partitions, kept, 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 _decompose_alone( + node: torch.fx.Node, decompositions: Dict[torch._ops.OperatorBase, Callable] +) -> torch.fx.GraphModule: + """Traces the decomposition of node by itself, from the values of its + inputs.""" + inputs = node.all_input_nodes + + def call(*args): + env = dict(zip(inputs, args)) + return node.target( + *torch.fx.node.map_arg(node.args, env.__getitem__), + **torch.fx.node.map_arg(node.kwargs, env.__getitem__), + ) + + # pre_dispatch keeps whole a composite op the table leaves out, such as + # matmul, as lowering does. + return make_fx( + call, + decomposition_table=decompositions, + tracing_mode="fake", + pre_dispatch=True, + )(*(arg.meta["val"] for arg in inputs)) + + +def _decomposing_may_leave_out( + graph_module: torch.fx.GraphModule, + decompositions: Dict[torch._ops.OperatorBase, Callable], + not_kept: Collection[torch._ops.OpOverload], + kept: Collection[torch._ops.OpOverload], + from_kept_inside: Collection[torch.fx.Node], + sizes: Collection[str], + is_supported: Callable[[torch.fx.Node], bool], +) -> bool: + """Whether decomposing the calls of the not_kept ops may make a call of a + kept op that Core ML rejects, or make partition() leave out a call of a + kept op it does not leave out in graph_module. + + That needs a decomposition that makes a call of a kept op that Core ML + rejects, or that reads or makes a value Core ML cannot take or return, or a + decomposed call that reads only the weights, constants and state the + partitioner moves into Core ML and makes a tensor or a call of a kept op, + since its decomposition can leave a partition without tensor inputs. Each + call is traced by itself, in milliseconds, while decomposing a large + program takes minutes. A call that cannot be traced by itself may matter. + """ + for node in graph_module.graph.nodes: + if node.target not in decompositions or node.target not in not_kept: + continue + from_inside = all( + arg in from_kept_inside + for arg in node.all_input_nodes + if _makes_tensor(arg) + ) + if from_inside and _makes_tensor(node): + return True + try: + calls = _decompose_alone(node, decompositions).graph.nodes + except Exception: + return True + if any( + call.target in kept + and ( + from_inside + or not is_supported(call) + or _cannot_be_coreml_output(call) + or any( + _cannot_be_coreml_input(arg, sizes) or _makes_symbolic_value(arg) + for arg in call.all_input_nodes + ) + ) + for call in calls + ): + return True + return False + + +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: @@ -249,23 +721,6 @@ def should_override_support(self, node) -> bool: ) return True - # TODO: enable this after bugs in ExecuTorch's partitioner are fixed - # # If lower_full_graph=False, do not partition nodes with symbolic args because it can result in symbolic args - # # in the placeholders due to partitioning, which CoreML does not support - # if not self.lower_full_graph and any( - # isinstance(arg, torch.fx.Node) - # and isinstance( - # arg.meta.get("val", None), - # (torch.SymInt, torch.SymBool, torch.SymFloat), - # ) - # for arg in node.args - # ): - # self.log_once( - # "Skipping op for CoreML delegation because it contains symbolic args: " - # + node_target_name - # ) - # return True - return False def is_node_supported(self, submodules, node: torch.fx.Node) -> bool: @@ -344,6 +799,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,16 +849,89 @@ def partition(self, exported_program: ExportedProgram) -> PartitionResult: logger.info("CoreMLPartitioner::partition") partition_tags = {} - capability_partitioner = CapabilityBasedPartitioner( + state_writers = self._state_writers(exported_program) + kept_inside = self._placeholders_kept_inside(exported_program, state_writers) + sizes = _dynamic_size_names(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, ), - allows_single_node_partition=True, + kept_inside, + sizes, + self._targets_ios18(), + state_writers, ) - partition_list = capability_partitioner.propose_partitions() + if left_out: + kept = set(kept_inside).difference( + _states_left_in_program(partition_list, state_writers) + ) + inside = {node for partition in partition_list for node in partition.nodes} + inputs = sorted( + f"{arg.name} ({arg.meta['val'].dtype}, {tuple(arg.meta['val'].shape)})" + for arg in _find_unsupported_inputs_of_left_out_ops( + left_out, inside, kept, sizes + ) + ) + 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) + ) + symbols = sorted( + node.name + for node in left_out + if ( + _makes_symbolic_value(node) + and any(user not in left_out for user in node.users) + ) + or any( + _makes_symbolic_value(arg) and arg not in inside + for arg in node.all_input_nodes + ) + ) + reasons = [] + if inputs: + reasons.append( + "reads a tensor CoreML cannot take as an input, of uint8 or " + "int8 or with a size computed from a dynamic size, from " + "outside a CoreML partition: " + ", ".join(inputs) + ) + if outputs: + reasons.append( + "makes an int8 tensor, which the CoreML delegate cannot " + "return, for use outside of CoreML: " + ", ".join(outputs) + ) + if symbols: + reasons.append( + "makes or reads a symbolic value, which CoreML can neither " + "take as an input nor return, outside of CoreML: " + + ", ".join(symbols) + ) + msg = ( + f"The model {', and '.join(reasons)}." + if reasons + else "The model uses a tensor or a symbolic value CoreML cannot handle." + ) + if self.lower_full_graph: + raise NotImplementedError( + f"{msg} Set lower_full_graph=False in the CoreMLPartitioner " + "to run these ops 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}" @@ -427,12 +956,70 @@ 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 _targets_ios18(self) -> bool: + """Whether the deployment target is iOS 18 or later.""" + 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, + state_writers: Dict[torch.fx.Node, List[torch.fx.Node]], + ) -> 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 = {state.name for state in state_writers} + 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_writers( + self, ep: ExportedProgram + ) -> Dict[torch.fx.Node, List[torch.fx.Node]]: + """The nodes that write each mutable buffer this partitioner may move + into Core ML as state, by the placeholder that reads it. A buffer that + no op reads, or whose new value is also returned, stays in the program. + """ + if not self.take_over_mutable_buffer: + return {} + signature = ep.graph_signature + nodes = {node.name: node for node in ep.graph.nodes} + returned = { + spec.arg.name + for spec in signature.output_specs + if spec.kind == OutputKind.USER_OUTPUT + } + writers: Dict[str, List[torch.fx.Node]] = {} + for name, buffer in signature.buffers_to_mutate.items(): + writers.setdefault(buffer, []).append(nodes[name]) + return { + nodes[name]: writers[buffer] + for name, buffer in signature.inputs_to_buffers.items() + if buffer in writers + and nodes[name].users + and not any(writer.name in returned for writer in writers[buffer]) + } + 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, @@ -452,21 +1039,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 +1064,105 @@ 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 value CoreML cannot take or return, 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. + + partition() sees the program once lowering decomposes the ops that are + not kept, which can make new calls of a kept op that CoreML rejects, + such as the matmul of a bfloat16 linear, or change the calls it moves + out. So when that may matter, this decomposes the program the same way + and finds the calls partition() leaves out, 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, self._state_writers(ep)) + skip = set(self.skip_ops_for_coreml_delegation) + support = create_op_support( + lambda submodules, node: is_supported(submodules, node) + and not (skip and _edge_op_names(node.target) & skip) + ) + from_kept_inside = _computed_from_kept_inside(ep.graph_module, kept_inside) + not_kept = set(rejected) + decompositions = _default_decomposition_table() + while keep: + table = dict(decompositions) + for op in keep: + table.pop(op, None) + kept = set(keep) + decomposed = _decomposing_may_leave_out( + ep.graph_module, + table, + not_kept, + kept, + from_kept_inside, + _dynamic_size_names(ep), + lambda node: is_supported(None, node), + ) + program = ep.run_decompositions(table) if decomposed else ep + graph_module = program.graph_module + sizes = _dynamic_size_names(program) + state_writers = self._state_writers(program) + calls = [node for node in graph_module.graph.nodes if node.target in kept] + left_out = _find_ops_left_out_of_any_partition( + graph_module, + support, + kept_inside, + sizes, + {writer for writers in state_writers.values() for writer in writers}, + ) + # The search is slow on a large graph, so run it only when it may + # leave out a call of an op lowering would otherwise decompose. + # Lowering keeps any other op whole either way, so whether such an + # op is kept changes nothing. + uncertain = _find_ops_partitioning_may_leave_out( + graph_module, + kept_inside, + sizes, + lambda node: support.is_node_supported(None, node), + {state.name for state in state_writers}, + ).difference(left_out) + if any( + node.target in decompositions and node in uncertain for node in calls + ): + _, left_out = _propose_partitions( + graph_module, + support, + kept_inside, + sizes, + self._targets_ios18(), + state_writers, + ) + else: + left_out |= uncertain + late = { + node.target + for node in calls + if node in left_out or (decomposed and not is_supported(None, node)) + } + if not late: + break + for op in late & decompositions.keys(): + 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..7e4e9cc82bf 100644 --- a/backends/apple/coreml/test/test_coreml_partitioner.py +++ b/backends/apple/coreml/test/test_coreml_partitioner.py @@ -7,6 +7,7 @@ import logging import sys import unittest +from unittest import mock import coremltools as ct @@ -16,8 +17,12 @@ import torchvision from executorch.backends.apple.coreml.compiler import CoreMLBackend -from executorch.backends.apple.coreml.partition import CoreMLPartitioner +from executorch.backends.apple.coreml.partition import ( + coreml_partitioner, + CoreMLPartitioner, +) from executorch.exir.backend.utils import format_delegated_graph +from executorch.exir.passes.init_mutable_pass import InitializedMutableBufferPass @torch.library.custom_op("unsupported::linear", mutates_args=()) @@ -50,6 +55,20 @@ def is_fbcode(): class TestCoreMLPartitioner(unittest.TestCase): edge_compile_config = executorch.exir.EdgeCompileConfig() + def _lower_with_one_search(self, ep, partitioner): + """Lowers ep, checking that partitions are searched for once, since + each search is slow on a large model.""" + with mock.patch.object( + coreml_partitioner, + "CapabilityBasedPartitioner", + wraps=coreml_partitioner.CapabilityBasedPartitioner, + ) as search: + edge = executorch.exir.to_edge_transform_and_lower( + ep, partitioner=[partitioner] + ) + self.assertEqual(search.call_count, 1) + return edge + def test_add_sub_skip_mm(self): class Model(torch.nn.Module): def forward(self, a, x, b): @@ -521,6 +540,908 @@ 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_int8_output(self): + """ + The Core ML delegate has no int8 data type to hand an int8 tensor back + to ExecuTorch. 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] + + # The cast to int8 reads a uint8 tensor that Core ML makes, which is no + # input from outside, so the error must not name it. + class FromUint8(torch.nn.Module): + def forward(self, x): + return (x.abs() * 2).to(torch.uint8).to(torch.int8) + + # 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(), FromUint8().eval()): + with self.subTest(model=type(model).__name__): + ep = torch.export.export(model, (x,), strict=True) + edge = self._lower_with_one_search(ep, 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" + ) as error: + executorch.exir.to_edge_transform_and_lower( + ep, partitioner=[CoreMLPartitioner(lower_full_graph=True)] + ) + self.assertNotIn("cannot take as an input", str(error.exception)) + + # Core ML makes this int8 tensor in one partition and reads it in + # another, past the asinh, which runs on the CPU, so the cast to int8 + # runs on the CPU too. + class AcrossPartitions(torch.nn.Module): + def forward(self, x): + a = (x * 2).to(torch.int8) + return a.to(torch.int32).float() * torch.asinh(a.float()) + + model = AcrossPartitions().eval() + x = torch.randint(-20, 20, (2, 4)).float() + ep = torch.export.export(model, (x,), strict=True) + edge = executorch.exir.to_edge_transform_and_lower( + ep, partitioner=[CoreMLPartitioner()] + ) + delegate_outputs = [ + n.meta["val"].dtype + for n in edge.exported_program().graph.nodes + if n.op == "call_function" and n.target.__name__ == "getitem" + ] + self.assertTrue(delegate_outputs) + self.assertNotIn(torch.int8, delegate_outputs) + 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, rtol=1e-3)) + + # Here the cast to int8 reads a uint8 tensor from another partition, + # and only ops on the CPU read its int8, so the log names that tensor. + class Uint8AcrossPartitions(torch.nn.Module): + def forward(self, x): + t = (x.abs() * 2).to(torch.uint8) + return t.to(torch.int8).to(torch.int32).float() * torch.asinh(t.float()) + + # It names that tensor too when the op that reads it, here maximum, + # reads an int8 tensor as well, made on the CPU only because maximum + # runs there. + class Int8MadeForTheReader(torch.nn.Module): + def forward(self, x): + t = (x.abs() * 2).to(torch.uint8) + q = (x * 2).to(torch.int8) + return torch.maximum(q, t).float() * torch.asinh(t.float()) + + for model in (Uint8AcrossPartitions().eval(), Int8MadeForTheReader().eval()): + with self.subTest(model=type(model).__name__): + ep = torch.export.export(model, (x,), strict=True) + with self.assertLogs(coreml_partitioner.logger, logging.INFO) as logs: + executorch.exir.to_edge_transform_and_lower( + ep, partitioner=[CoreMLPartitioner()] + ) + self.assertIn("cannot take as an input", "\n".join(logs.output)) + + def test_int8_state(self): + """ + Writing int8 state that Core ML keeps is no int8 output: Core ML keeps + the state as float16 state, so the whole model stays in Core ML. + When the partitioner does not take state over, Core ML cannot read the + int8 state, so the cast and the add that read it run on the CPU, with + the cast that makes the int8 tensor the add reads. When a CPU op reads + the state, or no op does, it stays in the program, and writing it is an + int8 output, so the cast that makes its new value runs on the CPU. + """ + + class Kept(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer( + "state", torch.tensor([[1, -2, 3, -4]], dtype=torch.int8) + ) + + def forward(self, x): + y = self.state.float() * x + self.state.add_((x > 0).to(torch.int8)) + return y + + class KeptFromZero(Kept): + # The Core ML delegate zeroes the state it holds when it loads. + def __init__(self): + super().__init__() + self.state.zero_() + + class NeverRead(Kept): + def forward(self, x): + self.state.copy_((x * 4).to(torch.int8)) + return x + 1 + + class ReadOnCpu(Kept): + def forward(self, x): + y = torch.asinh(self.state) * x + self.state.copy_((x > 0).to(torch.int8)) + return y + + x = torch.tensor([[0.3, -0.2, 0.4, -0.1]]) + compile_specs = CoreMLBackend.generate_compile_specs( + minimum_deployment_target=ct.target.iOS18 + ) + cast = "dim_order_ops._to_dim_order_copy.default" + for model_class, take_over, cpu_ops in ( + (KeptFromZero, True, set()), + (Kept, False, {cast, "aten.add.Tensor"}), + (ReadOnCpu, True, {cast, "aten.asinh.default"}), + (NeverRead, True, {cast}), + ): + with self.subTest(model=model_class.__name__, take_over=take_over): + ep = torch.export.export(model_class().eval(), (x,), strict=True) + partitioner = CoreMLPartitioner( + compile_specs=compile_specs, take_over_mutable_buffer=take_over + ) + edge = executorch.exir.to_edge_transform_and_lower( + ep, partitioner=[partitioner] + ) + 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"}, cpu_ops + ) + # Mutable buffers start uninitialized without this pass. + config = executorch.exir.ExecutorchBackendConfig( + passes=[InitializedMutableBufferPass(["state"])] + ) + buffer = edge.to_executorch(config).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_int8_input_is_read_outside_the_delegate(self): + """ + Core ML cannot take a uint8 or int8 input, so the ops that read an int8 + image stay on the CPU up to its conversion to float, and the float ops + after it are still delegated. With lower_full_graph=True, lowering + fails and names the input, but not a uint8 tensor that Core ML makes + and an op on the CPU reads. + """ + + 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() + x = torch.randint(-128, 128, (8, 8, 3), dtype=torch.int8) + 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\\.int8, "): + executorch.exir.to_edge_transform_and_lower( + ep, partitioner=[CoreMLPartitioner(lower_full_graph=True)] + ) + + class Mixed(torch.nn.Module): + def forward(self, image, x): + return (image + (x * 255).to(torch.uint8)).float() * 2 + + class Chain(torch.nn.Module): + def forward(self, image, x): + return ((image + 1) + (x * 255).to(torch.uint8)).float() * 2 + + inputs = (torch.randint(0, 100, (2, 4), dtype=torch.uint8), torch.rand(2, 4)) + for model in (Mixed(), Chain()): + with self.subTest(model=type(model).__name__): + ep = torch.export.export(model.eval(), inputs, strict=True) + with self.assertRaises(NotImplementedError) as error: + executorch.exir.to_edge_transform_and_lower( + ep, partitioner=[CoreMLPartitioner(lower_full_graph=True)] + ) + self.assertIn( + "partition: image (torch.uint8, (2, 4)).", str(error.exception) + ) + + def test_uint8_input_read_by_an_op_without_a_cpu_kernel(self): + """ + A reflect pad that reads a uint8 input runs on the CPU. Core ML keeps + pad whole, but pad is not a core ATen op and has no portable kernel, so + it is decomposed there. Below iOS 18 the index math it decomposes into + runs on the CPU too, since it would leave Core ML a partition with no + inputs. So does a pad that would be such a partition itself, also + once lowering decomposes the ops before it. + """ + + class Model(torch.nn.Module): + def forward(self, image): + image = torch.nn.functional.pad(image, (1, 1, 1, 1), mode="reflect") + return image.float() / 255.0 + + class PadOfBuffer(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("w", torch.rand(1, 3, 8)) + + def forward(self, image): + p = torch.nn.functional.pad(self.w, (1, 1), mode="reflect") + return torch.asinh(p) + image.float() + + class PadOfFilledBuffer(PadOfBuffer): + def __init__(self): + super().__init__() + self.register_buffer("cols", torch.tensor([0, 2])) + + def forward(self, image): + w = self.w.index_fill(2, self.cols, 0.0) + p = torch.nn.functional.pad(w, (1, 1), mode="reflect") + return torch.asinh(p) + image.float() + + image = torch.randint(0, 256, (1, 3, 10), dtype=torch.uint8) + for model, x in ( + (Model(), torch.randint(0, 256, (1, 3, 8, 8), dtype=torch.uint8)), + # The pad reads only a buffer Core ML takes over, and the asinh + # after it runs on the CPU, so it would be a partition with no + # inputs. + (PadOfBuffer(), image), + # Core ML rejects index_fill, so the pad reads only buffers once + # lowering decomposes it, into a scalar_tensor and an index_put. + (PadOfFilledBuffer(), image), + ): + 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.assertIn("aten.arange.start_step", op_names) + self.assertNotIn("aten.pad.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, rtol=1e-3)) + + def test_pad_of_buffer_next_to_uint8_input_stays_whole(self): + """ + A pad that reads only a buffer Core ML takes over would be a partition + without inputs if the ops after it moved to the CPU. In a model with a + uint8 input some ops do move, so the search for partitions decides + whether pad stays whole. Here the partition of the pad also reads x, + so pad stays whole. + """ + + class Model(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("w", torch.rand(1, 3, 4)) + + def forward(self, image, x): + return image.float() + torch.nn.functional.pad(self.w, (1, 1)) * x + + model = Model().eval() + inputs = ( + torch.randint(0, 256, (1, 3, 6), dtype=torch.uint8), + torch.rand(1, 3, 6), + ) + ep = torch.export.export(model, inputs, strict=True) + keep, _ = CoreMLPartitioner().ops_to_not_decompose(ep) + self.assertIn(torch.ops.aten.pad.default, keep) + edge = executorch.exir.to_edge_transform_and_lower( + ep, partitioner=[CoreMLPartitioner()] + ) + buffer = edge.to_executorch().buffer + if _TEST_RUNTIME: + program = Runtime.get().load_program(buffer) + out = program.load_method("forward").execute(inputs)[0] + self.assertTrue(torch.allclose(out, model(*inputs), atol=1e-2, rtol=1e-3)) + + 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, also from iOS 18, where a partition with no inputs at all + may stay in Core ML. So do a reshape to a size read from the uint8 + input, a scale computed from it, and the add after the reshape, since + Core ML cannot take a tensor sized by an expression of the height, such + as 3 * height, as an input. + """ + + 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) + 1 + + 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, + ) + compile_specs = CoreMLBackend.generate_compile_specs( + minimum_deployment_target=ct.target.iOS18 + ) + edge = self._lower_with_one_search( + ep, 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", + "aten.add.Tensor", + ): + 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_symbolic_value_read_outside_the_delegate(self): + """ + The Core ML delegate returns only tensors, so the item of a flag, a + symbolic value that sym_ite reads on the CPU, runs on the CPU too, + below and from iOS 18, where a partition with no inputs at all may + stay in Core ML. So does a size the model returns, and with + lower_full_graph=True lowering fails and names it. The reshape to a + size computed from it runs on the CPU too, and so does the op after it, + since Core ML cannot take a tensor sized by an expression of a dynamic + size, such as 3 * rows, as an input. The other ops are still + delegated. + """ + + class Model(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("flag", torch.tensor(True)) + + def forward(self, x): + k = torch.sym_ite(self.flag.item(), 2, 3) + return (torch.relu(x) + 1) * k + + model = Model().eval() + x = torch.rand(2, 3) + for target in (ct.target.iOS15, ct.target.iOS18): + with self.subTest(target=target): + ep = torch.export.export(model, (x,), strict=False) + compile_specs = CoreMLBackend.generate_compile_specs( + minimum_deployment_target=target + ) + edge = self._lower_with_one_search( + ep, 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("_local_scalar_dense.default", op_names) + self.assertIn("executorch_call_delegate", op_names) + for name in ("aten.relu.default", "aten.add.Tensor", "aten.mul.Tensor"): + 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)) + + class ReturnsSize(torch.nn.Module): + def forward(self, x): + rows = x.shape[0] + return (torch.relu(x) + 1).reshape(1, rows * 3) * 2, rows + + ep = torch.export.export( + ReturnsSize().eval(), + (x,), + dynamic_shapes=({0: torch.export.Dim("rows", min=2, max=8)},), + strict=True, + ) + edge = self._lower_with_one_search(ep, CoreMLPartitioner()) + op_names = [ + n.target.__name__ + for n in edge.exported_program().graph.nodes + if n.op == "call_function" + ] + for name in ("sym_size.int", "aten.view_copy.default", "aten.mul.Tensor"): + self.assertIn(name, op_names) + for name in ("aten.relu.default", "aten.add.Tensor"): + self.assertNotIn(name, op_names) + self.assertIn("executorch_call_delegate", op_names) + buffer = edge.to_executorch().buffer + if _TEST_RUNTIME: + y = torch.rand(5, 3) + program = Runtime.get().load_program(buffer) + out, rows = program.load_method("forward").execute((y,)) + expected = (torch.relu(y) + 1).reshape(1, 15) * 2 + self.assertTrue(torch.allclose(out, expected, atol=1e-2)) + self.assertEqual(rows, 5) + with self.assertRaisesRegex(NotImplementedError, "symbolic value.*sym_size"): + executorch.exir.to_edge_transform_and_lower( + ep, partitioner=[CoreMLPartitioner(lower_full_graph=True)] + ) + + # The size of a uint8 input runs on the CPU only because it reads that + # input, so the error names the input and not the size. + class SizeOfUint8(torch.nn.Module): + def forward(self, image): + return image.float().reshape(image.shape[0], -1) * 2 + + ep = torch.export.export( + SizeOfUint8().eval(), + (torch.randint(0, 100, (4, 6), dtype=torch.uint8),), + dynamic_shapes=({0: torch.export.Dim("rows", min=2, max=8)},), + strict=True, + ) + with self.assertRaises(NotImplementedError) as error: + executorch.exir.to_edge_transform_and_lower( + ep, partitioner=[CoreMLPartitioner(lower_full_graph=True)] + ) + self.assertIn("image (torch.uint8, ", str(error.exception)) + self.assertNotIn("sym_size", str(error.exception)) + + # Partitions on both sides of the asinh, which Core ML rejects, read the + # same size, so it runs on the CPU, and so do the reshapes, which the + # log explains with the symbolic value they read. + class SharedSize(torch.nn.Module): + def forward(self, x): + rows = x.shape[0] + return torch.asinh((x * 2).reshape(rows, 2, 2)).reshape(rows, 4) + 1 + + ep = torch.export.export( + SharedSize().eval(), + (torch.rand(4, 4),), + dynamic_shapes=({0: torch.export.Dim("rows", min=2, max=8)},), + strict=True, + ) + with self.assertLogs(coreml_partitioner.logger, logging.INFO) as logs: + executorch.exir.to_edge_transform_and_lower( + ep, partitioner=[CoreMLPartitioner()] + ) + self.assertIn("makes or reads a symbolic value", "\n".join(logs.output)) + + def test_input_sized_by_a_declared_expression_stays_in_the_delegate(self): + """ + coremltools takes an input sized by an expression of a dynamic size + when the exported program declares it, such as 2 * n here, so the ops + that read it stay in Core ML. + """ + + class Model(torch.nn.Module): + def forward(self, x, y): + return x.reshape(1, -1, 2).sum(-1) + y + + size = torch.export.Dim("n", min=2, max=32) + ep = torch.export.export( + Model().eval(), + (torch.rand(1, 8), torch.rand(1, 4)), + dynamic_shapes={"x": {1: 2 * size}, "y": {1: size}}, + 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: + x, y = torch.rand(1, 12), torch.rand(1, 6) + program = Runtime.get().load_program(buffer) + out = program.load_method("forward").execute((x, y))[0] + self.assertTrue(torch.allclose(out, Model()(x, y), atol=1e-2)) + + def test_program_is_decomposed_only_when_it_may_matter(self): + """ + Finding the ops partition() runs on the CPU decomposes the whole + program, as lowering will, only when that may change the answer, since + that takes minutes on a large model. Here index_fill, which Core ML + rejects, decomposes into ops the program does not keep. The item of a + buffer runs on the CPU, since the model returns its value, but it makes + no tensor, so decomposing it leaves every partition as it is. + """ + + class Fill(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("cols", torch.tensor([0, 2])) + + def forward(self, x): + return x.index_fill(1, self.cols, 0.0) + 1 + + class ItemOfBuffer(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("n", torch.tensor(3)) + + def forward(self, x): + return x + 1, self.n.item() + + decompose = torch.export.ExportedProgram.run_decompositions + for model in (Fill(), ItemOfBuffer()): + with self.subTest(model=type(model).__name__): + ep = torch.export.export(model.eval(), (torch.rand(2, 4),), strict=True) + with mock.patch.object( + torch.export.ExportedProgram, + "run_decompositions", + autospec=True, + side_effect=decompose, + ) as run_decompositions: + keep, _ = CoreMLPartitioner().ops_to_not_decompose(ep) + self.assertIn(torch.ops.aten.add.Tensor, keep) + self.assertEqual(run_decompositions.call_count, 0) + + def test_uint8_input_read_by_a_kept_op_once_decomposed(self): + """ + isposinf of an int tensor decomposes into zeros_like, and the portable + kernels have no zeros_like. The zeros_like that Core ML runs keeps + zeros_like whole, so the one that isposinf of a uint8 input decomposes + into would run whole on the CPU. Tracing the decomposition of + isposinf by itself shows that it reads the uint8 input, so the program + is decomposed to find that call, and zeros_like is decomposed too. + """ + + class Model(torch.nn.Module): + def forward(self, image, x): + return torch.zeros_like(x) + x, torch.isposinf(image) + + model = Model().eval() + inputs = ( + torch.randint(0, 256, (1, 3, 4), dtype=torch.uint8), + torch.rand(1, 3, 4), + ) + 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.zeros_like.default", op_names) + buffer = edge.to_executorch().buffer + if _TEST_RUNTIME: + program = Runtime.get().load_program(buffer) + outputs = program.load_method("forward").execute(inputs) + for out, ref in zip(outputs, model(*inputs)): + self.assertTrue(torch.allclose(out.float(), ref.float(), atol=1e-2)) + + def test_uint8_input_with_state(self): + """ + A CPU op that reads a uint8 input can sit between the Core ML ops that + read and write state, or make its new value. Such state stays outside + of Core ML, which reads it as an input. + """ + + 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 + + x = torch.randint(0, 3, (1, 3, 2, 2), dtype=torch.uint8) + for model_class in (Split, CopyIn): + 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=[CoreMLPartitioner()] + ) + op_names = [ + n.target.__name__ + for n in edge.exported_program().graph.nodes + if n.op == "call_function" + ] + # Below iOS 18 Core ML rejects a partition without inputs, but + # the state is an input here, so Core ML takes state * 2 too. + self.assertNotIn("aten.mul.Tensor", op_names) + self.assertIn( + "state", + edge.exported_program().graph_signature.buffers_to_mutate.values(), + ) + 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_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. So is uint8 state that stays in + the program, because a CPU op reads it too, also through its new value, + or because the model returns its new value. A pad of such state runs on + the CPU decomposed, since the portable kernels have no pad. With + lower_full_graph=True, lowering fails and names that state. + """ + + class Model(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("lut", torch.arange(0, 256, 32, dtype=torch.uint8)) + + def forward(self, x): + return x * (self.lut.float() / 255.0) + + class ReadOnCpu(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer( + "state", torch.tensor([[1, 2, 3, 4]], dtype=torch.uint8) + ) + + def forward(self, image, x): + y = torch.nn.functional.pad(self.state, (1, 1)).float() * x + y = y + (image + self.state).float().sum() + self.state.add_(1) + return y + + class NewValueOnCpu(ReadOnCpu): + def forward(self, x): + y = self.state.float() * x + self.state.copy_((x * 10).to(torch.uint8)) + return y, torch.diagonal(self.state).float() + + class Returned(ReadOnCpu): + def forward(self, x): + self.state.add_(1) + return self.state.float() * x, self.state + + state_inputs = ( + torch.randint(0, 3, (1, 4), dtype=torch.uint8), + torch.rand(1, 6), + ) + x = torch.rand(1, 4) + for model_class, inputs, partitioner in ( + ( + Model, + (torch.randn(2, 8),), + CoreMLPartitioner(take_over_constant_data=False), + ), + (ReadOnCpu, state_inputs, CoreMLPartitioner()), + (NewValueOnCpu, (x,), CoreMLPartitioner()), + (Returned, (x,), CoreMLPartitioner()), + ): + with self.subTest(model=model_class.__name__): + ep = torch.export.export(model_class().eval(), inputs, strict=True) + edge = executorch.exir.to_edge_transform_and_lower( + ep, partitioner=[partitioner] + ) + 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) + # Mutable buffers start uninitialized without this pass. + config = executorch.exir.ExecutorchBackendConfig( + passes=[InitializedMutableBufferPass(["state"])] + ) + buffer = edge.to_executorch(config).buffer + if _TEST_RUNTIME: + method = Runtime.get().load_program(buffer).load_method("forward") + eager = model_class().eval() + for _ in range(2): + outputs = method.execute(inputs) + refs = eager(*inputs) + refs = refs if isinstance(refs, tuple) else (refs,) + for out, ref in zip(outputs, refs, strict=True): + self.assertTrue( + torch.allclose( + out.float(), ref.float(), atol=1e-2, rtol=1e-3 + ) + ) + with self.assertRaisesRegex(NotImplementedError, "b_state \\(torch\\.uint8, "): + executorch.exir.to_edge_transform_and_lower( + torch.export.export(ReadOnCpu().eval(), state_inputs, strict=True), + partitioner=[CoreMLPartitioner(lower_full_graph=True)], + ) + + def test_uint8_made_inside_the_delegate_keeps_ops_whole(self): + """ + A uint8 tensor made inside the delegate is not an input, so the pad + that reads it stays in the delegate, whole. + """ + + class Model(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" + ) + + model = Model().eval() + x = torch.rand(1, 3, 8) + ep = torch.export.export(model, (x,), strict=True) + edge = executorch.exir.to_edge_transform_and_lower( + ep, partitioner=[CoreMLPartitioner()] + ) + op_names = { + n.target.__name__ + for n in edge.exported_program().graph.nodes + if n.op == "call_function" + } + self.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 a skipped 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. The log names the + uint8 tensors that ops read from the CPU or from another partition. + """ + + 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 ( + (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) + with self.assertLogs(coreml_partitioner.logger, logging.INFO) as logs: + edge = executorch.exir.to_edge_transform_and_lower( + ep, + partitioner=[ + CoreMLPartitioner(skip_ops_for_coreml_delegation=skip) + ], + ) + self.assertEqual( + "cannot take as an input" in "\n".join(logs.output), + not isinstance(model, PadInt8), + ) + 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 39d64a3f01e..249942958e8 100644 --- a/backends/apple/coreml/test/test_torch_ops.py +++ b/backends/apple/coreml/test/test_torch_ops.py @@ -777,6 +777,43 @@ def forward(self, ids, x): [types.int32] * 4, ) + def test_fills_of_a_tensor_use_fill_like(self): + """ + A fill the shape of a tensor becomes a constant as big as that tensor + in the model when the shape is known at conversion, unless it is a + fill_like, which needs iOS 16, the target this test uses. + """ + + class Fills(torch.nn.Module): + def forward(self, x): + ones = torch.ones_like(x) + torch.zeros_like(x) + return ones + torch.full_like(x, 2) + x.clone().fill_(3) + + model = Fills().eval() + x = torch.rand(4, 64) + delegated_program = executorch.exir.to_edge_transform_and_lower( + torch.export.export(model, (x,)), + partitioner=[ + self._coreml_partitioner(minimum_deployment_target=ct.target.iOS16) + ], + ) + self.assertEqual( + [ + node.target.__name__ + for node in delegated_program.exported_program().graph.nodes + if node.op == "call_function" + ], + ["executorch_call_delegate", "getitem"], + ) + lowered = delegated_program.exported_program().graph_module.lowered_module_0 + program = ct.convert( + lowered.original_module, + convert_to="milinternal", + minimum_deployment_target=ct.target.iOS16, + ) + self.assertEqual(len(program.find_ops(op_type="fill_like")), 4) + self._compare_outputs(delegated_program.to_executorch(), model, (x,)) + @unittest.skipIf(is_fbcode(), "The coremltools in fbcode has no rms_norm.") def test_rms_norm_without_eps(self): """ diff --git a/docs/source/backends/coreml/coreml-partitioner.md b/docs/source/backends/coreml/coreml-partitioner.md index fa7a1de8277..9d5631b0903 100644 --- a/docs/source/backends/coreml/coreml-partitioner.md +++ b/docs/source/backends/coreml/coreml-partitioner.md @@ -6,8 +6,12 @@ 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. -- `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. +- `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 or gathers from, such as a lookup table, fails to lower with the default settings, since coremltools cannot cast or gather from 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, or an op has to run on the CPU as described below, 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. + +The Core ML delegate cannot take a uint8 or int8 tensor as an input, so the ops that read one, such as the cast of a uint8 camera frame, run on the CPU. For camera frames, converting each frame with the `ImageProcessor` in [`extension/image`](https://github.com/pytorch/executorch/tree/main/extension/image) and exporting the model with a float input lets those ops run in Core ML too. + +A few other values cannot cross the edge of a Core ML part either, so the ops on the Core ML side of that edge run on the CPU. Core ML cannot take a symbolic value, such as a size known only at run time, as an input. It also cannot take a tensor sized by an expression of a dynamic size, such as `2 * height`, unless the model was exported with that expression as a derived `Dim`. And the delegate cannot return an int8 tensor or a symbolic value. #### Core ML CompileSpec