Skip to content

Commit f680a4a

Browse files
authored
Merge branch 'main' into export-D68109702
2 parents f26bca2 + 9f47380 commit f680a4a

File tree

2 files changed

+189
-154
lines changed

2 files changed

+189
-154
lines changed

examples/xnnpack/aot_compiler.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -92,6 +92,7 @@
9292
logging.info("Quantizing Model...")
9393
# TODO(T165162973): This pass shall eventually be folded into quantizer
9494
model = quantize(model, example_inputs)
95+
ep = torch.export.export_for_training(model, example_inputs)
9596

9697
edge = to_edge_transform_and_lower(
9798
ep,

install_requirements.py

Lines changed: 188 additions & 154 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,9 @@
66
# LICENSE file in the root directory of this source tree.
77

88

9+
import argparse
910
import glob
11+
import itertools
1012
import os
1113
import platform
1214
import re
@@ -63,69 +65,23 @@ def python_is_compatible():
6365
return True
6466

6567

66-
if not python_is_compatible():
67-
sys.exit(1)
68+
def clean():
69+
print("Cleaning build artifacts...")
70+
print("Cleaning pip-out/...")
71+
shutil.rmtree("pip-out/", ignore_errors=True)
72+
dirs = glob.glob("cmake-out*/") + glob.glob("cmake-android-out/")
73+
for d in dirs:
74+
print(f"Cleaning {d}...")
75+
shutil.rmtree(d, ignore_errors=True)
76+
print("Done cleaning build artifacts.")
6877

69-
# Parse options.
7078

71-
EXECUTORCH_BUILD_PYBIND = ""
72-
CMAKE_ARGS = os.getenv("CMAKE_ARGS", "")
73-
CMAKE_BUILD_ARGS = os.getenv("CMAKE_BUILD_ARGS", "")
74-
USE_PYTORCH_NIGHTLY = True
79+
VALID_PYBINDS = ["coreml", "mps", "xnnpack"]
7580

76-
args = sys.argv[1:]
77-
for arg in args:
78-
if arg == "--pybind":
79-
pass
80-
elif arg in ["coreml", "mps", "xnnpack"]:
81-
if "--pybind" in args:
82-
arg_upper = arg.upper()
83-
EXECUTORCH_BUILD_PYBIND = "ON"
84-
CMAKE_ARGS += f" -DEXECUTORCH_BUILD_{arg_upper}=ON"
85-
else:
86-
print(f"Error: {arg} must follow --pybind")
87-
sys.exit(1)
88-
elif arg == "off":
89-
if "--pybind" in args:
90-
if EXECUTORCH_BUILD_PYBIND == "ON":
91-
print("Cannot turnoff pybind option as it is already set.")
92-
sys.exit(1)
93-
EXECUTORCH_BUILD_PYBIND = "OFF"
94-
else:
95-
print(f"Error: {arg} must follow --pybind")
96-
sys.exit(1)
97-
98-
elif arg == "--clean":
99-
print("Cleaning build artifacts...")
100-
print("Cleaning pip-out/...")
101-
shutil.rmtree("pip-out/", ignore_errors=True)
102-
dirs = glob.glob("cmake-out*/") + glob.glob("cmake-android-out/")
103-
for d in dirs:
104-
print(f"Cleaning {d}...")
105-
shutil.rmtree(d, ignore_errors=True)
106-
print("Done cleaning build artifacts.")
107-
sys.exit(0)
108-
elif arg == "--use-pt-pinned-commit":
109-
# This option is used in CI to make sure that PyTorch build from the pinned commit
110-
# is used instead of nightly. CI jobs wouldn't be able to catch regression from the
111-
# latest PT commit otherwise
112-
USE_PYTORCH_NIGHTLY = False
113-
else:
114-
print(f"Error: Unknown option {arg}")
115-
sys.exit(1)
11681

117-
# If --pybind is not set explicitly for backends (e.g., --pybind xnnpack)
118-
# or is not turned off explicitly (--pybind off)
119-
# then install XNNPACK by default.
120-
if EXECUTORCH_BUILD_PYBIND == "":
121-
EXECUTORCH_BUILD_PYBIND = "ON"
122-
CMAKE_ARGS += " -DEXECUTORCH_BUILD_XNNPACK=ON"
82+
# The pip repository that hosts nightly torch packages.
83+
TORCH_NIGHTLY_URL = "https://download.pytorch.org/whl/nightly/cpu"
12384

124-
# Use ClangCL on Windows.
125-
# ClangCL is an alias to Clang that configures it to work in an MSVC-compatible
126-
# mode. Using it on Windows to avoid compiler compatibility issues for MSVC.
127-
if os.name == "nt":
128-
CMAKE_ARGS += " -T ClangCL"
12985

13086
# Since ExecuTorch often uses main-branch features of pytorch, only the nightly
13187
# pip versions will have the required features.
@@ -135,102 +91,180 @@ def python_is_compatible():
13591
# package versions.
13692
NIGHTLY_VERSION = "dev20250104"
13793

138-
# The pip repository that hosts nightly torch packages.
139-
TORCH_NIGHTLY_URL = "https://download.pytorch.org/whl/nightly/cpu"
14094

141-
# pip packages needed by exir.
142-
EXIR_REQUIREMENTS = [
143-
# Setting USE_PYTORCH_NIGHTLY to false to test the pinned PyTorch commit. Note
144-
# that we don't need to set any version number there because they have already
145-
# been installed on CI before this step, so pip won't reinstall them
146-
f"torch==2.6.0.{NIGHTLY_VERSION}" if USE_PYTORCH_NIGHTLY else "torch",
147-
(
148-
f"torchvision==0.22.0.{NIGHTLY_VERSION}"
149-
if USE_PYTORCH_NIGHTLY
150-
else "torchvision"
151-
), # For testing.
152-
"typing-extensions",
153-
]
154-
155-
# pip packages needed to run examples.
156-
# TODO: Make each example publish its own requirements.txt
157-
EXAMPLES_REQUIREMENTS = [
158-
"timm==1.0.7",
159-
f"torchaudio==2.6.0.{NIGHTLY_VERSION}" if USE_PYTORCH_NIGHTLY else "torchaudio",
160-
"torchsr==1.0.4",
161-
"transformers==4.47.1",
162-
]
163-
164-
# pip packages needed for development.
165-
DEVEL_REQUIREMENTS = [
166-
"cmake", # For building binary targets.
167-
"pip>=23", # For building the pip package.
168-
"pyyaml", # Imported by the kernel codegen tools.
169-
"setuptools>=63", # For building the pip package.
170-
"tomli", # Imported by extract_sources.py when using python < 3.11.
171-
"wheel", # For building the pip package archive.
172-
"zstd", # Imported by resolve_buck.py.
173-
]
174-
175-
# Assemble the list of requirements to actually install.
176-
# TODO: Add options for reducing the number of requirements.
177-
REQUIREMENTS_TO_INSTALL = EXIR_REQUIREMENTS + DEVEL_REQUIREMENTS + EXAMPLES_REQUIREMENTS
178-
179-
# Install the requirements. `--extra-index-url` tells pip to look for package
180-
# versions on the provided URL if they aren't available on the default URL.
181-
subprocess.run(
182-
[
183-
sys.executable,
184-
"-m",
185-
"pip",
186-
"install",
187-
*REQUIREMENTS_TO_INSTALL,
188-
"--extra-index-url",
189-
TORCH_NIGHTLY_URL,
190-
],
191-
check=True,
192-
)
193-
194-
LOCAL_REQUIREMENTS = [
195-
"third-party/ao", # We need the latest kernels for fast iteration, so not relying on pypi.
196-
]
197-
198-
# Install packages directly from local copy instead of pypi.
199-
# This is usually not recommended.
200-
subprocess.run(
201-
[
202-
sys.executable,
203-
"-m",
204-
"pip",
205-
"install",
206-
*LOCAL_REQUIREMENTS,
207-
],
208-
check=True,
209-
)
95+
def install_requirements(use_pytorch_nightly):
96+
# pip packages needed by exir.
97+
EXIR_REQUIREMENTS = [
98+
# Setting use_pytorch_nightly to false to test the pinned PyTorch commit. Note
99+
# that we don't need to set any version number there because they have already
100+
# been installed on CI before this step, so pip won't reinstall them
101+
f"torch==2.6.0.{NIGHTLY_VERSION}" if use_pytorch_nightly else "torch",
102+
(
103+
f"torchvision==0.22.0.{NIGHTLY_VERSION}"
104+
if use_pytorch_nightly
105+
else "torchvision"
106+
), # For testing.
107+
"typing-extensions",
108+
]
109+
110+
# pip packages needed to run examples.
111+
# TODO: Make each example publish its own requirements.txt
112+
EXAMPLES_REQUIREMENTS = [
113+
"timm==1.0.7",
114+
f"torchaudio==2.6.0.{NIGHTLY_VERSION}" if use_pytorch_nightly else "torchaudio",
115+
"torchsr==1.0.4",
116+
"transformers==4.47.1",
117+
]
118+
119+
# pip packages needed for development.
120+
DEVEL_REQUIREMENTS = [
121+
"cmake", # For building binary targets.
122+
"pip>=23", # For building the pip package.
123+
"pyyaml", # Imported by the kernel codegen tools.
124+
"setuptools>=63", # For building the pip package.
125+
"tomli", # Imported by extract_sources.py when using python < 3.11.
126+
"wheel", # For building the pip package archive.
127+
"zstd", # Imported by resolve_buck.py.
128+
]
129+
130+
# Assemble the list of requirements to actually install.
131+
# TODO: Add options for reducing the number of requirements.
132+
REQUIREMENTS_TO_INSTALL = (
133+
EXIR_REQUIREMENTS + DEVEL_REQUIREMENTS + EXAMPLES_REQUIREMENTS
134+
)
135+
136+
# Install the requirements. `--extra-index-url` tells pip to look for package
137+
# versions on the provided URL if they aren't available on the default URL.
138+
subprocess.run(
139+
[
140+
sys.executable,
141+
"-m",
142+
"pip",
143+
"install",
144+
*REQUIREMENTS_TO_INSTALL,
145+
"--extra-index-url",
146+
TORCH_NIGHTLY_URL,
147+
],
148+
check=True,
149+
)
150+
151+
LOCAL_REQUIREMENTS = [
152+
"third-party/ao", # We need the latest kernels for fast iteration, so not relying on pypi.
153+
]
154+
155+
# Install packages directly from local copy instead of pypi.
156+
# This is usually not recommended.
157+
subprocess.run(
158+
[
159+
sys.executable,
160+
"-m",
161+
"pip",
162+
"install",
163+
*LOCAL_REQUIREMENTS,
164+
],
165+
check=True,
166+
)
167+
168+
169+
def main(args):
170+
if not python_is_compatible():
171+
sys.exit(1)
172+
173+
# Parse options.
174+
175+
EXECUTORCH_BUILD_PYBIND = ""
176+
CMAKE_ARGS = os.getenv("CMAKE_ARGS", "")
177+
CMAKE_BUILD_ARGS = os.getenv("CMAKE_BUILD_ARGS", "")
178+
use_pytorch_nightly = True
179+
180+
parser = argparse.ArgumentParser()
181+
parser.add_argument(
182+
"--pybind",
183+
action="append",
184+
nargs="+",
185+
help="one or more of coreml/mps/xnnpack, or off",
186+
)
187+
parser.add_argument(
188+
"--clean",
189+
action="store_true",
190+
help="clean build artifacts and pip-out instead of installing",
191+
)
192+
parser.add_argument(
193+
"--use-pt-pinned-commit",
194+
action="store_true",
195+
help="build from the pinned PyTorch commit instead of nightly",
196+
)
197+
args = parser.parse_args(args)
198+
if args.pybind:
199+
# Flatten list of lists.
200+
args.pybind = list(itertools.chain(*args.pybind))
201+
if "off" in args.pybind:
202+
if len(args.pybind) != 1:
203+
raise Exception(
204+
f"Cannot combine `off` with other pybinds: {args.pybind}"
205+
)
206+
EXECUTORCH_BUILD_PYBIND = "OFF"
207+
else:
208+
for pybind_arg in args.pybind:
209+
if pybind_arg not in VALID_PYBINDS:
210+
raise Exception(
211+
f"Unrecognized pybind argument {pybind_arg}; valid options are: {', '.join(VALID_PYBINDS)}"
212+
)
213+
EXECUTORCH_BUILD_PYBIND = "ON"
214+
CMAKE_ARGS += f" -DEXECUTORCH_BUILD_{pybind_arg.upper()}=ON"
215+
216+
if args.clean:
217+
clean()
218+
return
219+
220+
if args.use_pt_pinned_commit:
221+
# This option is used in CI to make sure that PyTorch build from the pinned commit
222+
# is used instead of nightly. CI jobs wouldn't be able to catch regression from the
223+
# latest PT commit otherwise
224+
use_pytorch_nightly = False
225+
226+
install_requirements(use_pytorch_nightly)
227+
228+
# If --pybind is not set explicitly for backends (e.g., --pybind xnnpack)
229+
# or is not turned off explicitly (--pybind off)
230+
# then install XNNPACK by default.
231+
if EXECUTORCH_BUILD_PYBIND == "":
232+
EXECUTORCH_BUILD_PYBIND = "ON"
233+
CMAKE_ARGS += " -DEXECUTORCH_BUILD_XNNPACK=ON"
234+
235+
# Use ClangCL on Windows.
236+
# ClangCL is an alias to Clang that configures it to work in an MSVC-compatible
237+
# mode. Using it on Windows to avoid compiler compatibility issues for MSVC.
238+
if os.name == "nt":
239+
CMAKE_ARGS += " -T ClangCL"
240+
241+
#
242+
# Install executorch pip package. This also makes `flatc` available on the path.
243+
# The --extra-index-url may be necessary if pyproject.toml has a dependency on a
244+
# pre-release or nightly version of a torch package.
245+
#
246+
247+
# Set environment variables
248+
os.environ["EXECUTORCH_BUILD_PYBIND"] = EXECUTORCH_BUILD_PYBIND
249+
os.environ["CMAKE_ARGS"] = CMAKE_ARGS
250+
os.environ["CMAKE_BUILD_ARGS"] = CMAKE_BUILD_ARGS
251+
252+
# Run the pip install command
253+
subprocess.run(
254+
[
255+
sys.executable,
256+
"-m",
257+
"pip",
258+
"install",
259+
".",
260+
"--no-build-isolation",
261+
"-v",
262+
"--extra-index-url",
263+
TORCH_NIGHTLY_URL,
264+
],
265+
check=True,
266+
)
210267

211-
#
212-
# Install executorch pip package. This also makes `flatc` available on the path.
213-
# The --extra-index-url may be necessary if pyproject.toml has a dependency on a
214-
# pre-release or nightly version of a torch package.
215-
#
216268

217-
# Set environment variables
218-
os.environ["EXECUTORCH_BUILD_PYBIND"] = EXECUTORCH_BUILD_PYBIND
219-
os.environ["CMAKE_ARGS"] = CMAKE_ARGS
220-
os.environ["CMAKE_BUILD_ARGS"] = CMAKE_BUILD_ARGS
221-
222-
# Run the pip install command
223-
subprocess.run(
224-
[
225-
sys.executable,
226-
"-m",
227-
"pip",
228-
"install",
229-
".",
230-
"--no-build-isolation",
231-
"-v",
232-
"--extra-index-url",
233-
TORCH_NIGHTLY_URL,
234-
],
235-
check=True,
236-
)
269+
if __name__ == "__main__":
270+
main(sys.argv[1:])

0 commit comments

Comments
 (0)