bloqade-circuit 0.7.8__tar.gz → 0.7.9__tar.gz

{bloqade_circuit-0.7.8 → bloqade_circuit-0.7.9}/PKG-INFO

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

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

{bloqade_circuit-0.7.8 → bloqade_circuit-0.7.9}/src/bloqade/pyqrack/reg.py

@@ -2,8 +2,8 @@ import enum
 from typing import TYPE_CHECKING
 from dataclasses import dataclass
 
-from bloqade.types import MeasurementResult
 from bloqade.qasm2.types import Qubit
+from bloqade.squin.types import MeasurementResult
 
 if TYPE_CHECKING:
     from pyqrack import QrackSimulator

bloqade_circuit-0.7.9/src/bloqade/qasm2/passes/fold.py

@@ -0,0 +1,41 @@
+from dataclasses import field, dataclass
+
+from kirin import ir
+from kirin.passes import Pass
+from kirin.ir.method import Method
+from kirin.rewrite.abc import RewriteResult
+
+from bloqade.qasm2.dialects import expr
+from bloqade.rewrite.passes import AggressiveUnroll
+
+from .unroll_if import UnrollIfs
+
+
+@dataclass
+class QASM2Fold(Pass):
+    """Fold pass for qasm2.extended"""
+
+    inline_gate_subroutine: bool = True
+    unroll_ifs: bool = True
+    aggressive_unroll: AggressiveUnroll = field(init=False)
+
+    def __post_init__(self):
+        def inline_simple(node: ir.Statement):
+            if isinstance(node, expr.GateFunction):
+                return self.inline_gate_subroutine
+
+            return True
+
+        self.aggressive_unroll = AggressiveUnroll(
+            self.dialects, inline_simple, no_raise=self.no_raise
+        )
+
+    def unsafe_run(self, mt: Method) -> RewriteResult:
+        result = RewriteResult()
+
+        if self.unroll_ifs:
+            result = UnrollIfs(mt.dialects).unsafe_run(mt).join(result)
+
+        result = self.aggressive_unroll.unsafe_run(mt).join(result)
+
+        return result

{bloqade_circuit-0.7.8 → bloqade_circuit-0.7.9}/src/bloqade/rewrite/passes/__init__.py

@@ -1 +1,2 @@
+from .aggressive_unroll import AggressiveUnroll as AggressiveUnroll
 from .canonicalize_ilist import CanonicalizeIList as CanonicalizeIList

bloqade_circuit-0.7.9/src/bloqade/rewrite/passes/aggressive_unroll.py

@@ -0,0 +1,93 @@
+from typing import Callable
+from dataclasses import field, dataclass
+
+from kirin import ir
+from kirin.passes import Pass, HintConst, TypeInfer
+from kirin.rewrite import (
+    Walk,
+    Chain,
+    Inline,
+    Fixpoint,
+    Call2Invoke,
+    ConstantFold,
+    CFGCompactify,
+    InlineGetItem,
+    InlineGetField,
+    DeadCodeElimination,
+)
+from kirin.dialects import scf, ilist
+from kirin.ir.method import Method
+from kirin.rewrite.abc import RewriteResult
+from kirin.rewrite.cse import CommonSubexpressionElimination
+from kirin.passes.aggressive import UnrollScf
+
+
+@dataclass
+class Fold(Pass):
+    hint_const: HintConst = field(init=False)
+
+    def __post_init__(self):
+        self.hint_const = HintConst(self.dialects, no_raise=self.no_raise)
+
+    def unsafe_run(self, mt: Method) -> RewriteResult:
+        result = RewriteResult()
+        result = self.hint_const.unsafe_run(mt).join(result)
+        rule = Chain(
+            ConstantFold(),
+            Call2Invoke(),
+            InlineGetField(),
+            InlineGetItem(),
+            ilist.rewrite.InlineGetItem(),
+            ilist.rewrite.HintLen(),
+        )
+        result = Fixpoint(Walk(rule)).rewrite(mt.code).join(result)
+
+        return result
+
+
+@dataclass
+class AggressiveUnroll(Pass):
+    """A pass to unroll structured control flow"""
+
+    additional_inline_heuristic: Callable[[ir.Statement], bool] = lambda node: True
+
+    fold: Fold = field(init=False)
+    typeinfer: TypeInfer = field(init=False)
+    scf_unroll: UnrollScf = field(init=False)
+
+    def __post_init__(self):
+        self.fold = Fold(self.dialects, no_raise=self.no_raise)
+        self.typeinfer = TypeInfer(self.dialects, no_raise=self.no_raise)
+        self.scf_unroll = UnrollScf(self.dialects, no_raise=self.no_raise)
+
+    def unsafe_run(self, mt: Method) -> RewriteResult:
+        result = RewriteResult()
+        result = self.scf_unroll.unsafe_run(mt).join(result)
+        result = (
+            Walk(Chain(ilist.rewrite.ConstList2IList(), ilist.rewrite.Unroll()))
+            .rewrite(mt.code)
+            .join(result)
+        )
+        result = self.typeinfer.unsafe_run(mt).join(result)
+        result = self.fold.unsafe_run(mt).join(result)
+        result = Walk(Inline(self.inline_heuristic)).rewrite(mt.code).join(result)
+        result = Walk(Fixpoint(CFGCompactify())).rewrite(mt.code).join(result)
+
+        rule = Chain(
+            CommonSubexpressionElimination(),
+            DeadCodeElimination(),
+        )
+        result = Fixpoint(Walk(rule)).rewrite(mt.code).join(result)
+
+        return result
+
+    def inline_heuristic(self, node: ir.Statement) -> bool:
+        """The heuristic to decide whether to inline a function call or not.
+        inside loops and if-else, only inline simple functions, i.e.
+        functions with a single block
+        """
+        return not isinstance(
+            node.parent_stmt, (scf.For, scf.IfElse)
+        ) and self.additional_inline_heuristic(
+            node
+        )  # always inline calls outside of loops and if-else

bloqade_circuit-0.7.9/src/bloqade/rewrite/passes/canonicalize_ilist.py

@@ -0,0 +1,32 @@
+from dataclasses import field, dataclass
+
+from kirin import ir, passes
+from kirin.rewrite import (
+    Walk,
+    Chain,
+    Fixpoint,
+)
+from kirin.dialects.ilist import rewrite
+
+
+@dataclass
+class CanonicalizeIList(passes.Pass):
+
+    fold_pass: passes.Fold = field(init=False)
+
+    def __post_init__(self):
+        self.fold_pass = passes.Fold(self.dialects, no_raise=self.no_raise)
+
+    def unsafe_run(self, mt: ir.Method):
+        result = Fixpoint(
+            Walk(
+                Chain(
+                    rewrite.InlineGetItem(),
+                    rewrite.FlattenAdd(),
+                    rewrite.HintLen(),
+                )
+            )
+        ).rewrite(mt.code)
+
+        result = self.fold_pass(mt).join(result)
+        return result

{bloqade_circuit-0.7.8 → bloqade_circuit-0.7.9}/test/analysis/measure_id/test_measure_id.py

@@ -4,7 +4,6 @@ from kirin.dialects import scf
 from bloqade.squin import op, qubit, kernel
 from bloqade.analysis.measure_id import MeasurementIDAnalysis
 from bloqade.analysis.measure_id.lattice import (
-    NotMeasureId,
     MeasureIdBool,
     MeasureIdTuple,
     InvalidMeasureId,
@@ -111,10 +110,15 @@ def test_scf_cond_true():
     HintConst(dialects=test.dialects).unsafe_run(test)
     frame, _ = MeasurementIDAnalysis(test.dialects).run_analysis(test)
 
-    assert [frame.entries[result] for result in results_at(test, 0, 7)] == [
-        NotMeasureId(),
-        MeasureIdTuple((MeasureIdBool(idx=1),)),
+    # MeasureIdTuple(data=MeasureIdBool(idx=1),) should occur twice:
+    # First from the measurement in the true branch, then
+    # the result of the scf.IfElse itself
+    analysis_results = [
+        val
+        for val in frame.entries.values()
+        if val == MeasureIdTuple(data=(MeasureIdBool(idx=1),))
     ]
+    assert len(analysis_results) == 2
 
 
 def test_scf_cond_false():
@@ -136,10 +140,13 @@ def test_scf_cond_false():
     HintConst(dialects=test.dialects).unsafe_run(test)
     frame, _ = MeasurementIDAnalysis(test.dialects).run_analysis(test)
 
-    assert [frame.entries[result] for result in results_at(test, 0, 7)] == [
-        NotMeasureId(),
-        MeasureIdBool(idx=1),
+    # MeasureIdBool(idx=1) should occur twice:
+    # First from the measurement in the false branch, then
+    # the result of the scf.IfElse itself
+    analysis_results = [
+        val for val in frame.entries.values() if val == MeasureIdBool(idx=1)
     ]
+    assert len(analysis_results) == 2
 
 
 def test_slice():

{bloqade_circuit-0.7.8 → bloqade_circuit-0.7.9}/test/qasm2/emit/test_qasm2_emit.py

@@ -288,3 +288,30 @@ def test_if():
 
     target = qasm2.emit.QASM2()
     target.emit(nested_kernel)
+
+
+def test_loop_unroll():
+    n_qubits = 4
+
+    @qasm2.extended
+    def ghz_linear():
+        q = qasm2.qreg(n_qubits)
+        qasm2.h(q[0])
+        for i in range(1, n_qubits):
+            qasm2.cx(q[i - 1], q[i])
+
+    target = qasm2.emit.QASM2(
+        allow_parallel=True,
+    )
+    qasm2_str = target.emit_str(ghz_linear)
+
+    assert qasm2_str == (
+        """KIRIN {func,lowering.call,lowering.func,py.ilist,qasm2.core,qasm2.expr,qasm2.indexing,qasm2.noise,qasm2.parallel,qasm2.uop,scf};
+include "qelib1.inc";
+qreg q[4];
+h q[0];
+CX q[0], q[1];
+CX q[1], q[2];
+CX q[2], q[3];
+"""
+    )

{bloqade_circuit-0.7.8 → bloqade_circuit-0.7.9}/uv.lock

@@ -103,7 +103,7 @@ wheels = [
 
 [[package]]
 name = "bloqade-circuit"
-version = "0.7.8"
+version = "0.7.9"
 source = { editable = "." }
 dependencies = [
     { name = "kirin-toolchain" },

bloqade_circuit-0.7.8/src/bloqade/qasm2/passes/fold.py

@@ -1,100 +0,0 @@
-from dataclasses import field, dataclass
-
-from kirin import ir
-from kirin.passes import Pass, TypeInfer
-from kirin.rewrite import (
-    Walk,
-    Chain,
-    Inline,
-    Fixpoint,
-    WrapConst,
-    Call2Invoke,
-    ConstantFold,
-    CFGCompactify,
-    InlineGetItem,
-    InlineGetField,
-    DeadCodeElimination,
-    CommonSubexpressionElimination,
-)
-from kirin.analysis import const
-from kirin.dialects import scf, ilist
-from kirin.ir.method import Method
-from kirin.rewrite.abc import RewriteResult
-
-from bloqade.qasm2.dialects import expr
-
-from .unroll_if import UnrollIfs
-
-
-@dataclass
-class QASM2Fold(Pass):
-    """Fold pass for qasm2.extended"""
-
-    constprop: const.Propagate = field(init=False)
-    inline_gate_subroutine: bool = True
-    unroll_ifs: bool = True
-
-    def __post_init__(self):
-        self.constprop = const.Propagate(self.dialects)
-        self.typeinfer = TypeInfer(self.dialects)
-
-    def unsafe_run(self, mt: Method) -> RewriteResult:
-        result = RewriteResult()
-        frame, _ = self.constprop.run_analysis(mt)
-        result = Walk(WrapConst(frame)).rewrite(mt.code).join(result)
-        rule = Chain(
-            ConstantFold(),
-            Call2Invoke(),
-            InlineGetField(),
-            InlineGetItem(),
-            DeadCodeElimination(),
-            CommonSubexpressionElimination(),
-        )
-        result = Fixpoint(Walk(rule)).rewrite(mt.code).join(result)
-
-        result = (
-            Walk(
-                Chain(
-                    scf.unroll.PickIfElse(),
-                    scf.unroll.ForLoop(),
-                    scf.trim.UnusedYield(),
-                )
-            )
-            .rewrite(mt.code)
-            .join(result)
-        )
-
-        if self.unroll_ifs:
-            UnrollIfs(mt.dialects).unsafe_run(mt).join(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)
-        result = (
-            Walk(Chain(ilist.rewrite.ConstList2IList(), ilist.rewrite.Unroll()))
-            .rewrite(mt.code)
-            .join(result)
-        )
-
-        def inline_simple(node: ir.Statement):
-            if isinstance(node, expr.GateFunction):
-                return self.inline_gate_subroutine
-
-            if not isinstance(node.parent_stmt, (scf.For, scf.IfElse)):
-                return True  # always inline calls outside of loops and if-else
-
-            # inside loops and if-else, only inline simple functions, i.e. functions with a single block
-            if (trait := node.get_trait(ir.CallableStmtInterface)) is None:
-                return False  # not a callable, don't inline to be safe
-            region = trait.get_callable_region(node)
-            return len(region.blocks) == 1
-
-        result = (
-            Walk(
-                Inline(inline_simple),
-            )
-            .rewrite(mt.code)
-            .join(result)
-        )
-        result = Walk(Fixpoint(CFGCompactify())).rewrite(mt.code).join(result)
-        return result

bloqade_circuit-0.7.8/src/bloqade/rewrite/passes/canonicalize_ilist.py

@@ -1,28 +0,0 @@
-from dataclasses import dataclass
-
-from kirin import ir
-from kirin.passes import Pass
-from kirin.rewrite import (
-    Walk,
-    Chain,
-    Fixpoint,
-)
-from kirin.analysis import const
-
-from ..rules.flatten_ilist import FlattenAddOpIList
-from ..rules.inline_getitem_ilist import InlineGetItemFromIList
-
-
-@dataclass
-class CanonicalizeIList(Pass):
-
-    def unsafe_run(self, mt: ir.Method):
-
-        cp_result_frame, _ = const.Propagate(dialects=mt.dialects).run_analysis(mt)
-
-        return Fixpoint(
-            Chain(
-                Walk(InlineGetItemFromIList(constprop_result=cp_result_frame.entries)),
-                Walk(FlattenAddOpIList()),
-            )
-        ).rewrite(mt.code)

bloqade_circuit-0.7.8/src/bloqade/rewrite/rules/flatten_ilist.py

@@ -1,51 +0,0 @@
-from dataclasses import dataclass
-
-from kirin import ir
-from kirin.dialects import py, ilist
-from kirin.rewrite.abc import RewriteRule, RewriteResult
-
-
-@dataclass
-class FlattenAddOpIList(RewriteRule):
-
-    def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
-        if not isinstance(node, py.binop.Add):
-            return RewriteResult()
-
-        # check if we are adding two ilist.New objects
-        new_data = ()
-
-        # lhs:
-        if not isinstance(node.lhs.owner, ilist.New):
-            if not (
-                isinstance(node.lhs.owner, py.Constant)
-                and isinstance(
-                    const_ilist := node.lhs.owner.value.unwrap(), ilist.IList
-                )
-                and len(const_ilist.data) == 0
-            ):
-                return RewriteResult()
-
-        else:
-            new_data += node.lhs.owner.values
-
-        # rhs:
-        if not isinstance(node.rhs.owner, ilist.New):
-            if not (
-                isinstance(node.rhs.owner, py.Constant)
-                and isinstance(
-                    const_ilist := node.rhs.owner.value.unwrap(), ilist.IList
-                )
-                and len(const_ilist.data) == 0
-            ):
-                return RewriteResult()
-
-        else:
-            new_data += node.rhs.owner.values
-
-        new_stmt = ilist.New(values=new_data)
-        node.replace_by(new_stmt)
-
-        return RewriteResult(
-            has_done_something=True,
-        )

bloqade_circuit-0.7.8/src/bloqade/rewrite/rules/inline_getitem_ilist.py

@@ -1,31 +0,0 @@
-from dataclasses import dataclass
-
-from kirin import ir
-from kirin.analysis import const
-from kirin.dialects import py, ilist
-from kirin.rewrite.abc import RewriteRule, RewriteResult
-
-
-@dataclass
-class InlineGetItemFromIList(RewriteRule):
-    constprop_result: dict[ir.SSAValue, const.Result]
-
-    def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
-        if not isinstance(node, py.indexing.GetItem):
-            return RewriteResult()
-
-        if not isinstance(node.obj.owner, ilist.New):
-            return RewriteResult()
-
-        if not isinstance(
-            index_value := self.constprop_result.get(node.index), const.Value
-        ):
-            return RewriteResult()
-
-        elem_ssa = node.obj.owner.values[index_value.data]
-
-        node.result.replace_by(elem_ssa)
-
-        return RewriteResult(
-            has_done_something=True,
-        )