From a189810df6c4b0492463d1ddb68993c9abc32c7f Mon Sep 17 00:00:00 2001
From: peanut256 <peanut256@users.noreply.github.com>
Date: Mon, 26 Feb 2024 00:16:45 +0100
Subject: [PATCH] Determine max VRAM on macOS using
 `recommendedMaxWorkingSetSize` (#2354)

* read iogpu.wired_limit_mb on macOS

Fix for https://github.com/ollama/ollama/issues/1826

* improved determination of available vram on macOS

read the recommended maximal vram on macOS via Metal API

* Removed macOS-specific logging

* Remove logging from gpu_darwin.go

* release Core Foundation object

fixes a possible memory leak
---
 gpu/gpu_darwin.go     | 23 +++++++----------------
 gpu/gpu_info_darwin.h |  3 +++
 gpu/gpu_info_darwin.m | 11 +++++++++++
 3 files changed, 21 insertions(+), 16 deletions(-)
 create mode 100644 gpu/gpu_info_darwin.h
 create mode 100644 gpu/gpu_info_darwin.m

diff --git a/gpu/gpu_darwin.go b/gpu/gpu_darwin.go
index 97907511..9a418c0b 100644
--- a/gpu/gpu_darwin.go
+++ b/gpu/gpu_darwin.go
@@ -1,12 +1,14 @@
 //go:build darwin
 
 package gpu
-
+/*
+#cgo CFLAGS: -x objective-c
+#cgo LDFLAGS: -framework Foundation -framework CoreGraphics -framework Metal
+#include "gpu_info_darwin.h"
+*/
 import "C"
 import (
 	"runtime"
-
-	"github.com/pbnjay/memory"
 )
 
 // CheckVRAM returns the free VRAM in bytes on Linux machines with NVIDIA GPUs
@@ -15,19 +17,8 @@ func CheckVRAM() (int64, error) {
 		// gpu not supported, this may not be metal
 		return 0, nil
 	}
-
-	// on macOS, there's already buffer for available vram (see below) so just return the total
-	systemMemory := int64(memory.TotalMemory())
-
-	// macOS limits how much memory is available to the GPU based on the amount of system memory
-	// TODO: handle case where iogpu.wired_limit_mb is set to a higher value
-	if systemMemory <= 36*1024*1024*1024 {
-		systemMemory = systemMemory * 2 / 3
-	} else {
-		systemMemory = systemMemory * 3 / 4
-	}
-
-	return systemMemory, nil
+	recommendedMaxVRAM := int64(C.getRecommendedMaxVRAM())
+	return recommendedMaxVRAM, nil
 }
 
 func GetGPUInfo() GpuInfo {
diff --git a/gpu/gpu_info_darwin.h b/gpu/gpu_info_darwin.h
new file mode 100644
index 00000000..6ba30c0a
--- /dev/null
+++ b/gpu/gpu_info_darwin.h
@@ -0,0 +1,3 @@
+#import <Metal/Metal.h>
+#include <stdint.h>
+uint64_t getRecommendedMaxVRAM();
diff --git a/gpu/gpu_info_darwin.m b/gpu/gpu_info_darwin.m
new file mode 100644
index 00000000..06d7b69b
--- /dev/null
+++ b/gpu/gpu_info_darwin.m
@@ -0,0 +1,11 @@
+//go:build darwin
+#include "gpu_info_darwin.h"
+
+uint64_t getRecommendedMaxVRAM()
+{
+	id<MTLDevice> device = MTLCreateSystemDefaultDevice();
+	uint64_t result = device.recommendedMaxWorkingSetSize;
+	CFRelease(device);
+	return result;
+}
+