experimental support for CUDA

p_kit/solver/annealing.py

@@ -12,6 +12,10 @@ def job():
         _, time_points, _ = solver.copy().solve(circuit.copy(), annealing_func=annealing)
         return time_points[-n_last_samples-1:-1, :]
 
-    samples = Parallel(n_jobs=n_jobs)(delayed(job)() for _ in range(n_shots))
+    if getattr(solver, 'device', 'cpu') == 'cuda':
+        all_m = solver.solve(circuit, annealing_func=annealing, n_shots=n_shots)
+        samples = all_m[-n_last_samples-1:-1]  # (n_last_samples, n_shots, n_pbits)
+        return samples.transpose(1, 0, 2).reshape(n_shots * n_last_samples, circuit.n_pbits)
 
+    samples = Parallel(n_jobs=n_jobs)(delayed(job)() for _ in range(n_shots))
     return np.array(samples).reshape((n_shots * n_last_samples, circuit.n_pbits))

p_kit/solver/base_solver.py

@@ -1,21 +1,32 @@
-import random
-import time
 import numpy as np
 
+try:
+    import cupy as cp
+except ImportError:
+    cp = None
+
+
 class Solver:
-    def __init__(self, Nt, dt, i0, expected_mean=0, seed=None) -> None:
+    def __init__(self, Nt, dt, i0, expected_mean=0, seed=None, device='cpu') -> None:
         self.Nt = Nt
         self.dt = dt
         self.i0 = i0
         self.expected_mean = expected_mean
         self.seed = seed
-        self._random_gen = np.random.default_rng(self.seed)
-
+        self.device = device
+        if device == 'cuda':
+            if cp is None:
+                raise ImportError("cupy is required for device='cuda'. Install with: pip install cupy-cuda13x")
+            self.xp = cp
+        else:
+            self.xp = np
+        self._random_gen = self.xp.random.default_rng(self.seed)
+
     def random(self, n_pbits):
         return self._random_gen.random(n_pbits)
 
-    def solve(self, annealing_func):
+    def solve(self, c, annealing_func=None, n_shots=1):
         raise NotImplementedError()
-    
+
     def copy(self):
         raise NotImplementedError()

p_kit/solver/csd_solver.py

@@ -5,38 +5,38 @@
 
 
 class CaSuDaSolver(Solver):
-    # K. Y. Camsari, B. M. Sutton, and S. Datta, ‘p-bits for probabilistic spin logic’, Applied Physics Reviews, vol. 6, no. 1, p. 011305, Mar. 2019, doi: 10.1063/1.5055860.
+    # K. Y. Camsari, B. M. Sutton, and S. Datta, 'p-bits for probabilistic spin logic', Applied Physics Reviews, vol. 6, no. 1, p. 011305, Mar. 2019, doi: 10.1063/1.5055860.
+
+    def solve(self, c: PCircuit, annealing_func=constant, n_shots=1):
 
-    def solve(self, c: PCircuit, annealing_func=constant):
-
         # credit: https://www.purdue.edu/p-bit/blog.html
+        xp = self.xp
         n_pbits = c.n_pbits
 
-        all_I = np.zeros((self.Nt, n_pbits))
-        all_m = np.zeros((self.Nt, n_pbits))
-        E = np.zeros(self.Nt)
-
-        m = np.sign(0.5 - self.random(n_pbits))
+        J = xp.asarray(c.J)
+        h = xp.asarray(c.h)
+        threshold = float(np.arctanh(self.expected_mean))
 
-        threshold = np.arctanh(self.expected_mean)
+        # m is (n_shots, n_pbits) — works for n_shots=1 too
+        all_m = xp.zeros((self.Nt, n_shots, n_pbits))
+        all_I = xp.zeros((self.Nt, n_pbits))
+        E = xp.zeros(self.Nt)
+        m = xp.sign(0.5 - self.random((n_shots, n_pbits)))
 
         for run in range(self.Nt):
-
-            # compute input biases
-            I = annealing_func(self, run) * (np.dot(m, c.J) + c.h)
-
-            # apply S(input)
-            s = np.exp(-self.dt * np.exp(-m * (I + threshold)))
-
-            # compute new output
-            m = m * np.sign(s - self.random(n_pbits))
-
-            all_I[run] = I
+            I = annealing_func(self, run) * (m @ J + h)
+            s = xp.exp(-self.dt * xp.exp(-m * (I + threshold)))
+            m = m * xp.sign(s - self.random((n_shots, n_pbits)))
             all_m[run] = m
+            all_I[run] = I[0]
+            E[run] = self.i0 * (xp.dot(m[0], h) + 0.5 * xp.dot(xp.dot(m[0], J), m[0]))
 
-            E[run] = self.i0 * (np.dot(m, c.h) + 0.5 * np.dot(np.dot(m, c.J), m))
+        if self.device == 'cuda':
+            all_I, all_m, E = all_I.get(), all_m.get(), E.get()
 
-        return all_I, all_m, E
+        if n_shots == 1:
+            return all_I, all_m[:, 0, :], E
+        return all_m
 
     def copy(self):
-        return CaSuDaSolver(self.Nt, self.dt, self.i0, self.expected_mean, self.seed)
+        return CaSuDaSolver(self.Nt, self.dt, self.i0, self.expected_mean, self.seed, self.device)

setup.py

@@ -44,6 +44,7 @@
                         ],
       extras_require={
                       'tests': ['pytest', 'seaborn', 'flake8'],
+                      'gpu': ['cupy-cuda13x'],
                       },
       zip_safe=False,
 )