Add apple silicon GPU Acceleration Support ("mps") (#335)

* Add apple silicon GPU Acceleration Support * format fix * update * update * update * Fix module not found error --------- Co-authored-by: rusty1s <[email protected]>
rusty1s · Jul 19, 2023 · 14781cb · 14781cb
1 parent 40693ab
commit 14781cb
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 5 deletions.
diff --git a/benchmark/main.py b/benchmark/main.py
@@ -1,13 +1,13 @@
-import time
-import os.path as osp
+import argparse
 import itertools
+import os.path as osp
+import time
 
-import argparse
-import wget
 import torch
+import wget
 from scipy.io import loadmat
-
 from torch_scatter import scatter_add
+
 from torch_sparse.tensor import SparseTensor
 
 short_rows = [
@@ -62,6 +62,9 @@ def time_func(func, x):
     try:
         if torch.cuda.is_available():
             torch.cuda.synchronize()
+        elif torch.backends.mps.is_available():
+            import torch.mps
+            torch.mps.synchronize()
         t = time.perf_counter()
 
         if not args.with_backward:
@@ -77,6 +80,9 @@ def time_func(func, x):
 
         if torch.cuda.is_available():
             torch.cuda.synchronize()
+        elif torch.backends.mps.is_available():
+            import torch.mps
+            torch.mps.synchronize()
         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')]
+if torch.backends.mps.is_available():
+    devices += [torch.device('mps')]
 
 
 def tensor(x: Any, dtype: torch.dtype, device: torch.device):