Pulse Compiler - Scheduling pass (#11981)

* Basic implementation

* Add align right, align sequential (sequence+schedule)

* Add draw, flatten

* Corrections

* Split into separate IR PR (temporary remove tests which rely on passes)

* Update qiskit/pulse/ir/ir.py

Co-authored-by: Naoki Kanazawa <nkanazawa1989@gmail.com>

* Corrections.

* Add to do.

* Fixes

* Disable lint

* MapMixedFrame + SetSequence passes

* Doc fixes

* Add SchedulePass, restore tests depending on scheduling.

* doc fix

* Corrections

* add todo

* Move schedule logic to pass.

* Corrections.

* blank line

* Update qiskit/pulse/compiler/passes/schedule.py

Co-authored-by: Naoki Kanazawa <nkanazawa1989@gmail.com>

* Corrections

* Minor fixes

---------

Co-authored-by: Naoki Kanazawa <nkanazawa1989@gmail.com>

qiskit/pulse/compiler/__init__.py

@@ -13,4 +13,4 @@
 """Pass-based Qiskit pulse program compiler."""
 
 from .passmanager import BlockTranspiler, BlockToIrCompiler
-from .passes import MapMixedFrame, SetSequence
+from .passes import MapMixedFrame, SetSequence, SetSchedule

qiskit/pulse/compiler/passes/__init__.py

@@ -14,3 +14,4 @@
 
 from .map_mixed_frames import MapMixedFrame
 from .set_sequence import SetSequence
+from .schedule import SetSchedule

qiskit/pulse/compiler/passes/schedule.py

@@ -0,0 +1,189 @@
+# This code is part of Qiskit.
+#
+# (C) Copyright IBM 2024.
+#
+# This code is licensed under the Apache License, Version 2.0. You may
+# obtain a copy of this license in the LICENSE.txt file in the root directory
+# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
+#
+# Any modifications or derivative works of this code must retain this
+# copyright notice, and modified files need to carry a notice indicating
+# that they have been altered from the originals.
+
+"""A scheduling pass for Qiskit PulseIR compilation."""
+
+from __future__ import annotations
+from functools import singledispatchmethod
+from collections import defaultdict
+from rustworkx import PyDAG, topological_sort, number_weakly_connected_components
+
+from qiskit.pulse.compiler.basepasses import TransformationPass
+from qiskit.pulse.ir import SequenceIR
+from qiskit.pulse.transforms import AlignmentKind, AlignLeft, AlignRight, AlignSequential
+from qiskit.pulse.exceptions import PulseCompilerError
+
+
+class SetSchedule(TransformationPass):
+    """Concretely schedule ``SequenceIR`` object.
+
+    The pass traverses the ``SequenceIR``, and recursively sets initial time for every
+    node in the sequence (and sub-sequences). The scheduling is done according to the
+    alignment strategy, and requires that the ``sequence`` property is already sequenced,
+    typically with the pass :class:`~qiskit.pulse.compiler.passes.SetSequence`."""
+
+    def __init__(self):
+        """Create new SetSchedule pass"""
+        super().__init__(target=None)
+
+    def run(
+        self,
+        passmanager_ir: SequenceIR,
+    ) -> SequenceIR:
+
+        self._schedule_recursion(passmanager_ir)
+        return passmanager_ir
+
+    def _schedule_recursion(self, prog: SequenceIR) -> None:
+        """Recursively schedule the IR.
+
+        Nested IR objects must be scheduled first, so we traverse the IR,
+        and recursively schedule the IR objects.
+        After all nested IR objects are scheduled, we apply the scheduling strategy to the
+        current object.
+
+        Arguments:
+            prog: The IR object to be scheduled.
+        """
+        for elem in prog.elements():
+            if isinstance(elem, SequenceIR):
+                self._schedule_recursion(elem)
+
+        self._schedule_single_program(prog.alignment, prog.sequence, prog.time_table)
+
+    @singledispatchmethod
+    def _schedule_single_program(
+        self, alignment: AlignmentKind, sequence: PyDAG, time_table: defaultdict
+    ) -> None:
+        """Concretely schedule the IR object.
+
+        The ``time_table`` argument is mutated to include the initial time of each element of
+        ``sequence``, according to the structure of ``sequence`` and the alignment.
+        The function assumes that nested IR objects are already scheduled.
+
+        ``sequence`` is assumed to have the following structure - node 0 marks the beginning of the
+        sequence, while node 1 marks the end of it. All branches of the graph originate from node 0
+        and end at node 1.
+
+        Arguments:
+            alignment: The alignment of the IR object.
+            sequence: The sequence of the IR object.
+            time_table: The time_table of the IR object.
+        """
+        raise NotImplementedError
+
+    # pylint: disable=unused-argument
+    @_schedule_single_program.register(AlignLeft)
+    @_schedule_single_program.register(AlignSequential)
+    def _schedule_asap(
+        self, alignment: AlignmentKind, sequence: PyDAG, time_table: defaultdict
+    ) -> None:
+        """Concretely schedule the IR object, aligning to the left.
+
+        The ``time_table`` argument is mutated to include the initial time of each element of
+        ``sequence``, according to the structure of ``sequence`` and aligning to the left.
+        The function assumes that nested IR objects are already scheduled.
+
+        ``sequence`` is assumed to have the following structure - node 0 marks the beginning of the
+        sequence, while node 1 marks the end of it. All branches of the graph originate from node 0
+        and end at node 1.
+
+        Arguments:
+            alignment: The alignment of the IR object.
+            sequence: The sequence of the IR object.
+            time_table: The time_table of the IR object.
+
+        Raises:
+            PulseCompilerError: If the sequence is not sequenced as expected.
+        """
+        nodes = topological_sort(sequence)
+        if number_weakly_connected_components(sequence) != 1 or nodes[0] != 0 or nodes[-1] != 1:
+            raise PulseCompilerError(
+                "The pulse program is not sequenced as expected. "
+                "Insert SetSequence pass in advance of SchedulePass."
+            )
+
+        for node_index in nodes:
+            if node_index in (0, 1):
+                # in,out nodes
+                continue
+            preds = sequence.predecessor_indices(node_index)
+            if preds == [0]:
+                time_table[node_index] = 0
+            else:
+                time_table[node_index] = max(
+                    time_table[pred] + sequence.get_node_data(pred).duration for pred in preds
+                )
+
+    # pylint: disable=unused-argument
+    @_schedule_single_program.register(AlignRight)
+    def _schedule_alap(
+        self, alignment: AlignmentKind, sequence: PyDAG, time_table: defaultdict
+    ) -> None:
+        """Concretely schedule the IR object, aligning to the right.
+
+        The ``time_table`` argument is mutated to include the initial time of each element of
+        ``sequence``, according to the structure of ``sequence`` and aligning to the right.
+        The function assumes that nested IR objects are already scheduled.
+
+        ``sequence`` is assumed to have the following structure - node 0 marks the beginning of the
+        sequence, while node 1 marks the end of it. All branches of the graph originate from node 0
+        and end at node 1.
+
+        Arguments:
+            alignment: The alignment of the IR object.
+            sequence: The sequence of the IR object.
+            time_table: The time_table of the IR object.
+
+        Raises:
+            PulseCompilerError: If the sequence is not sequenced as expected.
+        """
+        # We reverse the sequence, schedule to the left, then reverse the timings.
+        reversed_sequence = sequence.copy()
+        reversed_sequence.reverse()
+
+        nodes = topological_sort(reversed_sequence)
+
+        if number_weakly_connected_components(sequence) != 1 or nodes[0] != 1 or nodes[-1] != 0:
+            raise PulseCompilerError(
+                "The pulse program is not sequenced as expected. "
+                "Insert SetSequence pass in advance of SchedulePass."
+            )
+
+        for node_index in nodes:
+            if node_index in (0, 1):
+                # in,out nodes
+                continue
+            preds = reversed_sequence.predecessor_indices(node_index)
+            if preds == [1]:
+                time_table[node_index] = 0
+            else:
+                time_table[node_index] = max(
+                    time_table[pred] + sequence.get_node_data(pred).duration for pred in preds
+                )
+
+        total_duration = max(
+            time_table[i] + sequence.get_node_data(i).duration
+            for i in reversed_sequence.predecessor_indices(0)
+        )
+
+        for node in sequence.node_indices():
+            if node not in (0, 1):
+                time_table[node] = (
+                    total_duration - time_table[node] - sequence.get_node_data(node).duration
+                )
+
+    def __hash__(self):
+        return hash((self.__class__.__name__,))
+
+    def __eq__(self, other):
+        return self.__class__.__name__ == other.__class__.__name__

qiskit/pulse/ir/ir.py

@@ -77,6 +77,11 @@ def sequence(self) -> PyDAG:
         """Return the DAG sequence of the SequenceIR"""
         return self._sequence
 
+    @property
+    def time_table(self) -> defaultdict:
+        """Return the timetable of the SequenceIR"""
+        return self._time_table
+
     def append(self, element: SequenceIR | Instruction) -> int:
         """Append element to the SequenceIR
 
@@ -217,11 +222,19 @@ def _draw_nodes(n):
     def flatten(self, inplace: bool = False) -> SequenceIR:
         """Recursively flatten the SequenceIR.
 
+        The flattening process includes breaking up nested IRs until only instructions remain.
+        The flattened object will contain all instructions, timing information, and the
+        complete sequence graph. However, the alignment of nested IRs will be lost. Because of
+        this, flattening an unscheduled IR is not allowed.
+
         Args:
             inplace: If ``True`` flatten the object itself. If ``False`` return a flattened copy.
 
         Returns:
             A flattened ``SequenceIR`` object.
+
+        Raises:
+            PulseError: If the IR (or nested IRs) are not scheduled.
         """
         # TODO : Verify that the block\sub blocks are sequenced correctly.
         if inplace:
@@ -230,7 +243,8 @@ def flatten(self, inplace: bool = False) -> SequenceIR:
             block = copy.deepcopy(self)
             block._sequence[0] = SequenceIR._InNode
             block._sequence[1] = SequenceIR._OutNode
-
+        # TODO : Consider replacing the alignment to "NullAlignment", as the original alignment
+        #  has no meaning.
         # TODO : Create a dedicated half shallow copier.
 
         def edge_map(_x, _y, _node):
@@ -240,18 +254,23 @@ def edge_map(_x, _y, _node):
                 return 1
             return None
 
+        if any(
+            block.time_table[x] is None for x in block.sequence.node_indices() if x not in (0, 1)
+        ):
+            raise PulseError("Can not flatten unscheduled IR")
+
         for ind in block.sequence.node_indices():
             if isinstance(sub_block := block.sequence.get_node_data(ind), SequenceIR):
                 sub_block.flatten(inplace=True)
-                initial_time = block._time_table[ind]
+                initial_time = block.time_table[ind]
                 nodes_mapping = block._sequence.substitute_node_with_subgraph(
                     ind, sub_block.sequence, lambda x, y, _: edge_map(x, y, ind)
                 )
                 if initial_time is not None:
                     for old_node in nodes_mapping.keys():
                         if old_node not in (0, 1):
                             block._time_table[nodes_mapping[old_node]] = (
-                                initial_time + sub_block._time_table[old_node]
+                                initial_time + sub_block.time_table[old_node]
                             )
 
                 del block._time_table[ind]