Merge pull request #19 from qiboteam/multigpuops

Add ops required by qibo multi-GPU

src/qibojit/custom_operators/__init__.py

@@ -93,8 +93,19 @@ def collapse_state(state, qubits, result, nqubits, normalize=True):
     return backend.collapse_state(state, qubits, result, nqubits, normalize)
 
 
+def transpose_state(pieces, state, nqubits, order):
+    # always fall back to numba CPU backend because for ops not implemented on GPU
+    numba_backend = backend.get("numba")
+    return numba_backend.transpose_state(pieces, state, nqubits, order)
+
+
+def swap_pieces(piece0, piece1, new_global, nlocal):
+    # always fall back to numba CPU backend because for ops not implemented on GPU
+    numba_backend = backend.get("numba")
+    return numba_backend.swap_pieces(piece0, piece1, new_global, nlocal)
+
+
 def measure_frequencies(frequencies, probs, nshots, nqubits, seed=1234, nthreads=None):
-    # always fall back to numba CPU backend because this op is not implemented
-    # on GPU
+    # always fall back to numba CPU backend because for ops not implemented on GPU
     numba_backend = backend.get("numba")
     return numba_backend.measure_frequencies(frequencies, probs, nshots, nqubits, seed, nthreads)

src/qibojit/custom_operators/backends.py

@@ -32,6 +32,14 @@ def initial_state(self, nqubits, dtype, is_matrix=False): # pragma: no cover
     def collapse_state(self, state, qubits, result, nqubits, normalize): # pragma: no cover
         raise NotImplementedError
 
+    @abstractmethod
+    def transpose_state(self, pieces, state, nqubits, order): # pragma: no cover
+        raise NotImplementedError
+
+    @abstractmethod
+    def swap_pieces(self, piece0, piece1, new_global, nlocal): # pragma: no cover
+        raise NotImplementedError
+
     @abstractmethod
     def measure_frequencies(self, frequencies, probs, nshots, nqubits, seed=1234): # pragma: no cover
         raise NotImplementedError
@@ -101,6 +109,12 @@ def collapse_state(self, state, qubits, result, nqubits, normalize=True):
             return self.ops.collapse_state_normalized(state, qubits, result, nqubits)
         return self.ops.collapse_state(state, qubits, result, nqubits)
 
+    def transpose_state(self, pieces, state, nqubits, order):
+        return self.ops.transpose_state(tuple(pieces), state, nqubits, tuple(order))
+
+    def swap_pieces(self, piece0, piece1, new_global, nlocal):
+        return self.ops.swap_pieces(piece0, piece1, new_global, nlocal)
+
     def measure_frequencies(self, frequencies, probs, nshots, nqubits, seed=1234, nthreads=None):
         if nthreads is None:
             import psutil
@@ -274,6 +288,14 @@ def collapse_state(self, state, qubits, result, nqubits, normalize=True):
             state = state / norm
         return state
 
+    def transpose_state(self, pieces, state, nqubits, order):
+        raise NotImplementedError("`transpose_state` method is not "
+                                  "implemented for GPU.")
+
+    def swap_pieces(self, piece0, piece1, new_global, nlocal):
+        raise NotImplementedError("`swap_pieces` method is not "
+                                  "implemented for GPU.")
+
     def measure_frequencies(self, frequencies, probs, nshots, nqubits, seed=1234):
         raise NotImplementedError("`measure_frequencies` method is not "
                                   "implemented for GPU.")

src/qibojit/custom_operators/ops.py

@@ -90,3 +90,29 @@ def measure_frequencies(frequencies, probs, nshots, nqubits, seed, nthreads):
             frequencies_private[shot] += 1
     frequencies += thread_frequencies.sum(axis=0)
     return frequencies
+
+
+@njit(cache=True, parallel=True)
+def transpose_state(pieces, state, nqubits, order):
+    nstates = 1 << nqubits
+    ndevices = len(pieces)
+    npiece = nstates // ndevices
+    qubit_exponents = [1 << (nqubits - x - 1) for x in order[::-1]]
+
+    for g in prange(nstates): # pylint: disable=not-an-iterable
+        k = 0
+        for q in range(nqubits):
+            if ((g >> q) % 2):
+                k += qubit_exponents[q]
+        state[g] = pieces[k // npiece][k % npiece]
+    return state
+
+
+@njit(cache=True, parallel=True)
+def swap_pieces(piece0, piece1, new_global, nlocal):
+    m = nlocal - new_global - 1
+    tk = 1 << m
+    nstates = 1 << (nlocal - 1)
+    for g in prange(nstates): # pylint: disable=not-an-iterable
+        i = ((g >> m) << (m + 1)) + (g & (tk - 1))
+        piece0[i + tk], piece1[i] = piece1[i], piece0[i + tk]

src/qibojit/tests/test_ops.py

@@ -2,6 +2,7 @@
 import itertools
 import numpy as np
 from qibojit import custom_operators as op
+from qibojit.tests.test_gates import qubits_tensor, random_state
 
 
 @pytest.mark.parametrize("is_matrix", [False, True])
@@ -23,9 +24,7 @@ def test_initial_state(backend, dtype, is_matrix):
 @pytest.mark.parametrize("normalize", [False, True])
 def test_collapse_state(backend, nqubits, targets, results, normalize, dtype):
     atol = 1e-7 if dtype == "complex64" else 1e-14
-    shape = (2 ** nqubits,)
-    state = np.random.random(shape) + 1j * np.random.random(shape)
-    state = state.astype(dtype)
+    state = random_state(nqubits, dtype)
     slicer = nqubits * [slice(None)]
     for t, r in zip(targets, results):
         slicer[t] = r
@@ -46,6 +45,77 @@ def test_collapse_state(backend, nqubits, targets, results, normalize, dtype):
     np.testing.assert_allclose(state, target_state, atol=atol)
 
 
+def generate_transpose_qubits(nqubits):
+    """Generates global qubits randomly."""
+    qubits = np.arange(nqubits)
+    np.random.shuffle(qubits)
+    return qubits
+
+
+CONFIG = ((n, generate_transpose_qubits(n))
+          for _ in range(5) for n in range(3, 11))
+@pytest.mark.parametrize("nqubits,qubits", CONFIG)
+@pytest.mark.parametrize("ndevices", [2, 4, 8])
+def test_transpose_state(nqubits, qubits, ndevices, dtype):
+    qubit_order = list(qubits)
+    state = random_state(nqubits, dtype)
+    state_tensor = np.reshape(state, nqubits * (2,))
+    target_state = np.transpose(state_tensor, qubit_order).flatten()
+    new_state = np.zeros_like(state)
+    state = np.reshape(state, (ndevices, int(state.shape[0]) // ndevices))
+    pieces = [state[i] for i in range(ndevices)]
+    new_state = op.transpose_state(pieces, new_state, nqubits, qubit_order)
+    np.testing.assert_allclose(new_state, target_state)
+
+
+CONFIG = ((n, np.random.randint(1, n)) for _ in range(10) for n in range(4, 11))
+@pytest.mark.parametrize("nqubits,local", CONFIG)
+def test_swap_pieces_zero_global(nqubits, local, dtype):
+    state = random_state(nqubits, dtype)
+    target_state = np.copy(state)
+    shape = (2, int(state.shape[0]) // 2)
+    state = np.reshape(state, shape)
+
+    qubits = qubits_tensor(nqubits, [0, local])
+    target_state = op.apply_swap(target_state, nqubits, 0, local, qubits)
+    target_state = np.reshape(op.to_numpy(target_state), shape)
+    piece0, piece1 = state[0], state[1]
+    op.swap_pieces(piece0, piece1, local - 1, nqubits - 1)
+    np.testing.assert_allclose(piece0, target_state[0])
+    np.testing.assert_allclose(piece1, target_state[1])
+
+
+CONFIG = ((n, np.random.randint(0, n), np.random.randint(0, n))
+          for _ in range(10) for n in range(5, 11))
+@pytest.mark.parametrize("nqubits,qlocal,qglobal", CONFIG)
+def test_swap_pieces(nqubits, qlocal, qglobal, dtype):
+    state = random_state(nqubits, dtype)
+    target_state = np.copy(state)
+    shape = (2, int(state.shape[0]) // 2)
+
+    while qlocal == qglobal:
+        qlocal = np.random.randint(0, nqubits)
+
+    transpose_order = ([qglobal] + list(range(qglobal)) +
+                        list(range(qglobal + 1, nqubits)))
+
+    qubits = qubits_tensor(nqubits, [qglobal, qlocal])
+    target_state = op.apply_swap(target_state, nqubits, qglobal, qlocal, qubits)
+    target_state = op.to_numpy(target_state)
+    target_state = np.reshape(target_state, nqubits * (2,))
+    target_state = np.transpose(target_state, transpose_order)
+    target_state = np.reshape(target_state, shape)
+
+    state = np.reshape(state, nqubits * (2,))
+    state = np.transpose(state, transpose_order)
+    state = np.reshape(state, shape)
+    piece0, piece1 = state[0], state[1]
+    new_global = qlocal - int(qglobal < qlocal)
+    op.swap_pieces(piece0, piece1, new_global, nqubits - 1)
+    np.testing.assert_allclose(piece0, target_state[0])
+    np.testing.assert_allclose(piece1, target_state[1])
+
+
 @pytest.mark.parametrize("realtype", ["float32", "float64"])
 @pytest.mark.parametrize("inttype", ["int32", "int64"])
 @pytest.mark.parametrize("nthreads", [None, 4])