Implement pass that splits multiline ifs in qasm2

Author: david-pl

src/bloqade/qasm2/passes/__init__.py

@@ -3,3 +3,4 @@
 from .py2qasm import Py2QASM as Py2QASM
 from .qasm2py import QASM2Py as QASM2Py
 from .parallel import UOpToParallel as UOpToParallel
+from .unroll_if import UnrollIfs as UnrollIfs

src/bloqade/qasm2/passes/fold.py

@@ -23,13 +23,16 @@
 
 from bloqade.qasm2.dialects import expr
 
+from .unroll_if import UnrollIfs
+
 
 @dataclass
 class QASM2Fold(Pass):
     """Fold pass for qasm2.extended"""
 
     constprop: const.Propagate = field(init=False)
     inline_gate_subroutine: bool = True
+    unroll_ifs: bool = True
 
     def __post_init__(self):
         self.constprop = const.Propagate(self.dialects)
@@ -61,6 +64,9 @@ def unsafe_run(self, mt: Method) -> RewriteResult:
             .join(result)
         )
 
+        if self.unroll_ifs:
+            UnrollIfs(mt.dialects).unsafe_run(mt).join(result)
+
         # run typeinfer again after unroll etc. because we now insert
         # a lot of new nodes, which might have more precise types
         self.typeinfer.unsafe_run(mt)

src/bloqade/qasm2/passes/unroll_if.py

@@ -0,0 +1,25 @@
+from kirin import ir
+from kirin.passes import Pass
+from kirin.rewrite import (
+    Walk,
+    Chain,
+    Fixpoint,
+    ConstantFold,
+    CommonSubexpressionElimination,
+)
+
+from ..rewrite.split_ifs import LiftThenBody, SplitIfStmts
+
+
+class UnrollIfs(Pass):
+    """This pass lifts statements that are not UOP out of the if body and then splits whatever is left into multiple if statements so you obtain valid QASM2"""
+
+    def unsafe_run(self, mt: ir.Method):
+        result = Walk(LiftThenBody()).rewrite(mt.code)
+        result = Walk(SplitIfStmts()).rewrite(mt.code).join(result)
+        result = (
+            Fixpoint(Walk(Chain(ConstantFold(), CommonSubexpressionElimination())))
+            .rewrite(mt.code)
+            .join(result)
+        )
+        return result

src/bloqade/qasm2/rewrite/split_ifs.py

@@ -0,0 +1,65 @@
+from kirin import ir
+from kirin.dialects import scf, func
+from kirin.rewrite.abc import RewriteRule, RewriteResult
+
+from ..dialects.uop.stmts import SingleQubitGate, TwoQubitCtrlGate
+from ..dialects.core.stmts import Reset, Measure
+from ..dialects.expr.stmts import GateFunction
+
+# TODO: unify with PR #248
+AllowedThenType = SingleQubitGate | TwoQubitCtrlGate | Measure | GateFunction | Reset
+
+DontLiftType = AllowedThenType | scf.Yield | func.Return
+
+
+class LiftThenBody(RewriteRule):
+    """Lifts anything that's not a UOP or a yield/return out of the then body"""
+
+    def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
+        if not isinstance(node, scf.IfElse):
+            return RewriteResult()
+
+        then_stmts = node.then_body.stmts()
+
+        # TODO: should we leave QRegGet in?
+        lift_stmts = [stmt for stmt in then_stmts if not isinstance(stmt, DontLiftType)]
+
+        for stmt in lift_stmts:
+            stmt.detach()
+            stmt.insert_before(node)
+
+        return RewriteResult(has_done_something=True)
+
+
+class SplitIfStmts(RewriteRule):
+    """Splits the then body of an if-else statement into multiple if statements"""
+
+    def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
+        if not isinstance(node, scf.IfElse):
+            return RewriteResult()
+
+        *stmts, yield_or_return = node.then_body.stmts()
+
+        if len(stmts) == 1:
+            return RewriteResult()
+
+        is_yield = isinstance(yield_or_return, scf.Yield)
+
+        for stmt in stmts:
+            stmt.detach()
+
+            yield_or_return = scf.Yield() if is_yield else func.Return()
+
+            then_block = ir.Block((stmt, yield_or_return), argtypes=(node.cond.type,))
+            then_body = ir.Region(then_block)
+            else_body = node.else_body.clone()
+            else_body.detach()
+            new_if = scf.IfElse(
+                cond=node.cond, then_body=then_body, else_body=else_body
+            )
+
+            new_if.insert_after(node)
+
+        node.delete()
+
+        return RewriteResult(has_done_something=True)

test/qasm2/passes/test_unroll_if.py

@@ -0,0 +1,47 @@
+from bloqade import qasm2
+from bloqade.qasm2.emit import QASM2
+
+
+def test_unrolling_ifs():
+    @qasm2.main
+    def main():
+        q = qasm2.qreg(2)
+        c = qasm2.creg(2)
+
+        qasm2.h(q[0])
+        qasm2.measure(q[0], c[0])
+
+        if c[0] == 1:
+            qasm2.x(q[0])
+            qasm2.x(q[1])
+
+        return q
+
+    main.print()
+
+    target = QASM2()
+    ast = target.emit(main)
+
+    qasm2.parse.pprint(ast)
+
+    @qasm2.main
+    def main_unrolled():
+        q = qasm2.qreg(2)
+        c = qasm2.creg(2)
+
+        qasm2.h(q[0])
+        qasm2.measure(q[0], c[0])
+
+        if c[0] == 1:
+            qasm2.x(q[0])
+        if c[0] == 1:
+            qasm2.x(q[1])
+
+        return q
+
+    main_unrolled.print()
+
+    target = QASM2()
+    ast_unrolled = target.emit(main_unrolled)
+
+    qasm2.parse.pprint(ast_unrolled)