From 9aef0acdaa7571f9415bbfcdcfd5aa026942eb77 Mon Sep 17 00:00:00 2001 From: Aleksander Wennersteen Date: Thu, 7 Mar 2024 11:33:47 +0100 Subject: [PATCH 1/6] [FEAT] Debug logging and NVTX annotations --- pyproject.toml | 2 +- pyqtorch/__init__.py | 10 +++++++++ pyqtorch/analog.py | 33 +++++++++++++++++++++++++++++ pyqtorch/circuit.py | 49 ++++++++++++++++++++++++++++++++++++------- pyqtorch/primitive.py | 27 ++++++++++++++++++++++++ pyqtorch/utils.py | 13 +++++++++++- 6 files changed, 124 insertions(+), 10 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 5f13d939..d54b983e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -16,7 +16,7 @@ authors = [ ] requires-python = ">=3.8,<3.13" license = {text = "Apache 2.0"} -version = "1.0.6" +version = "1.0.7" classifiers=[ "License :: OSI Approved :: Apache Software License", "Programming Language :: Python", diff --git a/pyqtorch/__init__.py b/pyqtorch/__init__.py index faef5ad1..3cb68f7a 100644 --- a/pyqtorch/__init__.py +++ b/pyqtorch/__init__.py @@ -9,6 +9,8 @@ # # limitations under the License. from __future__ import annotations +import logging + import torch from .analog import HamiltonianEvolution @@ -43,3 +45,11 @@ ) torch.set_default_dtype(torch.float64) + +if __name__ not in logging.Logger.manager.loggerDict.keys(): + import sys + from logging import getLogger + + _logger = getLogger(__name__) + _logger.setLevel(logging.INFO) + _logger.addHandler(logging.StreamHandler(sys.stderr)) diff --git a/pyqtorch/analog.py b/pyqtorch/analog.py index 72f087b1..46531b0a 100644 --- a/pyqtorch/analog.py +++ b/pyqtorch/analog.py @@ -1,5 +1,7 @@ from __future__ import annotations +import logging +from logging import getLogger from typing import Tuple import torch @@ -10,6 +12,29 @@ BATCH_DIM = 2 +logger = getLogger(__name__) + + +def forward_hook(*args, **kwargs) -> None: # type: ignore[no-untyped-def] + logger.debug("Forward complete") + torch.cuda.nvtx.range_pop() + + +def pre_forward_hook(*args, **kwargs) -> None: # type: ignore[no-untyped-def] + logger.debug("Executing forward") + torch.cuda.nvtx.range_push("HamiltonianEvolution.forward") + + +def backward_hook(*args, **kwargs) -> None: # type: ignore[no-untyped-def] + logger.debug("Backward complete") + torch.cuda.nvtx.range_pop() + + +def pre_backward_hook(*args, **kwargs) -> None: # type: ignore[no-untyped-def] + logger.debug("Executed backward") + torch.cuda.nvtx.range_push("Hamiltonian Evolution.backward") + + class HamiltonianEvolution(torch.nn.Module): def __init__( self, @@ -36,6 +61,14 @@ def _matrixexp_operator(hamiltonian: Operator, time_evolution: torch.Tensor) -> self._evolve_diag_operator = _diag_operator self._evolve_matrixexp_operator = _matrixexp_operator + logger.debug("Hamiltonian Evolution initialized") + if logger.isEnabledFor(logging.DEBUG): + # When Debugging let's add logging and NVTX markers + # WARNING: incurs performance penalty + self.register_forward_hook(forward_hook, always_call=True) + self.register_full_backward_hook(backward_hook) + self.register_forward_pre_hook(pre_forward_hook) + self.register_full_backward_pre_hook(pre_backward_hook) def forward( self, diff --git a/pyqtorch/circuit.py b/pyqtorch/circuit.py index 1835ecc6..04518f4c 100644 --- a/pyqtorch/circuit.py +++ b/pyqtorch/circuit.py @@ -1,8 +1,10 @@ from __future__ import annotations +import logging from logging import getLogger from typing import Any, Iterator +import torch from torch import Tensor from torch import device as torch_device from torch.nn import Module, ModuleList @@ -10,6 +12,27 @@ from pyqtorch.utils import DiffMode, State, inner_prod, zero_state logger = getLogger(__name__) +logger.setLevel(logging.DEBUG) + + +def forward_hook(*args, **kwargs) -> None: # type: ignore[no-untyped-def] + logger.debug("Forward complete") + torch.cuda.nvtx.range_pop() + + +def pre_forward_hook(*args, **kwargs) -> None: # type: ignore[no-untyped-def] + logger.debug("Executing forward") + torch.cuda.nvtx.range_push("QuantumCircuit.forward") + + +def backward_hook(*args, **kwargs) -> None: # type: ignore[no-untyped-def] + logger.debug("Backward complete") + torch.cuda.nvtx.range_pop() + + +def pre_backward_hook(*args, **kwargs) -> None: # type: ignore[no-untyped-def] + logger.debug("Executed backward") + torch.cuda.nvtx.range_push("QuantumCircuit.backward") class QuantumCircuit(Module): @@ -23,6 +46,16 @@ def __init__(self, n_qubits: int, operations: list[Module]): self._device = next(iter(set((op.device for op in self.operations)))) except StopIteration: pass + logger.debug("QuantumCircuit initialized") + logger.info("Hello") + logger.exception("hello") + if logger.isEnabledFor(logging.DEBUG): + # When Debugging let's add logging and NVTX markers + # WARNING: incurs performance penalty + self.register_forward_hook(forward_hook, always_call=True) + self.register_full_backward_hook(backward_hook) + self.register_forward_pre_hook(pre_forward_hook) + self.register_full_backward_pre_hook(pre_backward_hook) def __mul__(self, other: Module | QuantumCircuit) -> QuantumCircuit: n_qubits = max(self.n_qubits, other.n_qubits) @@ -50,15 +83,15 @@ def __eq__(self, other: Any) -> bool: def __hash__(self) -> int: return hash(self.__key()) - def run(self, state: State = None, values: dict[str, Tensor] = {}) -> State: - if state is None: - state = self.init_state() + def forward(self, state: State, values: dict[str, Tensor] = {}) -> State: for op in self.operations: state = op(state, values) return state - def forward(self, state: State, values: dict[str, Tensor] = {}) -> State: - return self.run(state, values) + def run(self, state: State = None, values: dict[str, Tensor] = {}) -> State: + if state is None: + state = self.init_state() + return self.forward(state, values) @property def device(self) -> torch_device: @@ -95,14 +128,14 @@ def expectation( """ if observable is None: raise ValueError("Please provide an observable to compute expectation.") - if state is None: - state = circuit.init_state(batch_size=1) if diff_mode == DiffMode.AD: - state = circuit.run(state, values) + state = circuit.forward(state, values) return inner_prod(state, observable.forward(state, values)).real elif diff_mode == DiffMode.ADJOINT: from pyqtorch.adjoint import AdjointExpectation + if state is None: + state = circuit.init_state(batch_size=1) return AdjointExpectation.apply(circuit, observable, state, values.keys(), *values.values()) else: raise ValueError(f"Requested diff_mode '{diff_mode}' not supported.") diff --git a/pyqtorch/primitive.py b/pyqtorch/primitive.py index 3254bcb3..00674ed8 100644 --- a/pyqtorch/primitive.py +++ b/pyqtorch/primitive.py @@ -1,5 +1,7 @@ from __future__ import annotations +import logging +from logging import getLogger from math import log2 from typing import Tuple @@ -9,6 +11,24 @@ from pyqtorch.matrices import OPERATIONS_DICT, _controlled, _dagger from pyqtorch.utils import Operator, State, product_state +logger = getLogger(__name__) + + +def forward_hook(*args, **kwargs) -> None: # type: ignore[no-untyped-def] + torch.cuda.nvtx.range_pop() + + +def pre_forward_hook(*args, **kwargs) -> None: # type: ignore[no-untyped-def] + torch.cuda.nvtx.range_push("Primitive.forward") + + +def backward_hook(*args, **kwargs) -> None: # type: ignore[no-untyped-def] + torch.cuda.nvtx.range_pop() + + +def pre_backward_hook(*args, **kwargs) -> None: # type: ignore[no-untyped-def] + torch.cuda.nvtx.range_push("Primitive.backward") + class Primitive(torch.nn.Module): def __init__(self, pauli: torch.Tensor, target: int) -> None: @@ -19,6 +39,13 @@ def __init__(self, pauli: torch.Tensor, target: int) -> None: self.register_buffer("pauli", pauli) self._param_type = None self._device = self.pauli.device + if logger.isEnabledFor(logging.DEBUG): + # When Debugging let's add logging and NVTX markers + # WARNING: incurs performance penalty + self.register_forward_hook(forward_hook, always_call=True) + self.register_full_backward_hook(backward_hook) + self.register_forward_pre_hook(pre_forward_hook) + self.register_full_backward_pre_hook(pre_backward_hook) def __key(self) -> tuple: return self.qubit_support diff --git a/pyqtorch/utils.py b/pyqtorch/utils.py index aae18c60..d23eb0f3 100644 --- a/pyqtorch/utils.py +++ b/pyqtorch/utils.py @@ -1,6 +1,8 @@ from __future__ import annotations +import logging from enum import Enum +from logging import getLogger from typing import Sequence import torch @@ -10,12 +12,21 @@ State = torch.Tensor Operator = torch.Tensor +logger = getLogger(__name__) + def inner_prod(bra: torch.Tensor, ket: torch.Tensor) -> torch.Tensor: + if logger.isEnabledFor(logging.DEBUG): + logger.debug("Inner prod calculation") + torch.cuda.nvtx.range_push("inner_prod") n_qubits = len(bra.size()) - 1 bra = bra.reshape((2**n_qubits, bra.size(-1))) ket = ket.reshape((2**n_qubits, ket.size(-1))) - return torch.einsum("ib,ib->b", bra.conj(), ket) + res = torch.einsum("ib,ib->b", bra.conj(), ket) + if logger.isEnabledFor(logging.DEBUG): + torch.cuda.nvtx.range_pop() + logger.debug("Inner prod complete") + return res def overlap(bra: torch.Tensor, ket: torch.Tensor) -> torch.Tensor: From f8bb025d358530df50cb2627d49b0063465fde03 Mon Sep 17 00:00:00 2001 From: Aleksander Wennersteen Date: Thu, 7 Mar 2024 13:55:51 +0100 Subject: [PATCH 2/6] Setup logging and torch defaults before importing --- pyproject.toml | 2 +- pyqtorch/__init__.py | 49 +++++++++++++++++++++++++++++++++++--------- 2 files changed, 40 insertions(+), 11 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index d54b983e..3a6db84a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -70,7 +70,7 @@ line-length = 100 required-imports = ["from __future__ import annotations"] [tool.ruff.per-file-ignores] -"__init__.py" = ["F401"] +"__init__.py" = ["F401", "E402"] [tool.ruff.mccabe] max-complexity = 15 diff --git a/pyqtorch/__init__.py b/pyqtorch/__init__.py index 3cb68f7a..bb215f09 100644 --- a/pyqtorch/__init__.py +++ b/pyqtorch/__init__.py @@ -10,9 +10,48 @@ from __future__ import annotations import logging +import os +import sys import torch +torch.set_default_dtype(torch.float64) + +logging_levels = { + "DEBUG": logging.DEBUG, + "INFO": logging.INFO, + "WARNING": logging.WARNING, + "ERROR": logging.ERROR, + "CRITICAL": logging.CRITICAL, +} + +LOG_BASE_LEVEL = os.environ.get("PYQ_LOG_LEVEL", None) +QADENCE_LOG_LEVEL = os.environ.get("QADENCE_LOG_LEVEL", None) +LOG_LEVEL: str | None | int = QADENCE_LOG_LEVEL if not LOG_BASE_LEVEL else LOG_BASE_LEVEL + + +if LOG_LEVEL: + LOG_LEVEL: int = logging_levels.get(LOG_LEVEL, logging.INFO) # type: ignore[arg-type, no-redef] + # If logger not setup, add handler to stderr + # else use setup presumably from Qadence + handle = None + if __name__ not in logging.Logger.manager.loggerDict.keys(): + handle = logging.StreamHandler(sys.stderr) + handle.set_name("console") + + logger = logging.getLogger(__name__) + if handle: + logger.addHandler(handle) + + logger.setLevel(LOG_LEVEL) + [ + h.setLevel(LOG_LEVEL) # type: ignore[func-returns-value] + for h in logger.handlers + if h.get_name() == "console" + ] + logger.debug("PyQTorch logger successfully setup") + + from .analog import HamiltonianEvolution from .apply import apply_operator from .circuit import QuantumCircuit, expectation @@ -43,13 +82,3 @@ uniform_state, zero_state, ) - -torch.set_default_dtype(torch.float64) - -if __name__ not in logging.Logger.manager.loggerDict.keys(): - import sys - from logging import getLogger - - _logger = getLogger(__name__) - _logger.setLevel(logging.INFO) - _logger.addHandler(logging.StreamHandler(sys.stderr)) From 6dc54466b18280cb66f439c3e2fa7cd3e8d2ee92 Mon Sep 17 00:00:00 2001 From: Aleksander Wennersteen Date: Fri, 8 Mar 2024 08:06:01 +0100 Subject: [PATCH 3/6] Remove errnoeous logging and convert level to uppercase automatically --- pyqtorch/__init__.py | 6 +++--- pyqtorch/circuit.py | 2 -- 2 files changed, 3 insertions(+), 5 deletions(-) diff --git a/pyqtorch/__init__.py b/pyqtorch/__init__.py index bb215f09..0673ec65 100644 --- a/pyqtorch/__init__.py +++ b/pyqtorch/__init__.py @@ -25,9 +25,9 @@ "CRITICAL": logging.CRITICAL, } -LOG_BASE_LEVEL = os.environ.get("PYQ_LOG_LEVEL", None) -QADENCE_LOG_LEVEL = os.environ.get("QADENCE_LOG_LEVEL", None) -LOG_LEVEL: str | None | int = QADENCE_LOG_LEVEL if not LOG_BASE_LEVEL else LOG_BASE_LEVEL +LOG_BASE_LEVEL = os.environ.get("PYQ_LOG_LEVEL", "").upper() +QADENCE_LOG_LEVEL = os.environ.get("QADENCE_LOG_LEVEL", "").upper() +LOG_LEVEL: str | int = QADENCE_LOG_LEVEL if not LOG_BASE_LEVEL else LOG_BASE_LEVEL if LOG_LEVEL: diff --git a/pyqtorch/circuit.py b/pyqtorch/circuit.py index 04518f4c..823b14a1 100644 --- a/pyqtorch/circuit.py +++ b/pyqtorch/circuit.py @@ -47,8 +47,6 @@ def __init__(self, n_qubits: int, operations: list[Module]): except StopIteration: pass logger.debug("QuantumCircuit initialized") - logger.info("Hello") - logger.exception("hello") if logger.isEnabledFor(logging.DEBUG): # When Debugging let's add logging and NVTX markers # WARNING: incurs performance penalty From 22db75701b7481382dbe6d9fea06df6c30e44f26 Mon Sep 17 00:00:00 2001 From: Aleksander Wennersteen Date: Wed, 15 May 2024 20:54:18 +0200 Subject: [PATCH 4/6] Finalise --- pyqtorch/__init__.py | 40 +++++++++++++++++++--------------------- 1 file changed, 19 insertions(+), 21 deletions(-) diff --git a/pyqtorch/__init__.py b/pyqtorch/__init__.py index d0418e54..522d54b0 100644 --- a/pyqtorch/__init__.py +++ b/pyqtorch/__init__.py @@ -25,31 +25,29 @@ "CRITICAL": logging.CRITICAL, } -LOG_BASE_LEVEL = os.environ.get("PYQ_LOG_LEVEL", "").upper() -QADENCE_LOG_LEVEL = os.environ.get("QADENCE_LOG_LEVEL", "").upper() -LOG_LEVEL: str | int = QADENCE_LOG_LEVEL if not LOG_BASE_LEVEL else LOG_BASE_LEVEL +LOG_LEVEL: str = os.environ.get("PYQ_LOG_LEVEL", "").upper() -if LOG_LEVEL: - LOG_LEVEL: int = logging_levels.get(LOG_LEVEL, logging.INFO) # type: ignore[arg-type, no-redef] - # If logger not setup, add handler to stderr - # else use setup presumably from Qadence - handle = None - if __name__ not in logging.Logger.manager.loggerDict.keys(): - handle = logging.StreamHandler(sys.stderr) - handle.set_name("console") +# if LOG_LEVEL: +LOG_LEVEL: int = logging_levels.get(LOG_LEVEL, logging.INFO) # type: ignore[arg-type, no-redef] +# If logger not setup, add handler to stderr +# else use setup presumably from Qadence +handle = None +if __name__ not in logging.Logger.manager.loggerDict.keys(): + handle = logging.StreamHandler(sys.stderr) + handle.set_name("console") - logger = logging.getLogger(__name__) - if handle: - logger.addHandler(handle) +logger = logging.getLogger(__name__) +if handle: + logger.addHandler(handle) +[ + h.setLevel(LOG_LEVEL) # type: ignore[func-returns-value] + for h in logger.handlers + if h.get_name() == "console" +] +logger.setLevel(LOG_LEVEL) - logger.setLevel(LOG_LEVEL) - [ - h.setLevel(LOG_LEVEL) # type: ignore[func-returns-value] - for h in logger.handlers - if h.get_name() == "console" - ] - logger.debug("PyQTorch logger successfully setup") +logger.info(f"PyQTorch logger successfully setup with log level {LOG_LEVEL}") from .adjoint import expectation From 3a3f7e44794979653e5bad72937726972410ef79 Mon Sep 17 00:00:00 2001 From: Aleksander Wennersteen Date: Wed, 15 May 2024 20:58:06 +0200 Subject: [PATCH 5/6] Remove dead code, possibly added by automation? --- pyqtorch/primitive.py | 9 --------- 1 file changed, 9 deletions(-) diff --git a/pyqtorch/primitive.py b/pyqtorch/primitive.py index 730d49bd..81d9e669 100644 --- a/pyqtorch/primitive.py +++ b/pyqtorch/primitive.py @@ -47,15 +47,6 @@ def __init__(self, pauli: Tensor, target: int) -> None: self.register_forward_pre_hook(pre_forward_hook) self.register_full_backward_pre_hook(pre_backward_hook) - def __key(self) -> tuple: - return self.qubit_support - - def __eq__(self, other: object) -> bool: - if isinstance(other, type(self)): - return self.__key() == other.__key() - else: - return False - def __hash__(self) -> int: return hash(self.qubit_support) From 9323296fea37ce9c029e2f3b879fca3e0d79d42e Mon Sep 17 00:00:00 2001 From: seitzdom Date: Tue, 28 May 2024 17:48:47 +0200 Subject: [PATCH 6/6] add docs --- docs/index.md | 17 +++++++++++++++++ pyqtorch/__init__.py | 4 ---- pyqtorch/circuit.py | 1 - 3 files changed, 17 insertions(+), 5 deletions(-) diff --git a/docs/index.md b/docs/index.md index d99c2f69..1e0203c0 100644 --- a/docs/index.md +++ b/docs/index.md @@ -378,3 +378,20 @@ def fig_to_html(fig: Figure) -> str: # markdown-exec: hide return buffer.getvalue() # markdown-exec: hide print(fig_to_html(plt.gcf())) # markdown-exec: hide ``` + +## CUDA Profiling and debugging + +To debug your quantum programs on `CUDA` devices, `pyqtorch` offers a `DEBUG` mode, which can be activated via +setting the `PYQ_LOG_LEVEL` environment variable. + +```bash +export PYQ_LOG_LEVEL=DEBUG +``` + +Before running your script, make sure to install the following packages: + +```bash +pip install nvidia-pyindex +pip install nvidia-dlprof[pytorch] +``` +For more information, check [the dlprof docs](https://docs.nvidia.com/deeplearning/frameworks/dlprof-user-guide/index.html). diff --git a/pyqtorch/__init__.py b/pyqtorch/__init__.py index 522d54b0..d65e93e6 100644 --- a/pyqtorch/__init__.py +++ b/pyqtorch/__init__.py @@ -13,10 +13,6 @@ import os import sys -import torch - -torch.set_default_dtype(torch.float64) - logging_levels = { "DEBUG": logging.DEBUG, "INFO": logging.INFO, diff --git a/pyqtorch/circuit.py b/pyqtorch/circuit.py index 47b68af5..f85563a3 100644 --- a/pyqtorch/circuit.py +++ b/pyqtorch/circuit.py @@ -18,7 +18,6 @@ from pyqtorch.utils import State, add_batch_dim, zero_state logger = getLogger(__name__) -logger.setLevel(logging.DEBUG) def forward_hook(*args, **kwargs) -> None: # type: ignore[no-untyped-def]