From 9fb35cbd899553e2554f49a74a231478a732b469 Mon Sep 17 00:00:00 2001
From: RyanOnTheInside <7623207+ryanontheinside@users.noreply.github.com>
Date: Sun, 8 Feb 2026 08:08:09 -0500
Subject: [PATCH] feat: mps compatability

Supersedes #10
Co-Authored-By: maybites <registered@maybites.ch>
---
 metrics.py | 68 ++++++++++++++++++++++++++++++++++++------------------
 1 file changed, 46 insertions(+), 22 deletions(-)

diff --git a/metrics.py b/metrics.py
index dc6335d..2b4cf67 100644
--- a/metrics.py
+++ b/metrics.py
@@ -1,38 +1,62 @@
-"""Memory measurement utilities for ProfilerX"""
+"""Memory measurement utilities for ProfilerX — supports CUDA, MPS, and CPU-only."""
 import psutil
 import os
 
 _process = psutil.Process(os.getpid())
 
-def _has_cuda():
-    try:
-        import torch
-        return torch.cuda.is_available()
-    except Exception:
-        return False
+# Detect GPU backend once at import time
+_GPU_BACKEND = None  # 'cuda', 'mps', or None
+try:
+    import torch
+    if torch.cuda.is_available():
+        _GPU_BACKEND = 'cuda'
+    elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
+        _GPU_BACKEND = 'mps'
+except Exception:
+    pass
+
 
 def snapshot_memory():
     """Return (vram_bytes, ram_bytes) current allocation."""
     ram = _process.memory_info().rss
-    if _has_cuda():
-        import torch
-        vram = torch.cuda.memory_allocated()
-    else:
-        vram = 0
+    vram = 0
+    try:
+        if _GPU_BACKEND == 'cuda':
+            import torch
+            vram = torch.cuda.memory_allocated()
+        elif _GPU_BACKEND == 'mps':
+            import torch
+            vram = torch.mps.current_allocated_memory()
+    except Exception:
+        pass
     return vram, ram
 
+
 def reset_peak():
-    """Reset CUDA peak memory stats."""
-    if _has_cuda():
-        import torch
-        torch.cuda.reset_peak_memory_stats()
+    """Reset peak memory stats for the active GPU backend."""
+    try:
+        if _GPU_BACKEND == 'cuda':
+            import torch
+            torch.cuda.reset_peak_memory_stats()
+        elif _GPU_BACKEND == 'mps':
+            # MPS has no peak reset API; this is a no-op.
+            pass
+    except Exception:
+        pass
+
 
 def get_peak():
-    """Return (vram_peak_bytes, ram_bytes). RAM peak is just current RSS (no true peak without sampling)."""
+    """Return (vram_peak_bytes, ram_bytes). RAM peak is just current RSS."""
     ram = _process.memory_info().rss
-    if _has_cuda():
-        import torch
-        vram = torch.cuda.max_memory_allocated()
-    else:
-        vram = 0
+    vram = 0
+    try:
+        if _GPU_BACKEND == 'cuda':
+            import torch
+            vram = torch.cuda.max_memory_allocated()
+        elif _GPU_BACKEND == 'mps':
+            # MPS has no peak tracking; return current allocation as best estimate.
+            import torch
+            vram = torch.mps.current_allocated_memory()
+    except Exception:
+        pass
     return vram, ram