# Repository: https://gitlab.com/qblox/packages/software/qblox-scheduler
# Licensed according to the LICENSE file on the main branch
#
# Copyright 2020-2025, Quantify Consortium
# Copyright 2025, Qblox B.V.
"""
Utility class for dynamically allocating
Qblox acquisition indices and bins and for Qblox sequencers.
"""
from __future__ import annotations
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any
from qblox_scheduler.backends.qblox import constants
from qblox_scheduler.enums import BinMode
from qblox_scheduler.helpers.collections import make_hash
from qblox_scheduler.helpers.importers import export_python_object_to_path_string
from qblox_scheduler.json_utils import JSONSerializable
if TYPE_CHECKING:
from collections.abc import Hashable
from qblox_scheduler.backends.types.common import ThresholdedTriggerCountMetadata
from qblox_scheduler.helpers.generate_acq_channels_data import AcquisitionIndices
[docs]
QbloxAcquisitionIndex = int
@dataclass
"""
Acquisition hardware mapping for timetag trace acquisitions.
Each value maps the acquisition index to a hardware bin,
which is specified by the Qblox acquisition index, and the Qblox acquisition bin.
"""
@dataclass
[docs]
class BinnedAcqInfo:
"""Acquisition info for binned acquisition."""
[docs]
acq_index: AcquisitionIndices
def __getstate__(self) -> dict:
return {
"deserialization_type": export_python_object_to_path_string(self.__class__),
"data": {
"acq_channel": self.acq_channel,
"acq_index": self.acq_index,
"thresholded_trigger_count_metadata": self.thresholded_trigger_count_metadata,
},
}
def __setstate__(self, state: dict) -> None:
self.__init__(**state["data"])
@dataclass
[docs]
class BinnedAcqControlFlowNode:
"""
Node to represent all acquisitions which are within the same loop tree structure.
If and only if repetitions and bin_modes are not None, then it is a loop node.
"""
[docs]
parent: BinnedAcqControlFlowNode | None
[docs]
children: list[BinnedAcqControlFlowNode | BinnedAcqInfo]
[docs]
repetitions: int | None
[docs]
bin_mode: BinMode | None
def __eq__(self, other: object) -> bool:
return hash(self) == hash(other)
def __hash__(self) -> int:
return make_hash(
dict(children=self.children, repetitions=self.repetitions, bin_mode=self.bin_mode)
)
[docs]
def add_control_flow_child(
self, child_repetitions: int | None, bin_mode: BinMode | None
) -> BinnedAcqControlFlowNode:
"""
Adds a new control flow as a child to the current node,
and returns it.
"""
new_child = BinnedAcqControlFlowNode(
parent=self, children=[], repetitions=child_repetitions, bin_mode=bin_mode
)
self.children.append(new_child)
return new_child
[docs]
def return_control_flow_child(self) -> BinnedAcqControlFlowNode:
"""
Returns the parent, and
if the current node is empty, it will remove it from the tree.
Only call this function after any add_control_flow_child was called.
"""
if self.parent is None:
# While iterating the control flow tree,
# this can only be called after we return from a control flow,
# which means that the parent cannot be None.
raise RuntimeError("Compilation error: add_control_flow_child was not called before.")
if len(self.children) == 0:
# Removing the current branch of the tree
# if there are no children in it.
self.parent.children.pop()
return self.parent
@staticmethod
[docs]
def get_nearest_loop_node(
node: BinnedAcqControlFlowNode,
) -> BinnedAcqControlFlowNode:
"""
Returns argument if it's a loop node, otherwise returns
the parent's nearest loop node or argument if there's no parent (root node).
"""
if node.repetitions is not None or node.parent is None:
return node
else:
return BinnedAcqControlFlowNode.get_nearest_loop_node(node.parent)
def __getstate__(self) -> dict:
# Remove parents to not cause any circular references in serialization.
return {
"deserialization_type": export_python_object_to_path_string(self.__class__),
"data": {
"parent": None,
"children": self.children,
"repetitions": self.repetitions,
"bin_mode": self.bin_mode,
},
}
def __setstate__(self, state: dict) -> None:
# Parents in the serialized data were removed,
# to not cause circular references.
# We manually readd the parents in deserialization.
data = state["data"]
self.parent = None
self.children = data["children"]
for child in self.children:
if isinstance(child, BinnedAcqControlFlowNode):
child.parent = self
self.repetitions = data["repetitions"]
self.bin_mode = BinMode(data["bin_mode"]) if data["bin_mode"] is not None else None
@dataclass
[docs]
class QbloxAcquisitionHardwareMappingBinned:
"""Acquisition hardware mapping for binned acquisitions."""
[docs]
qblox_acq_index: QbloxAcquisitionIndex
[docs]
qblox_acq_bin_offset: int
"""
The starting bin where all acquisition data is stored for this mapping.
"""
[docs]
tree: BinnedAcqControlFlowNode
"""
Root node for the whole tree of the loop and acquisition tree.
"""
def __getstate__(self) -> dict:
return {
"deserialization_type": export_python_object_to_path_string(self.__class__),
"data": {
"qblox_acq_index": self.qblox_acq_index,
"qblox_acq_bin_offset": self.qblox_acq_bin_offset,
"tree": self.tree,
},
}
def __setstate__(self, state: dict) -> None:
self.__init__(**state["data"])
@dataclass
[docs]
class QbloxAcquisitionHardwareMapping(JSONSerializable):
"""Acquisition hardware mapping for all acquisitions."""
[docs]
non_binned: dict[Hashable, QbloxAcquisitionIndex]
[docs]
binned: list[QbloxAcquisitionHardwareMappingBinned]
def __getstate__(self) -> dict:
return {
"deserialization_type": export_python_object_to_path_string(self.__class__),
"data": {
"non_binned": self.non_binned,
"binned": self.binned,
"timetagtrace": self.timetagtrace,
},
}
def __setstate__(self, state: dict) -> None:
self.__init__(**state["data"])
[docs]
class AcquisitionMemoryError(ValueError):
"""Raised when there is an error in allocating acquisition memory."""
[docs]
class QbloxAcquisitionModuleResourceManager:
"""Utility class that keeps track of all the reserved acquisition resources for a module."""
def __init__(self, maximum_bins: int) -> None:
[docs]
self._maximum_bins = maximum_bins
@property
[docs]
def total_used_bins(self) -> int:
"""Total amount of in-use bins in module."""
return self._used_bins
@total_used_bins.setter
def total_used_bins(self, value: int) -> None:
self._used_bins = value
@property
[docs]
def total_remaining_free_bins(self) -> int:
"""Total amount of remaining available bins in module."""
return self._maximum_bins - self._used_bins
[docs]
class _SequencerAcquisitionModel:
def __init__(
self,
module_resource_manager: QbloxAcquisitionModuleResourceManager,
) -> None:
[docs]
self._num_bins: list[int] = []
[docs]
self._module_resource_manager = module_resource_manager
[docs]
self._maximum_qblox_acq_indices = constants.NUMBER_OF_QBLOX_ACQ_INDICES
[docs]
def reserve_new_qblox_acq_index(self, num_bins: int = 0) -> int:
if len(self._num_bins) >= self._maximum_qblox_acq_indices:
raise AcquisitionMemoryError("Out of Qblox acquisition indices.")
self._num_bins.append(0)
acq_index = len(self._num_bins) - 1
if num_bins:
self.reserve_bins(acq_index, num_bins)
return acq_index
[docs]
def reserve_bins(self, qblox_acq_index: int, num_bins: int) -> int:
"""
Reserves num_bins number of bins of the qblox_acq_index.
Returns the bin number of the first reserved bin.
"""
if num_bins > self._module_resource_manager.total_remaining_free_bins:
raise AcquisitionMemoryError("Out of Qblox acquisition bins.")
next_free_bin = self._num_bins[qblox_acq_index]
self._num_bins[qblox_acq_index] += num_bins
self._module_resource_manager.total_used_bins += num_bins
return next_free_bin
[docs]
def to_acq_declaration_dict(self) -> dict[str, Any]:
"""
Acquisition declaration dictionary.
This data is used in :class:`qblox_instruments.qcodes_drivers.Sequencer`
`sequence` parameter's `"acquisitions"`.
"""
return {
str(qblox_acq_index): {"num_bins": num_bins, "index": qblox_acq_index}
for qblox_acq_index, num_bins in enumerate(self._num_bins)
}
[docs]
class QbloxAcquisitionIndexManager:
"""
Utility class that keeps track of all the reserved indices, bins for a sequencer.
Each acquisition channel is mapped to a unique Qblox acquisition index.
For binned acquisitions, each new allocation request reserves
the Qblox acquisition bins in order (incrementing the bin index by one).
For trace and ttl and other acquisitions, the whole Qblox acquisition index is reserved,
there, the bin index has no relevance.
"""
def __init__(self, module_resource_manager: QbloxAcquisitionModuleResourceManager) -> None:
[docs]
self._acq_hardware_mapping_not_binned: dict[Hashable, QbloxAcquisitionIndex] = {}
"""
Acquisition hardware mapping for not binned acquisitions.
"""
[docs]
self._sequencer_acquisition_model = _SequencerAcquisitionModel(module_resource_manager)
"""
Data model of sequencer acquisition memory, which keeps track of the allocated
amount of bins for each Qblox acquisition index.
"""
[docs]
self._binned_qblox_acq_index: int | None = None
"""
Qblox acquisition index used by the binned acquisitions.
"""
[docs]
self._trace_allocated: bool = False
"""
Specifying whether a Trace or TimetagTrace have already been allocated.
"""
[docs]
self._acq_hardware_mapping_binned: list[QbloxAcquisitionHardwareMappingBinned] = []
"""
Acquisition hardware mapping for time tag trace.
"""
[docs]
def allocate_bins_binned(
self,
number_of_acq_indices: int,
tree: BinnedAcqControlFlowNode,
repetitions: int | None,
) -> tuple[int, int]:
"""
Allocates Qblox acquisition bins for binned acquisitions.
Parameters
----------
number_of_acq_indices
Number of acquisition indices to allocate.
tree
The loop tree structure for all of the acquisitions.
repetitions
Repetitions of the schedule when using append bin mode.
Returns
-------
The Qblox acquisition index, and the Qblox acquisition bin offset as integers.
Raises
------
AcquisitionMemoryError
When the QbloxAcquisitionBinManager runs out of bins to allocate.
"""
# Currently only repetitions=1 is implemented.
assert repetitions == 1
self._binned_qblox_acq_index = (
self._sequencer_acquisition_model.reserve_new_qblox_acq_index()
if (self._binned_qblox_acq_index is None)
else self._binned_qblox_acq_index
)
next_free_qblox_bin = self._sequencer_acquisition_model.reserve_bins(
self._binned_qblox_acq_index, number_of_acq_indices * repetitions
)
self._acq_hardware_mapping_binned.append(
QbloxAcquisitionHardwareMappingBinned(
self._binned_qblox_acq_index, next_free_qblox_bin, tree
)
)
return self._binned_qblox_acq_index, next_free_qblox_bin
[docs]
def allocate_qblox_index(self, acq_channel: Hashable) -> int:
"""
Allocates a whole Qblox acquisition index for ttl, other acquisition
for the given acquisition channel.
Parameters
----------
acq_channel
Acquisition channel.
Returns
-------
The Qblox acquisition index.
Raises
------
AcquisitionMemoryError
When the QbloxAcquisitionBinManager runs out of acquisition indices to allocate.
"""
if acq_channel in self._acq_hardware_mapping_not_binned:
return self._acq_hardware_mapping_not_binned[acq_channel]
qblox_acq_index: int = self._sequencer_acquisition_model.reserve_new_qblox_acq_index()
self._acq_hardware_mapping_not_binned[acq_channel] = qblox_acq_index
self._sequencer_acquisition_model.reserve_bins(
qblox_acq_index, constants.MAX_NUMBER_OF_RUNTIME_ALLOCATED_QBLOX_ACQ_BINS
)
return qblox_acq_index
[docs]
def allocate_trace(self, acq_channel: Hashable) -> tuple[int, int]:
"""
Allocates a whole Qblox acquisition index for trace
for the given acquisition channel.
Parameters
----------
acq_channel
Acquisition channel.
Returns
-------
The Qblox acquisition index, and the Qblox acquisition bin offset as integers.
Raises
------
AcquisitionMemoryError
When the QbloxAcquisitionBinManager runs out of acquisition indices to allocate.
"""
if acq_channel in self._acq_hardware_mapping_not_binned:
return self._acq_hardware_mapping_not_binned[acq_channel], 0
elif self._trace_allocated:
raise AcquisitionMemoryError(
f"Only one acquisition channel per port-clock can be specified, "
f"if the 'Trace' acquisition protocol is used. "
f"Attempted to compile for acquisition channel '{acq_channel}'."
)
qblox_acq_index: int = self._sequencer_acquisition_model.reserve_new_qblox_acq_index()
self._acq_hardware_mapping_not_binned[acq_channel] = qblox_acq_index
self._sequencer_acquisition_model.reserve_bins(qblox_acq_index, 1)
self._trace_allocated = True
return qblox_acq_index, 0
[docs]
def acq_declaration_dict(self) -> dict[str, Any]:
"""
Returns the acquisition declaration dict, which is needed for the qblox-instruments.
This data is used in :class:`qblox_instruments.qcodes_drivers.Sequencer`
`sequence` parameter's `"acquisitions"`.
Returns
-------
The acquisition declaration dict.
"""
return self._sequencer_acquisition_model.to_acq_declaration_dict()
[docs]
def acq_hardware_mapping(
self,
) -> QbloxAcquisitionHardwareMapping:
"""
Returns the acquisition hardware mapping, which is needed for
qblox-scheduler instrument coordinator to figure out which hardware index, bin needs
to be mapped to which output acquisition data.
Returns
-------
The acquisition hardware mapping.
"""
return QbloxAcquisitionHardwareMapping(
non_binned=self._acq_hardware_mapping_not_binned,
binned=self._acq_hardware_mapping_binned,
timetagtrace=self._acq_hardware_mapping_timetagtrace,
)
