Skip to content

Commit 5d3eb99

Browse files
forforever73lvyichen
forforever73
and
lvyichen
authored
metal : per-op source split + parallel compile (ggml-org#24021)
* preliminary extract common header * op source split * split metallib into 8 libs && load in parallel * derive kernel->library routing from functionNames * x-macro lib list + underscore filenames, dedup QK_NL, MRC fixes * op source split 8 to 20 * improve robustness of source fallback * clean up * change bool -> atomic_bool * only prepend headers that source actually includes * no semaphore, use GCD global queue * dedup library compile path, fix NSError lifetime, rename gla * relocate upstream concat/rope_back/repeat kernel changes into split files * move ggml-common.h from common.h into dequantize.h to shrink binary size --------- Co-authored-by: lvyichen <lvyichen@stepfun.com>
1 parent 5fd2dc2 commit 5d3eb99

27 files changed

Lines changed: 11322 additions & 10928 deletions

‎ggml/CMakeLists.txt‎

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -341,9 +341,6 @@ set(GGML_PUBLIC_HEADERS
341341
include/gguf.h)
342342

343343
set_target_properties(ggml PROPERTIES PUBLIC_HEADER "${GGML_PUBLIC_HEADERS}")
344-
#if (GGML_METAL)
345-
# set_target_properties(ggml PROPERTIES RESOURCE "${CMAKE_CURRENT_SOURCE_DIR}/src/ggml-metal.metal")
346-
#endif()
347344
install(TARGETS ggml LIBRARY PUBLIC_HEADER)
348345
install(TARGETS ggml-base LIBRARY)
349346

‎ggml/src/ggml-metal/CMakeLists.txt‎

Lines changed: 120 additions & 52 deletions
Original file line numberDiff line numberDiff line change
@@ -24,62 +24,119 @@ if (GGML_METAL_NDEBUG)
2424
endif()
2525

2626
set(METALLIB_COMMON "${CMAKE_CURRENT_SOURCE_DIR}/../ggml-common.h")
27+
set(METALLIB_KERNELS_COMMON "${CMAKE_CURRENT_SOURCE_DIR}/kernels/common.h")
28+
set(METALLIB_KERNELS_DEQUANTIZE "${CMAKE_CURRENT_SOURCE_DIR}/kernels/dequantize.h")
29+
set(METALLIB_KERNELS_QUANTIZE "${CMAKE_CURRENT_SOURCE_DIR}/kernels/quantize.h")
30+
31+
set(METALLIB_KERNEL_SOURCES
32+
kernels/fa.metal
33+
kernels/mul_mv.metal
34+
kernels/mul_mm.metal
35+
kernels/quantize.metal
36+
kernels/softmax.metal
37+
kernels/norm.metal
38+
kernels/unary.metal
39+
kernels/binbcast.metal
40+
kernels/reduce.metal
41+
kernels/tri.metal
42+
kernels/ssm.metal
43+
kernels/wkv.metal
44+
kernels/gated_delta_net.metal
45+
kernels/solve_tri.metal
46+
kernels/rope.metal
47+
kernels/conv.metal
48+
kernels/upscale.metal
49+
kernels/argsort.metal
50+
kernels/pool.metal
51+
kernels/misc.metal
52+
)
53+
2754
if (GGML_METAL_EMBED_LIBRARY)
2855
enable_language(ASM)
2956

3057
add_compile_definitions(GGML_METAL_EMBED_LIBRARY)
3158

32-
set(METALLIB_SOURCE "${CMAKE_CURRENT_SOURCE_DIR}/ggml-metal.metal")
33-
set(METALLIB_IMPL "${CMAKE_CURRENT_SOURCE_DIR}/ggml-metal-impl.h")
59+
set(METALLIB_IMPL "${CMAKE_CURRENT_SOURCE_DIR}/ggml-metal-impl.h")
3460

3561
file(MAKE_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}/autogenerated")
3662

37-
# merge ggml-common.h and ggml-metal.metal into a single file
38-
set(METALLIB_EMBED_ASM "${CMAKE_CURRENT_BINARY_DIR}/autogenerated/ggml-metal-embed.s")
39-
set(METALLIB_SOURCE_EMBED "${CMAKE_CURRENT_BINARY_DIR}/autogenerated/ggml-metal-embed.metal")
40-
set(METALLIB_SOURCE_EMBED_TMP "${CMAKE_CURRENT_BINARY_DIR}/autogenerated/ggml-metal-embed.metal.tmp")
63+
set(METALLIB_EMBED_ASM_FILES "")
64+
foreach(src ${METALLIB_KERNEL_SOURCES})
65+
get_filename_component(kind ${src} NAME_WE)
66+
# symbol names must be valid C identifiers ('-' is not allowed)
67+
string(REPLACE "-" "_" kind_sym ${kind})
68+
69+
set(SRC "${CMAKE_CURRENT_SOURCE_DIR}/kernels/${kind}.metal")
70+
set(EMBED "${CMAKE_CURRENT_BINARY_DIR}/autogenerated/ggml-metal-embed-${kind}.metal")
71+
set(ASM "${CMAKE_CURRENT_BINARY_DIR}/autogenerated/ggml-metal-embed-${kind}.s")
72+
73+
# only prepend headers that this source actually includes
74+
set(HEADERS_FOR_SRC ${METALLIB_KERNELS_COMMON})
75+
file(STRINGS ${SRC} _has_dequantize REGEX "#include \"dequantize\\.h\"")
76+
file(STRINGS ${SRC} _has_quantize REGEX "#include \"quantize\\.h\"")
77+
if(_has_dequantize)
78+
list(APPEND HEADERS_FOR_SRC ${METALLIB_KERNELS_DEQUANTIZE})
79+
endif()
80+
if(_has_quantize)
81+
list(APPEND HEADERS_FOR_SRC ${METALLIB_KERNELS_QUANTIZE})
82+
endif()
83+
84+
add_custom_command(
85+
OUTPUT "${ASM}"
86+
# Step 1: concatenate shared headers + this kernel source
87+
COMMAND cat ${HEADERS_FOR_SRC} ${SRC} > "${EMBED}.tmp1"
88+
# Step 2: remove internal #include and #pragma once
89+
COMMAND sed -e "/\#include \"common.h\"/d" -e "/\#include \"dequantize.h\"/d" -e "/\#include \"quantize.h\"/d" -e "/\#pragma once/d" < "${EMBED}.tmp1" > "${EMBED}.tmp2"
90+
# Step 3: inline ggml-common.h (replacing __embed_ggml-common.h__ sentinel)
91+
COMMAND sed -e "/__embed_ggml-common.h__/r ${METALLIB_COMMON}" -e "/__embed_ggml-common.h__/d" < "${EMBED}.tmp2" > "${EMBED}.tmp3"
92+
# Step 4: inline ggml-metal-impl.h
93+
COMMAND sed -e "/\#include \"ggml-metal-impl.h\"/r ${METALLIB_IMPL}" -e "/\#include \"ggml-metal-impl.h\"/d" < "${EMBED}.tmp3" > "${EMBED}"
94+
# Step 5: emit an asm chunk with kind-specific start/end symbols
95+
# note: '-' is illegal in C symbols, so we use kind_sym; the macOS
96+
# section name is limited to 16 chars so we keep it shared
97+
# across kinds (__ggml_metallib) and only vary the global symbols.
98+
COMMAND echo ".section __DATA,__ggml_metallib" > "${ASM}"
99+
COMMAND echo ".globl _ggml_metallib_${kind_sym}_start" >> "${ASM}"
100+
COMMAND echo "_ggml_metallib_${kind_sym}_start:" >> "${ASM}"
101+
COMMAND echo .incbin "\"${EMBED}\"" >> "${ASM}"
102+
COMMAND echo ".globl _ggml_metallib_${kind_sym}_end" >> "${ASM}"
103+
COMMAND echo "_ggml_metallib_${kind_sym}_end:" >> "${ASM}"
104+
DEPENDS ../ggml-common.h ggml-metal-impl.h
105+
kernels/common.h kernels/dequantize.h kernels/quantize.h
106+
kernels/${kind}.metal
107+
COMMENT "Generate embedded Metal library for ${kind}"
108+
VERBATIM
109+
)
41110

42-
add_custom_command(
43-
OUTPUT "${METALLIB_EMBED_ASM}"
44-
COMMAND echo "Embedding Metal library"
45-
COMMAND sed -e "/__embed_ggml-common.h__/r ${METALLIB_COMMON}" -e "/__embed_ggml-common.h__/d" < "${METALLIB_SOURCE}" > "${METALLIB_SOURCE_EMBED_TMP}"
46-
COMMAND sed -e "/\#include \"ggml-metal-impl.h\"/r ${METALLIB_IMPL}" -e "/\#include \"ggml-metal-impl.h\"/d" < "${METALLIB_SOURCE_EMBED_TMP}" > "${METALLIB_SOURCE_EMBED}"
47-
COMMAND echo ".section __DATA,__ggml_metallib" > "${METALLIB_EMBED_ASM}"
48-
COMMAND echo ".globl _ggml_metallib_start" >> "${METALLIB_EMBED_ASM}"
49-
COMMAND echo "_ggml_metallib_start:" >> "${METALLIB_EMBED_ASM}"
50-
COMMAND echo .incbin "\"${METALLIB_SOURCE_EMBED}\"" >> "${METALLIB_EMBED_ASM}"
51-
COMMAND echo ".globl _ggml_metallib_end" >> "${METALLIB_EMBED_ASM}"
52-
COMMAND echo "_ggml_metallib_end:" >> "${METALLIB_EMBED_ASM}"
53-
DEPENDS ../ggml-common.h ggml-metal.metal ggml-metal-impl.h
54-
COMMENT "Generate assembly for embedded Metal library"
55-
VERBATIM
56-
)
111+
list(APPEND METALLIB_EMBED_ASM_FILES "${ASM}")
112+
endforeach()
57113

58-
target_sources(ggml-metal PRIVATE "${METALLIB_EMBED_ASM}")
114+
target_sources(ggml-metal PRIVATE ${METALLIB_EMBED_ASM_FILES})
59115
else()
60-
# copy metal files to bin directory
116+
# copy header files to bin directory
61117
configure_file(../ggml-common.h ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/ggml-common.h COPYONLY)
62-
configure_file(ggml-metal.metal ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/ggml-metal.metal COPYONLY)
63118
configure_file(ggml-metal-impl.h ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/ggml-metal-impl.h COPYONLY)
64119

120+
file(MAKE_DIRECTORY "${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/kernels")
121+
configure_file(kernels/common.h ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/kernels/common.h COPYONLY)
122+
configure_file(kernels/dequantize.h ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/kernels/dequantize.h COPYONLY)
123+
configure_file(kernels/quantize.h ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/kernels/quantize.h COPYONLY)
124+
125+
foreach(src ${METALLIB_KERNEL_SOURCES})
126+
configure_file(${src} ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/${src} COPYONLY)
127+
endforeach()
128+
65129
if (GGML_METAL_SHADER_DEBUG)
66-
# custom command to do the following:
67-
# xcrun -sdk macosx metal -fno-fast-math -c ggml-metal.metal -o ggml-metal.air
68-
# xcrun -sdk macosx metallib ggml-metal.air -o default.metallib
69-
#
70-
# note: this is the only way I found to disable fast-math in Metal. it's ugly, but at least it works
71-
# disabling fast math is needed in order to pass tests/test-backend-ops
130+
# note: disabling fast math is needed in order to pass tests/test-backend-ops
72131
# note: adding -fno-inline fixes the tests when using MTL_SHADER_VALIDATION=1
73132
# note: unfortunately, we have to call it default.metallib instead of ggml.metallib
74133
# ref: https://github.com/ggml-org/whisper.cpp/issues/1720
75134
# note: adding -g causes segmentation fault during compile
76-
#set(XC_FLAGS -fno-fast-math -fno-inline -g)
77135
set(XC_FLAGS -fno-fast-math -fno-inline)
78136
else()
79137
set(XC_FLAGS -O3)
80138
endif()
81139

82-
# Append macOS metal versioning flags
83140
if (GGML_METAL_MACOSX_VERSION_MIN)
84141
message(STATUS "Adding -mmacosx-version-min=${GGML_METAL_MACOSX_VERSION_MIN} flag to metal compilation")
85142
list (APPEND XC_FLAGS -mmacosx-version-min=${GGML_METAL_MACOSX_VERSION_MIN})
@@ -90,35 +147,46 @@ else()
90147
list (APPEND XC_FLAGS -std=${GGML_METAL_STD})
91148
endif()
92149

150+
# Compile each kernel source to .air, then link into default.metallib
151+
set(AIR_FILES "")
152+
foreach(src ${METALLIB_KERNEL_SOURCES})
153+
get_filename_component(name ${src} NAME_WE)
154+
set(AIR "${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/${name}.air")
155+
list(APPEND AIR_FILES ${AIR})
156+
add_custom_command(
157+
OUTPUT ${AIR}
158+
COMMAND xcrun -sdk macosx metal ${XC_FLAGS} -I ${CMAKE_RUNTIME_OUTPUT_DIRECTORY} -c ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/${src} -o ${AIR}
159+
DEPENDS ${src} kernels/common.h kernels/dequantize.h kernels/quantize.h ${METALLIB_COMMON} ggml-metal-impl.h
160+
COMMENT "Compiling ${src}"
161+
VERBATIM
162+
)
163+
endforeach()
164+
93165
add_custom_command(
94166
OUTPUT ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/default.metallib
95-
COMMAND xcrun -sdk macosx metal ${XC_FLAGS} -c ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/ggml-metal.metal -o - |
96-
xcrun -sdk macosx metallib - -o ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/default.metallib
167+
COMMAND xcrun -sdk macosx metallib ${AIR_FILES} -o ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/default.metallib
97168
COMMAND rm -f ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/ggml-common.h
98-
COMMAND rm -f ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/ggml-metal.metal
99-
DEPENDS ggml-metal.metal ${METALLIB_COMMON}
100-
COMMENT "Compiling Metal kernels"
101-
)
169+
COMMAND rm -f ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/ggml-metal-impl.h
170+
COMMAND rm -rf ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/kernels
171+
DEPENDS ${AIR_FILES}
172+
COMMENT "Linking Metal kernels into default.metallib"
173+
)
102174

103-
# FIXME: only add to the ggml-metal target?
104175
add_custom_target(
105176
ggml-metal-lib ALL
106177
DEPENDS ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/default.metallib
107-
)
178+
)
108179
endif() # GGML_METAL_EMBED_LIBRARY
109180

110181
if (NOT GGML_METAL_EMBED_LIBRARY)
111182
install(
112-
FILES src/ggml-metal/ggml-metal.metal
113-
PERMISSIONS
114-
OWNER_READ
115-
OWNER_WRITE
116-
GROUP_READ
117-
WORLD_READ
118-
DESTINATION ${CMAKE_INSTALL_BINDIR})
119-
120-
install(
121-
FILES ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/default.metallib
122-
DESTINATION ${CMAKE_INSTALL_BINDIR}
123-
)
183+
DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/kernels/
184+
DESTINATION ${CMAKE_INSTALL_BINDIR}/kernels
185+
FILES_MATCHING PATTERN "*.metal" PATTERN "*.h"
186+
)
187+
188+
install(
189+
FILES ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/default.metallib
190+
DESTINATION ${CMAKE_INSTALL_BINDIR}
191+
)
124192
endif()

0 commit comments

Comments
 (0)