diff --git a/examples/models/llama/export_llama_lib.py b/examples/models/llama/export_llama_lib.py index c22bdac4ed2..5cdc9a40940 100644 --- a/examples/models/llama/export_llama_lib.py +++ b/examples/models/llama/export_llama_lib.py @@ -1805,9 +1805,7 @@ def _export_llama(llm_config: LlmConfig) -> LLMEdgeManager: # noqa: C901 verbose=llm_config.debug.verbose, ) elif llm_config.backend.coreml.enabled and not ( - llm_config.backend.vulkan.enabled - or llm_config.backend.mps.enabled - or llm_config.backend.qnn.enabled + llm_config.backend.vulkan.enabled or llm_config.backend.qnn.enabled ): builder = _to_edge_and_lower_llama_coreml( builder_exported, diff --git a/examples/models/llama/tests/test_export_llama_lib.py b/examples/models/llama/tests/test_export_llama_lib.py index 40ba701f84e..a803a59db52 100644 --- a/examples/models/llama/tests/test_export_llama_lib.py +++ b/examples/models/llama/tests/test_export_llama_lib.py @@ -6,6 +6,7 @@ # LICENSE file in the root directory of this source tree. import unittest +from unittest.mock import patch from executorch.devtools.backend_debug import get_delegation_info @@ -23,6 +24,7 @@ TOSAQuantizer = None VgfQuantizer = None +from executorch.examples.models.llama import export_llama_lib from executorch.examples.models.llama.export_llama_lib import ( _export_llama, build_args_parser, @@ -41,6 +43,52 @@ class ExportLlamaLibTest(unittest.TestCase): + def _assert_routes_to(self, lowering, coreml=False, vulkan=False, qnn=False): + """Assert which lowering an export routes to, without running one.""" + llm_config = LlmConfig() + llm_config.backend.coreml.enabled = coreml + llm_config.backend.vulkan.enabled = vulkan + llm_config.backend.qnn.enabled = qnn + # _validate_args rejects dynamic shapes when Core ML or QNN is enabled. + llm_config.model.enable_dynamic_shape = False + + class Reached(Exception): + pass + + with patch.object(export_llama_lib, lowering, side_effect=Reached) as target: + with self.assertRaises(Reached): + _export_llama(llm_config) + target.assert_called_once() + return target + + def test_core_ml_alone_routes_to_the_core_ml_lowering(self): + """Core ML on its own must reach the Core ML lowering. + + The guard read a backend config field that no longer exists, so this raised AttributeError + before reaching any lowering. + """ + self._assert_routes_to("_to_edge_and_lower_llama_coreml", coreml=True) + + def test_core_ml_with_qnn_routes_to_the_combined_lowering(self): + """Core ML with QNN must keep the QNN partitioner, so it takes the combined lowering.""" + target = self._assert_routes_to( + "_to_edge_and_lower_llama", coreml=True, qnn=True + ) + self.assertTrue(target.call_args.kwargs["coreml"]) + self.assertTrue(target.call_args.kwargs["qnn"]) + + def test_core_ml_with_vulkan_routes_to_the_combined_lowering(self): + """Core ML with Vulkan must keep the Vulkan partitioner the same way. + + This is what pins the Vulkan half of the exclusion clause: the Core ML lowering takes no + Vulkan argument, so routing there would drop the partitioner silently. + """ + target = self._assert_routes_to( + "_to_edge_and_lower_llama", coreml=True, vulkan=True + ) + self.assertTrue(target.call_args.kwargs["coreml"]) + self.assertTrue(target.call_args.kwargs["vulkan"]) + def test_has_expected_ops_and_op_counts(self): """ Checks the presence of unwanted expensive ops.