from kosmos.circuit_runner.qiskit_runner import QiskitRunner
from kosmos.dqc_scheduling.event import EventId
from kosmos.dqc_scheduling.event_queue import EventQueue
from kosmos.dqc_scheduling.execution_scheduler import ExecutionStep
from kosmos.protocols.circuit_execution_protocol import CircuitExecutionProtocol
from kosmos.protocols.config.protocol import (
CircuitExecutionProtocolConfig,
RemoteOperationProtocolConfig,
)
from kosmos.protocols.eg_protocol import EGProtocol, EGProtocolConfig
from kosmos.protocols.remote_operation_protocol import RemoteOperationProtocol
from kosmos.quantum_logic.quantum_register_manager import QuantumRegisterManager
from kosmos.topology.net import Network
from kosmos.topology.node import NodeId, QuantumNode
[docs]
def schedule_circuit_execution(
step: ExecutionStep,
event_queue: EventQueue,
network: Network,
circuit_runner: QiskitRunner,
) -> None:
"""Schedule a circuit execution protocol.
Args:
step (ExecutionStep): Execution step containing circuit and dependencies.
event_queue (EventQueue): Event queue to add the protocol to.
network (Network): The quantum network topology.
circuit_runner (QiskitRunner): Circuit runner for execution.
"""
execution_protocol = CircuitExecutionProtocol(
config=CircuitExecutionProtocolConfig(),
network=network,
circuit_runner=circuit_runner,
circuit=step.circuit,
)
event_queue.add_event_with_dependencies(
protocol=execution_protocol,
dependencies=list(step.dependencies),
event_id=step.event_id,
)
[docs]
def schedule_remote_operation( # noqa: PLR0913
step: ExecutionStep,
event_queue: EventQueue,
network: Network,
quantum_manager: QuantumRegisterManager,
eg_protocol_cls: type[EGProtocol],
eg_protocol_config: EGProtocolConfig,
last_remote_op_per_node_pair: dict[tuple[NodeId, NodeId], EventId],
node_comm_usage: dict[NodeId, dict[int, list[EventId]]],
) -> None:
"""Schedule entanglement generation and remote operation protocols.
Args:
step (ExecutionStep): Execution step containing remote operation info.
event_queue (EventQueue): Event queue to add protocols to.
network (Network): The quantum network topology.
quantum_manager (QuantumRegisterManager): Quantum register manager reference.
eg_protocol_cls (type[EGProtocol]): EG protocol class to instantiate.
eg_protocol_config (EGProtocolConfig): Configuration for EG protocol.
last_remote_op_per_node_pair (dict[tuple[NodeId, NodeId], EventId]): Tracks last remote op
between each node pair.
node_comm_usage (dict[NodeId, dict[int, list[EventId]]]): Tracks communication qubit usage
per node and gate_time.
"""
node_pair = (
step.remote_operation_info.source_node,
step.remote_operation_info.target_node,
)
gate_time = step.remote_operation_info.gate_time
eg_dependencies = compute_eg_dependencies(
node_pair, gate_time, last_remote_op_per_node_pair, node_comm_usage
)
eg_event_id = schedule_eg_protocol(
node_pair,
step.event_id,
eg_dependencies,
event_queue,
network,
quantum_manager,
eg_protocol_cls,
eg_protocol_config,
)
remote_op_event_id = schedule_remote_op_protocol(
step, eg_event_id, event_queue, network, quantum_manager
)
node_pair_key = get_ordered_node_pair_key(node_pair)
last_remote_op_per_node_pair[node_pair_key] = remote_op_event_id
for node in node_pair:
node_comm_usage[node.id][gate_time].append(eg_event_id)
[docs]
def compute_eg_dependencies(
node_pair: tuple[QuantumNode, QuantumNode],
gate_time: int,
last_remote_op_per_node_pair: dict[tuple[NodeId, NodeId], EventId],
node_comm_usage: dict[NodeId, dict[int, list[EventId]]],
) -> list[EventId]:
"""Compute dependencies for entanglement generation.
Args:
node_pair (tuple[NodeId, NodeId]): Source and target nodes.
gate_time (int): Gate-timestep in space-time matrix.
last_remote_op_per_node_pair (dict[tuple[NodeId, NodeId], EventId]): Last operation between
each node pair.
node_comm_usage (dict[NodeId, dict[int, list[EventId]]]): Current comm qubit usage.
Returns:
list[EventId]: Dependencies that must complete before EG can start.
"""
node_pair_key = get_ordered_node_pair_key(node_pair)
eg_dependencies = []
if node_pair_key in last_remote_op_per_node_pair:
eg_dependencies.append(last_remote_op_per_node_pair[node_pair_key])
for node in node_pair:
if node.id not in node_comm_usage:
node_comm_usage[node.id] = {}
if gate_time not in node_comm_usage[node.id]:
node_comm_usage[node.id][gate_time] = []
concurrent_ops = node_comm_usage[node.id][gate_time]
if len(concurrent_ops) >= node.communication_qubits:
eg_dependencies.append(concurrent_ops[0])
return eg_dependencies
[docs]
def schedule_eg_protocol( # noqa: PLR0913
node_pair: tuple[QuantumNode, QuantumNode],
step_event_id: EventId,
eg_dependencies: list[EventId],
event_queue: EventQueue,
network: Network,
quantum_manager: QuantumRegisterManager,
eg_protocol_cls: type[EGProtocol],
eg_protocol_config: EGProtocolConfig,
) -> EventId:
"""Schedule entanglement generation protocol.
Args:
node_pair (tuple[QuantumNode, QuantumNode]): Source and target nodes.
step_event_id (EventId): Event ID of the remote operation step.
eg_dependencies (list[EventId]): Dependencies for the EG protocol.
event_queue (EventQueue): Event queue to add protocol to.
network (Network): The quantum network topology.
quantum_manager (QuantumRegisterManager): Quantum register manager.
eg_protocol_cls (type[EGProtocol]): EG protocol class to instantiate.
eg_protocol_config (EGProtocolConfig): Configuration for EG protocol.
Returns:
EventId: Event ID of the scheduled EG protocol.
"""
eg_protocol = eg_protocol_cls(
config=eg_protocol_config,
network=network,
quantum_manager=quantum_manager,
source_node=node_pair[0],
target_node=node_pair[1],
)
return event_queue.add_event_with_dependencies(
protocol=eg_protocol,
dependencies=eg_dependencies,
event_id=EventId(f"eg_{node_pair[0].id}_{node_pair[1].id}_{step_event_id}"),
)
[docs]
def schedule_remote_op_protocol(
step: ExecutionStep,
eg_event_id: EventId,
event_queue: EventQueue,
network: Network,
quantum_manager: QuantumRegisterManager,
) -> EventId:
"""Schedule remote operation protocol (gate or teleportation).
Args:
step (ExecutionStep): Execution step containing remote operation info.
eg_event_id (EventId): Event ID of the EG protocol dependency.
event_queue (EventQueue): Event queue to add protocol to.
network (Network): The quantum network topology.
quantum_manager (QuantumRegisterManager): Quantum register manager.
Returns:
EventId: Event ID of the scheduled remote operation.
"""
if step.is_remote_gate:
config = RemoteOperationProtocolConfig(operation_type="remote_cnot")
elif step.is_teleportation:
config = RemoteOperationProtocolConfig(operation_type="teleportation")
else:
msg = "Neither remote gate nor teleportation. Something went wrong."
raise ValueError(msg)
remote_protocol = RemoteOperationProtocol(
config=config,
network=network,
quantum_manager=quantum_manager,
source_node=step.remote_operation_info.source_node,
target_node=step.remote_operation_info.target_node,
)
return event_queue.add_event_with_dependencies(
protocol=remote_protocol,
dependencies=[*list(step.dependencies), eg_event_id],
event_id=step.event_id,
)
[docs]
def get_ordered_node_pair_key(node_pair: tuple[QuantumNode, QuantumNode]) -> tuple[NodeId, NodeId]:
"""Ensure consistent ordering for node pair keys.
Args:
node_pair (tuple[QuantumNode, QuantumNode]): Tuple of source_node and target_node.
Returns:
tuple[NodeId, NodeId]: Ordered node IDs for use as dict key.
"""
if node_pair[0].id.value <= node_pair[1].id.value:
return node_pair[0].id, node_pair[1].id
return node_pair[1].id, node_pair[0].id