bloqade-circuit 0.2.3__py3-none-any.whl → 0.4.0__py3-none-any.whl

bloqade/qasm2/rewrite/{heuristic_noise.py → noise/heuristic_noise.py}

@@ -5,9 +5,8 @@ from kirin import ir
 from kirin.rewrite import abc as rewrite_abc
 from kirin.dialects import ilist
 
-from bloqade.noise import native
 from bloqade.analysis import address
-from bloqade.qasm2.dialects import uop, glob, parallel
+from bloqade.qasm2.dialects import uop, glob, noise, parallel
 
 
 @dataclass
@@ -18,21 +17,8 @@ class NoiseRewriteRule(rewrite_abc.RewriteRule):
     """
 
     address_analysis: Dict[ir.SSAValue, address.Address]
-    gate_noise_params: native.GateNoiseParams = field(
-        default_factory=native.GateNoiseParams
-    )
-    noise_model: native.MoveNoiseModelABC = field(
-        default_factory=native.TwoRowZoneModel
-    )
-
-    def __post_init__(self):
-        self.qubit_ssa_value: Dict[int, ir.SSAValue] = {}
-        for ssa_value, addr in self.address_analysis.items():
-            if (
-                isinstance(addr, address.AddressQubit)
-                and ssa_value not in self.qubit_ssa_value
-            ):
-                self.qubit_ssa_value[addr.data] = ssa_value
+    qubit_ssa_value: Dict[int, ir.SSAValue]
+    noise_model: noise.MoveNoiseModelABC = field(default_factory=noise.TwoRowZoneModel)
 
     def rewrite_Statement(self, node: ir.Statement) -> rewrite_abc.RewriteResult:
         if isinstance(node, uop.SingleQubitGate):
@@ -54,22 +40,18 @@ class NoiseRewriteRule(rewrite_abc.RewriteRule):
         qargs: ir.SSAValue,
         probs: Tuple[float, float, float, float],
     ):
-        native.PauliChannel(qargs, px=probs[0], py=probs[1], pz=probs[2]).insert_before(
+        noise.PauliChannel(qargs, px=probs[0], py=probs[1], pz=probs[2]).insert_before(
             node
         )
-        native.AtomLossChannel(qargs, prob=probs[3]).insert_before(node)
+        noise.AtomLossChannel(qargs, prob=probs[3]).insert_before(node)
 
         return rewrite_abc.RewriteResult(has_done_something=True)
 
     def rewrite_single_qubit_gate(self, node: uop.SingleQubitGate):
-        probs = (
-            self.gate_noise_params.local_px,
-            self.gate_noise_params.local_py,
-            self.gate_noise_params.local_pz,
-            self.gate_noise_params.local_loss_prob,
-        )
         (qargs := ilist.New(values=(node.qarg,))).insert_before(node)
-        return self.insert_single_qubit_noise(node, qargs.result, probs)
+        return self.insert_single_qubit_noise(
+            node, qargs.result, self.noise_model.local_errors
+        )
 
     def rewrite_global_single_qubit_gate(self, node: glob.UGate):
         addrs = self.address_analysis[node.registers]
@@ -85,14 +67,10 @@ class NoiseRewriteRule(rewrite_abc.RewriteRule):
             for qid in addr.data:
                 qargs.append(self.qubit_ssa_value[qid])
 
-        probs = (
-            self.gate_noise_params.global_px,
-            self.gate_noise_params.global_py,
-            self.gate_noise_params.global_pz,
-            self.gate_noise_params.global_loss_prob,
-        )
         (qargs := ilist.New(values=tuple(qargs))).insert_before(node)
-        return self.insert_single_qubit_noise(node, qargs.result, probs)
+        return self.insert_single_qubit_noise(
+            node, qargs.result, self.noise_model.global_errors
+        )
 
     def rewrite_parallel_single_qubit_gate(self, node: parallel.RZ | parallel.UGate):
         addrs = self.address_analysis[node.qargs]
@@ -102,15 +80,11 @@ class NoiseRewriteRule(rewrite_abc.RewriteRule):
         if not all(isinstance(addr, address.AddressQubit) for addr in addrs.data):
             return rewrite_abc.RewriteResult()
 
-        probs = (
-            self.gate_noise_params.local_px,
-            self.gate_noise_params.local_py,
-            self.gate_noise_params.local_pz,
-            self.gate_noise_params.local_loss_prob,
-        )
         assert isinstance(node.qargs, ir.ResultValue)
         assert isinstance(node.qargs.stmt, ilist.New)
-        return self.insert_single_qubit_noise(node, node.qargs, probs)
+        return self.insert_single_qubit_noise(
+            node, node.qargs, self.noise_model.local_errors
+        )
 
     def move_noise_stmts(
         self,
@@ -126,9 +100,9 @@ class NoiseRewriteRule(rewrite_abc.RewriteRule):
             nodes.append(
                 qargs := ilist.New(tuple(self.qubit_ssa_value[q] for q in qubits))
             )
-            nodes.append(native.AtomLossChannel(qargs.result, prob=probs[3]))
+            nodes.append(noise.AtomLossChannel(qargs.result, prob=probs[3]))
             nodes.append(
-                native.PauliChannel(qargs.result, px=probs[0], py=probs[1], pz=probs[2])
+                noise.PauliChannel(qargs.result, px=probs[0], py=probs[1], pz=probs[2])
             )
 
         return nodes
@@ -139,34 +113,30 @@ class NoiseRewriteRule(rewrite_abc.RewriteRule):
         qargs: ir.SSAValue,
     ) -> list[ir.Statement]:
         return [
-            native.CZPauliChannel(
+            noise.CZPauliChannel(
                 ctrls,
                 qargs,
-                px_ctrl=self.gate_noise_params.cz_paired_gate_px,
-                py_ctrl=self.gate_noise_params.cz_paired_gate_py,
-                pz_ctrl=self.gate_noise_params.cz_paired_gate_pz,
-                px_qarg=self.gate_noise_params.cz_paired_gate_px,
-                py_qarg=self.gate_noise_params.cz_paired_gate_py,
-                pz_qarg=self.gate_noise_params.cz_paired_gate_pz,
+                px_ctrl=self.noise_model.cz_paired_gate_px,
+                py_ctrl=self.noise_model.cz_paired_gate_py,
+                pz_ctrl=self.noise_model.cz_paired_gate_pz,
+                px_qarg=self.noise_model.cz_paired_gate_px,
+                py_qarg=self.noise_model.cz_paired_gate_py,
+                pz_qarg=self.noise_model.cz_paired_gate_pz,
                 paired=True,
             ),
-            native.CZPauliChannel(
+            noise.CZPauliChannel(
                 ctrls,
                 qargs,
-                px_ctrl=self.gate_noise_params.cz_unpaired_gate_px,
-                py_ctrl=self.gate_noise_params.cz_unpaired_gate_py,
-                pz_ctrl=self.gate_noise_params.cz_unpaired_gate_pz,
-                px_qarg=self.gate_noise_params.cz_unpaired_gate_px,
-                py_qarg=self.gate_noise_params.cz_unpaired_gate_py,
-                pz_qarg=self.gate_noise_params.cz_unpaired_gate_pz,
+                px_ctrl=self.noise_model.cz_unpaired_gate_px,
+                py_ctrl=self.noise_model.cz_unpaired_gate_py,
+                pz_ctrl=self.noise_model.cz_unpaired_gate_pz,
+                px_qarg=self.noise_model.cz_unpaired_gate_px,
+                py_qarg=self.noise_model.cz_unpaired_gate_py,
+                pz_qarg=self.noise_model.cz_unpaired_gate_pz,
                 paired=False,
             ),
-            native.AtomLossChannel(
-                ctrls, prob=self.gate_noise_params.cz_gate_loss_prob
-            ),
-            native.AtomLossChannel(
-                qargs, prob=self.gate_noise_params.cz_gate_loss_prob
-            ),
+            noise.AtomLossChannel(ctrls, prob=self.noise_model.cz_gate_loss_prob),
+            noise.AtomLossChannel(qargs, prob=self.noise_model.cz_gate_loss_prob),
         ]
 
     def rewrite_cz_gate(self, node: uop.CZ):

bloqade/{noise/native/rewrite.py → qasm2/rewrite/noise/remove_noise.py}

@@ -4,8 +4,8 @@ from kirin import ir
 from kirin.rewrite import abc, dce, walk, fixpoint
 from kirin.passes.abc import Pass
 
-from .stmts import PauliChannel, CZPauliChannel, AtomLossChannel
-from ._dialect import dialect
+from ...dialects.noise.stmts import PauliChannel, CZPauliChannel, AtomLossChannel
+from ...dialects.noise._dialect import dialect
 
 
 class RemoveNoiseRewrite(abc.RewriteRule):

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)

bloqade/qbraid/lowering.py

@@ -4,13 +4,13 @@ from dataclasses import field, dataclass
 from kirin import ir, types, passes
 from kirin.dialects import func, ilist
 
-from bloqade import noise, qasm2
+from bloqade import qasm2
 from bloqade.qbraid import schema
-from bloqade.qasm2.dialects import glob, parallel
+from bloqade.qasm2.dialects import glob, noise, parallel
 
 
 @ir.dialect_group(
-    [func, qasm2.core, qasm2.uop, parallel, glob, qasm2.expr, noise.native, ilist]
+    [func, qasm2.core, qasm2.uop, parallel, glob, qasm2.expr, noise, ilist]
 )
 def qbraid_noise(
     self,
@@ -192,7 +192,7 @@ class Lowering:
                 qargs := ilist.New(values=tuple(self.qubit_id_map[q] for q in qubits))
             )
             self.block_list.append(
-                noise.native.PauliChannel(px=px, py=py, pz=pz, qargs=qargs.result)
+                noise.PauliChannel(px=px, py=py, pz=pz, qargs=qargs.result)
             )
 
         for (p_ctrl, p_qarg), qubits in paired_layers.items():
@@ -204,7 +204,7 @@ class Lowering:
                 qargs := ilist.New(values=tuple(self.qubit_id_map[q] for q in qargs))
             )
             self.block_list.append(
-                noise.native.CZPauliChannel(
+                noise.CZPauliChannel(
                     paired=True,
                     px_ctrl=p_ctrl[0],
                     py_ctrl=p_ctrl[1],
@@ -226,7 +226,7 @@ class Lowering:
                 qargs := ilist.New(values=tuple(self.qubit_id_map[q] for q in qargs))
             )
             self.block_list.append(
-                noise.native.CZPauliChannel(
+                noise.CZPauliChannel(
                     paired=False,
                     px_ctrl=p_ctrl[0],
                     py_ctrl=p_ctrl[1],
@@ -285,7 +285,7 @@ class Lowering:
                 qargs := ilist.New(values=tuple(self.qubit_id_map[q] for q in qubits))
             )
             self.block_list.append(
-                noise.native.PauliChannel(px=px, py=py, pz=pz, qargs=qargs.result)
+                noise.PauliChannel(px=px, py=py, pz=pz, qargs=qargs.result)
             )
 
     def lower_measurement(self, operation: schema.Measurement):
@@ -303,7 +303,7 @@ class Lowering:
         for survival_prob, qubits in layers.items():
             self.block_list.append(qargs := ilist.New(values=qubits))
             self.block_list.append(
-                noise.native.AtomLossChannel(prob=survival_prob, qargs=qargs.result)
+                noise.AtomLossChannel(prob=survival_prob, qargs=qargs.result)
             )
 
     def lower_number(self, value: float | int) -> ir.SSAValue:

bloqade/squin/__init__.py

@@ -1,2 +1,8 @@
-from . import op as op, wire as wire, noise as noise, qubit as qubit
+from . import (
+    op as op,
+    wire as wire,
+    noise as noise,
+    qubit as qubit,
+    lowering as lowering,
+)
 from .groups import wired as wired, kernel as kernel

bloqade/squin/analysis/nsites/__init__.py

@@ -1,6 +1,7 @@
 # Need this for impl registration to work properly!
 from . import impls as impls
 from .lattice import (
+    Sites as Sites,
     NoSites as NoSites,
     AnySites as AnySites,
     NumberSites as NumberSites,

bloqade/squin/analysis/nsites/impls.py

@@ -1,6 +1,6 @@
 from kirin import interp
 
-from bloqade.squin import op
+from bloqade.squin import op, wire
 
 from .lattice import (
     NoSites,
@@ -9,6 +9,21 @@ from .lattice import (
 from .analysis import NSitesAnalysis
 
 
+@wire.dialect.register(key="op.nsites")
+class SquinWire(interp.MethodTable):
+
+    @interp.impl(wire.Apply)
+    @interp.impl(wire.Broadcast)
+    def apply(
+        self,
+        interp: NSitesAnalysis,
+        frame: interp.Frame,
+        stmt: wire.Apply | wire.Broadcast,
+    ):
+
+        return tuple(frame.get(input) for input in stmt.inputs)
+
+
 @op.dialect.register(key="op.nsites")
 class SquinOp(interp.MethodTable):
 

bloqade/squin/groups.py

@@ -1,11 +1,11 @@
 from kirin import ir, passes
 from kirin.prelude import structural_no_opt
+from kirin.rewrite import Walk, Chain
 from kirin.dialects import ilist
-from kirin.rewrite.walk import Walk
 
 from . import op, wire, qubit
 from .op.rewrite import PyMultToSquinMult
-from .rewrite.measure_desugar import MeasureDesugarRule
+from .rewrite.desugar import ApplyDesugarRule, MeasureDesugarRule
 
 
 @ir.dialect_group(structural_no_opt.union([op, qubit]))
@@ -13,7 +13,7 @@ def kernel(self):
     fold_pass = passes.Fold(self)
     typeinfer_pass = passes.TypeInfer(self)
     ilist_desugar_pass = ilist.IListDesugar(self)
-    measure_desugar_pass = Walk(MeasureDesugarRule())
+    desugar_pass = Walk(Chain(MeasureDesugarRule(), ApplyDesugarRule()))
     py_mult_to_mult_pass = PyMultToSquinMult(self)
 
     def run_pass(method: ir.Method, *, fold=True, typeinfer=True):
@@ -25,7 +25,7 @@ def kernel(self):
 
         if typeinfer:
             typeinfer_pass(method)
-            measure_desugar_pass.rewrite(method.code)
+            desugar_pass.rewrite(method.code)
 
         ilist_desugar_pass(method)
 

bloqade/squin/lowering.py

@@ -0,0 +1,27 @@
+import ast
+from dataclasses import dataclass
+
+from kirin import lowering
+
+from . import qubit
+
+
+@dataclass(frozen=True)
+class ApplyAnyCallLowering(lowering.FromPythonCall["qubit.ApplyAny"]):
+    """
+    Custom lowering for ApplyAny that collects vararg qubits into a single tuple argument
+    """
+
+    def lower(
+        self, stmt: type["qubit.ApplyAny"], state: lowering.State, node: ast.Call
+    ):
+        if len(node.args) < 2:
+            raise lowering.BuildError(
+                "Apply requires at least one operator and one qubit as arguments!"
+            )
+        op, *qubits = node.args
+        op_ssa = state.lower(op).expect_one()
+        qubits_lowered = [state.lower(qbit).expect_one() for qbit in qubits]
+
+        s = stmt(op_ssa, tuple(qubits_lowered))
+        return state.current_frame.push(s)

bloqade/squin/noise/__init__.py

@@ -1,27 +1,8 @@
-# Put all the proper wrappers here
-
-from kirin.lowering import wraps as _wraps
-
-from bloqade.squin.op.types import Op
-
 from . import stmts as stmts
-
-
-@_wraps(stmts.PauliError)
-def pauli_error(basis: Op, p: float) -> Op: ...
-
-
-@_wraps(stmts.PPError)
-def pp_error(op: Op, p: float) -> Op: ...
-
-
-@_wraps(stmts.Depolarize)
-def depolarize(n_qubits: int, p: float) -> Op: ...
-
-
-@_wraps(stmts.PauliChannel)
-def pauli_channel(n_qubits: int, params: tuple[float, ...]) -> Op: ...
-
-
-@_wraps(stmts.QubitLoss)
-def qubit_loss(p: float) -> Op: ...
+from ._dialect import dialect as dialect
+from ._wrapper import (
+    pp_error as pp_error,
+    depolarize as depolarize,
+    qubit_loss as qubit_loss,
+    pauli_channel as pauli_channel,
+)

bloqade/squin/noise/_wrapper.py

@@ -0,0 +1,25 @@
+from kirin.lowering import wraps
+
+from bloqade.squin.op.types import Op
+
+from . import stmts
+
+
+@wraps(stmts.PauliError)
+def pauli_error(basis: Op, p: float) -> Op: ...
+
+
+@wraps(stmts.PPError)
+def pp_error(op: Op, p: float) -> Op: ...
+
+
+@wraps(stmts.Depolarize)
+def depolarize(n_qubits: int, p: float) -> Op: ...
+
+
+@wraps(stmts.PauliChannel)
+def pauli_channel(n_qubits: int, params: tuple[float, ...]) -> Op: ...
+
+
+@wraps(stmts.QubitLoss)
+def qubit_loss(p: float) -> Op: ...

bloqade/squin/op/__init__.py

@@ -1,162 +1,37 @@
-from kirin import ir as _ir
-from kirin.prelude import structural_no_opt as _structural_no_opt
-from kirin.lowering import wraps as _wraps
-
 from . import stmts as stmts, types as types, rewrite as rewrite
+from .stdlib import (
+    ch as ch,
+    cx as cx,
+    cy as cy,
+    cz as cz,
+    rx as rx,
+    ry as ry,
+    rz as rz,
+    cphase as cphase,
+)
 from .traits import Unitary as Unitary, MaybeUnitary as MaybeUnitary
 from ._dialect import dialect as dialect
-
-
-@_wraps(stmts.Kron)
-def kron(lhs: types.Op, rhs: types.Op) -> types.Op: ...
-
-
-@_wraps(stmts.Mult)
-def mult(lhs: types.Op, rhs: types.Op) -> types.Op: ...
-
-
-@_wraps(stmts.Scale)
-def scale(op: types.Op, factor: complex) -> types.Op: ...
-
-
-@_wraps(stmts.Adjoint)
-def adjoint(op: types.Op) -> types.Op: ...
-
-
-@_wraps(stmts.Control)
-def control(op: types.Op, *, n_controls: int) -> types.Op:
-    """
-    Create a controlled operator.
-
-    Note, that when considering atom loss, the operator will not be applied if
-    any of the controls has been lost.
-
-    Args:
-        operator: The operator to apply under the control.
-        n_controls: The number qubits to be used as control.
-
-    Returns:
-        Operator
-    """
-    ...
-
-
-@_wraps(stmts.Identity)
-def identity(*, sites: int) -> types.Op: ...
-
-
-@_wraps(stmts.Rot)
-def rot(axis: types.Op, angle: float) -> types.Op: ...
-
-
-@_wraps(stmts.ShiftOp)
-def shift(theta: float) -> types.Op: ...
-
-
-@_wraps(stmts.PhaseOp)
-def phase(theta: float) -> types.Op: ...
-
-
-@_wraps(stmts.X)
-def x() -> types.Op: ...
-
-
-@_wraps(stmts.Y)
-def y() -> types.Op: ...
-
-
-@_wraps(stmts.Z)
-def z() -> types.Op: ...
-
-
-@_wraps(stmts.H)
-def h() -> types.Op: ...
-
-
-@_wraps(stmts.S)
-def s() -> types.Op: ...
-
-
-@_wraps(stmts.T)
-def t() -> types.Op: ...
-
-
-@_wraps(stmts.P0)
-def p0() -> types.Op: ...
-
-
-@_wraps(stmts.P1)
-def p1() -> types.Op: ...
-
-
-@_wraps(stmts.Sn)
-def spin_n() -> types.Op: ...
-
-
-@_wraps(stmts.Sp)
-def spin_p() -> types.Op: ...
-
-
-@_wraps(stmts.U3)
-def u(theta: float, phi: float, lam: float) -> types.Op: ...
-
-
-@_wraps(stmts.PauliString)
-def pauli_string(*, string: str) -> types.Op: ...
-
-
-# stdlibs
-@_ir.dialect_group(_structural_no_opt.add(dialect))
-def op(self):
-    def run_pass(method):
-        pass
-
-    return run_pass
-
-
-@op
-def rx(theta: float) -> types.Op:
-    """Rotation X gate."""
-    return rot(x(), theta)
-
-
-@op
-def ry(theta: float) -> types.Op:
-    """Rotation Y gate."""
-    return rot(y(), theta)
-
-
-@op
-def rz(theta: float) -> types.Op:
-    """Rotation Z gate."""
-    return rot(z(), theta)
-
-
-@op
-def cx() -> types.Op:
-    """Controlled X gate."""
-    return control(x(), n_controls=1)
-
-
-@op
-def cy() -> types.Op:
-    """Controlled Y gate."""
-    return control(y(), n_controls=1)
-
-
-@op
-def cz() -> types.Op:
-    """Control Z gate."""
-    return control(z(), n_controls=1)
-
-
-@op
-def ch() -> types.Op:
-    """Control H gate."""
-    return control(h(), n_controls=1)
-
-
-@op
-def cphase(theta: float) -> types.Op:
-    """Control Phase gate."""
-    return control(phase(theta), n_controls=1)
+from ._wrapper import (
+    h as h,
+    s as s,
+    t as t,
+    u as u,
+    x as x,
+    y as y,
+    z as z,
+    p0 as p0,
+    p1 as p1,
+    rot as rot,
+    kron as kron,
+    mult as mult,
+    phase as phase,
+    reset as reset,
+    scale as scale,
+    shift as shift,
+    spin_n as spin_n,
+    spin_p as spin_p,
+    adjoint as adjoint,
+    control as control,
+    identity as identity,
+    pauli_string as pauli_string,
+)