diff --git a/backends/apple/coreml/compiler/coreml_preprocess.py b/backends/apple/coreml/compiler/coreml_preprocess.py index 931d1552606..8d44661f35e 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,46 @@ def _gathers_sharing_a_weight(mlmodel: ct.models.MLModel) -> List[str]: return shared +@register_pass(namespace="executorch", override=True, 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, and keeps a tile that + changes nothing, such as one from a same-shape expand, only when it makes + an output. So a partition can return one of its own inputs, or one value + twice, directly or through such a tile. Core ML then hands back one of + those outputs 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, tiles = [], [] + for output in function.outputs: + value = output + if ( + output.op is not None + and output.op.op_type == "tile" + and output.op.x.sym_type == output.sym_type + ): + tiles.append(output.op) + value = output.op.x + if value is not output or value in seen: + with function: + copy = mb.identity(x=value) + # coremltools wants every op to carry the source program's + # scopes, which set_outputs, unlike a var replacement, skips. + # The copy takes those of the output it replaces. + function._copy_scope_info(output, copy) + value = copy + seen.add(value) + outputs.append(value) + function.set_outputs(outputs) + # A pass list without dead_code_elimination would keep such a tile. + function.remove_ops([op for op in tiles if not op.outputs[0].child_ops]) + + class COMPILE_SPEC_KEYS(Enum): COMPUTE_UNITS = "compute_units" MODEL_TYPE = "model_type" @@ -565,6 +608,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..ed9cf0d5e18 100644 --- a/backends/apple/coreml/compiler/torch_ops.py +++ b/backends/apple/coreml/compiler/torch_ops.py @@ -8,13 +8,19 @@ # coremltools than is used by ExecuTorch. Each op registered here should have a link to a PR in coremltools that adds # the op to the coremltools library. +import copy +import functools + 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, @@ -25,6 +31,7 @@ register_torch_op, ) from coremltools.converters.mil.mil import types +from coremltools.converters.mil.mil.types.symbolic import is_symbolic from executorch.exir.dim_order_utils import get_memory_format @@ -49,6 +56,155 @@ 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. https://github.com/apple/coremltools/pull/2865 +# fixes zeros_like; no upstream issue covers the others yet. For example, +# torch.full((1, 16), 0) is int64 in torch, and coremltools makes it fp32. 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) + + +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) + + +@register_torch_op(override=True) +def full_like(context, node): + _fill_like(context, node, _ct_ops.full_like) + + +@register_torch_op(override=True) +def fill(context, node): + _fill_like(context, node, _ct_ops.fill) + + +# coremltools converts empty, new_empty and empty_like as zeros, new_zeros and +# zeros_like, so these overrides take their names too. +@register_torch_op(torch_alias=["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)) + + +# coremltools writes y[:, i] = v to y[:, i:i + 1] without first wrapping a +# negative i, so for i = -1 it writes to an empty slice. No upstream issue +# covers this yet. For example, y[:, torch.tensor([-1])] = v fails to lower +# with "The update tensor should have shape [3, 0]. Got (3, 1)", and from +# iOS 18 Core ML skips the write when i is an input. A TorchScript node, which +# passes its indices as one list, keeps the coremltools converter. +@register_torch_op(override=True) +def index_put(context, node): + names = node.inputs[1] + if not isinstance(names, tuple) or None not in names: + return _ct_ops.index_put(context, node) + x = context[node.inputs[0]] + wrapped = list(names) + for dim, name in enumerate(names): + if name is None or not types.is_int(context[name].dtype): + continue + index = context[name] + size = x.shape[dim] + if is_symbolic(size): + size = _utils.pymil_value_at(mb.shape(x=x), dim) + index = mb.select( + cond=mb.less(x=index, y=0), + a=mb.add(x=index, y=size), + b=index, + name=f"{node.name}_index_{dim}", + ) + context.add(index) + wrapped[dim] = index.name + node = copy.copy(node) + node.inputs = (node.inputs[0], tuple(wrapped), *node.inputs[2:]) + return _ct_ops.index_put(context, node) + + def is_fbcode(): return not hasattr(_torch.version, "git_version") @@ -58,6 +214,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 +228,19 @@ 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 = copy.copy(node) + node.inputs = inputs + _rms_norm(context, node) + else: # TODO: remove this case when fbcode updates to coremltools 9.0 @register_torch_op( @@ -238,3 +408,162 @@ def dequantize_codebook(context, node): except ValueError as e: _raise_with_executorch_hint(e) context.add(output, node.name) + + +# Torch computes these ops in float for an int or bool input, while coremltools +# passes such an input on to Core ML ops that take only float, except for an int +# input of log and log1p. https://github.com/apple/coremltools/pull/2683 and +# https://github.com/apple/coremltools/pull/2735 cast an int input of +# reciprocal, log2 and log10 after coremltools 9.0; no upstream change covers +# the other ops or a bool input yet. For example, lowering x.long().sqrt() fails +# with "expects tensor or scalar of dtype from type domain ['fp16', 'fp32'] but +# got tensor[2,3,int32]". +def _with_float_input(convert): + @functools.wraps(convert) + def converter(context, node): + x = context[node.inputs[0]] + if types.is_int(x.dtype) or types.is_bool(x.dtype): + context.add(mb.cast(x=x, dtype="fp32", name=f"{node.name}_input_fp32")) + node = copy.copy(node) + node.inputs = [f"{node.name}_input_fp32", *node.inputs[1:]] + convert(context, node) + + return converter + + +for _convert in ( + _ct_ops.acos, + _ct_ops.asin, + _ct_ops.atan, + _ct_ops.atanh, + _ct_ops.cos, + _ct_ops.cosh, + _ct_ops.erf, + _ct_ops.exp, + _ct_ops.expm1, + _ct_ops.log, + _ct_ops.log1p, + _ct_ops.log10, + _ct_ops.log2, + _ct_ops.reciprocal, + _ct_ops.rsqrt, + _ct_ops.sigmoid, + _ct_ops.sin, + _ct_ops.sinh, + _ct_ops.sqrt, + _ct_ops.tan, + _ct_ops.tanh, +): + register_torch_op(_with_float_input(_convert), override=True) + + +# For an int input, torch returns floor, ceil and round of it unchanged, and +# computes relu and relu6 in int, while coremltools passes the input on to Core +# ML ops that take only float. No upstream issue covers this yet. For example, +# lowering x.long().floor() fails with "expects tensor or scalar of dtype from +# type domain ['fp16', 'fp32'] but got tensor[2,3,int32]". +def _with_int_input(convert, convert_int): + @functools.wraps(convert) + def converter(context, node): + x = context[node.inputs[0]] + if types.is_int(x.dtype): + context.add(convert_int(x=x, name=node.name)) + else: + convert(context, node) + + return converter + + +for _convert, _convert_int in ( + (_ct_ops.ceil, mb.identity), + (_ct_ops.floor, mb.identity), + (_ct_ops.relu, lambda x, name: mb.maximum(x=x, y=0, name=name)), + (_ct_ops.relu6, lambda x, name: mb.minimum(x=mb.maximum(x=x, y=0), y=6, name=name)), +): + register_torch_op(_with_int_input(_convert, _convert_int), override=True) +# coremltools registers round as _round with the alias round. +register_torch_op( + _with_int_input(_ct_ops._round, mb.identity), torch_alias=["round"], override=True +) + + +# Torch clips the input of hardtanh at its bounds in the dtype of the input, +# while coremltools passes the bounds as they are to clip, which takes only +# float and bounds of the dtype of the input. So nn.ReLU6 on an int tensor, and +# F.relu6(x, inplace=True), which passes int bounds, fail to lower with "the +# named input `alpha` must have the same data type as the named input `x`". +# https://github.com/apple/coremltools/pull/2894 casts the bounds to the dtype of +# a float input; no upstream change covers an int input yet. The dtype comes from +# torch, since coremltools gives some tensors another one, such as int32 to +# torch.arange(4, dtype=torch.float32). A TorchScript node, which carries no +# dtype, keeps the coremltools converter. +@register_torch_op(override=True) +def hardtanh(context, node): + if _torch_dtype(node) is None: + return _ct_ops.hardtanh(context, node) + x, low, high = _get_inputs(context, node, expected=3) + dtype = types.numpy_type_to_builtin_type(_core_ml_type(node)) + if x.dtype != dtype: + x = mb.cast(x=x, dtype=types.builtin_to_string(dtype)) + if types.is_int(x.dtype): + # Torch drops the fraction of each bound for an int input, as int() does. + clipped = mb.maximum(x=x, y=int(low.val)) + context.add(mb.minimum(x=clipped, y=int(high.val), name=node.name)) + else: + dtype = types.nptype_from_builtin(x.dtype) + context.add( + mb.clip(x=x, alpha=dtype(low.val), beta=dtype(high.val), name=node.name) + ) + + +# Torch computes clamp_min and clamp_max in the dtype it promotes the tensor and +# the bound to, while coremltools requires the two to share a dtype. So +# x.clamp_min(0) of a float x, and a float bound on an int tensor, fail to lower +# with an AssertionError. No upstream issue covers this yet. A TorchScript node, +# which carries no dtype, keeps the coremltools converter. +def _with_inputs_in_torch_dtype(op, convert): + @functools.wraps(convert) + def converter(context, node): + if _torch_dtype(node) is None: + return convert(context, node) + dtype = types.numpy_type_to_builtin_type(_core_ml_type(node)) + x, bound = _get_inputs(context, node, expected=2) + if x.dtype != dtype: + x = mb.cast(x=x, dtype=types.builtin_to_string(dtype)) + if bound.dtype != dtype: + bound = mb.cast(x=bound, dtype=types.builtin_to_string(dtype)) + context.add(op(x=x, y=bound, name=node.name)) + + return converter + + +for _op, _convert in ( + (mb.maximum, _ct_ops.clamp_min), + (mb.minimum, _ct_ops.clamp_max), +): + register_torch_op(_with_inputs_in_torch_dtype(_op, _convert), override=True) + + +# Torch copies src into the dtype of self, while coremltools' copy keeps the +# dtype of src. Lowering turns y[:, i] = v into a copy too, so v kept its dtype +# in y. No upstream issue covers this yet. For example, lowering +# F.gelu(x.clone().copy_(x.long())) fails with "expects tensor or scalar of +# dtype from type domain ['fp16', 'fp32'] but got tensor[2,3,int32]". The dtype +# comes from torch where the node carries it, since coremltools gives some +# tensors another one, such as int32 to torch.arange(4, dtype=torch.float32). +def _copy_into_dtype_of_self(context, node): + x, src = context[node.inputs[0]], context[node.inputs[1]] + dtype = x.dtype + if _torch_dtype(node) is not None: + dtype = types.numpy_type_to_builtin_type(_core_ml_type(node)) + if src.dtype != dtype: + name = f"{node.name}_src" + context.add(mb.cast(x=src, dtype=types.builtin_to_string(dtype), name=name)) + node = copy.copy(node) + node.inputs = [node.inputs[0], name, *node.inputs[2:]] + _ct_ops.copy(context, node) + + +register_torch_op( + functools.wraps(_ct_ops.copy)(_copy_into_dtype_of_self), override=True +) diff --git a/backends/apple/coreml/test/test_torch_ops.py b/backends/apple/coreml/test/test_torch_ops.py index 10c3f01a585..39d64a3f01e 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.scope import ScopeSource 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() @@ -93,6 +102,222 @@ def test_dequantize_affine_b4w_linear(self): et_prog = delegated_program.to_executorch() self._compare_outputs(et_prog, model, example_inputs) + def test_int_input_of_float_ops(self): + """ + Torch computes these ops in float for an int or bool input, and + coremltools passed the input on, so lowering failed. + """ + + class Model(torch.nn.Module): + def forward(self, x): + i = x.to(torch.int64) + j = i + 1 + k = (i > 1).to(torch.int64) + return torch.stack( + [ + i.sqrt(), + i.exp(), + i.expm1(), + i.sigmoid(), + i.sin(), + i.cos(), + i.tan(), + i.sinh(), + i.cosh(), + i.tanh(), + i.atan(), + i.erf(), + j.rsqrt(), + j.reciprocal(), + j.log2(), + j.log10(), + k.asin(), + k.acos(), + (i - i).atanh(), + (i > 1).sqrt(), + (i >= 0).log(), + (i > 1).log1p(), + torch.full((2, 3), 4).sqrt(), + ] + ) + + # Small integers, so that each op reads the same ints in Core ML. + model, example_inputs = Model(), (torch.randint(0, 4, (2, 3)).float(),) + ep = torch.export.export(model, example_inputs) + delegated_program = executorch.exir.to_edge_transform_and_lower( + ep, + partitioner=[self._coreml_partitioner()], + ) + self.assertEqual( + [ + node.target.__name__ + for node in delegated_program.exported_program().graph.nodes + if node.op == "call_function" + ], + ["executorch_call_delegate", "getitem"], + ) + self._compare_outputs(delegated_program.to_executorch(), model, example_inputs) + + def test_int_input_of_ops_that_keep_it_int(self): + """ + Torch keeps an int input of floor, ceil, round, relu, relu6 and + hardtanh int, and coremltools passed it on to Core ML ops that take only + float, so lowering failed. Int bounds of hardtanh failed on a float16 + input too. + """ + + class Model(torch.nn.Module): + def forward(self, x): + # -4, 0, 4 and 8, so that relu, relu6 and hardtanh clip on both + # sides. + i = x.to(torch.int64) * 4 - 4 + ops = [i.floor(), i.ceil(), i.round(), i.relu()] + ops += [ + torch.nn.functional.relu6(i), + # Torch drops the fraction of each bound, so this clips at + # -1 and 2. + torch.nn.functional.hardtanh(i, -1.5, 2.7), + # Int bounds on a float16 input. + torch.nn.functional.hardtanh(x.half() * 4 - 4, 0, 6).long(), + torch.full((2, 3), 4).floor(), + ] + return torch.stack(ops).float() + + model, example_inputs = Model(), ( + torch.tensor([[0.0, 1.0, 2.0], [3.0, 2.0, 1.0]]), + ) + delegated_program = executorch.exir.to_edge_transform_and_lower( + torch.export.export(model, example_inputs), + partitioner=[self._coreml_partitioner()], + ) + self.assertEqual( + [ + node.target.__name__ + for node in delegated_program.exported_program().graph.nodes + if node.op == "call_function" + ], + ["executorch_call_delegate", "getitem"], + ) + self._compare_outputs(delegated_program.to_executorch(), model, example_inputs) + + def test_clamp_min_max_with_bound_of_other_dtype(self): + """ + Torch computes clamp_min and clamp_max in the dtype it promotes the + tensor and the bound to, while coremltools required the two to share a + dtype. So lowering an int bound on a float tensor failed, and so did a + float bound on an int tensor. + """ + + class Model(torch.nn.Module): + def forward(self, x): + i = (x * 2).to(torch.int64) + ops = [x.clamp_min(1), x.clamp_max(1)] + ops += [i.clamp_min(0.5), i.clamp_max(2.5)] + ops += [x.clamp_min(i[:1]), i.clamp_max(x[:1] * 2)] + ops.append(torch.full((2, 3), 0).clamp_min(0.5)) + return torch.stack(ops) + + # Values that float16 holds exactly, so that Core ML drops the same + # fractions as eager. + model, example_inputs = Model(), ((torch.rand(2, 3) * 2).half().float(),) + delegated_program = executorch.exir.to_edge_transform_and_lower( + torch.export.export(model, example_inputs), + partitioner=[self._coreml_partitioner()], + ) + self.assertEqual( + [ + node.target.__name__ + for node in delegated_program.exported_program().graph.nodes + if node.op == "call_function" + ], + ["executorch_call_delegate", "getitem"], + ) + self._compare_outputs(delegated_program.to_executorch(), model, example_inputs) + + def test_hardtanh_of_float_arange(self): + """ + coremltools gives torch.arange(6, dtype=torch.float32) the dtype int32, so + its hardtanh failed to lower. Torch clips it in float, at bounds with a + fraction. A TorchScript node carries no dtype and keeps the coremltools + converter. + """ + + class Model(torch.nn.Module): + def forward(self, x): + a = torch.arange(6, dtype=torch.float32).reshape(2, 3) + return torch.nn.functional.hardtanh(a, 0.5, 2.5) + x + + model, example_inputs = Model(), (torch.zeros(2, 3),) + ep = torch.export.export(model, example_inputs) + # With no passes, Core ML keeps the clip instead of folding it. + program = ct.convert( + ep.run_decompositions({}), + convert_to="milinternal", + pass_pipeline=ct.PassPipeline.EMPTY, + ) + self.assertEqual( + [(op.alpha.val, op.beta.val) for op in program.find_ops(op_type="clip")], + [(0.5, 2.5)], + ) + ct.convert( + torch.jit.trace(torch.nn.Hardtanh(0.5, 2.5), example_inputs), + inputs=[ct.TensorType(shape=(2, 3))], + convert_to="milinternal", + ) + delegated_program = executorch.exir.to_edge_transform_and_lower( + ep, partitioner=[self._coreml_partitioner()] + ) + self.assertEqual( + [ + node.target.__name__ + for node in delegated_program.exported_program().graph.nodes + if node.op == "call_function" + ], + ["executorch_call_delegate", "getitem"], + ) + self._compare_outputs(delegated_program.to_executorch(), model, example_inputs) + + def test_copy_into_other_dtype(self): + """ + Torch copies a value into a tensor of another dtype in the dtype of + the tensor, also for y[:, i] = v, and coremltools kept the dtype of the + value, so lowering failed. + """ + + class Model(torch.nn.Module): + def forward(self, x): + y = x * 2 + y[:, 0] = (x[:, 1] * 4).to(torch.int64) + y[:, 1] = torch.full((2,), 7) + y[:, 2] = x[:, 1] > 0.5 + z = x.clone().copy_((x * 4).to(torch.int64)) + # Torch drops the fraction of a float written into an int. + i = (x * 4).to(torch.int64) + i[:, 0] = x[:, 1] * 4 + # coremltools makes this arange int32, so the copy has to take + # the dtype from torch. + a = torch.arange(6, dtype=torch.float32).reshape(2, 3).copy_(x) + w = torch.full((2, 3), 4) + w[:, 0] = x[:, 0] * 4 + return y + torch.nn.functional.gelu(z) + i + w + a.softmax(-1) + + # Values that float16 holds exactly, so that the comparison with 0.5 in + # Core ML agrees with eager. + model, example_inputs = Model(), (torch.rand(2, 3).half().float(),) + delegated_program = executorch.exir.to_edge_transform_and_lower( + torch.export.export(model, example_inputs), + partitioner=[self._coreml_partitioner()], + ) + self.assertEqual( + [ + node.target.__name__ + for node in delegated_program.exported_program().graph.nodes + if node.op == "call_function" + ], + ["executorch_call_delegate", "getitem"], + ) + self._compare_outputs(delegated_program.to_executorch(), model, example_inputs) + def test_dequantize_affine_c4w_embedding(self): model, example_inputs = self._get_test_model() quantize_( @@ -296,14 +521,30 @@ 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) + mlmodel = CoreMLBackend._convert_to_mlmodel( + executorch.exir.to_edge(ep).exported_program(), + CoreMLBackend.generate_compile_specs(), + ) + description = mlmodel.get_spec().description + outputs = [output.name for output in description.output] + self.assertEqual(len(set(outputs)), len(outputs)) + self.assertFalse(set(outputs) & {i.name for i in description.input}) delegated_program = executorch.exir.to_edge_transform_and_lower( ep, partitioner=[self._coreml_partitioner()], @@ -317,6 +558,337 @@ def forward(self, x): et_prog = delegated_program.to_executorch() self._compare_outputs(et_prog, model, example_inputs) + def test_same_shape_expand_output(self): + """ + coremltools keeps a same-shape expand as a tile when it makes an + output, and Core ML on the CPU then hands back the value it expands + wrong, as for the scale_tril that MultivariateNormal returns next to + its expand. An expand that changes the shape stays a tile. + """ + + class Model(torch.nn.Module): + def __init__(self, size, expand_first): + super().__init__() + self.size = size + self.expand_first = expand_first + + def forward(self, x): + y = x * 2 + outputs = (y, y.expand(self.size, -1, -1)) + return outputs[::-1] if self.expand_first else outputs + + # Without the fix, Core ML returns y wrong on the CPU in float32. + compile_specs = CoreMLBackend.generate_compile_specs( + compute_unit=ct.ComputeUnit.CPU_ONLY, + compute_precision=ct.precision.FLOAT32, + ) + example_inputs = (torch.randn(1, 3, 8),) + for size, expand_first, op_types in ( + (1, False, ["mul", "identity"]), + (1, True, ["identity", "mul"]), + (2, False, ["mul", "tile"]), + ): + with self.subTest(size=size, expand_first=expand_first): + model = Model(size, expand_first) + ep = torch.export.export(model, example_inputs) + edge_program = executorch.exir.to_edge(ep).exported_program() + mlmodel = CoreMLBackend._convert_to_mlmodel(edge_program, compile_specs) + outputs = mlmodel._mil_program.functions["main"].outputs + self.assertEqual([output.op.op_type for output in outputs], op_types) + # Debug info attributes each output to the node that makes it. + self.assertEqual( + [ + output.op.scopes[ScopeSource.EXIR_DEBUG_HANDLE] + for output in outputs + ], + [ + [node.meta["debug_handle"]] + for node in edge_program.graph.output_node().args[0] + ], + ) + delegated_program = executorch.exir.to_edge_transform_and_lower( + ep, partitioner=[CoreMLPartitioner(compile_specs=compile_specs)] + ) + self._compare_outputs( + delegated_program.to_executorch(), model, example_inputs + ) + # The pass removes the tile it replaces, also when the pass list has + # no dead_code_elimination to do it. + mlmodel = CoreMLBackend._convert_to_mlmodel( + executorch.exir.to_edge( + torch.export.export(Model(1, False), example_inputs) + ).exported_program(), + CoreMLBackend.generate_compile_specs( + compute_precision=ct.precision.FLOAT32, + pass_names=["common::const_elimination"], + ), + ) + self.assertNotIn( + "tile", + [op.op_type for op in mlmodel._mil_program.functions["main"].operations], + ) + + 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_fills_keep_their_dtype(self): + """ + coremltools makes full without a dtype, zeros and new_zeros as float, + ones_like as float from iOS 16, zeros_like with the dtype of its input, + and full_like, new_full, fill and scalar_tensor with the type of their + fill value, instead of the dtype torch gives them. nn.MultiheadAttention + turns a bool mask into a float one with zeros_like(mask, dtype=...). + """ + + class Fills(torch.nn.Module): + def forward(self, ids, half): + 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.full_like(ids, 0.0) + count = torch.scalar_tensor(1001.0, dtype=torch.long) + ints = (ids + ones + twos + zeros + torch.full((2, 3), 1)) * count + flags = torch.where(ids > 0, torch.full((2, 3), True), ids > 1) + return ( + ints, + flags, + half * torch.ones_like(half) + torch.zeros_like(half), + ) + + class Attention(torch.nn.MultiheadAttention): + def forward(self, x, mask): + return super().forward(x, x, x, attn_mask=mask, need_weights=False)[0] + + fills_inputs = (torch.arange(6).reshape(2, 3), torch.rand(2, 3).half()) + compile_specs = CoreMLBackend.generate_compile_specs( + compute_precision=ct.precision.FLOAT32, + minimum_deployment_target=ct.target.iOS16, + ) + delegated_program = executorch.exir.to_edge_transform_and_lower( + torch.export.export(Fills(), fills_inputs), + partitioner=[CoreMLPartitioner(compile_specs=compile_specs)], + ) + self.assertEqual( + { + node.target.__name__ + for node in delegated_program.exported_program().graph.nodes + if node.op == "call_function" + }, + {"executorch_call_delegate", "getitem"}, + ) + # The runtime casts each Core ML output to the dtype ExecuTorch expects, + # so the dtypes are read from the program the delegate converts. + lowered = delegated_program.exported_program().graph_module.lowered_module_0 + program = ct.convert( + lowered.original_module, + convert_to="milinternal", + compute_precision=ct.precision.FLOAT32, + minimum_deployment_target=ct.target.iOS16, + ) + self.assertEqual( + [output.dtype for output in program.functions["main"].outputs], + [types.int32, types.bool, types.fp16], + ) + self._compare_outputs(delegated_program.to_executorch(), Fills(), fills_inputs) + model = Attention(16, 2, batch_first=True).eval() + example_inputs = ( + torch.randn(1, 6, 16), + torch.ones(6, 6, dtype=torch.bool).triu(1), + ) + attention = torch.export.export(model, example_inputs) + program = ct.convert(attention.run_decompositions({}), convert_to="milinternal") + masks = [ + op.outputs[0].dtype + for op in program.functions["main"].operations + if op.name.startswith("zeros_like") + ] + self.assertTrue(masks and all(types.is_float(dtype) for dtype in masks)) + delegated_program = executorch.exir.to_edge_transform_and_lower( + attention, partitioner=[CoreMLPartitioner()] + ) + self.assertEqual( + [ + node.target.__name__ + for node in delegated_program.exported_program().graph.nodes + if node.op == "call_function" + ], + ["executorch_call_delegate", "getitem"], + ) + self._compare_outputs(delegated_program.to_executorch(), model, example_inputs) + + def test_empty_keeps_its_dtype(self): + """ + coremltools converts empty, new_empty and empty_like as zeros, new_zeros + and zeros_like, so they need the dtype torch gives them too. With dim + order ops skipped, as the Core ML export example does, empty reaches + Core ML as well. + """ + + class Empties(torch.nn.Module): + def forward(self, ids, x): + return ( + torch.empty(2, 3, dtype=torch.long), + ids.new_empty((2, 3)), + torch.empty_like(x, dtype=torch.long), + ids + 1, + ) + + inputs = (torch.arange(6).reshape(2, 3), torch.rand(2, 3)) + compile_specs = CoreMLBackend.generate_compile_specs( + compute_precision=ct.precision.FLOAT32 + ) + delegated_program = executorch.exir.to_edge_transform_and_lower( + torch.export.export(Empties(), inputs), + partitioner=[CoreMLPartitioner(compile_specs=compile_specs)], + compile_config=executorch.exir.EdgeCompileConfig(_skip_dim_order=True), + ) + lowered = delegated_program.exported_program().graph_module.lowered_module_0 + program = ct.convert( + lowered.original_module, + convert_to="milinternal", + compute_precision=ct.precision.FLOAT32, + ) + self.assertEqual( + [output.dtype for output in program.functions["main"].outputs], + [types.int32] * 4, + ) + + @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, and without elementwise_affine the weight too. + 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 + model = torch.nn.RMSNorm(65536, elementwise_affine=False).eval() + ep = torch.export.export(model, (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], + ) + compile_specs = CoreMLBackend.generate_compile_specs( + compute_precision=ct.precision.FLOAT32 + ) + 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_negative_index_next_to_slice(self): + """ + coremltools writes y[:, i] = v to y[:, i:i + 1] without wrapping a + negative i, so for i = -1 the slice is empty: lowering failed for a + constant i, and from iOS 18 Core ML skipped the write for an input i. + """ + + class Write(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("last", torch.tensor([-1])) + + def forward(self, x, v, i, rows, cols): + # Two copies, so that a skipped write shows in the output. + y = x.clone() + y[:, self.last] = v + z = x.clone() + z[:, i] = v * 2 + z[rows, cols] = v[0] + return y, z + + class Functional(torch.nn.Module): + def forward(self, x, v, i): + return torch.ops.aten.index_put(x, [None, i], v) + + model = Write() + example_inputs = ( + torch.rand(3, 4), + torch.rand(3, 1), + torch.tensor([-1]), + torch.tensor([0, 1]), + torch.tensor([-1, 2]), + ) + program = ct.convert( + torch.export.export(model, example_inputs).run_decompositions({}), + convert_to="milinternal", + minimum_deployment_target=ct.target.iOS16, + ) + # The constant is wrapped while converting, and the write without a + # slice is left to coremltools, so only the input next to a slice + # needs a select at run time. + self.assertEqual(len(program.find_ops(op_type="select")), 1) + # TorchScript passes the indices as one list, which keeps the + # coremltools converter. + ct.convert( + torch.jit.trace(Functional(), example_inputs[:3]), + inputs=[ + ct.TensorType(shape=(3, 4)), + ct.TensorType(shape=(3, 1)), + ct.TensorType(shape=(1,), dtype=types.int32), + ], + convert_to="milinternal", + ) + width = torch.export.Dim("width", min=2, max=8) + for dynamic_shapes in (None, ({1: width}, None, None, None, None)): + with self.subTest(dynamic=dynamic_shapes is not None): + ep = torch.export.export( + model, example_inputs, dynamic_shapes=dynamic_shapes + ) + delegated_program = executorch.exir.to_edge_transform_and_lower( + ep, partitioner=[self._coreml_partitioner()] + ) + self.assertEqual( + [ + node.target.__name__ + for node in delegated_program.exported_program().graph.nodes + if node.op == "call_function" + ], + ["executorch_call_delegate", "getitem", "getitem"], + ) + et_program = delegated_program.to_executorch() + self._compare_outputs(et_program, model, example_inputs) + if dynamic_shapes is not None: + # Another width catches a wrap by the exported width. + wide_inputs = (torch.rand(3, 6), *example_inputs[1:]) + self._compare_outputs(et_program, model, wide_inputs) + def test_dequantize_affine_below_ios18_raises_with_hint(self): """ Regression test for https://github.com/pytorch/executorch/issues/13122.