metal : build metallib + fix embed path

.github/workflows/build.yml

@@ -333,6 +333,7 @@ jobs:
  mkdir build
  cd build
  cmake -G Xcode .. \
+ -DLLAMA_METAL_EMBED_LIBRARY=ON \
  -DLLAMA_BUILD_EXAMPLES=OFF \
  -DLLAMA_BUILD_TESTS=OFF \
  -DLLAMA_BUILD_SERVER=OFF \
@@ -361,6 +362,7 @@ jobs:
  mkdir build
  cd build
  cmake -G Xcode .. \
+ -DLLAMA_METAL_EMBED_LIBRARY=ON \
  -DLLAMA_BUILD_EXAMPLES=OFF \
  -DLLAMA_BUILD_TESTS=OFF \
  -DLLAMA_BUILD_SERVER=OFF \

.gitignore

@@ -25,6 +25,8 @@
 .vscode/
 .idea/
 
+ggml-metal-embed.metal
+
 lcov-report/
 gcovr-report/
 

CMakeLists.txt

@@ -197,9 +197,6 @@ if (LLAMA_METAL)
  add_compile_definitions(GGML_METAL_NDEBUG)
  endif()
 
- # get full path to the file
- #add_compile_definitions(GGML_METAL_DIR_KERNELS="${CMAKE_CURRENT_SOURCE_DIR}/")
-
  # copy ggml-common.h and ggml-metal.metal to bin directory
  configure_file(ggml-common.h ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/ggml-common.h COPYONLY)
  configure_file(ggml-metal.metal ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/ggml-metal.metal COPYONLY)
@@ -208,53 +205,62 @@ if (LLAMA_METAL)
  enable_language(ASM)
  add_compile_definitions(GGML_METAL_EMBED_LIBRARY)
 
+ set(METALLIB_COMMON "${CMAKE_CURRENT_SOURCE_DIR}/ggml-common.h")
  set(METALLIB_SOURCE "${CMAKE_CURRENT_SOURCE_DIR}/ggml-metal.metal")
+
  file(MAKE_DIRECTORY "${CMAKE_BINARY_DIR}/autogenerated")
- set(EMBED_METALLIB_ASSEMBLY "${CMAKE_BINARY_DIR}/autogenerated/ggml-embed-metallib.s")
+
+ # merge ggml-common.h and ggml-metal.metal into a single file
+ set(METALLIB_EMBED_ASM "${CMAKE_BINARY_DIR}/autogenerated/ggml-metal-embed.s")
+ set(METALLIB_SOURCE_EMBED "${CMAKE_BINARY_DIR}/autogenerated/ggml-metal-embed.metal")
 
  add_custom_command(
- OUTPUT ${EMBED_METALLIB_ASSEMBLY}
- COMMAND echo ".section __DATA,__ggml_metallib" > ${EMBED_METALLIB_ASSEMBLY}
- COMMAND echo ".globl _ggml_metallib_start" >> ${EMBED_METALLIB_ASSEMBLY}
- COMMAND echo "_ggml_metallib_start:" >> ${EMBED_METALLIB_ASSEMBLY}
- COMMAND echo ".incbin \\\"${METALLIB_SOURCE}\\\"" >> ${EMBED_METALLIB_ASSEMBLY}
- COMMAND echo ".globl _ggml_metallib_end" >> ${EMBED_METALLIB_ASSEMBLY}
- COMMAND echo "_ggml_metallib_end:" >> ${EMBED_METALLIB_ASSEMBLY}
- DEPENDS ${METALLIB_SOURCE}
+ OUTPUT ${METALLIB_EMBED_ASM}
+ COMMAND echo "Embedding Metal library"
+ COMMAND sed -e '/\#include \"ggml-common.h\"/r ${METALLIB_COMMON}' -e '/\#include \"ggml-common.h\"/d' < ${METALLIB_SOURCE} > ${METALLIB_SOURCE_EMBED}
+ COMMAND echo ".section __DATA,__ggml_metallib" > ${METALLIB_EMBED_ASM}
+ COMMAND echo ".globl _ggml_metallib_start" >> ${METALLIB_EMBED_ASM}
+ COMMAND echo "_ggml_metallib_start:" >> ${METALLIB_EMBED_ASM}
+ COMMAND echo ".incbin \\\"${METALLIB_SOURCE_EMBED}\\\"" >> ${METALLIB_EMBED_ASM}
+ COMMAND echo ".globl _ggml_metallib_end" >> ${METALLIB_EMBED_ASM}
+ COMMAND echo "_ggml_metallib_end:" >> ${METALLIB_EMBED_ASM}
+ DEPENDS ggml-metal.metal ggml-common.h
  COMMENT "Generate assembly for embedded Metal library"
  )
 
- set(GGML_SOURCES_METAL ${GGML_SOURCES_METAL} ${EMBED_METALLIB_ASSEMBLY})
- endif()
-
- if (LLAMA_METAL_SHADER_DEBUG)
- # custom command to do the following:
- # xcrun -sdk macosx metal -fno-fast-math -c ggml-metal.metal -o ggml-metal.air
- # xcrun -sdk macosx metallib ggml-metal.air -o default.metallib
- #
- # note: this is the only way I found to disable fast-math in Metal. it's ugly, but at least it works
- # disabling fast math is needed in order to pass tests/test-backend-ops
- # note: adding -fno-inline fixes the tests when using MTL_SHADER_VALIDATION=1
- # note: unfortunately, we have to call it default.metallib instead of ggml.metallib
- # ref: https://github.com/ggerganov/whisper.cpp/issues/1720
- set(XC_FLAGS -fno-fast-math -fno-inline -g)
- if (LLAMA_QKK_64)
- set(XC_FLAGS ${XC_FLAGS} -DQK_K=64)
+ set(GGML_SOURCES_METAL ${GGML_SOURCES_METAL} ${METALLIB_EMBED_ASM})
+ else()
+ if (LLAMA_METAL_SHADER_DEBUG)
+ # custom command to do the following:
+ # xcrun -sdk macosx metal -fno-fast-math -c ggml-metal.metal -o ggml-metal.air
+ # xcrun -sdk macosx metallib ggml-metal.air -o default.metallib
+ #
+ # note: this is the only way I found to disable fast-math in Metal. it's ugly, but at least it works
+ # disabling fast math is needed in order to pass tests/test-backend-ops
+ # note: adding -fno-inline fixes the tests when using MTL_SHADER_VALIDATION=1
+ # note: unfortunately, we have to call it default.metallib instead of ggml.metallib
+ # ref: https://github.com/ggerganov/whisper.cpp/issues/1720
+ set(XC_FLAGS -fno-fast-math -fno-inline -g)
+ else()
+ set(XC_FLAGS -O3)
  endif()
 
  add_custom_command(
  OUTPUT ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/default.metallib
  COMMAND xcrun -sdk macosx metal ${XC_FLAGS} -c ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/ggml-metal.metal -o ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/ggml-metal.air
  COMMAND xcrun -sdk macosx metallib ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/ggml-metal.air -o ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/default.metallib
- DEPENDS ggml-metal.metal
+ COMMAND rm -f ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/ggml-metal.air
+ COMMAND rm -f ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/ggml-common.h
+ COMMAND rm -f ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/ggml-metal.metal
+ DEPENDS ggml-metal.metal ggml-common.h
  COMMENT "Compiling Metal kernels"
- )
+  )
 
  add_custom_target(
  ggml-metal ALL
  DEPENDS ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/default.metallib
- )
- endif()
+  )
+ endif() # LLAMA_METAL_EMBED_LIBRARY
 
  set(LLAMA_EXTRA_LIBS ${LLAMA_EXTRA_LIBS}
  ${FOUNDATION_LIBRARY}

Makefile

@@ -553,15 +553,16 @@ ggml-metal.o: ggml-metal.m ggml-metal.h
  $(CC) $(CFLAGS) -c $< -o $@
 
 ifdef LLAMA_METAL_EMBED_LIBRARY
-ggml-metal-embed.o: ggml-metal.metal
+ggml-metal-embed.o: ggml-metal.metal ggml-common.h
  @echo "Embedding Metal library"
+ @sed -e '/#include "ggml-common.h"/r ggml-common.h' -e '/#include "ggml-common.h"/d' < ggml-metal.metal > ggml-metal-embed.metal
  $(eval TEMP_ASSEMBLY=$(shell mktemp))
- @echo ".section __DATA, __ggml_metallib" > $(TEMP_ASSEMBLY)
- @echo ".globl _ggml_metallib_start" >> $(TEMP_ASSEMBLY)
- @echo "_ggml_metallib_start:" >> $(TEMP_ASSEMBLY)
- @echo ".incbin \"$<\"" >> $(TEMP_ASSEMBLY)
- @echo ".globl _ggml_metallib_end" >> $(TEMP_ASSEMBLY)
- @echo "_ggml_metallib_end:" >> $(TEMP_ASSEMBLY)
+ @echo ".section __DATA, __ggml_metallib"  >  $(TEMP_ASSEMBLY)
+ @echo ".globl _ggml_metallib_start"  >> $(TEMP_ASSEMBLY)
+ @echo "_ggml_metallib_start:"  >> $(TEMP_ASSEMBLY)
+ @echo ".incbin \"ggml-metal-embed.metal\"" >> $(TEMP_ASSEMBLY)
+ @echo ".globl _ggml_metallib_end"  >> $(TEMP_ASSEMBLY)
+ @echo "_ggml_metallib_end:"  >> $(TEMP_ASSEMBLY)
  @$(AS) $(TEMP_ASSEMBLY) -o $@
  @rm -f ${TEMP_ASSEMBLY}
 endif

ggml-metal.m

@@ -280,19 +280,33 @@ static void ggml_metal_log(enum ggml_log_level level, const char * format, ...){
  id<MTLLibrary> metal_library;
 
  // load library
+ //
+ // - first check if the library is embedded
+ // - then check if the library is in the bundle
+ // - if not found, load the source and compile it
+ // - if that fails, return NULL
  {
  NSBundle * bundle = nil;
 #ifdef SWIFT_PACKAGE
  bundle = SWIFTPM_MODULE_BUNDLE;
 #else
  bundle = [NSBundle bundleForClass:[GGMLMetalClass class]];
 #endif
+
  NSError * error = nil;
- NSString * libPath = [bundle pathForResource:@"default" ofType:@"metallib"];
- if (libPath != nil) {
+
+#if GGML_METAL_EMBED_LIBRARY
+ const bool try_metallib = false;
+#else
+ const bool try_metallib = true;
+#endif
+
+ NSString * path_lib = [bundle pathForResource:@"default" ofType:@"metallib"];
+ if (try_metallib && path_lib != nil) {
  // pre-compiled library found
- NSURL * libURL = [NSURL fileURLWithPath:libPath];
- GGML_METAL_LOG_INFO("%s: loading '%s'\n", __func__, [libPath UTF8String]);
+ NSURL * libURL = [NSURL fileURLWithPath:path_lib];
+ GGML_METAL_LOG_INFO("%s: loading '%s'\n", __func__, [path_lib UTF8String]);
+
  metal_library = [ctx->device newLibraryWithURL:libURL error:&error];
  if (error) {
  GGML_METAL_LOG_ERROR("%s: error: %s\n", __func__, [[error description] UTF8String]);
@@ -305,31 +319,34 @@ static void ggml_metal_log(enum ggml_log_level level, const char * format, ...){
  extern const char ggml_metallib_start[];
  extern const char ggml_metallib_end[];
 
- NSString * src  = [[NSString alloc] initWithBytes:ggml_metallib_start length:(ggml_metallib_end-ggml_metallib_start) encoding:NSUTF8StringEncoding];
+ NSString * src = [[NSString alloc] initWithBytes:ggml_metallib_start length:(ggml_metallib_end-ggml_metallib_start) encoding:NSUTF8StringEncoding];
 #else
  GGML_METAL_LOG_INFO("%s: default.metallib not found, loading from source\n", __func__);
 
- NSString * sourcePath;
- NSString * ggmlMetalPathResources = [[NSProcessInfo processInfo].environment objectForKey:@"GGML_METAL_PATH_RESOURCES"];
+ NSString * path_source;
+ NSString * path_resource = [[NSProcessInfo processInfo].environment objectForKey:@"GGML_METAL_PATH_RESOURCES"];
 
- GGML_METAL_LOG_INFO("%s: GGML_METAL_PATH_RESOURCES = %s\n", __func__, ggmlMetalPathResources ? [ggmlMetalPathResources UTF8String] : "nil");
+ GGML_METAL_LOG_INFO("%s: GGML_METAL_PATH_RESOURCES = %s\n", __func__, path_resource ? [path_resource UTF8String] : "nil");
 
- if (ggmlMetalPathResources) {
- sourcePath = [ggmlMetalPathResources stringByAppendingPathComponent:@"ggml-metal.metal"];
+ if (path_resource) {
+ path_source = [path_resource stringByAppendingPathComponent:@"ggml-metal.metal"];
  } else {
- sourcePath = [bundle pathForResource:@"ggml-metal" ofType:@"metal"];
+ path_source = [bundle pathForResource:@"ggml-metal" ofType:@"metal"];
  }
- if (sourcePath == nil) {
+
+ if (path_source == nil) {
  GGML_METAL_LOG_WARN("%s: error: could not use bundle path to find ggml-metal.metal, falling back to trying cwd\n", __func__);
- sourcePath = @"ggml-metal.metal";
+ path_source = @"ggml-metal.metal";
  }
- GGML_METAL_LOG_INFO("%s: loading '%s'\n", __func__, [sourcePath UTF8String]);
- NSString * src = [NSString stringWithContentsOfFile:sourcePath encoding:NSUTF8StringEncoding error:&error];
+
+ GGML_METAL_LOG_INFO("%s: loading '%s'\n", __func__, [path_source UTF8String]);
+
+ NSString * src = [NSString stringWithContentsOfFile:path_source encoding:NSUTF8StringEncoding error:&error];
  if (error) {
  GGML_METAL_LOG_ERROR("%s: error: %s\n", __func__, [[error description] UTF8String]);
  return NULL;
  }
-#endif
+#endif // GGML_METAL_EMBED_LIBRARY
 
  @autoreleasepool {
  // dictionary of preprocessor macros

ggml-metal.metal

@@ -4,9 +4,6 @@
 
 #include <metal_stdlib>
 
-#define GGML_COMMON_IMPL_METAL
-#include "ggml-common.h"
-
 using namespace metal;
 
 #define MAX(x, y) ((x) > (y) ? (x) : (y))