Using MLIR to optimize Quantum Programs

I’ve been able to use the pattern-match-rewrite framework to perform some local and basic optimizations, but I need to do a few global analyses and optimizations. I have following questions wrt MLIR:

  1. Is it possible to create a DAG of quantum program where operations are nodes and edges are qubits?
  2. If yes, then is it possible to create a pass where a piece of information is propagated throughout the nodes of the graph?

Could someone point me to the places in the MLIR source that could be helpful for the above objectives?

Thank you.

There are several ways to encode any DAG in MLIR. The most common one is to leverage SSA use-def chains to represent the edges of a graph, and operations to represent nodes.

A -> B -> C
     ---> D

Can become something like

%a = node() : !type.t
%b = node(%a) : !type.t
%c = node(%b) : !type.t
%d = node(%b) : !type.t

with the definition of an SSA value being the source of the edge and its uses being the sinks. Use-def chains (within a block) are acyclic. It is only the matter of giving proper semantics to operations and proper types to the values they produce.

Other ways to encode DAGs can be based on using CFG through custom terminators or the symbol-based mechanism through named references, but neither of those guarantees acyclicity by construction.

I am not aware of any work on using MLIR for modeling quantum computations, so if you could elaborate how do you represent them, it might be interesting and helpful.
One encoding I can see is to represent quantum gates as operations that use and produce values of “quantum state”. The entire structure can be put into a region of an operation with extra verifiers that could ensure, for example, that each value is used at most once to, e.g., avoid pretending a quantum state is cloned. Then one can use something like register assignment algorithms to assign “quantum state” values to actual qubits.

I’ve hit a roadblock while working with the SSA form. It stems from the fact that operations in Quantum are reversible and therefore, there’s no single output from operations but instead the state of the qubits changes (or system of qubits if they are entangled).

Here’s an example to aid discussion:

Let’s consider a backend with a universal gate set: <Single-qubit gates, CNOT>
H(q1); CNOT(q1, q2); H(q1);
Equivalent Circuit:

In some backends, there is a native equivalent gate operation:

which results in following code:
CZ(q1, q2);

If I visualize the graph of the circuit as follows, then each operation would be a node and edges would indicate the qubits:

So, can MLIR be used for this:

  1. Representation?
  2. Optimization using DAG passes?

[Note: circuits compiled using qiskit and DAG concept is already implemented in Qiskit SDK]

That’s why my suggested encoding was dealing with “quantum state” not qubits. You can see each value as state of the i-th qubit at time t_j (discretizing over time), or as state of the i-th qubit after j gates applied to it. For example:

quantum.circuit inputs(%0: !quantum.state) -> !quantum.state {
  %1 = gate.hadamard (%0) : !quantum.state -> !quantum.state
  quantum.outputs (%1)

could “take” |0> as %0 and “assign” (|0> + |1>) / sqrt(2) to %1. You can trace the fact that %0 and %1 are on the same line in the circuit by following use-def chains of gates (uses = gate inputs, defs = gate outputs).

The circuit you describe would be then modeled as

quantum.circuit inputs(%0: !quantum.state, %1: !quantum.state) -> (!quantum.state, !quantum.state) {
  %2 = gate.H (%1) : !quantum.state -> !quantum.state
  %3:2 = gate.CNOT(%0, %2) : (!quantum.state, !quantum.state) -> (!quantum.state, !quantum.state)
  %4 = gate.H (%3#1) : !quantum.state -> !quantum.state
  quantum.output %3#0, %4 : !quantum.state, !quantum.state

and can be rewritten by a pattern. It may be awkward, if at all possible, to express in DRR though (@antiagainst to the rescue; I need to make sure two different ops use the results of the same multi-result op). You will also need some logic to trace use-def chains to recover the qubits, but nothing fancy from what I can see (the value %1 defined by the user of %0 stays in the first qubit, and so does the k-th value defined by a following op that uses %1 as its k-th operand).

Now, what you actually struggle with is the difference between value semantics and memory semantics. SSA is about values stored in memory. For that sake, qubits are memory. Their state is the SSA value. You can, of course, model the memory. With a new type or by trying to abuse memref.

quantum.circuit inputs(%0: !quantum.bit, %1: !quantum.bit) {
  gate.H(%1) : !quantum.bit
  gate.CNOT(%0, %1) : !quantum.bit, !quantum.bit
  gate.H(%1) : !quantum.bit

This type would have memory semantics, i.e. each gate mutates the state of the qubits it operates on. In general, this representation is harder for a compiler to reason about because it needs to understand what gets mutated and in which order (that’s one of the reasons why SSA exists in the first place). However, it is relatively straightforward in the quantum domain as long as you keep each qubit as a separate SSA value: just look at all uses of each value and rely on the lexical order of operations to reconstruct the chain. The latter is possible because you wouldn’t have “classical” control flow with blocks and branches. This is only possible because you have a reasonable number of qubits (I’m not aware of even models that use millions of qubits, whereas we routinely consume megabytes of memory and taking each bit of memory as a separate operation argument sounds impractical) so the addressing/aliasing problems that make memory-based semantics hard to reason about go away. The moment you decide you have a “set of qubits” where you don’t care about which one specifically to use, or to have some classical constructs around (e.g. repeat a circuit multiple times), these problems will bite you back.

There is little support for such memory-related operations in the pattern rewriting infrastructure, but it may still be feasible to do the rewriting with some matching logic based on value use-lists. The closest analogy we have is Linalg buffers.