Add support for pulse reference to QPY (Qiskit#9890)

* Add support pulse reference to QPY

* Review comments

Co-authored-by: Matthew Treinish <mtreinish@kortar.org>

---------

Co-authored-by: Matthew Treinish <mtreinish@kortar.org>

qiskit/qpy/__init__.py

@@ -126,6 +126,46 @@
 by ``num_circuits`` in the file header). There is no padding between the
 circuits in the data.
 
+.. _qpy_version_7:
+
+Version 7
+=========
+
+Version 7 adds support for :class:`.~Reference` instruction and serialization of
+a :class:`.~ScheduleBlock` program while keeping its reference to subroutines::
+
+    from qiskit import pulse
+    from qiskit import qpy
+
+    with pulse.build() as schedule:
+        pulse.reference("cr45p", "q0", "q1")
+        pulse.reference("x", "q0")
+        pulse.reference("cr45p", "q0", "q1")
+
+    with open('template_ecr.qpy', 'wb') as fd:
+        qpy.dump(schedule, fd)
+
+The conventional :ref:`qpy_schedule_block` data model is preserved, but in
+version 7 it is immediately followed by an extra :ref:`qpy_mapping` utf8 bytes block
+representing the data of the referenced subroutines.
+
+New type key character is added to the :ref:`qpy_schedule_instructions` group
+for the :class:`.~Reference` instruction.
+
+- ``y``: :class:`~qiskit.pulse.instructions.Reference` instruction
+
+New type key character is added to the :ref:`qpy_schedule_operands` group
+for the operands of :class:`.~Reference` instruction,
+which is a tuple of strings, e.g. ("cr45p", "q0", "q1").
+
+- ``o``: string (operand string)
+
+Note that this is the same encoding with the built-in Python string, however,
+the standard value encoding in QPY uses ``s`` type character for string data,
+which conflicts with the :class:`~qiskit.pulse.library.SymbolicPulse` in the scope of
+pulse instruction operands. A special type character ``o`` is reserved for
+the string data that appears in the pulse instruction operands.
+
 .. _qpy_version_6:
 
 Version 6
@@ -213,6 +253,8 @@
 the same QPY interface. Input data type is implicitly analyzed and
 no extra option is required to save the schedule block.
 
+.. _qpy_schedule_block_header:
+
 SCHEDULE_BLOCK_HEADER
 ---------------------
 
@@ -230,6 +272,11 @@
 ``metadata_size`` utf8 bytes of the JSON serialized metadata dictionary
 attached to the schedule.
 
+.. _qpy_schedule_alignments:
+
+SCHEDULE_BLOCK_ALIGNMENTS
+-------------------------
+
 Then, alignment context of the schedule block starts with ``char``
 representing the supported context type followed by the :ref:`qpy_sequence` block representing
 the parameters associated with the alignment context :attr:`AlignmentKind._context_params`.
@@ -243,6 +290,11 @@
 Note that :class:`~.AlignFunc` context is not supported becasue of the callback function
 stored in the context parameters.
 
+.. _qpy_schedule_instructions:
+
+SCHEDULE_BLOCK_INSTRUCTIONS
+---------------------------
+
 This alignment block is further followed by ``num_element`` length of block elements which may
 consist of nested schedule blocks and schedule instructions.
 Each schedule instruction starts with ``char`` representing the instruction type
@@ -261,6 +313,12 @@
 - ``r``: :class:`~qiskit.pulse.instructions.ShiftPhase` instruction
 - ``b``: :class:`~qiskit.pulse.instructions.RelativeBarrier` instruction
 - ``t``: :class:`~qiskit.pulse.instructions.TimeBlockade` instruction
+- ``y``: :class:`~qiskit.pulse.instructions.Reference` instruction (new in version 0.7)
+
+.. _qpy_schedule_operands:
+
+SCHEDULE_BLOCK_OPERANDS
+-----------------------
 
 The operands of these instances can be serialized through the standard QPY value serialization
 mechanism, however there are special object types that only appear in the schedule operands.
@@ -272,6 +330,7 @@
 - ``c``: :class:`~qiskit.pulse.channels.Channel`
 - ``w``: :class:`~qiskit.pulse.library.Waveform`
 - ``s``: :class:`~qiskit.pulse.library.SymbolicPulse`
+- ``o``: string (operand string, new in version 0.7)
 
 .. _qpy_schedule_channel:
 

qiskit/qpy/binary_io/schedules.py

@@ -19,10 +19,11 @@
 import numpy as np
 
 from qiskit.exceptions import QiskitError
-from qiskit.pulse import library, channels
+from qiskit.pulse import library, channels, instructions
 from qiskit.pulse.schedule import ScheduleBlock
 from qiskit.qpy import formats, common, type_keys
 from qiskit.qpy.binary_io import value
+from qiskit.qpy.exceptions import QpyError
 from qiskit.utils import optionals as _optional
 
 if _optional.HAS_SYMENGINE:
@@ -238,6 +239,8 @@ def _loads_operand(type_key, data_bytes, version):
             return common.data_from_binary(data_bytes, _read_symbolic_pulse_v6, version=version)
     if type_key == type_keys.ScheduleOperand.CHANNEL:
         return common.data_from_binary(data_bytes, _read_channel, version=version)
+    if type_key == type_keys.ScheduleOperand.OPERAND_STR:
+        return data_bytes.decode(common.ENCODE)
 
     return value.loads_value(type_key, data_bytes, version, {})
 
@@ -259,6 +262,24 @@ def _read_element(file_obj, version, metadata_deserializer):
     return instance
 
 
+def _loads_reference_item(type_key, data_bytes, version, metadata_deserializer):
+    if type_key == type_keys.Value.NULL:
+        return None
+    if type_key == type_keys.Program.SCHEDULE_BLOCK:
+        return common.data_from_binary(
+            data_bytes,
+            deserializer=read_schedule_block,
+            version=version,
+            metadata_deserializer=metadata_deserializer,
+        )
+
+    raise QpyError(
+        f"Loaded schedule reference item is neither None nor ScheduleBlock. "
+        f"Type key {type_key} is not valid data type for a reference items. "
+        "This data cannot be loaded. Please check QPY version."
+    )
+
+
 def _write_channel(file_obj, data):
     type_key = type_keys.ScheduleChannel.assign(data)
     common.write_type_key(file_obj, type_key)
@@ -340,6 +361,9 @@ def _dumps_operand(operand):
     elif isinstance(operand, channels.Channel):
         type_key = type_keys.ScheduleOperand.CHANNEL
         data_bytes = common.data_to_binary(operand, _write_channel)
+    elif isinstance(operand, str):
+        type_key = type_keys.ScheduleOperand.OPERAND_STR
+        data_bytes = operand.encode(common.ENCODE)
     else:
         type_key, data_bytes = value.dumps_value(operand)
 
@@ -361,6 +385,20 @@ def _write_element(file_obj, element, metadata_serializer):
         value.write_value(file_obj, element.name)
 
 
+def _dumps_reference_item(schedule, metadata_serializer):
+    if schedule is None:
+        type_key = type_keys.Value.NULL
+        data_bytes = b""
+    else:
+        type_key = type_keys.Program.SCHEDULE_BLOCK
+        data_bytes = common.data_to_binary(
+            obj=schedule,
+            serializer=write_schedule_block,
+            metadata_serializer=metadata_serializer,
+        )
+    return type_key, data_bytes
+
+
 def read_schedule_block(file_obj, version, metadata_deserializer=None):
     """Read a single ScheduleBlock from the file like object.
 
@@ -382,7 +420,6 @@ def read_schedule_block(file_obj, version, metadata_deserializer=None):
         TypeError: If any of the instructions is invalid data format.
         QiskitError: QPY version is earlier than block support.
     """
-
     if version < 5:
         QiskitError(f"QPY version {version} does not support ScheduleBlock.")
 
@@ -406,6 +443,22 @@ def read_schedule_block(file_obj, version, metadata_deserializer=None):
         block_elm = _read_element(file_obj, version, metadata_deserializer)
         block.append(block_elm, inplace=True)
 
+    # Load references
+    if version >= 7:
+        flat_key_refdict = common.read_mapping(
+            file_obj=file_obj,
+            deserializer=_loads_reference_item,
+            version=version,
+            metadata_deserializer=metadata_deserializer,
+        )
+        ref_dict = {}
+        for key_str, schedule in flat_key_refdict.items():
+            if schedule is not None:
+                composite_key = tuple(key_str.split(instructions.Reference.key_delimiter))
+                ref_dict[composite_key] = schedule
+        if ref_dict:
+            block.assign_references(ref_dict, inplace=True)
+
     return block
 
 
@@ -440,5 +493,20 @@ def write_schedule_block(file_obj, block, metadata_serializer=None):
     file_obj.write(metadata)
 
     _write_alignment_context(file_obj, block.alignment_context)
-    for block_elm in block.blocks:
+    for block_elm in block._blocks:
+        # Do not call block.blocks. This implicitly assigns references to instruction.
+        # This breaks original reference structure.
         _write_element(file_obj, block_elm, metadata_serializer)
+
+    # Write references
+    flat_key_refdict = {}
+    for ref_keys, schedule in block._reference_manager.items():
+        # Do not call block.reference. This returns the reference of most outer program by design.
+        key_str = instructions.Reference.key_delimiter.join(ref_keys)
+        flat_key_refdict[key_str] = schedule
+    common.write_mapping(
+        file_obj=file_obj,
+        mapping=flat_key_refdict,
+        serializer=_dumps_reference_item,
+        metadata_serializer=metadata_serializer,
+    )

qiskit/qpy/common.py

@@ -20,7 +20,7 @@
 
 from qiskit.qpy import formats
 
-QPY_VERSION = 6
+QPY_VERSION = 7
 ENCODE = "utf8"
 
 

qiskit/qpy/type_keys.py

@@ -45,6 +45,7 @@
     ShiftPhase,
     RelativeBarrier,
     TimeBlockade,
+    Reference,
 )
 from qiskit.pulse.library import Waveform, SymbolicPulse
 from qiskit.pulse.schedule import ScheduleBlock
@@ -233,6 +234,7 @@ class ScheduleInstruction(TypeKeyBase):
     SHIFT_PHASE = b"r"
     BARRIER = b"b"
     TIME_BLOCKADE = b"t"
+    REFERENCE = b"y"
 
     # 's' is reserved by ScheduleBlock, i.e. block can be nested as an element.
     # Call instructon is not supported by QPY.
@@ -261,6 +263,8 @@ def assign(cls, obj):
             return cls.BARRIER
         if isinstance(obj, TimeBlockade):
             return cls.TIME_BLOCKADE
+        if isinstance(obj, Reference):
+            return cls.REFERENCE
 
         raise exceptions.QpyError(
             f"Object type '{type(obj)}' is not supported in {cls.__name__} namespace."
@@ -286,6 +290,8 @@ def retrieve(cls, type_key):
             return RelativeBarrier
         if type_key == cls.TIME_BLOCKADE:
             return TimeBlockade
+        if type_key == cls.REFERENCE:
+            return Reference
 
         raise exceptions.QpyError(
             f"A class corresponding to type key '{type_key}' is not found in {cls.__name__} namespace."
@@ -303,6 +309,12 @@ class ScheduleOperand(TypeKeyBase):
     # Data format of these object is somewhat opaque and not defiend well.
     # It's rarely used in the Qiskit experiements. Of course these can be added later.
 
+    # We need to have own string type definition for operands of schedule instruction.
+    # Note that string type is already defined in the Value namespace,
+    # but its key "s" conflicts with the SYMBOLIC_PULSE in the ScheduleOperand namespace.
+    # New in QPY version 7.
+    OPERAND_STR = b"o"
+
     @classmethod
     def assign(cls, obj):
         if isinstance(obj, Waveform):
@@ -311,6 +323,8 @@ def assign(cls, obj):
             return cls.SYMBOLIC_PULSE
         if isinstance(obj, Channel):
             return cls.CHANNEL
+        if isinstance(obj, str):
+            return cls.OPERAND_STR
 
         raise exceptions.QpyError(
             f"Object type '{type(obj)}' is not supported in {cls.__name__} namespace."

releasenotes/notes/add-support-for-qpy-reference-70478baa529fff8c.yaml

@@ -0,0 +1,19 @@
+---
+features:
+  - |
+    QPY supports pulse :class:`~.ScheduleBlock` with unassigned reference,
+    and preserves the data structure for the reference to subroutines.
+    This feature allows to save a template pulse program for tasks such as pulse calibration.
+
+    .. code-block:: python
+
+      from qiskit import pulse
+      from qiskit import qpy
+
+      with pulse.build() as schedule:
+          pulse.reference("cr45p", "q0", "q1")
+          pulse.reference("x", "q0")
+          pulse.reference("cr45p", "q0", "q1")
+
+      with open('template_ecr.qpy', 'wb') as fd:
+          qpy.dump(schedule, fd)

test/python/qpy/test_block_load_from_qpy.py

@@ -207,6 +207,50 @@ def test_called_schedule(self):
             builder.call(refsched, name="test_ref")
         self.assert_roundtrip_equal(test_sched)
 
+    def test_unassigned_reference(self):
+        """Test schedule with unassigned reference."""
+        with builder.build() as test_sched:
+            builder.reference("custom1", "q0")
+            builder.reference("custom1", "q1")
+
+        self.assert_roundtrip_equal(test_sched)
+
+    def test_partly_assigned_reference(self):
+        """Test schedule with partly assigned reference."""
+        with builder.build() as test_sched:
+            builder.reference("custom1", "q0")
+            builder.reference("custom1", "q1")
+
+        with builder.build() as sub_q0:
+            builder.delay(Parameter("duration"), DriveChannel(0))
+
+        test_sched.assign_references(
+            {("custom1", "q0"): sub_q0},
+            inplace=True,
+        )
+
+        self.assert_roundtrip_equal(test_sched)
+
+    def test_nested_assigned_reference(self):
+        """Test schedule with assigned reference for nested schedule."""
+        with builder.build() as test_sched:
+            with builder.align_left():
+                builder.reference("custom1", "q0")
+            builder.reference("custom1", "q1")
+
+        with builder.build() as sub_q0:
+            builder.delay(Parameter("duration"), DriveChannel(0))
+
+        with builder.build() as sub_q1:
+            builder.delay(Parameter("duration"), DriveChannel(1))
+
+        test_sched.assign_references(
+            {("custom1", "q0"): sub_q0, ("custom1", "q1"): sub_q1},
+            inplace=True,
+        )
+
+        self.assert_roundtrip_equal(test_sched)
+
     def test_bell_schedule(self):
         """Test complex schedule to create a Bell state."""
         with builder.build() as test_sched: