jaunatisblue/final-qibotn
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge pull request #17 from qiboteam/cuQuantum_mps

Browse Source 
Cu quantum mps
This commit is contained in:
liwei
2023-10-17 10:29:08 +08:00
committed by GitHub
parent 6455825050 e91045f5a3
commit cb0a6f0d19
8 changed files with 306 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
.github/workflows/rules.yml vendored
View File

@@ -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]

6
setup.py
View File

@@ -42,8 +42,6 @@ setup(
"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 @@ setup(
"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"),

82
src/qibotn/MPSUtils.py Normal file
View File

@@ -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")

5
src/qibotn/QiboCircuitConvertor.py
View File

@@ -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(
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

38
src/qibotn/QiboCircuitToMPS.py Normal file
View File

@@ -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)

13
src/qibotn/cutn.py
View File

@@ -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}
)

129
src/qibotn/mps_contraction_helper.py Normal file
View File

@@ -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})

38
tests/test_cuquantum_cutensor_backend.py
View File

@@ -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)