Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Cu quantum mps #17

Merged
merged 31 commits into from
Oct 17, 2023
Merged
Show file tree
Hide file tree
Changes from 27 commits
Commits
Show all changes
31 commits
Select commit Hold shift + click to select a range
9eb93cf
Initial commit
Tankya2 Jul 12, 2023
76f61bc
Updated cuquantum requirement
Tankya2 Jul 14, 2023
3cb0fec
Added MPS codes
Tankya2 Jul 14, 2023
a17d8e6
free handle
Tankya2 Jul 17, 2023
b043e6a
Add the pytest function for MPS in cuquantum.
liweintu Jul 24, 2023
4da70db
remove MPO operation
Tankya2 Jul 24, 2023
ac712d2
removed additional helper file
Tankya2 Jul 24, 2023
6a22db0
remove __main__
Tankya2 Jul 24, 2023
73d65a6
Removed unused imports
Tankya2 Aug 3, 2023
fc3e0c2
Return function output directly
Tankya2 Aug 3, 2023
ca6d579
Removed duplicated self.options
Tankya2 Aug 3, 2023
ce018ab
Removed default dtype
Tankya2 Aug 3, 2023
32611cd
Update the mps initialization function name
Tankya2 Aug 3, 2023
c15e4ca
Update function input as options attribute removed
Tankya2 Aug 3, 2023
32bee13
Removed return of tensor as it is modified in place
Tankya2 Aug 3, 2023
139748d
change to literal for dict
Tankya2 Aug 15, 2023
696f052
change list conversion method
Tankya2 Aug 15, 2023
7abb82c
Change list conversion method
Tankya2 Aug 15, 2023
a42ba15
use **kwargs
Tankya2 Aug 16, 2023
c3809d3
remove return in apply_gate()
Tankya2 Aug 16, 2023
cb21d1d
Update fucntions to use **kwargs
Tankya2 Aug 17, 2023
3fafe2b
Remove path cache
Tankya2 Aug 17, 2023
89bdbfb
Format with Black
Tankya2 Aug 17, 2023
78246cf
Change filename to that of naming convention
Tankya2 Aug 17, 2023
1dba4c3
Update lib name
Tankya2 Aug 17, 2023
d81cce5
Change list conversion using constructor
Tankya2 Aug 17, 2023
d558609
Minor fix of typo and format.
liweintu Aug 29, 2023
ba442c9
Format with black
Tankya2 Aug 30, 2023
4902195
Move cuquantum-python-cu11 package to extras to avoid the issue of mi…
liweintu Aug 30, 2023
b84ccad
Move cupy package to extras to avoid the issue of missing CUDA
liweintu Aug 30, 2023
e91045f
Check CUDA_PATH to determine if CUDA is installed and if the build jo…
liweintu Sep 12, 2023
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -43,7 +43,7 @@ def version():
"qibojit>=0.0.7",
"quimb[tensor]>=1.4.0",
"cupy>=11.6.0",
liweintu marked this conversation as resolved.
Show resolved Hide resolved
"cuquantum-python-cu11",
"cuquantum-python-cu11>=23.3.0",
liweintu marked this conversation as resolved.
Show resolved Hide resolved
],
extras_require={
"docs": [],
Expand Down
82 changes: 82 additions & 0 deletions src/qibotn/MPSUtils.py
Original file line number Diff line number Diff line change
@@ -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>
"""
Comment on lines +7 to +9
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sooner or later we'll start using Pydocstring, and it will check that every first line will end with a full stop.

Suggested change
"""
Generate the MPS with an initial state of |00...00>
"""
"""Generate the MPS with an initial state of |00...00>."""

(this is just an irrelevant remark, because this PR is already pretty good, so we can improve even some details).

state_tensor = cp.asarray([1, 0], dtype=dtype).reshape(1, 2, 1)
mps_tensors = [state_tensor] * num_qubits
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I believe CuPy arrays to be immutable, so it should not be a problem.
But try to confirm it, because, otherwise, you'd be generating many references to the same object, that might have counter-intuitive effects (that is even if you modify just one, all the others will change as well.

return mps_tensors
Comment on lines +11 to +12
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It's not necessarily bad, but it is just simpler to avoid the assignment, since you're immediately returning (unless you're very motivated to give it a name, but the name will usually be very close to the function one)

Suggested change
mps_tensors = [state_tensor] * num_qubits
return mps_tensors
return [state_tensor] * num_qubits



def mps_site_right_swap(mps_tensors, i, **kwargs):
"""
Perform the swap operation between the ith and i+1th MPS tensors.
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Tell the user that the operation will be performed in-place.

"""
# 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
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Since the operation is in-place, you do not need to return the modified object, since it's the one passed by the user.
(and there should be only one way of doing something, so it would be confusing to have the two references around)



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")
2 changes: 1 addition & 1 deletion src/qibotn/QiboCircuitConvertor.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.

Expand Down
38 changes: 38 additions & 0 deletions src/qibotn/QiboCircuitToMPS.py
Original file line number Diff line number Diff line change
@@ -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)
Comment on lines +20 to +23
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why do you need all these assignments to self?

If possible, keep everything local (in this case local to the __init__ method).

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 changes: 12 additions & 1 deletion src/qibotn/cutn.py
Original file line number Diff line number Diff line change
@@ -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 changes: 129 additions & 0 deletions src/qibotn/mps_contraction_helper.py
Original file line number Diff line number Diff line change
@@ -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):
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Just a proposal to simplify the method names: since the class is already named MPSContractionHelper, you could avoid prefixing methods with contract_ (since it is reasonable that an MPSContractionHelper.norm method will make a contraction to give you the norm).

However, it is just a proposal and definitely subject to opinions. So, feel free to keep as it is.

"""
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):
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
for i, o in enumerate(mps_tensors):
for i, op in enumerate(mps_tensors):

To slightly improve the name's expresiveness, at very little expense ^^

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):
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Same of above

Suggested change
for i, o in enumerate(mps_tensors):
for i, op 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 changes: 38 additions & 0 deletions tests/test_cuquantum_cutensor_backend.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@

import config
import numpy as np
import cupy as cp
import pytest
import qibo
from qibo.models import QFT
Expand Down Expand Up @@ -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)