Merge branch 'main' into david/229-pyqrack-qubit-reset

Author: david-pl

src/bloqade/noise/native/_wrappers.py

@@ -8,15 +8,28 @@
 
 
 @wraps(native.AtomLossChannel)
-def atom_loss_channel(
-    qargs: ilist.IList[Qubit, Any] | list, *, prob: float
-) -> None: ...
+def atom_loss_channel(qargs: ilist.IList[Qubit, Any] | list, *, prob: float) -> None:
+    """Apply an atom loss channel to a list of qubits.
+
+    Args:
+        qargs (ilist.IList[Qubit, Any] | list): List of qubits to apply the noise to.
+        prob (float): The loss probability.
+    """
+    ...
 
 
 @wraps(native.PauliChannel)
 def pauli_channel(
     qargs: ilist.IList[Qubit, Any] | list, *, px: float, py: float, pz: float
-) -> None: ...
+) -> None:
+    """Apply a Pauli channel to a list of qubits.
+
+    Args:
+        qargs (ilist.IList[Qubit, Any] | list): List of qubits to apply the noise to.
+        px (float): Probability of X error.
+        py (float): Probability of Y error.
+        pz (float): Probability of Z error.
+    """
 
 
 @wraps(native.CZPauliChannel)
@@ -31,4 +44,20 @@ def cz_pauli_channel(
     py_qarg: float,
     pz_qarg: float,
     paired: bool,
-) -> None: ...
+) -> None:
+    """Insert noise for a CZ gate with a Pauli channel on qubits.
+
+    Args:
+        ctrls: List of control qubits.
+        qarg2: List of target qubits.
+        px_ctrl: Probability of X error on control qubits.
+        py_ctrl: Probability of Y error on control qubits.
+        pz_ctrl: Probability of Z error on control qubits.
+        px_qarg: Probability of X error on target qubits.
+        py_qarg: Probability of Y error on target qubits.
+        pz_qarg: Probability of Z error on target qubits.
+        paired: If True, the noise is applied to both control and target qubits
+            are not lost otherwise skip this error. If False Apply the noise on
+            the whatever qubit is not lost.
+    """
+    ...

src/bloqade/qasm2/_qasm_loading.py

@@ -3,7 +3,7 @@
 import pathlib
 from typing import Any
 
-from kirin import ir, types
+from kirin import ir, lowering
 from kirin.dialects import func
 
 from . import parse
@@ -16,6 +16,7 @@ def loads(
     *,
     kernel_name: str = "main",
     dialects: ir.DialectGroup | None = None,
+    returns: str | None = None,
     globals: dict[str, Any] | None = None,
     file: str | None = None,
     lineno_offset: int = 0,
@@ -54,7 +55,7 @@ def loads(
     # TODO: add source info
     stmt = parse.loads(qasm)
     qasm2_lowering = QASM2(dialects or main)
-    body = qasm2_lowering.run(
+    frame = qasm2_lowering.get_frame(
         stmt,
         source=qasm,
         file=file,
@@ -63,13 +64,21 @@ def loads(
         col_offset=col_offset,
         compactify=compactify,
     )
-    return_value = func.ConstantNone()
-    body.blocks[0].stmts.append(return_value)
-    body.blocks[0].stmts.append(func.Return(value_or_stmt=return_value))
 
+    if returns is not None:
+        return_value = frame.get(returns)
+        if return_value is None:
+            raise lowering.BuildError(f"Cannot find return value {returns}")
+    else:
+        return_value = func.ConstantNone()
+        frame.push(return_value)
+
+    return_node = frame.push(func.Return(value_or_stmt=return_value))
+
+    body = frame.curr_region
     code = func.Function(
         sym_name=kernel_name,
-        signature=func.Signature((), types.NoneType),
+        signature=func.Signature((), return_node.value.type),
         body=body,
     )
 
@@ -88,6 +97,7 @@ def loadfile(
     *,
     kernel_name: str = "main",
     dialects: ir.DialectGroup | None = None,
+    returns: str | None = None,
     globals: dict[str, Any] | None = None,
     file: str | None = None,
     lineno_offset: int = 0,
@@ -132,6 +142,7 @@ def loadfile(
         source,
         kernel_name=kernel_name,
         dialects=dialects,
+        returns=returns,
         globals=globals,
         file=file,
         lineno_offset=lineno_offset,

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/groups.py

@@ -2,6 +2,7 @@
 from kirin.prelude import structural_no_opt
 from kirin.dialects import scf, func, ilist, lowering
 
+from bloqade.noise import native
 from bloqade.qasm2.dialects import (
     uop,
     core,
@@ -90,6 +91,7 @@ def run_pass(
             noise,
             parallel,
             core,
+            native,
         ]
     )
 )

src/bloqade/qasm2/parse/lowering.py

@@ -28,6 +28,28 @@ def run(
         col_offset: int = 0,
         compactify: bool = True,
     ) -> ir.Region:
+
+        frame = self.get_frame(
+            stmt,
+            source=source,
+            globals=globals,
+            file=file,
+            lineno_offset=lineno_offset,
+            col_offset=col_offset,
+        )
+
+        return frame.curr_region
+
+    def get_frame(
+        self,
+        stmt: ast.Node,
+        source: str | None = None,
+        globals: dict[str, Any] | None = None,
+        file: str | None = None,
+        lineno_offset: int = 0,
+        col_offset: int = 0,
+        compactify: bool = True,
+    ) -> lowering.Frame:
         # TODO: add source info
         state = lowering.State(
             self,
@@ -41,13 +63,13 @@ def run(
             finalize_next=False,
         ) as frame:
             self.visit(state, stmt)
-            region = frame.curr_region
 
-        if compactify:
-            from kirin.rewrite import Walk, CFGCompactify
+            if compactify:
+                from kirin.rewrite import Walk, CFGCompactify
+
+                Walk(CFGCompactify()).rewrite(frame.curr_region)
 
-            Walk(CFGCompactify()).rewrite(region)
-        return region
+            return frame
 
     def visit(self, state: lowering.State[ast.Node], node: ast.Node) -> lowering.Result:
         name = node.__class__.__name__

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)