Merge pull request #199 from qiboteam/cupy_matrices

Cupy matrices

.github/workflows/selfhosted.yml

@@ -0,0 +1,23 @@
+# CI workflow that runs on selfhosted GPU
+name: Tests with gpu
+
+on:
+  pull_request:
+    types: [labeled]
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.head_ref }}
+  cancel-in-progress: true
+
+jobs:
+  build:
+    if: contains(join(github.event.pull_request.labels.*.name), 'run-on')
+    uses: qiboteam/workflows/.github/workflows/selfhosted.yml@v1
+    with:
+      used-labels: ${{ toJSON(github.event.pull_request.labels.*.name) }}
+      python-version: "3.10"
+      artifact-url: ${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }}
+      poetry-extras: "--with cuda12,test"
+
+    secrets:
+      repo_token: ${{ secrets.GITHUB_TOKEN }}

pyproject.toml

@@ -22,7 +22,7 @@ classifiers = [
 [tool.poetry.dependencies]
 python = "^3.9,<3.13"
 numba = ">=0.59.0"
-qibo = "^0.2.13"
+qibo = "^0.2.14"
 scipy = "^1.10.1"
 psutil = "^5.9.5"
 cupy-cuda12x = { version = "^13.1.0", optional = true }
@@ -40,9 +40,19 @@ pylint = "^3.0.3"
 pytest-cov = "^4.1.0"
 pytest-env = "^0.8.1"
 
-[tool.poetry.extras]
-cupy = ["cupy-cuda12x"]
-cuquantum = ["cuquantum-python-cu12"]
+[tool.poetry.group.cuda11]
+optional = true
+
+[tool.poetry.group.cuda11.dependencies]
+cupy-cuda11x = "^13.1.0"
+cuquantum-python-cu11 = "^23.3.0"
+
+[tool.poetry.group.cuda12]
+optional = true
+
+[tool.poetry.group.cuda12.dependencies]
+cupy-cuda12x = "^13.1.0"
+cuquantum-python-cu12 = "^23.3.0"
 
 [tool.poe.tasks]
 test = "pytest"
@@ -60,5 +70,5 @@ omit = ["src/qibojit/backends/clifford_operations*"]
 
 [tool.pytest.ini_options]
 testpaths = ['src/qibojit/tests/']
-addopts = ['--cov=qibojit', '--cov-report=xml']
+addopts = ['--cov=qibojit', '--cov-report=xml', '--cov-report=html']
 env = ["D:NUMBA_DISABLE_JIT=1"]

selfhosted

@@ -0,0 +1,19 @@
+#!/bin/bash
+
+# Script for running the selfhosted tests on QPUs directly from GitHub
+# Tests need to be copied to /tmp/ because coverage does not work with NFS
+
+cp -r src/qibojit/tests /tmp/
+cp pyproject.toml /tmp/
+cd /tmp/tests
+source /nfs/users/github/actions-runner/_work/qibojit/qibojit/testenv/bin/activate
+pytest
+pytest_status=$?
+if [[ $pytest_status -ne 0 ]]
+    then
+        exit $pytest_status
+fi
+cd -
+mv /tmp/tests/coverage.xml .
+mv /tmp/tests/htmlcov .
+rm -r /tmp/tests

src/qibojit/backends/clifford_operations_cpu.py

@@ -3,7 +3,7 @@
 import numpy as np
 from numba import njit, prange, uint64
 
-PARALLEL = False
+PARALLEL = True
 
 
 @njit("Tuple((u1[:], u1[:,:], u1[:,:]))(u1[:,:], u8)", parallel=PARALLEL, cache=True)

src/qibojit/backends/cpu.py

@@ -64,6 +64,7 @@ def __init__(self):
             4: self.gates.apply_four_qubit_gate_kernel,
             5: self.gates.apply_five_qubit_gate_kernel,
         }
+
         if sys.platform == "darwin":  # pragma: no cover
             self.set_threads(psutil.cpu_count(logical=False))
         else:

src/qibojit/backends/gpu.py

@@ -5,7 +5,11 @@
 from qibo.config import log, raise_error
 
 from qibojit.backends.cpu import NumbaBackend
-from qibojit.backends.matrices import CupyMatrices, CuQuantumMatrices, CustomMatrices
+from qibojit.backends.matrices import (
+    CupyMatrices,
+    CustomCuQuantumMatrices,
+    CustomMatrices,
+)
 
 
 class CupyBackend(NumbaBackend):  # pragma: no cover
@@ -574,7 +578,7 @@ def __init__(self):
         self.versions["cuquantum"] = self.cuquantum.__version__
         self.supports_multigpu = True
         self.handle = self.cusv.create()
-        self.custom_matrices = CuQuantumMatrices(self.dtype)
+        self.custom_matrices = CustomCuQuantumMatrices(self.dtype)
 
     def __del__(self):
         if hasattr(self, "cusv"):
@@ -584,7 +588,7 @@ def set_precision(self, precision):
         if precision != self.precision:
             super().set_precision(precision)
             if self.custom_matrices:
-                self.custom_matrices = CuQuantumMatrices(self.dtype)
+                self.custom_matrices = CustomCuQuantumMatrices(self.dtype)
 
     def get_cuda_type(self, dtype="complex64"):
         if dtype == "complex128":
@@ -696,7 +700,7 @@ def two_qubit_base(
             nBitSwaps = 1
             bitSwaps = [(target1, target2)]
             maskLen = ncontrols
-            maskBitString = self.np.ones(ncontrols)
+            maskBitString = self.cp.ones(ncontrols)
             maskOrdering = controls
 
             self.cusv.swap_index_bits(

src/qibojit/backends/matrices.py

@@ -4,39 +4,11 @@
 from qibo.backends.npmatrices import NumpyMatrices
 
 
-class CupyMatrices(NumpyMatrices):  # pragma: no cover
-    # Necessary to avoid https://github.com/qiboteam/qibo/issues/928
-    def Unitary(self, u):
-        import cupy as cp  # pylint: disable=import-error
-
-        if isinstance(u, cp.ndarray):
-            u = u.get()
-        return super().Unitary(u)
-
-
-class CuQuantumMatrices(NumpyMatrices):
-    # These matrices are used by the custom operators and may
-    # not correspond to the mathematical representation of each gate
-
-    @cached_property
-    def CNOT(self):
-        return self.X
-
-    @cached_property
-    def CY(self):
-        return self.Y
-
-    @cached_property
-    def CZ(self):
-        return self.Z
-
-    @cached_property
-    def CSX(self):
-        return self.SX
+class CustomMatrices(NumpyMatrices):
+    """Matrices used by custom operators.
 
-    @cached_property
-    def CSXDG(self):
-        return self.SXDG
+    They may not correspond to the mathematical representation of each gate.
+    """
 
     def CRX(self, theta):
         return self.RX(theta)
@@ -47,34 +19,41 @@ def CRY(self, theta):
     def CRZ(self, theta):
         return self.RZ(theta)
 
-    def CU1(self, theta):
-        return self.U1(theta)
-
     def CU2(self, phi, lam):
         return self.U2(phi, lam)
 
     def CU3(self, theta, phi, lam):
         return self.U3(theta, phi, lam)
 
+    def U1(self, theta):
+        dtype = getattr(np, self.dtype)
+        return self._cast(np.exp(1j * theta), dtype=dtype)
+
+    def CU1(self, theta):
+        return self.U1(theta)
+
     @cached_property
-    def TOFFOLI(self):
-        return self.X
+    def CZ(self):
+        return self.Z
 
     @cached_property
     def CCZ(self):
         return self.Z
 
-    def DEUTSCH(self, theta):
-        return 1j * self.RX(2 * theta)
+    @cached_property
+    def CY(self):
+        return self.Y
 
+    @cached_property
+    def CSX(self):
+        return self.SX
 
-class CustomMatrices(CuQuantumMatrices):
-    # These matrices are used by the custom operators and may
-    # not correspond to the mathematical representation of each gate
+    @cached_property
+    def CSXDG(self):
+        return self.SXDG
 
-    def U1(self, theta):
-        dtype = getattr(np, self.dtype)
-        return dtype(np.exp(1j * theta))
+    def DEUTSCH(self, theta):
+        return 1j * self.RX(2 * theta)
 
     def fSim(self, theta, phi):
         cost = np.cos(theta) + 0j
@@ -85,3 +64,69 @@ def fSim(self, theta, phi):
     def GeneralizedfSim(self, u, phi):
         phase = np.exp(-1j * phi)
         return np.array([u[0, 0], u[0, 1], u[1, 0], u[1, 1], phase], dtype=self.dtype)
+
+
+class CupyMatrices(NumpyMatrices):  # pragma: no cover
+    """Casting NumpyMatrices to Cupy arrays."""
+
+    def __init__(self, dtype):
+        super().__init__(dtype)
+        import cupy as cp  # pylint: disable=E0401
+
+        self.cp = cp
+
+    def _cast(self, x, dtype):
+        is_cupy = [
+            isinstance(item, self.cp.ndarray) for sublist in x for item in sublist
+        ]
+        if any(is_cupy) and not all(is_cupy):
+            # for parametrized gates x is a mixed list of cp.arrays and floats
+            # thus a simple cp.array(x) fails
+            # first convert the cp.arrays to numpy, then build the numpy array and move it
+            # back to GPU
+            dim = len(x)
+            return self.cp.array(
+                np.array(
+                    [
+                        item.get() if isinstance(item, self.cp.ndarray) else item
+                        for sublist in x
+                        for item in sublist
+                    ]
+                ).reshape(dim, dim),
+                dtype=dtype,
+            )
+        return self.cp.array(x, dtype=dtype)
+
+    # Necessary to avoid https://github.com/qiboteam/qibo/issues/928
+    def Unitary(self, u):
+        dtype = getattr(np, self.dtype)
+        return self._cast(u, dtype=dtype)
+
+
+class CustomCuQuantumMatrices(CustomMatrices):  # pragma: no cover
+    """Matrices used by CuQuantum custom operators."""
+
+    @cached_property
+    def CNOT(self):
+        return self.X
+
+    @cached_property
+    def TOFFOLI(self):
+        return self.X
+
+    @cached_property
+    def CY(self):
+        return self.Y
+
+    @cached_property
+    def CZ(self):
+        return self.Z
+
+    def U1(self, theta):
+        return NumpyMatrices.U1(self, theta)
+
+    def fSim(self, theta, phi):
+        return NumpyMatrices.fSim(self, theta, phi)
+
+    def GeneralizedfSim(self, u, phi):
+        return NumpyMatrices.GeneralizedfSim(self, u, phi)

src/qibojit/tests/test_backends.py

@@ -3,6 +3,8 @@
 
 from qibojit.backends import MetaBackend
 
+from .conftest import AVAILABLE_BACKENDS, BACKENDS
+
 
 def test_device_setter(backend):
     if backend.platform == "numba":
@@ -20,6 +22,7 @@ def test_thread_setter(backend):
     original_threads = numba.get_num_threads()
     backend.set_threads(1)
     assert numba.get_num_threads() == 1
+    backend.set_threads(original_threads)
 
 
 @pytest.mark.parametrize("array_type", [None, "float32", "float64"])
@@ -119,5 +122,7 @@ def test_backend_eigh_sparse(backend, sparse_type, k):
 
 
 def test_metabackend_list_available():
-    available_backends = dict(zip(("numba", "cupy", "cuquantum"), (True, False, False)))
+    available_backends = {
+        backend: backend in AVAILABLE_BACKENDS for backend in BACKENDS
+    }
     assert MetaBackend().list_available() == available_backends

src/qibojit/tests/test_gates.py

@@ -204,6 +204,24 @@ def test_apply_csx(backend, nqubits, targets, dtype):
     backend.assert_allclose(state, target_state, atol=ATOL.get(dtype))
 
 
+@pytest.mark.parametrize(
+    ("nqubits", "target", "controls"),
+    [
+        (3, 2, [0, 1]),
+        (4, 2, [1, 3]),
+    ],
+)
+def test_apply_ccz(backend, nqubits, target, controls, dtype):
+    tbackend = NumpyBackend()
+    state = random_statevector(2**nqubits, backend=tbackend).astype(dtype)
+    gate = gates.CCZ(*controls, target)
+
+    set_precision(dtype, backend, tbackend)
+    target_state = tbackend.apply_gate(gate, np.copy(state), nqubits)
+    state = backend.apply_gate(gate, np.copy(state), nqubits)
+    backend.assert_allclose(state, target_state, atol=ATOL.get(dtype))
+
+
 @pytest.mark.parametrize(
     ("nqubits", "targets"),
     [