Loop unrolling

[julfarn]

We want to support loop unrolling for backands uncapable of (backward) control flow. In the first versions, loop unrolling will be mandatory.
Going forward, we will gradually increase support for loops throughout the whole pipeline, considering vectorization as well.

• Add loop unrolling to the pipeline.
• Ensure that all loops have been unrolled, otherwise emit error. In particular, all loops must have a static number of iterations.
• Implement test cases.

[julfarn]

@rainij , see also #47.

[julfarn]

I propose to add a new dedicated unrolling pass instead. Here's the reasoning

• Builtin loop unrolling exists for the affine dialect only, some tools exist for linalg.
  • We could change the pipeline to have affine as an intermediate step, but for technical reasons some instructions are lowered directly to scf within qrisp.
• It's relatively easy to implement.
• We have full control over the behavior.

[julfarn]

After having another look at the qrisp side of the problem, here are some insights:

• The python for loops I tried out are translated to stablehlo.while. This seems to be necessary at least for the case of dependent iterations. Stablehlo seems to have capabilities for parallelism, but only stablehlo.while exists as a general purpose loop (in particular, there is no stablehlo.for). I believe we have to accept that with a jax basis, we first get some form of while loops.
• Stablehlo has only very limited tooling for dealing with its own stablehlo.while ops. Mostly elimination in case the condition is always false. It does not get converted in the conversion pass to linalg. This is most likely because linalg does not have a while op. I believe the stablehlo.while operations need to be lowered manually in some way.
• Currently, qrisp converts stablehlo.while to scf.while using xDSL. This is semantically straightforward, and there's no obvious alternative (affine doesn't have a while loop).
• Because only for-loops can be unrolled (be it by a custom or builtin pass), at some point, while loops need to be analyzed and "raised back" to for-loops. The natural candidates are affine.for or scf.for. I am not aware of builtin tools to perform this task. The advantage of raising to affine.for is that there is a builtin loop unrolling pass. In addition, affine seems to be the correct starting environment for parallelism and vectorization @EnriqueNB.
• It is unfortunate that the for-loop semantics get lost by jasp and have to be reconstructed later. It would be nice if we got the correct semantics directly from qrisp output. However, it may be that we want to have the analysis and raising capabilities anyway. So I think it makes sense to implement a pass that transforms scf.while to affine.for if possible.

[burgholzer]

Just a comment from the sidelines: Have you checked whether the situation is similar when using PennyLane/Catalyst as a frontend/compiler? They also use StableHLO, but I have not yet perceived this kind of limitation to while loops when we interacted with the Catalyst compiler in the past. Maybe there is a way to get around this in Qrisp.

What you are proposing here sounds like a whole lot of patch work glue code to "fix" something that shouldn't even be a problem or should be the frontend's responsibility.
But that's just my random opinion from following the discussions.

[julfarn]

What you are proposing here sounds like a whole lot of patch work glue code to "fix" something that shouldn't even be a problem or should be the frontend's responsibility.

I agree with this sentiment, though I hope it will turn out not to be that much work. Good point about Catalyst, I've been meaning to have a closer look since the last time you suggested it, but haven't done so yet.

[julfarn]

Just a quick WIP update on Catalyst, so that it's written down somewhere:

This is the syntax for writing a for loop in Catalyst that will compile to scf.for:

@for_loop(1, n, 1)
def loop_fn(i): 
    qp.CNOT(wires=(i - 1, i))
loop_fn()

You can tell this is a very customized workaround, which was probably not the first choice of the developers.
The following also works (in this small example), but is discouraged. The loop was unrolled in this case during compilation.

for i in range(1, 5):
    qp.CNOT(wires=(i - 1, i))

For reference, here's what the corresponding loop looks like in qrisp:

for i in jrange(1, 5):
    cx(q_var[i-1], q_var[i])

Here, jrange is a qrisp function. So qrisp may have some control over the lowering here.

Still open questions:

• How does the unrolling happen in Catalyst?
• Could qrisp use jrange to better control the lowering of for-loops?

Also tagging @RiccardoRomanello @ctfle because you were interested.