Implement pass that splits multiline ifs in qasm2 (#252)

Closes #189 

Also, it will probably break the tests in #248

---------

Co-authored-by: Phillip Weinberg <[email protected]>

Author: david-pl

src/bloqade/qasm2/emit/target.py

@@ -27,6 +27,7 @@ def __init__(
         allow_parallel: bool = False,
         allow_global: bool = False,
         custom_gate: bool = True,
+        unroll_ifs: bool = True,
     ) -> None:
         """Initialize the QASM2 target.
 
@@ -43,9 +44,14 @@ def __init__(
             qelib1 (bool):
                 Include the `include "qelib1.inc"` line in the resulting QASM2 AST that's
                 submitted to qBraid. Defaults to `True`.
+
             custom_gate (bool):
                 Include the custom gate definitions in the resulting QASM2 AST. Defaults to `True`. If `False`, all the qasm2.gate will be inlined.
 
+            unroll_ifs (bool):
+                Unrolls if statements with multiple qasm2 statements in the body in order to produce valid qasm2 output, which only allows a single
+                operation in an if body. Defaults to `True`.
+
 
 
         """
@@ -58,6 +64,7 @@ def __init__(
         self.custom_gate = custom_gate
         self.allow_parallel = allow_parallel
         self.allow_global = allow_global
+        self.unroll_ifs = unroll_ifs
 
         if allow_parallel:
             self.main_target = self.main_target.add(qasm2.dialects.parallel)
@@ -87,9 +94,11 @@ def emit(self, entry: ir.Method) -> ast.MainProgram:
 
         # make a cloned instance of kernel
         entry = entry.similar()
-        QASM2Fold(entry.dialects, inline_gate_subroutine=not self.custom_gate).fixpoint(
-            entry
-        )
+        QASM2Fold(
+            entry.dialects,
+            inline_gate_subroutine=not self.custom_gate,
+            unroll_ifs=self.unroll_ifs,
+        ).fixpoint(entry)
 
         if not self.allow_global:
             # rewrite global to parallel

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,66 @@
+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
+
+# TODO: unify with PR #248
+AllowedThenType = SingleQubitGate | TwoQubitCtrlGate | Measure | Reset
+
+DontLiftType = AllowedThenType | scf.Yield | func.Return | func.Invoke
+
+
+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()
+
+        lift_stmts = [stmt for stmt in then_stmts if not isinstance(stmt, DontLiftType)]
+
+        if len(lift_stmts) == 0:
+            return RewriteResult()
+
+        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_before(node)
+
+        node.delete()
+
+        return RewriteResult(has_done_something=True)

test/qasm2/emit/test_qasm2_emit.py

@@ -257,8 +257,7 @@ def multiline_then():
 
         return q
 
-    target = qasm2.emit.QASM2()
-
+    target = qasm2.emit.QASM2(unroll_ifs=False)
     with pytest.raises(InterpreterError):
         target.emit(multiline_then)
 

test/qasm2/passes/test_unroll_if.py

@@ -0,0 +1,105 @@
+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)
+
+
+def test_nested_kernels():
+    @qasm2.main
+    def nested(q: qasm2.QReg, c: qasm2.CReg):
+        qasm2.h(q[0])
+
+        qasm2.measure(q, c)
+        if c[0] == 1:
+            qasm2.x(q[0])
+            qasm2.x(q[1])
+
+        return q
+
+    @qasm2.main
+    def main():
+        q = qasm2.qreg(2)
+        c = qasm2.creg(2)
+
+        nested(q, c)
+
+        return c
+
+    target = QASM2()
+    ast = target.emit(main)
+
+    qasm2.parse.pprint(ast)
+
+
+def test_conditional_nested_kernel():
+    @qasm2.main
+    def nested(q: qasm2.QReg, c: qasm2.CReg):
+        qasm2.h(q[0])
+
+        qasm2.measure(q, c)
+
+        qasm2.x(q[0])
+        qasm2.x(q[1])
+
+        return q
+
+    @qasm2.main
+    def main():
+        q = qasm2.qreg(2)
+        c = qasm2.creg(2)
+
+        qasm2.h(q[0])
+        qasm2.measure(q, c)
+
+        if c[0] == 1:
+            nested(q, c)
+
+        return c
+
+    target = QASM2(unroll_ifs=True)
+    ast = target.emit(main)
+
+    qasm2.parse.pprint(ast)