Expose kwarg and fix insertion order

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/rewrite/split_ifs.py

@@ -4,12 +4,11 @@
 
 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 | Reset
 
-DontLiftType = AllowedThenType | scf.Yield | func.Return
+DontLiftType = AllowedThenType | scf.Yield | func.Return | func.Invoke
 
 
 class LiftThenBody(RewriteRule):
@@ -21,16 +20,10 @@ def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
 
         then_stmts = node.then_body.stmts()
 
-        # TODO: should we leave QRegGet in?
-        lift_stmts: list[ir.Statement] = []
-        for stmt in then_stmts:
-            if isinstance(stmt, DontLiftType):
-                continue
+        lift_stmts = [stmt for stmt in then_stmts if not isinstance(stmt, DontLiftType)]
 
-            if isinstance(stmt, func.Invoke) and isinstance(stmt.callee, GateFunction):
-                continue
-
-            lift_stmts.append(stmt)
+        if len(lift_stmts) == 0:
+            return RewriteResult()
 
         for stmt in lift_stmts:
             stmt.detach()
@@ -66,7 +59,7 @@ def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
                 cond=node.cond, then_body=then_body, else_body=else_body
             )
 
-            new_if.insert_after(node)
+            new_if.insert_before(node)
 
         node.delete()
 

test/qasm2/passes/test_unroll_if.py

@@ -72,3 +72,34 @@ def main():
     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)