diff --git a/.github/workflows/rules.yml b/.github/workflows/rules.yml index c0f29179..d1c06323 100644 --- a/.github/workflows/rules.yml +++ b/.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] diff --git a/setup.py b/setup.py index 2b1184ce..e7819456 100644 --- a/setup.py +++ b/setup.py @@ -42,8 +42,6 @@ def version(): "qibo>=0.1.10", "qibojit>=0.0.7", "quimb[tensor]>=1.4.0", - "cupy>=11.6.0", - "cuquantum-python-cu11", ], extras_require={ "docs": [], @@ -55,6 +53,10 @@ def version(): "analysis": [ "pylint>=2.16.0", ], + "cuda": [ + "cupy>=11.6.0", + "cuquantum-python-cu11>=23.3.0", + ], }, python_requires=">=3.8.0", long_description=(HERE / "README.md").read_text(encoding="utf-8"), diff --git a/src/qibotn/MPSUtils.py b/src/qibotn/MPSUtils.py new file mode 100644 index 00000000..fd1b4c73 --- /dev/null +++ b/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") diff --git a/src/qibotn/QiboCircuitConvertor.py b/src/qibotn/QiboCircuitConvertor.py index c30cfb64..9212f70a 100644 --- a/src/qibotn/QiboCircuitConvertor.py +++ b/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 @@ def state_vector_operands(self): for key in qubits_frontier: out_list.append(qubits_frontier[key]) - operand_exp_interleave = [x for y in zip( - operands, mode_labels) for x in y] + operand_exp_interleave = [x for y in zip(operands, mode_labels) for x in y] operand_exp_interleave.append(out_list) return operand_exp_interleave diff --git a/src/qibotn/QiboCircuitToMPS.py b/src/qibotn/QiboCircuitToMPS.py new file mode 100644 index 00000000..d51093f5 --- /dev/null +++ b/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) diff --git a/src/qibotn/cutn.py b/src/qibotn/cutn.py index e6f3e8c3..ae36497e 100644 --- a/src/qibotn/cutn.py +++ b/src/qibotn/cutn.py @@ -1,8 +1,19 @@ -# from qibotn import quimb as qiboquimb from qibotn.QiboCircuitConvertor import QiboCircuitToEinsum from cuquantum import contract +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_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} + ) diff --git a/src/qibotn/mps_contraction_helper.py b/src/qibotn/mps_contraction_helper.py new file mode 100644 index 00000000..ee8e4e42 --- /dev/null +++ b/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}) diff --git a/tests/test_cuquantum_cutensor_backend.py b/tests/test_cuquantum_cutensor_backend.py index e7f28043..1df3d0f9 100644 --- a/tests/test_cuquantum_cutensor_backend.py +++ b/tests/test_cuquantum_cutensor_backend.py @@ -2,6 +2,7 @@ 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)