Skip to content
Open
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
12 changes: 10 additions & 2 deletions CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -602,15 +602,23 @@ if(EXECUTORCH_BUILD_CORTEX_M)
list(APPEND _executorch_backends coretex_m_backend)
endif()

if(EXECUTORCH_BUILD_CUDA)
# Build common AOTI functionality (required for CUDA)
# Build common AOTI functionality if needed by CUDA or Metal backends
if(EXECUTORCH_BUILD_CUDA OR EXECUTORCH_BUILD_METAL)
add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/backends/aoti)
endif()

if(EXECUTORCH_BUILD_CUDA)
# Build CUDA-specific AOTI functionality
add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/backends/cuda)
# Add aoti_cuda to backends - it already depends on aoti_common
list(APPEND _executorch_backends aoti_cuda)
endif()

if(EXECUTORCH_BUILD_METAL)
add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/backends/apple/metal)
list(APPEND _executorch_backends metal_backend)
endif()

if(EXECUTORCH_BUILD_EXTENSION_APPLE)
add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/extension/apple)
endif()
Expand Down
6 changes: 5 additions & 1 deletion backends/aoti/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,11 @@ target_include_directories(
)
target_compile_options(aoti_common PUBLIC -fexceptions -frtti -fPIC)
# Ensure symbols are exported properly
target_link_options(aoti_common PUBLIC -Wl,--export-dynamic)
if(APPLE)
target_link_options(aoti_common PUBLIC -Wl,-export_dynamic)
else()
target_link_options(aoti_common PUBLIC -Wl,--export-dynamic)
endif()

# Link against ExecuTorch libraries and standard libraries
target_link_libraries(aoti_common PUBLIC extension_tensor ${CMAKE_DL_LIBS})
Expand Down
111 changes: 111 additions & 0 deletions backends/apple/metal/CMakeLists.txt
Original file line number Diff line number Diff line change
@@ -0,0 +1,111 @@
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.
#
# Build AOTI Metal backend for runtime.
#
# ### Editing this file ###
#
# This file should be formatted with
# ~~~
# cmake-format -i CMakeLists.txt
# ~~~
# It should also be cmake-lint clean.
#
set(CMAKE_EXPORT_COMPILE_COMMANDS ON)

# Source root directory for executorch.
if(NOT EXECUTORCH_ROOT)
set(EXECUTORCH_ROOT ${CMAKE_CURRENT_SOURCE_DIR}/../..)
endif()

include(${EXECUTORCH_ROOT}/tools/cmake/Utils.cmake)
# Use full torch package to get library paths, but only link specific libraries
find_package_torch()

set(_aoti_metal_sources
runtime/metal_backend.cpp
runtime/shims/memory.cpp
runtime/shims/et_metal.mm
runtime/shims/et_metal_ops.mm
runtime/shims/shim_mps.mm
runtime/shims/tensor_attribute.cpp
runtime/shims/utils.cpp
)

add_library(metal_backend STATIC ${_aoti_metal_sources})
target_include_directories(
metal_backend
PUBLIC $<BUILD_INTERFACE:${EXECUTORCH_ROOT}> $<INSTALL_INTERFACE:include>
# PyTorch AOTI headers from ExecutorTorch's torch detection
${TORCH_INCLUDE_DIRS}
)

# Link Metal framework
find_library(METAL_LIBRARY Metal REQUIRED)
find_library(FOUNDATION_LIBRARY Foundation REQUIRED)
find_library(METALPERFORMANCESHADERS_LIBRARY MetalPerformanceShaders REQUIRED)
find_library(
METALPERFORMANCESHADERSGRAPH_LIBRARY MetalPerformanceShadersGraph REQUIRED
)
target_link_libraries(
metal_backend
PUBLIC ${METAL_LIBRARY} ${FOUNDATION_LIBRARY}
${METALPERFORMANCESHADERS_LIBRARY}
${METALPERFORMANCESHADERSGRAPH_LIBRARY}
)

target_compile_options(metal_backend PUBLIC -fexceptions -frtti -fPIC)

target_link_options(metal_backend PUBLIC -Wl,-export_dynamic)

# Find PyTorch's OpenMP library specifically for libtorch-less AOTI
get_torch_base_path(TORCH_BASE_PATH)
find_library(
TORCH_OMP_LIBRARY
NAMES omp libomp
PATHS "${TORCH_BASE_PATH}/lib"
NO_DEFAULT_PATH
)

if(TORCH_OMP_LIBRARY)
message(STATUS "Found PyTorch OpenMP library: ${TORCH_OMP_LIBRARY}")
# Get the directory containing the OpenMP library for rpath
get_filename_component(TORCH_OMP_LIB_DIR ${TORCH_OMP_LIBRARY} DIRECTORY)
message(STATUS "OpenMP library directory: ${TORCH_OMP_LIB_DIR}")
else()
message(
WARNING "PyTorch OpenMP library not found, may cause runtime linking issues"
)
endif()

# Link against appropriate backends and standard libraries
target_link_libraries(
metal_backend PUBLIC aoti_common extension_tensor ${CMAKE_DL_LIBS}
${TORCH_OMP_LIBRARY}
)

# Set rpath for OpenMP library to avoid runtime linking issues
if(TORCH_OMP_LIBRARY AND TORCH_OMP_LIB_DIR)
# Add the OpenMP library directory to the rpath
set_target_properties(
metal_backend PROPERTIES BUILD_RPATH "${TORCH_OMP_LIB_DIR}"
INSTALL_RPATH "${TORCH_OMP_LIB_DIR}"
)
# Also try common OpenMP library locations
target_link_options(
metal_backend PUBLIC -Wl,-rpath,${TORCH_OMP_LIB_DIR}
-Wl,-rpath,/usr/local/opt/libomp/lib
-Wl,-rpath,/opt/homebrew/opt/libomp/lib
)
message(STATUS "Added rpath for OpenMP library: ${TORCH_OMP_LIB_DIR}")
endif()

executorch_target_link_options_shared_lib(metal_backend)
install(
TARGETS metal_backend
EXPORT ExecuTorchTargets
DESTINATION lib
)
Loading
Loading