Interface to "high-level" API (State, NetworkOperator and Expectation)

docs/api.rst

@@ -2,5 +2,5 @@ API documentation
 -----------------
 
 .. toctree::
-    modules/fullTN.rst
+    modules/general_state.rst
     modules/structured_state.rst

docs/modules/general_state.rst

@@ -0,0 +1,32 @@
+Full tensor network (general state) contraction
+===============================================
+
+.. automodule:: pytket.extensions.cutensornet.general_state
+
+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
+
+cuQuantum `high-level` API interface
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: pytket.extensions.cutensornet.general_state.GeneralState
+
+.. autoclass:: pytket.extensions.cutensornet.general_state.GeneralOperator
+
+.. autoclass:: pytket.extensions.cutensornet.general_state.GeneralExpectationValue
+
+Pytket backend
+~~~~~~~~~~~~~~
+
+.. automodule:: pytket.extensions.cutensornet
+    :members: CuTensorNetBackend

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

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,

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,

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 conversion from tket primitives to cuQuantum primitives."""
+
+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

pytket/extensions/cutensornet/tensor_network_convert.py → 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