From ca6b16eaf2b50f445b70c49e3d9b174f11573451 Mon Sep 17 00:00:00 2001
From: Snider <snider@host.uk.com>
Date: Thu, 19 Feb 2026 20:39:51 +0000
Subject: [PATCH] feat(metal): bind memory diagnostics and device info
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

New bindings from mlx-c memory.h and metal.h:
- GetCacheMemory() — current allocator cache size
- ResetPeakMemory() — reset high-water mark
- SetWiredLimit() — control wired memory limit
- GetDeviceInfo() — GPU architecture, max buffer, memory size

All exposed at root package level via register_metal.go delegates.

Co-Authored-By: Virgil <virgil@lethean.io>
---
 internal/metal/stream.go | 38 ++++++++++++++++++++++++++++++++++++++
 register_metal.go        | 15 +++++++++++++++
 2 files changed, 53 insertions(+)

diff --git a/internal/metal/stream.go b/internal/metal/stream.go
index 4b3afd9..e8c5b23 100644
--- a/internal/metal/stream.go
+++ b/internal/metal/stream.go
@@ -77,3 +77,41 @@ func GetPeakMemory() uint64 {
 func ClearCache() {
 	C.mlx_clear_cache()
 }
+
+// GetCacheMemory returns the current Metal cache memory in bytes.
+func GetCacheMemory() uint64 {
+	var mem C.size_t
+	C.mlx_get_cache_memory(&mem)
+	return uint64(mem)
+}
+
+// ResetPeakMemory resets the peak memory high-water mark.
+func ResetPeakMemory() {
+	C.mlx_reset_peak_memory()
+}
+
+// SetWiredLimit sets the Metal wired memory limit. Returns the previous limit.
+func SetWiredLimit(limit uint64) uint64 {
+	var prev C.size_t
+	C.mlx_set_wired_limit(&prev, C.size_t(limit))
+	return uint64(prev)
+}
+
+// DeviceInfo holds Metal GPU hardware information.
+type DeviceInfo struct {
+	Architecture                string
+	MaxBufferLength             uint64
+	MaxRecommendedWorkingSetSize uint64
+	MemorySize                  uint64
+}
+
+// GetDeviceInfo returns Metal GPU hardware information.
+func GetDeviceInfo() DeviceInfo {
+	info := C.mlx_metal_device_info()
+	return DeviceInfo{
+		Architecture:                C.GoString(&info.architecture[0]),
+		MaxBufferLength:             uint64(info.max_buffer_length),
+		MaxRecommendedWorkingSetSize: uint64(info.max_recommended_working_set_size),
+		MemorySize:                  uint64(info.memory_size),
+	}
+}
diff --git a/register_metal.go b/register_metal.go
index a785aaa..754c2aa 100644
--- a/register_metal.go
+++ b/register_metal.go
@@ -35,6 +35,21 @@ func GetPeakMemory() uint64 { return metal.GetPeakMemory() }
 // ClearCache clears the Metal memory cache.
 func ClearCache() { metal.ClearCache() }
 
+// GetCacheMemory returns the current Metal cache memory in bytes.
+func GetCacheMemory() uint64 { return metal.GetCacheMemory() }
+
+// ResetPeakMemory resets the peak memory high-water mark.
+func ResetPeakMemory() { metal.ResetPeakMemory() }
+
+// SetWiredLimit sets the Metal wired memory limit. Returns the previous value.
+func SetWiredLimit(limit uint64) uint64 { return metal.SetWiredLimit(limit) }
+
+// DeviceInfo holds Metal GPU hardware information.
+type DeviceInfo = metal.DeviceInfo
+
+// GetDeviceInfo returns Metal GPU hardware information.
+func GetDeviceInfo() DeviceInfo { return metal.GetDeviceInfo() }
+
 // metalBackend implements inference.Backend for native Metal inference.
 type metalBackend struct{}