Core ML: stop the gather opt-out from failing tied embeddings

backends/apple/coreml/compiler/coreml_preprocess.py

@@ -579,14 +579,52 @@ def _convert_to_mlmodel(
             logger.warning(
                 "Core ML Backend op_linear_quantizer_config API is experimental"
             )
+            # Gathers are left alone so that embedding tables are not quantized by this
+            # config, except where the table is also some other op's weight. A tied
+            # embedding is one constant feeding both a gather and a linear, and coremltools
+            # refuses to compress a constant its consumers disagree about, so opting the
+            # gather out there does not skip the table, it fails the whole lowering.
+            tied = CoreMLBackend.gathers_sharing_a_weight(mlmodel)
             config = cto.coreml.OptimizationConfig(
                 global_config=op_linear_quantizer_config,
                 op_type_configs={"gather": None},
+                op_name_configs={name: op_linear_quantizer_config for name in tied},
             )
             mlmodel = cto.coreml.linear_quantize_weights(mlmodel, config=config)
 
         return mlmodel
 
+    @staticmethod
+    def gathers_sharing_a_weight(mlmodel: ct.models.MLModel) -> List[str]:
+        """
+        Names of gather ops whose table is also consumed by an op of another type.
+
+        Such a constant reaches the compressor under two different configurations at once,
+        which it rejects rather than resolves. Naming these gathers lets them be configured
+        the same way as the op they share the constant with; every other gather keeps the
+        opt-out. Returns nothing for a model whose program is not available, which leaves
+        the behaviour as it was.
+        """
+        program = getattr(mlmodel, "_mil_program", None)
+        if program is None:
+            return []
+
+        consumers: Dict[str, List[Any]] = {}
+        for function in program.functions.values():
+            for op in function.operations:
+                for value in op.inputs.values():
+                    for var in value if isinstance(value, (list, tuple)) else [value]:
+                        source = getattr(var, "op", None)
+                        if source is not None and source.op_type == "const":
+                            consumers.setdefault(var.name, []).append(op)
+
+        shared: List[str] = []
+        for ops in consumers.values():
+            gathers = [op for op in ops if op.op_type == "gather"]
+            if gathers and len(gathers) < len(ops):
+                shared += [op.name for op in gathers]
+        return shared
+
     @staticmethod
     def preprocess_model(
         mlmodel: ct.models.MLModel, model_type: MODEL_TYPE

backends/apple/coreml/test/test_coreml_partitioner.py

@@ -433,6 +433,46 @@ def forward(self, x):
         self.assertIn("executorch_call_delegate", op_names)
         self.assertNotIn("aten.argmax.default", op_names)
 
+    def test_tied_embedding_is_quantized_with_its_linear(self):
+        """A weight that is both an embedding table and an output projection lowers.
+
+        The quantizer config opts gathers out so that embedding tables are not compressed
+        by it. When the table is tied to a linear's weight, that opt-out leaves one
+        constant with two configurations, which coremltools rejects — so the opt-out has
+        to stop at the tie rather than fail the whole lowering.
+        """
+
+        class Tied(torch.nn.Module):
+            def __init__(self, vocab=512, dim=64, tie=True):
+                super().__init__()
+                self.embedding = torch.nn.Embedding(vocab, dim)
+                self.output = torch.nn.Linear(dim, vocab, bias=False)
+                if tie:
+                    self.output.weight = self.embedding.weight
+
+            def forward(self, ids):
+                return self.output(self.embedding(ids))
+
+        ids = torch.zeros(1, 4, dtype=torch.long)
+        for tie in (True, False):
+            compile_specs = CoreMLBackend.generate_compile_specs(
+                minimum_deployment_target=ct.target.iOS18,
+                compute_precision=ct.precision(ct.precision.FLOAT16.value),
+                op_linear_quantizer_config={
+                    "mode": "linear_symmetric",
+                    "dtype": "int4",
+                    "granularity": "per_channel",
+                },
+            )
+            exported = torch.export.export(Tied(tie=tie).eval(), (ids,), strict=True)
+            delegated = executorch.exir.to_edge_transform_and_lower(
+                exported,
+                partitioner=[CoreMLPartitioner(compile_specs=compile_specs)],
+            )
+            # Reaching a program at all is the assertion: the failure this covers is
+            # raised during lowering, not carried in the result.
+            self.assertIsNotNone(delegated.to_executorch())
+
     def test_deprecation_warning_for_to_backend_workflow(self):
         """
         Test that the deprecated to_edge + to_backend workflow shows a deprecation warning.
@@ -532,5 +572,6 @@ 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_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()