cuStateVec bugfixes and slight refactor (#64)

* Updated argument names for new API (cuStateVec 1.6.0)

* Bumped up version requirements

* Fixing a bug in CuStateVec in the implementation of the RYY gate

* Refactor of state.vector so that it returns a numpy array for easier testing and integration

* Adding some of the changes from Dominic's PR #58

* Added more bindings requested by @dlucchetti

pyproject.toml

@@ -59,7 +59,8 @@ simulators = [
     "projectq>=0.5.0,<0.9.0",
 ]
 cuda = [
-    "cuquantum-python>=23.6.0",
+    "cuquantum-python>=24.03.0",
+    "custatevec>=1.6.0",
     "cupy>=10.4.0",
 ]
 wasmtime = [

python/pecos/simulators/custatevec/bindings.py

@@ -11,15 +11,23 @@
 
 import pecos.simulators.custatevec.gates_one_qubit as one_q
 import pecos.simulators.custatevec.gates_two_qubit as two_q
-from pecos.simulators.custatevec.gates_init import init_zero
+from pecos.simulators.custatevec.gates_init import init_one, init_zero
 from pecos.simulators.custatevec.gates_meas import meas_z
 
 # Supporting gates from table:
 #   https://github.com/CQCL/phir/blob/main/spec.md#table-ii---quantum-operations
 
 gate_dict = {
     "Init": init_zero,
+    "init |0>": init_zero,
+    "init |1>": init_one,
     "Measure": meas_z,
+    "measure Z": meas_z,
+    "leak": init_zero,
+    "leak |0>": init_zero,
+    "leak |1>": init_one,
+    "unleak |0>": init_zero,
+    "unleak |1>": init_one,
     "I": one_q.identity,
     "X": one_q.X,
     "Y": one_q.Y,
@@ -28,12 +36,19 @@
     "RY": one_q.RY,
     "RZ": one_q.RZ,
     "R1XY": one_q.R1XY,
+    "RXY1Q": one_q.R1XY,
     "SX": one_q.SX,
     "SXdg": one_q.SXdg,
+    "SqrtX": one_q.SX,
+    "SqrtXd": one_q.SXdg,
     "SY": one_q.SY,
     "SYdg": one_q.SYdg,
+    "SqrtY": one_q.SY,
+    "SqrtYd": one_q.SYdg,
     "SZ": one_q.SZ,
     "SZdg": one_q.SZdg,
+    "SqrtZ": one_q.SZ,
+    "SqrtZd": one_q.SZdg,
     "H": one_q.H,
     "F": one_q.F,
     "Fdg": one_q.Fdg,
@@ -51,6 +66,7 @@
     "SYY": two_q.SYY,
     "SYYdg": two_q.SYYdg,
     "SZZ": two_q.SZZ,
+    "SqrtZZ": two_q.SZZ,
     "SZZdg": two_q.SZZdg,
     "SWAP": two_q.SWAP,
 }

python/pecos/simulators/custatevec/gates_init.py

@@ -26,3 +26,16 @@ def init_zero(state, qubit: int, **params: Any) -> None:
 
     if result:
         X(state, qubit)
+
+
+def init_one(state, qubit: int, **params: Any) -> None:
+    """Initialise or reset the qubit to state |1>
+
+    Args:
+        state: An instance of CuStateVec
+        qubit: The index of the qubit to be initialised
+    """
+    result = meas_z(state, qubit)
+
+    if not result:
+        X(state, qubit)

python/pecos/simulators/custatevec/gates_meas.py

@@ -36,12 +36,12 @@ def meas_z(state, qubit: int, **params: Any) -> int:
 
     result = cusv.measure_on_z_basis(
         handle=state.libhandle,
-        sv=state.vector.data.ptr,
+        sv=state.cupy_vector.data.ptr,
         sv_data_type=state.cuda_type,
         n_index_bits=state.num_qubits,  # Number of qubits in the statevector
         basis_bits=[target],  # The index of the qubit being measured
         n_basis_bits=1,  # Number of qubits being measured
-        rand_num=np.random.random(),  # Source of randomness for the measurement
+        randnum=np.random.random(),  # Source of randomness for the measurement
         collapse=cusv.Collapse.NORMALIZE_AND_ZERO,  # Collapse and normalise
     )
     state.stream.synchronize()

python/pecos/simulators/custatevec/gates_one_qubit.py

@@ -34,7 +34,7 @@ def _apply_one_qubit_matrix(state, qubit: int, matrix: cp.ndarray) -> None:
 
     cusv.apply_matrix(
         handle=state.libhandle,
-        sv=state.vector.data.ptr,
+        sv=state.cupy_vector.data.ptr,
         sv_data_type=state.cuda_type,
         n_index_bits=state.num_qubits,
         matrix=matrix.data.ptr,
@@ -47,8 +47,8 @@ def _apply_one_qubit_matrix(state, qubit: int, matrix: cp.ndarray) -> None:
         control_bit_values=[],  # No value of control bit assigned
         n_controls=0,
         compute_type=state.compute_type,
-        workspace=0,  # Let cuQuantum use the mempool we configured
-        workspace_size=0,  # Let cuQuantum use the mempool we configured
+        extra_workspace=0,  # Let cuQuantum use the mempool we configured
+        extra_workspace_size_in_bytes=0,  # Let cuQuantum use the mempool we configured
     )
     state.stream.synchronize()
 

python/pecos/simulators/custatevec/gates_two_qubit.py

@@ -39,7 +39,7 @@ def _apply_controlled_matrix(state, control: int, target: int, matrix: cp.ndarra
 
     cusv.apply_matrix(
         handle=state.libhandle,
-        sv=state.vector.data.ptr,
+        sv=state.cupy_vector.data.ptr,
         sv_data_type=state.cuda_type,
         n_index_bits=state.num_qubits,
         matrix=matrix.data.ptr,
@@ -52,8 +52,8 @@ def _apply_controlled_matrix(state, control: int, target: int, matrix: cp.ndarra
         control_bit_values=[],  # No value of control bit assigned
         n_controls=1,
         compute_type=state.compute_type,
-        workspace=0,  # Let cuQuantum use the mempool we configured
-        workspace_size=0,  # Let cuQuantum use the mempool we configured
+        extra_workspace=0,  # Let cuQuantum use the mempool we configured
+        extra_workspace_size_in_bytes=0,  # Let cuQuantum use the mempool we configured
     )
     state.stream.synchronize()
 
@@ -142,7 +142,7 @@ def _apply_two_qubit_matrix(state, qubits: tuple[int, int], matrix: cp.ndarray)
 
     cusv.apply_matrix(
         handle=state.libhandle,
-        sv=state.vector.data.ptr,
+        sv=state.cupy_vector.data.ptr,
         sv_data_type=state.cuda_type,
         n_index_bits=state.num_qubits,
         matrix=matrix.data.ptr,
@@ -155,8 +155,8 @@ def _apply_two_qubit_matrix(state, qubits: tuple[int, int], matrix: cp.ndarray)
         control_bit_values=[],  # No value of control bit assigned
         n_controls=0,
         compute_type=state.compute_type,
-        workspace=0,  # Let cuQuantum use the mempool we configured
-        workspace_size=0,  # Let cuQuantum use the mempool we configured
+        extra_workspace=0,  # Let cuQuantum use the mempool we configured
+        extra_workspace_size_in_bytes=0,  # Let cuQuantum use the mempool we configured
     )
     state.stream.synchronize()
 
@@ -218,7 +218,7 @@ def RYY(state, qubits: tuple[int, int], angles: tuple[float], **params: Any) ->
             math.cos(theta / 2),
             0,
             0,
-            math.sin(theta / 2),
+            1j * math.sin(theta / 2),
             0,
             math.cos(theta / 2),
             -1j * math.sin(theta / 2),
@@ -227,7 +227,7 @@ def RYY(state, qubits: tuple[int, int], angles: tuple[float], **params: Any) ->
             -1j * math.sin(theta / 2),
             math.cos(theta / 2),
             0,
-            math.sin(theta / 2),
+            1j * math.sin(theta / 2),
             0,
             0,
             math.cos(theta / 2),
@@ -380,7 +380,7 @@ def SWAP(state, qubits: tuple[int, int], **params: Any) -> None:
     # Possibly faster since it may just be an internal qubit relabelling or sv reshape
     cusv.apply_generalized_permutation_matrix(
         handle=state.libhandle,
-        sv=state.vector.data.ptr,
+        sv=state.cupy_vector.data.ptr,
         sv_data_type=state.cuda_type,
         n_index_bits=state.num_qubits,
         permutation=[0, 2, 1, 3],  # Leave |00> and |11> where they are, swap the other two
@@ -392,7 +392,7 @@ def SWAP(state, qubits: tuple[int, int], **params: Any) -> None:
         controls=[],
         control_bit_values=[],  # No value of control bit assigned
         n_controls=0,
-        workspace=0,  # Let cuQuantum use the mempool we configured
-        workspace_size=0,  # Let cuQuantum use the mempool we configured
+        extra_workspace=0,  # Let cuQuantum use the mempool we configured
+        extra_workspace_size_in_bytes=0,  # Let cuQuantum use the mempool we configured
     )
     state.stream.synchronize()

python/pecos/simulators/custatevec/state.py

@@ -12,6 +12,7 @@
 import random
 
 import cupy as cp
+import numpy as np
 from cuquantum import ComputeType, cudaDataType
 from cuquantum import custatevec as cusv
 
@@ -52,8 +53,8 @@ def __init__(self, num_qubits, seed=None) -> None:
         self.compute_type = ComputeType.COMPUTE_64F
 
         # Allocate the statevector in GPU and initialize it to |0>
-        self.vector = cp.zeros(shape=2**num_qubits, dtype=self.cp_type)
-        self.vector[0] = 1
+        self.cupy_vector = cp.zeros(shape=2**num_qubits, dtype=self.cp_type)
+        self.cupy_vector[0] = 1
 
         ####################################################
         # Set up cuStateVec library and GPU memory handles #
@@ -97,7 +98,18 @@ def free(ptr, size, stream):
         mem_handler = (malloc, free, "GPU memory handler")
         cusv.set_device_mem_handler(self.libhandle, mem_handler)
 
+    def reset(self):
+        """Reset the quantum state for another run without reinitializing."""
+        # Initialize all qubits in the zero state
+        self.vector = cp.zeros(shape=2**self.num_qubits, dtype=self.cp_type)
+        self.vector[0] = 1
+        return self
+
     def __del__(self) -> None:
-        # CuPy will release GPU memory when the variable ``self.vector`` is no longer
+        # CuPy will release GPU memory when the variable ``self.cupy_vector`` is no longer
         # reachable. However, we need to manually destroy the library handle.
         cusv.destroy(self.libhandle)
+
+    @property
+    def vector(self) -> np.ndarray:
+        return self.cupy_vector.get()

tests/integration/state_sim_tests/test_statevec.py

@@ -21,7 +21,7 @@
 try:
     import cuquantum
 
-    imported_cuquantum = vparse(cuquantum._version.__version__) >= vparse("23.6.0")  # noqa: SLF001
+    imported_cuquantum = vparse(cuquantum._version.__version__) >= vparse("24.03.0")  # noqa: SLF001
     import cupy as cp
 
     imported_cupy = vparse(version("cupy")) >= vparse("10.4.0")