[CMake] Find HomeBrew OpenMP on MacOS (#145870)

Either via `OMP_PREFIX` envvar or just searching in that folder
Pull Request resolved: https://github.com/pytorch/pytorch/pull/145870
Approved by: https://github.com/Skylion007
This commit is contained in:
Nikita Shulga 2025-01-28 13:14:14 -08:00 committed by PyTorch MergeBot
parent f388ba5986
commit c26bb9ba5b

View File

@ -79,6 +79,14 @@ cmake_policy(SET CMP0012 NEW) # if() recognizes numbers and booleans
cmake_policy(SET CMP0054 NEW) # if() quoted variables not dereferenced
cmake_policy(SET CMP0057 NEW) # if IN_LIST
if(NOT "$ENV{OMP_PREFIX}" STREQUAL "")
set(OpenMP_PREFIX "$ENV{OMP_PREFIX}")
elseif(${CMAKE_SYSTEM_NAME} STREQUAL "Darwin" AND IS_DIRECTORY /opt/homebrew/Cellar/libomp)
file(GLOB HOMEBREW_OMP_VERSIONS "/opt/homebrew/Cellar/libomp/*")
list(GET HOMEBREW_OMP_VERSIONS 0 OpenMP_PREFIX)
endif()
function(_OPENMP_FLAG_CANDIDATES LANG)
if(NOT OpenMP_${LANG}_FLAG)
unset(OpenMP_FLAG_CANDIDATES)
@ -93,7 +101,7 @@ function(_OPENMP_FLAG_CANDIDATES LANG)
else()
# AppleClang may need a header file, search for omp.h with hints to brew
# default include dir
find_path(__header_dir "omp.h" HINTS "/usr/local/include")
find_path(__header_dir "omp.h" HINTS "/usr/local/include" "${OpenMP_PREFIX}/include")
endif()
set(OMP_FLAG_AppleClang "-Xpreprocessor -fopenmp" "-Xpreprocessor -fopenmp -I${__header_dir}")
@ -273,6 +281,16 @@ function(_OPENMP_GET_FLAGS LANG FLAG_MODE OPENMP_FLAG_VAR OPENMP_LIB_NAMES_VAR)
mark_as_advanced(OpenMP_libomp_LIBRARY)
endif()
# Use OpenMP_PREFIX if defined
if (NOT OpenMP_libomp_LIBRARY AND NOT "${OpenMP_PREFIX}" STREQUAL "")
find_library(OpenMP_libomp_LIBRARY
NAMES omp gomp iomp5
HINTS "${OpenMP_PREFIX}/lib"
DOC "libomp location for OpenMP"
)
mark_as_advanced(OpenMP_libomp_LIBRARY)
endif()
if(OpenMP_libomp_LIBRARY MATCHES "iomp5")
set(OpenMP_libiomp5_LIBRARY "${MKL_OPENMP_LIBRARY}" CACHE STRING "libiomp5 location for OpenMP")
if("-fopenmp=libiomp5" IN_LIST OpenMP_${LANG}_FLAG_CANDIDATES)
@ -666,5 +684,6 @@ unset(OpenMP_Fortran_TEST_SOURCE)
unset(OpenMP_C_CXX_CHECK_VERSION_SOURCE)
unset(OpenMP_Fortran_CHECK_VERSION_SOURCE)
unset(OpenMP_Fortran_INCLUDE_LINE)
unset(OpenMP_PREFIX)
cmake_policy(POP)