Merge branch 'main' into mps-for-quimb

.github/workflows/rules.yml

@@ -8,7 +8,7 @@ on:
 
 jobs:
   build:
-    if: contains(github.event.pull_request.labels.*.name, 'run-workflow') || github.event_name == 'push'
+    if: contains(github.event.pull_request.labels.*.name, 'run-workflow') || github.event_name == 'push' && {{ $CUDA_PATH != '' }}
     strategy:
       matrix:
         os: [ubuntu-latest]

src/qibotn/MPSUtils.py

@@ -0,0 +1,82 @@
+import cupy as cp
+from cuquantum.cutensornet.experimental import contract_decompose
+from cuquantum import contract
+
+
+def initial(num_qubits, dtype):
+    """
+    Generate the MPS with an initial state of |00...00>
+    """
+    state_tensor = cp.asarray([1, 0], dtype=dtype).reshape(1, 2, 1)
+    mps_tensors = [state_tensor] * num_qubits
+    return mps_tensors
+
+
+def mps_site_right_swap(mps_tensors, i, **kwargs):
+    """
+    Perform the swap operation between the ith and i+1th MPS tensors.
+    """
+    # contraction followed by QR decomposition
+    a, _, b = contract_decompose(
+        "ipj,jqk->iqj,jpk",
+        *mps_tensors[i : i + 2],
+        algorithm=kwargs.get("algorithm", None),
+        options=kwargs.get("options", None)
+    )
+    mps_tensors[i : i + 2] = (a, b)
+    return mps_tensors
+
+
+def apply_gate(mps_tensors, gate, qubits, **kwargs):
+    """
+    Apply the gate operand to the MPS tensors in-place.
+
+    Args:
+        mps_tensors: A list of rank-3 ndarray-like tensor objects.
+            The indices of the ith tensor are expected to be the bonding index to the i-1 tensor,
+            the physical mode, and then the bonding index to the i+1th tensor.
+        gate: A ndarray-like tensor object representing the gate operand.
+            The modes of the gate is expected to be output qubits followed by input qubits, e.g,
+            ``A, B, a, b`` where ``a, b`` denotes the inputs and ``A, B`` denotes the outputs.
+        qubits: A sequence of integers denoting the qubits that the gate is applied onto.
+        algorithm: The contract and decompose algorithm to use for gate application.
+            Can be either a `dict` or a `ContractDecomposeAlgorithm`.
+        options: Specify the contract and decompose options.
+
+    Returns:
+        The updated MPS tensors.
+    """
+
+    n_qubits = len(qubits)
+    if n_qubits == 1:
+        # single-qubit gate
+        i = qubits[0]
+        mps_tensors[i] = contract(
+            "ipj,qp->iqj", mps_tensors[i], gate, options=kwargs.get("options", None)
+        )  # in-place update
+    elif n_qubits == 2:
+        # two-qubit gate
+        i, j = qubits
+        if i > j:
+            # swap qubits order
+            return apply_gate(mps_tensors, gate.transpose(1, 0, 3, 2), (j, i), **kwargs)
+        elif i + 1 == j:
+            # two adjacent qubits
+            a, _, b = contract_decompose(
+                "ipj,jqk,rspq->irj,jsk",
+                *mps_tensors[i : i + 2],
+                gate,
+                algorithm=kwargs.get("algorithm", None),
+                options=kwargs.get("options", None)
+            )
+            mps_tensors[i : i + 2] = (a, b)  # in-place update
+        else:
+            # non-adjacent two-qubit gate
+            # step 1: swap i with i+1
+            mps_site_right_swap(mps_tensors, i, **kwargs)
+            # step 2: apply gate to (i+1, j) pair. This amounts to a recursive swap until the two qubits are adjacent
+            apply_gate(mps_tensors, gate, (i + 1, j), **kwargs)
+            # step 3: swap back i and i+1
+            mps_site_right_swap(mps_tensors, i, **kwargs)
+    else:
+        raise NotImplementedError("Only one- and two-qubit gates supported")

src/qibotn/QiboCircuitConvertor.py

@@ -6,7 +6,7 @@ class QiboCircuitToEinsum:
     """Convert a circuit to a Tensor Network (TN) representation.
     The circuit is first processed to an intermediate form by grouping each gate
     matrix with its corresponding qubit it is acting on to a list. It is then
-    converted it to an equivalent TN expression through the class function
+    converted to an equivalent TN expression through the class function
     state_vector_operands() following the Einstein summation convention in the
     interleave format.
 
@@ -44,8 +44,7 @@ class QiboCircuitToEinsum:
         for key in qubits_frontier:
             out_list.append(qubits_frontier[key])
 
-        operand_exp_interleave = [x for y in zip(
+        operand_exp_interleave = [x for y in zip(operands, mode_labels) for x in y]
-            operands, mode_labels) for x in y]
         operand_exp_interleave.append(out_list)
         return operand_exp_interleave
 
@@ -95,7 +94,8 @@ class QiboCircuitToEinsum:
             required_shape = self.op_shape_from_qubits(len(gate_qubits))
             self.gate_tensors.append(
                 (
-                    cp.asarray(gate.matrix).reshape(required_shape),
+                    cp.asarray(gate.matrix(), dtype=self.dtype).reshape(
+                        required_shape),
                     gate_qubits,
                 )
             )

src/qibotn/QiboCircuitToMPS.py

@@ -0,0 +1,38 @@
+import cupy as cp
+import numpy as np
+
+from cuquantum import cutensornet as cutn
+from qibotn.QiboCircuitConvertor import QiboCircuitToEinsum
+from qibotn.MPSUtils import initial, apply_gate
+
+
+class QiboCircuitToMPS:
+    def __init__(
+        self,
+        circ_qibo,
+        gate_algo,
+        dtype="complex128",
+        rand_seed=0,
+    ):
+        np.random.seed(rand_seed)
+        cp.random.seed(rand_seed)
+
+        self.num_qubits = circ_qibo.nqubits
+        self.handle = cutn.create()
+        self.dtype = dtype
+        self.mps_tensors = initial(self.num_qubits, dtype=dtype)
+        circuitconvertor = QiboCircuitToEinsum(circ_qibo)
+
+        for gate, qubits in circuitconvertor.gate_tensors:
+            # mapping from qubits to qubit indices
+            # apply the gate in-place
+            apply_gate(
+                self.mps_tensors,
+                gate,
+                qubits,
+                algorithm=gate_algo,
+                options={"handle": self.handle},
+            )
+
+    def __del__(self):
+        cutn.destroy(self.handle)

src/qibotn/cutn.py

@@ -1,8 +1,59 @@
-# from qibotn import quimb as qiboquimb
 from qibotn.QiboCircuitConvertor import QiboCircuitToEinsum
 from cuquantum import contract
+from cuquantum import cutensornet as cutn
+import multiprocessing
+from cupy.cuda.runtime import getDeviceCount
+import cupy as cp
+
+from qibotn.QiboCircuitToMPS import QiboCircuitToMPS
+from qibotn.mps_contraction_helper import MPSContractionHelper
 
 
 def eval(qibo_circ, datatype):
     myconvertor = QiboCircuitToEinsum(qibo_circ, dtype=datatype)
     return contract(*myconvertor.state_vector_operands())
+
+
+def eval_tn_MPI(qibo_circ, datatype, n_samples=8):
+    """Convert qibo circuit to tensornet (TN) format and perform contraction using multi node and multi GPU through MPI.
+    The conversion is performed by QiboCircuitToEinsum(), after which it goes through 2 steps: pathfinder and execution.
+    The pathfinder looks at user defined number of samples (n_samples) iteratively to select the least costly contraction path. This is sped up with multi thread.
+    After pathfinding the optimal path is used in the actual contraction to give a dense vector representation of the TN.
+    """
+
+    from mpi4py import MPI  # this line initializes MPI
+
+    ncpu_threads = multiprocessing.cpu_count() // 2
+
+    comm = MPI.COMM_WORLD
+    rank = comm.Get_rank()
+    device_id = rank % getDeviceCount()
+    cp.cuda.Device(device_id).use()
+
+    handle = cutn.create()
+    cutn.distributed_reset_configuration(handle, *cutn.get_mpi_comm_pointer(comm))
+    network_opts = cutn.NetworkOptions(handle=handle, blocking="auto")
+
+    # Perform circuit conversion
+    myconvertor = QiboCircuitToEinsum(qibo_circ, dtype=datatype)
+    operands_interleave = myconvertor.state_vector_operands()
+
+    # Pathfinder: To search for the optimal path. Optimal path are assigned to path and info attribute of the network object.
+    network = cutn.Network(*operands_interleave, options=network_opts)
+    network.contract_path(optimize={"samples": n_samples, "threads": ncpu_threads})
+
+    # Execution: To execute the contraction using the optimal path found previously
+    result = network.contract()
+
+    cutn.destroy(handle)
+
+    return result, rank
+
+  
+def eval_mps(qibo_circ, gate_algo, datatype):
+    myconvertor = QiboCircuitToMPS(qibo_circ, gate_algo, dtype=datatype)
+    mps_helper = MPSContractionHelper(myconvertor.num_qubits)
+
+    return mps_helper.contract_state_vector(
+        myconvertor.mps_tensors, {"handle": myconvertor.handle}
+    )

src/qibotn/mps_contraction_helper.py

@@ -0,0 +1,129 @@
+from cuquantum import contract, contract_path, CircuitToEinsum, tensor
+
+
+class MPSContractionHelper:
+    """
+    A helper class to compute various quantities for a given MPS.
+
+    Interleaved format is used to construct the input args for `cuquantum.contract`.
+    A concrete example on how the modes are populated for a 7-site MPS is provided below:
+
+          0     2     4     6     8    10     12    14
+    bra -----A-----B-----C-----D-----E-----F-----G-----
+             |     |     |     |     |     |     |
+            1|    3|    5|    7|    9|   11|   13|
+             |     |     |     |     |     |     |
+    ket -----a-----b-----c-----d-----e-----f-----g-----
+          15    16    17    18    19    20    21    22
+
+
+    The follwing compute quantities are supported:
+
+        - the norm of the MPS.
+        - the equivalent state vector from the MPS.
+        - the expectation value for a given operator.
+        - the equivalent state vector after multiplying an MPO to an MPS.
+
+    Note that for the nth MPS tensor (rank-3), the modes of the tensor are expected to be `(i,p,j)`
+    where i denotes the bonding mode with the (n-1)th tensor, p denotes the physical mode for the qubit and
+    j denotes the bonding mode with the (n+1)th tensor.
+
+    Args:
+        num_qubits: The number of qubits for the MPS.
+    """
+
+    def __init__(self, num_qubits):
+        self.num_qubits = num_qubits
+        self.bra_modes = [(2 * i, 2 * i + 1, 2 * i + 2) for i in range(num_qubits)]
+        offset = 2 * num_qubits + 1
+        self.ket_modes = [
+            (i + offset, 2 * i + 1, i + 1 + offset) for i in range(num_qubits)
+        ]
+
+    def contract_norm(self, mps_tensors, options=None):
+        """
+        Contract the corresponding tensor network to form the norm of the MPS.
+
+        Args:
+            mps_tensors: A list of rank-3 ndarray-like tensor objects.
+                The indices of the ith tensor are expected to be bonding index to the i-1 tensor,
+                the physical mode, and then the bonding index to the i+1th tensor.
+            options: Specify the contract and decompose options.
+
+        Returns:
+            The norm of the MPS.
+        """
+        interleaved_inputs = []
+        for i, o in enumerate(mps_tensors):
+            interleaved_inputs.extend(
+                [o, self.bra_modes[i], o.conj(), self.ket_modes[i]]
+            )
+        interleaved_inputs.append([])  # output
+        return self._contract(interleaved_inputs, options=options).real
+
+    def contract_state_vector(self, mps_tensors, options=None):
+        """
+        Contract the corresponding tensor network to form the state vector representation of the MPS.
+
+        Args:
+            mps_tensors: A list of rank-3 ndarray-like tensor objects.
+                The indices of the ith tensor are expected to be bonding index to the i-1 tensor,
+                the physical mode, and then the bonding index to the i+1th tensor.
+            options: Specify the contract and decompose options.
+
+        Returns:
+            An ndarray-like object as the state vector.
+        """
+        interleaved_inputs = []
+        for i, o in enumerate(mps_tensors):
+            interleaved_inputs.extend([o, self.bra_modes[i]])
+        output_modes = tuple([bra_modes[1] for bra_modes in self.bra_modes])
+        interleaved_inputs.append(output_modes)  # output
+        return self._contract(interleaved_inputs, options=options)
+
+    def contract_expectation(
+        self, mps_tensors, operator, qubits, options=None, normalize=False
+    ):
+        """
+        Contract the corresponding tensor network to form the state vector representation of the MPS.
+
+        Args:
+            mps_tensors: A list of rank-3 ndarray-like tensor objects.
+                The indices of the ith tensor are expected to be bonding index to the i-1 tensor,
+                the physical mode, and then the bonding index to the i+1th tensor.
+            operator: A ndarray-like tensor object.
+                The modes of the operator are expected to be output qubits followed by input qubits, e.g,
+                ``A, B, a, b`` where `a, b` denotes the inputs and `A, B'` denotes the outputs.
+            qubits: A sequence of integers specifying the qubits that the operator is acting on.
+            options: Specify the contract and decompose options.
+            normalize: Whether to scale the expectation value by the normalization factor.
+
+        Returns:
+            An ndarray-like object as the state vector.
+        """
+
+        interleaved_inputs = []
+        extra_mode = 3 * self.num_qubits + 2
+        operator_modes = [None] * len(qubits) + [self.bra_modes[q][1] for q in qubits]
+        qubits = list(qubits)
+        for i, o in enumerate(mps_tensors):
+            interleaved_inputs.extend([o, self.bra_modes[i]])
+            k_modes = self.ket_modes[i]
+            if i in qubits:
+                k_modes = (k_modes[0], extra_mode, k_modes[2])
+                q = qubits.index(i)
+                operator_modes[q] = extra_mode  # output modes
+                extra_mode += 1
+            interleaved_inputs.extend([o.conj(), k_modes])
+        interleaved_inputs.extend([operator, tuple(operator_modes)])
+        interleaved_inputs.append([])  # output
+        if normalize:
+            norm = self.contract_norm(mps_tensors, options=options)
+        else:
+            norm = 1
+        return self._contract(interleaved_inputs, options=options) / norm
+
+    def _contract(self, interleaved_inputs, options=None):
+        path = contract_path(*interleaved_inputs, options=options)[0]
+
+        return contract(*interleaved_inputs, options=options, optimize={"path": path})

tests/test_cuquantum_cutensor_backend.py

@@ -2,6 +2,7 @@ from timeit import default_timer as timer
 
 import config
 import numpy as np
+import cupy as cp
 import pytest
 import qibo
 from qibo.models import QFT
@@ -46,3 +47,40 @@ def test_eval(nqubits: int, dtype="complex128"):
     assert 1e-2 * qibo_time < cutn_time < 1e2 * qibo_time
     assert np.allclose(
         result_sv, result_tn), "Resulting dense vectors do not match"
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize("nqubits", [2, 5, 10])
+def test_mps(nqubits: int, dtype="complex128"):
+    """Evaluate MPS with cuQuantum.
+
+    Args:
+        nqubits (int): Total number of qubits in the system.
+        dtype (str): The data type for precision, 'complex64' for single,
+            'complex128' for double.
+    """
+    import qibotn.cutn
+
+    # Test qibo
+    qibo.set_backend(backend=config.qibo.backend,
+                     platform=config.qibo.platform)
+
+    qibo_time, (circ_qibo, result_sv) = time(
+        lambda: qibo_qft(nqubits, swaps=True))
+
+    result_sv_cp = cp.asarray(result_sv)
+
+    # Test of MPS
+    gate_algo = {'qr_method': False,
+                 'svd_method': {
+                     'partition': 'UV',
+                     'abs_cutoff': 1e-12,
+                 }}
+
+    cutn_time, result_tn = time(
+        lambda: qibotn.cutn.eval_mps(circ_qibo, gate_algo, dtype).flatten())
+
+    print(
+        f"State vector difference: {abs(result_tn - result_sv_cp).max():0.3e}")
+
+    assert cp.allclose(result_tn, result_sv_cp)