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.

Sign up for GitHub

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

Feature/cheaper 2q gates #83

Closed
wants to merge 4 commits into from
Closed

Feature/cheaper 2q gates #83

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
5573dd5
Always canonicalise MPS before 2q gate
PabloAndresCQ Mar 4, 2024
aec03e4
Both MPS and TTN now apply SVD on a two-qubit gate tensor before cont…
PabloAndresCQ Mar 4, 2024
4fb73e8
Cheeky bug fix
PabloAndresCQ Mar 20, 2024
a531f04
Another cheeky bugfix. I will split these two commits in a different …
PabloAndresCQ Mar 22, 2024
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
8 changes: 5 additions & 3 deletions pytket/extensions/cutensornet/structured_state/general.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -54,7 +54,6 @@ class CuTensorNetHandle:
"""

def __init__(self, device_id: Optional[int] = None):
self.handle = cutn.create()
self._is_destroyed = False

# Make sure CuPy uses the specified device
Expand All @@ -63,6 +62,8 @@ def __init__(self, device_id: Optional[int] = None):
dev = cp.cuda.Device()
self.device_id = int(dev)

self.handle = cutn.create()

def __enter__(self) -> CuTensorNetHandle:
return self

Expand Down Expand Up @@ -128,14 +129,15 @@ def __init__(
ValueError: If the value of ``chi`` is set below 2.
ValueError: If the value of ``truncation_fidelity`` is not in [0,1].
"""
_CHI_LIMIT = 2**60
if (
chi is not None
chi is not None and chi < _CHI_LIMIT
and truncation_fidelity is not None
and truncation_fidelity != 1.0
):
raise ValueError("Cannot fix both chi and truncation_fidelity.")
if chi is None:
chi = 2**60 # In practice, this is like having it be unbounded
chi = _CHI_LIMIT # In practice, this is like having it be unbounded
if truncation_fidelity is None:
truncation_fidelity = 1

Expand Down
40 changes: 19 additions & 21 deletions pytket/extensions/cutensornet/structured_state/mps_gate.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -96,27 +96,16 @@ def _apply_2q_gate(self, positions: tuple[int, int], gate: Op) -> MPSxGate:
l_pos = min(positions)
r_pos = max(positions)

# Always canonicalise. Even in the case of exact simulation (no truncation)
# canonicalisation may reduce the bond dimension (thanks to reduced QR).
self.canonicalise(l_pos, r_pos)

# Figure out the new dimension of the shared virtual bond
new_dim = 2 * min(
self.get_virtual_dimensions(l_pos)[0],
self.get_virtual_dimensions(r_pos)[1],
)

# Canonicalisation may be required if `new_dim` is larger than `chi`
# or if set by `truncation_fidelity`
if new_dim > self._cfg.chi or self._cfg.truncation_fidelity < 1:
# If truncation required, convert to canonical form before
# contracting. Avoids the need to apply gauge transformations
# to the larger tensor resulting from the contraction.
self.canonicalise(l_pos, r_pos)

# Since canonicalisation may change the dimension of the bonds,
# we need to recalculate the value of `new_dim`
new_dim = 2 * min(
self.get_virtual_dimensions(l_pos)[0],
self.get_virtual_dimensions(r_pos)[1],
)

# Load the gate's unitary to the GPU memory
gate_unitary = gate.get_unitary().astype(dtype=self._cfg._complex_t, copy=False)
gate_tensor = cp.asarray(gate_unitary, dtype=self._cfg._complex_t)
Expand All @@ -129,26 +118,35 @@ def _apply_2q_gate(self, positions: tuple[int, int], gate: Op) -> MPSxGate:
# r -> physical bond of the right tensor in the MPS
# L -> left bond of the outcome of the gate
# R -> right bond of the outcome of the gate
# S -> shared bond of the gate tensor's SVD
# a,b,c -> the virtual bonds of the tensors

if l_pos == positions[0]:
gate_bonds = "LRlr"
else: # Implicit swap
gate_bonds = "RLrl"

left_bonds = "abl"
right_bonds = "bcr"
result_bonds = "acLR"
# Apply SVD on the gate tensor to remove any zero singular values ASAP
svd_method = tensor.SVDMethod(
abs_cutoff=self._cfg.zero,
partition="U", # Contract S directly into U
)
# Apply the SVD decomposition using the configuration defined above
U, S, V = tensor.decompose(
f"{gate_bonds}->SLl,SRr", gate_tensor, method=svd_method, options=options
)
assert S is None # Due to "partition" option in SVDMethod

# Contract
self._logger.debug("Contracting the two-qubit gate with its site tensors...")
T = cq.contract(
gate_bonds + "," + left_bonds + "," + right_bonds + "->" + result_bonds,
gate_tensor,
f"SLl,abl,SRr,bcr->acLR",
U,
self.tensors[l_pos],
V,
self.tensors[r_pos],
options=options,
optimize={"path": [(0, 1), (0, 1)]},
optimize={"path": [(0, 1), (0, 1), (0, 1)]},
)
self._logger.debug(f"Intermediate tensor of size (MiB)={T.nbytes / 2**20}")

Expand Down
72 changes: 43 additions & 29 deletions pytket/extensions/cutensornet/structured_state/ttn_gate.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -115,6 +115,7 @@ def _apply_2q_gate(self, q0: Qubit, q1: Qubit, gate: Op) -> TTNxGate:
# b -> the input bond of the gate on q1
# A -> the output bond of the gate on q0
# B -> the output bond of the gate on q1
# S -> the shared bond of the gate tensor's SVD
# l -> left child bond of the TTN node
# r -> right child bond of the TTN node
# p -> the parent bond of the TTN node
Expand Down Expand Up @@ -172,19 +173,29 @@ def _apply_2q_gate(self, q0: Qubit, q1: Qubit, gate: Op) -> TTNxGate:
# of `canonicalise()`.
self.canonicalise(center=(*common_path, DirTTN.LEFT))

# The overall strategy is to connect the `a` bond of the gate tensor to
# the corresponding bond for `q0` in the TTN (so that its bond `A`) becomes
# the new physical bond for `q0`. However, bonds `b` and `B` corresponding to
# `q1` are left open. We combine this `gate_tensor` with the leaf node of `q0`
# and QR-decompose the result; where the Q tensor will be the new
# (canonicalised) leaf node and R becomes our `msg_tensor`. The latter contains
# the open bonds `b` and `B` and our objective is to "push" this `msg_tensor`
# Apply SVD on the gate tensor to remove any zero singular values ASAP
svd_method = tensor.SVDMethod(
abs_cutoff=self._cfg.zero,
partition="U", # Contract S directly into U
)
# Apply the SVD decomposition using the configuration defined above
U, S, V = tensor.decompose(
f"{gate_bonds}->SAa,SBb", gate_tensor, method=svd_method, options=options
)
assert S is None # Due to "partition" option in SVDMethod

# The overall strategy is to connect the `U` tensor above with the physical bond
# for `q0` in the TTN, so that its bond `A` becomes the new physical bond and
# the bond `S` is left dangling (open). We combine this `gate_tensor` with the
# leaf node of `q0` and QR-decompose the result; where the Q tensor will be the
# new (canonicalised) leaf node and R becomes our `msg_tensor`. The latter
# contains the open bond `S` and our objective is to "push" this `msg_tensor`
# through the TTN towards the leaf node of `q1`. Here, "push through" means
# contract with the next tensor, and apply QR decomposition, so that the
# `msg_tensor` carrying `b` and `B` ends up one bond closer to `q1`.
# Once `msg_tensor` is directly connected to the leaf node containing `q1`, we
# just need to contract them, connecting `b` to `q1`, with `B` becoming the
# new physical bond.
# just need to apply the `V` tensor above to `q1` and connect its `S` bond with
# that of the `msg_tensor`.
bonds_to_q0 = [ # Bonds in the "arc" from the common ancestor to `q0`
path_q0[:i] for i in range(len(common_path) + 1, len(path_q0) + 1)
]
Expand All @@ -205,14 +216,14 @@ def _apply_2q_gate(self, q0: Qubit, q1: Qubit, gate: Op) -> TTNxGate:
assert len(bonds_to_q1) == 1 or len(bonds_to_q1[0]) < len(bonds_to_q1[1])
assert len(bonds_to_q1[-1]) == len(path_q1)

# The `msg_tensor` has four bonds. Our convention will be that the first bond
# always corresponds to `B`, the second bond is `b`, the third bond connects
# it to the TTN in the child direction and the fourth connects it to the TTN
# in the `DirTTN.PARENT` direction. If we label the third bond with `l`, then
# the fourth bond will be labelled `L` (and vice versa). Same for `r` and `p`.
# The `msg_tensor` has three bonds. Our convention will be that the first bond
# always corresponds to `S`, the second bond connects the `msg_tensor`
# to the TTN in the child direction and the third connects it to the TTN
# in the `DirTTN.PARENT` direction. If we label the second bond with `l`, then
# the third bond will be labelled `L` (and vice versa). Same for `r` and `p`.

# We begin applying the gate to the TTN by contracting `gate_tensor` into the
# leaf node containing `q0`, with the `b` and `B` bonds of the latter left open.
# We begin applying the gate to the TTN by contracting `U` into the
# leaf node containing `q0`, with the `S` bond of the former left open.
# We immediately QR-decompose the resulting tensor, so that Q becomes the new
# (canonicalised) leaf node and R becomes the `msg_tensor` that we will be
# "pushing" through the rest of the arc towards `q1`.
Expand All @@ -223,13 +234,14 @@ def _apply_2q_gate(self, q0: Qubit, q1: Qubit, gate: Op) -> TTNxGate:
leaf_bonds = "".join(aux_bonds) + "p"
aux_bonds[bond_q0] = "A"
Q_bonds = "".join(aux_bonds) + "s"
R_bonds = "Bbsp" # The `msg_tensor`
R_bonds = "Ssp" # The `msg_tensor`
U_bonds = "SAa"

# Apply the contraction followed by a QR decomposition
leaf_node.tensor, msg_tensor = contract_decompose(
f"{leaf_bonds},{gate_bonds}->{Q_bonds},{R_bonds}",
f"{leaf_bonds},{U_bonds}->{Q_bonds},{R_bonds}",
leaf_node.tensor,
gate_tensor,
U,
algorithm={"qr_method": tensor.QRMethod()},
options=options,
optimize={"path": [(0, 1)]},
Expand All @@ -248,9 +260,9 @@ def _apply_2q_gate(self, q0: Qubit, q1: Qubit, gate: Op) -> TTNxGate:
node = self.nodes[parent_bond]

node_bonds = "lrp"
msg_bonds = "BbLl" if child_dir == DirTTN.LEFT else "BbRr"
msg_bonds = "SLl" if child_dir == DirTTN.LEFT else "SRr"
Q_bonds = "Lrs" if child_dir == DirTTN.LEFT else "lRs"
R_bonds = "Bbsp" # The new `msg_tensor`
R_bonds = "Ssp" # The new `msg_tensor`

self._logger.debug(
f"Pushing msg_tensor ({msg_tensor.nbytes // 2**20} MiB) through node "
Expand All @@ -277,9 +289,9 @@ def _apply_2q_gate(self, q0: Qubit, q1: Qubit, gate: Op) -> TTNxGate:
common_ancestor_node = self.nodes[parent_bond]

node_bonds = "lrp"
msg_bonds = "BbLl" if child_dir == DirTTN.LEFT else "BbRr"
msg_bonds = "SLl" if child_dir == DirTTN.LEFT else "SRr"
Q_bonds = "Lsp" if child_dir == DirTTN.LEFT else "sRp"
R_bonds = "Bbrs" if child_dir == DirTTN.LEFT else "Bbls" # The new `msg_tensor`
R_bonds = "Srs" if child_dir == DirTTN.LEFT else "Sls" # The new `msg_tensor`

self._logger.debug(
f"Pushing msg_tensor ({msg_tensor.nbytes // 2**20} MiB) through node "
Expand Down Expand Up @@ -312,9 +324,9 @@ def _apply_2q_gate(self, q0: Qubit, q1: Qubit, gate: Op) -> TTNxGate:
node = self.nodes[parent_bond]

node_bonds = "lrp"
msg_bonds = "BbpP"
msg_bonds = "SpP"
Q_bonds = "srP" if child_dir == DirTTN.LEFT else "lsP"
R_bonds = "Bbls" if child_dir == DirTTN.LEFT else "Bbrs" # New `msg_tensor`
R_bonds = "Sls" if child_dir == DirTTN.LEFT else "Srs" # New `msg_tensor`

self._logger.debug(
f"Pushing msg_tensor ({msg_tensor.nbytes // 2**20} MiB) through node "
Expand All @@ -334,23 +346,25 @@ def _apply_2q_gate(self, q0: Qubit, q1: Qubit, gate: Op) -> TTNxGate:
node.canonical_form = child_dir

# Finally, the `msg_tensor` is in the parent bond of the leaf node of `q1`.
# All we need to do is contract the `msg_tensor` into the leaf.
# All we need to do is contract the `msg_tensor` and `V` into the leaf.
leaf_node = self.nodes[path_q1]
n_qbonds = len(leaf_node.tensor.shape) - 1 # Num of qubit bonds
aux_bonds = [chr(x) for x in range(n_qbonds)]
aux_bonds[bond_q1] = "b" # Connect `b` to `q1`
leaf_bonds = "".join(aux_bonds) + "p"
msg_bonds = "BbpP"
msg_bonds = "SpP"
V_bonds = "SBb"
aux_bonds[bond_q1] = "B" # `B` becomes the new physical bond `q1`
result_bonds = "".join(aux_bonds) + "P"

# Apply the contraction
leaf_node.tensor = cq.contract(
f"{leaf_bonds},{msg_bonds}->{result_bonds}",
f"{leaf_bonds},{V_bonds},{msg_bonds}->{result_bonds}",
leaf_node.tensor,
V,
msg_tensor,
options=options,
optimize={"path": [(0, 1)]},
optimize={"path": [(0, 1), (0, 1)]},
)
# The leaf node lost its canonical form
leaf_node.canonical_form = None
Expand Down
Loading