From 3cb0fec99c2b54a2d48c7401de224736d04091f9 Mon Sep 17 00:00:00 2001 From: tankya2 Date: Fri, 14 Jul 2023 09:51:06 +0800 Subject: [PATCH] Added MPS codes --- src/qibotn/MPSContractionHelper.py | 155 +++++++++++++++++++++++++++++ src/qibotn/MPSUtils.py | 78 +++++++++++++++ src/qibotn/QiboCircuitToMPS.py | 24 +++++ src/qibotn/cutn.py | 78 ++++----------- 4 files changed, 278 insertions(+), 57 deletions(-) create mode 100644 src/qibotn/MPSContractionHelper.py create mode 100644 src/qibotn/MPSUtils.py create mode 100644 src/qibotn/QiboCircuitToMPS.py diff --git a/src/qibotn/MPSContractionHelper.py b/src/qibotn/MPSContractionHelper.py new file mode 100644 index 0000000..cd0905e --- /dev/null +++ b/src/qibotn/MPSContractionHelper.py @@ -0,0 +1,155 @@ +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.path_cache = dict() + 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('norm', 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('sv', 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(f'exp{qubits}', interleaved_inputs, options=options) / norm + + def contract_mps_mpo_to_state_vector(self, mps_tensors, mpo_tensors, options=None): + """ + Contract the corresponding tensor network to form the output state vector from applying the MPO to the MPS. + + 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. + mpo_tensors: A list of rank-4 ndarray-like tensor objects. + The indics of the ith tensor are expected to be the bonding index to the i-1 tensor, + the output physical mode, the bonding index to the i+1th tensor and then the inputput physical mode. + options: Specify the contract and decompose options. + + Returns: + An ndarray-like object as the output state vector. + """ + interleaved_inputs = [] + for i, o in enumerate(mps_tensors): + interleaved_inputs.extend([o, self.bra_modes[i]]) + output_modes = [] + offset = 2 * self.num_qubits + 1 + for i, o in enumerate(mpo_tensors): + mpo_modes = (2*i+offset, 2*i+offset+1, 2*i+offset+2, 2*i+1) + output_modes.append(2*i+offset+1) + interleaved_inputs.extend([o, mpo_modes]) + interleaved_inputs.append(output_modes) + return self._contract('mps_mpo', interleaved_inputs, options=options) + + def _contract(self, key, interleaved_inputs, options=None): + """ + Perform the contraction task given interleaved inputs. Path will be cached. + """ + if key not in self.path_cache: + self.path_cache[key] = contract_path(*interleaved_inputs, options=options)[0] + path = self.path_cache[key] + return contract(*interleaved_inputs, options=options, optimize={'path':path}) diff --git a/src/qibotn/MPSUtils.py b/src/qibotn/MPSUtils.py new file mode 100644 index 0000000..cb77e4b --- /dev/null +++ b/src/qibotn/MPSUtils.py @@ -0,0 +1,78 @@ +import cupy as cp +from cuquantum.cutensornet.experimental import contract_decompose +from cuquantum import contract + +def get_initial_mps(num_qubits, dtype='complex128'): + """ + 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, + algorithm=None, + options=None +): + """ + 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=algorithm, options=options) + mps_tensors[i:i+2] = (a, b) + return mps_tensors + +def apply_gate( + mps_tensors, + gate, + qubits, + algorithm=None, + options=None +): + """ + 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=options) # 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), algorithm=algorithm, options=options) + elif i+1 == j: + # two adjacent qubits + a, _, b = contract_decompose('ipj,jqk,rspq->irj,jsk', *mps_tensors[i:i+2], gate, algorithm=algorithm, options=options) + 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, algorithm=algorithm, options=options) + # 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), algorithm=algorithm, options=options) + # step 3: swap back i and i+1 + mps_site_right_swap(mps_tensors, i, algorithm=algorithm, options=options) + else: + raise NotImplementedError("Only one- and two-qubit gates supported") + return mps_tensors \ No newline at end of file diff --git a/src/qibotn/QiboCircuitToMPS.py b/src/qibotn/QiboCircuitToMPS.py new file mode 100644 index 0000000..f66804d --- /dev/null +++ b/src/qibotn/QiboCircuitToMPS.py @@ -0,0 +1,24 @@ +import cupy as cp +import numpy as np + +from cuquantum import cutensornet as cutn +from QiboCircuitConvertor import QiboCircuitToEinsum +from MPSUtils import get_initial_mps, 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.options = {'handle': self.handle} + self.dtype = dtype + self.mps_tensors = get_initial_mps(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=self.options) + diff --git a/src/qibotn/cutn.py b/src/qibotn/cutn.py index 6e4a98b..f29a62e 100644 --- a/src/qibotn/cutn.py +++ b/src/qibotn/cutn.py @@ -5,68 +5,32 @@ import cuquantum from cuquantum import cutensornet as cutn import cupy as cp import numpy as np +from qibo.models import QFT +from QiboCircuitToMPS import QiboCircuitToMPS +from MPSContractionHelper import MPSContractionHelper + def eval(qibo_circ, datatype): myconvertor = QiboCircuitToEinsum(qibo_circ, dtype=datatype) return contract(*myconvertor.state_vector_operands()) -def eval_mps(qibo_circ, datatype): - #Create MPS - cutensornet.create() - return contract() - +def eval_mps(qibo_circ, gate_algo, datatype): + myconvertor = QiboCircuitToMPS(qibo_circ, gate_algo, datatype) + mps_helper = MPSContractionHelper(myconvertor.num_qubits) + sv_mps = mps_helper.contract_state_vector(myconvertor.mps_tensors,myconvertor.options) + return sv_mps if __name__ == "__main__": - print("cuTensorNet-vers:", cutn.get_version()) - dev = cp.cuda.Device() # get current device - props = cp.cuda.runtime.getDeviceProperties(dev.id) - print("===== device info ======") - print("GPU-name:", props["name"].decode()) - print("GPU-clock:", props["clockRate"]) - print("GPU-memoryClock:", props["memoryClockRate"]) - print("GPU-nSM:", props["multiProcessorCount"]) - print("GPU-major:", props["major"]) - print("GPU-minor:", props["minor"]) - print("========================") - - data_type = cuquantum.cudaDataType.CUDA_C_64F - compute_type = cuquantum.ComputeType.COMPUTE_64F - num_sites = 16 - phys_extent = 2 - max_virtual_extent = 12 - - ## we initialize the MPS state as a product state |000...000> - initial_state = [] - for i in range(num_sites): - # we create dummpy indices for MPS tensors on the boundary for easier bookkeeping - # we'll use Fortran layout throughout this example - tensor = cp.zeros((1,2,1), dtype=np.complex128, order="F") - tensor[0,0,0] = 1.0 - initial_state.append(tensor) - - mps_helper = MPSHelper(num_sites, phys_extent, max_virtual_extent, initial_state, data_type, compute_type) - - ################################## - # Setup options for gate operation - ################################## - - abs_cutoff = 1e-2 - rel_cutoff = 1e-2 - renorm = cutn.TensorSVDNormalization.L2 - partition = cutn.TensorSVDPartition.UV_EQUAL - mps_helper.set_svd_config(abs_cutoff, rel_cutoff, renorm, partition) - - gate_algo = cutn.GateSplitAlgo.REDUCED - mps_helper.set_gate_algorithm(gate_algo) - - ##################################### - # Workspace estimation and allocation - ##################################### - - free_mem, total_mem = dev.mem_info - worksize = free_mem *.7 - required_workspace_size = mps_helper.compute_max_workspace_sizes() - work = cp.cuda.alloc(worksize) - print(f"Maximal workspace size requried: {required_workspace_size / 1024 ** 3:.3f} GB") - mps_helper.set_workspace(work, required_workspace_size) + num_qubits = 25 + swaps = True + circ_qibo = QFT(num_qubits, swaps) + exact_gate_algorithm = {'qr_method': False, + 'svd_method':{'partition': 'UV', 'abs_cutoff':1e-12}} + dtype="complex128" + sv_mps = eval_mps(circ_qibo, exact_gate_algorithm, dtype) + sv_reference = eval(circ_qibo, dtype) + state_vec = np.array(circ_qibo()) + print(f"State vector difference: {abs(sv_mps-sv_reference).max():0.3e}") + assert cp.allclose(sv_mps, sv_reference) + assert cp.allclose(sv_mps.flatten(), state_vec)