Skip to content

Commit 06ea3d6

Browse files
Add Metal backend build system and runtime integration (pytorch#15024)
Build system changes: - Add EXECUTORCH_BUILD_METAL build option in default.cmake - Fix AOTI linker flags for Apple platforms (use -export_dynamic) - Extend AOTI build support to Metal backend - Register Metal backend and configure Metal framework linkage - Add PyTorch AOTI headers and OpenMP library detection with rpath Runtime implementation: - Add main Metal backend runtime providing entry point for AOTI-compiled model execution on Metal devices This commit ties together all Metal backend components and enables building the complete backend.
1 parent 2ff74d0 commit 06ea3d6

File tree

7 files changed

+696
-5
lines changed

7 files changed

+696
-5
lines changed

CMakeLists.txt

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -605,15 +605,23 @@ if(EXECUTORCH_BUILD_CORTEX_M)
605605
list(APPEND _executorch_backends coretex_m_backend)
606606
endif()
607607

608-
if(EXECUTORCH_BUILD_CUDA)
609-
# Build common AOTI functionality (required for CUDA)
608+
# Build common AOTI functionality if needed by CUDA or Metal backends
609+
if(EXECUTORCH_BUILD_CUDA OR EXECUTORCH_BUILD_METAL)
610610
add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/backends/aoti)
611+
endif()
612+
613+
if(EXECUTORCH_BUILD_CUDA)
611614
# Build CUDA-specific AOTI functionality
612615
add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/backends/cuda)
613616
# Add aoti_cuda to backends - it already depends on aoti_common
614617
list(APPEND _executorch_backends aoti_cuda)
615618
endif()
616619

620+
if(EXECUTORCH_BUILD_METAL)
621+
add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/backends/apple/metal)
622+
list(APPEND _executorch_backends metal_backend)
623+
endif()
624+
617625
if(EXECUTORCH_BUILD_EXTENSION_APPLE)
618626
add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/extension/apple)
619627
endif()

backends/aoti/CMakeLists.txt

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -42,9 +42,13 @@ target_compile_options(
4242
$<$<NOT:$<CXX_COMPILER_ID:MSVC>>:-fexceptions -frtti -fPIC>
4343
)
4444
# Ensure symbols are exported properly
45-
target_link_options(
46-
aoti_common PUBLIC $<$<NOT:$<CXX_COMPILER_ID:MSVC>>:-Wl,--export-dynamic>
47-
)
45+
if(APPLE)
46+
target_link_options(aoti_common PUBLIC -Wl,-export_dynamic)
47+
else()
48+
target_link_options(
49+
aoti_common PUBLIC $<$<NOT:$<CXX_COMPILER_ID:MSVC>>:-Wl,--export-dynamic>
50+
)
51+
endif()
4852

4953
# Link against ExecuTorch libraries and standard libraries
5054
target_link_libraries(aoti_common PUBLIC extension_tensor ${CMAKE_DL_LIBS})
Lines changed: 120 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,120 @@
1+
# Copyright (c) Meta Platforms, Inc. and affiliates.
2+
# All rights reserved.
3+
#
4+
# This source code is licensed under the BSD-style license found in the
5+
# LICENSE file in the root directory of this source tree.
6+
#
7+
# Build AOTI Metal backend for runtime.
8+
#
9+
# ### Editing this file ###
10+
#
11+
# This file should be formatted with
12+
# ~~~
13+
# cmake-format -i CMakeLists.txt
14+
# ~~~
15+
# It should also be cmake-lint clean.
16+
#
17+
cmake_minimum_required(VERSION 3.29)
18+
19+
set(CMAKE_CXX_STANDARD 17)
20+
set(CMAKE_CXX_STANDARD_REQUIRED ON)
21+
22+
set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
23+
24+
if(NOT APPLE)
25+
message(FATAL_ERROR "Metal backend requires macOS")
26+
endif()
27+
28+
# Source root directory for executorch.
29+
if(NOT EXECUTORCH_ROOT)
30+
set(EXECUTORCH_ROOT ${CMAKE_CURRENT_SOURCE_DIR}/../..)
31+
endif()
32+
33+
include(${EXECUTORCH_ROOT}/tools/cmake/Utils.cmake)
34+
# Use full torch package to get library paths, but only link specific libraries
35+
find_package_torch()
36+
37+
set(_aoti_metal_sources
38+
runtime/metal_backend.cpp
39+
runtime/shims/memory.cpp
40+
runtime/shims/et_metal.mm
41+
runtime/shims/et_metal_ops.mm
42+
runtime/shims/shim_mps.mm
43+
runtime/shims/tensor_attribute.cpp
44+
runtime/shims/utils.cpp
45+
)
46+
47+
add_library(metal_backend STATIC ${_aoti_metal_sources})
48+
target_include_directories(
49+
metal_backend
50+
PUBLIC $<BUILD_INTERFACE:${EXECUTORCH_ROOT}> $<INSTALL_INTERFACE:include>
51+
# PyTorch AOTI headers from ExecutorTorch's torch detection
52+
${TORCH_INCLUDE_DIRS}
53+
)
54+
55+
# Link Metal framework
56+
find_library(METAL_LIBRARY Metal REQUIRED)
57+
find_library(FOUNDATION_LIBRARY Foundation REQUIRED)
58+
find_library(METALPERFORMANCESHADERS_LIBRARY MetalPerformanceShaders REQUIRED)
59+
find_library(
60+
METALPERFORMANCESHADERSGRAPH_LIBRARY MetalPerformanceShadersGraph REQUIRED
61+
)
62+
target_link_libraries(
63+
metal_backend
64+
PUBLIC ${METAL_LIBRARY} ${FOUNDATION_LIBRARY}
65+
${METALPERFORMANCESHADERS_LIBRARY}
66+
${METALPERFORMANCESHADERSGRAPH_LIBRARY}
67+
)
68+
69+
target_compile_options(metal_backend PUBLIC -fexceptions -frtti -fPIC)
70+
71+
target_link_options(metal_backend PUBLIC -Wl,-export_dynamic)
72+
73+
# Find PyTorch's OpenMP library specifically for libtorch-less AOTI
74+
get_torch_base_path(TORCH_BASE_PATH)
75+
find_library(
76+
TORCH_OMP_LIBRARY
77+
NAMES omp libomp
78+
PATHS "${TORCH_BASE_PATH}/lib"
79+
NO_DEFAULT_PATH
80+
)
81+
82+
if(TORCH_OMP_LIBRARY)
83+
message(STATUS "Found PyTorch OpenMP library: ${TORCH_OMP_LIBRARY}")
84+
# Get the directory containing the OpenMP library for rpath
85+
get_filename_component(TORCH_OMP_LIB_DIR ${TORCH_OMP_LIBRARY} DIRECTORY)
86+
message(STATUS "OpenMP library directory: ${TORCH_OMP_LIB_DIR}")
87+
else()
88+
message(
89+
WARNING "PyTorch OpenMP library not found, may cause runtime linking issues"
90+
)
91+
endif()
92+
93+
# Link against appropriate backends and standard libraries
94+
target_link_libraries(
95+
metal_backend PUBLIC aoti_common extension_tensor ${CMAKE_DL_LIBS}
96+
${TORCH_OMP_LIBRARY}
97+
)
98+
99+
# Set rpath for OpenMP library to avoid runtime linking issues
100+
if(TORCH_OMP_LIBRARY AND TORCH_OMP_LIB_DIR)
101+
# Add the OpenMP library directory to the rpath
102+
set_target_properties(
103+
metal_backend PROPERTIES BUILD_RPATH "${TORCH_OMP_LIB_DIR}"
104+
INSTALL_RPATH "${TORCH_OMP_LIB_DIR}"
105+
)
106+
# Also try common OpenMP library locations
107+
target_link_options(
108+
metal_backend PUBLIC -Wl,-rpath,${TORCH_OMP_LIB_DIR}
109+
-Wl,-rpath,/usr/local/opt/libomp/lib
110+
-Wl,-rpath,/opt/homebrew/opt/libomp/lib
111+
)
112+
message(STATUS "Added rpath for OpenMP library: ${TORCH_OMP_LIB_DIR}")
113+
endif()
114+
115+
executorch_target_link_options_shared_lib(metal_backend)
116+
install(
117+
TARGETS metal_backend
118+
EXPORT ExecuTorchTargets
119+
DESTINATION lib
120+
)

0 commit comments

Comments
 (0)