bloqade-circuit 0.4.5__py3-none-any.whl → 0.5.1__py3-none-any.whl

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/auxiliary/interp.py

@@ -1,7 +1,6 @@
 from kirin import interp
 
 from . import stmts
-from .types import RecordResult
 from ._dialect import dialect
 
 
@@ -28,12 +27,3 @@ class StimAuxMethods(interp.MethodTable):
         stmt: stmts.Neg,
     ):
         return (-frame.get(stmt.operand),)
-
-    @interp.impl(stmts.GetRecord)
-    def get_rec(
-        self,
-        interpreter: interp.Interpreter,
-        frame: interp.Frame,
-        stmt: stmts.GetRecord,
-    ):
-        return (RecordResult(value=frame.get(stmt.id)),)

bloqade/stim/dialects/auxiliary/stmts/annotate.py

@@ -10,7 +10,7 @@ PyNum = types.Union(types.Int, types.Float)
 @statement(dialect=dialect)
 class GetRecord(ir.Statement):
     name = "get_rec"
-    traits = frozenset({lowering.FromPythonCall()})
+    traits = frozenset({ir.Pure(), lowering.FromPythonCall()})
     id: ir.SSAValue = info.argument(type=types.Int)
     result: ir.ResultValue = info.result(type=RecordType)
 

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/__init__.py

@@ -1 +1 @@
-from .squin_to_stim import SquinToStim as SquinToStim
+from .squin_to_stim import SquinToStimPass as SquinToStimPass

bloqade/stim/passes/simplify_ifs.py

@@ -0,0 +1,32 @@
+from dataclasses import dataclass
+
+from kirin import ir
+from kirin.passes import Pass
+from kirin.rewrite import (
+    Walk,
+    Chain,
+    Fixpoint,
+    ConstantFold,
+    CommonSubexpressionElimination,
+)
+
+from ..rewrite.ifs_to_stim import StimLiftThenBody, StimSplitIfStmts
+
+
+@dataclass
+class StimSimplifyIfs(Pass):
+
+    def unsafe_run(self, mt: ir.Method):
+
+        result = Chain(
+            Fixpoint(Walk(StimLiftThenBody())),
+            Walk(StimSplitIfStmts()),
+        ).rewrite(mt.code)
+
+        result = (
+            Fixpoint(Walk(Chain(ConstantFold(), CommonSubexpressionElimination())))
+            .rewrite(mt.code)
+            .join(result)
+        )
+
+        return result

bloqade/stim/passes/squin_to_stim.py

@@ -5,58 +5,153 @@ from kirin.rewrite import (
     Walk,
     Chain,
     Fixpoint,
+    CFGCompactify,
+    InlineGetItem,
+    InlineGetField,
     DeadCodeElimination,
     CommonSubexpressionElimination,
 )
+from kirin.analysis import const
+from kirin.dialects import scf, ilist
 from kirin.ir.method import Method
 from kirin.passes.abc import Pass
 from kirin.rewrite.abc import RewriteResult
+from kirin.passes.inline import InlinePass
 
-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,
+    WrapAddressAnalysis,
+)
+from bloqade.rewrite.passes import CanonicalizeIList
+from bloqade.analysis.address import AddressAnalysis
+from bloqade.analysis.measure_id import MeasurementIDAnalysis
+
+from .simplify_ifs import StimSimplifyIfs
+from ..rewrite.ifs_to_stim import IfToStim
 
 
 @dataclass
-class SquinToStim(Pass):
+class SquinToStimPass(Pass):
 
     def unsafe_run(self, mt: Method) -> RewriteResult:
-        fold_pass = Fold(mt.dialects)
-        # 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!
+        # inline aggressively:
+        rewrite_result = InlinePass(
+            dialects=mt.dialects, no_raise=self.no_raise
+        ).unsafe_run(mt)
+
+        rule = Chain(
+            InlineGetField(),
+            InlineGetItem(),
+            scf.unroll.ForLoop(),
+            scf.trim.UnusedYield(),
+        )
+        rewrite_result = Fixpoint(Walk(rule)).rewrite(mt.code).join(rewrite_result)
+        # fold_pass = Fold(mt.dialects, no_raise=self.no_raise)
+        # rewrite_result = fold_pass(mt)
+        rewrite_result = (
+            Walk(Fixpoint(CFGCompactify())).rewrite(mt.code).join(rewrite_result)
+        )
+        rewrite_result = (
+            StimSimplifyIfs(mt.dialects, no_raise=self.no_raise)
+            .unsafe_run(mt)
+            .join(rewrite_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)
+        rewrite_result = (
+            Walk(Chain(ilist.rewrite.ConstList2IList(), ilist.rewrite.Unroll()))
+            .rewrite(mt.code)
+            .join(rewrite_result)
+        )
+        rewrite_result = Fold(mt.dialects, no_raise=self.no_raise)(mt)
+
+        rewrite_result = (
+            CanonicalizeIList(dialects=mt.dialects, no_raise=self.no_raise)
+            .unsafe_run(mt)
+            .join(rewrite_result)
+        )
+
+        # after this the program should be in a state where it is analyzable
+        # -------------------------------------------------------------------
+
+        mia = MeasurementIDAnalysis(dialects=mt.dialects)
+        meas_analysis_frame, _ = mia.run_analysis(mt, no_raise=self.no_raise)
+
+        aa = AddressAnalysis(dialects=mt.dialects)
+        address_analysis_frame, _ = aa.run_analysis(mt, no_raise=self.no_raise)
+
+        # wrap the address analysis result
+        rewrite_result = (
+            Walk(WrapAddressAnalysis(address_analysis=address_analysis_frame.entries))
+            .rewrite(mt.code)
+            .join(rewrite_result)
+        )
+
+        # 2. rewrite
+        rewrite_result = (
+            Walk(
+                IfToStim(
+                    measure_analysis=meas_analysis_frame.entries,
+                    measure_count=mia.measure_count,
+                )
+            )
+            .rewrite(mt.code)
+            .join(rewrite_result)
+        )
+
+        # 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(
                     SquinQubitToStim(),
                     SquinWireToStim(),
-                    SquinMeasureToStim(),  # reduce duplicated logic, can split out even more rules later
+                    SquinMeasureToStim(
+                        measure_id_result=meas_analysis_frame.entries,
+                        total_measure_count=mia.measure_count,
+                    ),  # reduce duplicated logic, can split out even more rules later
                     SquinWireIdentityElimination(),
                 )
             )
             .rewrite(mt.code)
             .join(rewrite_result)
         )
-
-        # Convert all PyConsts to Stim Constants
         rewrite_result = (
-            Walk(Chain(PyConstantToStim())).rewrite(mt.code).join(rewrite_result)
+            CanonicalizeIList(dialects=mt.dialects, no_raise=self.no_raise)
+            .unsafe_run(mt)
+            .join(rewrite_result)
         )
 
-        # remove any squin.qubit.new that's left around
-        ## Not considered pure so DCE won't touch it but
-        ## it isn't being used anymore considering everything is a
-        ## stim dialect statement
+        # Convert all PyConsts to Stim Constants
+        rewrite_result = Walk(PyConstantToStim()).rewrite(mt.code).join(rewrite_result)
+
+        # clear up leftover stmts
+        # - remove any squin.qubit.new that's left around
         rewrite_result = (
             Fixpoint(
                 Walk(
@@ -71,16 +166,4 @@ class SquinToStim(Pass):
             .join(rewrite_result)
         )
 
-        # do program verification here,
-        # ideally use built-in .verify() to catch any
-        # incompatible statements as the full rewrite process should not
-        # leave statements from any other dialects (other than the stim main group)
-        mt_verification_clone = mt.similar(stim_main_group)
-
-        # suggested by Kai, will work for now
-        for stmt in mt_verification_clone.code.walk():
-            assert (
-                stmt.dialect in stim_main_group
-            ), "Statements detected that are not part of the stim dialect, please verify the original code is valid for rewrite!"
-
         return rewrite_result

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