bloqade-circuit 0.7.12__py3-none-any.whl → 0.8.0__py3-none-any.whl

bloqade/squin/noise/__init__.py

@@ -1,10 +1,2 @@
-from . import stmts as stmts
+from . import stmts as stmts, _interface as _interface
 from ._dialect import dialect as dialect
-from ._wrapper import (
-    depolarize as depolarize,
-    qubit_loss as qubit_loss,
-    depolarize2 as depolarize2,
-    pauli_error as pauli_error,
-    two_qubit_pauli_channel as two_qubit_pauli_channel,
-    single_qubit_pauli_channel as single_qubit_pauli_channel,
-)

bloqade/squin/noise/_dialect.py

@@ -1,3 +1,3 @@
 from kirin import ir
 
-dialect = ir.Dialect(name="squin.noise")
+dialect = ir.Dialect("squin.noise")

bloqade/squin/noise/_interface.py

@@ -0,0 +1,45 @@
+from typing import Any, Literal, TypeVar
+
+from kirin.dialects import ilist
+from kirin.lowering import wraps
+
+from bloqade.types import Qubit
+
+from . import stmts
+
+
+@wraps(stmts.Depolarize)
+def depolarize(p: float, qubits: ilist.IList[Qubit, Any]) -> None: ...
+
+
+N = TypeVar("N", bound=int)
+
+
+@wraps(stmts.Depolarize2)
+def depolarize2(
+    p: float, controls: ilist.IList[Qubit, N], targets: ilist.IList[Qubit, N]
+) -> None: ...
+
+
+@wraps(stmts.SingleQubitPauliChannel)
+def single_qubit_pauli_channel(
+    px: float, py: float, pz: float, qubits: ilist.IList[Qubit, Any]
+) -> None: ...
+
+
+@wraps(stmts.TwoQubitPauliChannel)
+def two_qubit_pauli_channel(
+    probabilities: ilist.IList[float, Literal[15]],
+    controls: ilist.IList[Qubit, N],
+    targets: ilist.IList[Qubit, N],
+) -> None: ...
+
+
+@wraps(stmts.QubitLoss)
+def qubit_loss(p: float, qubits: ilist.IList[Qubit, Any]) -> None: ...
+
+
+@wraps(stmts.CorrelatedQubitLoss)
+def correlated_qubit_loss(
+    p: float, qubits: ilist.IList[ilist.IList[Qubit, N], Any]
+) -> None: ...

bloqade/squin/noise/stmts.py

@@ -2,74 +2,110 @@ from kirin import ir, types, lowering
 from kirin.decl import info, statement
 from kirin.dialects import ilist
 
+from bloqade.types import QubitType
+
 from ._dialect import dialect
-from ..op.types import OpType, NumOperators, MultiQubitPauliOpType
 
 
 @statement
 class NoiseChannel(ir.Statement):
-    traits = frozenset({ir.Pure(), lowering.FromPythonCall()})
-    result: ir.ResultValue = info.result(OpType)
+    traits = frozenset({lowering.FromPythonCall()})
 
 
-@statement(dialect=dialect)
-class PauliError(NoiseChannel):
-    basis: ir.SSAValue = info.argument(MultiQubitPauliOpType)
-    p: ir.SSAValue = info.argument(types.Float)
+@statement
+class SingleQubitNoiseChannel(NoiseChannel):
+    # NOTE: we are not adding e.g. qubits here, since inheriting then will
+    # change the order of the wrapper arguments
+    pass
+
+
+@statement
+class TwoQubitNoiseChannel(NoiseChannel):
+    pass
 
 
 @statement(dialect=dialect)
-class Depolarize(NoiseChannel):
+class SingleQubitPauliChannel(SingleQubitNoiseChannel):
     """
-    Apply depolarize error to single qubit
+    This will apply one of the randomly chosen Pauli operators according to the
+    given probabilities (p_x, p_y, p_z).
     """
 
-    p: ir.SSAValue = info.argument(types.Float)
+    px: ir.SSAValue = info.argument(types.Float)
+    py: ir.SSAValue = info.argument(types.Float)
+    pz: ir.SSAValue = info.argument(types.Float)
+    qubits: ir.SSAValue = info.argument(ilist.IListType[QubitType, types.Any])
+
+
+N = types.TypeVar("N", bound=types.Int)
 
 
 @statement(dialect=dialect)
-class Depolarize2(NoiseChannel):
+class TwoQubitPauliChannel(TwoQubitNoiseChannel):
     """
-    Apply correlated depolarize error to two qubits
-
     This will apply one of the randomly chosen Pauli products:
 
     {IX, IY, IZ, XI, XX, XY, XZ, YI, YX, YY, YZ, ZI, ZX, ZY, ZZ}
+
+    but the choice is weighed with the given probability.
+
+    NOTE: the given parameters are ordered as given in the list above!
     """
 
-    p: ir.SSAValue = info.argument(types.Float)
+    probabilities: ir.SSAValue = info.argument(
+        ilist.IListType[QubitType, types.Literal(15)]
+    )
+    controls: ir.SSAValue = info.argument(ilist.IListType[QubitType, N])
+    targets: ir.SSAValue = info.argument(ilist.IListType[QubitType, N])
 
 
 @statement(dialect=dialect)
-class SingleQubitPauliChannel(NoiseChannel):
-    params: ir.SSAValue = info.argument(ilist.IListType[types.Float, types.Literal(3)])
+class Depolarize(SingleQubitNoiseChannel):
+    """
+    Apply depolarize error to single qubit.
+
+    This randomly picks one of the three Pauli operators to apply. Each Pauli
+    operator has the probability `p / 3` to be selected. No operator is applied
+    with the probability `1 - p`.
+    """
+
+    p: ir.SSAValue = info.argument(types.Float)
+    qubits: ir.SSAValue = info.argument(ilist.IListType[QubitType, types.Any])
 
 
 @statement(dialect=dialect)
-class TwoQubitPauliChannel(NoiseChannel):
+class Depolarize2(TwoQubitNoiseChannel):
     """
-    This will apply one of the randomly chosen Pauli products:
+    Apply correlated depolarize error to two qubits
 
-    {IX, IY, IZ, XI, XX, XY, XZ, YI, YX, YY, YZ, ZI, ZX, ZY, ZZ}
+    This will apply one of the randomly chosen Pauli products each with probability `p / 15`:
 
-    but the choice is weighed with the given probability.
-
-    NOTE: the given parameters are ordered as given in the list above!
+    `{IX, IY, IZ, XI, XX, XY, XZ, YI, YX, YY, YZ, ZI, ZX, ZY, ZZ}`
     """
 
-    params: ir.SSAValue = info.argument(ilist.IListType[types.Float, types.Literal(15)])
+    p: ir.SSAValue = info.argument(types.Float)
+    controls: ir.SSAValue = info.argument(ilist.IListType[QubitType, N])
+    targets: ir.SSAValue = info.argument(ilist.IListType[QubitType, N])
 
 
 @statement(dialect=dialect)
-class QubitLoss(NoiseChannel):
+class QubitLoss(SingleQubitNoiseChannel):
+    """
+    Apply an atom loss with channel.
+    """
+
     # NOTE: qubit loss error (not supported by Stim)
     p: ir.SSAValue = info.argument(types.Float)
+    qubits: ir.SSAValue = info.argument(ilist.IListType[QubitType, types.Any])
 
 
 @statement(dialect=dialect)
-class StochasticUnitaryChannel(ir.Statement):
-    operators: ir.SSAValue = info.argument(ilist.IListType[OpType, NumOperators])
-    probabilities: ir.SSAValue = info.argument(
-        ilist.IListType[types.Float, NumOperators]
+class CorrelatedQubitLoss(NoiseChannel):
+    """
+    Apply a correlated atom loss channel.
+    """
+
+    p: ir.SSAValue = info.argument(types.Float)
+    qubits: ir.SSAValue = info.argument(
+        ilist.IListType[ilist.IListType[QubitType, N], types.Any]
     )
-    result: ir.ResultValue = info.result(OpType)

bloqade/squin/rewrite/U3_to_clifford.py

@@ -1,47 +1,56 @@
 # create rewrite rule name SquinMeasureToStim using kirin
 import math
-from typing import List, Tuple, Callable
 
 import numpy as np
 from kirin import ir
 from kirin.dialects import py
 from kirin.rewrite.abc import RewriteRule, RewriteResult
 
-from bloqade.squin import op, qubit
+from bloqade.squin import gate
 
 
-def sdag() -> list[ir.Statement]:
-    return [_op := op.stmts.S(), op.stmts.Adjoint(op=_op.result, is_unitary=True)]
+# Placeholder type, swap in an actual S statement with adjoint=True
+# during the rewrite method
+class Sdag(ir.Statement):
+    pass
+
+
+class SqrtXdag(ir.Statement):
+    pass
+
+
+class SqrtYdag(ir.Statement):
+    pass
 
 
 # (theta, phi, lam)
 U3_HALF_PI_ANGLE_TO_GATES: dict[
-    tuple[int, int, int], Callable[[], Tuple[List[ir.Statement], ...]]
+    tuple[int, int, int], list[type[ir.Statement]] | list[None]
 ] = {
-    (0, 0, 0): lambda: ([op.stmts.Identity(sites=1)],),
-    (0, 0, 1): lambda: ([op.stmts.S()],),
-    (0, 0, 2): lambda: ([op.stmts.Z()],),
-    (0, 0, 3): lambda: (sdag(),),
-    (1, 0, 0): lambda: ([op.stmts.SqrtY()],),
-    (1, 0, 1): lambda: ([op.stmts.S()], [op.stmts.SqrtY()]),
-    (1, 0, 2): lambda: ([op.stmts.H()],),
-    (1, 0, 3): lambda: (sdag(), [op.stmts.SqrtY()]),
-    (1, 1, 0): lambda: ([op.stmts.SqrtY()], [op.stmts.S()]),
-    (1, 1, 1): lambda: ([op.stmts.S()], [op.stmts.SqrtY()], [op.stmts.S()]),
-    (1, 1, 2): lambda: ([op.stmts.Z()], [op.stmts.SqrtY()], [op.stmts.S()]),
-    (1, 1, 3): lambda: (sdag(), [op.stmts.SqrtY()], [op.stmts.S()]),
-    (1, 2, 0): lambda: ([op.stmts.SqrtY()], [op.stmts.Z()]),
-    (1, 2, 1): lambda: ([op.stmts.S()], [op.stmts.SqrtY()], [op.stmts.Z()]),
-    (1, 2, 2): lambda: ([op.stmts.Z()], [op.stmts.SqrtY()], [op.stmts.Z()]),
-    (1, 2, 3): lambda: (sdag(), [op.stmts.SqrtY()], [op.stmts.Z()]),
-    (1, 3, 0): lambda: ([op.stmts.SqrtY()], sdag()),
-    (1, 3, 1): lambda: ([op.stmts.S()], [op.stmts.SqrtY()], sdag()),
-    (1, 3, 2): lambda: ([op.stmts.Z()], [op.stmts.SqrtY()], sdag()),
-    (1, 3, 3): lambda: (sdag(), [op.stmts.SqrtY()], sdag()),
-    (2, 0, 0): lambda: ([op.stmts.Y()],),
-    (2, 0, 1): lambda: ([op.stmts.S()], [op.stmts.Y()]),
-    (2, 0, 2): lambda: ([op.stmts.Z()], [op.stmts.Y()]),
-    (2, 0, 3): lambda: (sdag(), [op.stmts.Y()]),
+    (0, 0, 0): [None],
+    (0, 0, 1): [gate.stmts.S],
+    (0, 0, 2): [gate.stmts.Z],
+    (0, 0, 3): [Sdag],
+    (1, 0, 0): [gate.stmts.SqrtY],
+    (1, 0, 1): [gate.stmts.S, gate.stmts.SqrtY],
+    (1, 0, 2): [gate.stmts.H],
+    (1, 0, 3): [Sdag, gate.stmts.SqrtY],
+    (1, 1, 0): [gate.stmts.S, SqrtXdag],
+    (1, 1, 1): [gate.stmts.Z, SqrtXdag],
+    (1, 1, 2): [Sdag, SqrtXdag],
+    (1, 1, 3): [SqrtXdag],
+    (1, 2, 0): [gate.stmts.Z, SqrtYdag],
+    (1, 2, 1): [Sdag, SqrtYdag],
+    (1, 2, 2): [SqrtYdag],
+    (1, 2, 3): [gate.stmts.S, SqrtYdag],
+    (1, 3, 0): [Sdag, gate.stmts.SqrtX],
+    (1, 3, 1): [gate.stmts.SqrtX],
+    (1, 3, 2): [gate.stmts.S, gate.stmts.SqrtX],
+    (1, 3, 3): [gate.stmts.Z, gate.stmts.SqrtX],
+    (2, 0, 0): [gate.stmts.Y],
+    (2, 0, 1): [gate.stmts.S, gate.stmts.Y],
+    (2, 0, 2): [gate.stmts.X],
+    (2, 0, 3): [Sdag, gate.stmts.Y],
 }
 
 
@@ -61,8 +70,8 @@ class SquinU3ToClifford(RewriteRule):
     """
 
     def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
-        if isinstance(node, (qubit.Apply, qubit.Broadcast)):
-            return self.rewrite_ApplyOrBroadcast_onU3(node)
+        if isinstance(node, gate.stmts.U3):
+            return self.rewrite_U3(node)
         else:
             return RewriteResult()
 
@@ -87,35 +96,39 @@ class SquinU3ToClifford(RewriteRule):
         else:
             return round((angle / math.tau) % 1 * 4) % 4
 
-    def rewrite_ApplyOrBroadcast_onU3(
-        self, node: qubit.Apply | qubit.Broadcast
-    ) -> RewriteResult:
+    def rewrite_U3(self, node: gate.stmts.U3) -> RewriteResult:
         """
         Rewrite Apply and Broadcast nodes to their clifford equivalent statements.
         """
-        if not isinstance(node.operator.owner, op.stmts.U3):
-            return RewriteResult()
 
-        gates = self.decompose_U3_gates(node.operator.owner)
+        gates = self.decompose_U3_gates(node)
 
         if len(gates) == 0:
             return RewriteResult()
 
-        for stmt_list in gates:
-            for gate_stmt in stmt_list[:-1]:
-                gate_stmt.insert_before(node)
-
-            oper = stmt_list[-1]
-            oper.insert_before(node)
-            new_node = node.__class__(operator=oper.result, qubits=node.qubits)
-            new_node.insert_before(node)
+        # Get rid of the U3 gate altogether if it's identity
+        if len(gates) == 1 and gates[0] is None:
+            node.delete()
+            return RewriteResult(has_done_something=True)
+
+        for gate_stmt in gates:
+            if gate_stmt is Sdag:
+                new_stmt = gate.stmts.S(adjoint=True, qubits=node.qubits)
+            elif gate_stmt is SqrtXdag:
+                new_stmt = gate.stmts.SqrtX(adjoint=True, qubits=node.qubits)
+            elif gate_stmt is SqrtYdag:
+                new_stmt = gate.stmts.SqrtY(adjoint=True, qubits=node.qubits)
+            else:
+                new_stmt = gate_stmt(qubits=node.qubits)
+            new_stmt.insert_before(node)
 
         node.delete()
 
-        # rewrite U3 to clifford gates
         return RewriteResult(has_done_something=True)
 
-    def decompose_U3_gates(self, node: op.stmts.U3) -> Tuple[List[ir.Statement], ...]:
+    def decompose_U3_gates(
+        self, node: gate.stmts.U3
+    ) -> list[type[ir.Statement]] | list[None]:
         """
         Rewrite U3 statements to clifford gates if possible.
         """
@@ -124,7 +137,13 @@ class SquinU3ToClifford(RewriteRule):
         lam = self.get_constant(node.lam)
 
         if theta is None or phi is None or lam is None:
-            return ()
+            return []
+
+        # Angles will be in units of turns, we convert to radians
+        # to allow for the old logic to work
+        theta = theta * math.tau
+        phi = phi * math.tau
+        lam = lam * math.tau
 
         # For U3(2*pi*n, phi, lam) = U3(0, 0, lam + phi) which is a Z rotation.
         if np.isclose(np.mod(theta, math.tau), 0):
@@ -139,13 +158,13 @@ class SquinU3ToClifford(RewriteRule):
         lam_half_pi: int | None = self.resolve_angle(lam)
 
         if theta_half_pi is None or phi_half_pi is None or lam_half_pi is None:
-            return ()
+            return []
 
         angles_key = (theta_half_pi, phi_half_pi, lam_half_pi)
         if angles_key not in U3_HALF_PI_ANGLE_TO_GATES:
             angles_key = equivalent_u3_para(*angles_key)
             if angles_key not in U3_HALF_PI_ANGLE_TO_GATES:
-                return ()
+                return []
 
         gates_stmts = U3_HALF_PI_ANGLE_TO_GATES.get(angles_key)
 
@@ -154,4 +173,4 @@ class SquinU3ToClifford(RewriteRule):
             gates_stmts is not None
         ), "internal error, U3 gates not found for angles: {}".format(angles_key)
 
-        return gates_stmts()
+        return gates_stmts

bloqade/squin/rewrite/__init__.py

@@ -1,7 +1,5 @@
 from .wrap_analysis import (
-    SitesAttribute as SitesAttribute,
     AddressAttribute as AddressAttribute,
-    WrapOpSiteAnalysis as WrapOpSiteAnalysis,
     WrapAddressAnalysis as WrapAddressAnalysis,
 )
 from .U3_to_clifford import SquinU3ToClifford as SquinU3ToClifford

bloqade/squin/rewrite/remove_dangling_qubits.py

@@ -1,14 +1,14 @@
 from kirin import ir
 from kirin.rewrite.abc import RewriteRule, RewriteResult
 
-from bloqade.squin import qubit
+from bloqade import qubit
 
 
 class RemoveDeadRegister(RewriteRule):
 
     def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
 
-        if not isinstance(node, qubit.New):
+        if not isinstance(node, qubit.stmts.New):
             return RewriteResult()
 
         if bool(node.result.uses):

bloqade/squin/rewrite/wrap_analysis.py

@@ -5,12 +5,11 @@ from kirin import ir
 from kirin.rewrite.abc import RewriteRule, RewriteResult
 from kirin.print.printer import Printer
 
-from bloqade.squin import op, wire
+from bloqade import qubit
 from bloqade.analysis.address import Address
-from bloqade.squin.analysis.nsites import Sites
 
 
-@wire.dialect.register
+@qubit.dialect.register
 @dataclass
 class AddressAttribute(ir.Attribute):
 
@@ -25,21 +24,6 @@ class AddressAttribute(ir.Attribute):
         printer.print(self.address)
 
 
-@op.dialect.register
-@dataclass
-class SitesAttribute(ir.Attribute):
-
-    name = "Sites"
-    sites: Sites
-
-    def __hash__(self) -> int:
-        return hash(self.sites)
-
-    def print_impl(self, printer: Printer) -> None:
-        # Can return to implementing this later
-        printer.print(self.sites)
-
-
 @dataclass
 class WrapAnalysis(RewriteRule):
 
@@ -61,7 +45,8 @@ class WrapAddressAnalysis(WrapAnalysis):
     address_analysis: dict[ir.SSAValue, Address]
 
     def wrap(self, value: ir.SSAValue) -> bool:
-        address_analysis_result = self.address_analysis[value]
+        if (address_analysis_result := self.address_analysis.get(value)) is None:
+            return False
 
         if value.hints.get("address") is not None:
             return False
@@ -69,19 +54,3 @@ class WrapAddressAnalysis(WrapAnalysis):
         value.hints["address"] = AddressAttribute(address_analysis_result)
 
         return True
-
-
-@dataclass
-class WrapOpSiteAnalysis(WrapAnalysis):
-
-    op_site_analysis: dict[ir.SSAValue, Sites]
-
-    def wrap(self, value: ir.SSAValue) -> bool:
-        op_site_analysis_result = self.op_site_analysis[value]
-
-        if value.hints.get("sites") is not None:
-            return False
-
-        value.hints["sites"] = SitesAttribute(op_site_analysis_result)
-
-        return True

bloqade/squin/stdlib/broadcast/__init__.py

@@ -0,0 +1,34 @@
+from .gate import (
+    h as h,
+    s as s,
+    t as t,
+    x as x,
+    y as y,
+    z as z,
+    cx as cx,
+    cy as cy,
+    cz as cz,
+    rx as rx,
+    ry as ry,
+    rz as rz,
+    u3 as u3,
+    s_adj as s_adj,
+    shift as shift,
+    t_adj as t_adj,
+    sqrt_x as sqrt_x,
+    sqrt_y as sqrt_y,
+    sqrt_z as sqrt_z,
+    sqrt_x_adj as sqrt_x_adj,
+    sqrt_y_adj as sqrt_y_adj,
+    sqrt_z_adj as sqrt_z_adj,
+)
+from .noise import (
+    bit_flip as bit_flip,
+    depolarize as depolarize,
+    qubit_loss as qubit_loss,
+    depolarize2 as depolarize2,
+    correlated_qubit_loss as correlated_qubit_loss,
+    two_qubit_pauli_channel as two_qubit_pauli_channel,
+    single_qubit_pauli_channel as single_qubit_pauli_channel,
+)
+from ._qubit import reset as reset, measure as measure

bloqade/squin/stdlib/broadcast/_qubit.py

@@ -0,0 +1,4 @@
+from bloqade.qubit.stdlib.broadcast import (
+    reset as reset,
+    measure as measure,
+)