bloqade-circuit 0.4.4__py3-none-any.whl → 0.5.0__py3-none-any.whl

bloqade/squin/op/stmts.py

@@ -98,6 +98,15 @@ class ConstantUnitary(ConstantOp):
 
 @statement(dialect=dialect)
 class U3(PrimitiveOp):
+    """
+    The rotation operator U3(theta, phi, lam).
+    Note that we use the convention from the QASM2 specification, namely
+
+    $$
+    U_3(\theta, \phi, \lambda) = R_z(\phi) R_y(\theta) R_z(\lambda)
+    $$
+    """
+
     traits = frozenset({ir.Pure(), lowering.FromPythonCall(), Unitary(), FixedSites(1)})
     theta: ir.SSAValue = info.argument(types.Float)
     phi: ir.SSAValue = info.argument(types.Float)
@@ -110,7 +119,7 @@ class PhaseOp(PrimitiveOp):
     A phase operator.
 
     $$
-    PhaseOp(theta) = e^{i \theta} I
+    PhaseOp(\theta) = e^{i \theta} I
     $$
     """
 
@@ -124,7 +133,7 @@ class ShiftOp(PrimitiveOp):
     A phase shift operator.
 
     $$
-    Shift(theta) = \\begin{bmatrix} 1 & 0 \\\\ 0 & e^{i \\theta} \\end{bmatrix}
+    Shift(\theta) = \\begin{bmatrix} 1 & 0 \\\\ 0 & e^{i \\theta} \\end{bmatrix}
     $$
     """
 
@@ -141,6 +150,15 @@ class Reset(PrimitiveOp):
     traits = frozenset({ir.Pure(), lowering.FromPythonCall(), FixedSites(1)})
 
 
+@statement(dialect=dialect)
+class ResetToOne(PrimitiveOp):
+    """
+    Reset qubits to the one state. Mainly needed to accommodate cirq's GeneralizedAmplitudeDampingChannel
+    """
+
+    traits = frozenset({ir.Pure(), lowering.FromPythonCall(), FixedSites(1)})
+
+
 @statement
 class CliffordOp(ConstantUnitary):
     pass

bloqade/squin/qubit.py

@@ -17,6 +17,7 @@ from kirin.lowering import wraps
 from bloqade.types import Qubit, QubitType
 from bloqade.squin.op.types import Op, OpType
 
+from .types import MeasurementResult, MeasurementResultType
 from .lowering import ApplyAnyCallLowering
 
 dialect = ir.Dialect("squin.qubit")
@@ -65,8 +66,8 @@ class MeasureQubit(ir.Statement):
     name = "measure.qubit"
 
     traits = frozenset({lowering.FromPythonCall()})
-    qubit: ir.SSAValue = info.argument(ilist.IListType[QubitType])
-    result: ir.ResultValue = info.result(ilist.IListType[types.Bool])
+    qubit: ir.SSAValue = info.argument(QubitType)
+    result: ir.ResultValue = info.result(MeasurementResultType)
 
 
 @statement(dialect=dialect)
@@ -75,7 +76,7 @@ class MeasureQubitList(ir.Statement):
 
     traits = frozenset({lowering.FromPythonCall()})
     qubits: ir.SSAValue = info.argument(ilist.IListType[QubitType])
-    result: ir.ResultValue = info.result(ilist.IListType[types.Bool])
+    result: ir.ResultValue = info.result(ilist.IListType[MeasurementResultType])
 
 
 # NOTE: no dependent types in Python, so we have to mark it Any...
@@ -131,9 +132,11 @@ def apply(operator: Op, *qubits) -> None: ...
 
 
 @overload
-def measure(input: Qubit) -> bool: ...
+def measure(input: Qubit) -> MeasurementResult: ...
 @overload
-def measure(input: ilist.IList[Qubit, Any] | list[Qubit]) -> ilist.IList[bool, Any]: ...
+def measure(
+    input: ilist.IList[Qubit, Any] | list[Qubit],
+) -> ilist.IList[MeasurementResult, Any]: ...
 
 
 @wraps(MeasureAny)

bloqade/squin/rewrite/__init__.py

@@ -4,4 +4,5 @@ from .wrap_analysis import (
     WrapOpSiteAnalysis as WrapOpSiteAnalysis,
     WrapAddressAnalysis as WrapAddressAnalysis,
 )
+from .U3_to_clifford import SquinU3ToClifford as SquinU3ToClifford
 from .remove_dangling_qubits import RemoveDeadRegister as RemoveDeadRegister

bloqade/squin/rewrite/canonicalize.py

@@ -0,0 +1,60 @@
+from typing import cast
+
+from kirin import ir
+from kirin.rewrite import abc
+from kirin.dialects import cf
+
+from .. import wire
+
+
+class CanonicalizeWired(abc.RewriteRule):
+    def rewrite_Statement(self, node: ir.Statement) -> abc.RewriteResult:
+
+        if (
+            not isinstance(node, wire.Wired)
+            or len(node.qubits) != 0
+            or (parent_region := node.parent_region) is None
+        ):
+            return abc.RewriteResult()
+
+        parent_block = cast(ir.Block, node.parent_block)
+
+        # the body doesn't contain any quantum operations so we can safely inline the
+        # body into the parent block
+
+        # move all statements after `node` in the current block into another block
+        after_block = ir.Block()
+
+        stmt = node.next_stmt
+        while stmt is not None:
+            stmt.detach()
+            after_block.stmts.append(stmt)
+            stmt = node.next_stmt
+
+        # remap all results of the node to the arguments of the after_block
+        for result in node.results:
+            arg = after_block.args.append_from(result.type, result.name)
+            result.replace_by(arg)
+
+        parent_block_idx = parent_region._block_idx[parent_block]
+        # insert goto of parent block to the body block of the node.
+        parent_region.blocks.insert(parent_block_idx + 1, after_block)
+        # insert all blocks of the body of the node after the parent region
+        # making sure to convert any yield statements to jump statements to the after_block
+        parent_block.stmts.append(
+            cf.Branch(
+                arguments=(),
+                successor=node.body.blocks[0],
+            )
+        )
+        for block in reversed(node.body.blocks):
+            block.detach()
+            if isinstance((yield_stmt := block.last_stmt), wire.Yield):
+                yield_stmt.replace_by(
+                    cf.Branch(yield_stmt.values, successor=after_block)
+                )
+
+            parent_region.blocks.insert(parent_block_idx + 1, block)
+
+        node.delete()
+        return abc.RewriteResult(has_done_something=True)

bloqade/squin/rewrite/desugar.py

@@ -1,5 +1,5 @@
 from kirin import ir, types
-from kirin.dialects import ilist
+from kirin.dialects import py, ilist
 from kirin.rewrite.abc import RewriteRule, RewriteResult
 
 from bloqade.squin.qubit import (
@@ -53,12 +53,59 @@ class ApplyDesugarRule(RewriteRule):
         op = node.operator
         qubits = node.qubits
 
-        if len(qubits) == 1 and qubits[0].type.is_subseteq(ilist.IListType):
-            # NOTE: already calling with just a single argument that is already an ilist
+        if len(qubits) > 1 and all(q.type.is_subseteq(QubitType) for q in qubits):
+            (qubits_ilist_stmt := ilist.New(qubits)).insert_before(node)
+            qubits_ilist = qubits_ilist_stmt.result
+
+        elif len(qubits) == 1 and qubits[0].type.is_subseteq(QubitType):
+            (qubits_ilist_stmt := ilist.New(qubits)).insert_before(node)
+            qubits_ilist = qubits_ilist_stmt.result
+
+        elif len(qubits) == 1 and qubits[0].type.is_subseteq(
+            ilist.IListType[QubitType, types.Any]
+        ):
             qubits_ilist = qubits[0]
-        else:
-            (qubits_ilist_stmt := ilist.New(values=qubits)).insert_before(node)
+
+        elif len(qubits) == 1:
+            # TODO: remove this elif clause once we're at kirin v0.18
+            # NOTE: this is a temporary workaround for kirin#408
+            # currently type inference fails here in for loops since the loop var
+            # is an IList for some reason
+
+            if not isinstance(qubits[0], ir.ResultValue):
+                return RewriteResult()
+
+            is_ilist = isinstance(qbit_stmt := qubits[0].stmt, ilist.New)
+            if is_ilist:
+                if len(qbit_stmt.values) != 1:
+                    return RewriteResult()
+
+                if not isinstance(
+                    qbit_getindex_result := qbit_stmt.values[0], ir.ResultValue
+                ):
+                    return RewriteResult()
+
+                qbit_getindex = qbit_getindex_result.stmt
+            else:
+                qbit_getindex = qubits[0].stmt
+
+            if not isinstance(qbit_getindex, py.indexing.GetItem):
+                return RewriteResult()
+
+            if not qbit_getindex.obj.type.is_subseteq(
+                ilist.IListType[QubitType, types.Any]
+            ):
+                return RewriteResult()
+
+            if is_ilist:
+                values = qbit_stmt.values
+            else:
+                values = [qubits[0]]
+
+            (qubits_ilist_stmt := ilist.New(values=values)).insert_before(node)
             qubits_ilist = qubits_ilist_stmt.result
+        else:
+            return RewriteResult()
 
         stmt = Apply(operator=op, qubits=qubits_ilist)
         node.replace_by(stmt)

bloqade/squin/types.py

@@ -0,0 +1,8 @@
+from kirin import types
+
+
+class MeasurementResult:
+    pass
+
+
+MeasurementResultType = types.PyClass(MeasurementResult)

bloqade/squin/wire.py

@@ -6,12 +6,15 @@ circuits. Thus we do not define wrapping functions for the statements in this
 dialect.
 """
 
-from kirin import ir, types, lowering
+from kirin import ir, types, lowering, exception
 from kirin.decl import info, statement
+from kirin.dialects import func
 from kirin.lowering import wraps
+from kirin.ir.attrs.types import TypeAttribute
 
 from bloqade.types import Qubit, QubitType
 
+from .types import MeasurementResultType
 from .op.types import Op, OpType
 
 # from kirin.lowering import wraps
@@ -49,11 +52,87 @@ class Unwrap(ir.Statement):
     result: ir.ResultValue = info.result(WireType)
 
 
+@statement(dialect=dialect)
+class Wired(ir.Statement):
+    traits = frozenset()
+
+    qubits: tuple[ir.SSAValue, ...] = info.argument(QubitType)
+    memory_zone: str = info.attribute()
+    body: ir.Region = info.region(multi=True)
+
+    def __init__(
+        self,
+        body: ir.Region,
+        *qubits: ir.SSAValue,
+        memory_zone: str,
+        result_types: tuple[TypeAttribute, ...] | None = None,
+    ):
+        if result_types is None:
+            for block in body.blocks:
+                if isinstance(block.last_stmt, Yield):
+                    result_types = tuple(arg.type for arg in block.last_stmt.values)
+                    break
+
+        if result_types is None:
+            result_types = ()
+
+        super().__init__(
+            args=qubits,
+            args_slice={
+                "qubits": slice(0, None),
+            },
+            regions=[body],
+            attributes={
+                "memory_zone": ir.PyAttr(memory_zone)
+            },  # body of the wired statement
+            result_types=result_types,
+        )
+
+    def check(self):
+        entry_block = self.body.blocks[0]
+
+        if len(entry_block.args) != len(self.qubits):
+            raise exception.StaticCheckError(
+                f"Expected {len(self.qubits)} arguments, got {len(entry_block.args)}."
+            )
+        for arg in entry_block.args:
+            if not arg.type.is_subseteq(WireType):
+                raise exception.StaticCheckError(
+                    f"Expected argument of type {WireType}, got {arg.type}."
+                )
+        for block in self.body.blocks:
+            last_stmt = block.last_stmt
+            if isinstance(last_stmt, func.Return):
+                raise exception.StaticCheckError(
+                    "Return statements are not allowed in the body of a Wired statement."
+                )
+            elif isinstance(last_stmt, Yield) and len(last_stmt.values) != len(
+                self.results
+            ):
+                raise exception.StaticCheckError(
+                    f"Expected {len(self.results)} return values, got {len(last_stmt.values)}."
+                )
+
+
+@statement(dialect=dialect)
+class Yield(ir.Statement):
+    traits = frozenset({})
+    values: tuple[ir.SSAValue, ...] = info.argument(WireType)
+
+    def __init__(self, *args: ir.SSAValue):
+        super().__init__(
+            args=args,
+            args_slice={
+                "values": slice(0, None),
+            },
+        )
+
+
 # In Quake, you put a wire in and get a wire out when you "apply" an operator
 # In this case though we just need to indicate that an operator is applied to list[wires]
 @statement(dialect=dialect)
 class Apply(ir.Statement):  # apply(op, w1, w2, ...)
-    traits = frozenset({lowering.FromPythonCall(), ir.Pure()})
+    traits = frozenset({lowering.FromPythonCall()})
     operator: ir.SSAValue = info.argument(OpType)
     inputs: tuple[ir.SSAValue, ...] = info.argument(WireType)
 
@@ -88,6 +167,13 @@ class Broadcast(ir.Statement):
         )  # custom lowering required for wrapper to work here
 
 
+@statement(dialect=dialect)
+class RegionMeasure(ir.Statement):
+    traits = frozenset({lowering.FromPythonCall(), WireTerminator()})
+    wire: ir.SSAValue = info.argument(WireType)
+    result: ir.ResultValue = info.result(MeasurementResultType)
+
+
 # NOTE: measurement cannot be pure because they will collapse the state
 #       of the qubit. The state is a hidden state that is not visible to
 #      the user in the wire dialect.
@@ -96,14 +182,14 @@ class Measure(ir.Statement):
     traits = frozenset({lowering.FromPythonCall(), WireTerminator()})
     wire: ir.SSAValue = info.argument(WireType)
     qubit: ir.SSAValue = info.argument(QubitType)
-    result: ir.ResultValue = info.result(types.Int)
+    result: ir.ResultValue = info.result(MeasurementResultType)
 
 
 @statement(dialect=dialect)
-class NonDestructiveMeasure(ir.Statement):
+class LossResolvingMeasure(ir.Statement):
     traits = frozenset({lowering.FromPythonCall()})
     input_wire: ir.SSAValue = info.argument(WireType)
-    result: ir.ResultValue = info.result(types.Int)
+    result: ir.ResultValue = info.result(MeasurementResultType)
     out_wire: ir.ResultValue = info.result(WireType)
 
 

bloqade/stim/__init__.py

@@ -31,6 +31,7 @@ from ._wrappers import (
     z_error as z_error,
     detector as detector,
     identity as identity,
+    qubit_loss as qubit_loss,
     depolarize1 as depolarize1,
     depolarize2 as depolarize2,
     pauli_string as pauli_string,

bloqade/stim/_wrappers.py

@@ -190,3 +190,7 @@ def y_error(p: float, targets: tuple[int, ...]) -> None: ...
 
 @wraps(noise.ZError)
 def z_error(p: float, targets: tuple[int, ...]) -> None: ...
+
+
+@wraps(noise.QubitLoss)
+def qubit_loss(probs: tuple[float, ...], targets: tuple[int, ...]) -> None: ...

bloqade/stim/dialects/noise/emit.py

@@ -66,6 +66,7 @@ class EmitStimNoiseMethods(MethodTable):
         return ()
 
     @impl(stmts.TrivialError)
+    @impl(stmts.QubitLoss)
     def non_stim_error(
         self, emit: EmitStimMain, frame: EmitStrFrame, stmt: stmts.TrivialError
     ):

bloqade/stim/dialects/noise/stmts.py

@@ -104,3 +104,8 @@ class TrivialCorrelatedError(NonStimCorrelatedError):
 @statement(dialect=dialect)
 class TrivialError(NonStimError):
     name = "TRIV_ERROR"
+
+
+@statement(dialect=dialect)
+class QubitLoss(NonStimError):
+    name = "loss"

bloqade/stim/passes/squin_to_stim.py

@@ -8,6 +8,7 @@ from kirin.rewrite import (
     DeadCodeElimination,
     CommonSubexpressionElimination,
 )
+from kirin.analysis import const
 from kirin.ir.method import Method
 from kirin.passes.abc import Pass
 from kirin.rewrite.abc import RewriteResult
@@ -16,11 +17,12 @@ from bloqade.stim.groups import main as stim_main_group
 from bloqade.stim.rewrite import (
     SquinWireToStim,
     PyConstantToStim,
+    SquinNoiseToStim,
     SquinQubitToStim,
     SquinMeasureToStim,
     SquinWireIdentityElimination,
 )
-from bloqade.squin.rewrite import RemoveDeadRegister
+from bloqade.squin.rewrite import SquinU3ToClifford, RemoveDeadRegister
 
 
 @dataclass
@@ -31,10 +33,23 @@ class SquinToStim(Pass):
         # propagate constants
         rewrite_result = fold_pass(mt)
 
+        cp_frame, _ = const.Propagate(dialects=mt.dialects).run_analysis(mt)
+        cp_results = cp_frame.entries
+
         # Assume that address analysis and
         # wrapping has been done before this pass!
 
+        # Rewrite the noise statements first.
+        rewrite_result = (
+            Walk(SquinNoiseToStim(cp_results=cp_results))
+            .rewrite(mt.code)
+            .join(rewrite_result)
+        )
+
         # Wrap Rewrite + SquinToStim can happen w/ standard walk
+
+        rewrite_result = Walk(SquinU3ToClifford()).rewrite(mt.code).join(rewrite_result)
+
         rewrite_result = (
             Walk(
                 Chain(

bloqade/stim/rewrite/__init__.py

@@ -1,3 +1,4 @@
+from .squin_noise import SquinNoiseToStim as SquinNoiseToStim
 from .wire_to_stim import SquinWireToStim as SquinWireToStim
 from .qubit_to_stim import SquinQubitToStim as SquinQubitToStim
 from .squin_measure import SquinMeasureToStim as SquinMeasureToStim

bloqade/stim/rewrite/qubit_to_stim.py

@@ -1,12 +1,13 @@
 from kirin import ir
 from kirin.rewrite.abc import RewriteRule, RewriteResult
 
-from bloqade.squin import op, qubit
+from bloqade.squin import op, noise, qubit
 from bloqade.squin.rewrite import AddressAttribute
 from bloqade.stim.dialects import gate
 from bloqade.stim.rewrite.util import (
-    SQUIN_STIM_GATE_MAPPING,
+    SQUIN_STIM_OP_MAPPING,
     rewrite_Control,
+    rewrite_QubitLoss,
     insert_qubit_idx_from_address,
 )
 
@@ -30,11 +31,17 @@ class SquinQubitToStim(RewriteRule):
 
         # this is an SSAValue, need it to be the actual operator
         applied_op = stmt.operator.owner
+
+        if isinstance(applied_op, noise.stmts.QubitLoss):
+            return rewrite_QubitLoss(stmt)
+
         assert isinstance(applied_op, op.stmts.Operator)
 
         if isinstance(applied_op, op.stmts.Control):
             return rewrite_Control(stmt)
 
+        # need to handle Control through separate means
+
         # check if its adjoint, assume its canonicalized so no nested adjoints.
         is_conj = False
         if isinstance(applied_op, op.stmts.Adjoint):
@@ -44,9 +51,7 @@ class SquinQubitToStim(RewriteRule):
             is_conj = True
             applied_op = applied_op.op.owner
 
-        # need to handle Control through separate means
-        # but we can handle X, Y, Z, H, and S here just fine
-        stim_1q_op = SQUIN_STIM_GATE_MAPPING.get(type(applied_op))
+        stim_1q_op = SQUIN_STIM_OP_MAPPING.get(type(applied_op))
         if stim_1q_op is None:
             return RewriteResult()
 
@@ -55,7 +60,6 @@ class SquinQubitToStim(RewriteRule):
         if address_attr is None:
             return RewriteResult()
 
-        # sometimes you can get a whole AddressReg...
         assert isinstance(address_attr, AddressAttribute)
         qubit_idx_ssas = insert_qubit_idx_from_address(
             address=address_attr, stmt_to_insert_before=stmt

bloqade/stim/rewrite/squin_noise.py

@@ -0,0 +1,120 @@
+from typing import Dict, Tuple
+from dataclasses import dataclass
+
+from kirin.ir import SSAValue, Statement
+from kirin.analysis import const
+from kirin.dialects import py
+from kirin.rewrite.abc import RewriteRule, RewriteResult
+
+from bloqade.squin import wire, noise as squin_noise, qubit
+from bloqade.stim.dialects import noise as stim_noise
+from bloqade.stim.rewrite.util import (
+    create_wire_passthrough,
+    insert_qubit_idx_after_apply,
+)
+
+
+@dataclass
+class SquinNoiseToStim(RewriteRule):
+
+    cp_results: Dict[SSAValue, const.Result]
+
+    def rewrite_Statement(self, node: Statement) -> RewriteResult:
+        match node:
+            case qubit.Apply() | qubit.Broadcast():
+                return self.rewrite_Apply_and_Broadcast(node)
+            case _:
+                return RewriteResult()
+
+    def rewrite_Apply_and_Broadcast(
+        self, stmt: qubit.Apply | qubit.Broadcast
+    ) -> RewriteResult:
+        """Rewrite Apply and Broadcast to their stim statements."""
+
+        # this is an SSAValue, need it to be the actual operator
+        applied_op = stmt.operator.owner
+
+        if isinstance(applied_op, squin_noise.stmts.NoiseChannel):
+
+            qubit_idx_ssas = insert_qubit_idx_after_apply(stmt=stmt)
+            if qubit_idx_ssas is None:
+                return RewriteResult()
+
+            stim_stmt = None
+            if isinstance(applied_op, squin_noise.stmts.SingleQubitPauliChannel):
+                stim_stmt = self.rewrite_SingleQubitPauliChannel(stmt, qubit_idx_ssas)
+            elif isinstance(applied_op, squin_noise.stmts.TwoQubitPauliChannel):
+                stim_stmt = self.rewrite_TwoQubitPauliChannel(stmt, qubit_idx_ssas)
+
+            if isinstance(stmt, (wire.Apply, wire.Broadcast)):
+                create_wire_passthrough(stmt)
+
+            if stim_stmt is not None:
+                stmt.replace_by(stim_stmt)
+            if len(stmt.operator.owner.result.uses) == 0:
+                stmt.operator.owner.delete()
+
+            return RewriteResult(has_done_something=True)
+        return RewriteResult()
+
+    def rewrite_SingleQubitPauliChannel(
+        self,
+        stmt: qubit.Apply | qubit.Broadcast | wire.Broadcast | wire.Apply,
+        qubit_idx_ssas: Tuple[SSAValue],
+    ) -> Statement:
+        """Rewrite squin.noise.SingleQubitPauliChannel to stim.PauliChannel1."""
+
+        squin_channel = stmt.operator.owner
+        assert isinstance(squin_channel, squin_noise.stmts.SingleQubitPauliChannel)
+
+        params = self.cp_results.get(squin_channel.params).data
+        new_stmts = [
+            p_x := py.Constant(params[0]),
+            p_y := py.Constant(params[1]),
+            p_z := py.Constant(params[2]),
+        ]
+        for new_stmt in new_stmts:
+            new_stmt.insert_before(stmt)
+
+        stim_stmt = stim_noise.PauliChannel1(
+            targets=qubit_idx_ssas,
+            px=p_x.result,
+            py=p_y.result,
+            pz=p_z.result,
+        )
+        return stim_stmt
+
+    def rewrite_TwoQubitPauliChannel(
+        self,
+        stmt: qubit.Apply | qubit.Broadcast | wire.Broadcast | wire.Apply,
+        qubit_idx_ssas: Tuple[SSAValue],
+    ) -> Statement:
+        """Rewrite squin.noise.SingleQubitPauliChannel to stim.PauliChannel1."""
+
+        squin_channel = stmt.operator.owner
+        assert isinstance(squin_channel, squin_noise.stmts.TwoQubitPauliChannel)
+
+        params = self.cp_results.get(squin_channel.params).data
+        param_stmts = [py.Constant(p) for p in params]
+        for param_stmt in param_stmts:
+            param_stmt.insert_before(stmt)
+
+        stim_stmt = stim_noise.PauliChannel2(
+            targets=qubit_idx_ssas,
+            pix=param_stmts[0].result,
+            piy=param_stmts[1].result,
+            piz=param_stmts[2].result,
+            pxi=param_stmts[3].result,
+            pxx=param_stmts[4].result,
+            pxy=param_stmts[5].result,
+            pxz=param_stmts[6].result,
+            pyi=param_stmts[7].result,
+            pyx=param_stmts[8].result,
+            pyy=param_stmts[9].result,
+            pyz=param_stmts[10].result,
+            pzi=param_stmts[11].result,
+            pzx=param_stmts[12].result,
+            pzy=param_stmts[13].result,
+            pzz=param_stmts[14].result,
+        )
+        return stim_stmt