[pybind] Fix rpath for custom ops

Summary: This is a followup to #7096. That PR was aiming to rely on
`CMAKE_INSTALL_RPATH_USE_LINK_PATH` to set RPATH automatically. However
due to the caveat that we are not installing `_portable_lib.so` into the
correct path in CMake stage (we move the .so in setup.py after it's
installed), we are not able to add the correct RPATH to
libcustom_ops_aot_lib.so.

This PR still set the INSTALL_RPATH manually and adds TODOs to fix it
later.

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:

Author: larryliu0820

CMakeLists.txt

@@ -56,6 +56,21 @@ if(NOT CMAKE_BUILD_TYPE)
   set(CMAKE_BUILD_TYPE Debug)
 endif()
 
+# Setup RPATH.
+# See https://gitlab.kitware.com/cmake/community/-/wikis/doc/cmake/RPATH-handling
+# Use separate rpaths during build and install phases
+set(CMAKE_SKIP_BUILD_RPATH OFF)
+# Don't use the install-rpath during the build phase
+set(CMAKE_BUILD_WITH_INSTALL_RPATH ON)
+# Automatically add all linked folders that are NOT in the build directory to
+# the rpath (per library?)
+# TODO: Doesn't work for us right now because we are not installing .so's into the
+# correct locations. For example we have libcustom_ops_aot_lib.so depending on
+# _portable_lib.so, which was eventually put under <site-packages>/executorch/extension/pybindings/
+# but this rpath is not automatically added because at build time it seems `portable_lib`
+# is being built under the same directory, so no extra rpath is being added. To
+# properly fix this we need to install `portable_lib` into the correct path.
+set(CMAKE_INSTALL_RPATH_USE_LINK_PATH ON)
 # ------------------------------ OPTIONS -------------------------------------
 # WARNING: Please don't add example specific options in this CMakeLists.txt.
 # Instead please use `find_package(executorch REQUIRED)` in the example
@@ -682,22 +697,6 @@ if(EXECUTORCH_BUILD_PTHREADPOOL
 endif()
 
 if(EXECUTORCH_BUILD_PYBIND)
-  # Setup RPATH.
-  # See https://gitlab.kitware.com/cmake/community/-/wikis/doc/cmake/RPATH-handling
-  if(APPLE)
-    set(CMAKE_MACOSX_RPATH ON)
-    set(_rpath_portable_origin "@loader_path")
-  else()
-    set(_rpath_portable_origin $ORIGIN)
-  endif(APPLE)
-  # Use separate rpaths during build and install phases
-  set(CMAKE_SKIP_BUILD_RPATH  FALSE)
-  # Don't use the install-rpath during the build phase
-  set(CMAKE_BUILD_WITH_INSTALL_RPATH FALSE)
-  set(CMAKE_INSTALL_RPATH "${_rpath_portable_origin}")
-  # Automatically add all linked folders that are NOT in the build directory to
-  # the rpath (per library?)
-  set(CMAKE_INSTALL_RPATH_USE_LINK_PATH TRUE)
   add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/third-party/pybind11)
 
   if(NOT EXECUTORCH_BUILD_EXTENSION_DATA_LOADER)

extension/llm/custom_ops/CMakeLists.txt

@@ -84,13 +84,21 @@ if(EXECUTORCH_BUILD_KERNELS_CUSTOM_AOT)
   target_include_directories(
     custom_ops_aot_lib PRIVATE "${CMAKE_CURRENT_BINARY_DIR}/../../../include"
   )
+  # TODO: This only works if we install portable_lib.so to
+  # <site-packages>/executorch/extension/pybindings/.
+  if(APPLE)
+    set(RPATH "@loader_path/../../pybindings")
+  else()
+    set(RPATH "$ORIGIN/../../pybindings")
+  endif()
+  set_target_properties(custom_ops_aot_lib PROPERTIES INSTALL_RPATH ${RPATH})
   if(TARGET portable_lib)
     # If we have portable_lib built, custom_ops_aot_lib gives the ability to use
     # the ops in PyTorch and ExecuTorch through pybind
     target_link_libraries(custom_ops_aot_lib PUBLIC portable_lib)
   else()
     # If no portable_lib, custom_ops_aot_lib still gives the ability to use the
-    # ops in PyTorch
+    # ops in PyTorchk
     target_link_libraries(custom_ops_aot_lib PUBLIC executorch_core)
   endif()
 
@@ -109,5 +117,7 @@ if(EXECUTORCH_BUILD_KERNELS_CUSTOM_AOT)
            ${_common_compile_options} -DET_USE_THREADPOOL
   )
 
-  install(TARGETS custom_ops_aot_lib DESTINATION lib)
+  install(TARGETS custom_ops_aot_lib
+          LIBRARY DESTINATION executorch/extension/llm/custom_ops
+  )
 endif()

kernels/quantized/CMakeLists.txt

@@ -126,16 +126,14 @@ if(NOT CMAKE_GENERATOR STREQUAL "Xcode"
       # 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"
-        )
+        set(RPATH "@loader_path/../../extensions/pybindings")
+      else()
+        set(RPATH "$ORIGIN/../../extensions/pybindings")
       endif()
+      set_target_properties(
+        quantized_ops_aot_lib PROPERTIES BUILD_RPATH ${RPATH} INSTALL_RPATH
+                                                              ${RPATH}
+      )
     endif()
   endif()
 endif()