Support for several models and more customizations

create_ai_layer.py

@@ -32,6 +32,7 @@
 from __future__ import annotations
 
 import os
+import inspect
 import re
 import shlex
 import subprocess
@@ -181,39 +182,102 @@ def model_params(mlops: dict) -> dict:
     return params
 
 
-def export_model(spec, params: dict) -> bytes:
-    """Quantize model/model.py, delegate it to the Ethos-U and return the .pte."""
+def export_models(spec, params: dict):
+    """Export every PTE described by model/model.py."""
     import torch
+    from torch import nn
     from executorch.backends.arm.ethosu import EthosUPartitioner
     from executorch.backends.arm.quantizer import EthosUQuantizer, get_symmetric_quantization_config
     from executorch.exir import EdgeCompileConfig, ExecutorchBackendConfig, to_edge_transform_and_lower
     from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e
 
     sys.path.insert(0, str(HERE / "model"))
-    from model import get_calibration_inputs, get_model
-
-    model = get_model(**{k: v for k, v in params.items() if k == "input_shape"})
-    samples = get_calibration_inputs(**params)
-    example = (samples[0],)
-    graph = torch.export.export(model, example).module()
-
-    # Quantize the whole graph so the partitioner can move every node into the
-    # Ethos-U delegate; only a float<->int8 boundary stays on the CPU.
-    quantizer = EthosUQuantizer(spec)
-    quantizer.set_global(get_symmetric_quantization_config(is_per_channel=True))
-    prepared = prepare_pt2e(graph, quantizer)
-    with torch.no_grad():
-        for sample in samples:
-            prepared(sample)  # calibrate
-    quantized = convert_pt2e(prepared)
-
-    edge = to_edge_transform_and_lower(
-        torch.export.export(quantized, example),
-        partitioner=[EthosUPartitioner(spec)],
-        compile_config=EdgeCompileConfig(_check_ir_validity=False),
+    import model as model_module
+    from description import Model, Pte
+
+    shape_arg = inspect.signature(model_module.get_model).parameters.get("input_shape")
+    shape_required = shape_arg is not None and shape_arg.default is inspect.Parameter.empty
+    source = (
+        model_module.get_model(input_shape=params["input_shape"])
+        if shape_required and "input_shape" in params
+        else model_module.get_model()
     )
-    program = edge.to_executorch(ExecutorchBackendConfig(extract_delegate_segments=False))
-    return bytes(program.buffer)
+    if isinstance(source, nn.Module):
+        if "input_shape" in params and not shape_required:
+            source = model_module.get_model(input_shape=params["input_shape"])
+        source = [Pte(name="model", models=[Model(name="forward", model=source,
+            calibration_inputs=model_module.get_calibration_inputs,
+            input_shape=params.get("input_shape", (1, 16, 16, 3)),
+            calibration_samples=params.get("calibration_samples", 2))])]
+    elif isinstance(source, Model):
+        source = [Pte(name="model", models=[source])]
+    elif isinstance(source, Pte):
+        source = [source]
+    elif isinstance(source, list) and source:
+        if all(isinstance(item, Model) for item in source):
+            source = [Pte(name="model", models=source)]
+        elif not all(isinstance(item, Pte) for item in source):
+            raise TypeError("get_model() must return a homogeneous list of Model or Pte objects")
+    else:
+        raise TypeError("get_model() must return an nn.Module, Model, Pte, or a nonempty list of Model or Pte")
+
+    results = []
+    for pte in source:
+        if not isinstance(pte, Pte) or not pte.models:
+            raise ValueError("each Pte must contain at least one Model")
+        programs = {}
+        for item in pte.models:
+            if not isinstance(item, Model) or item.model is None:
+                raise TypeError("each Pte.models entry must be a Model with an nn.Module")
+            if item.name in programs:
+                raise ValueError(f"duplicate model method name: {item.name}")
+            use_default_quantization = (
+                item.calibration_inputs is not None and item.custom_quantization is None
+            )
+            samples = []
+            if use_default_quantization:
+                if item.input_shape is None or item.calibration_samples is None:
+                    raise ValueError(
+                        f"{item.name}: input_shape and calibration_samples are required "
+                        "for default quantization"
+                    )
+                samples = item.calibration_inputs(item.input_shape, item.calibration_samples)
+                if not samples:
+                    raise ValueError(f"{item.name}: calibration_inputs returned no samples")
+            if item.example_inputs is not None:
+                example = item.example_inputs
+            elif samples:
+                example = (samples[0],)
+            elif item.input_shape is not None:
+                example = (torch.zeros(item.input_shape),)
+            else:
+                raise ValueError(
+                    f"{item.name}: example_inputs or input_shape is required to export the model"
+                )
+            graph = torch.export.export(item.model, example).module()
+            if item.custom_quantization is not None:
+                quantized = item.custom_quantization(spec, graph)
+            elif use_default_quantization:
+                quantizer = EthosUQuantizer(spec)
+                quantizer.set_global(get_symmetric_quantization_config(is_per_channel=True))
+                prepared = prepare_pt2e(graph, quantizer)
+                with torch.no_grad():
+                    for sample in samples:
+                        prepared(sample)
+                quantized = convert_pt2e(prepared)
+            else:
+                quantized = graph
+            programs[item.name] = torch.export.export(quantized, example)
+
+        edge = to_edge_transform_and_lower(
+            programs,
+            partitioner=[EthosUPartitioner(spec)],
+            constant_methods=pte.metadata.details if pte.metadata else None,
+            compile_config=EdgeCompileConfig(_check_ir_validity=False),
+        )
+        program = edge.to_executorch(ExecutorchBackendConfig(extract_delegate_segments=False))
+        results.append((pte.name, bytes(program.buffer)))
+    return results
 
 
 def components(pte: bytes, pack: Path) -> tuple[list[str], list[str]]:
@@ -244,13 +308,13 @@ def components(pte: bytes, pack: Path) -> tuple[list[str], list[str]]:
     return ["Runtime", "Kernel Utils", "Kernel Registration", "Backend EthosU"], sorted(selected)
 
 
-def c_array(pte: bytes) -> str:
+def c_array(pte: bytes, symbol: str = SYMBOL) -> str:
     rows = [", ".join(f"0x{b:02x}" for b in pte[i : i + 16]) for i in range(0, len(pte), 16)]
     return (
         "// Generated by create_ai_layer.py -- do not edit.\n"
-        f"__attribute__((aligned(16))) const unsigned char {SYMBOL}[] = {{\n  "
+        f"__attribute__((aligned(16))) const unsigned char {symbol}[] = {{\n  "
         + ",\n  ".join(rows)
-        + f"\n}};\nconst unsigned long {SYMBOL}_size = sizeof({SYMBOL});\n"
+        + f"\n}};\nconst unsigned long {symbol}_size = sizeof({symbol});\n"
     )
 
 
@@ -267,7 +331,12 @@ def c_array(pte: bytes) -> str:
 """
 
 
-def clayer(mlops: dict, runtime: list[str], operators: list[str], mlops_file: Path, version: str) -> str:
+def header(symbol: str) -> str:
+    return HEADER.replace(SYMBOL, symbol)
+
+
+def clayer(mlops: dict, runtime: list[str], operators: list[str], mlops_file: Path, version: str,
+           names: list[str]) -> str:
     lines = [
         f"# Generated by create_ai_layer.py from {mlops_file.name} -- do not edit.",
         f"# Re-run `python create_ai_layer.py {mlops_file.name}` after changing",
@@ -290,10 +359,16 @@ def clayer(mlops: dict, runtime: list[str], operators: list[str], mlops_file: Pa
         *[f"    - component: Machine Learning:ExecuTorch Operators:{c}" for c in operators],
         "",
         "  groups:",
-        f"    - group: {mlops['model'].get('name', 'Model')}",
-        "      files:",
-        f"        - file: ./{SYMBOL}.c",
-        f"        - file: ./{SYMBOL}.h",
+        *[
+            line
+            for name in names
+            for line in (
+                f"    - group: {mlops['model'].get('name', 'Model') if names == ['model'] else name}",
+                "      files:",
+                f"        - file: ./{name}_pte.c",
+                f"        - file: ./{name}_pte.h",
+            )
+        ],
         "",
     ]
     return "\n".join(lines)
@@ -311,17 +386,25 @@ def main() -> None:
     layer_file = mlops_file.parent / mlops["model"]["clayer"]
     layer_dir = layer_file.parent
 
-    pte = export_model(compile_spec(mlops, mlops_file.parent), model_params(mlops))
+    exports = export_models(compile_spec(mlops, mlops_file.parent), model_params(mlops))
+    names = [name for name, _ in exports]
+    if len(names) != len(set(names)) or any(not re.fullmatch(r"[A-Za-z_][A-Za-z_0-9]*", name) for name in names):
+        raise ValueError("Pte names must be unique C identifiers")
     version = executorch_version(mlops_file)
-    runtime, operators = components(pte, executorch_pack(version))
+    component_sets = [components(pte, executorch_pack(version)) for _, pte in exports]
+    runtime = sorted(set().union(*(set(item[0]) for item in component_sets)))
+    operators = sorted(set().union(*(set(item[1]) for item in component_sets)))
 
     layer_dir.mkdir(parents=True, exist_ok=True)
-    (layer_dir / "model.pte").write_bytes(pte)
-    (layer_dir / f"{SYMBOL}.c").write_text(c_array(pte), newline="\n")
-    (layer_dir / f"{SYMBOL}.h").write_text(HEADER, newline="\n")
-    layer_file.write_text(clayer(mlops, runtime, operators, mlops_file, version), newline="\n")
-
-    print(f"[ai_layer] {len(pte)} byte program, operators: {operators}")
+    for name, pte in exports:
+        symbol = f"{name}_pte"
+        (layer_dir / f"{name}.pte").write_bytes(pte)
+        (layer_dir / f"{symbol}.c").write_text(c_array(pte, symbol), newline="\n")
+        (layer_dir / f"{symbol}.h").write_text(header(symbol), newline="\n")
+        print(f"[ai_layer] {name}: {len(pte)} byte program")
+    layer_file.write_text(clayer(mlops, runtime, operators, mlops_file, version, names), newline="\n")
+
+    print(f"[ai_layer] operators: {operators}")
     print(f"[ai_layer] wrote {layer_file}")
 
 

model/description.py

@@ -0,0 +1,56 @@
+from __future__ import annotations
+
+import torch
+from torch import nn
+from dataclasses import dataclass
+from typing import Callable, TYPE_CHECKING, TypeAlias
+from torch.fx import GraphModule
+
+if TYPE_CHECKING:
+    from executorch.backends.arm.ethosu import EthosUCompileSpec
+
+MetadataValue: TypeAlias = (
+    bool
+    | int
+    | float
+    | str
+    | torch.Tensor
+    | torch.dtype
+    | torch.layout
+    | list["MetadataValue"]
+    | dict[str, "MetadataValue"]
+)
+
+# Description of a model
+@dataclass
+class Model:
+    # Identifier for the model (used as a unique name)
+    name: str
+    # None means the model is already quantized; otherwise provide calibration samples.
+    # Function has get_calibration_inputs(
+    # input_shape: tuple[int, ...] = INPUT_SHAPE, calibration_samples: int = 2
+    # ) -> list[torch.Tensor]
+    calibration_inputs: Callable[[tuple[int, ...], int], list[torch.Tensor]] | None = None
+    # If the model do not use the default Ethos-U quantization scheme 
+    # For instance may use : get_symmetric_a16w8_quantization_config
+    custom_quantization: Callable[[EthosUCompileSpec, GraphModule], GraphModule] | None = None
+    model: nn.Module | None = None
+    # Parameters for this model's default calibration. They are unused when
+    # calibration_inputs is None or custom_quantization is set.
+    input_shape: tuple[int, ...] | None = None
+    calibration_samples: int | None = None
+    # Required for an already quantized or custom quantized model when
+    # input_shape is None; torch.export still needs example inputs.
+    example_inputs: tuple[torch.Tensor, ...] | None = None
+
+# Description of metadata for a particular model or experiment
+@dataclass
+class Metadata:
+    details: dict[str, MetadataValue]
+
+# When a Pte file contains metadata or several models
+@dataclass
+class Pte:
+    name: str
+    models: list[Model]
+    metadata: Metadata | None = None

model/model.py

@@ -9,6 +9,7 @@
 
 import torch
 from torch import nn
+import description as de
 
 # NHWC, the layout of a camera frame, which the Ethos-U reads as it is. An NCHW
 # input costs one more NPU operation that reorders it before the first