diff --git a/backends/apple/coreml/partition/coreml_partitioner.py b/backends/apple/coreml/partition/coreml_partitioner.py index 4453818f956..a7e59c914eb 100644 --- a/backends/apple/coreml/partition/coreml_partitioner.py +++ b/backends/apple/coreml/partition/coreml_partitioner.py @@ -26,6 +26,7 @@ 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.experimental.symbolic_shapes import statically_known_true, sym_eq from torch.fx.passes.infra.partitioner import CapabilityBasedPartitioner from torch.fx.passes.operator_support import OperatorSupportBase @@ -162,18 +163,141 @@ def _is_view_op(op: torch._ops.OpOverload) -> bool: return (alias_info is not None) and (not alias_info.is_write) +_MIN_MAX_TARGETS = ( + torch.ops.aten.min.default, + torch.ops.aten.min.dim, + torch.ops.aten.min.other, + torch.ops.aten.max.default, + torch.ops.aten.max.dim, + torch.ops.aten.max.other, + torch.ops.aten.amin.default, + torch.ops.aten.amax.default, + torch.ops.aten.minimum.default, + torch.ops.aten.maximum.default, + exir_ops.edge.aten.min.default, + exir_ops.edge.aten.min.dim, + exir_ops.edge.aten.min.other, + exir_ops.edge.aten.max.default, + exir_ops.edge.aten.max.dim, + exir_ops.edge.aten.max.other, + exir_ops.edge.aten.amin.default, + exir_ops.edge.aten.amax.default, + exir_ops.edge.aten.minimum.default, + exir_ops.edge.aten.maximum.default, +) + + +# No upstream issue yet. For example, converting x.amax() of a bool x raises +# 'Op "amax" (op_type: reduce_max) Input x="x" expects tensor or scalar of +# dtype from type domain ['fp16', 'fp32', 'int32'] but got tensor[3,2,bool]'. +def _is_min_max_of_bool(node: torch.fx.Node) -> bool: + """coremltools 9.0 cannot compute a min or max of bool tensors. + + It raises ``ValueError`` when every input is bool, so that the result is + bool, such as for the all-true check that torch.distributions exports, + instead of reporting the op as unsupported. A bool next to a float or int + tensor works, and so do int tensors. + """ + if node.target not in _MIN_MAX_TARGETS: + return False + val = node.meta.get("val") + if isinstance(val, (list, tuple)): + val = val[0] + return isinstance(val, torch.Tensor) and val.dtype == torch.bool + + +_INDEX_PUT_TARGETS = ( + torch.ops.aten.index_put.default, + exir_ops.edge.aten.index_put.default, +) + + +# https://github.com/apple/coremltools/issues/2141 covers a bool index, which +# coremltools 9.0 converts only on its own. No upstream issue covers the other +# forms yet. For example, y[:, idx] = v with idx = torch.tensor([0, 2]) raises +# "Input squeeze_0 has rank 1 != other inputs rank 0". +def _index_put_refusal( + node: torch.fx.Node, minimum_deployment_target: Optional[ct.target] +) -> Optional[str]: + """The form of an ``index_put`` that coremltools 9.0 cannot convert, such + as "with accumulate=True next to a slice", or None for any other node. + + With a whole slice, such as ``y[:, idx] = v``, it writes ``v`` into the + part of ``y`` the indices select, keeping each indexed dim as size 1. That + needs ``accumulate=False``, int indices of shape ``()`` or ``(1,)``, a + static size in every other dim, and a non-bool ``y`` below iOS 18. For a + size known only at run time, it raises "Cannot add const". With a 0-d + index or several indices, torch lays out that part differently, so ``v`` + may only vary in the dims after the last index. Without a whole slice, + several indices need the same shape and no bool mask. It raises or writes + the wrong place for any other form, such as ``idx`` with two entries, + instead of reporting the op as unsupported. + """ + if node.target not in _INDEX_PUT_TARGETS: + return None + y, values = node.args[0].meta.get("val"), node.args[2].meta.get("val") + indices = [i if i is None else i.meta.get("val") for i in node.args[1]] + given = {dim: index for dim, index in enumerate(indices) if index is not None} + if not given or not all( + isinstance(t, torch.Tensor) for t in (y, values, *given.values()) + ): + return None + if len(given) == len(indices): + if len(indices) > 1 and any( + index.dtype == torch.bool + or not statically_known_true(sym_eq(index.shape, indices[0].shape)) + for index in indices + ): + return "with indices that differ in shape or include a bool mask" + return None + accumulate = ( + node.args[3] if len(node.args) > 3 else node.kwargs.get("accumulate", False) + ) + if accumulate: + return "with accumulate=True next to a slice" + if any( + index.dtype == torch.bool + or index.dim() > 1 + or not statically_known_true(index.numel() == 1) + for index in given.values() + ): + return ( + "with an index other than an int tensor of shape () or (1,) next to a slice" + ) + selected = [1 if dim in given else size for dim, size in enumerate(y.shape)] + if any(isinstance(size, torch.SymInt) for size in selected): + return "with a size known only at run time next to a slice" + reorders = len(given) > 1 or any(index.dim() == 0 for index in given.values()) + after = len(selected) - 1 - max(given) + if reorders and not all( + statically_known_true(v == 1) + for v in values.shape[: max(0, values.dim() - after)] + ): + return ( + "with a 0-d index or several indices, and values that vary before " + "the last index" + ) + if y.dtype == torch.bool and ( + minimum_deployment_target is None or minimum_deployment_target < ct.target.iOS18 + ): + return "on a bool tensor with a minimum_deployment_target below iOS 18" + return None + + class _OperatorsSupportedForCoreMLBackend(OperatorSupportBase): def __init__( self, skip_ops_for_coreml_delegation: Optional[List[str]] = None, lower_full_graph: bool = False, log: bool = False, + minimum_deployment_target: Optional[ct.target] = None, ) -> None: if skip_ops_for_coreml_delegation is None: skip_ops_for_coreml_delegation = [] super().__init__() self.skip_ops_for_coreml_delegation = skip_ops_for_coreml_delegation self.lower_full_graph = lower_full_graph + self.minimum_deployment_target = minimum_deployment_target self._logged_msgs = set() self._log = log @@ -195,7 +319,9 @@ def should_skip_op_for_delegation(self, node_target_name: str) -> bool: return True return False - def should_override_support(self, node) -> bool: + def override_reason(self, node) -> Optional[str]: + """Why CoreML cannot take node, even when coremltools converts its op.""" + name = getattr(node.target, "__name__", node.target) # https://github.com/apple/coremltools/issues/2573 if ( node.target @@ -208,20 +334,14 @@ def should_override_support(self, node) -> bool: and "alpha" in node.kwargs and node.kwargs["alpha"] != 1 ): - self.log_once( - "torch.ops.aten.{sub, add}.Tensor with alpha != 1 is not supported by CoreML. Overriding support." - ) - return True + return "torch.ops.aten.{sub, add}.Tensor with alpha != 1 is not supported by CoreML." # Ops that are unsupported by CoreML purely on the basis of their # target — no per-arg conditions to check. Grouped by upstream issue # so the comment trail still points at the underlying coremltools / # executorch bug for each entry. if node.target in _UNSUPPORTED_OP_TARGETS: - self.log_once( - f"{node.target} is not supported by CoreML. Overriding op support." - ) - return True + return f"{name} is not supported by CoreML." if _touches_unsupported_high_rank_tensor( node, @@ -229,25 +349,23 @@ def should_override_support(self, node) -> bool: getattr(n.target, "__name__", "").lower() ), ): - self.log_once( - f"{node.target} on a tensor above rank 5 is not supported by " - "CoreML. Overriding op support." - ) - return True + return f"{name} on a tensor above rank 5 is not supported by CoreML." if _is_arg_min_max_over_flattened_input(node): - self.log_once( + return ( "torch.ops.aten.{argmax, argmin}.default with dim=None is " - "not supported by CoreML. Overriding op support." + "not supported by CoreML." ) - return True if _reads_or_produces_bfloat16(node): - self.log_once( - f"{node.target} on a bfloat16 tensor is not supported by CoreML. " - "Overriding op support." - ) - return True + return f"{name} on a bfloat16 tensor is not supported by CoreML." + + if _is_min_max_of_bool(node): + return f"{name} on bool tensors is not supported by CoreML." + + refusal = _index_put_refusal(node, self.minimum_deployment_target) + if refusal is not None: + return f"{name} {refusal} is not supported by CoreML." # TODO: enable this after bugs in ExecuTorch's partitioner are fixed # # If lower_full_graph=False, do not partition nodes with symbolic args because it can result in symbolic args @@ -266,7 +384,7 @@ def should_override_support(self, node) -> bool: # ) # return True - return False + return None def is_node_supported(self, submodules, node: torch.fx.Node) -> bool: # get_attr node can always be supported on any backend @@ -284,10 +402,17 @@ def is_node_supported(self, submodules, node: torch.fx.Node) -> bool: is_supported = ct.converters.mil.frontend.torch.is_torch_fx_node_supported( node ) - if self.should_override_support(node): + reason = self.override_reason(node) + if reason is not None: + self.log_once(f"{reason} Overriding op support.") is_supported = False if not is_supported: + if self.lower_full_graph and reason is not None: + raise NotImplementedError( + f"{reason} Set lower_full_graph=False in the " + "CoreMLPartitioner to run it outside of CoreML." + ) if self.lower_full_graph: raise NotImplementedError( f"""CoreML does not support the op {node_target_name}, but you have set lower_full_graph=True in the CoreMLPartitioner. @@ -343,6 +468,11 @@ def __init__( self.take_over_mutable_buffer = take_over_mutable_buffer self.lower_full_graph = lower_full_graph self.take_over_constant_data = take_over_constant_data + self._minimum_deployment_target = ( + CoreMLBackend.min_deployment_target_from_compile_specs( + self.delegation_spec.compile_specs + ) + ) self._logged_msgs = set() if self.lower_full_graph: @@ -399,6 +529,7 @@ def partition(self, exported_program: ExportedProgram) -> PartitionResult: self.skip_ops_for_coreml_delegation, self.lower_full_graph, log=True, + minimum_deployment_target=self._minimum_deployment_target, ), allows_single_node_partition=True, ) @@ -438,6 +569,7 @@ def ops_to_not_decompose( self.skip_ops_for_coreml_delegation, self.lower_full_graph, log=False, + minimum_deployment_target=self._minimum_deployment_target, ) # CoreML prevents certain ops (like triu) from lowering to CoreML when put in the ExecuTorch op namespace diff --git a/backends/apple/coreml/test/test_coreml_partitioner.py b/backends/apple/coreml/test/test_coreml_partitioner.py index 0129eb1d33f..64f1308786d 100644 --- a/backends/apple/coreml/test/test_coreml_partitioner.py +++ b/backends/apple/coreml/test/test_coreml_partitioner.py @@ -17,6 +17,9 @@ 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_VALID_TEST_RUNTIME, +) from executorch.exir.backend.utils import format_delegated_graph @@ -521,6 +524,170 @@ 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_min_max_of_bool_is_skipped(self): + """ + coremltools raises on a min or max of bool tensors, such as the + all-true check torch.distributions exports. The partitioner must leave + those ops to the portable backend and still delegate the rest, + including a min or max with a float input. + """ + + class Model(torch.nn.Module): + def forward(self, x): + ok = (x.tril() == x).view(1, -1) + flags = ok.min(-1)[0].to(x.dtype) + ok.amax(-1).to(x.dtype) + both = torch.maximum(ok, ok.flip(-1)).view(x.shape).to(x.dtype) + floats = x.amax(-1, keepdim=True) + torch.maximum(x > 0.5, x) + return x * 2 + flags + both + floats + + model = Model().eval() + x = torch.rand(14, 14) + ep = torch.export.export(model, (x,), strict=True) + edge = executorch.exir.to_edge_transform_and_lower( + ep, partitioner=[CoreMLPartitioner()] + ) + op_names = [ + n.target.__name__ + for n in edge.exported_program().graph.nodes + if n.op == "call_function" + ] + self.assertIn("executorch_call_delegate", op_names) + self.assertIn("aten.min.dim", op_names) + self.assertEqual(op_names.count("aten.amax.default"), 1) + self.assertEqual(op_names.count("aten.maximum.default"), 1) + if _TEST_RUNTIME: + program = Runtime.get().load_program(edge.to_executorch().buffer) + out = program.load_method("forward").execute((x,))[0] + self.assertTrue(torch.allclose(out, model(x), atol=1e-2)) + + def test_index_put_forms_coremltools_gets_wrong(self): + """ + coremltools fails some forms of ``index_put``, such as + ``y[:, idx] = v`` with two entries in ``idx``, or writes them to the + wrong place, so the partitioner refuses those, and they run on the + portable backend. With lower_full_graph=True, lowering fails and says + why. + """ + + class Scatter(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("cols", torch.tensor([0, 2])) + self.register_buffer("col", torch.tensor([1])) + self.register_buffer("last", torch.tensor(3)) + self.register_buffer("rows", torch.tensor([[1], [2]])) + self.register_buffer("flag", torch.tensor([2])) + self.register_buffer("rows_mask", torch.tensor([True, True, True])) + self.register_buffer("trio", torch.tensor([0, 2, 3])) + self.register_buffer("grid", torch.tensor([[1]])) + self.register_buffer("yes", torch.tensor([True])) + + def forward(self, x, a, b, c): + y = x * 2 + y[:, self.cols] = a + y[:, self.grid] = a[:, :, None] + y[self.rows, self.cols] = b + y[self.rows_mask, self.trio] = c + torch.ops.aten.index_put_(y, [None, self.col], a, True) + torch.ops.aten.index_put_(y, [None, self.last], c) + mask = y > 0.5 + mask[:, self.flag] = a > 0.5 + z = a * 2 + z[:, self.yes] = c[:, None] + return y + mask + z + + class Scatter3d(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("last", torch.tensor(2)) + self.register_buffer("mid", torch.tensor(1)) + + def forward(self, x, v): + y = x.clone() + torch.ops.aten.index_put_(y, [None, None, self.last], v) + # Values that vary only after the index stay in Core ML. + torch.ops.aten.index_put_(y, [None, self.mid], v[0] * torch.arange(3.0)) + return y * 2 + + ios18 = CoreMLBackend.generate_compile_specs( + minimum_deployment_target=ct.target.iOS18 + ) + scatter_inputs = ( + torch.rand(3, 4), + torch.rand(3, 1), + torch.rand(2, 2), + torch.rand(3), + ) + scatter3d_inputs = (torch.rand(2, 2, 3), torch.rand(2, 1)) + cases = [ + (Scatter(), scatter_inputs, 8, None), + # From iOS 18 the bool write stays in Core ML. + (Scatter(), scatter_inputs, 7, ios18), + (Scatter3d(), scatter3d_inputs, 1, None), + ] + for model, inputs, cpu_calls, compile_specs in cases: + with self.subTest( + model=type(model).__name__, ios18=compile_specs is not None + ): + ep = torch.export.export(model.eval(), inputs, strict=True) + edge = executorch.exir.to_edge_transform_and_lower( + ep, partitioner=[CoreMLPartitioner(compile_specs=compile_specs)] + ) + op_names = [ + n.target.__name__ + for n in edge.exported_program().graph.nodes + if n.op == "call_function" + ] + self.assertIn("executorch_call_delegate", op_names) + self.assertEqual(op_names.count("aten.index_put.default"), cpu_calls) + # An iOS 18 program needs macOS 15 to run. + if IS_VALID_TEST_RUNTIME: + program = Runtime.get().load_program(edge.to_executorch().buffer) + out = program.load_method("forward").execute(inputs)[0] + self.assertTrue(torch.allclose(out, model(*inputs), atol=1e-2)) + + class BoolWrite(torch.nn.Module): + def __init__(self, index): + super().__init__() + self.register_buffer("index", index) + + def forward(self, x, a): + return torch.ops.aten.index_put(x > 0.5, [None, self.index], a > 0.5) + + for model, inputs, reason in ( + (Scatter3d(), scatter3d_inputs, "values that vary before the last index"), + ( + BoolWrite(torch.tensor([2])), + scatter_inputs[:2], + "minimum_deployment_target below iOS 18", + ), + # A bool write that a later target would not fix gives its own reason. + (BoolWrite(torch.tensor([0, 2])), scatter_inputs[:2], "other than an int"), + ): + with self.assertRaisesRegex(NotImplementedError, reason): + executorch.exir.to_edge_transform_and_lower( + torch.export.export(model, inputs, strict=True), + partitioner=[CoreMLPartitioner(lower_full_graph=True)], + ) + # With a batch size known only at run time, the write that stays in + # Core ML above goes to the CPU too. + batch = torch.export.Dim("batch", min=2, max=8) + ep = torch.export.export( + Scatter3d(), + scatter3d_inputs, + dynamic_shapes=({0: batch}, {0: batch}), + 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.count("aten.index_put.default"), 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_coreml_recipes.py b/backends/apple/coreml/test/test_coreml_recipes.py index 98d240d74b5..47b21207c2a 100644 --- a/backends/apple/coreml/test/test_coreml_recipes.py +++ b/backends/apple/coreml/test/test_coreml_recipes.py @@ -136,6 +136,7 @@ def test_int4_weight_only_per_channel(self): def test_int4_weight_only_per_group(self): """Test INT4 weight-only per-group quantization with different group sizes""" + torch.manual_seed(0) class CustomTwoLinearModel(nn.Module): def __init__(self): @@ -149,7 +150,7 @@ def forward(self, x): return x model = CustomTwoLinearModel().eval() - example_inputs = [(torch.randn(1, 32),)] + example_inputs = [(torch.randn(16, 32),)] # Test with different group sizes for group_size in [8, 16, 32]: with self.subTest(group_size=group_size): @@ -165,7 +166,7 @@ def forward(self, x): self.check_fully_delegated(session) self._compare_eager_quantized_model_outputs( - session, example_inputs, atol=1e-3 + session, example_inputs, atol=1e-2 ) self._compare_eager_unquantized_model_outputs( session, model, example_inputs, sqnr_threshold=15