From 0d5f0a81c5766c18970c9b3019e5d3165a2b05f4 Mon Sep 17 00:00:00 2001 From: Nikita Shulga Date: Wed, 29 Jan 2025 17:18:46 -0800 Subject: [PATCH] [CMake] Find HomeBrew OpenMP on MacOS (#145870) Either via `OMP_PREFIX` envvar or by searching in `/opt/homebrew/opt/libomp` folder Modify libomp bundling logic in setup.py to change absolute path to libomp.dylib to a relative one if necessary Pull Request resolved: https://github.com/pytorch/pytorch/pull/145870 Approved by: https://github.com/Skylion007, https://github.com/atalman ghstack dependencies: #145871 --- cmake/Modules/FindOpenMP.cmake | 20 ++++++++++++- setup.py | 51 ++++++++++++++++++++++------------ 2 files changed, 53 insertions(+), 18 deletions(-) diff --git a/cmake/Modules/FindOpenMP.cmake b/cmake/Modules/FindOpenMP.cmake index f465aa63ee7..01cc3a48122 100644 --- a/cmake/Modules/FindOpenMP.cmake +++ b/cmake/Modules/FindOpenMP.cmake @@ -79,6 +79,13 @@ 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 EXISTS /opt/homebrew/opt/libomp) + set(OpenMP_PREFIX "/opt/homebrew/opt/libomp") +endif() + function(_OPENMP_FLAG_CANDIDATES LANG) if(NOT OpenMP_${LANG}_FLAG) unset(OpenMP_FLAG_CANDIDATES) @@ -93,7 +100,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 +280,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 +683,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) diff --git a/setup.py b/setup.py index 3886865acf3..23b201ca9ad 100644 --- a/setup.py +++ b/setup.py @@ -572,35 +572,52 @@ class build_ext(setuptools.command.build_ext.build_ext): assert rpath.startswith("path ") rpaths.append(rpath.split(" ", 1)[1].rsplit("(", 1)[0][:-1]) - omp_lib_name = ( - "libomp.dylib" if os.uname().machine == "arm64" else "libiomp5.dylib" - ) - omp_rpath_lib_path = os.path.join("@rpath", omp_lib_name) - if omp_rpath_lib_path not in libs: + omplib_path = get_cmake_cache_vars()["OpenMP_libomp_LIBRARY"] + omplib_name = get_cmake_cache_vars()["OpenMP_C_LIB_NAMES"] + ".dylib" + omplib_rpath_path = os.path.join("@rpath", omplib_name) + + # This logic is fragile and checks only two cases: + # - libtorch_cpu depends on `@rpath/libomp.dylib`e (happens when built inside miniconda environment) + # - libtorch_cpu depends on `/abs/path/to/libomp.dylib` (happens when built with libomp from homebrew) + if not any(c in libs for c in [omplib_path, omplib_rpath_path]): return # Copy libomp/libiomp5 from rpath locations + target_lib = os.path.join(self.build_lib, "torch", "lib", omplib_name) + libomp_relocated = False for rpath in rpaths: - source_lib = os.path.join(rpath, omp_lib_name) + source_lib = os.path.join(rpath, omplib_name) if not os.path.exists(source_lib): continue - target_lib = os.path.join(self.build_lib, "torch", "lib", omp_lib_name) self.copy_file(source_lib, target_lib) # Delete old rpath and add @loader_lib to the rpath # This should prevent delocate from attempting to package another instance # of OpenMP library in torch wheel as well as loading two libomp.dylib into # the address space, as libraries are cached by their unresolved names - subprocess.check_call( - [ - "install_name_tool", - "-rpath", - rpath, - "@loader_path", - libtorch_cpu_path, - ] - ) + install_name_tool_args = [ + "-rpath", + rpath, + "@loader_path", + ] + libomp_relocated = True break - + if not libomp_relocated and os.path.exists(omplib_path): + self.copy_file(omplib_path, target_lib) + install_name_tool_args = [ + "-change", + omplib_path, + omplib_rpath_path, + ] + if "@loader_path" not in rpaths: + install_name_tool_args += [ + "-add_rpath", + "@loader_path", + ] + libomp_relocated = True + if libomp_relocated: + install_name_tool_args.insert(0, "install_name_tool") + install_name_tool_args.append(libtorch_cpu_path) + subprocess.check_call(install_name_tool_args) # Copy omp.h from OpenMP_C_FLAGS and copy it into include folder omp_cflags = get_cmake_cache_vars()["OpenMP_C_FLAGS"] if not omp_cflags: