Skip to content

Commit 82d725a

Browse files
leeclemnetclaude
andcommitted
feat: add ExecuTorch runtime
Add ExecuTorchRuntime (XNNPACK, cpu) on the Runtime contract. It loads a .pte program and times only Method.execute(). Inputs are cast to the program's dtypes before the timed call, and outputs are cloned after it, because ExecuTorch can reuse its output buffers. No script rows use the runtime yet: a .pte comes from torch.export on the PyTorch model, not from ONNX, and the exported artifacts are not in the bucket. The tests build tiny .pte programs with torch.export. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
1 parent 2fd7dbd commit 82d725a

7 files changed

Lines changed: 1040 additions & 3 deletions

File tree

‎.github/workflows/tests.yml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,6 @@ jobs:
1818
with:
1919
enable-cache: true
2020
- name: Install
21-
run: uv sync --python 3.12 --extra onnx-cpu --extra openvino --group dev
21+
run: uv sync --python 3.12 --extra onnx-cpu --extra openvino --extra executorch --group dev
2222
- name: Test
2323
run: uv run pytest -q

‎README.md‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -52,11 +52,12 @@ Other hosts install the extra for each runtime they need:
5252
| Extra | Adds | Host |
5353
|---|---|---|
5454
| `openvino` | `OpenVINORuntime` | Any CPU host. |
55+
| `executorch` | `ExecuTorchRuntime` | Any CPU host. |
5556

5657
The `nvidia` extra now includes `openvino`.
5758

5859
```bash
59-
uv sync --python 3.12 --extra openvino
60+
uv sync --python 3.12 --extra openvino --extra executorch
6061
```
6162

6263
### Options
@@ -84,6 +85,7 @@ Each runtime times the smallest call that runs the full graph. Preprocessing and
8485
| ONNX Runtime (GPU) | `run_with_iobinding` | Not included. IOBinding binds GPU buffers. |
8586
| ONNX Runtime (CPU) | `run_with_iobinding` | None (CPU). |
8687
| OpenVINO | `InferRequest.infer()` | None (CPU). |
88+
| ExecuTorch | `Method.execute()` | None (CPU). |
8789

8890
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.
8991

‎pyproject.toml‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@ tensorrt = [
3131
onnx-cpu = ["onnxruntime>=1.22.0"]
3232
onnx-gpu = ["onnxruntime-gpu>=1.22.0"]
3333
openvino = ["openvino>=2025.0"]
34+
executorch = ["executorch>=1.0"]
3435
# One extra for each host.
3536
nvidia = ["single-artifact-benchmarking[tensorrt,onnx-gpu,openvino]"]
3637

@@ -58,6 +59,7 @@ markers = [
5859
"tensorrt: needs TensorRT and an NVIDIA GPU",
5960
"onnxruntime: needs onnxruntime",
6061
"openvino: needs openvino",
62+
"executorch: needs executorch",
6163
]
6264

6365
# PyPI holds only stub source packages for TensorRT 10; the wheels are on NVIDIA's index.

‎sab/runtimes/__init__.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,13 @@
11
from sab.runtimes.base import DEVICES, PRECISIONS, InputSpec, Runtime, RuntimeFactory, runtime_class
2+
from sab.runtimes.executorch import ExecuTorchRuntime
23
from sab.runtimes.onnxruntime import ONNXRuntime
34
from sab.runtimes.openvino import OpenVINORuntime
45
from sab.runtimes.tensorrt import TRTRuntime
56

67
__all__ = [
78
"DEVICES",
89
"PRECISIONS",
10+
"ExecuTorchRuntime",
911
"InputSpec",
1012
"ONNXRuntime",
1113
"OpenVINORuntime",

‎sab/runtimes/executorch.py‎

Lines changed: 108 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,108 @@
1+
import importlib.metadata
2+
from typing import ClassVar
3+
4+
import torch
5+
6+
from sab.profiler import CPUProfiler
7+
from sab.runtimes.base import InputSpec, Runtime
8+
9+
# ExecuTorch ScalarType value -> torch dtype
10+
_DTYPES = {
11+
0: torch.uint8,
12+
1: torch.int8,
13+
2: torch.int16,
14+
3: torch.int32,
15+
4: torch.int64,
16+
5: torch.float16,
17+
6: torch.float32,
18+
7: torch.float64,
19+
11: torch.bool,
20+
}
21+
_FLOAT_DTYPES = (torch.float16, torch.float32, torch.float64)
22+
23+
WARMUP_ITERATIONS = 10
24+
25+
26+
def _input_name(index: int) -> str:
27+
return f"input{index}"
28+
29+
30+
class ExecuTorchRuntime(Runtime):
31+
"""Runs a `.pte` artifact. The timed call is `Method.execute`.
32+
33+
A `.pte` has no input names, so inputs are named by position: `input0`, `input1`, and so on.
34+
Outputs are named `output0`, `output1`, and so on. The artifact fixes the backend when it is
35+
exported. The CPU artifacts use XNNPACK.
36+
"""
37+
38+
name: ClassVar[str] = "executorch"
39+
devices: ClassVar[frozenset[str]] = frozenset({"cpu"})
40+
41+
def __init__(self, artifact_path: str, device: str, precision: str, *, method_name: str = "forward"):
42+
super().__init__(artifact_path, device, precision)
43+
from executorch.runtime import Runtime as ExecuTorch
44+
45+
self.input_device = "cpu"
46+
self.profiler = CPUProfiler()
47+
program = ExecuTorch.get().load_program(artifact_path)
48+
self.method = program.load_method(method_name)
49+
50+
metadata = self.method.metadata
51+
input_infos = [metadata.input_tensor_meta(i) for i in range(metadata.num_inputs())]
52+
self._input_shapes = [tuple(info.sizes()) for info in input_infos]
53+
self._input_dtypes = [_DTYPES[info.dtype()] for info in input_infos]
54+
55+
self._image_index = self._find_image_input()
56+
self.input_spec = InputSpec(name=_input_name(self._image_index), shape=self._input_shapes[self._image_index])
57+
58+
self.warmup()
59+
60+
@classmethod
61+
def is_available(cls, device: str) -> bool:
62+
if device not in cls.devices:
63+
return False
64+
try:
65+
import executorch.runtime # noqa: F401
66+
except ImportError:
67+
return False
68+
return True
69+
70+
@classmethod
71+
def version(cls) -> str:
72+
return f"executorch {importlib.metadata.version('executorch')}"
73+
74+
def run(self, inputs: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
75+
arguments = self._to_artifact_inputs(inputs)
76+
77+
with self.profiler.profile():
78+
raw_outputs = self.method.execute(arguments)
79+
80+
# ExecuTorch can reuse its output buffers on the next call, so each caller gets its own copy.
81+
return {f"output{i}": output.detach().clone() for i, output in enumerate(raw_outputs)}
82+
83+
def warmup(self, num_iterations: int = WARMUP_ITERATIONS):
84+
"""Run zeros through the method to warm the caches before real measurements begin."""
85+
zeros = [
86+
torch.zeros(shape, dtype=dtype)
87+
for shape, dtype in zip(self._input_shapes, self._input_dtypes)
88+
]
89+
for _ in range(num_iterations):
90+
self.method.execute(zeros)
91+
self.profiler.reset()
92+
93+
def _find_image_input(self) -> int:
94+
"""Input 0, unless input 0 is not a 4-D float and exactly one other input is."""
95+
image_like = [
96+
index
97+
for index, (shape, dtype) in enumerate(zip(self._input_shapes, self._input_dtypes))
98+
if len(shape) == 4 and dtype in _FLOAT_DTYPES
99+
]
100+
if 0 in image_like or len(image_like) != 1:
101+
return 0
102+
return image_like[0]
103+
104+
def _to_artifact_inputs(self, inputs: dict[str, torch.Tensor]) -> list[torch.Tensor]:
105+
return [
106+
inputs[_input_name(index)].detach().to("cpu", dtype).contiguous()
107+
for index, dtype in enumerate(self._input_dtypes)
108+
]

‎tests/test_executorch.py‎

Lines changed: 125 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,125 @@
1+
import sys
2+
3+
import pytest
4+
import torch
5+
6+
pytest.importorskip("executorch.runtime")
7+
8+
from executorch.backends.xnnpack.partition.xnnpack_partitioner import XnnpackPartitioner # noqa: E402
9+
from executorch.exir import to_edge_transform_and_lower # noqa: E402
10+
11+
from sab.runtimes.executorch import ExecuTorchRuntime # noqa: E402
12+
13+
pytestmark = pytest.mark.executorch
14+
15+
16+
class ConvModel(torch.nn.Module):
17+
"""One NCHW image input."""
18+
19+
def __init__(self):
20+
super().__init__()
21+
self.conv = torch.nn.Conv2d(3, 4, 3, padding=1)
22+
23+
def forward(self, images):
24+
return self.conv(images)
25+
26+
27+
class ImageAndSizesModel(ConvModel):
28+
"""Image first, then an int64 input, like the `orig_target_sizes` input of some detectors."""
29+
30+
def forward(self, images, sizes):
31+
return self.conv(images), sizes + 1
32+
33+
34+
class SizesAndImageModel(ConvModel):
35+
def forward(self, sizes, images):
36+
return sizes + 1, self.conv(images)
37+
38+
39+
CASES = {
40+
"conv": (ConvModel, (torch.rand(1, 3, 8, 8),)),
41+
"image_and_sizes": (ImageAndSizesModel, (torch.rand(1, 3, 8, 8), torch.tensor([[10, 20]]))),
42+
"sizes_and_image": (SizesAndImageModel, (torch.tensor([[10, 20]]), torch.rand(1, 3, 8, 8))),
43+
}
44+
45+
46+
@pytest.fixture(scope="session")
47+
def model_dir(tmp_path_factory):
48+
return tmp_path_factory.mktemp("pte")
49+
50+
51+
@pytest.fixture(scope="session")
52+
def build_pte(model_dir):
53+
"""Export each tiny model once per session. Returns (module, example inputs, path)."""
54+
built = {}
55+
56+
def build(case: str):
57+
if case not in built:
58+
torch.manual_seed(0)
59+
model_class, example_inputs = CASES[case]
60+
module = model_class().eval()
61+
exported = torch.export.export(module, example_inputs)
62+
program = to_edge_transform_and_lower(exported, partitioner=[XnnpackPartitioner()]).to_executorch()
63+
path = model_dir / f"{case}.pte"
64+
path.write_bytes(program.buffer)
65+
built[case] = (module, example_inputs, str(path))
66+
return built[case]
67+
68+
return build
69+
70+
71+
@pytest.mark.parametrize("platform", ["linux", "darwin"])
72+
def test_cpu_is_available_on_linux_and_macos(monkeypatch, platform):
73+
monkeypatch.setattr(sys, "platform", platform)
74+
assert ExecuTorchRuntime.is_available("cpu")
75+
76+
77+
def test_run_matches_the_torch_model_and_outputs_own_their_memory(build_pte):
78+
module, _, path = build_pte("conv")
79+
runtime = ExecuTorchRuntime(path, "cpu", "fp32")
80+
images = torch.rand(1, 3, 8, 8)
81+
82+
first = runtime.run({"input0": images})
83+
runtime.run({"input0": torch.rand(1, 3, 8, 8)})
84+
85+
assert runtime.input_spec.name == "input0"
86+
assert set(first) == {"output0"}
87+
torch.testing.assert_close(first["output0"], module(images), atol=1e-4, rtol=1e-4)
88+
89+
90+
def test_run_passes_inputs_in_positional_order(build_pte):
91+
module, _, path = build_pte("image_and_sizes")
92+
runtime = ExecuTorchRuntime(path, "cpu", "fp32")
93+
images = torch.rand(1, 3, 8, 8)
94+
sizes = torch.tensor([[10, 20]])
95+
96+
# The dict order is the reverse of the positional order.
97+
outputs = runtime.run({"input1": sizes, "input0": images})
98+
99+
assert runtime.input_spec.name == "input0"
100+
torch.testing.assert_close(outputs["output0"], module(images, sizes)[0], atol=1e-4, rtol=1e-4)
101+
assert outputs["output1"].tolist() == [[11.0, 21.0]]
102+
103+
104+
def test_run_finds_the_image_at_a_later_position(build_pte):
105+
module, _, path = build_pte("sizes_and_image")
106+
runtime = ExecuTorchRuntime(path, "cpu", "fp32")
107+
images = torch.rand(1, 3, 8, 8)
108+
sizes = torch.tensor([[10, 20]])
109+
110+
outputs = runtime.run({"input0": sizes, "input1": images})
111+
112+
assert runtime.input_spec.name == "input1"
113+
assert runtime.input_spec.shape == (1, 3, 8, 8)
114+
assert outputs["output0"].tolist() == [[11.0, 21.0]]
115+
torch.testing.assert_close(outputs["output1"], module(sizes, images)[1], atol=1e-4, rtol=1e-4)
116+
117+
118+
def test_run_records_one_timing_per_call_after_warmup(build_pte):
119+
*_, path = build_pte("conv")
120+
runtime = ExecuTorchRuntime(path, "cpu", "fp32")
121+
assert runtime.profiler.timings == []
122+
123+
for expected in (1, 2, 3):
124+
runtime.run({"input0": torch.rand(1, 3, 8, 8)})
125+
assert len(runtime.profiler.timings) == expected

0 commit comments

Comments
 (0)