Add rpath to libcustom_ops_aot_lib.dylib and libquantized_ops_aot_lib… (#6499)

* 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:

* Add smoke test

* Lint

Author: larryliu0820

build/packaging/smoke_test.py

@@ -15,6 +15,14 @@
 # will fail and the process will exit.
 from executorch.extension.pybindings import portable_lib  # usort: skip
 
+# Import custom ops. This requires portable_lib to be loaded first.
+from executorch.extension.llm.custom_ops import (  # noqa: F401, F403
+    sdpa_with_kv_cache,
+)  # usort: skip
+
+# Import quantized ops. This requires portable_lib to be loaded first.
+from executorch.kernels import quantized  # usort: skip # noqa: F401, F403
+
 # Import this after importing the ExecuTorch pybindings. If the pybindings
 # links against a different torch.so than this uses, there will be a set of
 # symbol comflicts; the process will either exit now, or there will be issues
@@ -75,6 +83,15 @@ def main():
     assert len(ops) > 0, "Empty operator list"
     print(f"Found {len(ops)} operators; first element '{ops[0]}'")
 
+    # Make sure custom ops are registered.
+    assert (
+        "llama::sdpa_with_kv_cache" in ops
+    ), f"sdpa_with_kv_cache not registered, Got ops: {ops}"
+
+    # Make sure quantized ops are registered.
+    assert (
+        "quantized_decomposed::add.out" in ops
+    ), f"quantized_decomposed::add.out not registered, Got ops: {ops}"
     # Export LinearModel to .pte data.
     pte_data: bytes = export_linear_model()
 

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_aot_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()