diff --git a/docs/api.rst b/docs/api.rst index f669c940..1e91e2c8 100644 --- a/docs/api.rst +++ b/docs/api.rst @@ -1,6 +1,11 @@ API documentation ----------------- +.. autoclass:: pytket.extensions.cutensornet.CuTensorNetHandle + + .. automethod:: destroy + + .. toctree:: - modules/fullTN.rst + modules/general_state.rst modules/structured_state.rst diff --git a/docs/changelog.rst b/docs/changelog.rst index 11cc6d38..a53343c6 100644 --- a/docs/changelog.rst +++ b/docs/changelog.rst @@ -4,6 +4,7 @@ Changelog Unreleased ---------- +* New API: ``GeneralState`` for exact simulation of circuits via contraction-path optimisation. Currently supports ``get_statevector()`` and ``expectation_value()``. * New feature: ``add_qubit`` to add fresh qubits at specified positions in an ``MPS``. * New feature: added an option to ``measure`` to toggle destructive measurement on/off. Currently only supported for ``MPS``. * New feature: a seed can now be provided to ``Config`` objects, providing reproducibility across ``StructuredState`` simulations. diff --git a/docs/modules/fullTN.rst b/docs/modules/fullTN.rst deleted file mode 100644 index 48f4ee13..00000000 --- a/docs/modules/fullTN.rst +++ /dev/null @@ -1,5 +0,0 @@ -Full tensor network contraction -=============================== - -.. automodule:: pytket.extensions.cutensornet - :members: TensorNetwork, PauliOperatorTensorNetwork, ExpectationValueTensorNetwork, measure_qubits_state, tk_to_tensor_network, CuTensorNetBackend \ No newline at end of file diff --git a/docs/modules/general_state.rst b/docs/modules/general_state.rst new file mode 100644 index 00000000..4ee45838 --- /dev/null +++ b/docs/modules/general_state.rst @@ -0,0 +1,31 @@ +General state (exact) simulation +================================ + +.. automodule:: pytket.extensions.cutensornet.general_state + +.. autoclass:: pytket.extensions.cutensornet.general_state.GeneralState() + + .. automethod:: __init__ + .. automethod:: get_statevector + .. automethod:: expectation_value + .. automethod:: destroy + +cuQuantum `contract` API interface +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +.. autoclass:: pytket.extensions.cutensornet.general_state.TensorNetwork + +.. autoclass:: pytket.extensions.cutensornet.general_state.PauliOperatorTensorNetwork + +.. autoclass:: pytket.extensions.cutensornet.general_state.ExpectationValueTensorNetwork + +.. autofunction:: pytket.extensions.cutensornet.general_state.measure_qubits_state + +.. autofunction:: pytket.extensions.cutensornet.general_state.tk_to_tensor_network + + +Pytket backend +~~~~~~~~~~~~~~ + +.. automodule:: pytket.extensions.cutensornet + :members: CuTensorNetBackend diff --git a/docs/modules/structured_state.rst b/docs/modules/structured_state.rst index f360d196..e6b658d7 100644 --- a/docs/modules/structured_state.rst +++ b/docs/modules/structured_state.rst @@ -16,10 +16,6 @@ Simulation .. automethod:: __init__ -.. autoclass:: pytket.extensions.cutensornet.structured_state.CuTensorNetHandle - - .. automethod:: destroy - Classes ~~~~~~~ diff --git a/pytket/extensions/cutensornet/__init__.py b/pytket/extensions/cutensornet/__init__.py index 5d12f155..ba720477 100644 --- a/pytket/extensions/cutensornet/__init__.py +++ b/pytket/extensions/cutensornet/__init__.py @@ -13,17 +13,8 @@ # limitations under the License. """Module for conversion from tket primitives to cuQuantum primitives.""" -# _metadata.py is copied to the folder after installation. -from ._metadata import __extension_version__, __extension_name__ # type: ignore - from .backends import CuTensorNetBackend +from .general import CuTensorNetHandle -from .tensor_network_convert import ( - TensorNetwork, - PauliOperatorTensorNetwork, - ExpectationValueTensorNetwork, - tk_to_tensor_network, - measure_qubits_state, -) - -from .utils import circuit_statevector_postselect +# _metadata.py is copied to the folder after installation. +from ._metadata import __extension_version__, __extension_name__ # type: ignore diff --git a/pytket/extensions/cutensornet/backends/cutensornet_backend.py b/pytket/extensions/cutensornet/backends/cutensornet_backend.py index b6e13757..5c1af286 100644 --- a/pytket/extensions/cutensornet/backends/cutensornet_backend.py +++ b/pytket/extensions/cutensornet/backends/cutensornet_backend.py @@ -30,7 +30,7 @@ from pytket.backends.backend import KwargTypes, Backend, BackendResult from pytket.backends.backendinfo import BackendInfo from pytket.backends.resulthandle import _ResultIdTuple -from pytket.extensions.cutensornet.tensor_network_convert import ( +from pytket.extensions.cutensornet.general_state import ( TensorNetwork, ExpectationValueTensorNetwork, tk_to_tensor_network, diff --git a/pytket/extensions/cutensornet/general.py b/pytket/extensions/cutensornet/general.py index f39ed406..3f33882e 100644 --- a/pytket/extensions/cutensornet/general.py +++ b/pytket/extensions/cutensornet/general.py @@ -11,9 +11,85 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +from __future__ import annotations # type: ignore +import warnings import logging from logging import Logger +from typing import Any, Optional + +try: + import cupy as cp # type: ignore +except ImportError: + warnings.warn("local settings failed to import cupy", ImportWarning) +try: + import cuquantum.cutensornet as cutn # type: ignore +except ImportError: + warnings.warn("local settings failed to import cutensornet", ImportWarning) + + +class CuTensorNetHandle: + """Initialise the cuTensorNet library with automatic workspace memory + management. + + Note: + Always use as ``with CuTensorNetHandle() as libhandle:`` so that cuTensorNet + handles are automatically destroyed at the end of execution. + + Attributes: + handle (int): The cuTensorNet library handle created by this initialisation. + device_id (int): The ID of the device (GPU) where cuTensorNet is initialised. + If not provided, defaults to ``cp.cuda.Device()``. + """ + + def __init__(self, device_id: Optional[int] = None): + self._is_destroyed = False + + # Make sure CuPy uses the specified device + dev = cp.cuda.Device(device_id) + dev.use() + + self.dev = dev + self.device_id = dev.id + + self._handle = cutn.create() + + @property + def handle(self) -> Any: + if self._is_destroyed: + raise RuntimeError( + "The cuTensorNet library handle is out of scope.", + "See the documentation of CuTensorNetHandle.", + ) + return self._handle + + def destroy(self) -> None: + """Destroys the memory handle, releasing memory. + + Only call this method if you are initialising a ``CuTensorNetHandle`` outside + a ``with CuTensorNetHandle() as libhandle`` statement. + """ + cutn.destroy(self._handle) + self._is_destroyed = True + + def __enter__(self) -> CuTensorNetHandle: + return self + + def __exit__(self, exc_type: Any, exc_value: Any, exc_tb: Any) -> None: + self.destroy() + + def print_device_properties(self, logger: Logger) -> None: + """Prints local GPU properties.""" + device_props = cp.cuda.runtime.getDeviceProperties(self.dev.id) + logger.debug("===== device info ======") + logger.debug("GPU-name:", device_props["name"].decode()) + logger.debug("GPU-clock:", device_props["clockRate"]) + logger.debug("GPU-memoryClock:", device_props["memoryClockRate"]) + logger.debug("GPU-nSM:", device_props["multiProcessorCount"]) + logger.debug("GPU-major:", device_props["major"]) + logger.debug("GPU-minor:", device_props["minor"]) + logger.debug("========================") + def set_logger( logger_name: str, diff --git a/pytket/extensions/cutensornet/general_state/__init__.py b/pytket/extensions/cutensornet/general_state/__init__.py new file mode 100644 index 00000000..ea7a0775 --- /dev/null +++ b/pytket/extensions/cutensornet/general_state/__init__.py @@ -0,0 +1,26 @@ +# Copyright 2019-2024 Quantinuum +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Module for simulating circuits with no predetermined tensor network structure.""" + +from .utils import circuit_statevector_postselect + +from .tensor_network_convert import ( + TensorNetwork, + PauliOperatorTensorNetwork, + ExpectationValueTensorNetwork, + tk_to_tensor_network, + measure_qubits_state, +) + +from .tensor_network_state import GeneralState diff --git a/pytket/extensions/cutensornet/tensor_network_convert.py b/pytket/extensions/cutensornet/general_state/tensor_network_convert.py similarity index 99% rename from pytket/extensions/cutensornet/tensor_network_convert.py rename to pytket/extensions/cutensornet/general_state/tensor_network_convert.py index 23a715a1..c4ab04a4 100644 --- a/pytket/extensions/cutensornet/tensor_network_convert.py +++ b/pytket/extensions/cutensornet/general_state/tensor_network_convert.py @@ -22,10 +22,9 @@ from networkx.classes.reportviews import OutMultiEdgeView, OutMultiEdgeDataView # type: ignore import numpy as np from numpy.typing import NDArray -from pytket import Qubit # type: ignore from pytket.utils import Graph -from pytket.pauli import QubitPauliString # type: ignore -from pytket.circuit import Circuit, Qubit # type: ignore +from pytket.pauli import QubitPauliString +from pytket.circuit import Circuit, Qubit from pytket.utils import permute_rows_cols_in_unitary from pytket.extensions.cutensornet.general import set_logger diff --git a/pytket/extensions/cutensornet/general_state/tensor_network_state.py b/pytket/extensions/cutensornet/general_state/tensor_network_state.py new file mode 100644 index 00000000..d4dc5570 --- /dev/null +++ b/pytket/extensions/cutensornet/general_state/tensor_network_state.py @@ -0,0 +1,424 @@ +# Copyright 2019-2024 Quantinuum +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations +import logging +from typing import Union, Optional +import warnings + +try: + import cupy as cp # type: ignore +except ImportError: + warnings.warn("local settings failed to import cupy", ImportWarning) +import numpy as np +from sympy import Expr # type: ignore +from numpy.typing import NDArray +from pytket.circuit import Circuit +from pytket.extensions.cutensornet.general import CuTensorNetHandle, set_logger +from pytket.utils.operators import QubitPauliOperator + +try: + import cuquantum as cq # type: ignore + from cuquantum import cutensornet as cutn # type: ignore +except ImportError: + warnings.warn("local settings failed to import cuquantum", ImportWarning) + + +class GeneralState: + """Wrapper of cuTensorNet object for exact simulations via path optimisation.""" + + def __init__( + self, + circuit: Circuit, + libhandle: CuTensorNetHandle, + loglevel: int = logging.INFO, + ) -> None: + """Constructs a tensor network representating a pytket circuit. + + The resulting object stores the *uncontracted* tensor network. + + Note: + A ``libhandle`` is created via a ``with CuTensorNetHandle() as libhandle:`` + statement. The device where the ``GeneralState`` is stored will match the + one specified by the library handle. + + Note: + The ``circuit`` must not contain any ``CircBox`` or non-unitary command. + Internally, implicit wire swaps are replaced with explicit SWAP gates. + + Args: + circuit: A pytket circuit to be converted to a tensor network. + libhandle: An instance of a ``CuTensorNetHandle``. + loglevel: Internal logger output level. + """ + self._logger = set_logger("GeneralState", loglevel) + self._circuit = circuit.copy() + # TODO: This is not strictly necessary; implicit SWAPs could be resolved by + # qubit relabelling, but it's only worth doing so if there are clear signs + # of inefficiency due to this. + self._circuit.replace_implicit_wire_swaps() + self._lib = libhandle + + libhandle.print_device_properties(self._logger) + + num_qubits = self._circuit.n_qubits + dim = 2 # We are always dealing with qubits, not qudits + qubits_dims = (dim,) * num_qubits # qubit size + self._logger.debug(f"Converting a quantum circuit with {num_qubits} qubits.") + data_type = cq.cudaDataType.CUDA_C_64F # for now let that be hard-coded + + self._state = cutn.create_state( + self._lib.handle, cutn.StatePurity.PURE, num_qubits, qubits_dims, data_type + ) + self._gate_tensors = [] + + # Append all gates to the TN + for com in self._circuit.get_commands(): + try: + gate_unitary = com.op.get_unitary() + except: + raise ValueError( + "All commands in the circuit must be unitary gates. The circuit " + f"contains {com}; no unitary matrix could be retrived for it." + ) + self._gate_tensors.append(_formatted_tensor(gate_unitary, com.op.n_qubits)) + gate_qubit_indices = tuple( + self._circuit.qubits.index(qb) for qb in com.qubits + ) + + cutn.state_apply_tensor_operator( + handle=self._lib.handle, + tensor_network_state=self._state, + num_state_modes=com.op.n_qubits, + state_modes=gate_qubit_indices, + tensor_data=self._gate_tensors[-1].data.ptr, + tensor_mode_strides=0, + immutable=1, + adjoint=0, + unitary=1, + ) + + def get_statevector( + self, + attributes: Optional[dict] = None, + scratch_fraction: float = 0.5, + on_host: bool = True, + ) -> Union[cp.ndarray, np.ndarray]: + """Contracts the circuit and returns the final statevector. + + Args: + attributes: Optional. A dict of cuTensorNet `StateAttribute` keys and + their values. + scratch_fraction: Optional. Fraction of free memory on GPU to allocate as + scratch space. + on_host: Optional. If ``True``, converts cupy ``ndarray`` to numpy + ``ndarray``, copying it to host device (CPU). + Raises: + MemoryError: If there is insufficient workspace on GPU. + Returns: + Either a ``cupy.ndarray`` on a GPU, or a ``numpy.ndarray`` on a + host device (CPU). Arrays are returned in a 1D shape. + """ + + #################################### + # Configure the TN for contraction # + #################################### + if attributes is None: + attributes = dict() + attribute_pairs = [ + (getattr(cutn.StateAttribute, k), v) for k, v in attributes.items() + ] + + for attr, val in attribute_pairs: + attr_dtype = cutn.state_get_attribute_dtype(attr) + attr_arr = np.asarray(val, dtype=attr_dtype) + cutn.state_configure( + self._lib.handle, + self._state, + attr, + attr_arr.ctypes.data, + attr_arr.dtype.itemsize, + ) + + try: + ###################################### + # Allocate workspace for contraction # + ###################################### + stream = cp.cuda.Stream() + free_mem = self._lib.dev.mem_info[0] + scratch_size = int(scratch_fraction * free_mem) + scratch_space = cp.cuda.alloc(scratch_size) + self._logger.debug( + f"Allocated {scratch_size} bytes of scratch memory on GPU" + ) + work_desc = cutn.create_workspace_descriptor(self._lib.handle) + + cutn.state_prepare( + self._lib.handle, + self._state, + scratch_size, + work_desc, + stream.ptr, + ) + workspace_size_d = cutn.workspace_get_memory_size( + self._lib.handle, + work_desc, + cutn.WorksizePref.RECOMMENDED, + cutn.Memspace.DEVICE, + cutn.WorkspaceKind.SCRATCH, + ) + + if workspace_size_d <= scratch_size: + cutn.workspace_set_memory( + self._lib.handle, + work_desc, + cutn.Memspace.DEVICE, + cutn.WorkspaceKind.SCRATCH, + scratch_space.ptr, + workspace_size_d, + ) + self._logger.debug( + f"Set {workspace_size_d} bytes of workspace memory out of the" + f" allocated scratch space." + ) + + else: + raise MemoryError( + f"Insufficient workspace size on the GPU device {self._lib.dev.id}" + ) + + ################### + # Contract the TN # + ################### + state_vector = cp.empty( + (2,) * self._circuit.n_qubits, dtype="complex128", order="F" + ) + cutn.state_compute( + self._lib.handle, + self._state, + work_desc, + (state_vector.data.ptr,), + stream.ptr, + ) + stream.synchronize() + sv = state_vector.flatten() + if on_host: + sv = cp.asnumpy(sv) + # Apply the phase from the circuit + sv *= np.exp(1j * np.pi * self._circuit.phase) + return sv + + finally: + cutn.destroy_workspace_descriptor(work_desc) # type: ignore + del scratch_space + + def expectation_value( + self, + operator: QubitPauliOperator, + attributes: Optional[dict] = None, + scratch_fraction: float = 0.5, + ) -> complex: + """Calculates the expectation value of the given operator. + + Args: + operator: The operator whose expectation value is to be measured. + attributes: Optional. A dict of cuTensorNet `ExpectationAttribute` keys + and their values. + scratch_fraction: Optional. Fraction of free memory on GPU to allocate as + scratch space. + + Raises: + ValueError: If the operator acts on qubits not present in the circuit. + + Returns: + The expectation value. + """ + + ############################################ + # Generate the cuTensorNet operator object # + ############################################ + pauli_tensors = { + "X": _formatted_tensor(np.asarray([[0, 1], [1, 0]]), 1), + "Y": _formatted_tensor(np.asarray([[0, -1j], [1j, 0]]), 1), + "Z": _formatted_tensor(np.asarray([[1, 0], [0, -1]]), 1), + "I": _formatted_tensor(np.asarray([[1, 0], [0, 1]]), 1), + } + num_qubits = self._circuit.n_qubits + qubits_dims = (2,) * num_qubits + data_type = cq.cudaDataType.CUDA_C_64F + + tn_operator = cutn.create_network_operator( + self._lib.handle, num_qubits, qubits_dims, data_type + ) + + self._logger.debug("Adding operator terms:") + for pauli_string, coeff in operator._dict.items(): + if isinstance(coeff, Expr): + numeric_coeff = complex(coeff.evalf()) # type: ignore + else: + numeric_coeff = complex(coeff) # type: ignore + self._logger.debug(f" {numeric_coeff}, {pauli_string}") + + # Raise an error if the operator acts on qubits that are not in the circuit + if any(q not in self._circuit.qubits for q in pauli_string.map.keys()): + raise ValueError( + f"The operator is acting on qubits {pauli_string.map.keys()}, " + "but some of these are not present in the circuit, whose set of " + f"qubits is: {self._circuit.qubits}." + ) + + # Obtain the tensors corresponding to this operator + qubit_pauli_map = { + q: pauli_tensors[pauli.name] for q, pauli in pauli_string.map.items() + } + + num_pauli = len(qubit_pauli_map) + num_modes = (1,) * num_pauli + state_modes = tuple( + (self._circuit.qubits.index(qb),) for qb in qubit_pauli_map.keys() + ) + gate_data = tuple(tensor.data.ptr for tensor in qubit_pauli_map.values()) + + cutn.network_operator_append_product( + handle=self._lib.handle, + tensor_network_operator=tn_operator, + coefficient=numeric_coeff, + num_tensors=num_pauli, + num_state_modes=num_modes, + state_modes=state_modes, + tensor_mode_strides=0, + tensor_data=gate_data, + ) + + ###################################################### + # Configure the cuTensorNet expectation value object # + ###################################################### + expectation = cutn.create_expectation( + self._lib.handle, self._state, tn_operator + ) + + if attributes is None: + attributes = dict() + attribute_pairs = [ + (getattr(cutn.ExpectationAttribute, k), v) for k, v in attributes.items() + ] + + for attr, val in attribute_pairs: + attr_dtype = cutn.expectation_get_attribute_dtype(attr) + attr_arr = np.asarray(val, dtype=attr_dtype) + cutn.expectation_configure( + self._lib.handle, + expectation, + attr, + attr_arr.ctypes.data, + attr_arr.dtype.itemsize, + ) + + try: + ###################################### + # Allocate workspace for contraction # + ###################################### + stream = cp.cuda.Stream() + free_mem = self._lib.dev.mem_info[0] + scratch_size = int(scratch_fraction * free_mem) + scratch_space = cp.cuda.alloc(scratch_size) + + self._logger.debug( + f"Allocated {scratch_size} bytes of scratch memory on GPU" + ) + work_desc = cutn.create_workspace_descriptor(self._lib.handle) + cutn.expectation_prepare( + self._lib.handle, + expectation, + scratch_size, + work_desc, + stream.ptr, + ) + workspace_size_d = cutn.workspace_get_memory_size( + self._lib.handle, + work_desc, + cutn.WorksizePref.RECOMMENDED, + cutn.Memspace.DEVICE, + cutn.WorkspaceKind.SCRATCH, + ) + + if workspace_size_d <= scratch_size: + cutn.workspace_set_memory( + self._lib.handle, + work_desc, + cutn.Memspace.DEVICE, + cutn.WorkspaceKind.SCRATCH, + scratch_space.ptr, + workspace_size_d, + ) + self._logger.debug( + f"Set {workspace_size_d} bytes of workspace memory out of the" + f" allocated scratch space." + ) + else: + raise MemoryError( + f"Insufficient workspace size on the GPU device {self._lib.dev.id}" + ) + + ################################# + # Compute the expectation value # + ################################# + expectation_value = np.empty(1, dtype="complex128") + state_norm = np.empty(1, dtype="complex128") + cutn.expectation_compute( + self._lib.handle, + expectation, + work_desc, + expectation_value.ctypes.data, + state_norm.ctypes.data, + stream.ptr, + ) + stream.synchronize() + + # Note: we can also return `state_norm.item()`, but this should be 1 since + # we are always running unitary circuits + assert np.isclose(state_norm.item(), 1.0) + + return expectation_value.item() # type: ignore + + finally: + ##################################################### + # Destroy the Operator and ExpectationValue objects # + ##################################################### + cutn.destroy_workspace_descriptor(work_desc) # type: ignore + cutn.destroy_expectation(expectation) + cutn.destroy_network_operator(tn_operator) + del scratch_space + + def destroy(self) -> None: + """Destroy the tensor network and free up GPU memory. + + Note: + Users are required to call `destroy()` when done using a + `GeneralState` object. GPU memory deallocation is not + guaranteed otherwise. + """ + cutn.destroy_state(self._state) + + +def _formatted_tensor(matrix: NDArray, n_qubits: int) -> cp.ndarray: + """Convert a matrix to the tensor format accepted by NVIDIA's API.""" + + # Transpose is needed because of the way cuTN stores tensors. + # See https://github.com/NVIDIA/cuQuantum/discussions/124 + # #discussioncomment-8683146 for details. + cupy_matrix = cp.asarray(matrix).T.astype(dtype="complex128", order="F") + # We also need to reshape since a matrix only has 2 bonds, but for an + # n-qubit gate we want 2^n bonds for input and another 2^n for output + return cupy_matrix.reshape([2] * (2 * n_qubits), order="F") diff --git a/pytket/extensions/cutensornet/utils.py b/pytket/extensions/cutensornet/general_state/utils.py similarity index 80% rename from pytket/extensions/cutensornet/utils.py rename to pytket/extensions/cutensornet/general_state/utils.py index 3069191f..d0b77927 100644 --- a/pytket/extensions/cutensornet/utils.py +++ b/pytket/extensions/cutensornet/general_state/utils.py @@ -1,6 +1,20 @@ +# Copyright 2019-2024 Quantinuum +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + from numpy.typing import NDArray from pytket.backends.backendresult import BackendResult -from pytket.circuit import Qubit, Circuit # type: ignore +from pytket.circuit import Qubit, Circuit def _reorder_qlist( diff --git a/pytket/extensions/cutensornet/structured_state/__init__.py b/pytket/extensions/cutensornet/structured_state/__init__.py index be9dd9f5..3daffd99 100644 --- a/pytket/extensions/cutensornet/structured_state/__init__.py +++ b/pytket/extensions/cutensornet/structured_state/__init__.py @@ -19,7 +19,9 @@ https://github.com/CQCL/pytket-cutensornet. """ -from .general import CuTensorNetHandle, Config, StructuredState +from pytket.extensions.cutensornet import CuTensorNetHandle + +from .general import Config, StructuredState from .simulation import SimulationAlgorithm, simulate, prepare_circuit_mps from .mps import DirMPS, MPS diff --git a/pytket/extensions/cutensornet/structured_state/general.py b/pytket/extensions/cutensornet/structured_state/general.py index d8de30f1..a458e723 100644 --- a/pytket/extensions/cutensornet/structured_state/general.py +++ b/pytket/extensions/cutensornet/structured_state/general.py @@ -26,11 +26,8 @@ import cupy as cp # type: ignore except ImportError: warnings.warn("local settings failed to import cupy", ImportWarning) -try: - import cuquantum.cutensornet as cutn # type: ignore -except ImportError: - warnings.warn("local settings failed to import cutensornet", ImportWarning) +from pytket.extensions.cutensornet import CuTensorNetHandle # An alias for the CuPy type used for tensors try: @@ -39,47 +36,6 @@ Tensor = Any -class CuTensorNetHandle: - """Initialise the cuTensorNet library with automatic workspace memory - management. - - Note: - Always use as ``with CuTensorNetHandle() as libhandle:`` so that cuTensorNet - handles are automatically destroyed at the end of execution. - - Attributes: - handle (int): The cuTensorNet library handle created by this initialisation. - device_id (int): The ID of the device (GPU) where cuTensorNet is initialised. - If not provided, defaults to ``cp.cuda.Device()``. - """ - - def __init__(self, device_id: Optional[int] = None): - self._is_destroyed = False - - # Make sure CuPy uses the specified device - cp.cuda.Device(device_id).use() - - dev = cp.cuda.Device() - self.device_id = int(dev) - - self.handle = cutn.create() - - def destroy(self) -> None: - """Destroys the memory handle, releasing memory. - - Only call this method if you are initialising a ``CuTensorNetHandle`` outside - a ``with CuTensorNetHandle() as libhandle`` statement. - """ - cutn.destroy(self.handle) - self._is_destroyed = True - - def __enter__(self) -> CuTensorNetHandle: - return self - - def __exit__(self, exc_type: Any, exc_value: Any, exc_tb: Any) -> None: - self.destroy() - - class Config: """Configuration class for simulation using ``StructuredState``.""" diff --git a/pytket/extensions/cutensornet/structured_state/mps.py b/pytket/extensions/cutensornet/structured_state/mps.py index dda12f34..c87aae5c 100644 --- a/pytket/extensions/cutensornet/structured_state/mps.py +++ b/pytket/extensions/cutensornet/structured_state/mps.py @@ -32,9 +32,9 @@ from pytket.circuit import Command, Op, OpType, Qubit from pytket.pauli import Pauli, QubitPauliString -from pytket.extensions.cutensornet.general import set_logger +from pytket.extensions.cutensornet.general import CuTensorNetHandle, set_logger -from .general import CuTensorNetHandle, Config, StructuredState, Tensor +from .general import Config, StructuredState, Tensor class DirMPS(Enum): diff --git a/pytket/extensions/cutensornet/structured_state/mps_mpo.py b/pytket/extensions/cutensornet/structured_state/mps_mpo.py index ad64b64a..ea29c6d8 100644 --- a/pytket/extensions/cutensornet/structured_state/mps_mpo.py +++ b/pytket/extensions/cutensornet/structured_state/mps_mpo.py @@ -29,7 +29,8 @@ warnings.warn("local settings failed to import cutensornet", ImportWarning) from pytket.circuit import Qubit -from .general import CuTensorNetHandle, Tensor, Config +from pytket.extensions.cutensornet import CuTensorNetHandle +from .general import Tensor, Config from .mps import ( DirMPS, MPS, diff --git a/pytket/extensions/cutensornet/structured_state/simulation.py b/pytket/extensions/cutensornet/structured_state/simulation.py index 0d68880e..88cc9002 100644 --- a/pytket/extensions/cutensornet/structured_state/simulation.py +++ b/pytket/extensions/cutensornet/structured_state/simulation.py @@ -32,8 +32,8 @@ from pytket.passes import DefaultMappingPass from pytket.predicates import CompilationUnit -from pytket.extensions.cutensornet.general import set_logger -from .general import CuTensorNetHandle, Config, StructuredState +from pytket.extensions.cutensornet.general import CuTensorNetHandle, set_logger +from .general import Config, StructuredState from .mps_gate import MPSxGate from .mps_mpo import MPSxMPO from .ttn_gate import TTNxGate @@ -60,7 +60,7 @@ def simulate( """Simulates the circuit and returns the ``StructuredState`` of the final state. Note: - A ``libhandle`` should be created via a ``with CuTensorNet() as libhandle:`` + A ``libhandle`` is created via a ``with CuTensorNetHandle() as libhandle:`` statement. The device where the ``StructuredState`` is stored will match the one specified by the library handle. diff --git a/pytket/extensions/cutensornet/structured_state/ttn.py b/pytket/extensions/cutensornet/structured_state/ttn.py index 59460190..9b36c6fa 100644 --- a/pytket/extensions/cutensornet/structured_state/ttn.py +++ b/pytket/extensions/cutensornet/structured_state/ttn.py @@ -33,9 +33,9 @@ from pytket.circuit import Command, Qubit from pytket.pauli import QubitPauliString -from pytket.extensions.cutensornet.general import set_logger +from pytket.extensions.cutensornet.general import CuTensorNetHandle, set_logger -from .general import CuTensorNetHandle, Config, StructuredState, Tensor +from .general import Config, StructuredState, Tensor class DirTTN(IntEnum): diff --git a/tests/test_cutensornet_postselect.py b/tests/test_cutensornet_postselect.py index faa379ff..99ef38c6 100644 --- a/tests/test_cutensornet_postselect.py +++ b/tests/test_cutensornet_postselect.py @@ -5,11 +5,13 @@ from pytket.pauli import Pauli, QubitPauliString # type: ignore from pytket.utils import QubitPauliOperator from pytket.extensions.cutensornet.backends import CuTensorNetBackend -from pytket.extensions.cutensornet.tensor_network_convert import ( # type: ignore +from pytket.extensions.cutensornet.general_state.tensor_network_convert import ( # type: ignore TensorNetwork, measure_qubits_state, ) -from pytket.extensions.cutensornet.utils import circuit_statevector_postselect +from pytket.extensions.cutensornet.general_state.utils import ( + circuit_statevector_postselect, +) @pytest.mark.parametrize( diff --git a/tests/test_general_state.py b/tests/test_general_state.py new file mode 100644 index 00000000..097f8c05 --- /dev/null +++ b/tests/test_general_state.py @@ -0,0 +1,206 @@ +import random +import numpy as np +import pytest +from pytket.circuit import ToffoliBox, Qubit +from pytket.passes import DecomposeBoxes, CnXPairwiseDecomposition +from pytket.transform import Transform +from pytket.pauli import QubitPauliString, Pauli +from pytket.utils.operators import QubitPauliOperator +from pytket.circuit import Circuit +from pytket.extensions.cutensornet.general_state import GeneralState +from pytket.extensions.cutensornet.structured_state import CuTensorNetHandle + + +@pytest.mark.parametrize( + "circuit", + [ + pytest.lazy_fixture("q5_empty"), # type: ignore + pytest.lazy_fixture("q8_empty"), # type: ignore + pytest.lazy_fixture("q2_x0"), # type: ignore + pytest.lazy_fixture("q2_x1"), # type: ignore + pytest.lazy_fixture("q2_v0"), # type: ignore + pytest.lazy_fixture("q2_x0cx01"), # type: ignore + pytest.lazy_fixture("q2_x1cx10x1"), # type: ignore + pytest.lazy_fixture("q2_x0cx01cx10"), # type: ignore + pytest.lazy_fixture("q2_v0cx01cx10"), # type: ignore + pytest.lazy_fixture("q2_hadamard_test"), # type: ignore + pytest.lazy_fixture("q2_lcu1"), # type: ignore + pytest.lazy_fixture("q2_lcu2"), # type: ignore + pytest.lazy_fixture("q2_lcu3"), # type: ignore + pytest.lazy_fixture("q3_v0cx02"), # type: ignore + pytest.lazy_fixture("q3_cx01cz12x1rx0"), # type: ignore + pytest.lazy_fixture("q3_toffoli_box_with_implicit_swaps"), # type: ignore + pytest.lazy_fixture("q4_lcu1"), # type: ignore + pytest.lazy_fixture("q4_multicontrols"), # type: ignore + pytest.lazy_fixture("q4_with_creates"), # type: ignore + pytest.lazy_fixture("q5_h0s1rz2ry3tk4tk13"), # type: ignore + pytest.lazy_fixture("q5_line_circ_30_layers"), # type: ignore + pytest.lazy_fixture("q6_qvol"), # type: ignore + pytest.lazy_fixture("q8_x0h2v5z6"), # type: ignore + ], +) +def test_get_statevec(circuit: Circuit) -> None: + with CuTensorNetHandle() as libhandle: + state = GeneralState(circuit, libhandle) + sv = state.get_statevector() + + sv_pytket = circuit.get_statevector() + assert np.allclose(sv, sv_pytket, atol=1e-10) + + op = QubitPauliOperator( + { + QubitPauliString({q: Pauli.I for q in circuit.qubits}): 1.0, + } + ) + + # Calculate the inner product as the expectation value + # of the identity operator: = + state = GeneralState(circuit, libhandle) + ovl = state.expectation_value(op) + assert ovl == pytest.approx(1.0) + + state.destroy() + + +def test_sv_toffoli_box_with_implicit_swaps() -> None: + # Using specific permutation here + perm = { + (False, False): (True, True), + (False, True): (False, False), + (True, False): (True, False), + (True, True): (False, True), + } + + # Create a circuit with more qubits and multiple applications of the permutation + # above + ket_circ = Circuit(3) + + # Create the circuit + ket_circ.add_toffolibox(ToffoliBox(perm), [Qubit(0), Qubit(1)]) # type: ignore + ket_circ.add_toffolibox(ToffoliBox(perm), [Qubit(1), Qubit(2)]) # type: ignore + + DecomposeBoxes().apply(ket_circ) + CnXPairwiseDecomposition().apply(ket_circ) + Transform.OptimiseCliffords().apply(ket_circ) + + # Convert and contract + with CuTensorNetHandle() as libhandle: + state = GeneralState(ket_circ, libhandle) + ket_net_vector = state.get_statevector() + state.destroy() + + # Compare to pytket statevector + ket_pytket_vector = ket_circ.get_statevector() + + assert np.allclose(ket_net_vector, ket_pytket_vector) + + +@pytest.mark.parametrize("n_qubits", [4, 5, 6]) +def test_sv_generalised_toffoli_box(n_qubits: int) -> None: + def to_bool_tuple(n_qubits: int, x: int) -> tuple: + bool_list = [] + for i in reversed(range(n_qubits)): + bool_list.append((x >> i) % 2 == 1) + return tuple(bool_list) + + random.seed(1) + + # Generate a random permutation + cycle = list(range(2**n_qubits)) + random.shuffle(cycle) + + perm = dict() + for orig, dest in enumerate(cycle): + perm[to_bool_tuple(n_qubits, orig)] = to_bool_tuple(n_qubits, dest) + + # Create a circuit implementing the permutation above + ket_circ = ToffoliBox(perm).get_circuit() # type: ignore + + DecomposeBoxes().apply(ket_circ) + CnXPairwiseDecomposition().apply(ket_circ) + Transform.OptimiseCliffords().apply(ket_circ) + + with CuTensorNetHandle() as libhandle: + state = GeneralState(ket_circ, libhandle) + ket_net_vector = state.get_statevector() + + ket_pytket_vector = ket_circ.get_statevector() + assert np.allclose(ket_net_vector, ket_pytket_vector) + + # Calculate the inner product as the expectation value + # of the identity operator: = + op = QubitPauliOperator( + { + QubitPauliString({q: Pauli.I for q in ket_circ.qubits}): 1.0, + } + ) + + state = GeneralState(ket_circ, libhandle) + ovl = state.expectation_value(op) + assert ovl == pytest.approx(1.0) + + state.destroy() + + +@pytest.mark.parametrize( + "circuit", + [ + pytest.lazy_fixture("q5_empty"), # type: ignore + pytest.lazy_fixture("q8_empty"), # type: ignore + pytest.lazy_fixture("q2_x0"), # type: ignore + pytest.lazy_fixture("q2_x1"), # type: ignore + pytest.lazy_fixture("q2_v0"), # type: ignore + pytest.lazy_fixture("q2_x0cx01"), # type: ignore + pytest.lazy_fixture("q2_x1cx10x1"), # type: ignore + pytest.lazy_fixture("q2_x0cx01cx10"), # type: ignore + pytest.lazy_fixture("q2_v0cx01cx10"), # type: ignore + pytest.lazy_fixture("q2_hadamard_test"), # type: ignore + pytest.lazy_fixture("q2_lcu1"), # type: ignore + pytest.lazy_fixture("q2_lcu2"), # type: ignore + pytest.lazy_fixture("q2_lcu3"), # type: ignore + pytest.lazy_fixture("q3_v0cx02"), # type: ignore + pytest.lazy_fixture("q3_cx01cz12x1rx0"), # type: ignore + pytest.lazy_fixture("q3_toffoli_box_with_implicit_swaps"), # type: ignore + pytest.lazy_fixture("q4_lcu1"), # type: ignore + pytest.lazy_fixture("q4_multicontrols"), # type: ignore + pytest.lazy_fixture("q4_with_creates"), # type: ignore + pytest.lazy_fixture("q5_h0s1rz2ry3tk4tk13"), # type: ignore + pytest.lazy_fixture("q5_line_circ_30_layers"), # type: ignore + pytest.lazy_fixture("q6_qvol"), # type: ignore + pytest.lazy_fixture("q8_x0h2v5z6"), # type: ignore + ], +) +@pytest.mark.parametrize( + "observable", + [ + QubitPauliOperator( + { + QubitPauliString({Qubit(0): Pauli.I, Qubit(1): Pauli.X}): 1.0, + } + ), + QubitPauliOperator( + { + QubitPauliString({Qubit(0): Pauli.X, Qubit(1): Pauli.Y}): 3.5 + 0.3j, + } + ), + QubitPauliOperator( + { + QubitPauliString({Qubit(0): Pauli.Z}): 0.25, + QubitPauliString({Qubit(1): Pauli.Y}): 0.33j, + QubitPauliString({Qubit(0): Pauli.X, Qubit(1): Pauli.X}): 0.42 + 0.1j, + } + ), + ], +) +def test_expectation_value(circuit: Circuit, observable: QubitPauliOperator) -> None: + # Note: not all qubits are acted on by the observable. The remaining qubits are + # interpreted to have I (identity) operators on them both by pytket and cutensornet. + exp_val_tket = observable.state_expectation(circuit.get_statevector()) + + # Calculate using GeneralState + with CuTensorNetHandle() as libhandle: + state = GeneralState(circuit, libhandle) + exp_val = state.expectation_value(observable) + + assert np.isclose(exp_val, exp_val_tket) + state.destroy() diff --git a/tests/test_structured_state.py b/tests/test_structured_state.py index c01e32f6..4a7ce396 100644 --- a/tests/test_structured_state.py +++ b/tests/test_structured_state.py @@ -21,7 +21,9 @@ SimulationAlgorithm, ) from pytket.extensions.cutensornet.structured_state.ttn import RootPath -from pytket.extensions.cutensornet.utils import circuit_statevector_postselect +from pytket.extensions.cutensornet.general_state.utils import ( + circuit_statevector_postselect, +) def test_libhandle_manager() -> None: diff --git a/tests/test_tensor_network_convert.py b/tests/test_tensor_network_convert.py index d221a264..43f4f95b 100644 --- a/tests/test_tensor_network_convert.py +++ b/tests/test_tensor_network_convert.py @@ -15,13 +15,13 @@ warnings.warn("local settings failed to import cutensornet", ImportWarning) from pytket.circuit import Circuit -from pytket.extensions.cutensornet.tensor_network_convert import ( # type: ignore +from pytket.extensions.cutensornet.general_state.tensor_network_convert import ( # type: ignore tk_to_tensor_network, TensorNetwork, ) -def state_contract(tn: List[Union[NDArray, List]], nqubit: int) -> NDArray: +def state_contract(tn: List[Union[NDArray, List]]) -> NDArray: """Calls cuQuantum contract function to contract an input state tensor network.""" state_tn = tn.copy() state: NDArray = cq.contract(*state_tn).flatten() @@ -59,7 +59,7 @@ def circuit_overlap_contract(circuit_ket: Circuit) -> float: ) def test_convert_statevec_overlap(circuit: Circuit) -> None: tn = tk_to_tensor_network(circuit) - result_cu = state_contract(tn, circuit.n_qubits).flatten().round(10) + result_cu = state_contract(tn).flatten().round(10) state_vector = np.array([circuit.get_statevector()]) assert np.allclose(result_cu, state_vector) ovl = circuit_overlap_contract(circuit) diff --git a/tests/test_utils.py b/tests/test_utils.py index 48453371..a4aad95f 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -1,5 +1,7 @@ import numpy -from pytket.extensions.cutensornet.utils import circuit_statevector_postselect +from pytket.extensions.cutensornet.general_state.utils import ( + circuit_statevector_postselect, +) from pytket import Circuit, Qubit # type: ignore