Skip to content

Commit

Permalink
[Wheel] Change libtorch_cpu OpenMP search path (#123417)
Browse files Browse the repository at this point in the history
To prevent delocate from double-packing it, which makes Torch wheels
unusable with torch.compile out of the box

Fixes #122705

Pull Request resolved: #123417
Approved by: https://github.com/atalman
  • Loading branch information
malfet authored and pytorchmergebot committed Apr 5, 2024
1 parent cfd06bd commit 5b0ce8f
Showing 1 changed file with 17 additions and 1 deletion.
18 changes: 17 additions & 1 deletion setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -548,7 +548,9 @@ def _embed_libomp(self):
omp_lib_name = (
"libomp.dylib" if os.uname().machine == "arm64" else "libiomp5.dylib"
)
if os.path.join("@rpath", omp_lib_name) not in libs:
omp_rpath_lib_path = os.path.join("@rpath", omp_lib_name)
omp_loader_lib_path = os.path.join("@loader_path", omp_lib_name)
if omp_rpath_lib_path not in libs:
return

# Copy libomp/libiomp5 from rpath locations
Expand All @@ -558,6 +560,20 @@ def _embed_libomp(self):
continue
target_lib = os.path.join(self.build_lib, "torch", "lib", omp_lib_name)
self.copy_file(source_lib, target_lib)
# Change OMP library load path to loader_path and delete old rpath
# This should prevent delocate from attempting to package another instance
# of OpenMP library in torch wheel
subprocess.check_call(
[
"install_name_tool",
"-change",
omp_rpath_lib_path,
omp_loader_lib_path,
"-delete_rpath",
rpath,
libtorch_cpu_path,
]
)
break

# Copy omp.h from OpenMP_C_FLAGS and copy it into include folder
Expand Down

0 comments on commit 5b0ce8f

Please sign in to comment.