qiskit icon indicating copy to clipboard operation
qiskit copied to clipboard

DAG based IR

Open TsafrirA opened this issue 1 year ago • 0 comments

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()

image

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()

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

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

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()

image 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()

image 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)

image 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.

TsafrirA avatar Feb 20 '24 04:02 TsafrirA