[Wheel] Change libtorch_cpu OpenMP search path (#123417) (#123442)

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

Co-authored-by: Nikita Shulga <nikita.shulga@gmail.com>

setup.py

@@ -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
@@ -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