bloqade-circuit 0.6.3__py3-none-any.whl → 0.6.5__py3-none-any.whl

bloqade/analysis/measure_id/__init__.py

@@ -1,2 +1,5 @@
 from . import impls as impls
-from .analysis import MeasurementIDAnalysis as MeasurementIDAnalysis
+from .analysis import (
+    MeasureIDFrame as MeasureIDFrame,
+    MeasurementIDAnalysis as MeasurementIDAnalysis,
+)

bloqade/analysis/measure_id/analysis.py

@@ -1,13 +1,19 @@
 from typing import TypeVar
+from dataclasses import field, dataclass
 
 from kirin import ir, interp
-from kirin.analysis import Forward, const
+from kirin.analysis import ForwardExtra, const
 from kirin.analysis.forward import ForwardFrame
 
 from .lattice import MeasureId, NotMeasureId
 
 
-class MeasurementIDAnalysis(Forward[MeasureId]):
+@dataclass
+class MeasureIDFrame(ForwardFrame[MeasureId]):
+    num_measures_at_stmt: dict[ir.Statement, int] = field(default_factory=dict)
+
+
+class MeasurementIDAnalysis(ForwardExtra[MeasureIDFrame, MeasureId]):
 
     keys = ["measure_id"]
     lattice = MeasureId
@@ -15,6 +21,11 @@ class MeasurementIDAnalysis(Forward[MeasureId]):
     # then use this to generate the negative values for target rec indices
     measure_count = 0
 
+    def initialize_frame(
+        self, code: ir.Statement, *, has_parent_access: bool = False
+    ) -> MeasureIDFrame:
+        return MeasureIDFrame(code, has_parent_access=has_parent_access)
+
     # Still default to bottom,
     # but let constants return the softer "NoMeasureId" type from impl
     def eval_stmt_fallback(

bloqade/analysis/measure_id/impls.py

@@ -1,4 +1,5 @@
 from kirin import types as kirin_types, interp
+from kirin.analysis import const
 from kirin.dialects import py, scf, func, ilist
 
 from bloqade.squin import wire, qubit
@@ -10,7 +11,7 @@ from .lattice import (
     MeasureIdTuple,
     InvalidMeasureId,
 )
-from .analysis import MeasurementIDAnalysis
+from .analysis import MeasureIDFrame, MeasurementIDAnalysis
 
 ## Can't do wire right now because of
 ## unresolved RFC on return type
@@ -113,6 +114,15 @@ class PyIndexing(interp.MethodTable):
             return (InvalidMeasureId(),)
 
 
+@py.assign.dialect.register(key="measure_id")
+class PyAssign(interp.MethodTable):
+    @interp.impl(py.Alias)
+    def alias(
+        self, interp: MeasurementIDAnalysis, frame: interp.Frame, stmt: py.assign.Alias
+    ):
+        return (frame.get(stmt.value),)
+
+
 @py.binop.dialect.register(key="measure_id")
 class PyBinOp(interp.MethodTable):
     @interp.impl(py.Add)
@@ -152,4 +162,33 @@ class Func(interp.MethodTable):
 # scf, particularly IfElse
 @scf.dialect.register(key="measure_id")
 class Scf(scf.absint.Methods):
-    pass
+
+    @interp.impl(scf.IfElse)
+    def if_else(
+        self,
+        interp_: MeasurementIDAnalysis,
+        frame: MeasureIDFrame,
+        stmt: scf.IfElse,
+    ):
+
+        frame.num_measures_at_stmt[stmt] = interp_.measure_count
+
+        # rest of the code taken directly from scf.absint.Methods base implementation
+
+        if isinstance(hint := stmt.cond.hints.get("const"), const.Value):
+            if hint.data:
+                return self._infer_if_else_cond(interp_, frame, stmt, stmt.then_body)
+            else:
+                return self._infer_if_else_cond(interp_, frame, stmt, stmt.else_body)
+        then_results = self._infer_if_else_cond(interp_, frame, stmt, stmt.then_body)
+        else_results = self._infer_if_else_cond(interp_, frame, stmt, stmt.else_body)
+
+        match (then_results, else_results):
+            case (interp.ReturnValue(then_value), interp.ReturnValue(else_value)):
+                return interp.ReturnValue(then_value.join(else_value))
+            case (interp.ReturnValue(then_value), _):
+                return then_results
+            case (_, interp.ReturnValue(else_value)):
+                return else_results
+            case _:
+                return interp_.join_results(then_results, else_results)

bloqade/pyqrack/base.py

@@ -48,7 +48,7 @@ def _default_pyqrack_args() -> PyQrackOptions:
         isSchmidtDecomposeMulti=True,
         isSchmidtDecompose=True,
         isStabilizerHybrid=False,
-        isBinaryDecisionTree=True,
+        isBinaryDecisionTree=False,
         isPaged=True,
         isCpuGpuHybrid=True,
         isOpenCL=True,

bloqade/squin/__init__.py

@@ -9,6 +9,10 @@ from . import (
 )
 from .groups import wired as wired, kernel as kernel
 
+# NOTE: it's important to keep these imports here since they import squin.kernel
+# we skip isort here
+from . import gate as gate, parallel as parallel  # isort: skip
+
 try:
     # NOTE: make sure optional cirq dependency is installed
     import cirq as cirq_package  # noqa: F401

bloqade/squin/gate.py

@@ -0,0 +1,193 @@
+from bloqade.types import Qubit
+
+from . import op as _op, qubit as _qubit
+from .groups import kernel
+
+
+@kernel
+def x(qubit: Qubit) -> None:
+    """x gate applied to qubit."""
+    op = _op.x()
+    _qubit.apply(op, qubit)
+
+
+@kernel
+def y(qubit: Qubit) -> None:
+    """y gate applied to qubit."""
+    op = _op.y()
+    _qubit.apply(op, qubit)
+
+
+@kernel
+def z(qubit: Qubit) -> None:
+    """z gate applied to qubit."""
+    op = _op.z()
+    _qubit.apply(op, qubit)
+
+
+@kernel
+def sqrt_x(qubit: Qubit) -> None:
+    """Square root x gate applied to qubit."""
+    op = _op.sqrt_x()
+    _qubit.apply(op, qubit)
+
+
+@kernel
+def sqrt_x_adj(qubit: Qubit) -> None:
+    """Adjoint sqrt_x gate applied to qubit."""
+    op = _op.sqrt_x()
+    _qubit.apply(_op.adjoint(op), qubit)
+
+
+@kernel
+def sqrt_y(qubit: Qubit) -> None:
+    """Square root y gate applied to qubit."""
+    op = _op.sqrt_y()
+    _qubit.apply(op, qubit)
+
+
+@kernel
+def sqrt_y_adj(qubit: Qubit) -> None:
+    """Adjoint sqrt_y gate applied to qubit."""
+    op = _op.sqrt_y()
+    _qubit.apply(_op.adjoint(op), qubit)
+
+
+@kernel
+def sqrt_z(qubit: Qubit) -> None:
+    """Square root z gate applied to qubit."""
+    op = _op.s()
+    _qubit.apply(op, qubit)
+
+
+@kernel
+def sqrt_z_adj(qubit: Qubit) -> None:
+    """Adjoint square root z gate applied to qubit."""
+    op = _op.s()
+    _qubit.apply(_op.adjoint(op), qubit)
+
+
+@kernel
+def h(qubit: Qubit) -> None:
+    """Hadamard gate applied to qubit."""
+    op = _op.h()
+    _qubit.apply(op, qubit)
+
+
+@kernel
+def s(qubit: Qubit) -> None:
+    """s gate applied to qubit."""
+    op = _op.s()
+    _qubit.apply(op, qubit)
+
+
+@kernel
+def s_adj(qubit: Qubit) -> None:
+    """Adjoint s gate applied to qubit."""
+    op = _op.s()
+    _qubit.apply(_op.adjoint(op), qubit)
+
+
+@kernel
+def t(qubit: Qubit) -> None:
+    """t gate applied to qubit."""
+    op = _op.t()
+    _qubit.apply(op, qubit)
+
+
+@kernel
+def t_adj(qubit: Qubit) -> None:
+    """Adjoint t gate applied to qubit."""
+    op = _op.t()
+    _qubit.apply(_op.adjoint(op), qubit)
+
+
+@kernel
+def p0(qubit: Qubit) -> None:
+    """Projector on 0 applied to qubit."""
+    op = _op.p0()
+    _qubit.apply(op, qubit)
+
+
+@kernel
+def p1(qubit: Qubit) -> None:
+    """Projector on 1 applied to qubit."""
+    op = _op.p1()
+    _qubit.apply(op, qubit)
+
+
+@kernel
+def spin_n(qubit: Qubit) -> None:
+    """Spin lowering gate applied to qubit."""
+    op = _op.spin_n()
+    _qubit.apply(op, qubit)
+
+
+@kernel
+def spin_p(qubit: Qubit) -> None:
+    """Spin raising gate applied to qubit."""
+    op = _op.spin_p()
+    _qubit.apply(op, qubit)
+
+
+@kernel
+def reset(qubit: Qubit) -> None:
+    """Reset qubit to 0."""
+    op = _op.reset()
+    _qubit.apply(op, qubit)
+
+
+@kernel
+def cx(control: Qubit, target: Qubit) -> None:
+    """Controlled x gate applied to control and target"""
+    op = _op.cx()
+    _qubit.apply(op, control, target)
+
+
+@kernel
+def cy(control: Qubit, target: Qubit) -> None:
+    """Controlled y gate applied to control and target"""
+    op = _op.cy()
+    _qubit.apply(op, control, target)
+
+
+@kernel
+def cz(control: Qubit, target: Qubit) -> None:
+    """Controlled z gate applied to control and target"""
+    op = _op.cz()
+    _qubit.apply(op, control, target)
+
+
+@kernel
+def ch(control: Qubit, target: Qubit) -> None:
+    """Controlled Hadamard gate applied to control and target"""
+    op = _op.ch()
+    _qubit.apply(op, control, target)
+
+
+@kernel
+def u(theta: float, phi: float, lam: float, qubit: Qubit) -> None:
+    """3D rotation gate applied to control and target"""
+    op = _op.u(theta, phi, lam)
+    _qubit.apply(op, qubit)
+
+
+@kernel
+def rx(theta: float, qubit: Qubit) -> None:
+    """Rotation X gate applied to qubit."""
+    op = _op.rot(_op.x(), theta)
+    _qubit.apply(op, qubit)
+
+
+@kernel
+def ry(theta: float, qubit: Qubit) -> None:
+    """Rotation Y gate applied to qubit."""
+    op = _op.rot(_op.y(), theta)
+    _qubit.apply(op, qubit)
+
+
+@kernel
+def rz(theta: float, qubit: Qubit) -> None:
+    """Rotation Z gate applied to qubit."""
+    op = _op.rot(_op.z(), theta)
+    _qubit.apply(op, qubit)

bloqade/squin/parallel.py

@@ -0,0 +1,200 @@
+from typing import Any, TypeVar
+
+from kirin.dialects import ilist
+
+from bloqade.types import Qubit
+
+from . import op as _op, qubit as _qubit
+from .groups import kernel
+
+
+@kernel
+def x(qubits: ilist.IList[Qubit, Any]) -> None:
+    """x gate applied to qubits in parallel."""
+    op = _op.x()
+    _qubit.broadcast(op, qubits)
+
+
+@kernel
+def y(qubits: ilist.IList[Qubit, Any]) -> None:
+    """y gate applied to qubits in parallel."""
+    op = _op.y()
+    _qubit.broadcast(op, qubits)
+
+
+@kernel
+def z(qubits: ilist.IList[Qubit, Any]) -> None:
+    """z gate applied to qubits in parallel."""
+    op = _op.z()
+    _qubit.broadcast(op, qubits)
+
+
+@kernel
+def sqrt_x(qubits: ilist.IList[Qubit, Any]) -> None:
+    """Square root x gate applied to qubits in parallel."""
+    op = _op.sqrt_x()
+    _qubit.broadcast(op, qubits)
+
+
+@kernel
+def sqrt_y(qubits: ilist.IList[Qubit, Any]) -> None:
+    """Square root y gate applied to qubits in parallel."""
+    op = _op.sqrt_y()
+    _qubit.broadcast(op, qubits)
+
+
+@kernel
+def sqrt_z(qubits: ilist.IList[Qubit, Any]) -> None:
+    """Square root gate applied to qubits in parallel."""
+    op = _op.s()
+    _qubit.broadcast(op, qubits)
+
+
+@kernel
+def h(qubits: ilist.IList[Qubit, Any]) -> None:
+    """Hadamard gate applied to qubits in parallel."""
+    op = _op.h()
+    _qubit.broadcast(op, qubits)
+
+
+@kernel
+def s(qubits: ilist.IList[Qubit, Any]) -> None:
+    """s gate applied to qubits in parallel."""
+    op = _op.s()
+    _qubit.broadcast(op, qubits)
+
+
+@kernel
+def t(qubits: ilist.IList[Qubit, Any]) -> None:
+    """t gate applied to qubits in parallel."""
+    op = _op.t()
+    _qubit.broadcast(op, qubits)
+
+
+@kernel
+def p0(qubits: ilist.IList[Qubit, Any]) -> None:
+    """Projector on 0 applied to qubits in parallel."""
+    op = _op.p0()
+    _qubit.broadcast(op, qubits)
+
+
+@kernel
+def p1(qubits: ilist.IList[Qubit, Any]) -> None:
+    """Projector on 1 applied to qubits in parallel."""
+    op = _op.p1()
+    _qubit.broadcast(op, qubits)
+
+
+@kernel
+def spin_n(qubits: ilist.IList[Qubit, Any]) -> None:
+    """Spin lowering gate applied to qubits in parallel."""
+    op = _op.spin_n()
+    _qubit.broadcast(op, qubits)
+
+
+@kernel
+def spin_p(qubits: ilist.IList[Qubit, Any]) -> None:
+    """Spin raising gate applied to qubits in parallel."""
+    op = _op.spin_p()
+    _qubit.broadcast(op, qubits)
+
+
+@kernel
+def reset(qubits: ilist.IList[Qubit, Any]) -> None:
+    """Reset qubit to 0."""
+    op = _op.reset()
+    _qubit.broadcast(op, qubits)
+
+
+N = TypeVar("N")
+
+
+@kernel
+def cx(controls: ilist.IList[Qubit, N], targets: ilist.IList[Qubit, N]) -> None:
+    """Controlled x gate applied to controls and targets in parallel."""
+    op = _op.cx()
+    _qubit.broadcast(op, controls, targets)
+
+
+@kernel
+def cy(controls: ilist.IList[Qubit, N], targets: ilist.IList[Qubit, N]) -> None:
+    """Controlled y gate applied to controls and targets in parallel."""
+    op = _op.cy()
+    _qubit.broadcast(op, controls, targets)
+
+
+@kernel
+def cz(controls: ilist.IList[Qubit, N], targets: ilist.IList[Qubit, N]) -> None:
+    """Controlled z gate applied to controls and targets in parallel."""
+    op = _op.cz()
+    _qubit.broadcast(op, controls, targets)
+
+
+@kernel
+def ch(controls: ilist.IList[Qubit, N], targets: ilist.IList[Qubit, N]) -> None:
+    """Controlled Hadamard gate applied to controls and targets in parallel."""
+    op = _op.ch()
+    _qubit.broadcast(op, controls, targets)
+
+
+@kernel
+def u(theta: float, phi: float, lam: float, qubits: ilist.IList[Qubit, Any]) -> None:
+    """3D rotation gate applied to controls and targets in parallel."""
+    op = _op.u(theta, phi, lam)
+    _qubit.broadcast(op, qubits)
+
+
+@kernel
+def rx(theta: float, qubits: ilist.IList[Qubit, Any]) -> None:
+    """Rotation X gate applied to qubits in parallel."""
+    op = _op.rot(_op.x(), theta)
+    _qubit.broadcast(op, qubits)
+
+
+@kernel
+def ry(theta: float, qubits: ilist.IList[Qubit, Any]) -> None:
+    """Rotation Y gate applied to qubits in parallel."""
+    op = _op.rot(_op.y(), theta)
+    _qubit.broadcast(op, qubits)
+
+
+@kernel
+def rz(theta: float, qubits: ilist.IList[Qubit, Any]) -> None:
+    """Rotation Z gate applied to qubits in parallel."""
+    op = _op.rot(_op.z(), theta)
+    _qubit.broadcast(op, qubits)
+
+
+@kernel
+def sqrt_x_adj(qubits: ilist.IList[Qubit, Any]) -> None:
+    """Adjoint sqrt_x gate applied to qubits in parallel."""
+    op = _op.sqrt_x()
+    _qubit.broadcast(_op.adjoint(op), qubits)
+
+
+@kernel
+def sqrt_y_adj(qubits: ilist.IList[Qubit, Any]) -> None:
+    """Adjoint sqrt_y gate applied to qubits in parallel."""
+    op = _op.sqrt_y()
+    _qubit.broadcast(_op.adjoint(op), qubits)
+
+
+@kernel
+def sqrt_z_adj(qubits: ilist.IList[Qubit, Any]) -> None:
+    """Adjoint square root z gate applied to qubits in parallel."""
+    op = _op.s()
+    _qubit.broadcast(_op.adjoint(op), qubits)
+
+
+@kernel
+def s_adj(qubits: ilist.IList[Qubit, Any]) -> None:
+    """Adjoint s gate applied to qubits in parallel."""
+    op = _op.s()
+    _qubit.broadcast(_op.adjoint(op), qubits)
+
+
+@kernel
+def t_adj(qubits: ilist.IList[Qubit, Any]) -> None:
+    """Adjoint t gate applied to qubits in parallel."""
+    op = _op.t()
+    _qubit.broadcast(_op.adjoint(op), qubits)

bloqade/stim/passes/__init__.py

@@ -1 +1,5 @@
-from .squin_to_stim import SquinToStimPass as SquinToStimPass
+from .squin_to_stim import (
+    SquinToStimPass as SquinToStimPass,
+    StimSimplifyIfs as StimSimplifyIfs,
+    AggressiveForLoopUnroll as AggressiveForLoopUnroll,
+)

bloqade/stim/passes/simplify_ifs.py

@@ -9,6 +9,7 @@ from kirin.rewrite import (
     ConstantFold,
     CommonSubexpressionElimination,
 )
+from kirin.dialects.ilist.passes import ConstList2IList
 
 from ..rewrite.ifs_to_stim import StimLiftThenBody, StimSplitIfStmts
 
@@ -23,8 +24,16 @@ class StimSimplifyIfs(Pass):
             Walk(StimSplitIfStmts()),
         ).rewrite(mt.code)
 
+        # because nested python lists don't have their
+        # member lists converted to ILists, ConstantFold
+        # can add python lists that can't be hashed, causing
+        # issues with CSE. ConstList2IList remedies that problem here.
         result = (
-            Fixpoint(Walk(Chain(ConstantFold(), CommonSubexpressionElimination())))
+            Chain(
+                Fixpoint(Walk(ConstantFold())),
+                Walk(ConstList2IList()),
+                Walk(CommonSubexpressionElimination()),
+            )
             .rewrite(mt.code)
             .join(result)
         )

bloqade/stim/passes/squin_to_stim.py

@@ -1,6 +1,6 @@
 from dataclasses import dataclass
 
-from kirin.passes import Fold
+from kirin.passes import Fold, HintConst, TypeInfer
 from kirin.rewrite import (
     Walk,
     Chain,
@@ -16,6 +16,7 @@ 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 kirin.rewrite.alias import InlineAlias
 
 from bloqade.stim.rewrite import (
     SquinWireToStim,
@@ -33,11 +34,43 @@ from bloqade.squin.rewrite import (
 from bloqade.rewrite.passes import CanonicalizeIList
 from bloqade.analysis.address import AddressAnalysis
 from bloqade.analysis.measure_id import MeasurementIDAnalysis
+from bloqade.squin.rewrite.desugar import ApplyDesugarRule
 
 from .simplify_ifs import StimSimplifyIfs
 from ..rewrite.ifs_to_stim import IfToStim
 
 
+@dataclass
+class AggressiveForLoopUnroll(Pass):
+    """
+    Aggressive unrolling of for loops, addresses cases where unroll
+    does not successfully handle nested loops because of a lack of constprop.
+
+    This should be invoked via fixpoint to let this be repeatedly applied until
+    no further rewrites are possible.
+    """
+
+    def unsafe_run(self, mt: Method) -> RewriteResult:
+        rule = Chain(
+            InlineGetField(),
+            InlineGetItem(),
+            scf.unroll.ForLoop(),
+            scf.trim.UnusedYield(),
+        )
+
+        # Intentionally only walk ONCE, let fixpoint happen with the WHOLE pass
+        # so that HintConst gets run right after, allowing subsequent unrolls to happen
+        rewrite_result = Walk(rule).rewrite(mt.code)
+
+        rewrite_result = (
+            HintConst(dialects=mt.dialects, no_raise=self.no_raise)
+            .unsafe_run(mt)
+            .join(rewrite_result)
+        )
+
+        return rewrite_result
+
+
 @dataclass
 class SquinToStimPass(Pass):
 
@@ -48,27 +81,24 @@ class SquinToStimPass(Pass):
             dialects=mt.dialects, no_raise=self.no_raise
         ).unsafe_run(mt)
 
-        rule = Chain(
-            InlineGetField(),
-            InlineGetItem(),
-            scf.unroll.ForLoop(),
-            scf.trim.UnusedYield(),
+        rewrite_result = (
+            AggressiveForLoopUnroll(dialects=mt.dialects, no_raise=self.no_raise)
+            .fixpoint(mt)
+            .join(rewrite_result)
         )
-        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)
         )
+
+        Walk(InlineAlias()).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)
@@ -82,6 +112,9 @@ class SquinToStimPass(Pass):
             .join(rewrite_result)
         )
 
+        TypeInfer(dialects=mt.dialects, no_raise=self.no_raise).unsafe_run(mt)
+        Walk(ApplyDesugarRule()).rewrite(mt.code)
+
         # after this the program should be in a state where it is analyzable
         # -------------------------------------------------------------------
 
@@ -99,12 +132,14 @@ class SquinToStimPass(Pass):
         )
 
         # 2. rewrite
+        ## Invoke DCE afterwards to eliminate any GetItems
+        ## that are no longer being used. This allows for
+        ## SquinMeasureToStim to safely eliminate
+        ## unused measure statements.
         rewrite_result = (
-            Walk(
-                IfToStim(
-                    measure_analysis=meas_analysis_frame.entries,
-                    measure_count=mia.measure_count,
-                )
+            Chain(
+                Walk(IfToStim(measure_frame=meas_analysis_frame)),
+                Fixpoint(Walk(DeadCodeElimination())),
             )
             .rewrite(mt.code)
             .join(rewrite_result)
@@ -120,17 +155,15 @@ class SquinToStimPass(Pass):
             Walk(
                 Chain(
                     SquinQubitToStim(),
+                    SquinMeasureToStim(),
                     SquinWireToStim(),
-                    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)
         )
+
         rewrite_result = (
             CanonicalizeIList(dialects=mt.dialects, no_raise=self.no_raise)
             .unsafe_run(mt)

bloqade/stim/rewrite/__init__.py

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

bloqade/stim/rewrite/ifs_to_stim.py

@@ -11,9 +11,9 @@ from bloqade.stim.rewrite.util import (
     SQUIN_STIM_CONTROL_GATE_MAPPING,
     insert_qubit_idx_from_address,
 )
+from bloqade.analysis.measure_id import MeasureIDFrame
 from bloqade.stim.dialects.auxiliary import GetRecord
 from bloqade.analysis.measure_id.lattice import (
-    MeasureId,
     MeasureIdBool,
 )
 
@@ -127,8 +127,7 @@ class IfToStim(IfElseSimplification, RewriteRule):
     Rewrite if statements to stim equivalent statements.
     """
 
-    measure_analysis: dict[ir.SSAValue, MeasureId]
-    measure_count: int
+    measure_frame: MeasureIDFrame
 
     def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
 
@@ -140,7 +139,7 @@ class IfToStim(IfElseSimplification, RewriteRule):
 
     def rewrite_IfElse(self, stmt: scf.IfElse) -> RewriteResult:
 
-        if not isinstance(self.measure_analysis[stmt.cond], MeasureIdBool):
+        if not isinstance(self.measure_frame.entries[stmt.cond], MeasureIdBool):
             return RewriteResult()
 
         # check that there is only qubit.Apply in the then-body,
@@ -161,12 +160,12 @@ class IfToStim(IfElseSimplification, RewriteRule):
             return RewriteResult()
 
         # get necessary measurement ID type from analysis
-        measure_id_bool = self.measure_analysis[stmt.cond]
+        measure_id_bool = self.measure_frame.entries[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
+            (measure_id_bool.idx - 1) - self.measure_frame.num_measures_at_stmt[stmt]
         )
         get_record_stmt = GetRecord(id=measure_id_idx_stmt.result)  # noqa: F841
 

bloqade/stim/rewrite/squin_measure.py

@@ -2,47 +2,16 @@
 from dataclasses import dataclass
 
 from kirin import ir
-from kirin.dialects import py, ilist
+from kirin.dialects import py
 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, auxiliary
+from bloqade.stim.dialects import collapse
 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
@@ -51,9 +20,6 @@ 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:
@@ -70,10 +36,6 @@ class SquinMeasureToStim(RewriteRule):
         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,
@@ -84,27 +46,6 @@ class SquinMeasureToStim(RewriteRule):
 
         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_circuit-0.6.3.dist-info → bloqade_circuit-0.6.5.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: bloqade-circuit
-Version: 0.6.3
+Version: 0.6.5
 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

{bloqade_circuit-0.6.3.dist-info → bloqade_circuit-0.6.5.dist-info}/RECORD

@@ -9,9 +9,9 @@ bloqade/analysis/address/impls.py,sha256=cWdq1CNV6HVDECtIgl9ogUa8nvFZm_Sy5Wcfk1f
 bloqade/analysis/address/lattice.py,sha256=dUq999feqPoBYkqEXe1hjHOn4TP_bkvKip8fyWQ-2-8,1755
 bloqade/analysis/fidelity/__init__.py,sha256=iJkhoHvCMU9bKxQqgxIWKQWvpqNFRgNBI5DK8-4RAB8,59
 bloqade/analysis/fidelity/analysis.py,sha256=G6JEYc8eeWJ9mwsbUAIzXuU2nrnTU4te41c04xE71gM,3218
-bloqade/analysis/measure_id/__init__.py,sha256=J9I58iIyt4IjB36xK6_q9PV1B-40VJ1Gu9hQ432A6gM,98
-bloqade/analysis/measure_id/analysis.py,sha256=93S3a_Wu7Bt1j217l1hZMY2tLlR00aUYN27o0IhZsqU,1748
-bloqade/analysis/measure_id/impls.py,sha256=ItS9bLMs8UdAkElYgrqD48Xlh9ZyoJq5xDPWcXjWH1c,4823
+bloqade/analysis/measure_id/__init__.py,sha256=r_R_br1e3H7ZzwkeQw4TnDAP4M_bUaRlRb7ZRdowvNI,145
+bloqade/analysis/measure_id/analysis.py,sha256=F22mLWeOLH_QHm25_4DZLLse4ljsiZr7UNThaIy4pzA,2149
+bloqade/analysis/measure_id/impls.py,sha256=ojFSQrrqj-jdcBklaLa3HHojHvw59Z7Dl2iNdfSuP1w,6376
 bloqade/analysis/measure_id/lattice.py,sha256=WPrn0R79umCH909BFWsUJ64qx9n_3KYimIW5UaXNuGU,1891
 bloqade/cirq_utils/__init__.py,sha256=1DRDCF3PpgJCOr0z7iULdrn3dqm7GLpRGs9AlqE7XA8,280
 bloqade/cirq_utils/lineprog.py,sha256=JosrhfeOHI9FycUT_sYFj8TBzLpo97TL8zK-Ap2U4eQ,11021
@@ -22,7 +22,7 @@ bloqade/cirq_utils/noise/conflict_graph.py,sha256=ZUwPWTknrb6SgtZUVPeICn3YA-nUeW
 bloqade/cirq_utils/noise/model.py,sha256=06Y_BLChOA-PhhAJcWLSgLVAAJoNjOrAujL1YCwcXA0,20590
 bloqade/cirq_utils/noise/transform.py,sha256=tvDt4WMLM8dKPME51y0_peSZk2-jKmjq0urOxm0lWuQ,2309
 bloqade/pyqrack/__init__.py,sha256=lonTS-luJkTVujCCtgdZRC12V7FQdoFcozAI-byXwN0,810
-bloqade/pyqrack/base.py,sha256=9z61PaaAFqCBBwkgsDZSr-qr9IQ5OJ_JUvltmJ7Bgls,4407
+bloqade/pyqrack/base.py,sha256=g0GRlEgyJ_P8z-lR8RK2CAuRTj6KPfglKX0iwrgg4DM,4408
 bloqade/pyqrack/device.py,sha256=40vduanEgA26GAW3buHoRpyqPA0xUt2tONY3w5JeH5s,7524
 bloqade/pyqrack/reg.py,sha256=uTL07CT1R0xUsInLmwU9YuuNdV6lV0lCs1zhdUz1qIs,1660
 bloqade/pyqrack/target.py,sha256=c78VtLWAiDNp_0sXwvVzhaEoeFsr1fUVsupxWuo6p3s,3661
@@ -125,10 +125,12 @@ bloqade/rewrite/rules/__init__.py,sha256=3e1Z5T3INqNtP6OU0Vivu_SdMOR_2KDixeA0Yjo
 bloqade/rewrite/rules/flatten_ilist.py,sha256=QoIxMaBXSlatpWzi5s_MAPnV3bV3GeoWc31RBw0WQ3s,1465
 bloqade/rewrite/rules/inline_getitem_ilist.py,sha256=uIXQRCsr3_GPMciDT4ghI-ezhQmkDcGcC6pguABPUVw,875
 bloqade/rewrite/rules/split_ifs.py,sha256=Nm4lpEUHZcnCeewIld0tt7UuGO69LiBGl7Uybuwissw,2119
-bloqade/squin/__init__.py,sha256=MH7i5gR9DhTjLMI6vsP_NT7_yoaEowYiQwsYhrrUEX0,454
+bloqade/squin/__init__.py,sha256=b7ZD69ql9GriIPxN6JhWxANJVzuIJh_r-gC24wsW1mM,621
 bloqade/squin/_typeinfer.py,sha256=bilWfC6whTMwewFCqDgB6vDHZsgXPr3azNOYqqnvtB4,780
+bloqade/squin/gate.py,sha256=tCnjfrSVsXHX7VxkEulZ2SQS5ydtmON8QlcGibM6c2I,4028
 bloqade/squin/groups.py,sha256=RXGJnNZUSXF_f5ljjhZ9At8UhaijayoxFoWvxEsUOWc,1310
 bloqade/squin/lowering.py,sha256=SR6q-IfV8WHPKT97M7UFu5KoRgAojfDno8Bft1mUSKM,1736
+bloqade/squin/parallel.py,sha256=X6Ps9kQIgnFMlZO14y2ntdxvivqbIP28PAWF8KmxByM,5172
 bloqade/squin/qubit.py,sha256=LgNJsm6qCyP7_O-lZg3YT8IiqzF5W5ff1VwQ79nXN4c,5148
 bloqade/squin/types.py,sha256=T3lkqid4HEWuAK_wRns_p-K5DbLDwlldoyZtVay7A3o,119
 bloqade/squin/wire.py,sha256=GZhF0EHCu7OU70zTV_N83yann-eQnYG_lM2u0QYFoAs,6596
@@ -203,14 +205,14 @@ bloqade/stim/emit/__init__.py,sha256=N2dPQY7OyqPwHAStDeOgYg2yfxqxMOz-N7pD5Z4JwlI
 bloqade/stim/emit/stim_str.py,sha256=JyEBoIhLQASogZcUWHI9tMD4JoXYrEqUr2qaZ30gZdc,1491
 bloqade/stim/parse/__init__.py,sha256=l2DjReB2KkgrDjP_4nP6RnoziiOewoSeZfTno1sVYTw,59
 bloqade/stim/parse/lowering.py,sha256=L-IcR_exlxsTVv4SQ0bhzIF4_L82P-GEdK6qRd6B86Y,23723
-bloqade/stim/passes/__init__.py,sha256=SO4OJLaRxq9Lt2AaUxNqiAz-eJnfWsJ6TijUAX64DM4,62
-bloqade/stim/passes/simplify_ifs.py,sha256=45lWKmc6ybt0KjEuE297IouPVJ7NwX_j9q4GNUsIkEc,690
-bloqade/stim/passes/squin_to_stim.py,sha256=pZvSd1rmOnLRfZONctQAb6ORjcoMd1Qjpbhyf-7m-MQ,5125
-bloqade/stim/rewrite/__init__.py,sha256=SHWryh7rZHXOlIz8BMNpj-w7-8VQCRMLt6PfzYFbBfw,434
-bloqade/stim/rewrite/ifs_to_stim.py,sha256=LXkmmTIgi8-Exrz1EGjG8QpdRn6EV3HZM-bZeAaf-cQ,6718
+bloqade/stim/passes/__init__.py,sha256=aysjOZyn0IrJQCQBEqiz8pwZ5u5t2s9TmEzA9Y9KG9w,167
+bloqade/stim/passes/simplify_ifs.py,sha256=zicqggWu_yzfrf2a7uUCt-ZenbYSEnFsyGxDfKw72qQ,1084
+bloqade/stim/passes/squin_to_stim.py,sha256=6TpnpvA6JA3dQyH6mxNpGn9-__sPlSmyUBIdThn9xVg,6018
+bloqade/stim/rewrite/__init__.py,sha256=zL5G73JEsXkehN7gCtUgGnmC2BJ3vKihOd1ohVwM68E,480
+bloqade/stim/rewrite/ifs_to_stim.py,sha256=jmCb6AwNqGXWuLR8BkAYIVvAlOptQzVPHWv52_v-Vsw,6755
 bloqade/stim/rewrite/py_constant_to_stim.py,sha256=PV8bHvn759-d_0JW4akaGSORW_oxigrlUBhAC51PJAU,1354
 bloqade/stim/rewrite/qubit_to_stim.py,sha256=oiKmi8BlBwXJq-8kGhN1nXgyxJ2UIt_9uouNkU1J8vs,2624
-bloqade/stim/rewrite/squin_measure.py,sha256=zPH2q_ciV2D615GK9l9LWGYmjv3dOju18jKMYERIL7c,4817
+bloqade/stim/rewrite/squin_measure.py,sha256=1zuILosGACN7rPYA87MYVwv0M4pPTala1YTe9owbhkw,2519
 bloqade/stim/rewrite/squin_noise.py,sha256=NafmAiByT4Y5895fZM4Od8arKjsJuW6F5wvRpAFFo70,6240
 bloqade/stim/rewrite/util.py,sha256=xnLDiEj45CBoG3mpG-ywE1Jjh1k_OVP4iI1A75VR6sw,7257
 bloqade/stim/rewrite/wire_identity_elimination.py,sha256=Cscu8yaSslPuW04HvbXx4HJ3JzdUZNUMyFqcvuc4sxY,795
@@ -228,7 +230,7 @@ bloqade/visual/animation/runtime/atoms.py,sha256=EmjxhujLiHHPS_HtH_B-7TiqeHgvW5u
 bloqade/visual/animation/runtime/ppoly.py,sha256=JB9IP53N1w6adBJEue6J5Nmj818Id9JvrlgrmiQTU1I,1385
 bloqade/visual/animation/runtime/qpustate.py,sha256=rlmxQeJSvaohXrTpXQL5y-NJcpvfW33xPaYM1slv7cc,4270
 bloqade/visual/animation/runtime/utils.py,sha256=ju9IzOWX-vKwfpqUjlUKu3Ssr_UFPFFq-tzH_Nqyo_c,1212
-bloqade_circuit-0.6.3.dist-info/METADATA,sha256=K2Me5w0S2bHxND1_JeYHn5CBC7HjLOdzfpacoATbqWo,3849
-bloqade_circuit-0.6.3.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-bloqade_circuit-0.6.3.dist-info/licenses/LICENSE,sha256=S5GIJwR6QCixPA9wryYb44ZEek0Nz4rt_zLUqP05UbU,13160
-bloqade_circuit-0.6.3.dist-info/RECORD,,
+bloqade_circuit-0.6.5.dist-info/METADATA,sha256=8Zb_1qZrO5V9v3JR32KxFQk9malYU1JPPJpkf1U9RPc,3849
+bloqade_circuit-0.6.5.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+bloqade_circuit-0.6.5.dist-info/licenses/LICENSE,sha256=S5GIJwR6QCixPA9wryYb44ZEek0Nz4rt_zLUqP05UbU,13160
+bloqade_circuit-0.6.5.dist-info/RECORD,,