From 21ac3a1503001020122db5dce6adf34b761675f5 Mon Sep 17 00:00:00 2001
From: kchro3 <62481661+kchro3@users.noreply.github.com>
Date: Sat, 9 Sep 2023 02:12:10 -0700
Subject: [PATCH] metal : support for Swift (#3078)

* Metal support for Swift

* update

* add a toggle for arm/arm64

* set minimum versions for all platforms

* update to use newLibraryWithURL

* bump version

Co-authored-by: Jhen-Jie Hong <iainst0409@gmail.com>

---------

Co-authored-by: Jhen-Jie Hong <iainst0409@gmail.com>
---
 Package.swift | 32 +++++++++++++++++++++++++++-----
 ggml-metal.m  | 14 +++++++++++---
 2 files changed, 38 insertions(+), 8 deletions(-)

diff --git a/Package.swift b/Package.swift
index 96f52c4f0..fb95ef7eb 100644
--- a/Package.swift
+++ b/Package.swift
@@ -2,8 +2,30 @@
 
 import PackageDescription
 
+#if arch(arm) || arch(arm64)
+let platforms: [SupportedPlatform]? = [
+    .macOS(.v11),
+    .iOS(.v14),
+    .watchOS(.v4),
+    .tvOS(.v14)
+]
+let exclude: [String] = []
+let additionalSources: [String] = ["ggml-metal.m"]
+let additionalSettings: [CSetting] = [
+    .unsafeFlags(["-fno-objc-arc"]),
+    .define("GGML_SWIFT"),
+    .define("GGML_USE_METAL")
+]
+#else
+let platforms: [SupportedPlatform]? = nil
+let exclude: [String] = ["ggml-metal.metal"]
+let additionalSources: [String] = []
+let additionalSettings: [CSetting] = []
+#endif
+
 let package = Package(
     name: "llama",
+    platforms: platforms,
     products: [
         .library(name: "llama", targets: ["llama"]),
     ],
@@ -11,23 +33,23 @@ let package = Package(
         .target(
             name: "llama",
             path: ".",
-            exclude: ["ggml-metal.metal"],
+            exclude: exclude,
             sources: [
                 "ggml.c",
                 "llama.cpp",
                 "ggml-alloc.c",
-                "k_quants.c"
-            ],
+                "k_quants.c",
+            ] + additionalSources,
             publicHeadersPath: "spm-headers",
             cSettings: [
                 .unsafeFlags(["-Wno-shorten-64-to-32"]),
                 .define("GGML_USE_K_QUANTS"),
                 .define("GGML_USE_ACCELERATE")
-            ],
+            ] + additionalSettings,
             linkerSettings: [
                 .linkedFramework("Accelerate")
             ]
-        ),
+        )
     ],
     cxxLanguageStandard: .cxx11
 )
diff --git a/ggml-metal.m b/ggml-metal.m
index 1e6845b97..b577d7f60 100644
--- a/ggml-metal.m
+++ b/ggml-metal.m
@@ -144,12 +144,20 @@ struct ggml_metal_context * ggml_metal_init(int n_cb) {
 
     ctx->d_queue = dispatch_queue_create("llama.cpp", DISPATCH_QUEUE_CONCURRENT);
 
-#if 0
-    // compile from source string and show compile log
+#ifdef GGML_SWIFT
+    // load the default.metallib file
     {
         NSError * error = nil;
 
-        ctx->library = [ctx->device newLibraryWithSource:msl_library_source options:nil error:&error];
+        NSBundle * bundle = [NSBundle bundleForClass:[GGMLMetalClass class]];
+        NSString * llamaBundlePath = [bundle pathForResource:@"llama_llama" ofType:@"bundle"];
+        NSBundle * llamaBundle = [NSBundle bundleWithPath:llamaBundlePath];
+        NSString * libPath = [llamaBundle pathForResource:@"default" ofType:@"metallib"];
+        NSURL * libURL = [NSURL fileURLWithPath:libPath];
+
+        // Load the metallib file into a Metal library
+        ctx->library = [ctx->device newLibraryWithURL:libURL error:&error];
+
         if (error) {
             metal_printf("%s: error: %s\n", __func__, [[error description] UTF8String]);
             return NULL;