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

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

bloqade/qasm2/rewrite/__init__.py

@@ -3,6 +3,7 @@ from .glob import (
     GlobalToParallelRule as GlobalToParallelRule,
 )
 from .register import RaiseRegisterRule as RaiseRegisterRule
+from .native_gates import RydbergGateSetRewriteRule as RydbergGateSetRewriteRule
 from .parallel_to_uop import ParallelToUOpRule as ParallelToUOpRule
 from .uop_to_parallel import (
     MergePolicyABC as MergePolicyABC,

bloqade/qasm2/rewrite/desugar.py

@@ -5,16 +5,17 @@ from kirin.passes import Pass
 from kirin.rewrite import abc, walk
 from kirin.dialects import py
 
+from bloqade.qasm2 import types
 from bloqade.qasm2.dialects import core
 
 
 class IndexingDesugarRule(abc.RewriteRule):
     def rewrite_Statement(self, node: ir.Statement) -> abc.RewriteResult:
         if isinstance(node, py.indexing.GetItem):
-            if node.obj.type.is_subseteq(core.QRegType):
+            if node.obj.type.is_subseteq(types.QRegType):
                 node.replace_by(core.QRegGet(reg=node.obj, idx=node.index))
                 return abc.RewriteResult(has_done_something=True)
-            elif node.obj.type.is_subseteq(core.CRegType):
+            elif node.obj.type.is_subseteq(types.CRegType):
                 node.replace_by(core.CRegGet(reg=node.obj, idx=node.index))
                 return abc.RewriteResult(has_done_something=True)
 

bloqade/qasm2/rewrite/heuristic_noise.py

@@ -18,6 +18,7 @@ class NoiseRewriteRule(rewrite_abc.RewriteRule):
     """
 
     address_analysis: Dict[ir.SSAValue, address.Address]
+    qubit_ssa_value: Dict[int, ir.SSAValue]
     gate_noise_params: native.GateNoiseParams = field(
         default_factory=native.GateNoiseParams
     )
@@ -25,15 +26,6 @@ class NoiseRewriteRule(rewrite_abc.RewriteRule):
         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
-
     def rewrite_Statement(self, node: ir.Statement) -> rewrite_abc.RewriteResult:
         if isinstance(node, uop.SingleQubitGate):
             return self.rewrite_single_qubit_gate(node)

bloqade/qasm2/rewrite/native_gates.py

@@ -173,6 +173,53 @@ class RydbergGateSetRewriteRule(abc.RewriteRule):
             cirq.XPowGate(exponent=0.5).on(self.cached_qubits[0]), node
         )
 
+    def rewrite_ccx(self, node: uop.CCX) -> abc.RewriteResult:
+        # from https://algassert.com/quirk#circuit=%7B%22cols%22:%5B%5B%22QFT3%22%5D,%5B%22inputA3%22,1,1,%22+=A3%22%5D,%5B1,1,1,%22%E2%80%A2%22,%22%E2%80%A2%22,%22X%22%5D,%5B1,1,1,%22%E2%80%A6%22,%22%E2%80%A6%22,%22%E2%80%A6%22%5D,%5B1,1,1,1,%22%E2%80%A2%22,%22Z%22%5D,%5B1,1,1,1,1,%22X%5E-%C2%BC%22%5D,%5B1,1,1,%22%E2%80%A2%22,1,%22Z%22%5D,%5B1,1,1,1,1,%22X%5E%C2%BC%22%5D,%5B1,1,1,1,%22%E2%80%A2%22,%22Z%22%5D,%5B1,1,1,1,1,%22X%5E-%C2%BC%22%5D,%5B1,1,1,%22Z%5E%C2%BC%22,%22Z%5E%C2%BC%22%5D,%5B1,1,1,1,%22H%22%5D,%5B1,1,1,%22%E2%80%A2%22,1,%22Z%22%5D,%5B1,1,1,%22%E2%80%A2%22,%22Z%22%5D,%5B1,1,1,1,%22X%5E-%C2%BC%22,%22X%5E%C2%BC%22%5D,%5B1,1,1,%22%E2%80%A2%22,%22Z%22%5D,%5B1,1,1,1,%22H%22%5D%5D%7D
+
+        # x^(1/4)
+        lam1, theta1, phi1 = map(
+            self.const_float,
+            map(around, (1.5707963267948966, 0.7853981633974483, -1.5707963267948966)),
+        )
+        lam1.insert_before(node)
+        theta1.insert_before(node)
+        phi1.insert_before(node)
+
+        lam1 = lam1.result
+        theta1 = theta1.result
+        phi1 = phi1.result
+
+        # x^(-1/4)
+        lam2, theta2, phi2 = map(
+            self.const_float,
+            map(around, (4.71238898038469, 0.7853981633974483, 1.5707963267948966)),
+        )
+        lam2.insert_before(node)
+        theta2.insert_before(node)
+        phi2.insert_before(node)
+        lam2 = lam2.result
+        theta2 = theta2.result
+        phi2 = phi2.result
+
+        uop.CZ(ctrl=node.ctrl1, qarg=node.qarg).insert_before(node)
+        uop.UGate(node.qarg, theta2, phi2, lam2).insert_before(node)
+        uop.CZ(ctrl=node.ctrl2, qarg=node.qarg).insert_before(node)
+        uop.UGate(node.qarg, theta1, phi1, lam1).insert_before(node)
+        uop.CZ(ctrl=node.ctrl1, qarg=node.qarg).insert_before(node)
+        uop.UGate(node.qarg, theta2, phi2, lam2).insert_before(node)
+        uop.T(node.ctrl1).insert_before(node)
+        uop.T(node.ctrl2).insert_before(node)
+        uop.H(node.ctrl1).insert_before(node)
+        uop.CZ(ctrl=node.ctrl2, qarg=node.qarg).insert_before(node)
+        uop.CZ(ctrl=node.ctrl2, qarg=node.ctrl1).insert_before(node)
+        uop.UGate(node.ctrl1, theta2, phi2, lam2).insert_before(node)
+        uop.UGate(node.qarg, theta2, phi2, lam2).insert_before(node)
+        uop.CZ(ctrl=node.ctrl2, qarg=node.ctrl1).insert_before(node)
+        uop.H(node.ctrl1).insert_before(node)
+        node.delete()  # delete the original CCX gate
+
+        return abc.RewriteResult(has_done_something=True)
+
     def rewrite_sxdg(self, node: uop.SXdag) -> abc.RewriteResult:
         return self._rewrite_1q_gates(
             cirq.XPowGate(exponent=-0.5).on(self.cached_qubits[0]), node
@@ -394,9 +441,12 @@ class RydbergGateSetRewriteRule(abc.RewriteRule):
         new_gate_stmts = self._generate_1q_gate_stmts(cirq_gate, node.qarg)
         return self._rewrite_gate_stmts(new_gate_stmts, node)
 
-    def _generate_2q_ctrl_gate_stmts(
+    def _generate_multi_ctrl_gate_stmts(
         self, cirq_gate: cirq.Operation, qubits_ssa: List[ir.SSAValue]
     ) -> list[ir.Statement]:
+        qubit_to_ssa_map = {
+            q: ssa for q, ssa in zip(self.cached_qubits[: len(qubits_ssa)], qubits_ssa)
+        }
         target_gates = self.gateset.decompose_to_target_gateset(cirq_gate, 0)
         new_stmts = []
         for new_gate in target_gates:
@@ -412,7 +462,7 @@ class RydbergGateSetRewriteRule(abc.RewriteRule):
                 new_stmts.append(phi2_stmt)
                 new_stmts.append(
                     uop.UGate(
-                        qarg=qubits_ssa[new_gate.qubits[0].x],
+                        qarg=qubit_to_ssa_map[new_gate.qubits[0]],
                         theta=phi0_stmt.result,
                         phi=phi1_stmt.result,
                         lam=phi2_stmt.result,
@@ -420,18 +470,31 @@ class RydbergGateSetRewriteRule(abc.RewriteRule):
                 )
             else:
                 # 2q
-                new_stmts.append(uop.CZ(ctrl=qubits_ssa[0], qarg=qubits_ssa[1]))
+                new_stmts.append(
+                    uop.CZ(
+                        ctrl=qubit_to_ssa_map[new_gate.qubits[0]],
+                        qarg=qubit_to_ssa_map[new_gate.qubits[1]],
+                    )
+                )
 
         return new_stmts
 
     def _rewrite_2q_ctrl_gates(
         self, cirq_gate: cirq.Operation, node: uop.TwoQubitCtrlGate
     ) -> abc.RewriteResult:
-        new_gate_stmts = self._generate_2q_ctrl_gate_stmts(
+        new_gate_stmts = self._generate_multi_ctrl_gate_stmts(
             cirq_gate, [node.ctrl, node.qarg]
         )
         return self._rewrite_gate_stmts(new_gate_stmts, node)
 
+    def _rewrite_3q_ctrl_gates(
+        self, cirq_gate: cirq.Operation, node: uop.CCX
+    ) -> abc.RewriteResult:
+        new_gate_stmts = self._generate_multi_ctrl_gate_stmts(
+            cirq_gate, [node.ctrl1, node.ctrl2, node.qarg]
+        )
+        return self._rewrite_gate_stmts(new_gate_stmts, node)
+
     def _rewrite_gate_stmts(
         self, new_gate_stmts: list[ir.Statement], node: ir.Statement
     ):

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/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,30 @@ 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)
+
+    @interp.impl(wire.MeasureAndReset)
+    def measure_and_reset(
+        self, interp: NSitesAnalysis, frame: interp.Frame, stmt: wire.MeasureAndReset
+    ):
+
+        # MeasureAndReset produces both a new wire
+        # and an integer which don't have any sites at all
+        return (NoSites(), NoSites())
+
+
 @op.dialect.register(key="op.nsites")
 class SquinOp(interp.MethodTable):
 

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,36 @@
-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,
+    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,
+)

bloqade/squin/op/_wrapper.py

@@ -0,0 +1,101 @@
+from kirin.lowering import wraps
+
+from . import stmts, types
+
+
+@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: ...