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

bloqade/squin/noise/stmts.py

@@ -10,7 +10,7 @@ from ..op.types import NumOperators
 
 @statement
 class NoiseChannel(ir.Statement):
-    traits = frozenset({lowering.FromPythonCall()})
+    traits = frozenset({ir.Pure(), lowering.FromPythonCall()})
     result: ir.ResultValue = info.result(OpType)
 
 

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/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,15 +5,19 @@ 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,
@@ -22,22 +26,95 @@ from bloqade.stim.rewrite import (
     SquinMeasureToStim,
     SquinWireIdentityElimination,
 )
-from bloqade.squin.rewrite import SquinU3ToClifford, 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 = (
@@ -47,7 +124,6 @@ class SquinToStim(Pass):
         )
 
         # Wrap Rewrite + SquinToStim can happen w/ standard walk
-
         rewrite_result = Walk(SquinU3ToClifford()).rewrite(mt.code).join(rewrite_result)
 
         rewrite_result = (
@@ -55,23 +131,27 @@ class SquinToStim(Pass):
                 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(
@@ -86,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/ifs_to_stim.py

@@ -0,0 +1,203 @@
+from dataclasses import field, dataclass
+
+from kirin import ir
+from kirin.dialects import py, scf, func
+from kirin.rewrite.abc import RewriteRule, RewriteResult
+
+from bloqade.squin import op, qubit
+from bloqade.rewrite.rules import LiftThenBody, SplitIfStmts
+from bloqade.squin.rewrite import AddressAttribute
+from bloqade.stim.rewrite.util import (
+    SQUIN_STIM_CONTROL_GATE_MAPPING,
+    insert_qubit_idx_from_address,
+)
+from bloqade.stim.dialects.auxiliary import GetRecord
+from bloqade.analysis.measure_id.lattice import (
+    MeasureId,
+    MeasureIdBool,
+)
+
+
+@dataclass
+class IfElseSimplification:
+
+    # Might be better to just do a rewrite_Region?
+    def is_rewriteable(self, node: scf.IfElse) -> bool:
+        return not (
+            self.contains_ifelse(node)
+            or self.is_nested_ifelse(node)
+            or self.has_else_body(node)
+        )
+
+    # A preliminary check to reject an IfElse from the "top down"
+    # use in conjunction with is_nested_ifelse
+    # to completely cover cases of nested IfElse statements
+    def contains_ifelse(self, stmt: scf.IfElse) -> bool:
+        """Check if the IfElse statement contains another IfElse statement."""
+        for child in stmt.walk(include_self=False):
+            if isinstance(child, scf.IfElse):
+                return True
+        return False
+
+    # because rewrite latches onto ANY scf.IfElse,
+    # you need a way to determine if you're touching an
+    # IfElse that's inside another IfElse
+    def is_nested_ifelse(self, stmt: scf.IfElse) -> bool:
+        """Check if the IfElse statement is nested."""
+        if stmt.parent_stmt is not None:
+            if isinstance(stmt.parent_stmt, scf.IfElse) or isinstance(
+                stmt.parent_stmt.parent_stmt, scf.IfElse
+            ):
+                return True
+            else:
+                return False
+        else:
+            return False
+
+    def has_else_body(self, stmt: scf.IfElse) -> bool:
+        """Check if the IfElse statement has an else body."""
+        if stmt.else_body.blocks and not (
+            len(stmt.else_body.blocks[0].stmts) == 1
+            and isinstance(else_term := stmt.else_body.blocks[0].last_stmt, scf.Yield)
+            and not else_term.values  # empty yield
+        ):
+            return True
+
+        return False
+
+
+DontLiftType = (
+    qubit.Apply,
+    qubit.Broadcast,
+    scf.Yield,
+    func.Return,
+    func.Invoke,
+    scf.IfElse,
+)
+
+
+@dataclass
+class StimLiftThenBody(IfElseSimplification, LiftThenBody):
+    exclude_stmts: tuple[type[ir.Statement], ...] = field(default=DontLiftType)
+
+    def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
+
+        if not isinstance(node, scf.IfElse):
+            return RewriteResult()
+
+        if not self.is_rewriteable(node):
+            return RewriteResult()
+
+        return super().rewrite_Statement(node)
+
+
+# Only run this after everything other than qubit.Apply/qubit.Broadcast has been
+# lifted out!
+class StimSplitIfStmts(IfElseSimplification, SplitIfStmts):
+    """Splits the then body of an if-else statement into multiple if statements
+
+    Given an IfElse with multiple valid statements in the then-body:
+
+    if measure_result:
+        squin.qubit.apply(op.X, q0)
+        squin.qubit.apply(op.Y, q1)
+
+    this should be rewritten to:
+
+    if measure_result:
+        squin.qubit.apply(op.X, q0)
+
+    if measure_result:
+        squin.qubit.apply(op.Y, q1)
+    """
+
+    def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
+        if not isinstance(node, scf.IfElse):
+            return RewriteResult()
+
+        if not self.is_rewriteable(node):
+            return RewriteResult()
+
+        return super().rewrite_Statement(node)
+
+
+@dataclass
+class IfToStim(IfElseSimplification, RewriteRule):
+    """
+    Rewrite if statements to stim equivalent statements.
+    """
+
+    measure_analysis: dict[ir.SSAValue, MeasureId]
+    measure_count: int
+
+    def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
+
+        match node:
+            case scf.IfElse():
+                return self.rewrite_IfElse(node)
+            case _:
+                return RewriteResult()
+
+    def rewrite_IfElse(self, stmt: scf.IfElse) -> RewriteResult:
+
+        if not isinstance(self.measure_analysis[stmt.cond], MeasureIdBool):
+            return RewriteResult()
+
+        # check that there is only qubit.Apply in the then-body,
+        # if there's more than that, we can't do a valid rewrite.
+        # Can reuse logic from SplitIf
+        *stmts, _ = stmt.then_body.stmts()
+        if len(stmts) != 1 or not isinstance(stmts[0], (qubit.Apply, qubit.Broadcast)):
+            return RewriteResult()
+
+        apply_or_broadcast = stmts[0]
+        # Check that the gate being applied/broadcasted can be converted to a stim
+        # controlled gate.
+        ctrl_op_target_gate = apply_or_broadcast.operator.owner
+        assert isinstance(ctrl_op_target_gate, op.stmts.Operator)
+
+        stim_gate = SQUIN_STIM_CONTROL_GATE_MAPPING.get(type(ctrl_op_target_gate))
+        if stim_gate is None:
+            return RewriteResult()
+
+        # get necessary measurement ID type from analysis
+        measure_id_bool = self.measure_analysis[stmt.cond]
+        assert isinstance(measure_id_bool, MeasureIdBool)
+
+        # generate get record statement
+        measure_id_idx_stmt = py.Constant(
+            (measure_id_bool.idx - 1) - self.measure_count
+        )
+        get_record_stmt = GetRecord(id=measure_id_idx_stmt.result)  # noqa: F841
+
+        # get address attribute and generate qubit idx statements
+        address_attr = apply_or_broadcast.qubits.hints.get("address")
+        if address_attr is None:
+            return RewriteResult()
+        assert isinstance(address_attr, AddressAttribute)
+
+        # note: insert things before (literally above/outside) the If
+        qubit_idx_ssas = insert_qubit_idx_from_address(
+            address=address_attr, stmt_to_insert_before=stmt
+        )
+        if qubit_idx_ssas is None:
+            return RewriteResult()
+
+        # Assemble the stim statement
+        # let GetRecord's SSA be repeated per each get qubit
+        ctrl_records = tuple(get_record_stmt.result for _ in qubit_idx_ssas)
+
+        stim_stmt = stim_gate(
+            targets=tuple(qubit_idx_ssas),
+            controls=ctrl_records,
+        )
+
+        # Insert the necessary SSA Values, then get rid of the scf.IfElse.
+        # The qubit indices have been successfully added,
+        # that just leaves the GetRecord statement and measurement ID index statement
+
+        measure_id_idx_stmt.insert_before(stmt)
+        get_record_stmt.insert_before(stmt)
+        stmt.replace_by(stim_stmt)
+
+        return RewriteResult(has_done_something=True)

bloqade/stim/rewrite/qubit_to_stim.py

@@ -13,6 +13,9 @@ from bloqade.stim.rewrite.util import (
 
 
 class SquinQubitToStim(RewriteRule):
+    """
+    NOTE this require address analysis result to be wrapped before using this rule.
+    """
 
     def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
 

bloqade/stim/rewrite/squin_measure.py

@@ -1,22 +1,59 @@
 # create rewrite rule name SquinMeasureToStim using kirin
+from dataclasses import dataclass
+
 from kirin import ir
-from kirin.dialects import py
+from kirin.dialects import py, ilist
 from kirin.rewrite.abc import RewriteRule, RewriteResult
 
 from bloqade.squin import wire, qubit
 from bloqade.squin.rewrite import AddressAttribute
-from bloqade.stim.dialects import collapse
+from bloqade.stim.dialects import collapse, auxiliary
 from bloqade.stim.rewrite.util import (
     is_measure_result_used,
     insert_qubit_idx_from_address,
 )
+from bloqade.analysis.measure_id.lattice import MeasureId, MeasureIdBool, MeasureIdTuple
+
+
+def replace_get_record(
+    node: ir.Statement, measure_id_bool: MeasureIdBool, meas_count: int
+):
+    assert isinstance(measure_id_bool, MeasureIdBool)
+    target_rec_idx = (measure_id_bool.idx - 1) - meas_count
+    idx_stmt = py.constant.Constant(target_rec_idx)
+    idx_stmt.insert_before(node)
+    get_record_stmt = auxiliary.GetRecord(idx_stmt.result)
+    node.replace_by(get_record_stmt)
+
+
+def insert_get_record_list(
+    node: ir.Statement, measure_id_tuple: MeasureIdTuple, meas_count: int
+):
+    """
+    Insert GetRecord statements before the given node
+    """
+    get_record_ssas = []
+    for measure_id_bool in measure_id_tuple.data:
+        assert isinstance(measure_id_bool, MeasureIdBool)
+        target_rec_idx = (measure_id_bool.idx - 1) - meas_count
+        idx_stmt = py.constant.Constant(target_rec_idx)
+        idx_stmt.insert_before(node)
+        get_record_stmt = auxiliary.GetRecord(idx_stmt.result)
+        get_record_stmt.insert_before(node)
+        get_record_ssas.append(get_record_stmt.result)
 
+    node.replace_by(ilist.New(values=get_record_ssas))
 
+
+@dataclass
 class SquinMeasureToStim(RewriteRule):
     """
     Rewrite squin measure-related statements to stim statements.
     """
 
+    measure_id_result: dict[ir.SSAValue, MeasureId]
+    total_measure_count: int
+
     def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
 
         match node:
@@ -28,20 +65,46 @@ class SquinMeasureToStim(RewriteRule):
     def rewrite_Measure(
         self, measure_stmt: qubit.MeasureQubit | qubit.MeasureQubitList | wire.Measure
     ) -> RewriteResult:
-        if is_measure_result_used(measure_stmt):
-            return RewriteResult()
 
         qubit_idx_ssas = self.get_qubit_idx_ssas(measure_stmt)
         if qubit_idx_ssas is None:
             return RewriteResult()
 
+        measure_id = self.measure_id_result[measure_stmt.result]
+        if not isinstance(measure_id, (MeasureIdBool, MeasureIdTuple)):
+            return RewriteResult()
+
         prob_noise_stmt = py.constant.Constant(0.0)
         stim_measure_stmt = collapse.MZ(
             p=prob_noise_stmt.result,
             targets=qubit_idx_ssas,
         )
         prob_noise_stmt.insert_before(measure_stmt)
-        measure_stmt.replace_by(stim_measure_stmt)
+        stim_measure_stmt.insert_before(measure_stmt)
+
+        if not is_measure_result_used(measure_stmt):
+            measure_stmt.delete()
+            return RewriteResult(has_done_something=True)
+
+        # replace dataflow with new stmt!
+        measure_id = self.measure_id_result[measure_stmt.result]
+        if isinstance(measure_id, MeasureIdBool):
+            replace_get_record(
+                node=measure_stmt,
+                measure_id_bool=measure_id,
+                meas_count=self.total_measure_count,
+            )
+        elif isinstance(measure_id, MeasureIdTuple):
+            insert_get_record_list(
+                node=measure_stmt,
+                measure_id_tuple=measure_id,
+                meas_count=self.total_measure_count,
+            )
+        else:
+            # already checked before, so this should not happen
+            raise ValueError(
+                f"Unexpected measure ID type: {type(measure_id)} for measure statement {measure_stmt}"
+            )
 
         return RewriteResult(has_done_something=True)
 

bloqade/stim/rewrite/util.py

@@ -182,10 +182,6 @@ def rewrite_QubitLoss(
         create_wire_passthrough(stmt)
 
     stmt.replace_by(stim_loss_stmt)
-    # NoiseChannels are not pure,
-    # need to manually delete because
-    # DCE won't touch them
-    stmt.operator.owner.delete()
 
     return RewriteResult(has_done_something=True)
 

bloqade/stim/upstream/__init__.py

@@ -0,0 +1 @@
+from .from_squin import squin_to_stim as squin_to_stim

bloqade/stim/upstream/from_squin.py

@@ -0,0 +1,10 @@
+from kirin import ir
+
+from ..groups import main
+from ..passes.squin_to_stim import SquinToStimPass
+
+
+def squin_to_stim(mt: ir.Method) -> ir.Method:
+    new_mt = mt.similar()
+    SquinToStimPass(mt.dialects, no_raise=False)(new_mt)
+    return new_mt.similar(dialects=main)

{bloqade_circuit-0.5.0.dist-info → bloqade_circuit-0.5.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: bloqade-circuit
-Version: 0.5.0
+Version: 0.5.2
 Summary: The software development toolkit for neutral atom arrays.
 Author-email: Roger-luo <rluo@quera.com>, kaihsin <khwu@quera.com>, weinbe58 <pweinberg@quera.com>, johnzl-777 <jlong@quera.com>
 License-File: LICENSE