diff --git a/backends/apple/coreml/partition/coreml_partitioner.py b/backends/apple/coreml/partition/coreml_partitioner.py index b4e550d8479..6e9dab32da9 100644 --- a/backends/apple/coreml/partition/coreml_partitioner.py +++ b/backends/apple/coreml/partition/coreml_partitioner.py @@ -83,6 +83,97 @@ def _is_arg_min_max_over_flattened_input(node: torch.fx.Node) -> bool: return dim is None +_MM_TARGETS = ( + torch.ops.aten.mm.default, + exir_ops.edge.aten.mm.default, +) + + +def _constant_input_names(ep: Optional[ExportedProgram]) -> frozenset: + """Names of the placeholders that hold parameters, buffers and constants.""" + if ep is None: + return frozenset() + sig = ep.graph_signature + return frozenset( + [ + *sig.inputs_to_parameters, + *sig.inputs_to_buffers, + *sig.inputs_to_lifted_tensor_constants, + ] + ) + + +def _is_integer_tensor(node: torch.fx.Node) -> bool: + val = node.meta.get("val") + return ( + isinstance(val, torch.Tensor) + and not val.dtype.is_floating_point + and not val.dtype.is_complex + and val.dtype != torch.bool + ) + + +def _constant_nodes(graph: torch.fx.Graph, constant_names: frozenset) -> set: + """Nodes whose value only depends on parameters, buffers and constants.""" + constant = set() + for node in graph.nodes: + if node.op == "get_attr" or ( + node.op == "placeholder" and node.name in constant_names + ): + constant.add(node) + elif node.op == "call_function" and node.all_input_nodes: + if all(n in constant for n in node.all_input_nodes): + constant.add(node) + return constant + + +def _is_integer_mm_with_constant_weight(node: torch.fx.Node, constant: set) -> bool: + if node.op != "call_function" or node.target not in _MM_TARGETS: + return False + if len(node.args) < 2 or not isinstance(node.args[1], torch.fx.Node): + return False + weight = node.args[1] + return weight in constant and _is_integer_tensor(weight) + + +def _integer_mm_with_constant_weight_nodes( + graph: torch.fx.Graph, constant_names: frozenset +) -> set: + """``mm`` nodes coremltools 9.0 cannot load, plus their constant-only feeders. + + coremltools 9.0 lowers ``mm`` with a constant weight to ``linear``, which + only takes floating point inputs, so an integer ``mm`` against a constant + lowers fine and then fails to load. An integer ``mm`` of two runtime + inputs becomes ``matmul`` and runs. See pytorch/executorch#11691. + + A constant-only op that only feeds such weights (e.g. a transpose of a + buffer) falls back too: left delegated on its own it would give Core ML a + partition with no runtime inputs. + + TODO: remove once ExecuTorch consumes a coremltools release that includes + apple/coremltools#2719, which keeps integer ``mm`` on ``matmul``. + """ + constant = _constant_nodes(graph, constant_names) + rejected = { + node + for node in graph.nodes + if _is_integer_mm_with_constant_weight(node, constant) + } + if not rejected: + return rejected + + for node in reversed(graph.nodes): + if node.op != "call_function" or node not in constant or not node.users: + continue + if all( + user in rejected + and (user.target not in _MM_TARGETS or user.args[0] is not node) + for user in node.users + ): + rejected.add(node) + return rejected + + def _is_view_op(op: torch._ops.OpOverload) -> bool: schema = op._schema if len(schema.arguments) == 0: @@ -97,12 +188,15 @@ def __init__( skip_ops_for_coreml_delegation: Optional[List[str]] = None, lower_full_graph: bool = False, log: bool = False, + exported_program: Optional[ExportedProgram] = 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._constant_names = _constant_input_names(exported_program) + self._integer_mm_fallback_cache = {} self._logged_msgs = set() self._log = log @@ -111,6 +205,13 @@ def log_once(self, msg: str) -> None: logger.info(msg) self._logged_msgs.add(msg) + def _integer_mm_fallback_nodes(self, graph: torch.fx.Graph) -> set: + if graph not in self._integer_mm_fallback_cache: + self._integer_mm_fallback_cache[graph] = ( + _integer_mm_with_constant_weight_nodes(graph, self._constant_names) + ) + return self._integer_mm_fallback_cache[graph] + def should_skip_op_for_delegation(self, node_target_name: str) -> bool: skipped_ops = self.skip_ops_for_coreml_delegation or [] if node_target_name in skipped_ops: @@ -159,6 +260,13 @@ def should_override_support(self, node) -> bool: ) return True + if node in self._integer_mm_fallback_nodes(node.graph): + self.log_once( + "torch.ops.aten.mm.default with an integer constant weight is " + "not supported by CoreML. Overriding op support." + ) + 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 @@ -309,6 +417,7 @@ def partition(self, exported_program: ExportedProgram) -> PartitionResult: self.skip_ops_for_coreml_delegation, self.lower_full_graph, log=True, + exported_program=exported_program, ), allows_single_node_partition=True, ) @@ -348,6 +457,7 @@ def ops_to_not_decompose( self.skip_ops_for_coreml_delegation, self.lower_full_graph, log=False, + exported_program=ep, ) # 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 78dd114eeda..f32db45be07 100644 --- a/backends/apple/coreml/test/test_coreml_partitioner.py +++ b/backends/apple/coreml/test/test_coreml_partitioner.py @@ -18,6 +18,7 @@ from executorch.backends.apple.coreml.compiler import CoreMLBackend from executorch.backends.apple.coreml.partition import CoreMLPartitioner from executorch.exir.backend.utils import format_delegated_graph +from executorch.exir.lowered_backend_module import get_lowered_submodules @torch.library.custom_op("unsupported::linear", mutates_args=()) @@ -433,6 +434,107 @@ def forward(self, x): self.assertIn("executorch_call_delegate", op_names) self.assertNotIn("aten.argmax.default", op_names) + def test_integer_mm_with_constant_weight_falls_back_to_portable(self): + """ + Regression test for https://github.com/pytorch/executorch/issues/11691. + + coremltools 9.0 lowers an ``mm`` against a constant weight to + ``linear``, which only takes floating point inputs, so an int32 ``mm`` + with a constant weight lowered fine but failed to load. The + partitioner must leave it (and any constant-only op feeding its + weight) to the portable backend, while still delegating the rest of + the graph, integer ``mm`` of two runtime inputs, and fp32 / fp16 + ``mm`` with a constant weight. + """ + + def int_mat(): + return torch.randint(-5, 5, (8, 8), dtype=torch.int32) + + class ConstWeight(torch.nn.Module): + def __init__(self): + super().__init__() + self.w = int_mat() + + def forward(self, x): + return torch.mm(torch.abs(x), self.w) + + class TransposedBuffer(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("w", int_mat()) + + def forward(self, x): + return torch.mm(torch.abs(x), self.w.t()) + + class TwoInputs(torch.nn.Module): + def forward(self, x, y): + return torch.mm(x, y) + + class FloatWeight(torch.nn.Module): + def __init__(self, dtype): + super().__init__() + self.w = torch.nn.Parameter(torch.randn(8, 8, dtype=dtype)) + + def forward(self, x): + return torch.mm(x, self.w) + + cases = [ + (ConstWeight(), (int_mat(),), True), + (TransposedBuffer(), (int_mat(),), True), + (TwoInputs(), (int_mat(), int_mat()), False), + (FloatWeight(torch.float32), (torch.randn(8, 8),), False), + ( + FloatWeight(torch.float16), + (torch.randn(8, 8, dtype=torch.float16),), + False, + ), + ] + for model, inputs, falls_back in cases: + model = model.eval() + ep = torch.export.export(model, inputs, strict=True) + compile_specs = None + if inputs[0].dtype == torch.float16: + # Core ML needs iOS16+ for fp16 inputs and outputs. + compile_specs = CoreMLBackend.generate_compile_specs( + minimum_deployment_target=ct.target.iOS16 + ) + 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" + ] + delegated_op_names = [ + n.target.__name__ + for _, lowered, _ in get_lowered_submodules( + edge.exported_program().graph_module + ) + for n in lowered.original_module.graph.nodes + if n.op == "call_function" + ] + # Every case keeps something delegated: either the mm itself, or + # the abs in front of an mm that falls back. + self.assertIn("executorch_call_delegate", op_names, model) + self.assertEqual("aten.mm.default" in op_names, falls_back, model) + self.assertEqual( + "aten.mm.default" in delegated_op_names, not falls_back, model + ) + if _TEST_RUNTIME: + et_prog = edge.to_executorch() + program = Runtime.get().load_program(et_prog.buffer) + et_outputs = program.load_method("forward").execute(inputs)[0] + self.assertTrue( + torch.allclose( + et_outputs.float(), + model(*inputs).float(), + atol=1e-02, + rtol=1e-02, + ), + model, + ) + def test_tied_embedding_is_quantized_with_its_linear(self): """A weight that is both an embedding table and an output projection lowers. @@ -572,6 +674,7 @@ def forward(self, x): test_runner.test_take_over_constant_data_false() test_runner.test_aten_rand_default_falls_back_to_portable() test_runner.test_aten_randn_is_still_delegated() + test_runner.test_integer_mm_with_constant_weight_falls_back_to_portable() test_runner.test_tied_embedding_is_quantized_with_its_linear() test_runner.test_deprecation_warning_for_to_backend_workflow() test_runner.test_no_warning_for_to_edge_transform_and_lower_workflow()