Skip to content

Let custom_ops_aot_lib use portable_lib #4024

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 1 commit into from
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
29 changes: 9 additions & 20 deletions CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -519,13 +519,6 @@ target_link_options_shared_lib(executorch)
#
add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/kernels/portable)

if(EXECUTORCH_BUILD_KERNELS_CUSTOM)
# TODO: move all custom kernels to ${CMAKE_CURRENT_SOURCE_DIR}/kernels/custom
add_subdirectory(
${CMAKE_CURRENT_SOURCE_DIR}/examples/models/llama2/custom_ops
)
endif()

if(EXECUTORCH_BUILD_KERNELS_OPTIMIZED)
add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/kernels/optimized)
endif()
Expand Down Expand Up @@ -688,17 +681,6 @@ if(EXECUTORCH_BUILD_PYBIND)
list(APPEND _dep_libs quantized_kernels quantized_ops_lib)
endif()

# TODO(larryliu): Fix macOS 2 dylibs having 2 sets of static variables issue
if(EXECUTORCH_BUILD_KERNELS_CUSTOM_AOT AND NOT APPLE)
list(APPEND _dep_libs custom_ops_aot_lib)
endif()
# TODO(laryliu): Fix linux duplicate registation problem. In GH CI worker
# libcustom_ops.a doesn't dedup with the one indirectly linked from
# libcustom_ops_aot_lib.a
if(EXECUTORCH_BUILD_KERNELS_CUSTOM AND APPLE)
target_link_options_shared_lib(custom_ops)
list(APPEND _dep_libs custom_ops)
endif()
# compile options for pybind
set(_pybind_compile_options
-Wno-deprecated-declarations
Expand All @@ -723,7 +705,7 @@ if(EXECUTORCH_BUILD_PYBIND)
target_link_libraries(util PRIVATE torch c10 executorch)

# pybind portable_lib
pybind11_add_module(portable_lib extension/pybindings/pybindings.cpp)
pybind11_add_module(portable_lib SHARED extension/pybindings/pybindings.cpp)
# The actual output file needs a leading underscore so it can coexist with
# portable_lib.py in the same python package.
set_target_properties(portable_lib PROPERTIES OUTPUT_NAME "_portable_lib")
Expand All @@ -732,7 +714,7 @@ if(EXECUTORCH_BUILD_PYBIND)
)
target_include_directories(portable_lib PRIVATE ${TORCH_INCLUDE_DIRS})
target_compile_options(portable_lib PUBLIC ${_pybind_compile_options})
target_link_libraries(portable_lib PUBLIC ${_dep_libs})
target_link_libraries(portable_lib PRIVATE ${_dep_libs})
if(APPLE)
# pip wheels will need to be able to find the torch libraries. On Linux, the
# .so has non-absolute dependencies on libs like "libtorch.so" without
Expand All @@ -758,5 +740,12 @@ if(EXECUTORCH_BUILD_PYBIND)
)
endif()

if(EXECUTORCH_BUILD_KERNELS_CUSTOM)
# TODO: move all custom kernels to ${CMAKE_CURRENT_SOURCE_DIR}/kernels/custom
add_subdirectory(
${CMAKE_CURRENT_SOURCE_DIR}/examples/models/llama2/custom_ops
)
endif()

# Print all summary
executorch_print_configuration_summary()
32 changes: 20 additions & 12 deletions examples/models/llama2/custom_ops/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -40,13 +40,7 @@ include(${EXECUTORCH_SRCS_FILE})
set(_common_include_directories ${EXECUTORCH_ROOT}/..)

# Custom op libraries
set(custom_ops_libs)
if(APPLE)
list(APPEND executorch_no_prim_ops_shared)
else()
list(APPEND executorch_no_prim_ops)
endif()
list(APPEND custom_ops_libs pthreadpool)
set(custom_ops_libs pthreadpool)
list(APPEND custom_ops_libs cpuinfo)
list(APPEND custom_ops_libs cpublas)
list(APPEND custom_ops_libs eigen_blas)
Expand All @@ -72,7 +66,9 @@ 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})
target_link_libraries(
custom_ops PUBLIC ${custom_ops_libs} executorch_no_prim_ops
)

target_compile_options(
custom_ops PUBLIC ${_common_compile_options} -DET_USE_THREADPOOL
Expand All @@ -84,7 +80,8 @@ if(EXECUTORCH_BUILD_KERNELS_CUSTOM_AOT)
# Add a AOT library
find_package(Torch CONFIG REQUIRED)
add_library(
custom_ops_aot_lib SHARED ${CMAKE_CURRENT_SOURCE_DIR}/op_sdpa_aot.cpp
custom_ops_aot_lib SHARED ${_custom_ops__srcs}
${CMAKE_CURRENT_SOURCE_DIR}/op_sdpa_aot.cpp
)
target_include_directories(
custom_ops_aot_lib PUBLIC "${_common_include_directories}"
Expand All @@ -93,10 +90,21 @@ if(EXECUTORCH_BUILD_KERNELS_CUSTOM_AOT)
custom_ops_aot_lib
PRIVATE "${CMAKE_CURRENT_BINARY_DIR}/../../../../include"
)
target_link_libraries(custom_ops_aot_lib PUBLIC custom_ops torch)
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
target_link_libraries(custom_ops_aot_lib PUBLIC executorch_no_prim_ops)
endif()

target_link_libraries(custom_ops_aot_lib PUBLIC cpublas torch)
target_compile_options(
custom_ops_aot_lib PUBLIC -Wno-deprecated-declarations -fPIC -frtti
-fexceptions
custom_ops_aot_lib
PUBLIC -Wno-deprecated-declarations -fPIC -frtti -fexceptions
${_common_compile_options} -DET_USE_THREADPOOL
)

install(TARGETS custom_ops_aot_lib DESTINATION lib)
Expand Down
Loading