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

Merge branch 'main' into mps-for-quimb

Browse Source
This commit is contained in:
nitinshivaraman
2023-11-02 17:29:46 +08:00
committed by GitHub
parent 2c15613e5e c617cf5659
commit b26360539c
7 changed files with 344 additions and 6 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]

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

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

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)

53
src/qibotn/cutn.py
View File

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

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)