frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-20 21:14:14 +08:00
Code Issues Packages Projects Releases Wiki Activity

Move prioritized text linker optimization code from setup.py to cmake (#160078)

Browse Source 
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 moves the prioritized text linker optimization from setup.py to cmake ( and enables by default on Linux aarch64 systems )

This change consolidates what was previously manual CI logic into a single location (cmake), 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.

Note:

Due to ninja/cmake graph generation issues we cannot apply the linker file globally to all targets to the targets must be manually defined. See CMakeLists.txt the main libraries torch_python, torch, torch_cpu, torch_cuda, torch_xpu have been targetted which should be enough to maintain the performance benefits outlined above.

Co-authored-by: Usamah Zaheer <usamah.zaheer@arm.com>

Pull Request resolved: https://github.com/pytorch/pytorch/pull/160078
Approved by: https://github.com/seemethere
This commit is contained in:
Robert Hardwick
2025-09-18 08:42:03 +00:00
committed by PyTorch MergeBot
parent 56893ca1f6
commit 1aeac304b8
8 changed files with 116 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
cmake/Summary.cmake
View File

@ -158,6 +158,7 @@ function(caffe2_print_configuration_summary)
if(${USE_KLEIDIAI})
message(STATUS " USE_KLEIDIAI : ${USE_KLEIDIAI}")
endif()
message(STATUS " USE_PRIORITIZED_TEXT_FOR_LD : ${USE_PRIORITIZED_TEXT_FOR_LD}")
message(STATUS " USE_UCC : ${USE_UCC}")
if(${USE_UCC})
message(STATUS " USE_SYSTEM_UCC : ${USE_SYSTEM_UCC}")

20
cmake/public/utils.cmake
View File

@ -482,6 +482,7 @@ function(torch_update_find_cuda_flags)
endfunction()
include(CheckCXXCompilerFlag)
include(CheckLinkerFlag)
##############################################################################
# CHeck if given flag is supported and append it to provided outputvar
@ -511,3 +512,22 @@ function(target_compile_options_if_supported target flag)
target_compile_options(${target} PRIVATE ${flag})
endif()
endfunction()
# Check if a global link option is supported
function(add_link_options_if_supported flag)
check_linker_flag(C "LINKER:${flag}" _supported)
if("${_supported}")
add_link_options("LINKER:${flag}")
else()
message(WARNING "Attempted to use unsupported link option : ${flag}.")
endif()
endfunction()
function(target_link_options_if_supported tgt flag)
check_linker_flag(C "LINKER:${flag}" _supported)
if("${_supported}")
target_link_options("${tgt}" PRIVATE "LINKER:${flag}")
else()
message(WARNING "Attempted to use unsupported link option : ${flag}.")
endif()
endfunction()