Build Native wheels in ExecuTorch release CI

.ci/scripts/tests/test_native_wheel_version.py

@@ -0,0 +1,56 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import importlib.util
+import os
+import unittest
+from pathlib import Path
+from unittest.mock import patch
+
+
+ROOT = Path(__file__).resolve().parents[3]
+if ROOT.name == "oss":
+    ROOT = ROOT.parent
+SPEC = importlib.util.spec_from_file_location(
+    "wheel_version", ROOT / "backends/native/python/wheel_version.py"
+)
+assert SPEC is not None and SPEC.loader is not None
+VERSION = importlib.util.module_from_spec(SPEC)
+SPEC.loader.exec_module(VERSION)
+
+
+class NativeWheelVersionTest(unittest.TestCase):
+    def test_invalid_metadata_configuration(self) -> None:
+        for field, settings in (
+            ("dependencies", {}),
+            ("version", {"input": "version.txt"}),
+        ):
+            with self.subTest(field=field, settings=settings):
+                with self.assertRaisesRegex(
+                    ValueError, "Only version metadata without settings is supported"
+                ):
+                    VERSION.dynamic_metadata(field, settings)
+
+    def test_build_version_overrides_source_version(self) -> None:
+        for version in ("1.6.0", "1.6.0rc2", "1.7.0.dev20261021"):
+            with self.subTest(version=version), patch.dict(
+                os.environ, {"BUILD_VERSION": f" {version}\n"}
+            ):
+                self.assertEqual(VERSION.dynamic_metadata("version", {}), version)
+
+    def test_local_build_uses_version_file(self) -> None:
+        for version in ("", " \n"):
+            with self.subTest(version=version), patch.dict(
+                os.environ, {"BUILD_VERSION": version}
+            ):
+                self.assertEqual(
+                    VERSION.dynamic_metadata("version", {}),
+                    (ROOT / "version.txt").read_text().strip(),
+                )
+
+
+if __name__ == "__main__":
+    unittest.main()

.ci/scripts/wheel/envvar_native.sh

@@ -0,0 +1,18 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+export NATIVE_WHEEL_ROOT="${GITHUB_WORKSPACE}/${REPOSITORY}"
+export NATIVE_WHEEL_DEPS="${GITHUB_WORKSPACE}/native-wheel-deps"
+export CMAKE_PREFIX_PATH="${NATIVE_WHEEL_DEPS}/libzip:${NATIVE_WHEEL_DEPS}:${NATIVE_WHEEL_DEPS}/shaderc${CMAKE_PREFIX_PATH:+:${CMAKE_PREFIX_PATH}}"
+export CMAKE_BUILD_PARALLEL_LEVEL="${CMAKE_BUILD_PARALLEL_LEVEL:-4}"
+export FLATC_EXECUTABLE="${NATIVE_WHEEL_DEPS}/flatbuffers/flatc"
+export CMAKE_ARGS="${CMAKE_ARGS:-} -DFLATC_EXECUTABLE=${FLATC_EXECUTABLE} -DGLSLC_PATH=${NATIVE_WHEEL_DEPS}/shaderc/bin/glslc -DEXECUTORCH_NATIVE_EXTRA_LICENSES=${NATIVE_WHEEL_DEPS}/libzip-1.11.4/LICENSE"
+export PIP_EXTRA_INDEX_URL="${PIP_EXTRA_INDEX_URL:+${PIP_EXTRA_INDEX_URL} }https://pypi.org/simple"
+
+# Nova sources this hook before invoking setup.py, allowing a subdirectory build.
+if [[ "$(pwd)" == "${NATIVE_WHEEL_ROOT}" ]]; then
+    cd "${NATIVE_WHEEL_ROOT}/backends/native/python" || exit 1
+fi

.ci/scripts/wheel/export_native.py

@@ -0,0 +1,60 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Generate Native wheel fixtures with the exporter from this checkout."""
+
+import argparse
+from pathlib import Path
+
+import numpy as np
+import torch
+from executorch.exir.native import to_native
+
+
+class Add(torch.nn.Module):
+    def forward(self, value: torch.Tensor) -> torch.Tensor:
+        return value + value
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--output", type=Path, required=True)
+    args = parser.parse_args()
+    args.output.mkdir(parents=True, exist_ok=True)
+    torch.manual_seed(0)
+    conv_bn = torch.nn.Sequential(
+        torch.nn.Conv2d(3, 4, kernel_size=3, padding=1),
+        torch.nn.BatchNorm2d(4),
+        torch.nn.ReLU(),
+    )
+    with torch.no_grad():
+        conv_bn[1].running_mean.copy_(torch.linspace(-1, 1, 4))
+        conv_bn[1].running_var.copy_(torch.linspace(0.5, 1.5, 4))
+        conv_bn[1].weight.copy_(torch.linspace(0.8, 1.2, 4))
+        conv_bn[1].bias.copy_(torch.linspace(-0.2, 0.2, 4))
+    models = {
+        "add": (Add(), torch.arange(1, 7, dtype=torch.float32).reshape(1, 6)),
+        "conv_bn": (
+            conv_bn,
+            torch.randn(1, 3, 8, 8),
+        ),
+    }
+    for name, (model, value) in models.items():
+        model.eval()
+        with torch.no_grad():
+            expected = model(value)
+        to_native(torch.export.export(model, (value,))).save(
+            str(args.output / f"{name}.ptn")
+        )
+        np.savez(
+            args.output / f"{name}.npz",
+            input=value.numpy(),
+            expected=expected.numpy(),
+        )
+
+
+if __name__ == "__main__":
+    main()

.ci/scripts/wheel/post_build_native.sh

@@ -0,0 +1,37 @@
+#!/bin/bash
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+set -euxo pipefail
+
+cd "${NATIVE_WHEEL_ROOT}/backends/native/python"
+wheels=(dist/executorch_native-*.whl)
+[[ ${#wheels[@]} == 1 && -f "${wheels[0]}" ]]
+python -m auditwheel repair "${wheels[0]}" --wheel-dir repaired
+rm "${wheels[0]}"
+mv repaired/*.whl dist/
+wheels=(dist/executorch_native-*.whl)
+[[ ${#wheels[@]} == 1 && -f "${wheels[0]}" ]]
+
+# Only the export environment needs PyTorch and the source checkout.
+torchao_version=$(cd "${NATIVE_WHEEL_ROOT}" && python -c \
+    'from install_requirements import TORCHAO_NIGHTLY_VERSION; print(TORCHAO_NIGHTLY_VERSION)')
+python -m pip install torch "torchao==${torchao_version}" numpy flatbuffers packaging pyyaml ruamel.yaml tabulate \
+    --index-url https://download.pytorch.org/whl/test/cpu \
+    --extra-index-url https://download.pytorch.org/whl/nightly/cpu \
+    --extra-index-url https://pypi.org/simple
+mkdir -p "${NATIVE_WHEEL_DEPS}/export-source"
+ln -s "${NATIVE_WHEEL_ROOT}" "${NATIVE_WHEEL_DEPS}/export-source/executorch"
+PYTHONPATH="${NATIVE_WHEEL_DEPS}/export-source" python \
+    "${NATIVE_WHEEL_ROOT}/.ci/scripts/wheel/export_native.py" \
+    --output "${NATIVE_WHEEL_DEPS}/fixtures"
+python "${NATIVE_WHEEL_ROOT}/.ci/scripts/wheel/test_native.py" \
+    --wheel "${wheels[0]}" --source "${NATIVE_WHEEL_ROOT}" \
+    --fixtures "${NATIVE_WHEEL_DEPS}/fixtures"
+
+# Nova's artifact steps also consume dist at the repository root.
+mkdir -p "${NATIVE_WHEEL_ROOT}/dist"
+cp "${wheels[0]}" "${NATIVE_WHEEL_ROOT}/dist/"

.ci/scripts/wheel/pre_build_native.sh

@@ -0,0 +1,30 @@
+#!/bin/bash
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+set -euxo pipefail
+
+python -m pip install build 'setuptools>=77' 'cmake>=3.26,<4' ninja auditwheel patchelf
+python -m unittest discover -s "${NATIVE_WHEEL_ROOT}/.ci/scripts/tests" \
+    -p test_native_wheel_version.py
+mkdir -p "${NATIVE_WHEEL_DEPS}"
+conda create -y -p "${NATIVE_WHEEL_DEPS}/shaderc" -c conda-forge shaderc zlib
+
+curl --fail --location --retry 3 https://libzip.org/download/libzip-1.11.4.tar.xz \
+    --output "${NATIVE_WHEEL_DEPS}/libzip.tar.xz"
+tar -xJf "${NATIVE_WHEEL_DEPS}/libzip.tar.xz" -C "${NATIVE_WHEEL_DEPS}"
+cmake -S "${NATIVE_WHEEL_DEPS}/libzip-1.11.4" -B "${NATIVE_WHEEL_DEPS}/libzip-build" \
+    -DCMAKE_BUILD_TYPE=Release -DCMAKE_INSTALL_PREFIX="${NATIVE_WHEEL_DEPS}/libzip" \
+    -DCMAKE_POSITION_INDEPENDENT_CODE=ON -DBUILD_SHARED_LIBS=OFF \
+    -DBUILD_TOOLS=OFF -DBUILD_REGRESS=OFF -DBUILD_EXAMPLES=OFF -DBUILD_DOC=OFF \
+    -DENABLE_BZIP2=OFF -DENABLE_LZMA=OFF -DENABLE_ZSTD=OFF \
+    -DENABLE_OPENSSL=OFF -DENABLE_GNUTLS=OFF -DENABLE_MBEDTLS=OFF
+cmake --build "${NATIVE_WHEEL_DEPS}/libzip-build" --target install --parallel
+
+cmake -S "${NATIVE_WHEEL_ROOT}/third-party/flatbuffers" \
+    -B "${NATIVE_WHEEL_DEPS}/flatbuffers" \
+    -DCMAKE_BUILD_TYPE=Release -DFLATBUFFERS_BUILD_TESTS=OFF
+cmake --build "${NATIVE_WHEEL_DEPS}/flatbuffers" --target flatc --parallel

.ci/scripts/wheel/test_native.py

@@ -0,0 +1,97 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Install a repaired Native wheel into a clean environment and test GGML."""
+
+import argparse
+import os
+import shutil
+import subprocess
+import tempfile
+import venv
+from pathlib import Path
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--wheel", type=Path, required=True)
+    parser.add_argument("--source", type=Path, required=True)
+    parser.add_argument("--fixtures", type=Path, required=True)
+    args = parser.parse_args()
+    wheel = args.wheel.resolve()
+    source = args.source.resolve()
+    fixtures = args.fixtures.resolve()
+    with tempfile.TemporaryDirectory(prefix="native-wheel-test-") as directory:
+        work = Path(directory)
+        venv.create(work / "venv", with_pip=True)
+        python = work / "venv/bin/python"
+        env = os.environ.copy()
+        for name in ("PYTHONPATH", "PYTHONHOME", "LD_LIBRARY_PATH", "LD_PRELOAD"):
+            env.pop(name, None)
+        env["PYTHONNOUSERSITE"] = "1"
+        subprocess.run(
+            [str(python), "-m", "pip", "install", str(wheel)],
+            cwd=work,
+            env=env,
+            check=True,
+            timeout=600,
+        )
+        subprocess.run(
+            [str(python), "-m", "pip", "check"],
+            cwd=work,
+            env=env,
+            check=True,
+            timeout=60,
+        )
+        shutil.copytree(fixtures, work / "fixtures")
+        shutil.copyfile(
+            source / "backends/native/pybindings/test/test_bindings.py",
+            work / "test_bindings.py",
+        )
+        subprocess.run(
+            [
+                str(python),
+                "-I",
+                "-c",
+                """
+import importlib.metadata
+import importlib.util
+from pathlib import Path
+import numpy as np
+from executorch_native import Model
+
+for name in ('torch', 'executorch'):
+    assert importlib.util.find_spec(name) is None, name
+installed = {d.metadata['Name'].lower() for d in importlib.metadata.distributions()}
+assert installed <= {'pip', 'setuptools', 'executorch-native', 'numpy'}, installed
+for name in ('add', 'conv_bn'):
+    package = Path('fixtures') / f'{name}.ptn'
+    with np.load(package.with_suffix('.npz')) as data:
+        session = Model.load(package).create_session('ggml')
+        (actual,) = session.run([data['input']])
+        np.testing.assert_allclose(actual, data['expected'], rtol=1e-4, atol=1e-4)
+    print(f'PASS: {package.name} with GGML')
+""",
+            ],
+            cwd=work,
+            env=env,
+            check=True,
+            timeout=120,
+        )
+        env["NATIVE_TEST_PACKAGE"] = str(work / "fixtures/add.ptn")
+        env["NATIVE_TEST_ENGINE"] = "ggml"
+        subprocess.run(
+            [str(python), "-m", "unittest", "-v", "test_bindings"],
+            cwd=work,
+            env=env,
+            check=True,
+            timeout=300,
+        )
+    print(f"PASS: clean install of {wheel.name}")
+
+
+if __name__ == "__main__":
+    main()

.github/workflows/build-wheels-native-aarch64-linux.yml

@@ -0,0 +1,73 @@
+# From https://github.com/pytorch/test-infra/wiki/Using-Nova-Reusable-Build-Workflows
+name: Build Aarch64 Linux Native Wheels
+
+on:
+  pull_request:
+    paths:
+      - .ci/**/*
+      - .github/workflows/build-wheels-native-aarch64-linux.yml
+      - '**/CMakeLists.txt'
+      - backends/native/**/*
+      - backends/vulkan/**/*
+      - backends/transforms/**/*
+      - exir/**/*
+      - schema/**/*
+      - examples/**/*
+      - version.txt
+      - .gitmodules
+      - pyproject.toml
+      - setup.py
+      - install_utils.py
+      - tools/cmake/**/*
+  push:
+    branches:
+      - nightly
+      - release/*
+    tags:
+      # NOTE: Binary build pipelines should only get triggered on release candidate builds
+      # Release candidate tags look like: v1.11.0-rc1
+      - v[0-9]+.[0-9]+.[0-9]+-rc[0-9]+
+      - ciflow/binaries/*
+  workflow_dispatch:
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref_name }}-${{ github.ref_type == 'branch' && github.sha }}-${{ github.event_name == 'workflow_dispatch' }}
+  cancel-in-progress: true
+
+jobs:
+  generate-matrix:
+    uses: pytorch/test-infra/.github/workflows/generate_binary_build_matrix.yml@main
+    with:
+      package-type: wheel
+      os: linux-aarch64
+      test-infra-repository: pytorch/test-infra
+      test-infra-ref: main
+      runner-fleet: osdc
+      with-cuda: disabled
+      with-rocm: disabled
+      python-versions: '["3.11", "3.12", "3.13", "3.14", "3.15"]'
+
+  build:
+    needs: generate-matrix
+    permissions:
+      id-token: write
+      contents: read
+    name: executorch-native
+    uses: pytorch/test-infra/.github/workflows/build_wheels_linux.yml@main
+    with:
+      repository: pytorch/executorch
+      ref: ""
+      test-infra-repository: pytorch/test-infra
+      test-infra-ref: main
+      build-matrix: ${{ needs.generate-matrix.outputs.matrix }}
+      submodules: recursive
+      env-var-script: .ci/scripts/wheel/envvar_native.sh
+      pre-script: ../../../.ci/scripts/wheel/pre_build_native.sh
+      post-script: ../../../.ci/scripts/wheel/post_build_native.sh
+      package-name: executorch-native
+      # The post-build script installs the wheel in a clean environment and
+      # tests it; Nova's smoke test requires a pinned torch dependency.
+      run-smoke-test: false
+      # Only "push" uploads; executorch-native is not published yet.
+      trigger-event: development
+      architecture: aarch64