bloqade-circuit 0.2.2__py3-none-any.whl → 0.3.0__py3-none-any.whl

bloqade/qasm2/dialects/noise.py

@@ -1,7 +1,9 @@
-from kirin import ir, types, lowering
+from kirin import ir, types, interp, lowering
 from kirin.decl import info, statement
 
+from bloqade.qasm2.parse import ast
 from bloqade.qasm2.types import QubitType
+from bloqade.qasm2.emit.gate import EmitQASM2Gate, EmitQASM2Frame
 
 dialect = ir.Dialect("qasm2.noise")
 
@@ -14,3 +16,33 @@ class Pauli1(ir.Statement):
     py: ir.SSAValue = info.argument(types.Float)
     pz: ir.SSAValue = info.argument(types.Float)
     qarg: ir.SSAValue = info.argument(QubitType)
+
+
+@dialect.register(key="emit.qasm2.gate")
+class NoiseEmit(interp.MethodTable):
+
+    @interp.impl(Pauli1)
+    def emit_pauli(
+        self,
+        emit: EmitQASM2Gate,
+        frame: EmitQASM2Frame,
+        stmt: Pauli1,
+    ):
+
+        px: ast.Number = frame.get(stmt.px)
+        py: ast.Number = frame.get(stmt.py)
+        pz: ast.Number = frame.get(stmt.pz)
+        qarg: ast.Bit | ast.Name = frame.get(stmt.qarg)
+
+        qarg_str = (
+            f"{qarg.name.id}[{qarg.addr}]"
+            if isinstance(qarg, ast.Bit)
+            else f"{qarg.id}"
+        )
+
+        frame.body.append(
+            ast.Comment(
+                text=f"noise.Pauli1({px.value}, {py.value}, {pz.value}) {qarg_str}]"
+            )
+        )
+        return ()

bloqade/qasm2/dialects/uop/__init__.py

@@ -1,4 +1,40 @@
-from . import _emit as _emit, stmts as stmts
-from .stmts import *  # noqa: F403
+from . import _emit as _emit, stmts as stmts, schedule as schedule
+from .stmts import (
+    CH as CH,
+    CU as CU,
+    CX as CX,
+    CY as CY,
+    CZ as CZ,
+    RX as RX,
+    RY as RY,
+    RZ as RZ,
+    SX as SX,
+    U1 as U1,
+    U2 as U2,
+    CCX as CCX,
+    CRX as CRX,
+    CRY as CRY,
+    CRZ as CRZ,
+    CSX as CSX,
+    CU1 as CU1,
+    CU3 as CU3,
+    RXX as RXX,
+    RZZ as RZZ,
+    H as H,
+    S as S,
+    T as T,
+    X as X,
+    Y as Y,
+    Z as Z,
+    Id as Id,
+    Sdag as Sdag,
+    Swap as Swap,
+    Tdag as Tdag,
+    CSwap as CSwap,
+    SXdag as SXdag,
+    UGate as UGate,
+    Barrier as Barrier,
+    SingleQubitGate as SingleQubitGate,
+    TwoQubitCtrlGate as TwoQubitCtrlGate,
+)
 from ._dialect import dialect as dialect
-from .schedule import *  # noqa: F403

bloqade/qasm2/dialects/uop/schedule.py

@@ -8,7 +8,7 @@ from ._dialect import dialect
 
 
 @dialect.register(key="qasm2.schedule.dag")
-class UOp(interp.MethodTable):
+class UOpSchedule(interp.MethodTable):
 
     @interp.impl(stmts.Id)
     @interp.impl(stmts.SXdag)

bloqade/qasm2/emit/impls/__init__.py

@@ -0,0 +1 @@
+from . import noise_native as noise_native

bloqade/qasm2/emit/impls/noise_native.py

@@ -0,0 +1,89 @@
+from typing import Any
+
+from kirin import interp
+from kirin.dialects import ilist
+
+from bloqade.noise import native
+from bloqade.qasm2.parse import ast
+from bloqade.qasm2.emit.gate import EmitQASM2Gate, EmitQASM2Frame
+
+
+@native.dialect.register(key="emit.qasm2.gate")
+class NativeNoise(interp.MethodTable):
+
+    def _convert(self, node: ast.Bit | ast.Name) -> str:
+        if isinstance(node, ast.Bit):
+            return f"{node.name.id}[{node.addr}]"
+        else:
+            return f"{node.id}"
+
+    @interp.impl(native.CZPauliChannel)
+    def emit_czp(
+        self,
+        emit: EmitQASM2Gate,
+        frame: EmitQASM2Frame,
+        stmt: native.CZPauliChannel,
+    ):
+        paired: bool = stmt.paired
+        px_ctrl: float = stmt.px_ctrl
+        py_ctrl: float = stmt.py_ctrl
+        pz_ctrl: float = stmt.pz_ctrl
+        px_qarg: float = stmt.pz_qarg
+        py_qarg: float = stmt.py_qarg
+        pz_qarg: float = stmt.pz_qarg
+        ctrls: ilist.IList[ast.Bit, Any] = frame.get(stmt.ctrls)
+        qargs: ilist.IList[ast.Bit, Any] = frame.get(stmt.qargs)
+        frame.body.append(
+            ast.Comment(
+                text=f"native.CZPauliChannel(paired={paired}, p_ctrl=[x:{px_ctrl}, y:{py_ctrl}, z:{pz_ctrl}], p_qarg[x:{px_qarg}, y:{py_qarg}, z:{pz_qarg}])"
+            )
+        )
+        frame.body.append(
+            ast.Comment(
+                text=f" -: ctrls: {', '.join([self._convert(q) for q in ctrls])}"
+            )
+        )
+        frame.body.append(
+            ast.Comment(
+                text=f" -: qargs: {', '.join([self._convert(q) for q in qargs])}"
+            )
+        )
+        return ()
+
+    @interp.impl(native.AtomLossChannel)
+    def emit_loss(
+        self,
+        emit: EmitQASM2Gate,
+        frame: EmitQASM2Frame,
+        stmt: native.AtomLossChannel,
+    ):
+        prob: float = stmt.prob
+        qargs: ilist.IList[ast.Bit, Any] = frame.get(stmt.qargs)
+        frame.body.append(ast.Comment(text=f"native.Atomloss(p={prob})"))
+        frame.body.append(
+            ast.Comment(
+                text=f" -: qargs: {', '.join([self._convert(q) for q in qargs])}"
+            )
+        )
+        return ()
+
+    @interp.impl(native.PauliChannel)
+    def emit_pauli(
+        self,
+        emit: EmitQASM2Gate,
+        frame: EmitQASM2Frame,
+        stmt: native.PauliChannel,
+    ):
+        px: float = stmt.px
+        py: float = stmt.py
+        pz: float = stmt.pz
+        qargs: ilist.IList[ast.Bit, Any] = frame.get(stmt.qargs)
+        frame.body.append(
+            ast.Comment(text=f"native.PauliChannel(px={px}, py={py}, pz={pz})")
+        )
+        frame.body.append(
+            ast.Comment(
+                text=f" -: qargs: {', '.join([self._convert(q) for q in qargs])}"
+            )
+        )
+        return ()

bloqade/qasm2/emit/main.py

@@ -5,8 +5,11 @@ from kirin.dialects import cf, scf, func
 from kirin.ir.dialect import Dialect as Dialect
 
 from bloqade.qasm2.parse import ast
+from bloqade.qasm2.dialects.uop import SingleQubitGate, TwoQubitCtrlGate
+from bloqade.qasm2.dialects.expr import GateFunction
 
 from .base import EmitQASM2Base, EmitQASM2Frame
+from ..dialects.core.stmts import Reset, Measure
 
 
 @dataclass
@@ -94,6 +97,24 @@ class Scf(interp.MethodTable):
 
         cond = emit.assert_node(ast.Cmp, frame.get(stmt.cond))
 
+        # NOTE: we need exactly one of those in the then body in order to emit valid QASM2
+        AllowedThenType = SingleQubitGate | TwoQubitCtrlGate | Measure | Reset
+
+        then_stmts = stmt.then_body.blocks[0].stmts
+        uop_stmts = 0
+        for s in then_stmts:
+            if isinstance(s, AllowedThenType):
+                uop_stmts += 1
+                continue
+
+            if isinstance(s, func.Invoke):
+                uop_stmts += isinstance(s.callee.code, GateFunction)
+
+        if uop_stmts != 1:
+            raise interp.InterpreterError(
+                "Cannot lower if-statement: QASM2 only allows exactly one quantum operation in the body."
+            )
+
         with emit.new_frame(stmt) as then_frame:
             then_frame.entries.update(frame.entries)
             emit.emit_block(then_frame, stmt.then_body.blocks[0])

bloqade/qasm2/emit/target.py

@@ -11,6 +11,7 @@ from bloqade.qasm2.passes.glob import GlobalToParallel
 from bloqade.qasm2.passes.py2qasm import Py2QASM
 from bloqade.qasm2.passes.parallel import ParallelToUOp
 
+from . import impls as impls  # register the tables
 from .gate import EmitQASM2Gate
 from .main import EmitQASM2Main
 
@@ -27,6 +28,8 @@ class QASM2:
         allow_parallel: bool = False,
         allow_global: bool = False,
         custom_gate: bool = True,
+        unroll_ifs: bool = True,
+        allow_noise: bool = True,
     ) -> None:
         """Initialize the QASM2 target.
 
@@ -43,13 +46,18 @@ class QASM2:
             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`.
+
 
 
         """
-        from bloqade import qasm2
+        from bloqade import noise, qasm2
 
         self.main_target = qasm2.main
         self.gate_target = qasm2.gate
@@ -58,6 +66,7 @@ class QASM2:
         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)
@@ -67,7 +76,11 @@ class QASM2:
             self.main_target = self.main_target.add(qasm2.dialects.glob)
             self.gate_target = self.gate_target.add(qasm2.dialects.glob)
 
-        if allow_global or allow_parallel:
+        if allow_noise:
+            self.main_target = self.main_target.add(noise.native)
+            self.gate_target = self.gate_target.add(noise.native)
+
+        if allow_global or allow_parallel or allow_noise:
             self.main_target = self.main_target.add(ilist)
             self.gate_target = self.gate_target.add(ilist)
 
@@ -87,9 +100,11 @@ class QASM2:
 
         # 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

bloqade/qasm2/groups.py

@@ -2,6 +2,7 @@ from kirin import ir, passes
 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 main(self):
             noise,
             parallel,
             core,
+            native,
         ]
     )
 )

bloqade/qasm2/parse/__init__.py

@@ -19,19 +19,22 @@ def loadfile(file: str | pathlib.Path):
         return loads(f.read())
 
 
-def pprint(node: ast.Node, *, console: Console | None = None):
+def pprint(node: ast.Node, *, console: Console | None = None, no_color: bool = False):
     if console:
-        return Printer(console).visit(node)
+        printer = Printer(console)
     else:
-        Printer().visit(node)
+        printer = Printer()
+    printer.console.no_color = no_color
+    printer.visit(node)
 
 
-def spprint(node: ast.Node, *, console: Console | None = None):
+def spprint(node: ast.Node, *, console: Console | None = None, no_color: bool = False):
     if console:
         printer = Printer(console)
     else:
         printer = Printer()
 
+    printer.console.no_color = no_color
     with printer.string_io() as stream:
         printer.visit(node)
         return stream.getvalue()

bloqade/qasm2/parse/lowering.py

@@ -1,5 +1,3 @@
-import os
-import pathlib
 from typing import Any
 from dataclasses import field, dataclass
 
@@ -19,131 +17,39 @@ class QASM2(lowering.LoweringABC[ast.Node]):
     hint_show_lineno: bool = field(default=True, kw_only=True)
     stacktrace: bool = field(default=True, kw_only=True)
 
-    def loads(
+    def run(
         self,
-        source: str,
-        kernel_name: str,
+        stmt: ast.Node,
         *,
-        returns: str | None = None,
+        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,
-    ) -> ir.Method:
-        from ..parse import loads
-
-        # TODO: add source info
-        stmt = loads(source)
+    ) -> ir.Region:
 
-        state = lowering.State(
-            self,
+        frame = self.get_frame(
+            stmt,
+            source=source,
+            globals=globals,
             file=file,
             lineno_offset=lineno_offset,
             col_offset=col_offset,
         )
-        with state.frame(
-            [stmt],
-            globals=globals,
-            finalize_next=False,
-        ) as frame:
-            try:
-                self.visit(state, stmt)
-                # append return statement with the return values
-                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))
-
-            except lowering.BuildError as e:
-                hint = state.error_hint(
-                    e,
-                    max_lines=self.max_lines,
-                    indent=self.hint_indent,
-                    show_lineno=self.hint_show_lineno,
-                )
-                if self.stacktrace:
-                    raise Exception(
-                        f"{e.args[0]}\n\n{hint}",
-                        *e.args[1:],
-                    ) from e
-                else:
-                    e.args = (hint,)
-                    raise e
-
-            region = frame.curr_region
-
-        if compactify:
-            from kirin.rewrite import Walk, CFGCompactify
-
-            Walk(CFGCompactify()).rewrite(region)
-
-        code = func.Function(
-            sym_name=kernel_name,
-            signature=func.Signature((), return_node.value.type),
-            body=region,
-        )
 
-        return ir.Method(
-            mod=None,
-            py_func=None,
-            sym_name=kernel_name,
-            arg_names=[],
-            dialects=self.dialects,
-            code=code,
-        )
+        return frame.curr_region
 
-    def loadfile(
-        self,
-        file: str | pathlib.Path,
-        *,
-        kernel_name: str | None = None,
-        returns: str | None = None,
-        globals: dict[str, Any] | None = None,
-        lineno_offset: int = 0,
-        col_offset: int = 0,
-        compactify: bool = True,
-    ) -> ir.Method:
-        if isinstance(file, str):
-            file = pathlib.Path(*os.path.split(file))
-
-        if not file.is_file() or not file.name.endswith(".qasm"):
-            raise ValueError("File must be a .qasm file")
-
-        kernel_name = (
-            file.name.replace(".qasm", "") if kernel_name is None else kernel_name
-        )
-
-        with file.open("r") as f:
-            source = f.read()
-
-        return self.loads(
-            source,
-            kernel_name,
-            returns=returns,
-            globals=globals,
-            file=str(file),
-            lineno_offset=lineno_offset,
-            col_offset=col_offset,
-            compactify=compactify,
-        )
-
-    def run(
+    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,
-    ) -> ir.Region:
+    ) -> lowering.Frame:
         # TODO: add source info
         state = lowering.State(
             self,
@@ -154,32 +60,16 @@ class QASM2(lowering.LoweringABC[ast.Node]):
         with state.frame(
             [stmt],
             globals=globals,
+            finalize_next=False,
         ) as frame:
-            try:
-                self.visit(state, stmt)
-            except lowering.BuildError as e:
-                hint = state.error_hint(
-                    e,
-                    max_lines=self.max_lines,
-                    indent=self.hint_indent,
-                    show_lineno=self.hint_show_lineno,
-                )
-                if self.stacktrace:
-                    raise Exception(
-                        f"{e.args[0]}\n\n{hint}",
-                        *e.args[1:],
-                    ) from e
-                else:
-                    e.args = (hint,)
-                    raise e
-
-            region = frame.curr_region
-
-        if compactify:
-            from kirin.rewrite import Walk, CFGCompactify
-
-            Walk(CFGCompactify()).rewrite(region)
-        return region
+            self.visit(state, stmt)
+
+            if compactify:
+                from kirin.rewrite import Walk, CFGCompactify
+
+                Walk(CFGCompactify()).rewrite(frame.curr_region)
+
+            return frame
 
     def visit(self, state: lowering.State[ast.Node], node: ast.Node) -> lowering.Result:
         name = node.__class__.__name__

bloqade/qasm2/parse/qasm2.lark

@@ -9,7 +9,7 @@ version: INT "." INT
 // stmts
 include: "include" STRING ";"
 ifstmt: "if" "(" expr "==" expr ")" ifbody
-ifbody: qop | "{" qop* "}" // allow multiple qops
+ifbody: qop
 opaque: "opaque" IDENTIFIER ["(" [params] ")"] qubits ";"
 barrier: "barrier" qubits ";"
 qreg: "qreg" IDENTIFIER "[" INT "]" ";"

bloqade/qasm2/passes/__init__.py

@@ -3,3 +3,4 @@ from .noise import NoisePass as NoisePass
 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

bloqade/qasm2/passes/fold.py

@@ -23,6 +23,8 @@ from kirin.rewrite.abc import RewriteResult
 
 from bloqade.qasm2.dialects import expr
 
+from .unroll_if import UnrollIfs
+
 
 @dataclass
 class QASM2Fold(Pass):
@@ -30,6 +32,7 @@ class QASM2Fold(Pass):
 
     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 @@ class QASM2Fold(Pass):
             .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)

bloqade/qasm2/passes/noise.py

@@ -21,6 +21,34 @@ class NoisePass(Pass):
     NOTE: This pass is not guaranteed to be supported long-term in bloqade. We will be
     moving towards a more general approach to noise modeling in the future.
 
+    ## Usage examples
+
+    ```
+    from bloqade import qasm2
+    from bloqade.noise import native
+    from bloqade.qasm2.passes.noise import NoisePass
+
+    noise_main = qasm2.extended.add(native.dialect)
+
+    @noise_main
+    def main():
+        q = qasm2.qreg(2)
+        qasm2.h(q[0])
+        qasm2.cx(q[0], q[1])
+        return q
+
+    # simple IR without any nosie
+    main.print()
+
+    noise_pass = NoisePass(noise_main)
+
+    # rewrite stuff in-place
+    noise_pass.unsafe_run(main)
+
+    # now, we do have noise channels in the IR
+    main.print()
+    ```
+
     """
 
     noise_model: native.MoveNoiseModelABC = field(
@@ -34,13 +62,32 @@ class NoisePass(Pass):
     def __post_init__(self):
         self.address_analysis = address.AddressAnalysis(self.dialects)
 
+    def get_qubit_values(self, mt: ir.Method):
+        frame, _ = self.address_analysis.run_analysis(mt, no_raise=self.no_raise)
+        qubit_ssa_values = {}
+        # Traverse statements in block order to fine the first SSA value for each qubit
+        for block in mt.callable_region.blocks:
+            for stmt in block.stmts:
+                if len(stmt.results) != 1:
+                    continue
+
+                addr = frame.entries.get(result := stmt.results[0])
+                if (
+                    isinstance(addr, address.AddressQubit)
+                    and (index := addr.data) not in qubit_ssa_values
+                ):
+                    qubit_ssa_values[index] = result
+
+        return qubit_ssa_values, frame.entries
+
     def unsafe_run(self, mt: ir.Method):
         result = LiftQubits(self.dialects).unsafe_run(mt)
-        frame, _ = self.address_analysis.run_analysis(mt, no_raise=self.no_raise)
+        qubit_ssa_value, address_analysis = self.get_qubit_values(mt)
         result = (
             Walk(
                 NoiseRewriteRule(
-                    address_analysis=frame.entries,
+                    qubit_ssa_value=qubit_ssa_value,
+                    address_analysis=address_analysis,
                     noise_model=self.noise_model,
                     gate_noise_params=self.gate_noise_params,
                 ),
@@ -49,5 +96,6 @@ class NoisePass(Pass):
             .rewrite(mt.code)
             .join(result)
         )
+
         result = Fixpoint(Walk(DeadCodeElimination())).rewrite(mt.code).join(result)
         return result

bloqade/qasm2/passes/parallel.py

@@ -27,6 +27,7 @@ from bloqade.qasm2.rewrite import (
     RaiseRegisterRule,
     UOpToParallelRule,
     SimpleOptimalMergePolicy,
+    RydbergGateSetRewriteRule,
 )
 from bloqade.squin.analysis import schedule
 
@@ -135,6 +136,7 @@ class UOpToParallel(Pass):
     """
 
     merge_policy_type: Type[MergePolicyABC] = SimpleOptimalMergePolicy
+    rewrite_to_native_first: bool = False
     constprop: const.Propagate = field(init=False)
 
     def __post_init__(self):
@@ -147,6 +149,13 @@ class UOpToParallel(Pass):
         if not result.has_done_something:
             return result
 
+        if self.rewrite_to_native_first:
+            result = (
+                Fixpoint(Walk(RydbergGateSetRewriteRule(self.dialects)))
+                .rewrite(mt.code)
+                .join(result)
+            )
+
         frame, _ = self.constprop.run_analysis(mt)
         result = Walk(WrapConst(frame)).rewrite(mt.code).join(result)