Enable prioritized linker optimization for AArch64 in setup.py and clean up CI script #160078
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/160078
Note: Links to docs will display an error until the docs builds have been completed. ✅ You can merge normally! (3 Unrelated Failures)As of commit 62e5e8c with merge base 9a680e1 ( UNSTABLE - The following jobs are marked as unstable, possibly due to flakiness on trunk:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
|
@pytorchbot label "ciflow/linux-aarch64" |
|
@pytorchbot label "topic: not user facing" |
|
@pytorchbot label "ciflow/linux-aarch64" |
| @@ -1243,12 +1243,27 @@ def main() -> None: | |||
| if BUILD_PYTHON_ONLY: | |||
| install_requires += [f"{LIBTORCH_PKG_NAME}=={TORCH_VERSION}"] | |||
|
|
|||
| if ( | |||
| os.getenv("USE_PRIORITIZED_TEXT_FOR_LD") is None | |||
There was a problem hiding this comment.
Shouldn't this logic live in the CMake?
There was a problem hiding this comment.
Hmm so we want to get CMake to call this python function
https://github.com/pytorch/pytorch/blob/main/tools/setup_helpers/generate_linker_script.py before the build if a certain platform and OS is set.
https://cmake.org/cmake/help/latest/command/add_custom_command.html#build-events
I'll investigate if it can be done.
| @@ -1258,7 +1273,7 @@ def main() -> None: | |||
| elif platform.system() == "Linux" and platform.processor() == "aarch64": | |||
There was a problem hiding this comment.
Just looking at this in more detail, this condition will never be True because of line 1251, so we can delete 1273-1279.
Note. This is a replica PR of #155901 which will be closed. I had to create a new PR in order to add it into my ghstack as there are some later commits which depend on it.
Summary
🚀 This PR enables prioritized text linker optimization by default on Linux aarch64 systems via logic added to setup.py.
If USE_PRIORITIZED_TEXT_FOR_LD is not already set, the build system will now automatically:
• Set it to 1 on Linux + aarch64
• Inject appropriate linker and compiler flags via LDFLAGS, CFLAGS, and CXXFLAGS
This change consolidates what was previously manual CI logic into a single location (setup.py), ensuring consistent behavior across local builds, CI pipelines, and developer environments.
Motivation
Prioritized text layout has measurable performance benefits on Arm systems by reducing code padding and improving cache utilization. This optimization was previously triggered manually via CI scripts (.ci/aarch64_linux/aarch64_ci_build.sh) or user-set environment variables. By detecting the target architecture within setup.py, this change enables the optimization automatically where applicable, improving maintainability and usability.
Stack from ghstack (oldest at bottom):
Co-authored-by: Usamah Zaheer usamah.zaheer@arm.com
cc @malfet @snadampal @milpuz01 @aditew01 @nikhil-arm @fadara01 @usamahz