Add rpath to libcustom_ops_aot_lib.dylib and libquantized_ops_aot_lib.dylib

Summary: As titled. This issue is from https://github.com/pytorch/torchchat/actions/runs/11523122333/job/32080481174?pr=1312

In that job when we try to load `libcustom_ops_aot_lib.dylib` into python,
it complains that it can't find `_portable_lib.cpython-310-darwin.so`.
This PR is trying to fix it by adding the relative path to
`_portable_lib.cpython-310-darwin.so` into `LC_RPATH`.

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:

Author: larryliu0820

extension/llm/custom_ops/CMakeLists.txt

@@ -59,9 +59,7 @@ target_include_directories(custom_ops PUBLIC "${_common_include_directories}")
 target_include_directories(
   custom_ops PRIVATE "${CMAKE_CURRENT_BINARY_DIR}/../../../include"
 )
-target_link_libraries(
-  custom_ops PUBLIC ${custom_ops_libs} executorch_core
-)
+target_link_libraries(custom_ops PUBLIC ${custom_ops_libs} executorch_core)
 
 target_compile_options(
   custom_ops PUBLIC ${_common_compile_options} -DET_USE_THREADPOOL
@@ -74,7 +72,8 @@ if(EXECUTORCH_BUILD_KERNELS_CUSTOM_AOT)
   find_package(Torch CONFIG REQUIRED)
   add_library(
     custom_ops_aot_lib SHARED
-    ${_custom_ops__srcs} ${CMAKE_CURRENT_SOURCE_DIR}/op_sdpa_aot.cpp
+    ${_custom_ops__srcs}
+    ${CMAKE_CURRENT_SOURCE_DIR}/op_sdpa_aot.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/op_fast_hadamard_transform_aten.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/op_tile_crop.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/op_tile_crop_aot.cpp
@@ -110,5 +109,26 @@ if(EXECUTORCH_BUILD_KERNELS_CUSTOM_AOT)
            ${_common_compile_options} -DET_USE_THREADPOOL
   )
 
+  # pip wheels will need to be able to find the dependent libraries. On Linux,
+  # the .so has non-absolute dependencies on libs like "_portable_lib.so"
+  # without paths; as long as we `import torch` first, those dependencies will
+  # work. But Apple dylibs do not support non-absolute dependencies, so we need
+  # to tell the loader where to look for its libraries. The LC_LOAD_DYLIB
+  # entries for the portable_lib libraries will look like
+  # "@rpath/_portable_lib.cpython-310-darwin.so", so we can add an LC_RPATH
+  # entry to look in a directory relative to the installed location of our
+  # _portable_lib.so file. To see these LC_* values, run `otool -l
+  # libcustom_ops_aot_lib.dylib`.
+  if(APPLE)
+    set_target_properties(
+      custom_ops_aot_lib
+      PROPERTIES # Assume this library will be installed in
+                 # <site-packages>/executorch/extension/llm/custom_ops/, and the
+                 # _portable_lib.so is installed in
+                 # <site-packages>/executorch/extension/pybindings/
+                 BUILD_RPATH "@loader_path/../../pybindings"
+                 INSTALL_RPATH "@loader_path/../../pybindings"
+    )
+  endif()
   install(TARGETS custom_ops_aot_lib DESTINATION lib)
 endif()

kernels/quantized/CMakeLists.txt

@@ -114,6 +114,28 @@ if(NOT CMAKE_GENERATOR STREQUAL "Xcode"
       target_link_libraries(
         quantized_ops_aot_lib PUBLIC quantized_ops_pybind_lib
       )
+
+      # pip wheels will need to be able to find the dependent libraries. On
+      # Linux, the .so has non-absolute dependencies on libs like
+      # "_portable_lib.so" without paths; as long as we `import torch` first,
+      # those dependencies will work. But Apple dylibs do not support
+      # non-absolute dependencies, so we need to tell the loader where to look
+      # for its libraries. The LC_LOAD_DYLIB entries for the portable_lib
+      # libraries will look like "@rpath/_portable_lib.cpython-310-darwin.so",
+      # so we can add an LC_RPATH entry to look in a directory relative to the
+      # installed location of our _portable_lib.so file. To see these LC_*
+      # values, run `otool -l libquantized_ops_lib.dylib`.
+      if(APPLE)
+        set_target_properties(
+          quantized_ops_lib
+          PROPERTIES # Assume this library will be installed in
+                     # <site-packages>/executorch/kernels/quantized/, and the
+                     # _portable_lib.so is installed in
+                     # <site-packages>/executorch/extension/pybindings/
+                     BUILD_RPATH "@loader_path/../../extensions/pybindings"
+                     INSTALL_RPATH "@loader_path/../../extensions/pybindings"
+        )
+      endif()
     endif()
   endif()
 endif()