bloqade-circuit 0.6.2__tar.gz → 0.6.4__tar.gz

{bloqade_circuit-0.6.2 → bloqade_circuit-0.6.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: bloqade-circuit
-Version: 0.6.2
+Version: 0.6.4
 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.2 → bloqade_circuit-0.6.4}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "bloqade-circuit"
-version = "0.6.2"
+version = "0.6.4"
 description = "The software development toolkit for neutral atom arrays."
 readme = "README.md"
 authors = [

bloqade_circuit-0.6.4/src/bloqade/analysis/measure_id/__init__.py

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

{bloqade_circuit-0.6.2 → bloqade_circuit-0.6.4}/src/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_circuit-0.6.2 → bloqade_circuit-0.6.4}/src/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_circuit-0.6.2 → bloqade_circuit-0.6.4}/src/bloqade/pyqrack/device.py

@@ -3,7 +3,9 @@ from dataclasses import field, dataclass
 
 import numpy as np
 from kirin import ir
+from kirin.passes import fold
 
+from bloqade.squin import noise as squin_noise
 from pyqrack.pauli import Pauli
 from bloqade.device import AbstractSimulatorDevice
 from bloqade.pyqrack.reg import Measurement, PyQrackQubit
@@ -16,6 +18,7 @@ from bloqade.pyqrack.base import (
     _default_pyqrack_args,
 )
 from bloqade.pyqrack.task import PyQrackSimulatorTask
+from bloqade.squin.noise.rewrite import RewriteNoiseStmts
 from bloqade.analysis.address.lattice import AnyAddress
 from bloqade.analysis.address.analysis import AddressAnalysis
 
@@ -47,14 +50,23 @@ class PyQrackSimulatorBase(AbstractSimulatorDevice[PyQrackSimulatorTask]):
         kwargs: dict[str, Any],
         memory: MemoryType,
     ) -> PyQrackSimulatorTask[Params, RetType, MemoryType]:
+
+        if squin_noise in mt.dialects:
+            # NOTE: rewrite noise statements
+            mt_ = mt.similar(mt.dialects)
+            RewriteNoiseStmts(mt_.dialects)(mt_)
+            fold.Fold(mt_.dialects)(mt_)
+        else:
+            mt_ = mt
+
         interp = PyQrackInterpreter(
-            mt.dialects,
+            mt_.dialects,
             memory=memory,
             rng_state=self.rng_state,
             loss_m_result=self.loss_m_result,
         )
         return PyQrackSimulatorTask(
-            kernel=mt, args=args, kwargs=kwargs, pyqrack_interp=interp
+            kernel=mt_, args=args, kwargs=kwargs, pyqrack_interp=interp
         )
 
     def state_vector(

{bloqade_circuit-0.6.2 → bloqade_circuit-0.6.4}/src/bloqade/pyqrack/squin/op.py

@@ -1,3 +1,5 @@
+import math
+
 from kirin import interp
 
 from bloqade.squin import op
@@ -122,6 +124,18 @@ class PyQrackMethods(interp.MethodTable):
     ) -> tuple[OperatorRuntimeABC]:
         return (OperatorRuntime(method_name=stmt.name.lower()),)
 
+    @interp.impl(op.stmts.SqrtX)
+    @interp.impl(op.stmts.SqrtY)
+    def sqrt(
+        self,
+        interp: PyQrackInterpreter,
+        frame: interp.Frame,
+        stmt: op.stmts.SqrtX | op.stmts.SqrtY,
+    ):
+        axis_name = "x" if isinstance(stmt, op.stmts.SqrtX) else "y"
+        axis = OperatorRuntime(method_name=axis_name)
+        return (RotRuntime(axis=axis, angle=-0.5 * math.pi),)
+
     @interp.impl(op.stmts.P0)
     @interp.impl(op.stmts.P1)
     def projector(

{bloqade_circuit-0.6.2 → bloqade_circuit-0.6.4}/src/bloqade/squin/cirq/__init__.py

@@ -157,6 +157,7 @@ def load_circuit(
 def emit_circuit(
     mt: ir.Method,
     qubits: Sequence[cirq.Qid] | None = None,
+    ignore_returns: bool = False,
 ) -> cirq.Circuit:
     """Converts a squin.kernel method to a cirq.Circuit object.
 
@@ -170,6 +171,9 @@ def emit_circuit(
             statement in the order they appear inside the kernel.
             **Note**: If a list of qubits is provided, make sure that there is a sufficient
             number of qubits for the resulting circuit.
+        ignore_returns (bool):
+            If `False`, emitting a circuit from a kernel that returns a value will error.
+            Set it to `True` in order to ignore the return value(s). Defaults to `False`.
 
     ## Examples:
 
@@ -228,11 +232,14 @@ def emit_circuit(
     and manipulate the qubits in other circuits directly written in cirq as well.
     """
 
-    if isinstance(mt.code, func.Function) and not mt.code.signature.output.is_subseteq(
-        types.NoneType
+    if (
+        not ignore_returns
+        and isinstance(mt.code, func.Function)
+        and not mt.code.signature.output.is_subseteq(types.NoneType)
     ):
         raise EmitError(
             "The method you are trying to convert to a circuit has a return value, but returning from a circuit is not supported."
+            " Set `ignore_returns = True` in order to simply ignore the return values and emit a circuit."
         )
 
     emitter = EmitCirq(qubits=qubits)
@@ -244,7 +251,12 @@ def emit_circuit(
     return emitter.run(mt_, args=())
 
 
-def dump_circuit(mt: ir.Method, qubits: Sequence[cirq.Qid] | None = None, **kwargs):
+def dump_circuit(
+    mt: ir.Method,
+    qubits: Sequence[cirq.Qid] | None = None,
+    ignore_returns: bool = False,
+    **kwargs,
+):
     """Converts a squin.kernel method to a cirq.Circuit object and dumps it as JSON.
 
     This just runs `emit_circuit` and calls the `cirq.to_json` function to emit a JSON.
@@ -259,7 +271,10 @@ def dump_circuit(mt: ir.Method, qubits: Sequence[cirq.Qid] | None = None, **kwar
             statement in the order they appear inside the kernel.
             **Note**: If a list of qubits is provided, make sure that there is a sufficient
             number of qubits for the resulting circuit.
+        ignore_returns (bool):
+            If `False`, emitting a circuit from a kernel that returns a value will error.
+            Set it to `True` in order to ignore the return value(s). Defaults to `False`.
 
     """
-    circuit = emit_circuit(mt, qubits=qubits)
+    circuit = emit_circuit(mt, qubits=qubits, ignore_returns=ignore_returns)
     return cirq.to_json(circuit, **kwargs)

bloqade_circuit-0.6.4/src/bloqade/stim/passes/__init__.py

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

{bloqade_circuit-0.6.2 → bloqade_circuit-0.6.4}/src/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_circuit-0.6.2 → bloqade_circuit-0.6.4}/src/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_circuit-0.6.2 → bloqade_circuit-0.6.4}/src/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_circuit-0.6.2 → bloqade_circuit-0.6.4}/src/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_circuit-0.6.2 → bloqade_circuit-0.6.4}/src/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)