Add Ethos-U to LLM extension (#17700)

Signed-off-by: Sebastian Larsson <sebastian.larsson@arm.com>

Author: Sebastian-Larsson

examples/models/llama/export_llama_lib.py

@@ -34,6 +34,7 @@
 from executorch.extension.llm.export.config.llm_config import LlmConfig
 from executorch.extension.llm.export.partitioner_lib import (
     get_coreml_partitioner,
+    get_ethosu_partitioner,
     get_mps_partitioner,
     get_openvino_partitioner,
     get_qnn_partitioner,
@@ -43,6 +44,7 @@
 )
 from executorch.extension.llm.export.quantizer_lib import (
     get_coreml_quantizer,
+    get_ethosu_quantizer,
     get_ov_quantizer,
     get_pt2e_quantization_params,
     get_pt2e_quantizers,
@@ -218,6 +220,7 @@ def build_args_parser() -> argparse.ArgumentParser:
             "coreml_baseline_8a_c4w",
             "vulkan_8w",
             "tosa_8a8w",
+            "ethosu_8a8w",
         ],
         help="Use PT2E quantization. Comma separated options. e.g. xnnpack_dynamic (for per channel 8 bit weight), xnnpack_dynamic_qc4 (for per channel 4 bit weight), embedding.",
     )
@@ -813,6 +816,14 @@ def get_quantizer_and_quant_params(llm_config):
             llm_config.backend.tosa.version, llm_config.quantization.pt2e_quantize.value
         )
         quantizers.append(tosa_quantizer)
+    if llm_config.backend.ethosu.enabled and llm_config.quantization.pt2e_quantize:
+        ethosu_quantizer = get_ethosu_quantizer(
+            llm_config.backend.ethosu.target,
+            llm_config.backend.ethosu.system_config,
+            llm_config.backend.ethosu.memory_mode,
+            llm_config.quantization.pt2e_quantize.value,
+        )
+        quantizers.append(ethosu_quantizer)
     if llm_config.backend.vulkan.enabled and llm_config.quantization.pt2e_quantize:
         assert (
             len(quantizers) == 0
@@ -984,20 +995,27 @@ def _to_edge_and_lower_llama_openvino(
     return builder.to_executorch(passes=additional_passes)
 
 
-def _to_edge_and_lower_llama_tosa(
+def _to_edge_and_lower_llama_arm(
     builder_exported,
     modelname,
     quantizers,
     additional_passes,
-    tosa_spec,
+    llm_config: LlmConfig,
     verbose: bool = False,
 ) -> LLMEdgeManager:
     logging.info("Lowering model using TOSA partitioner")
 
     partitioners = []
-    partitioners.append(get_tosa_partitioner(tosa_spec))
-
-    modelname = f"tosa_{modelname}"
+    if llm_config.backend.ethosu.enabled:
+        partitioners.append(
+            get_ethosu_partitioner(
+                llm_config.backend.ethosu.target,
+            )
+        )
+        modelname = f"ethosu_{modelname}"
+    elif llm_config.backend.tosa.enabled:
+        partitioners.append(get_tosa_partitioner(llm_config.backend.tosa.version))
+        modelname = f"tosa_{modelname}"
 
     builder = builder_exported.pt2e_quantize(quantizers).to_edge_transform_and_lower(
         partitioners
@@ -1365,13 +1383,13 @@ def _export_llama(llm_config: LlmConfig) -> LLMEdgeManager:  # noqa: C901
             openvino_device=llm_config.backend.openvino.device,
             verbose=llm_config.debug.verbose,
         )
-    elif llm_config.backend.tosa.enabled:
-        builder = _to_edge_and_lower_llama_tosa(
+    elif llm_config.backend.tosa.enabled or llm_config.backend.ethosu.enabled:
+        builder = _to_edge_and_lower_llama_arm(
             builder_exported,
             modelname,
             quantizers,
             additional_passes,
-            llm_config.backend.tosa.version,
+            llm_config,
             verbose=llm_config.debug.verbose,
         )
     else:

examples/models/llama/tests/test_export_llama_lib.py

@@ -10,12 +10,17 @@
 from executorch.devtools.backend_debug import get_delegation_info
 
 try:
-    from executorch.backends.arm.quantizer.arm_quantizer import TOSAQuantizer
+    from executorch.backends.arm.quantizer.arm_quantizer import (
+        EthosUQuantizer,
+        TOSAQuantizer,
+    )
 
     HAS_ARM_BACKEND = True
 except ImportError:
     HAS_ARM_BACKEND = False
+    EthosUQuantizer = None
     TOSAQuantizer = None
+
 from executorch.examples.models.llama.export_llama_lib import (
     _export_llama,
     build_args_parser,
@@ -73,3 +78,18 @@ def test_get_quantizer_and_quant_params_returns_tosa_quantizer(self):
         self.assertIsNone(quant_dtype)
         self.assertEqual(len(quantizers), 1)
         self.assertIsInstance(quantizers[0], TOSAQuantizer)
+
+    @unittest.skipUnless(HAS_ARM_BACKEND, "ARM backend not available")
+    def test_get_quantizer_and_quant_params_returns_ethosu_quantizer(self):
+        llm_config = LlmConfig()
+        llm_config.backend.ethosu.enabled = True
+        llm_config.quantization.pt2e_quantize = Pt2eQuantize.ethosu_8a8w
+
+        pt2e_quant_params, quantizers, quant_dtype = get_quantizer_and_quant_params(
+            llm_config
+        )
+
+        self.assertIsNone(pt2e_quant_params)
+        self.assertIsNone(quant_dtype)
+        self.assertEqual(len(quantizers), 1)
+        self.assertIsInstance(quantizers[0], EthosUQuantizer)

extension/llm/export/config/llm_config.py

@@ -348,6 +348,7 @@ class Pt2eQuantize(str, Enum):
     coreml_baseline_8a_c4w = "coreml_baseline_8a_c4w"
     vulkan_8w = "vulkan_8w"
     tosa_8a8w = "tosa_8a8w"
+    ethosu_8a8w = "ethosu_8a8w"
 
 
 class SpinQuant(str, Enum):
@@ -545,6 +546,18 @@ class TosaConfig:
     version: str = "TOSA-1.0+INT"
 
 
+@dataclass
+class EthosUConfig:
+    """
+    Configures the Ethos-U backend.
+    """
+
+    enabled: bool = False
+    target: str = "ethos-u85-128"  # Default target, can be overridden.
+    memory_mode: str = "default"
+    system_config: str = "default"
+
+
 @dataclass
 class BackendConfig:
     """
@@ -560,6 +573,7 @@ class BackendConfig:
     openvino: OpenvinoConfig = field(default_factory=OpenvinoConfig)
     torchao: TorchAOKernelsConfig = field(default_factory=TorchAOKernelsConfig)
     tosa: TosaConfig = field(default_factory=TosaConfig)
+    ethosu: EthosUConfig = field(default_factory=EthosUConfig)
 
 
 ################################################################################

extension/llm/export/partitioner_lib.py

@@ -246,3 +246,12 @@ def get_tosa_partitioner(version: str):
     compile_spec = TosaCompileSpec(version)
 
     return TOSAPartitioner(compile_spec)
+
+
+def get_ethosu_partitioner(target: str):
+    from executorch.backends.arm.ethosu.compile_spec import EthosUCompileSpec
+    from executorch.backends.arm.ethosu.partitioner import EthosUPartitioner
+
+    compile_spec = EthosUCompileSpec(target)
+
+    return EthosUPartitioner(compile_spec)

extension/llm/export/quantizer_lib.py

@@ -340,3 +340,24 @@ def get_tosa_quantizer(version: str, pt2e_quantize: str):
         raise ValueError(f"Unsupported quantizer specification {pt2e_quantize}")
 
     return quantizer
+
+
+def get_ethosu_quantizer(
+    target: str, system_config: str, memory_mode: str, pt2e_quantize: str
+):
+    from executorch.backends.arm.ethosu.compile_spec import EthosUCompileSpec
+    from executorch.backends.arm.quantizer.arm_quantizer import (
+        EthosUQuantizer,
+        get_symmetric_quantization_config,
+    )
+
+    compile_spec = EthosUCompileSpec(target, system_config, memory_mode)
+
+    quantizer = EthosUQuantizer(compile_spec)
+
+    if pt2e_quantize == "ethosu_8a8w":
+        quantizer.set_global(get_symmetric_quantization_config())
+    else:
+        raise ValueError(f"Unsupported quantizer specification {pt2e_quantize}")
+
+    return quantizer