Skip to content

Commit 0569909

Browse files
BB-fatBB-fat
committed
metal : Cache the Metal library at the device context level
1 parent 5e2d57b commit 0569909

File tree

1 file changed

+151
-130
lines changed

1 file changed

+151
-130
lines changed

‎ggml/src/ggml-metal/ggml-metal.m‎

Lines changed: 151 additions & 130 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,7 @@
4646
static struct ggml_backend_metal_device_context {
4747
id<MTLDevice> mtl_device;
4848
int mtl_device_ref_count;
49+
id<MTLLibrary> mtl_library;
4950

5051
bool has_simdgroup_reduction;
5152
bool has_simdgroup_mm;
@@ -57,6 +58,7 @@
5758
} g_ggml_ctx_dev_main = {
5859
/*.mtl_device =*/ nil,
5960
/*.mtl_device_ref_count =*/ 0,
61+
/*.mtl_library =*/ nil,
6062
/*.has_simdgroup_reduction =*/ false,
6163
/*.has_simdgroup_mm =*/ false,
6264
/*.has_residency_sets =*/ false,
@@ -65,6 +67,139 @@
6567
/*.name =*/ "",
6668
};
6769

70+
// load library
71+
//
72+
// - first check if the library is embedded
73+
// - then check if the library is in the bundle
74+
// - if not found, load the source and compile it
75+
// - if that fails, return NULL
76+
static id<MTLLibrary> ggml_metal_load_library(id<MTLDevice> device, bool use_bfloat) {
77+
id<MTLLibrary> metal_library = nil;
78+
NSError * error = nil;
79+
NSString * src = nil;
80+
81+
#if GGML_METAL_EMBED_LIBRARY
82+
GGML_LOG_INFO("%s: using embedded metal library\n", __func__);
83+
84+
extern const char ggml_metallib_start[];
85+
extern const char ggml_metallib_end[];
86+
87+
src = [[NSString alloc] initWithBytes:ggml_metallib_start length:(ggml_metallib_end-ggml_metallib_start) encoding:NSUTF8StringEncoding];
88+
89+
#else
90+
91+
#ifdef SWIFT_PACKAGE
92+
NSBundle * bundle = SWIFTPM_MODULE_BUNDLE;
93+
#else
94+
NSBundle * bundle = [NSBundle bundleForClass:[GGMLMetalClass class]];
95+
#endif
96+
97+
NSString * path_lib = [bundle pathForResource:@"default" ofType:@"metallib"];
98+
if (path_lib == nil) {
99+
// Try to find the resource in the directory where the current binary located.
100+
NSString * current_binary = [[NSProcessInfo processInfo] arguments][0];
101+
NSString * bin_dir = [current_binary stringByDeletingLastPathComponent];
102+
NSString * default_metallib_path = [NSString pathWithComponents:@[bin_dir, @"default.metallib"]];
103+
if ([[NSFileManager defaultManager] isReadableFileAtPath:default_metallib_path]) {
104+
GGML_LOG_INFO("%s: found '%s'\n", __func__, [default_metallib_path UTF8String]);
105+
NSDictionary * atts = [[NSFileManager defaultManager] attributesOfItemAtPath:default_metallib_path error:&error];
106+
if (atts && atts[NSFileType] == NSFileTypeSymbolicLink) {
107+
// Optionally, if this is a symlink, try to resolve it.
108+
default_metallib_path = [[NSFileManager defaultManager] destinationOfSymbolicLinkAtPath:default_metallib_path error:&error];
109+
if (default_metallib_path && [default_metallib_path length] > 0 && ![[default_metallib_path substringToIndex:1] isEqualToString:@"/"]) {
110+
// It is a relative path, adding the binary directory as directory prefix.
111+
default_metallib_path = [NSString pathWithComponents:@[bin_dir, default_metallib_path]];
112+
}
113+
if (!default_metallib_path || ![[NSFileManager defaultManager] isReadableFileAtPath:default_metallib_path]) {
114+
// Link to the resource could not be resolved.
115+
default_metallib_path = nil;
116+
} else {
117+
GGML_LOG_INFO("%s: symlink resolved '%s'\n", __func__, [default_metallib_path UTF8String]);
118+
}
119+
}
120+
} else {
121+
// The resource couldn't be found in the binary's directory.
122+
default_metallib_path = nil;
123+
}
124+
path_lib = default_metallib_path;
125+
}
126+
127+
if (path_lib != nil) {
128+
// pre-compiled library found
129+
NSURL * libURL = [NSURL fileURLWithPath:path_lib];
130+
GGML_LOG_INFO("%s: loading '%s'\n", __func__, [path_lib UTF8String]);
131+
132+
metal_library = [device newLibraryWithURL:libURL error:&error];
133+
if (error) {
134+
GGML_LOG_ERROR("%s: error: %s\n", __func__, [[error description] UTF8String]);
135+
return NULL;
136+
}
137+
} else {
138+
GGML_LOG_INFO("%s: default.metallib not found, loading from source\n", __func__);
139+
140+
NSString * path_source;
141+
NSString * path_resource = [[NSProcessInfo processInfo].environment objectForKey:@"GGML_METAL_PATH_RESOURCES"];
142+
143+
GGML_LOG_INFO("%s: GGML_METAL_PATH_RESOURCES = %s\n", __func__, path_resource ? [path_resource UTF8String] : "nil");
144+
145+
if (path_resource) {
146+
path_source = [path_resource stringByAppendingPathComponent:@"ggml-metal.metal"];
147+
} else {
148+
path_source = [bundle pathForResource:@"ggml-metal" ofType:@"metal"];
149+
}
150+
151+
if (path_source == nil) {
152+
GGML_LOG_WARN("%s: error: could not use bundle path to find ggml-metal.metal, falling back to trying cwd\n", __func__);
153+
path_source = @"ggml-metal.metal";
154+
}
155+
156+
GGML_LOG_INFO("%s: loading '%s'\n", __func__, [path_source UTF8String]);
157+
158+
src = [NSString stringWithContentsOfFile:path_source encoding:NSUTF8StringEncoding error:&error];
159+
if (error) {
160+
GGML_LOG_ERROR("%s: error: %s\n", __func__, [[error description] UTF8String]);
161+
return NULL;
162+
}
163+
}
164+
#endif
165+
166+
if (!metal_library) {
167+
@autoreleasepool {
168+
// dictionary of preprocessor macros
169+
NSMutableDictionary * prep = [NSMutableDictionary dictionary];
170+
171+
if (use_bfloat) {
172+
[prep setObject:@"1" forKey:@"GGML_METAL_USE_BF16"];
173+
}
174+
175+
#if GGML_METAL_EMBED_LIBRARY
176+
[prep setObject:@"1" forKey:@"GGML_METAL_EMBED_LIBRARY"];
177+
#endif
178+
179+
MTLCompileOptions * options = [MTLCompileOptions new];
180+
options.preprocessorMacros = prep;
181+
182+
//[options setFastMathEnabled:false];
183+
184+
metal_library = [device newLibraryWithSource:src options:options error:&error];
185+
if (error) {
186+
GGML_LOG_ERROR("%s: error: %s\n", __func__, [[error description] UTF8String]);
187+
return NULL;
188+
}
189+
190+
#if !__has_feature(objc_arc)
191+
[options release];
192+
#endif
193+
}
194+
}
195+
196+
#if GGML_METAL_EMBED_LIBRARY
197+
[src release];
198+
#endif // GGML_METAL_EMBED_LIBRARY
199+
200+
return metal_library;
201+
}
202+
68203
// acquire
69204
static id<MTLDevice> ggml_backend_metal_device_acq(struct ggml_backend_metal_device_context * ctx) {
70205
assert(ctx != NULL);
@@ -93,6 +228,13 @@
93228
#endif
94229

95230
strncpy(ctx->name, [[ctx->mtl_device name] UTF8String], sizeof(ctx->name) - 1);
231+
232+
if (ctx->mtl_library == nil) {
233+
ctx->mtl_library = ggml_metal_load_library(ctx->mtl_device, ctx->use_bfloat);
234+
if (ctx->mtl_library == nil) {
235+
GGML_LOG_ERROR("%s: error: failed to load metal library\n", __func__);
236+
}
237+
}
96238
}
97239

98240
ctx->mtl_device_ref_count++;
@@ -111,6 +253,11 @@ static void ggml_backend_metal_device_rel(struct ggml_backend_metal_device_conte
111253
if (ctx->mtl_device) {
112254
[ctx->mtl_device release];
113255
ctx->mtl_device = nil;
256+
257+
if (ctx->mtl_library) {
258+
[ctx->mtl_library release];
259+
ctx->mtl_library = nil;
260+
}
114261
}
115262
}
116263
}
@@ -522,136 +669,10 @@ @implementation GGMLMetalClass
522669

523670
ctx->d_queue = dispatch_queue_create("ggml-metal", DISPATCH_QUEUE_CONCURRENT);
524671

525-
id<MTLLibrary> metal_library = nil;
526-
527-
// load library
528-
//
529-
// - first check if the library is embedded
530-
// - then check if the library is in the bundle
531-
// - if not found, load the source and compile it
532-
// - if that fails, return NULL
533-
{
534-
NSError * error = nil;
535-
NSString * src = nil;
536-
537-
#if GGML_METAL_EMBED_LIBRARY
538-
GGML_LOG_INFO("%s: using embedded metal library\n", __func__);
539-
540-
extern const char ggml_metallib_start[];
541-
extern const char ggml_metallib_end[];
542-
543-
src = [[NSString alloc] initWithBytes:ggml_metallib_start length:(ggml_metallib_end-ggml_metallib_start) encoding:NSUTF8StringEncoding];
544-
545-
#else
546-
547-
#ifdef SWIFT_PACKAGE
548-
NSBundle * bundle = SWIFTPM_MODULE_BUNDLE;
549-
#else
550-
NSBundle * bundle = [NSBundle bundleForClass:[GGMLMetalClass class]];
551-
#endif
552-
553-
NSString * path_lib = [bundle pathForResource:@"default" ofType:@"metallib"];
554-
if (path_lib == nil) {
555-
// Try to find the resource in the directory where the current binary located.
556-
NSString * current_binary = [[NSProcessInfo processInfo] arguments][0];
557-
NSString * bin_dir = [current_binary stringByDeletingLastPathComponent];
558-
NSString * default_metallib_path = [NSString pathWithComponents:@[bin_dir, @"default.metallib"]];
559-
if ([[NSFileManager defaultManager] isReadableFileAtPath:default_metallib_path]) {
560-
GGML_LOG_INFO("%s: found '%s'\n", __func__, [default_metallib_path UTF8String]);
561-
NSDictionary * atts = [[NSFileManager defaultManager] attributesOfItemAtPath:default_metallib_path error:&error];
562-
if (atts && atts[NSFileType] == NSFileTypeSymbolicLink) {
563-
// Optionally, if this is a symlink, try to resolve it.
564-
default_metallib_path = [[NSFileManager defaultManager] destinationOfSymbolicLinkAtPath:default_metallib_path error:&error];
565-
if (default_metallib_path && [default_metallib_path length] > 0 && ![[default_metallib_path substringToIndex:1] isEqualToString:@"/"]) {
566-
// It is a relative path, adding the binary directory as directory prefix.
567-
default_metallib_path = [NSString pathWithComponents:@[bin_dir, default_metallib_path]];
568-
}
569-
if (!default_metallib_path || ![[NSFileManager defaultManager] isReadableFileAtPath:default_metallib_path]) {
570-
// Link to the resource could not be resolved.
571-
default_metallib_path = nil;
572-
} else {
573-
GGML_LOG_INFO("%s: symlink resolved '%s'\n", __func__, [default_metallib_path UTF8String]);
574-
}
575-
}
576-
} else {
577-
// The resource couldn't be found in the binary's directory.
578-
default_metallib_path = nil;
579-
}
580-
path_lib = default_metallib_path;
581-
}
582-
583-
if (path_lib != nil) {
584-
// pre-compiled library found
585-
NSURL * libURL = [NSURL fileURLWithPath:path_lib];
586-
GGML_LOG_INFO("%s: loading '%s'\n", __func__, [path_lib UTF8String]);
587-
588-
metal_library = [device newLibraryWithURL:libURL error:&error];
589-
if (error) {
590-
GGML_LOG_ERROR("%s: error: %s\n", __func__, [[error description] UTF8String]);
591-
return NULL;
592-
}
593-
} else {
594-
GGML_LOG_INFO("%s: default.metallib not found, loading from source\n", __func__);
595-
596-
NSString * path_source;
597-
NSString * path_resource = [[NSProcessInfo processInfo].environment objectForKey:@"GGML_METAL_PATH_RESOURCES"];
598-
599-
GGML_LOG_INFO("%s: GGML_METAL_PATH_RESOURCES = %s\n", __func__, path_resource ? [path_resource UTF8String] : "nil");
600-
601-
if (path_resource) {
602-
path_source = [path_resource stringByAppendingPathComponent:@"ggml-metal.metal"];
603-
} else {
604-
path_source = [bundle pathForResource:@"ggml-metal" ofType:@"metal"];
605-
}
606-
607-
if (path_source == nil) {
608-
GGML_LOG_WARN("%s: error: could not use bundle path to find ggml-metal.metal, falling back to trying cwd\n", __func__);
609-
path_source = @"ggml-metal.metal";
610-
}
611-
612-
GGML_LOG_INFO("%s: loading '%s'\n", __func__, [path_source UTF8String]);
613-
614-
src = [NSString stringWithContentsOfFile:path_source encoding:NSUTF8StringEncoding error:&error];
615-
if (error) {
616-
GGML_LOG_ERROR("%s: error: %s\n", __func__, [[error description] UTF8String]);
617-
return NULL;
618-
}
619-
}
620-
#endif
621-
622-
if (!metal_library) {
623-
@autoreleasepool {
624-
// dictionary of preprocessor macros
625-
NSMutableDictionary * prep = [NSMutableDictionary dictionary];
626-
627-
if (ctx_dev->use_bfloat) {
628-
[prep setObject:@"1" forKey:@"GGML_METAL_USE_BF16"];
629-
}
630-
631-
#if GGML_METAL_EMBED_LIBRARY
632-
[prep setObject:@"1" forKey:@"GGML_METAL_EMBED_LIBRARY"];
633-
#endif
634-
635-
MTLCompileOptions * options = [MTLCompileOptions new];
636-
options.preprocessorMacros = prep;
637-
638-
//[options setFastMathEnabled:false];
639-
640-
metal_library = [device newLibraryWithSource:src options:options error:&error];
641-
if (error) {
642-
GGML_LOG_ERROR("%s: error: %s\n", __func__, [[error description] UTF8String]);
643-
return NULL;
644-
}
645-
646-
#if !__has_feature(objc_arc)
647-
[options release];
648-
#endif
649-
}
650-
}
651-
652-
#if GGML_METAL_EMBED_LIBRARY
653-
[src release];
654-
#endif // GGML_METAL_EMBED_LIBRARY
672+
id<MTLLibrary> metal_library = ctx_dev->mtl_library;
673+
if (metal_library == nil) {
674+
GGML_LOG_ERROR("%s: error: metal library is nil\n", __func__);
675+
return NULL;
655676
}
656677

657678
// print MTL GPU family:

0 commit comments

Comments
 (0)