feat: add ExecuTorch runtime

.github/workflows/tests.yml

@@ -18,6 +18,6 @@ jobs:
         with:
           enable-cache: true
       - name: Install
-        run: uv sync --python 3.12 --extra onnx-cpu --extra openvino --group dev
+        run: uv sync --python 3.12 --extra onnx-cpu --extra openvino --extra executorch --group dev
       - name: Test
         run: uv run pytest -q

README.md

@@ -52,11 +52,12 @@ Other hosts install the extra for each runtime they need:
 | Extra | Adds | Host |
 |---|---|---|
 | `openvino` | `OpenVINORuntime` | Any CPU host. |
+| `executorch` | `ExecuTorchRuntime` | Any CPU host. |
 
 The `nvidia` extra now includes `openvino`, so `benchmark_all` also runs the OpenVINO CPU rows. These rows take most of the run time. To run only the TensorRT and ONNX Runtime rows, add `--runtimes=tensorrt,onnxruntime`.
 
 ```bash
-uv sync --python 3.12 --extra openvino
+uv sync --python 3.12 --extra openvino --extra executorch
 ```
 
 ### Options
@@ -84,6 +85,7 @@ Each runtime times the smallest call that runs the full graph. Preprocessing and
 | ONNX Runtime (GPU) | `run_with_iobinding` | Not included. IOBinding binds GPU buffers. |
 | ONNX Runtime (CPU) | `run_with_iobinding` | None (CPU). |
 | OpenVINO | `InferRequest.infer()` | None (CPU). |
+| ExecuTorch | `Method.execute()` | None (CPU). |
 
 OpenVINO rows read the existing `.onnx` files and compile them on the host, as TensorRT rows do. They set `INFERENCE_PRECISION_HINT` from the precision of the row. On a CPU with no native fp16, the fp16 rows are skipped, because OpenVINO compiles them in f32.
 

pyproject.toml

@@ -31,6 +31,7 @@ tensorrt = [
 onnx-cpu = ["onnxruntime>=1.22.0"]
 onnx-gpu = ["onnxruntime-gpu>=1.22.0"]
 openvino = ["openvino>=2025.0"]
+executorch = ["executorch>=1.0"]
 # One extra for each host.
 nvidia = ["single-artifact-benchmarking[tensorrt,onnx-gpu,openvino]"]
 
@@ -58,6 +59,7 @@ markers = [
     "tensorrt: needs TensorRT and an NVIDIA GPU",
     "onnxruntime: needs onnxruntime",
     "openvino: needs openvino",
+    "executorch: needs executorch",
 ]
 
 # PyPI holds only stub source packages for TensorRT 10; the wheels are on NVIDIA's index.

sab/runtimes/__init__.py

@@ -1,11 +1,13 @@
 from sab.runtimes.base import DEVICES, PRECISIONS, InputSpec, Runtime, RuntimeFactory, runtime_class
+from sab.runtimes.executorch import ExecuTorchRuntime
 from sab.runtimes.onnxruntime import ONNXRuntime
 from sab.runtimes.openvino import OpenVINORuntime
 from sab.runtimes.tensorrt import TRTRuntime
 
 __all__ = [
     "DEVICES",
     "PRECISIONS",
+    "ExecuTorchRuntime",
     "InputSpec",
     "ONNXRuntime",
     "OpenVINORuntime",

sab/runtimes/executorch.py

@@ -0,0 +1,112 @@
+import importlib.metadata
+from typing import ClassVar
+
+import torch
+
+from sab.profiler import CPUProfiler
+from sab.runtimes.base import InputSpec, Runtime
+
+# ExecuTorch ScalarType value -> torch dtype
+_DTYPES = {
+    0: torch.uint8,
+    1: torch.int8,
+    2: torch.int16,
+    3: torch.int32,
+    4: torch.int64,
+    5: torch.float16,
+    6: torch.float32,
+    7: torch.float64,
+    11: torch.bool,
+}
+_FLOAT_DTYPES = (torch.float16, torch.float32, torch.float64)
+
+WARMUP_ITERATIONS = 10
+
+
+def _input_name(index: int) -> str:
+    return f"input{index}"
+
+
+class ExecuTorchRuntime(Runtime):
+    """Runs a `.pte` artifact. The timed call is `Method.execute`.
+
+    A `.pte` has no input names, so inputs are named by position: `input0`, `input1`, and so on.
+    Outputs are named `output0`, `output1`, and so on. The artifact fixes the backend when it is
+    exported. The CPU artifacts use XNNPACK.
+    """
+
+    name: ClassVar[str] = "executorch"
+    devices: ClassVar[frozenset[str]] = frozenset({"cpu"})
+
+    def __init__(self, artifact_path: str, device: str, precision: str, *, method_name: str = "forward"):
+        super().__init__(artifact_path, device, precision)
+        from executorch.runtime import Runtime as ExecuTorch
+
+        self.input_device = "cpu"
+        self.profiler = CPUProfiler()
+        program = ExecuTorch.get().load_program(artifact_path)
+        self.method = program.load_method(method_name)
+
+        metadata = self.method.metadata
+        input_infos = [metadata.input_tensor_meta(i) for i in range(metadata.num_inputs())]
+        self._input_shapes = [tuple(info.sizes()) for info in input_infos]
+        self._input_dtypes = [_DTYPES[info.dtype()] for info in input_infos]
+
+        self._image_index = self._find_image_input()
+        self.input_spec = InputSpec(name=_input_name(self._image_index), shape=self._input_shapes[self._image_index])
+
+        self.warmup()
+
+    @classmethod
+    def is_available(cls, device: str) -> bool:
+        if device not in cls.devices:
+            return False
+        try:
+            import executorch.runtime  # noqa: F401
+        except ImportError:
+            return False
+        return True
+
+    @classmethod
+    def version(cls) -> str:
+        return f"executorch {importlib.metadata.version('executorch')}"
+
+    def run(self, inputs: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
+        arguments = self._to_artifact_inputs(inputs)
+
+        with self.profiler.profile():
+            raw_outputs = self.method.execute(arguments)
+
+        # ExecuTorch can reuse its output buffers on the next call, so each caller gets its own copy.
+        return {f"output{i}": output.detach().clone() for i, output in enumerate(raw_outputs)}
+
+    def warmup(self, num_iterations: int = WARMUP_ITERATIONS):
+        """Run zeros through the method to warm the caches before real measurements begin."""
+        zeros = [
+            torch.zeros(shape, dtype=dtype)
+            for shape, dtype in zip(self._input_shapes, self._input_dtypes)
+        ]
+        for _ in range(num_iterations):
+            self.method.execute(zeros)
+        self.profiler.reset()
+
+    def _find_image_input(self) -> int:
+        """Input 0, unless input 0 is not a 4-D float and exactly one other input is."""
+        image_like = [
+            index
+            for index, (shape, dtype) in enumerate(zip(self._input_shapes, self._input_dtypes))
+            if len(shape) == 4 and dtype in _FLOAT_DTYPES
+        ]
+        if 0 in image_like or len(image_like) != 1:
+            return 0
+        return image_like[0]
+
+    def _to_artifact_inputs(self, inputs: dict[str, torch.Tensor]) -> list[torch.Tensor]:
+        expected_names = {_input_name(index) for index in range(len(self._input_dtypes))}
+        if set(inputs) != expected_names:
+            raise ValueError(f"Inputs {sorted(inputs)} do not match the artifact inputs {sorted(expected_names)}")
+
+        return [
+            inputs[_input_name(index)].detach().to("cpu", dtype).contiguous()
+            for index, dtype in enumerate(self._input_dtypes)
+        ]

tests/test_executorch.py

@@ -0,0 +1,126 @@
+import pytest
+import torch
+
+pytest.importorskip("executorch.runtime")
+
+from executorch.backends.xnnpack.partition.xnnpack_partitioner import XnnpackPartitioner  # noqa: E402
+from executorch.exir import to_edge_transform_and_lower  # noqa: E402
+
+from sab.runtimes.executorch import ExecuTorchRuntime  # noqa: E402
+
+pytestmark = pytest.mark.executorch
+
+
+class ConvModel(torch.nn.Module):
+    """One NCHW image input."""
+
+    def __init__(self):
+        super().__init__()
+        self.conv = torch.nn.Conv2d(3, 4, 3, padding=1)
+
+    def forward(self, images):
+        return self.conv(images)
+
+
+class ImageAndSizesModel(ConvModel):
+    """Image first, then an int64 input, like the `orig_target_sizes` input of some detectors."""
+
+    def forward(self, images, sizes):
+        return self.conv(images), sizes + 1
+
+
+class SizesAndImageModel(ConvModel):
+    def forward(self, sizes, images):
+        return sizes + 1, self.conv(images)
+
+
+CASES = {
+    "conv": (ConvModel, (torch.rand(1, 3, 8, 8),)),
+    "image_and_sizes": (ImageAndSizesModel, (torch.rand(1, 3, 8, 8), torch.tensor([[10, 20]]))),
+    "sizes_and_image": (SizesAndImageModel, (torch.tensor([[10, 20]]), torch.rand(1, 3, 8, 8))),
+}
+
+
+@pytest.fixture(scope="session")
+def model_dir(tmp_path_factory):
+    return tmp_path_factory.mktemp("pte")
+
+
+@pytest.fixture(scope="session")
+def build_pte(model_dir):
+    """Export each tiny model once per session. Returns (module, example inputs, path)."""
+    built = {}
+
+    def build(case: str):
+        if case not in built:
+            torch.manual_seed(0)
+            model_class, example_inputs = CASES[case]
+            module = model_class().eval()
+            exported = torch.export.export(module, example_inputs)
+            program = to_edge_transform_and_lower(exported, partitioner=[XnnpackPartitioner()]).to_executorch()
+            path = model_dir / f"{case}.pte"
+            path.write_bytes(program.buffer)
+            built[case] = (module, example_inputs, str(path))
+        return built[case]
+
+    return build
+
+
+def test_run_matches_the_torch_model_and_outputs_own_their_memory(build_pte):
+    module, _, path = build_pte("conv")
+    runtime = ExecuTorchRuntime(path, "cpu", "fp32")
+    images = torch.rand(1, 3, 8, 8)
+
+    first = runtime.run({"input0": images})
+    runtime.run({"input0": torch.rand(1, 3, 8, 8)})
+
+    assert runtime.input_spec.name == "input0"
+    assert set(first) == {"output0"}
+    torch.testing.assert_close(first["output0"], module(images), atol=1e-4, rtol=1e-4)
+
+
+def test_run_passes_inputs_in_positional_order(build_pte):
+    module, _, path = build_pte("image_and_sizes")
+    runtime = ExecuTorchRuntime(path, "cpu", "fp32")
+    images = torch.rand(1, 3, 8, 8)
+    sizes = torch.tensor([[10, 20]])
+
+    # The dict order is the reverse of the positional order.
+    outputs = runtime.run({"input1": sizes, "input0": images})
+
+    assert runtime.input_spec.name == "input0"
+    torch.testing.assert_close(outputs["output0"], module(images, sizes)[0], atol=1e-4, rtol=1e-4)
+    assert outputs["output1"].tolist() == [[11.0, 21.0]]
+
+
+def test_run_finds_the_image_at_a_later_position(build_pte):
+    module, _, path = build_pte("sizes_and_image")
+    runtime = ExecuTorchRuntime(path, "cpu", "fp32")
+    images = torch.rand(1, 3, 8, 8)
+    sizes = torch.tensor([[10, 20]])
+
+    outputs = runtime.run({"input0": sizes, "input1": images})
+
+    assert runtime.input_spec.name == "input1"
+    assert runtime.input_spec.shape == (1, 3, 8, 8)
+    assert outputs["output0"].tolist() == [[11.0, 21.0]]
+    torch.testing.assert_close(outputs["output1"], module(sizes, images)[1], atol=1e-4, rtol=1e-4)
+
+
+def test_run_rejects_input_names_that_are_not_positional(build_pte):
+    *_, path = build_pte("conv")
+    runtime = ExecuTorchRuntime(path, "cpu", "fp32")
+
+    with pytest.raises(ValueError, match=r"\['images'\].*\['input0'\]"):
+        runtime.run({"images": torch.rand(1, 3, 8, 8)})
+    assert runtime.profiler.timings == []
+
+
+def test_run_records_one_timing_per_call_after_warmup(build_pte):
+    *_, path = build_pte("conv")
+    runtime = ExecuTorchRuntime(path, "cpu", "fp32")
+    assert runtime.profiler.timings == []
+
+    for expected in (1, 2, 3):
+        runtime.run({"input0": torch.rand(1, 3, 8, 8)})
+        assert len(runtime.profiler.timings) == expected