"""Abstractions for constructing QEC architectures."""
import json
from collections.abc import Mapping
from dataclasses import dataclass
from typing import Any, Literal, Protocol, overload
from guppylang.defs import GuppyFunctionDefinition
from hugr import Hugr
from hugr.metadata import Metadata
from hugr.ops import Module
from hugr.package import Package
from hugr.passes.composable import ComposablePass
from tket.passes import Normalize
from ._compile import LogicalCompiler, ReplacementCompiler, UncompilableError
from ._implement_ops import (
ImplementOps,
ImplementOpsSpec,
OpReplacements,
TyReplacements,
implement_ops,
)
__all__ = [
"EncodeSpec",
"EncoderParams",
"ImplementOps",
"ImplementOpsSpec",
"LogicalCompiler",
"OpReplacements",
"ReplacementCompiler",
"TyReplacements",
"UncompilableError",
"annotate_encoding",
"encode",
"implement_ops",
]
[docs]
@dataclass(frozen=True, kw_only=True)
class EncodeSpec:
"""A QEC-code-specific collection of passes that together fully encode a
computation."""
compile: LogicalCompiler | None = None
"""Pass to lower the computation to a logical level."""
logical_passes: list[ComposablePass] | None = None
"""Additional passes to run on the logical HUGR."""
implement_ops: ImplementOps | None = None
"""Lowers the logical computation to a physical level."""
@overload
def encode(
hugr: Package | GuppyFunctionDefinition[[], None],
spec: EncodeSpec,
*,
passes: list[ComposablePass] | None = None,
as_bytes: Literal[False] = False,
) -> Package: ...
@overload
def encode(
hugr: Package | GuppyFunctionDefinition[[], None],
spec: EncodeSpec,
*,
passes: list[ComposablePass] | None = None,
as_bytes: Literal[True],
) -> bytes: ...
[docs]
def encode(
hugr: Package | GuppyFunctionDefinition[[], None],
spec: EncodeSpec,
*,
passes: list[ComposablePass] | None = None,
as_bytes: bool = False,
) -> Package | bytes:
"""Encodes the given package (or Guppy function, directly compiled to a package for
convenience) by applying four stages.
1. running the given computational passes;
2. lowering the operations in the package to logical operations and
potentially performing static optimizations (e.g. resolving some qubit
address assignments statically);
3. running additional logical passes (e.g. inserting additional QEC cycles);
and
4. implementing the logical operations with physical gates.
The returned runnable package is guaranteed to be semantically equivalent to the
given one.
Args:
hugr: The package to encode, or a Guppy function for convenience.
spec: See :class:`EncodeSpec`.
passes: Computational passes to run on the given package. Defaults to one run
of :class:`Normalize`.
as_bytes: Whether to return the encoded package as bytes.
Returns:
The encoded runnable package, or its serialized representation when `as_bytes`
is `True`.
"""
pkg = hugr.compile_function() if isinstance(hugr, GuppyFunctionDefinition) else hugr
assert len(pkg.modules) == 1, "Given package contains more than one module"
assert not isinstance(pkg.modules[0].entrypoint_op(), Module), (
"Cannot process module-rooted HUGRs"
)
# 1. Passes with computational -> computational
for tket_pass in passes or [Normalize()]:
tket_pass(pkg.modules[0], inplace=True)
# 2. Lower computational -> logical
if spec.compile is not None:
pkg = spec.compile(pkg)
# 3. Passes with logical -> logical
for tket_pass in spec.logical_passes or []:
tket_pass(pkg.modules[0], inplace=True)
# 4. Lower logical -> physical
if spec.implement_ops is not None:
pkg_maybe_bytes = spec.implement_ops(pkg, as_bytes=as_bytes) # type: ignore[call-overload]
else:
pkg_maybe_bytes = pkg
if isinstance(pkg_maybe_bytes, Package) and as_bytes:
pkg_maybe_bytes = pkg_maybe_bytes.to_bytes()
return pkg_maybe_bytes
[docs]
class EncoderParams(Protocol):
"""Parameters used to annotate an encoded program."""
[docs]
def encoding(self) -> str:
"""The encoding to annotate on a program."""
[docs]
def params(self) -> Mapping[str, Any]:
"""The parameters to annotate on a program. Implementations should return values
that support serialization to JSON."""
class _MetadataEncoding(Metadata[Mapping[str, Any]]):
"""Metadata key for annotating parameters with which to encode a program."""
KEY = "guppyft.encoding"
[docs]
def annotate_encoding(hugr: Package | Hugr[Any], params: EncoderParams) -> None:
"""Annotate a HUGR package with encoding parameters.
Args:
hugr: The package or HUGR to annotate.
params: The serializable encoding parameters.
Raises:
ValueError: If `params` cannot be serialized as JSON.
"""
try:
json.dumps(params.params(), check_circular=True)
except TypeError as e:
raise ValueError("Could not serialize parameters") from e
for module in hugr.modules if isinstance(hugr, Package) else [hugr]:
module[module.module_root].metadata[_MetadataEncoding] = {
"encoding": params.encoding(),
"params": params.params(),
}