DAG based Pulse IR - Demo #11839

TsafrirA · 2024-02-20T04:57:53Z

Summary

This is an initial example for a DAG based PulseIR as an alternative to the option added in #11767.
This is based on a code example by @nkanazawa1989.

Details and comments

In #11767 we used a composite pattern and a list based tracking of the elements in the IR.
Here, we use instead a DAG representation. When initialized, the IR is comprised of graph of only nodes, and the edges are later added as part of a sequencing pass. This makes scheduling simple, as demonstrated by the scheduling pass implemented here.

As #11743 is still pending, the passes implemented here are kept in temporary files until the compiler will be added, and they could be sorted into their final homes.

Demo

A lot of the functionalities we need are already in place.
Quick set up:

from qiskit.pulse import Play, Qubit, QubitFrame, Constant, ShiftPhase
from qiskit.pulse.ir import IrBlock
from qiskit.pulse.ir.alignments import AlignRight, AlignSequential, AlignLeft
from qiskit.pulse.compiler.temp_passes import analyze_target_frame_pass, sequence_pass, schedule_pass
from matplotlib import pyplot as plt

And we can dive into a simple example with 3 play pulses - two of them sharing a mixed frame.

ir_example = IrBlock(AlignLeft())
ir_example.append(Play(Constant(100, 0.5), frame=QubitFrame(1), target=Qubit(1), name="q1qf1"))
ir_example.append(Play(Constant(100, 0.5), frame=QubitFrame(1), target=Qubit(1), name="q1qf1"))
ir_example.append(Play(Constant(100, 0.5), frame=QubitFrame(1), target=Qubit(2), name="q2qf1"))

ir_example.draw()

Initially, the IR is just a collection of nodes. To understand the dependencies, we first need to find all MixedFrames associated with each PulseTarget and Frame. We can use an analysis pass for that: (currently implemented as function, later we'll hook it into the compiler as a pass).

property_set = {}
analyze_target_frame_pass(ir_example, property_set)

print(property_set["target_frame_map"])

Examining property_set["target_frame_map"] we'll see that QubitFrame(1) is associated with two mixed frames, while all other objects are associated with just one of the two existing in the IR.

We can use this to sequence the instructions according to the alignment we choose.

sequence_pass(ir_example, property_set)
ir_example.draw()

Note how the instructions acting on the same mixed frame "block" each other.

If we repeat this procedure with AlignSequential() we'll get instead

The previous example didn't require any mapping or broadcasting. Let's see one with broadcasting:

ir_example = IrBlock(AlignLeft())
ir_example.append(ShiftPhase(0.1, frame=QubitFrame(1), name="qf1"))
ir_example.append(Play(Constant(100, 0.5), frame=QubitFrame(1), target=Qubit(1), name="q1qf1"))
ir_example.append(Play(Constant(100, 0.5), frame=QubitFrame(1), target=Qubit(1), name="q1qf1"))
ir_example.append(Play(Constant(100, 0.5), frame=QubitFrame(1), target=Qubit(2), name="q2qf1"))

property_set = {}
analyze_target_frame_pass(ir_example, property_set)
sequence_pass(ir_example, property_set)

ir_example.draw()

Note how the shift phase instruction is broadcasted to all mixed frames, thus "blocking" all of them.

Now we can take this and schedule it.

schedule_pass(ir_example, property_set)
print(ir_example.scheduled_elements())

[
[0, ShiftPhase(..., name='qf1')],
[0, Play(..., name='q1qf1')],
[100, Play(..., name='q1qf1')],
[0, Play(..., name='q2qf1')]
]

Note how q1qf2 is pushed to the left because it's not "blocked" by the other instructions. If we swapped this for right alignment we'll have that instruction pushed to the right:

[
[0, ShiftPhase(..., name='qf1')],
[0, Play(..., name='q1qf1')],
[100, Play(..., name='q1qf1')],
[*100*, Play(..., name='q2qf1')]
]

And we can do the same thing with sub blocks.

block = IrBlock(AlignLeft())
block.append(Play(Constant(100, 0.5), frame=QubitFrame(1), target=Qubit(1), name="q1qf1"))
block.append(Play(Constant(100, 0.5), frame=QubitFrame(2), target=Qubit(2), name="q2qf2"))

ir_example = IrBlock(AlignLeft())
ir_example.append(block)
ir_example.append(Play(Constant(100, 0.5), frame=QubitFrame(1), target=Qubit(1), name="q1qf1"))
ir_example.append(Play(Constant(100, 0.5), frame=QubitFrame(2), target=Qubit(2), name="q2qf2"))
ir_example.append(Play(Constant(100, 0.5), frame=QubitFrame(3), target=Qubit(3), name="q3qf3"))

property_set = {}
analyze_target_frame_pass(ir_example, property_set)
sequence_pass(ir_example, property_set)

ir_example.draw()

Note, how the sub block "blocks" all instructions which correspond to any of the mixed frames in it (but not others).

If we flatten the representation we get:

ir_example.draw(recursive=True)

Note how the instructions within the original sub-block now "block" the following instructions by themselves.

The PR is not in shape to be merged as is, and several issues still need addressing:

Validation of inputs, graphs etc.
Old model inputs - validation, conversion
Edge cases (instructions on one of Frame or PulseTarget with no associated MixedFrame will cause an error currently)
Alignment classes - how do they stack against the existing AlignmentKind?
Documentation
Should timing data be stored in the IR object (as done now) or in each node?

The PR does, however, provide a solid basis for discussion.

nkanazawa1989

I just checked IR update part. This looks good to me. Maybe you can split the PR into two parts so that I can review the sequence pass logic separately.

nkanazawa1989 · 2024-02-22T03:58:02Z

qiskit/pulse/ir/ir.py

            return None
-        else:
-            return min(elements_initial_times, default=None)
+        return min([self._time_table.get(node, None) for node in first_nodes], default=None)


I think you need try-except for TypeError because comparison operators don't support None.

default=None takes care of that.

nkanazawa1989 · 2024-02-22T04:03:18Z

qiskit/pulse/ir/ir.py

-            self.initial_time = initial_time
+        self.time_offset = 0
+        self._time_table = {}
+        self._children = []


Do you want to implement has_children method? Can be implemented later as needed.

I actually never used the _children attribute... So I am not sure it should even exist.

This might be convenient when we need to implement something like flattening pass, but you can remove for now.

nkanazawa1989 · 2024-02-22T04:42:22Z

qiskit/pulse/ir/ir.py

-
-class IrInstruction(IrElement):
-    """Pulse IR Instruction
+class IrBlock:


Can we give this a different name? I guess the original intention was "Intermediate representation of ScheduleBlock", and it makes me think this is still a collection of commands, i.e. code block + alignment. Since main purpose of this representation is sequencing, I prefer something like SequenceIR considering the lowering of the program; context alignment (lexical) -> sequence (ordered) -> schedule (t0 assigned).

I don't have a strong preference here.

nkanazawa1989 · 2024-02-22T04:50:22Z

qiskit/pulse/ir/alignments.py

+# that they have been altered from the originals.
+
+
+class Alignment:


This should be an abstract base class, right?

Yes. I didn't want to put too much work into this until we decide how to deal with the two alignments hierarchies.

nkanazawa1989 · 2024-02-22T05:02:03Z

qiskit/pulse/ir/ir.py

-        else:
-            self.initial_time = initial_time
+        self.time_offset = 0
+        self._time_table = {}


Suggested change

self._time_table = {}

self._time_table = defaultdict(lambda: None)

Maybe? Then you don't need to write self._time_table.get(index, None) everywhere.

Yeah I was thinking the same thing.

nkanazawa1989 · 2024-02-22T05:12:47Z

qiskit/pulse/ir/ir.py

+        return [
+            [self._time_table.get(ni, None), self._sequence.get_node_data(ni)]
+            for ni in self._sequence.node_indices()
+            if ni not in (0, 1)
+        ]


Suggested change

return [

[self._time_table.get(ni, None), self._sequence.get_node_data(ni)]

for ni in self._sequence.node_indices()

if ni not in (0, 1)

]

def _to_time_inst_tuple(n_index):

instruction = self._sequence.get_node_data(ni)

if t0 := self._time_table.get(n_index, None) is not None:

return t0 + self.time_offset, instruction

return None, instruction

return list(

map(

_to_time_inst_tuple,

filter(lambda i: i not in (0, 1), self._sequence.node_indices()),

)

)

You need to consider the time offset. I prefer tuple form for time instruction duo.

I opted not to do that, because it causes duplication where the initial time of sub blocks is stored in two places - the time table of the parent block, and the sub block itself. In general, I think it's better to have the initial time directly in the nodes. Eventually, this is where you want this the data to live - what is the instruction, and what time it should be applied.

(In the current implementation I took care of the offset in the flatten method. I left the sub-blocks unaware of their relative timing until that point).

the initial time of sub blocks is stored in two places - the time table of the parent block, and the sub block itself.

Hmm it depends how we calculate initial time. I assume an outer block just calls sub block's schedule_elements and extend, but if you assume sub block's element time t0 = outer block t0 + sub block element t0, then current logic looks reasonable.

(In the current implementation I took care of the offset in the flatten method. I left the sub-blocks unaware of their relative timing until that point).

I prefer this schedule_elements returns elements in flattened fashion, since lower level representation may not have notion of context and I don't think of any use case requiring time information but context must be preserved. You can also add flatten argument to this method (I prefer it defaults to True).

I think it's better to have the initial time directly in the nodes

I'm against this design, based on the following rule:

Same command (i.e. Instruction tied to channel) should be equivalent object.

When you attach time information directly to the command, for example, two play instructions issued at different time must become different objects. To avoid this, we need a wrapper object, e.g. scheduled command, and adopt composite pattern as you first implemented. However, this results in complex class dependency and class instantiation overhead as the number of instruction increases.

nkanazawa1989 · 2024-03-18T01:46:26Z

Close this PR in favor of #11980 and #11981

Basic implementation

7cd1396

nkanazawa1989 self-assigned this Feb 20, 2024

nkanazawa1989 added mod: pulse Related to the Pulse module experimental Experimental feature without API stability guarantee labels Feb 20, 2024

nkanazawa1989 linked an issue Feb 20, 2024 that may be closed by this pull request

Pulse Compiler and IR #10759

Open

9 tasks

nkanazawa1989 mentioned this pull request Feb 20, 2024

Pulse Compiler and IR #10759

Open

9 tasks

nkanazawa1989 changed the title ~~[pulse-branch] DAG based IR~~ DAG based IR Feb 20, 2024

Add align right, align sequential (sequence+schedule)

a35abed

nkanazawa1989 reviewed Feb 22, 2024

View reviewed changes

TsafrirA added 2 commits March 4, 2024 13:51

Add draw, flatten

592aba6

Corrections

8c83957

TsafrirA mentioned this pull request Mar 6, 2024

DAG Based Pulse IR #11951

Merged

TsafrirA changed the title ~~DAG based IR~~ DAG based Pulse IR - Demo Mar 6, 2024

nkanazawa1989 closed this Mar 18, 2024

TsafrirA deleted the DagBasedIR branch March 18, 2024 21:31

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

DAG based Pulse IR - Demo #11839

DAG based Pulse IR - Demo #11839

TsafrirA commented Feb 20, 2024 •

edited

Loading

nkanazawa1989 left a comment

nkanazawa1989 Feb 22, 2024

TsafrirA Feb 22, 2024

nkanazawa1989 Feb 22, 2024

TsafrirA Feb 22, 2024

nkanazawa1989 Feb 28, 2024

nkanazawa1989 Feb 22, 2024

TsafrirA Feb 22, 2024

nkanazawa1989 Feb 22, 2024

TsafrirA Feb 22, 2024

nkanazawa1989 Feb 22, 2024

TsafrirA Feb 22, 2024

nkanazawa1989 Feb 22, 2024

TsafrirA Feb 22, 2024

nkanazawa1989 Feb 28, 2024

nkanazawa1989 commented Mar 18, 2024

		# that they have been altered from the originals.


		class Alignment:

	self._time_table = {}
	self._time_table = defaultdict(lambda: None)

-        return [
-            [self._time_table.get(ni, None), self._sequence.get_node_data(ni)]
-            for ni in self._sequence.node_indices()
-            if ni not in (0, 1)
-        ]
+        def _to_time_inst_tuple(n_index):
+            instruction = self._sequence.get_node_data(ni)
+            if t0 := self._time_table.get(n_index, None) is not None:
+                return t0 + self.time_offset, instruction
+            return None, instruction
+        return list(
+            map(
+                _to_time_inst_tuple,
+                filter(lambda i: i not in (0, 1), self._sequence.node_indices()),
+            )
+        )

DAG based Pulse IR - Demo #11839

DAG based Pulse IR - Demo #11839

Conversation

TsafrirA commented Feb 20, 2024 • edited Loading

Summary

Details and comments

Demo

nkanazawa1989 left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

nkanazawa1989 commented Mar 18, 2024

TsafrirA commented Feb 20, 2024 •

edited

Loading