Add apple silicon GPU Acceleration Support

rusty1s · Jul 14, 2023 · 15970f3 · 15970f3
1 parent 40693ab
commit 15970f3
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 0 deletions.
diff --git a/benchmark/main.py b/benchmark/main.py
@@ -62,6 +62,12 @@ def time_func(func, x):
     try:
         if torch.cuda.is_available():
             torch.cuda.synchronize()
+        elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
+            try:
+                import torch.mps
+                torch.mps.synchronize()
+            except ImportError:
+                pass
         t = time.perf_counter()
 
         if not args.with_backward:
@@ -77,6 +83,12 @@ def time_func(func, x):
 
         if torch.cuda.is_available():
             torch.cuda.synchronize()
+        elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
+            try:
+                import torch.mps
+                torch.mps.synchronize()
+            except ImportError:
+                pass
         return time.perf_counter() - t
     except RuntimeError as e:
         if 'out of memory' not in str(e):

diff --git a/torch_sparse/testing.py b/torch_sparse/testing.py
@@ -16,6 +16,8 @@
 devices = [torch.device('cpu')]
 if torch.cuda.is_available():
     devices += [torch.device('cuda:0')]
+elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
+    devices += [torch.device('mps')]
 
 
 def tensor(x: Any, dtype: torch.dtype, device: torch.device):