diff --git a/create_ai_layer.py b/create_ai_layer.py index 5392de3b..ef2587ea 100644 --- a/create_ai_layer.py +++ b/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}") diff --git a/model/description.py b/model/description.py new file mode 100644 index 00000000..691be9ce --- /dev/null +++ b/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 diff --git a/model/model.py b/model/model.py index 4c63245d..31dd5921 100644 --- a/model/model.py +++ b/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