Skip to content

Commit b7d1e44

Browse files
JCNTHJCNTH
committed
Update
[ghstack-poisoned]
2 parents ded2b6f + 1ae2d09 commit b7d1e44

47 files changed

Lines changed: 1813 additions & 857 deletions

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎.ci/scripts/test_minimal_wheel.sh‎

Lines changed: 165 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,165 @@
1+
#!/bin/bash
2+
# Copyright (c) Meta Platforms, Inc. and affiliates.
3+
# All rights reserved.
4+
#
5+
# This source code is licensed under the BSD-style license found in the
6+
# LICENSE file in the root directory of this source tree.
7+
8+
set -euxo pipefail
9+
10+
PYTHON_EXECUTABLE="${PYTHON_EXECUTABLE:-python}"
11+
REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
12+
BUILD_VENV="${REPO_ROOT}/.venv-minimal-build"
13+
TEST_VENV="${REPO_ROOT}/.venv-minimal-test"
14+
15+
rm -rf "${BUILD_VENV}" "${TEST_VENV}" "${REPO_ROOT}/dist" "${REPO_ROOT}/pip-out"
16+
17+
"${PYTHON_EXECUTABLE}" -m venv "${BUILD_VENV}"
18+
source "${BUILD_VENV}/bin/activate"
19+
python -m pip install --upgrade pip
20+
python -m pip install \
21+
"cmake>=3.24,<4.0.0" \
22+
"numpy>=2.0.0" \
23+
packaging \
24+
pyyaml \
25+
setuptools \
26+
wheel \
27+
zstd \
28+
certifi \
29+
torch \
30+
torchvision \
31+
--index-url https://download.pytorch.org/whl/cpu \
32+
--extra-index-url https://pypi.org/simple
33+
34+
(
35+
cd "${REPO_ROOT}"
36+
EXECUTORCH_BUILD_MINIMAL=1 python setup.py bdist_wheel
37+
)
38+
39+
WHEEL_FILE="$(find "${REPO_ROOT}/dist" -maxdepth 1 -name 'executorch-*.whl' | head -1)"
40+
test -n "${WHEEL_FILE}"
41+
42+
python - "${WHEEL_FILE}" <<'PY'
43+
import re
44+
import sys
45+
import zipfile
46+
47+
wheel_file = sys.argv[1]
48+
with zipfile.ZipFile(wheel_file) as wheel:
49+
names = wheel.namelist()
50+
metadata_name = next(
51+
(name for name in names if name.endswith(".dist-info/METADATA")), None
52+
)
53+
if metadata_name is None:
54+
raise AssertionError(f"{wheel_file} has no METADATA")
55+
metadata_text = wheel.read(metadata_name).decode("utf-8")
56+
57+
for forbidden in (
58+
"executorch/backends/",
59+
"executorch/examples/",
60+
"executorch/kernels/",
61+
"executorch/runtime/",
62+
"executorch/devtools/",
63+
"executorch/extension/pybindings/",
64+
):
65+
matches = [name for name in names if name.startswith(forbidden)]
66+
if matches:
67+
raise AssertionError(f"{wheel_file} unexpectedly contains {matches[:5]}")
68+
69+
extensions = [
70+
name
71+
for name in names
72+
if name.endswith((".so", ".dylib", ".dll", ".pyd")) and "flatc" not in name
73+
]
74+
if extensions:
75+
raise AssertionError(f"{wheel_file} unexpectedly contains extensions: {extensions}")
76+
77+
78+
def _dist_name(requirement):
79+
name = re.split(r"[ ;\[<>=!~(]", requirement.strip(), maxsplit=1)[0]
80+
return re.sub(r"[-_.]+", "-", name).lower()
81+
82+
83+
# Only the core (non-extra) Requires-Dist entries define what a plain
84+
# "pip install" pulls; ignore the optional extras (cortex_m, vgf, ...).
85+
declared = {
86+
_dist_name(line.split(":", 1)[1])
87+
for line in metadata_text.splitlines()
88+
if line.startswith("Requires-Dist:") and "extra==" not in line.replace(" ", "")
89+
}
90+
# The minimal wheel must declare EXACTLY this core set and nothing else -- the
91+
# same names as `keep` in setup.py:_minimal_dependencies(). Exact match catches
92+
# both a heavy full-wheel dep leaking in (coremltools, pandas, or a re-added
93+
# mpmath/torch) and a required dep going missing.
94+
expected = {
95+
"flatbuffers",
96+
"numpy",
97+
"packaging",
98+
"pyyaml",
99+
"ruamel-yaml",
100+
"sympy",
101+
"tabulate",
102+
"typing-extensions",
103+
}
104+
if declared != expected:
105+
raise AssertionError(
106+
f"{wheel_file} minimal core deps mismatch: "
107+
f"unexpected={sorted(declared - expected)} missing={sorted(expected - declared)}"
108+
)
109+
PY
110+
111+
deactivate
112+
113+
"${PYTHON_EXECUTABLE}" -m venv "${TEST_VENV}"
114+
source "${TEST_VENV}/bin/activate"
115+
python -m pip install --upgrade pip
116+
# torch and torchvision are needed to export a model but are intentionally not
117+
# declared as wheel dependencies (consumers are expected to bring their own).
118+
python -m pip install \
119+
"torch" \
120+
"torchvision" \
121+
--index-url https://download.pytorch.org/whl/cpu \
122+
--extra-index-url https://pypi.org/simple
123+
# Install the minimal wheel WITHOUT --no-deps so pip resolves its declared
124+
# dependencies, confirming the slim set is correct and resolvable. (That no heavy
125+
# deps sneak in is guaranteed by the METADATA exact-match check above, which
126+
# covers the wheel's direct Requires-Dist.)
127+
python -m pip install \
128+
"${WHEEL_FILE}" \
129+
--index-url https://download.pytorch.org/whl/cpu \
130+
--extra-index-url https://pypi.org/simple
131+
132+
# flatc is the only compiled artifact in the minimal wheel and the reason it is
133+
# platform specific. Confirm it ships, resolves through _get_flatc_path() (the
134+
# executorch.data.bin lookup added for this build mode), and actually runs.
135+
python - <<'PY'
136+
import subprocess
137+
138+
from executorch.exir._serialize._flatbuffer import _get_flatc_path
139+
140+
flatc_path = _get_flatc_path()
141+
print(f"flatc resolved to: {flatc_path}")
142+
subprocess.run([flatc_path, "--version"], check=True)
143+
PY
144+
145+
python - <<'PY'
146+
from pathlib import Path
147+
148+
import torch
149+
from torch.export import export
150+
from torchvision.models import mobilenet_v2
151+
152+
from executorch.exir import to_edge_transform_and_lower
153+
154+
model = mobilenet_v2(weights=None).eval()
155+
example_inputs = (torch.randn(1, 3, 224, 224),)
156+
157+
edge_program = to_edge_transform_and_lower(export(model, example_inputs))
158+
executorch_program = edge_program.to_executorch()
159+
160+
output_path = Path("mv2_minimal.pte")
161+
with output_path.open("wb") as output_file:
162+
executorch_program.write_to_file(output_file)
163+
164+
assert output_path.stat().st_size > 0
165+
PY

‎.github/workflows/pull.yml‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,36 @@ jobs:
6161
6262
PYTHON_EXECUTABLE=python bash .ci/scripts/test_wheel_package_qnn.sh "${{ matrix.python-version }}"
6363
64+
test-minimal-wheel-linux:
65+
needs: changed-files
66+
if: |
67+
github.event_name != 'pull_request' ||
68+
contains(needs.changed-files.outputs.changed-files, '.ci/scripts/test_minimal_wheel.sh') ||
69+
contains(needs.changed-files.outputs.changed-files, '.github/workflows/pull.yml') ||
70+
contains(needs.changed-files.outputs.changed-files, 'exir/') ||
71+
contains(needs.changed-files.outputs.changed-files, 'extension/flat_tensor') ||
72+
contains(needs.changed-files.outputs.changed-files, 'extension/pytree') ||
73+
contains(needs.changed-files.outputs.changed-files, 'pyproject.toml') ||
74+
contains(needs.changed-files.outputs.changed-files, 'schema/') ||
75+
contains(needs.changed-files.outputs.changed-files, 'setup.py') ||
76+
contains(needs.changed-files.outputs.changed-files, 'tools/cmake/')
77+
name: test-minimal-wheel-linux
78+
uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main
79+
permissions:
80+
id-token: write
81+
contents: read
82+
with:
83+
runner: linux.2xlarge
84+
docker-image: ci-image:executorch-ubuntu-22.04-clang12
85+
submodules: 'recursive'
86+
ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
87+
timeout: 120
88+
script: |
89+
CONDA_ENV=$(conda env list --json | jq -r ".envs | .[-1]")
90+
conda activate "${CONDA_ENV}"
91+
92+
PYTHON_EXECUTABLE=python bash .ci/scripts/test_minimal_wheel.sh
93+
6494
test-setup-linux-gcc:
6595
name: test-setup-linux-gcc
6696
uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main

‎README-wheel.md‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,16 @@ The `executorch` pip package is in beta.
88
* Supported python versions: 3.10, 3.11, 3.12, 3.13
99
* Compatible systems: Linux x86_64, Linux aarch64, macOS aarch64
1010

11+
To build a minimal wheel from source, set
12+
`EXECUTORCH_BUILD_MINIMAL=1` when running `pip wheel` or `pip install`.
13+
That wheel contains the Python EXIR export path and `flatc` for `.pte`
14+
serialization, but omits runtime pybindings, kernels, backend packages, headers,
15+
examples, and devtools. It also declares only the Python dependencies the export
16+
path needs (no `coremltools`, `pandas`, `scikit-learn`, `hydra-core`, or
17+
`omegaconf`), so a normal install stays small. Like the full wheel it does not
18+
bundle PyTorch, so install a compatible `torch` separately. The wheel is still
19+
platform specific because it ships `flatc`.
20+
1121
The prebuilt `executorch.runtime` module included in this package provides a way
1222
to run ExecuTorch `.pte` files, with some restrictions:
1323
* Only [core ATen operators](docs/source/ir-ops-set-definition.md) are linked into the prebuilt module

‎backends/cadence/aot/BUCK‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -426,6 +426,7 @@ fbcode_target(_kind = runtime.python_library,
426426
typing = True,
427427
deps = [
428428
"//caffe2:torch",
429+
"//executorch/backends/transforms:permute_pass_utils",
429430
"//pytorch/ao:torchao",
430431
],
431432
)

‎backends/cadence/aot/compiler_funcs.py‎

Lines changed: 73 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,8 @@
1212
from typing import Any, cast, Optional, Union
1313

1414
import torch
15+
16+
from executorch.backends.transforms.permute_pass_utils import get_arg
1517
from torch._inductor.decomposition import remove_decompositions
1618
from torch.fx import GraphModule
1719
from torch.fx.passes.infra.pass_base import PassBase, PassResult
@@ -159,6 +161,40 @@ def extract_output_dequant_params(
159161
raise ValueError("Could not find dequantize_per_tensor at the output of the graph")
160162

161163

164+
def extract_all_output_dequant_params(
165+
module: torch.fx.GraphModule,
166+
) -> list[QuantArgs | None]:
167+
"""
168+
Extract per-output dequantization parameters from a multi-output model.
169+
170+
Returns a QuantArgs tuple for outputs ending in dequantize_per_tensor
171+
or None for outputs that aren't dequantized.
172+
"""
173+
output_nodes = module.graph.find_nodes(op="output")
174+
if not output_nodes:
175+
raise ValueError("No output node in graph")
176+
output_args = output_nodes[0].args[0]
177+
if not isinstance(output_args, (tuple, list)):
178+
output_args = (output_args,)
179+
180+
dequant_ops = _get_dequantize_ops()
181+
params: list[QuantArgs | None] = []
182+
for out in output_args:
183+
if not isinstance(out, torch.fx.Node) or out.target not in dequant_ops:
184+
params.append(None)
185+
continue
186+
params.append(
187+
(
188+
float(get_arg(out, "scale", float)),
189+
int(get_arg(out, "zero_point", int)),
190+
int(get_arg(out, "quant_min", int)),
191+
int(get_arg(out, "quant_max", int)),
192+
get_arg(out, "dtype", torch.dtype),
193+
)
194+
)
195+
return params
196+
197+
162198
def extract_output_dequant_params_through_permute(
163199
module: torch.fx.GraphModule,
164200
) -> QuantArgs:
@@ -400,33 +436,60 @@ def sink_dequants(program: torch.export.ExportedProgram) -> None:
400436

401437
class QuantizedOutputWrapper(torch.nn.Module):
402438
"""
403-
Wrapper that quantizes a model's output so it produces uint8 tensors.
439+
Wrapper that quantizes a model's output(s) so they produce quantized tensors.
404440
405441
Mirrors QuantizedInputWrapper: the wrapper adds a quantize_per_tensor after
406-
the model's output. When the graph is traced, the dequant (from the model) →
442+
each output. When the graph is traced, the dequant (from the model) →
407443
quant (from the wrapper) pair with matching parameters folds away, leaving
408444
the output in its quantized form.
409445
410446
Args:
411447
module: The module to wrap (may already be a QuantizedInputWrapper).
412-
output_quant_args: (scale, zero_point, qmin, qmax, dtype) for the output.
448+
output_quant_args: Quantization parameters — either a single QuantArgs
449+
tuple or a list with one entry per output.
413450
"""
414451

415452
def __init__(
416453
self,
417454
module: torch.nn.Module,
418-
output_quant_args: QuantArgs,
455+
output_quant_args: Union[QuantArgs, list[QuantArgs | None]],
419456
) -> None:
420457
super().__init__()
421458
self.module: torch.nn.Module = module
422-
self.output_quant_args: QuantArgs = output_quant_args
459+
if isinstance(output_quant_args, list):
460+
self._multi_output: bool = True
461+
self._per_output_args: list[QuantArgs | None] = output_quant_args
462+
else:
463+
self._multi_output = False
464+
self._per_output_args = [output_quant_args]
423465

424466
def forward(self, *args: torch.Tensor) -> Any:
425-
result = self.module(*args)
426-
scale, zp, qmin, qmax, dtype = self.output_quant_args
427-
return torch.ops.quantized_decomposed.quantize_per_tensor.default(
428-
result, scale, zp, qmin, qmax, dtype
429-
)
467+
model_output = self.module(*args)
468+
if not self._multi_output:
469+
quant_args = self._per_output_args[0]
470+
assert quant_args is not None
471+
scale, zero_point, quant_min, quant_max, dtype = quant_args
472+
return torch.ops.quantized_decomposed.quantize_per_tensor.default(
473+
model_output, scale, zero_point, quant_min, quant_max, dtype
474+
)
475+
476+
quantized_outputs: list[torch.Tensor] = []
477+
for output_index, output_tensor in enumerate(model_output):
478+
quant_args = (
479+
self._per_output_args[output_index]
480+
if output_index < len(self._per_output_args)
481+
else None
482+
)
483+
if quant_args is None:
484+
quantized_outputs.append(output_tensor)
485+
else:
486+
scale, zero_point, quant_min, quant_max, dtype = quant_args
487+
quantized_outputs.append(
488+
torch.ops.quantized_decomposed.quantize_per_tensor.default(
489+
output_tensor, scale, zero_point, quant_min, quant_max, dtype
490+
)
491+
)
492+
return tuple(quantized_outputs)
430493

431494

432495
def _get_transparent_ops() -> set[Any]:

‎backends/webgpu/runtime/WebGPUCompat.h‎

Lines changed: 6 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -10,26 +10,15 @@
1010

1111
#include <webgpu/webgpu.h>
1212

13-
#if defined(__EMSCRIPTEN__)
14-
#include <emscripten/emscripten.h>
15-
#else
16-
#include <chrono>
17-
#include <thread>
18-
#endif
13+
#include <cstdint>
1914

2015
namespace executorch::backends::webgpu {
2116

22-
// Make progress on pending WebGPU callbacks; callers loop until their done flag
23-
// is set. Native (Dawn): pump the event queue + brief yield (no busy-spin).
24-
// Browser: yield to the JS event loop.
25-
inline void webgpu_poll(WGPUInstance instance) {
26-
#if defined(__EMSCRIPTEN__)
27-
(void)instance;
28-
emscripten_sleep(0);
29-
#else
30-
wgpuInstanceProcessEvents(instance);
31-
std::this_thread::sleep_for(std::chrono::microseconds(50));
32-
#endif
17+
// Caller's instance must enable TimedWaitAny; returns the WaitAny status.
18+
inline WGPUWaitStatus webgpu_wait(WGPUInstance instance, WGPUFuture future) {
19+
WGPUFutureWaitInfo info = {};
20+
info.future = future;
21+
return wgpuInstanceWaitAny(instance, 1, &info, UINT64_MAX);
3322
}
3423

3524
} // namespace executorch::backends::webgpu

0 commit comments

Comments
 (0)